How to Implement Evolutionary Algorithms Using the EC Module in EvoRL (OpenES, CMA-ES)

EvoRL provides a plug-and-play EC (Evolutionary Computation) stack that lets you swap any evolutionary algorithm into a reinforcement learning workflow by wrapping it in an EvoOptimizer and attaching it to an ECWorkflowTemplate.

The EvoRL framework decouples evolutionary search from policy architecture, allowing you to evolve neural network controllers using modern ES methods with minimal boilerplate. The EC module is built around three layers: algorithm adapters that bridge external libraries, native optimizers that implement specific update rules, and workflow templates that handle distributed evaluation and metric logging.

Understanding the EvoRL EC Module Architecture

The architecture separates concerns into distinct layers so you can inject new algorithms without rewriting environment loops or evaluation logic.

Algorithm Adapters

The EvoXAlgorithmAdapter in evorl/ec/optimizers/evox_wrapper.py converts any evox.Algorithm into an EvoOptimizer compatible with EvoRL. It handles the ask/tell interface, flips the sign of fitness values (since evoX minimizes while EvoRL maximizes), and manages parameter vector flattening via ParamVectorSpec from evorl/utils/ec_utils.py.

Native Optimizers

For algorithms not requiring evoX, EvoRL provides native implementations. OpenES in evorl/ec/optimizers/openes.py implements the standard OpenAI-style evolution strategy with fitness shaping via compute_centered_ranks. CMAES and SepCMAES in evorl/ec/evox_algorithm/cmaes.py wrap the covariance matrix adaptation logic from evoX.

Workflow Templates

The ECWorkflowTemplate (inherited by OpenESWorkflow and CMAESWorkflow) orchestrates the training loop. It creates the environment, initializes a deterministic agent, builds the optimizer, and handles parallel evaluation across devices. The core logic resides in evorl/workflows/ec_workflow.py, which implements the generic step method for population evaluation and optimizer updates.

Implementing OpenES in EvoRL

OpenES is the simplest entry point for evolutionary RL. The workflow is defined in evorl/algorithms/ec/so/openes.py.

from evorl.algorithms.ec.so.openes import OpenESWorkflow
from evorl.utils.ec_utils import ParamVectorSpec

# The workflow builds the optimizer internally in _build_from_config

workflow = OpenESWorkflow.build_from_config(cfg, enable_jit=True)

# Initialize state

state = workflow.init(jax.random.PRNGKey(0))

# Run evolution

for i in range(100):
    state = workflow.step(state)

The OpenES optimizer in evorl/ec/optimizers/openes.py implements ask() to sample perturbations and tell() to update the center parameter using fitness-shaped rewards. It supports noise table sharing via OpenESNoiseTable for memory efficiency.

Implementing CMA-ES in EvoRL

CMA-ES requires wrapping the evoX algorithm. The implementation in evorl/algorithms/ec/so/cmaes.py demonstrates the pattern.

from evorl.algorithms.ec.so.cmaes import CMAESWorkflow
from evorl.ec.optimizers.evox_wrapper import EvoXAlgorithmAdapter
from evorl.ec.evox_algorithm.cmaes import CMAES
from evorl.utils.ec_utils import ParamVectorSpec

class CMAESWorkflow(ESWorkflowTemplate):
    @classmethod
    def _build_from_config(cls, cfg):
        # Create environment and deterministic agent

        env = create_env(cfg.env)
        agent = make_deterministic_ec_agent(cfg.agent, env)
        agent_state = agent.init(jax.random.PRNGKey(cfg.seed))
        
        # Prepare parameter vector specification

        param_vec = ParamVectorSpec(agent_state.params.policy_params)
        
        # Initialize CMA-ES from evoX

        cma_algo = CMAES(
            center_init=param_vec.to_vector(agent_state.params.policy_params),
            init_stdev=cfg.ec_optimizer.init_stdev,
            pop_size=cfg.ec_optimizer.pop_size,
            mu=cfg.ec_optimizer.mu
        )
        
        # Wrap with adapter to handle ask/tell and sign conversion

        ec_opt = EvoXAlgorithmAdapter(
            algorithm=cma_algo,
            param_vec_spec=param_vec
        )
        
        evaluator = Evaluator(env=env, action_fn=agent.evaluate_actions)
        
        return cls(
            config=cfg,
            env=env,
            agent=agent,
            ec_optimizer=ec_opt,
            ec_evaluator=evaluator,
            evaluator=evaluator,
            agent_state_vmap_axes=AgentState(params=0, obs_preprocessor_state=None),
        )

The EvoXAlgorithmAdapter in evorl/ec/optimizers/evox_wrapper.py handles the conversion between EvoRL's maximization objective and evoX's minimization convention by negating fitness values during tell.

Key Implementation Details

Ask/Tell Interface

All optimizers implement ask(state, key) returning perturbed parameters and tell(state, fitness) returning updated state. In evorl/ec/optimizers/openes.py, ask samples noise from a shared noise table or independent Gaussian, while tell computes weighted updates using fitness-shaped rewards.

Fitness Shaping

OpenES uses centered rank transformation via compute_centered_ranks in evorl/ec/optimizers/openes.py to normalize fitness values and reduce variance. This is applied before the tell update.

Distributed Evaluation

The ECWorkflow in evorl/workflows/ec_workflow.py automatically handles multi-device parallelism. It slices the population across available devices using state.distributed_info, evaluates subsets in parallel, and gathers fitnesses via all_gather before calling tell.

Observation Preprocessing

Workflows support observation normalization via obs_preprocessor configured in _postsetup. This is implemented in evorl/algorithms/ec/so/openes.py for OpenES and helps stabilize evolution in high-dimensional observation spaces.

Complete Working Example

Here is a minimal, runnable example using Hydra configuration:

import jax
from hydra import compose, initialize
from evorl.algorithms.ec.so.cmaes import CMAESWorkflow

# Initialize Hydra with the EvoRL config directory

with initialize(config_path="../configs"):
    cfg = compose(
        config_name="config",
        overrides=[
            "agent=ec/cmaes",
            "env=brax/ant",
            "num_iters=10",
            "ec_optimizer.pop_size=64",
        ]
    )

# Build and JIT-compile the workflow

workflow = CMAESWorkflow.build_from_config(cfg, enable_jit=True)

# Initialize random key and state

key = jax.random.PRNGKey(42)
state = workflow.init(key)

# Run evolution loop

for iteration in range(10):
    state = workflow.step(state)
    print(f"Iteration {iteration+1}: Best objective = {state.metrics.best_objective:.3f}")

This example mirrors the test suite in tests/test_ec_workflow.py, which validates the full pipeline by running 100 evolution steps and verifying metric shapes.

Summary

• EvoRL's EC module provides a three-layer architecture: algorithm adapters (EvoXAlgorithmAdapter), native optimizers (OpenES, CMAES), and workflow templates (ECWorkflowTemplate).
• OpenES is implemented natively in evorl/ec/optimizers/openes.py with fitness shaping and noise table support.
• CMA-ES requires wrapping the evoX implementation via EvoXAlgorithmAdapter in evorl/ec/optimizers/evox_wrapper.py, which handles sign conversion and parameter flattening.
• Workflows handle environment creation, deterministic agent setup, distributed evaluation, and metric logging automatically.
• Distributed training is supported out-of-the-box via population slicing and all_gather operations in evorl/workflows/ec_workflow.py.

Frequently Asked Questions

How do I add a custom evolutionary algorithm to EvoRL?

Implement the EvoOptimizer interface with ask and tell methods, or wrap an existing evox.Algorithm using EvoXAlgorithmAdapter. Then create a workflow class inheriting from ECWorkflowTemplate that instantiates your optimizer in _build_from_config. The workflow handles environment interaction and evaluation automatically.

Why does CMA-ES require an adapter while OpenES does not?

OpenES is implemented natively in EvoRL (evorl/ec/optimizers/openes.py) following the framework's maximization convention. CMA-ES is imported from the evoX library, which minimizes objectives. The EvoXAlgorithmAdapter in evorl/ec/optimizers/evox_wrapper.py negates fitness values and handles parameter vectorization to bridge this gap.

Can I use observation normalization with evolutionary algorithms?

Yes. Both OpenESWorkflow and CMAESWorkflow support observation preprocessing via the _postsetup method. You can configure a running mean/std normalizer (VBN) in the workflow config, which is applied to observations before they reach the policy during population evaluation.

How does EvoRL handle distributed evolution strategies?

The ECWorkflow in evorl/workflows/ec_workflow.py automatically distributes population evaluation across available devices. It slices the population based on distributed_info.rank and world_size, evaluates subsets in parallel, then gathers fitnesses using all_gather before calling the optimizer's tell method. This allows scaling to large populations across multiple GPUs without code changes.

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