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

> Learn to create custom environment wrappers in EvoRL by subclassing Wrapper. Override step and expose new spaces to transform actions, observations, and rewards effectively.

- Repository: [EMI-Group/evorl](https://github.com/emi-group/evorl)
- Tags: how-to-guide
- Published: 2026-03-01

---

**Subclass the `Wrapper` base class from [`evorl/envs/wrappers/wrapper.py`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.

```python

# 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.

```python
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.

```python
    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.

```python
    @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.

```python
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.

```python

# 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:**

```python
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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.