How to Implement Custom Evaluators for Agent Evaluation in EvoRL

To implement custom Evaluators in EvoRL, subclass the abstract Evaluator class from 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. 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: Shows multi-objective evaluation by overriding evaluate and adding custom metric_names.
• 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.

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:

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:

    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:

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 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 to safely split keys within scanned functions.

Summary

• Subclass Evaluator from 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 for multi-objective evaluation patterns and 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.

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