# How to Integrate Different Environment Backends (Brax, Gymnax, and MuJoCo Playground) in EvoRL

> Integrate Brax Gymnax and MuJoCo Playground in EvoRL using adapters training wrappers and a factory function Learn how EvoRL unifies diverse environment backends for streamlined reinforcement learning.

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

---

**EvoRL provides a unified `Env` abstraction that seamlessly integrates Brax, Gymnax, and MuJoCo Playground through a three-layer architecture of adapters, training wrappers, and a centralized factory function.**

Integrating different environment backends in EvoRL allows you to leverage high-performance physics simulators while maintaining a consistent API for evolutionary reinforcement learning algorithms. The `emi-group/evorl` repository implements a backend-agnostic design that abstracts away simulator-specific details through standardized adapters and configurable auto-reset wrappers.

## The Three-Layer Integration Architecture

EvoRL organizes environment integration into three distinct layers, each handling specific responsibilities from low-level API translation to high-level training semantics.

### Adapter Layer (Backend-Specific Wrappers)

The adapter layer converts third-party environment APIs into EvoRL's standardized `Env` interface. Each supported backend implements a dedicated adapter class:

- **`BraxAdapter`** in [`evorl/envs/brax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/brax.py) (lines 26-60) wraps Brax environments and extracts `obs`, `reward`, `done`, and `info` fields
- **`GymnaxAdapter`** in [`evorl/envs/gymnax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/gymnax.py) handles Gymnax environments and converts Gymnax spaces to EvoRL spaces via `gymnax_space_to_evorl_space`
- **`MjxEnvAdapter`** in [`evorl/envs/mujoco_playground.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/mujoco_playground.py) integrates MuJoCo Playground environments following the same pattern

These adapters inherit from `EnvAdapter` (itself a subclass of the abstract `Env`) and remain stateless except for storing the wrapped native environment, with all stochasticity driven by JAX PRNG keys supplied to `reset` and `step` calls.

### Wrapper Layer (Training-Oriented Behaviors)

After adapter instantiation, EvoRL applies a stack of wrappers defined in [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py) that provide RL-specific functionality:

- **`EpisodeWrapper`** – Tracks step counts, handles episode truncation and termination, and optionally computes episode returns
- **Auto-reset strategies** – Controlled by the `AutoresetMode` enum:
  - `NORMAL` – Uses `VmapAutoResetWrapper` for full reset each episode
  - `FAST` – Uses `FastVmapAutoResetWrapper` to reuse the first state and avoid extra random draws
  - `ENVPOOL` – Uses `VmapEnvPoolAutoResetWrapper` adding a separate reset step and `autoreset` flag
  - `DISABLED` – Leaves reset handling to the user
- **Vectorisation** – `VmapWrapper` and variants replicate environments `parallel` times for batched training
- **Utility wrappers** – `ActionSquashWrapper` scales actions to `[-1, 1]` and `ObsFlattenWrapper` flattens image observations

### Factory Layer (Unified Environment Creation)

The `create_env` function in [`evorl/envs/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/__init__.py) serves as the unified entry point, selecting the appropriate adapter and wrapper stack based on configuration parameters including `env_type`, `env_backend`, and `autoreset_mode`.

## Core Abstraction: The Env Interface

All environment integrations implement the abstract base class defined in [`evorl/envs/env.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/env.py):

```python
class Env(ABC):
    @abstractmethod
    def reset(self, key: chex.PRNGKey) -> EnvState: ...
    
    @abstractmethod
    def step(self, state: EnvState, action: Action) -> EnvState: ...
    
    @property
    @abstractmethod
    def action_space(self) -> Space: ...
    
    @property
    @abstractmethod
    def obs_space(self) -> Space: ...

```

The `EnvState` dataclass carries the raw environment state, current observation, reward, done flag, mutable `info` dictionary, and an internal `_internal` dictionary used by wrappers (e.g., storing auto-reset keys).

## Backend-Specific Implementation Details

### Brax Integration

The Brax adapter in [`evorl/envs/brax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/brax.py) creates a Brax `Env` instance and extracts observation, reward, done, and info fields during step transitions. The `create_wrapped_brax_env` function (lines 31-53) constructs the full wrapper stack including episode management and vectorisation.

### Gymnax Integration

Located in [`evorl/envs/gymnax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/gymnax.py), the Gymnax adapter uses `gymnax.make` to instantiate environments and calls `reset` and `step_env` for state transitions. It includes space conversion utilities to map Gymnax observation and action spaces to EvoRL's space definitions.

### MuJoCo Playground Integration

The MuJoCo Playground adapter in [`evorl/envs/mujoco_playground.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/mujoco_playground.py) follows the same stateless adapter pattern, wrapping MuJoCo Playground environments for compatibility with the EvoRL training pipeline.

## Training Wrappers and Auto-Reset Modes

EvoRL provides sophisticated auto-reset capabilities through [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py). The `AutoresetMode` enum determines how environments handle episode termination:

- **NORMAL mode** uses `VmapAutoResetWrapper` to perform full environment resets when episodes end
- **FAST mode** uses `FastVmapAutoResetWrapper` to optimize reset performance by reusing initial states
- **ENVPOOL mode** uses `VmapEnvPoolAutoResetWrapper` to add explicit reset steps and autoreset flags, compatible with EnvPool-style training loops
- **DISABLED mode** leaves reset handling entirely to the user code

The `EpisodeWrapper` tracks step counts, handles truncation based on `episode_length`, and computes episode returns when configured.

## Creating Environments with the Unified Factory

The recommended approach for environment creation uses the unified factory function. Configuration is typically supplied via Hydra/OmegaConf:

```yaml
env:
  env_type: brax          # Options: brax, gymnax, playground, jumanji, jaxmarl, envpool, gymnasium

  env_name: ant
  autoreset_mode: normal  # normal | fast | disabled | envpool

  episode_length: 1000
  parallel: 8

```

```python
from evorl.envs import create_env
import jax, hydra, omegaconf

@hydra.main(version_base=None, config_path=".", config_name="config")
def main(cfg: omegaconf.OmegaConf):
    env = create_env(cfg.env, seed=42)
    state = env.reset(jax.random.PRNGKey(0))
    # Training loop proceeds here

```

The `create_env` function in [`evorl/envs/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/__init__.py) matches the `env_type` parameter and dispatches to the appropriate `create_wrapped_*_env` function, handling all adapter instantiation and wrapper stacking automatically.

## Extending EvoRL with New Backends

To add support for a new simulation backend:

1. **Create an adapter class** inheriting from `EnvAdapter` (which extends `Env`). Implement `reset`, `step`, `action_space`, and `obs_space` methods following the interface in [`evorl/envs/env.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/env.py).

2. **Expose creator functions** implementing `create_mybackend_env` for raw access and `create_wrapped_mybackend_env` that constructs the standard wrapper stack (reuse logic from existing `create_wrapped_*_env` functions in [`evorl/envs/brax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/brax.py) or similar).

3. **Register in the factory** by adding imports to [`evorl/envs/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/__init__.py) and extending the `match` statement in `create_env` to handle your new `env_type`.

Because all wrappers operate on the abstract `Env` interface, agents, rollout workers, and evaluators require no modifications to work with new backends.

## Summary

- **EvoRL uses a three-layer architecture** consisting of backend-specific adapters, training-oriented wrappers, and a unified factory function to integrate Brax, Gymnax, and MuJoCo Playground.
- **Adapters** in [`evorl/envs/brax.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/brax.py), [`gymnax.py`](https://github.com/emi-group/evorl/blob/main/gymnax.py), and [`mujoco_playground.py`](https://github.com/emi-group/evorl/blob/main/mujoco_playground.py) translate third-party APIs into the standard `Env` interface defined in [`evorl/envs/env.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/env.py).
- **Wrappers** in [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py) provide episode management, four auto-reset modes (`NORMAL`, `FAST`, `ENVPOOL`, `DISABLED`), and vectorisation via `VmapWrapper` variants.
- **Unified creation** via `create_env` in [`evorl/envs/__init__.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/__init__.py) dispatches to backend-specific builders based on `env_type` configuration, supporting Hydra/OmegaConf configurations.
- **Extensibility** follows a clear pattern: implement `EnvAdapter`, create wrapped builder functions, and register in the factory match statement.

## Frequently Asked Questions

### What is the difference between Brax and Gymnax backends in EvoRL?

**Brax** environments are physics-based simulators optimized for massive parallelization on accelerators, typically used for continuous control tasks like locomotion. **Gymnax** provides classic control and toy text environments with a focus on simplicity and educational use. In EvoRL, both implement the same `Env` interface, but Brax adapters extract `obs`, `reward`, `done`, and `info` from Brax-specific state structures, while Gymnax adapters use `gymnax.make` and convert Gymnax spaces to EvoRL spaces.

### How does the auto-reset mechanism work in EvoRL environments?

EvoRL provides four auto-reset modes controlled by the `AutoresetMode` enum in [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py). **NORMAL** mode uses `VmapAutoResetWrapper` to perform full environment resets when episodes terminate. **FAST** mode uses `FastVmapAutoResetWrapper` to optimize performance by reusing the initial state instead of generating new random states. **ENVPOOL** mode uses `VmapEnvPoolAutoResetWrapper` to add explicit reset steps and autoreset flags, compatible with EnvPool-style training loops. **DISABLED** mode leaves reset handling to the user code.

### Can I use multiple different backends in the same EvoRL training script?

Yes, EvoRL's unified `Env` abstraction allows mixing backends within the same training workflow. Since all backends (Brax, Gymnax, MuJoCo Playground) implement the identical `Env` interface defined in [`evorl/envs/env.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/env.py), you can instantiate different environments using `create_env` or specific `create_wrapped_*_env` functions, and use them interchangeably in rollout workers or evaluation loops. The wrapper stack ensures consistent behavior for episode management and auto-reset across all backends.

### Where are the environment wrappers defined in the EvoRL codebase?

Environment wrappers are located in `evorl/envs/wrappers/`, with the primary training-oriented wrappers defined in [`evorl/envs/wrappers/training_wrapper.py`](https://github.com/emi-group/evorl/blob/main/evorl/envs/wrappers/training_wrapper.py). This file contains `EpisodeWrapper` for episode management, `VmapAutoResetWrapper` and `FastVmapAutoResetWrapper` for different auto-reset strategies, and `VmapEnvPoolAutoResetWrapper` for EnvPool compatibility. Additional utility wrappers like `ActionSquashWrapper` and `ObsFlattenWrapper` handle action scaling and observation flattening respectively.