# How to Implement Custom Evaluators for Agent Evaluation in EvoRL

> Learn to implement custom evaluators for agent evaluation in EvoRL. Subclass the Evaluator class, define metrics, and override the evaluate method for custom metrics.

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

---

**To implement custom Evaluators in EvoRL, subclass the abstract `Evaluator` class from [`evorl/evaluators/evaluator.py`](https://github.com/emi-group/evorl/blob/main/evorl/evaluators/evaluator.py), define your `metric_names` tuple, and override the `evaluate` method to extract custom metrics from `env_state.info.metrics` while reusing the library's rollout utilities.**

The EvoRL library (evorl) provides a flexible evaluation framework for evolutionary reinforcement learning experiments. By implementing custom Evaluators for agent evaluation in EvoRL, you can capture domain-specific metrics—such as energy consumption, safety violations, or multi-objective rewards—beyond standard episode returns. This guide walks you through the architecture and provides a complete implementation example based on the actual source code.

## Understanding the Evaluator Architecture

The evaluation system centers on the abstract `Evaluator` class defined in [`evorl/evaluators/evaluator.py`](https://github.com/emi-group/evorl/blob/main/evorl/evaluators/evaluator.py). This base class handles vectorized environment reset, parallel rollouts, and PRNG key management. It expects a vectorized environment (where `env.num_envs` is defined) and an action function conforming to the `AgentActionFn` signature.

Two reference implementations demonstrate extension patterns:

- **[`evorl/evaluators/mo_brax_evaluator.py`](https://github.com/emi-group/evorl/blob/main/evorl/evaluators/mo_brax_evaluator.py)**: Shows multi-objective evaluation by overriding `evaluate` and adding custom `metric_names`.
- **[`evorl/evaluators/ec_evaluator.py`](https://github.com/emi-group/evorl/blob/main/evorl/evaluators/ec_evaluator.py)**: Demonstrates Evolutionary-Computation style rollouts using `EpisodeCollector`.

The base class automatically handles the evaluation loop logic, including discounted return calculation via `compute_discount_return` and episode length tracking via `compute_episode_length`, both from [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py).

## Step-by-Step Guide to Building a Custom Evaluator

### Step 1: Create the Subclass

Begin by importing the base class and creating your implementation. The following example shows the minimal structure required:

```python
from evorl.evaluators.evaluator import Evaluator
from evorl.utils.jax_utils import rng_split
from evorl.utils.rl_toolkits import (
    compute_discount_return,
    compute_episode_length,
    fast_eval_rollout_episode,
    rollout,
)
import chex
import jax
import jax.numpy as jnp
import math
import logging

logger = logging.getLogger(__name__)

class MyCustomEvaluator(Evaluator):
    """Custom evaluator that records reward and a user-defined metric."""
    
    metric_names: tuple[str, ...] = ("reward", "my_metric", "episode_lengths")

```

### Step 2: Implement the evaluate Method

Override the `evaluate` method to handle parallel rollout execution and custom metric extraction. The implementation must account for both discounted (`discount < 1.0`) and undiscounted (`discount == 1.0`) cases using the library's optimized paths:

```python
    def evaluate(
        self,
        agent_state: chex.ArrayTree,
        key: chex.PRNGKey,
        num_episodes: int,
    ) -> chex.ArrayTree:
        """Run `num_episodes` rollouts and return a dict of stacked metrics."""
        num_envs = self.env.num_envs
        num_iters = math.ceil(num_episodes / num_envs)
        
        if num_episodes % num_envs:
            logger.warning(
                f"num_episodes ({num_episodes}) not divisible by envs ({num_envs}); "
                f"running {num_iters * num_envs} episodes instead."
            )

        action_fn = self.action_fn
        env_reset = self.env.reset
        env_step = self.env.step
        
        def _extract_metrics(env_state):
            """Extract custom metrics from environment state info."""
            raw = env_state.info.metrics
            return {
                name: raw.get(name, jnp.zeros_like(env_state.reward))
                for name in self.metric_names if name != "reward"
            }

        def _evaluate_one(key, _):
            next_key, init_key, rollout_key = rng_split(key, 3)
            state = env_reset(init_key)
            
            # Fast path for undiscounted evaluation

            if self.discount == 1.0:
                trajectory, _ = fast_eval_rollout_episode(
                    env_step,
                    action_fn,
                    state,
                    agent_state,
                    rollout_key,
                    self.max_episode_steps,
                )
                
                # Extract custom metrics and add to trajectory

                custom_metrics = _extract_metrics(state)
                for name in self.metric_names:
                    if name != "reward":
                        trajectory.rewards[name] = custom_metrics[name]
                
                returns = trajectory.rewards["reward"]
                lengths = compute_episode_length(trajectory.dones)
                
                # Compute discounted returns for custom metrics (treats as undiscounted when discount=1.0)

                extra = {
                    name: compute_discount_return(
                        trajectory.rewards[name], trajectory.dones, self.discount
                    )
                    for name in self.metric_names
                    if name not in ("reward", "episode_lengths")
                }
                result = {"reward": returns, "episode_lengths": lengths, **extra}
                
            else:
                # Standard rollout for discounted case

                trajectory, _ = rollout(
                    env_step,
                    action_fn,
                    state,
                    agent_state,
                    rollout_key,
                    self.max_episode_steps,
                )
                
                custom_metrics = _extract_metrics(state)
                for name in self.metric_names:
                    if name != "reward":
                        trajectory.rewards[name] = custom_metrics[name]
                
                returns = compute_discount_return(
                    trajectory.rewards["reward"], trajectory.dones, self.discount
                )
                lengths = compute_episode_length(trajectory.dones)
                
                extra = {
                    name: compute_discount_return(
                        trajectory.rewards[name], trajectory.dones, self.discount
                    )
                    for name in self.metric_names
                    if name not in ("reward", "episode_lengths")
                }
                result = {"reward": returns, "episode_lengths": lengths, **extra}
            
            return next_key, result

        # Run parallel evaluation using scan

        _, metrics = jax.lax.scan(_evaluate_one, key, (), length=num_iters)
        
        # Flatten the iteration dimension into the batch axis

        flat_metrics = {
            name: jax.lax.collapse(metrics[name], (0, 1))
            for name in metrics
        }
        return flat_metrics

```

### Step 3: Integrate with Your Training Pipeline

Instantiate your custom evaluator with a vectorized environment and action function:

```python
import jax
from evorl.envs import make_env

# Create vectorized environment

env = make_env("my_custom_env", num_envs=8)

# Define your policy's action function

def my_action_fn(agent_state, batch, key):
    # Your policy logic here

    return action, {}

# Initialize evaluator

evaluator = MyCustomEvaluator(
    env=env,
    action_fn=my_action_fn,
    max_episode_steps=1000,
    discount=0.99,
)

# Run evaluation

key = jax.random.PRNGKey(0)
metrics = evaluator.evaluate(agent_state, key, num_episodes=32)

print(f"Returns shape: {metrics['reward'].shape}")          # (32,)

print(f"Custom metric shape: {metrics['my_metric'].shape}") # (32,)

```

## Key Utility Functions

The EvoRL library provides several JAX-optimized utilities in [`evorl/utils/rl_toolkits.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/rl_toolkits.py) that your custom evaluator should leverage:

- **`compute_discount_return(rewards, dones, discount)`**: Calculates discounted returns across episode boundaries.
- **`compute_episode_length(dones)`**: Computes episode lengths from done signals.
- **`fast_eval_rollout_episode(...)`**: Optimized rollout for undiscounted evaluation (`discount == 1.0`).
- **`rollout(...)`**: Standard rollout function for gathering full trajectories.

For PRNG key management across JAX transformations, use `rng_split` from [`evorl/utils/jax_utils.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/jax_utils.py) to safely split keys within scanned functions.

## Summary

- Subclass **`Evaluator`** from [`evorl/evaluators/evaluator.py`](https://github.com/emi-group/evorl/blob/main/evorl/evaluators/evaluator.py) to create custom evaluation logic.
- Define **`metric_names`** as a class attribute to declare which metrics to extract from `env_state.info.metrics`.
- Override **`evaluate`** to implement parallel rollout logic using `jax.lax.scan`, handling both discounted and undiscounted cases.
- Extract custom metrics from **`env_state.info.metrics`** and aggregate them using `compute_discount_return`.
- Use **`jax.lax.collapse`** to flatten batch dimensions when returning results.
- Reference **[`mo_brax_evaluator.py`](https://github.com/emi-group/evorl/blob/main/mo_brax_evaluator.py)** for multi-objective evaluation patterns and **[`ec_evaluator.py`](https://github.com/emi-group/evorl/blob/main/ec_evaluator.py)** for evolutionary computation integration.

## Frequently Asked Questions

### How do I handle environments that don't expose metrics via `info.metrics`?

If your environment stores custom data in a different location, modify the `_extract_metrics` helper method in your evaluator to access the correct path. For example, if metrics are in `env_state.extras`, update the extraction logic to reference `env_state.extras.get(name)`. Ensure your environment wrapper consistently populates this field during the `step` and `reset` calls.

### What is the difference between `fast_eval_rollout_episode` and `rollout`?

According to the evorl source code, `fast_eval_rollout_episode` is optimized for undiscounted evaluation where `discount == 1.0` and only requires the final return, while `rollout` captures the full trajectory necessary for discounted return calculation. Your custom evaluator should check `self.discount` to select the appropriate path, as shown in the implementation example above.

### Can I implement custom aggregation logic beyond sum-based returns?

Yes. While the standard implementation uses `compute_discount_return` to sum rewards, you can replace this with custom JAX operations inside your overridden `evaluate` method. For example, you might calculate the maximum value reached during an episode, count specific events using `jnp.sum` on boolean masks, or compute variance across episode steps. Return these aggregations in the result dictionary alongside standard metrics.

### How do I ensure my evaluator works with different numbers of parallel environments?

The base class handles automatic batching via `math.ceil(num_episodes / num_envs)` and `jax.lax.scan`. Your implementation should respect `self.env.num_envs` and use `jax.lax.collapse` to flatten the leading dimensions before returning metrics. If `num_episodes` is not divisible by `num_envs`, the evaluator automatically runs extra episodes and logs a warning, ensuring you always receive exactly the requested statistical sample size.