How to Implement Cross-Entropy Method for RL (CEM-RL) with Covariance Adaptation in EvoRL

Implement CEM-RL in EvoRL by configuring the SepCEM optimizer with an ExponentialScheduleSpec for the cov_eps term, which adaptively decays the diagonal covariance during the population-based search.

EvoRL is a JAX-based evolutionary reinforcement learning framework developed by emi-group that provides modular implementations of hybrid algorithms. The CEM-RL implementation combines population-based Cross-Entropy Method optimization with TD3 actor-critic updates, featuring adaptive diagonal covariance through the cov_eps schedule to maintain stable exploration throughout training.

Core Architecture of CEM-RL

The CEM-RL system in EvoRL relies on three primary components that orchestrate the hybrid evolutionary and reinforcement learning process:

• SepCEM optimizer – Located in evorl/ec/optimizers/cem.py, this class implements separate-mean Cross-Entropy Method with diagonal covariance and optional mirror sampling. It maintains the population distribution parameters.
• ExponentialScheduleSpec – Defined in evorl/ec/optimizers/utils.py, this dataclass holds the exponential decay schedule (init, final, decay) for the cov_eps covariance floor.
• CEMRLWorkflow – Found in evorl/algorithms/erl/cemrl_td3/cemrl.py, this workflow orchestrates the RL-CEM loop by calling ask, performing TD3 updates, rolling out the population, and calling tell.

Covariance Adaptation Mechanism

The covariance adaptation prevents premature convergence while gradually focusing the search on high-performing regions. The mechanism operates through three distinct phases during each iteration.

Schedule Initialization

The covariance floor is configured via YAML and initialized within the optimizer state. In configs/agent/erl/cemrl.yaml, the schedule parameters define the exploration annealing:

cov_eps:
  init: 1e-2
  final: 1e-5
  decay: 0.001   # Polyak step size

During initialization in evorl/ec/optimizers/cem.py, the optimizer creates the full variance tree using the initial schedule value:

variance = jtu.tree_map(
    lambda x: jnp.full_like(x, self.cov_eps_schedule.init), mean
)

Population Sampling

During the ask phase in evorl/ec/optimizers/cem.py, the optimizer generates population noise using the current diagonal variance:

noise = jtu.tree_map(
    lambda x, var, k: jax.random.normal(k, (self.pop_size, *x.shape)) * jnp.sqrt(var),
    state.mean, state.variance, sample_keys,
)

This sampling draws from a Gaussian distribution where the standard deviation is derived from the adaptive variance state maintained across iterations.

Elite-Based Updates

The tell method in evorl/ec/optimizers/cem.py performs the core covariance adaptation after evaluating fitness scores. First, it computes the weighted variance of elite samples:

def var_update(m, x):
    x_norm = jnp.square(x[elite_indices] - m)
    return weight_sum(x_norm, self.elite_weights) + state.cov_eps

Then it updates the cov_eps floor using an exponential moving average toward the final value:

cov_eps = optax.incremental_update(
    self.cov_eps_schedule.final, state.cov_eps, self.cov_eps_schedule.decay
)

This dual update ensures the covariance reflects recent elite performance while gradually reducing the minimum exploration noise according to the exponential schedule.

Workflow Integration and RL Injection

The CEMRLWorkflow class manages the interaction between evolutionary search and policy gradients. During each training step in evorl/algorithms/erl/cemrl_td3/cemrl.py, the workflow:

1. Generates a population of actor parameters using ec_optimizer.ask()
2. Injects RL updates via _rl_injection() to refine a subset of actors with TD3 gradients
3. Evaluates the full population to obtain fitness scores from eval_metrics.episode_returns.mean(axis=-1)
4. Updates the distribution via ec_optimizer.tell() with the computed fitnesses

fitnesses = eval_metrics.episode_returns.mean(axis=-1)
ec_opt_state = self._rl_injection(ec_opt_state, pop_actor_params)
ec_metrics, ec_opt_state = self.ec_optimizer.tell(ec_opt_state, fitnesses)

The updated state.cov_eps is exposed through ec_info.cov_eps for monitoring purposes.

Configuration and Execution

To launch a CEM-RL experiment with covariance adaptation, execute the training script with the provided configuration:

python -m evorl.scripts.train \
    +config=agent/erl/cemrl.yaml \
    env_name=HalfCheetah-v4 \
    seed=42

You can customize the adaptation behavior by modifying the cov_eps block in configs/agent/erl/cemrl.yaml. For slower decay and higher initial exploration:

cov_eps:
  init: 5e-2
  final: 1e-6
  decay: 0.0005

Monitoring Covariance Adaptation

The workflow logs the current cov_eps value at each iteration, enabling analysis of the exploration schedule. Extract and visualize the adaptation trajectory:

import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv("logs/metrics.csv")
plt.plot(df["iteration"], df["ec/std"])
plt.xlabel("Iteration")
plt.ylabel("Mean Std (sqrt(variance))")
plt.title("Covariance Adaptation Over Time")
plt.show()

This monitoring confirms that the covariance properly decays from the initial value toward the floor while maintaining sufficient variance to escape local optima.

Summary

• SepCEM in evorl/ec/optimizers/cem.py implements diagonal covariance adaptation via the cov_eps schedule.
• The exponential schedule (init, final, decay) controls the minimum covariance floor to prevent premature collapse.
• Elite-based updates adjust the mean and variance using weighted samples from the top-performing population members.
• CEMRLWorkflow coordinates the hybrid TD3-CEM training loop through the standard ask and tell interface.
• Configuration occurs through YAML files with clear hyperparameters for population size, elite count, and covariance adaptation rates.

Frequently Asked Questions

What is the purpose of the cov_eps schedule in SepCEM?

The cov_eps schedule provides a time-varying floor for the diagonal covariance matrix, ensuring that exploration noise never collapses to zero prematurely. According to the implementation in evorl/ec/optimizers/cem.py, this term is added to the empirical elite variance during each tell update and decays exponentially from init toward final using the specified Polyak decay rate.

How does CEM-RL differ from standard CEM in EvoRL?

CEM-RL specifically integrates the SepCEM optimizer with a TD3 actor-critic learner through the CEMRLWorkflow class. Unlike standard CEM which only uses evolutionary updates, CEM-RL performs RL injection steps where sampled actors undergo gradient-based fine-tuning before the CEM update, combining the global search properties of evolution with the sample efficiency of policy gradients.

Can I modify the covariance adaptation parameters during training?

While the decay schedule is fixed at initialization from ExponentialScheduleSpec, you can adjust the init, final, and decay values in your YAML configuration to control the adaptation speed. The current implementation uses optax.incremental_update with a constant decay rate, so modifying parameters mid-training would require restarting from a checkpoint with altered config values.

Which environments work best with CEM-RL in EvoRL?

CEM-RL is designed for continuous control tasks compatible with TD3, such as MuJoCo environments (HalfCheetah, Walker2d, Ant). The covariance adaptation particularly benefits high-dimensional action spaces where maintaining appropriate exploration noise is critical. The default configs/agent/erl/cemrl.yaml provides tuned hyperparameters for standard MuJoCo benchmarks.

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