How to Create Custom Environment Wrappers in EvoRL: A Complete Guide

Subclass the Wrapper base class from evorl/envs/wrappers/wrapper.py, override the step() method to transform actions, observations, or rewards, and optionally expose new action_space or obs_space properties to alter the agent interface.

EvoRL provides a modular wrapper hierarchy for JAX-based reinforcement learning environments that allows you to compose transformations layer by layer. The emi-group/evorl repository implements a lightweight abstraction where all custom environment wrappers inherit from a common base class, enabling you to squash actions, scale rewards, or modify observations with minimal boilerplate.

Understanding the Wrapper Base Class

All environment wrappers in EvoRL inherit from the abstract Wrapper class defined in evorl/envs/wrappers/wrapper.py. This base class implements a transparent forwarding mechanism where every call (reset, step, obs_space, action_space) is delegated to the wrapped environment unless explicitly overridden.

The Wrapper constructor stores the inner environment via super().__init__(env), allowing you to chain transformations recursively. Because EnvState is implemented as a chex.dataclass, you can safely modify state fields using the .replace() method without breaking differentiability or functional purity.

Implementing ActionSquashWrapper

The ActionSquashWrapper in evorl/envs/wrappers/action_wrapper.py demonstrates the canonical pattern for rescaling continuous actions. It transforms actions from the standardized [-1, 1] range (expected by neural network policies) to the environment's native [low, high] bounds.

# evorl/envs/wrappers/action_wrapper.py

from evorl.envs.wrappers.wrapper import Wrapper
from evorl.envs import Env, EnvState, Action
from evx.space import Box
import jax.numpy as jnp

class ActionSquashWrapper(Wrapper):
    """Convert continuous action space from [-1, 1] to [low, high]."""

    def __init__(self, env: Env):
        super().__init__(env)

        # only Box actions are supported for now

        action_space = self.env.action_space
        assert isinstance(action_space, Box), "Only support Box action_space"

        # scale and bias to map [-1, 1] → [low, high]

        self.scale = (action_space.high - action_space.low) * 0.5
        self.bias = (action_space.high + action_space.low) * 0.5

    def step(self, state: EnvState, action: Action) -> EnvState:
        # transform the incoming action then forward it

        squashed_action = self.scale * action + self.bias
        return self.env.step(state, squashed_action)

    @property
    def action_space(self) -> Space:
        # the wrapper presents a standardized [-1, 1] space to agents

        return Box(low=-jnp.ones_like(self.scale), high=jnp.ones_like(self.scale))

Key implementation details from the source code:

• Type validation: The assert isinstance(action_space, Box) ensures the wrapper only operates on flat continuous spaces, preventing runtime shape mismatches.
• Pre-computed transformation: self.scale and self.bias are calculated once during initialization to avoid recomputing linear transformations at every step.
• API transparency: The action_space property returns a virtual [-1, 1] bounds, making the wrapper transparent to downstream policy networks while internally mapping to the physical environment's range.

Step-by-Step Guide to Creating Custom Wrappers

Follow this pattern to implement your own transformation logic:

Subclass Wrapper and Initialize State

Import the base class and call the parent constructor to establish the wrapper chain. Store any hyperparameters or stateful buffers as instance attributes.

from evorl.envs.wrappers.wrapper import Wrapper

class MyCustomWrapper(Wrapper):
    def __init__(self, env: Env, reward_bias: float = 0.0):
        super().__init__(env)
        self.reward_bias = reward_bias

Override the step Method

Intercept the action before it reaches the inner environment, or modify the reward/observation in the returned EnvState. Use .replace() to update immutable state objects.

    def step(self, state: EnvState, action: Action) -> EnvState:
        # Forward to inner environment

        next_state = self.env.step(state, action)
        
        # Modify reward

        adjusted_reward = next_state.reward + self.reward_bias
        
        # Update state immutably

        return next_state.replace(reward=adjusted_reward)

Expose Modified Spaces

If your transformation alters the shape or bounds of actions or observations, override the corresponding property to maintain consistency with the agent's network architecture.

    @property
    def action_space(self) -> Space:
        # Return transformed space if bounds changed

        return Box(low=-1.0, high=1.0, shape=self.env.action_space.shape)

Stacking Multiple Wrappers

EvoRL wrappers compose through nested instantiation. Each layer handles a single concern, creating a clean separation between action preprocessing, episode management, and vectorization.

from evorl.envs import create_brax_env
from evorl.envs.wrappers import ActionSquashWrapper, OneEpisodeWrapper, VmapWrapper

env = create_brax_env("ant")
env = ActionSquashWrapper(env)              # Map [-1, 1] to native range

env = OneEpisodeWrapper(env, 1000, 0.99)    # Handle termination and discounting

env = VmapWrapper(env, num_envs=4)          # Vectorize across 4 parallel environments

The resulting env object behaves as a single environment while applying transformations in the order: action squashing → episode tracking → batch vectorization.

Complete Working Example: Reward Bias Wrapper

Below is a production-ready template that adds a constant bias to rewards and logs the original value in the info dictionary.

# my_wrapper.py

from evorl.envs.wrappers.wrapper import Wrapper
from evorl.envs import Env, EnvState, Action
import jax.numpy as jnp

class RewardBiasWrapper(Wrapper):
    """Adds constant bias to rewards and stores original in info."""

    def __init__(self, env: Env, bias: float = 0.0):
        super().__init__(env)
        self.bias = bias

    def step(self, state: EnvState, action: Action) -> EnvState:
        next_state = self.env.step(state, action)
        
        # Apply bias

        new_reward = next_state.reward + self.bias
        
        # Store original reward for logging

        new_info = next_state.info.replace(ori_reward=next_state.reward)
        
        return next_state.replace(reward=new_reward, info=new_info)

Usage with rollout:

from evorl.envs import create_brax_env
from evorl.rollout import rollout
from my_wrapper import RewardBiasWrapper
import jax

env = create_brax_env("walker2d")
env = RewardBiasWrapper(env, bias=0.5)

key = jax.random.PRNGKey(42)
state = env.reset(key)

def policy(state, key):
    action = jax.random.uniform(key, env.action_space.shape, minval=-1, maxval=1)
    return action, key

traj, _ = rollout(env.step, policy, state, None, key, rollout_length=500)

# Access original rewards via traj.extras.env_extras.ori_reward

Summary

• Inherit from Wrapper: All custom wrappers must subclass the base class in evorl/envs/wrappers/wrapper.py and call super().__init__(env).
• Override step(): Transform actions before forwarding to the inner environment, or modify rewards/observations in the returned state using .replace().
• Update space properties: Override action_space or obs_space when your transformation changes valid bounds or shapes.
• Leverage state immutability: EnvState is a chex.dataclass; use state.replace(field=new_value) to create modified copies.
• Compose freely: Stack wrappers by nesting constructors—order matters, with the last wrapper applied being the outermost layer.

Frequently Asked Questions

What base class should I use for custom environment wrappers in EvoRL?

All custom wrappers must inherit from Wrapper located at evorl/envs/wrappers/wrapper.py. This base class provides the delegation mechanism that forwards method calls to the wrapped environment unless you explicitly override them.

How does ActionSquashWrapper handle different action space bounds?

ActionSquashWrapper calculates linear transformation parameters during initialization: scale = (high - low) * 0.5 and bias = (high + low) * 0.5. In the step() method, it applies squashed_action = scale * action + bias to map inputs from [-1, 1] to the environment's native [low, high] range, while exposing a standardized [-1, 1] space via the action_space property.

Can I modify observations using the same wrapper pattern?

Yes. Override the reset() and step() methods to transform the observation field in the returned EnvState. If the observation shape or bounds change, override the obs_space property to return the transformed space definition.

How do I access the original unwrapped environment?

The wrapped environment is stored as self.env within any wrapper class. You can access the original base environment by recursively accessing self.env.env through the wrapper stack, though direct access is rarely necessary since the Wrapper base class forwards all standard API calls automatically.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}