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

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 (lines 26-60) wraps Brax environments and extracts obs, reward, done, and info fields
• GymnaxAdapter in 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 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 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 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:

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 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, 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 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. 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:

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

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

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 or similar).

3. Register in the factory by adding imports to 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, gymnax.py, and mujoco_playground.py translate third-party APIs into the standard Env interface defined in evorl/envs/env.py.
• Wrappers in 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 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. 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, 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. 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.

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