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

> Implement Cross-Entropy Method for RL with covariance adaptation in EvoRL. Configure SepCEM optimizer with ExponentialScheduleSpec for adaptive diagonal covariance decay. Get started today.

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

---

**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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/configs/agent/erl/cemrl.yaml), the schedule parameters define the exploration annealing:

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

```

During initialization in [`evorl/ec/optimizers/cem.py`](https://github.com/emi-group/evorl/blob/main/evorl/ec/optimizers/cem.py), the optimizer creates the full variance tree using the initial schedule value:

```python
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`](https://github.com/emi-group/evorl/blob/main/evorl/ec/optimizers/cem.py), the optimizer generates population noise using the current diagonal variance:

```python
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`](https://github.com/emi-group/evorl/blob/main/evorl/ec/optimizers/cem.py) performs the core covariance adaptation after evaluating fitness scores. First, it computes the weighted variance of elite samples:

```python
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:

```python
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`](https://github.com/emi-group/evorl/blob/main/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

```python
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:

```bash
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`](https://github.com/emi-group/evorl/blob/main/configs/agent/erl/cemrl.yaml). For slower decay and higher initial exploration:

```yaml
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:

```python
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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/configs/agent/erl/cemrl.yaml) provides tuned hyperparameters for standard MuJoCo benchmarks.