How to Implement Evolution-guided Reinforcement Learning (ERL) Workflows in EvoRL

EvoRL implements Evolution-guided Reinforcement Learning through a dual-workflow architecture that combines gradient-based RL agents with evolutionary population optimization, orchestrated via the ERLWorkflowBase class which manages shared replay buffers, RL injection mechanisms, and distributed population evaluation.

EvoRL provides a flexible framework for hybrid neuroevolution and reinforcement learning research using JAX. This guide explains how to implement Evolution-guided Reinforcement Learning (ERL) workflows in EvoRL by leveraging its modular workflow system that glues together EC optimizers, RL algorithms like TD3, and distributed training infrastructure.

Core Architecture of EvoRL ERL Workflows

EvoRL structures ERL implementations around three core components: a gradient-based RL learner, an evolutionary population optimizer, and a unified workflow engine that coordinates both paradigms through shared state and replay buffers.

The Workflow Protocol

All workflows inherit from the abstract Workflow protocol defined in evorl/workflows/workflow.py. This base class defines the generic lifecycle methods that every workflow must implement:

# evorl/workflows/workflow.py

class Workflow(Protocol):
    def setup(self, key): ...
    def step(self, state): ...
    def evaluate(self, state): ...

The protocol ensures consistent interfaces for initialization (setup), training iterations (step), and evaluation across all EvoRL algorithms.

EC and RL Workflow Specializations

EvoRL provides two specialized workflow bases that handle distinct computational patterns:

ECWorkflow (evorl/workflows/ec_workflow.py) manages population-based evolution. It handles population initialization, distributed slicing across devices, and fitness evaluation through the EpisodeCollector. The step method implements the ask-evaluate-tell pattern:

# evorl/workflows/ec_workflow.py

class ECWorkflow(Workflow):
    def step(self, state):
        # 1️⃣ ask EC optimizer for population

        # 2️⃣ replace actor params in a copy of the RL state

        # 3️⃣ evaluate each slice of the population (EpisodeCollector)

        # 4️⃣ compute fitnesses → tell EC optimizer

RLWorkflow (evorl/workflows/rl_workflow.py) manages gradient-based learning. It handles replay buffer sampling, gradient computation via Optax, and parameter updates:

# evorl/workflows/rl_workflow.py

class RLWorkflow(Workflow):
    def step(self, state):
        # 1️⃣ sample from replay buffer

        # 2️⃣ compute gradients with optax

        # 3️⃣ apply updates

ERLWorkflowBase Integration

The ERLWorkflowBase class in evorl/algorithms/erl/erl_workflow.py inherits capabilities from both EC and RL workflows. It creates a single environment instance used for both EC and RL rollouts, manages a shared replay buffer, and implements population-mean evaluation. This base class provides the foundation for concrete implementations like ERLGAWorkflow (Genetic Algorithm + TD3) found in evorl/algorithms/erl/erl_td3/erl_ga.py.

Building an ERL Workflow from Configuration

The typical entry point for creating an ERL workflow is the ERLGAWorkflow._build_from_config method. This factory method receives an OmegaConf configuration and constructs all required components:

# evorl/algorithms/erl/erl_td3/erl_ga.py (excerpt)

env = create_env(...)

agent = make_mlp_td3_agent(...)

optimizer = optax.adam(config.optimizer.lr)           # RL optimizer

ec_optimizer = ERLGAMod(... )                        # Evolutionary optimizer

ec_collector = EpisodeCollector(env, action_fn=..., ...)   # EC rollouts

rl_collector = EpisodeCollector(env, action_fn=..., ...)   # RL rollouts

replay_buffer = ReplayBuffer(...)
evaluator = Evaluator(eval_env, action_fn=agent.evaluate_actions, ...)

agent_state_vmap_axes = AgentState(params=0, obs_preprocessor_state=None)

workflow = ERLGAWorkflow(
    env=env,
    agent=agent,
    agent_state_vmap_axes=agent_state_vmap_axes,
    optimizer=optimizer,
    ec_optimizer=ec_optimizer,
    ec_collector=ec_collector,
    rl_collector=rl_collector,
    evaluator=evaluator,
    replay_buffer=replay_buffer,
    config=config,
)

Key configuration parameters include:

• pop_size: Total population size for evolution
• num_elites: Number of top performers preserved each generation
• num_rl_agents: Number of RL agents contributing gradients
• rl_injection_interval: Frequency of injecting RL policies into the population

Running the ERL Training Loop

After construction, the workflow follows a standardized execution pattern using JAX PRNG keys:

import jax
from evorl.utils import random_key

key = random_key()
state = workflow.setup(key)          # initialise replay buffer, agent, EC pop.

final_state = workflow.learn(state)  # runs the full training loop

The learn method (implemented in ERLGAWorkflow.learn) repeatedly calls step, logs metrics via the built-in recorder, runs periodic evaluation, and checkpoints the complete State including the replay buffer. This continues until the total_episodes budget specified in the configuration is exhausted.

Key Mechanisms in ERL Workflows

RL Injection into the Evolutionary Population

ERL workflows periodically inject current RL agent parameters into the evolutionary population to combine gradient-based optimization with evolutionary search. This occurs in ERLGAWorkflow._rl_injection (lines 199-205), which copies the current RL actor parameters into designated slots of the EC population every rl_injection_interval steps. This mechanism enables the evolutionary process to exploit high-performing policies discovered through temporal-difference learning.

Shared Replay Buffer Architecture

Both EC and RL rollouts write trajectory data to a shared ReplayBuffer instance (evorl/replay_buffers/replay_buffer.py). The ECWorkflow._ec_rollout and RLWorkflow._rl_rollout methods both call self.replay_buffer.add() during execution. This shared memory allows evolutionary fitness evaluation and RL gradient updates to leverage the same environmental experience, improving sample efficiency compared to isolated training paradigms.

Distributed Population Slicing

When running across multiple devices via pmap, ECWorkflow.step handles distributed population slicing using state.distributed_info.rank and world_size. Each device processes a slice of the population independently, with metrics aggregated across devices via MetricBase.all_reduce(pmap_axis_name). This enables scaling ERL workflows to large population sizes across GPU/TPU clusters without modifying the core algorithm logic.

Customizing Your ERL Implementation

EvoRL's modular design supports several extension patterns:

Swap the RL algorithm: Replace make_mlp_td3_agent with any agent factory from evorl.algorithms.* (e.g., SAC) and adjust the optimizer signature accordingly.

Change the EC optimizer: Substitute ERLGAMod with alternative optimizers like ERLGA or ERLGAES from evorl/ec/optimizers/erl_ga.py to experiment with different selection and mutation strategies.

Define custom fitness: Override ECWorkflow._metrics_to_fitnesses to implement alternative fitness landscapes. For example, penalizing episode length:

from evorl.workflows.ec_workflow import ECWorkflow
import jax.numpy as jnp

class MyECWorkflow(ECWorkflow):
    def _metrics_to_fitnesses(self, metrics):
        reward = jnp.mean(metrics.episode_returns, axis=-1)
        length = jnp.mean(metrics.episode_lengths, axis=-1)
        return reward - 0.001 * length

Summary

• EvoRL implements ERL through composition: The framework combines ECWorkflow for evolutionary computation and RLWorkflow for gradient-based learning via the unified ERLWorkflowBase class.
• Configuration-driven instantiation: Use ERLGAWorkflow.build_from_config() to construct complete workflows from YAML specifications, handling agent creation, optimizer setup, and collector initialization.
• Shared resources maximize efficiency: Both paradigms share a single replay buffer and environment instance, with the RL agent periodically injecting parameters into the evolutionary population via _rl_injection.
• Native distributed support: Population evaluation automatically distributes across devices using JAX's pmap, with built-in metric aggregation via all_reduce operations.
• Extensible architecture: Swap RL algorithms, EC optimizers, or fitness functions by subclassing workflow components without rewriting the training loop.

Frequently Asked Questions

How does EvoRL handle the interaction between RL and evolutionary components?

EvoRL manages interaction through the ERLWorkflowBase class which orchestrates both paradigms via a shared replay buffer and periodic parameter injection. The _rl_injection method in ERLGAWorkflow copies current RL actor parameters into the EC population at intervals defined by rl_injection_interval, while both components read from the same replay buffer populated by EpisodeCollector rollouts.

Can I use a different RL algorithm than TD3 with ERL workflows?

Yes. While ERLGAWorkflow uses make_mlp_td3_agent by default, you can substitute any RL agent from the evorl.algorithms module. Modify the agent factory call in your workflow's _build_from_config method and ensure the optimizer configuration matches your chosen algorithm's requirements.

How does distributed training work with ERL populations?

When enable_multi_devices=True, EvoRL uses JAX's pmap to distribute the population across available devices. The ECWorkflow.step method slices the population by device rank, evaluates each slice independently, then aggregates fitness metrics using MetricBase.all_reduce. This occurs transparently without requiring changes to the EC optimizer or agent code.

What is the purpose of the shared replay buffer in ERL workflows?

The shared replay buffer (ReplayBuffer) allows both evolutionary rollouts and RL gradient updates to utilize the same experience data. When ec_collector and rl_collector execute environment interactions, both write transitions to self.replay_buffer. This design improves sample efficiency by ensuring TD3 updates learn from the diverse behaviors generated by the evolving population, not just the current RL policy.

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