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

To create a custom Agent in EvoRL, subclass the abstract Agent class from 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 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:

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.

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, 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, 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 patterns shows the minimal required structure:

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 and evorl/algorithms/ppo.py:

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) automatically handles the lifecycle:

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, 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 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.

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