# How to Create a Custom Agent in EvoRL with compute_actions and evaluate_actions Methods

> Learn to create a custom Agent in EvoRL by implementing compute_actions and evaluate_actions methods. Subclass the Agent class for flexible agent development.

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

---

**To create a custom Agent in EvoRL, subclass the abstract `Agent` class from [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) and implement the `init`, `compute_actions`, and `evaluate_actions` methods to handle initialization, stochastic exploration during rollouts, and deterministic evaluation respectively.**

EvoRL provides a minimal abstract base class for reinforcement learning agents that standardizes how policies interact with training workflows. To create a custom Agent in EvoRL, you must inherit from the `Agent` class and implement three core methods that define how the policy initializes its state, explores during data collection, and evaluates without exploration. This functional design ensures seamless integration with JAX-based workflows while maintaining compatibility with both on-policy and off-policy training loops.

## Understanding the Agent Interface in EvoRL

The EvoRL library defines a strict interface in [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) that all custom agents must follow. This interface separates the concerns of parameter initialization, stochastic action selection during training, and deterministic action selection during evaluation.

### The Abstract Agent Class

The `Agent` class inherits from `PyTreeNode` and defines three abstract methods that form the contract between your policy and EvoRL workflows:

```python
class Agent(PyTreeNode, metaclass=ABCMeta):
    @abstractmethod
    def init(self, obs_space: Space, action_space: Space, key: PRNGKey) -> AgentState: ...

    @abstractmethod
    def compute_actions(self, agent_state: AgentState,
                        sample_batch: SampleBatch,
                        key: PRNGKey) -> tuple[Action, PolicyExtraInfo]: ...

    @abstractmethod
    def evaluate_actions(self, agent_state: AgentState,
                         sample_batch: SampleBatch,
                         key: PRNGKey) -> tuple[Action, PolicyExtraInfo]: ...

```

### AgentState and PolicyExtraInfo

The `AgentState` object is a lightweight immutable container (sub-class of `PyTreeData`) that stores three critical components:

- **`params`** – Model weights, typically wrapped in a custom dataclass like `NetworkParams`.
- **`obs_preprocessor_state`** – Optional running statistics for observation normalization.
- **`extra_state`** – Additional static data such as action space bounds or network configuration.

`PolicyExtraInfo` is a `PyTreeDict` that carries auxiliary data from action computation, such as log-probabilities for PPO or hidden RNN states. The workflow stores these in `SampleBatch.extras.policy_extras` for later use by loss functions.

## Implementing the Core Methods

When you create a custom Agent in EvoRL, you must implement three distinct methods that correspond to different phases of the RL lifecycle.

### Initializing Agent State

The `init` method is called once at the start of a workflow to produce the initial `AgentState`. It receives the observation space, action space, and a JAX random key. For neural network agents, this method typically samples a dummy observation, initializes network parameters using `jax.random` or Flax, and packages them into `AgentState`.

```python
def init(self, obs_space, action_space, key):
    dummy_obs = obs_space.sample(key)
    dummy_obs = jax.tree_map(lambda x: x[None, ...], dummy_obs)  # Add batch dim

    params = self.policy_network.init(key, dummy_obs)
    return AgentState(params=Params(policy=params), extra_state=PyTreeDict())

```

### Stochastic Action Computation

The `compute_actions` method is invoked during **rollout** phases to generate actions for environment interaction. This is where you implement exploration strategies such as Gaussian noise, epsilon-greedy, or sampling from a categorical distribution. The method must return a tuple of `(action, policy_extra_info)`.

According to the source code in [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py), exploration noise is added in `compute_actions` while `evaluate_actions` remains deterministic. For continuous actions, typical implementations split network outputs into `mean` and `log_std`, sample noise using `jax.random.normal`, and apply bounds via `jnp.tanh`.

### Deterministic Action Evaluation

The `evaluate_actions` method is used during **evaluation** or testing phases where exploration is undesirable. This method should return the deterministic best action—typically the mode of the distribution for stochastic policies or the direct network output for deterministic policies.

As implemented in [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py), `evaluate_actions` often computes the same forward pass but selects the distribution mode rather than sampling, ensuring reproducible evaluation results.

## Practical Code Examples

The following examples demonstrate how to create a custom Agent in EvoRL, ranging from a minimal non-learnable policy to a full neural network implementation.

### Minimal Deterministic Agent

For debugging or baseline comparisons, you can implement a constant-action agent that requires no neural network. This example from [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) patterns shows the minimal required structure:

```python
from evorl.agent import Agent, AgentState
from evorl.sample_batch import SampleBatch
from evorl.types import Action, PolicyExtraInfo, PyTreeDict
import chex
import jax
import jax.numpy as jnp

class ConstantAgent(Agent):
    """Returns a fixed action for all observations."""
    
    def __init__(self, const_action: Action):
        self._const_action = const_action

    def init(self, obs_space, action_space, key):
        # No learnable parameters

        return AgentState(params={}, extra_state=PyTreeDict())

    def compute_actions(
        self,
        agent_state: AgentState,
        sample_batch: SampleBatch,
        key: chex.PRNGKey,
    ) -> tuple[Action, PolicyExtraInfo]:
        # Batch the constant action to match observation batch size

        batch_shape = sample_batch.obs.shape[:-len(self._const_action.shape)]
        actions = jax.vmap(lambda _: self._const_action)(jnp.arange(jnp.prod(jnp.array(batch_shape))))
        actions = actions.reshape(*batch_shape, *self._const_action.shape)
        return actions, PyTreeDict()

    def evaluate_actions(
        self,
        agent_state: AgentState,
        sample_batch: SampleBatch,
        key: chex.PRNGKey,
    ) -> tuple[Action, PolicyExtraInfo]:
        # Deterministic policies return the same action as stochastic mode

        return self.compute_actions(agent_state, sample_batch, key)

```

### Learnable Neural Network Agent

For a learnable policy, you typically define a Flax module and initialize parameters in `init`. This example follows the pattern from [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) and [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py):

```python
from evorl.agent import Agent, AgentState
from evorl.sample_batch import SampleBatch
from evorl.types import Action, PolicyExtraInfo, PyTreeDict, Params
import chex
import flax.linen as nn
import jax
import jax.numpy as jnp

class MLPPolicy(nn.Module):
    hidden_dims: tuple[int, ...] = (256, 256)
    action_dim: int = 1

    @nn.compact
    def __call__(self, obs):
        x = obs
        for dim in self.hidden_dims:
            x = nn.relu(nn.Dense(dim)(x))
        # Output mean and log_std for continuous control

        mean = nn.tanh(nn.Dense(self.action_dim)(x))
        log_std = nn.tanh(nn.Dense(self.action_dim)(x))
        return jnp.concatenate([mean, log_std], axis=-1)

class MLPAgent(Agent):
    policy_network: nn.Module = MLPPolicy()
    
    def init(self, obs_space, action_space, key):
        dummy_obs = obs_space.sample(key)
        dummy_obs = jax.tree_map(lambda x: x[None, ...], dummy_obs)
        params = self.policy_network.init(key, dummy_obs)
        return AgentState(
            params=Params(policy=params), 
            extra_state=PyTreeDict()
        )

    def compute_actions(self, agent_state, sample_batch, key):
        obs = sample_batch.obs
        raw = self.policy_network.apply(agent_state.params.policy, obs)
        mean, log_std = jnp.split(raw, 2, axis=-1)
        std = jnp.exp(log_std)
        
        # Gaussian exploration noise

        eps = jax.random.normal(key, mean.shape)
        actions = jnp.tanh(mean + eps * std)
        
        extras = PyTreeDict(mean=mean, log_std=log_std)
        return actions, extras

    def evaluate_actions(self, agent_state, sample_batch, key):
        obs = sample_batch.obs
        raw = self.policy_network.apply(agent_state.params.policy, obs)
        mean, _ = jnp.split(raw, 2, axis=-1)
        actions = jnp.tanh(mean)  # Deterministic mode

        return actions, PyTreeDict()

```

## Integrating with EvoRL Workflows

Once you create a custom Agent in EvoRL, you can plug it into any workflow. The `OnPolicyWorkflow` (defined in [`evorl/workflows/onpolicy_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/onpolicy_workflow.py)) automatically handles the lifecycle:

```python
from evorl.workflows import OnPolicyWorkflow
from evorl.envs import create_env

env = create_env(name="CartPole-v1")
agent = MLPAgent()  # Your custom agent

workflow = OnPolicyWorkflow(env=env, agent=agent, num_iterations=1000)
workflow.run()

```

During execution, the workflow calls `agent.init` once, then repeatedly invokes `compute_actions` for rollouts to populate `SampleBatch` data, and uses `evaluate_actions` during validation phases.

## Summary

To create a custom Agent in EvoRL with `compute_actions` and `evaluate_actions`, follow these key steps:

- **Inherit from `Agent`**: Subclass the abstract base class defined in [`evorl/agent.py`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py), which inherits from `PyTreeNode`.
- **Implement `init`**: Return an `AgentState` containing initialized network `params`, optional `obs_preprocessor_state`, and `extra_state`.
- **Implement `compute_actions`**: Use this for rollout exploration; add noise (Gaussian, epsilon-greedy) and return `(action, policy_extra_info)`.
- **Implement `evaluate_actions`**: Use this for deterministic evaluation; return the mode or mean action without exploration noise.
- **Return `PolicyExtraInfo`**: Package auxiliary data like log-probabilities or hidden states in a `PyTreeDict` for algorithm-specific loss calculations.

## Frequently Asked Questions

### What is the difference between compute_actions and evaluate_actions in EvoRL?

The `compute_actions` method is called during training rollouts to collect experience and should implement **exploration** strategies such as Gaussian noise or epsilon-greedy sampling. In contrast, `evaluate_actions` is used during testing or validation phases and must return **deterministic** actions, typically by selecting the distribution mode or mean without adding noise. This separation ensures reproducible evaluation metrics while maintaining adequate exploration during training.

### Do I need to implement loss functions to create a custom Agent in EvoRL?

No, implementing `critic_loss` or `actor_loss` methods is **optional** for a minimal custom Agent. These methods are only required if you are building an off-policy algorithm like TD3 or SAC that needs custom gradient updates. For basic policy evaluation or when using existing workflow templates, only `init`, `compute_actions`, and `evaluate_actions` are mandatory.

### How do I add observation normalization to my custom EvoRL Agent?

You can add observation normalization by utilizing the `obs_preprocessor` field and the `obs_preprocessor_state` attribute within `AgentState`. Initialize a running statistics filter from [`evorl/utils/running_statistics.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/running_statistics.py) in your `init` method, store its state in `AgentState.obs_preprocessor_state`, and apply the preprocessor to observations at the start of `compute_actions` and `evaluate_actions` before passing them to your network.

### Can I use Haiku instead of Flax for my custom Agent in EvoRL?

Yes, you can use Haiku or any JAX-compatible neural network library. The `Agent` interface is framework-agnostic; it only requires that your `init` method returns parameters in a format compatible with `jax.tree_map` operations (typically PyTrees). Replace the Flax `nn.Module` with a Haiku `hk.Module`, initialize parameters using `hk.transform`, and store them in `AgentState.params` following the same pattern shown in the Flax examples.