How to Add Custom Loss Functions to an Agent in EvoRL for Gradient-Based Updates

To add a custom loss function in EvoRL, implement a new loss method in your Agent subclass that returns a LossDict, then register the loss weight in your configuration to include it in the gradient update.

EvoRL is a JAX-based evolutionary reinforcement learning framework that cleanly separates policy logic from training workflows. When you need to add custom loss functions to an Agent in EvoRL for gradient-based updates, you must follow the framework's type contracts and hook into the workflow's loss aggregation mechanism.

Understanding the EvoRL Loss Architecture

The Agent and Workflow Separation

EvoRL separates policy logic (the Agent subclass) from the training loop (the workflow). The Agent defines how to compute losses, while the workflow handles gradient aggregation and optimization. This means you inject custom losses by extending the Agent class rather than modifying the workflow directly.

Key Types and Contracts

All loss methods must follow the LossFn signature defined in evorl/agent.py. A custom loss method must accept:

• agent_state: AgentState – the current parameters and optimizer state
• sample_batch: SampleBatch – the training data batch
• key: chex.PRNGKey – a JAX random key

The method must return a LossDict, which is a Mapping[str, chex.Array] defined in evorl/types.py. Each key in this dictionary represents a distinct loss term that the workflow can weight independently.

Implementing a Custom Loss Function

Step 1: Define the Loss Method in Your Agent

Create a subclass of an existing agent (e.g., PPOAgent or TD3Agent) and implement your custom loss logic. The method must return a dictionary with your new loss term.

def my_custom_loss(
    self,
    agent_state: AgentState,
    sample_batch: SampleBatch,
    key: chex.PRNGKey,
) -> LossDict:
    # Custom loss computation

    ...
    return {"custom_term": loss_value}

Step 2: Configure Loss Weights

The workflow multiplies each entry of the loss dict by a coefficient from config.loss_weights. Add an entry for your custom key in the YAML config or programmatically extend the dictionary:

loss_weights:
  actor_loss: 1.0
  critic_loss: 0.5
  custom_term: 0.001  # <-- new entry

Step 3: Integrate with the Training Workflow

For on-policy algorithms like PPO, the workflow calls self.agent.loss() (see evorl/algorithms/ppo.py and evorl/workflows/rl_workflow.py). Override the loss method to combine the base loss with your custom term.

For off-policy algorithms like TD3, the workflow calls separate critic_loss and actor_loss methods (see evorl/algorithms/td3.py). Add a new method (e.g., regularizer_loss) and modify the workflow's loss aggregation to include it, or simply add the term to the existing loss methods.

The gradient update is performed by agent_gradient_update in evorl/distributed/gradients.py. As long as the loss dict contains the new key, the update routine will include it automatically.

Practical Examples

Example: L2 Regularization for PPO

This example adds an L2 weight penalty to the PPO agent by extending PPOAgent and overriding the loss method.

# file: evorl/algorithms/custom_ppo.py

import chex
import jax.numpy as jnp
from evorl.agent import PPOAgent
from evorl.types import LossDict, AgentState, SampleBatch

class CustomPPOAgent(PPOAgent):
    """PPO agent with an extra L2 weight regularizer."""

    def loss(
        self, agent_state: AgentState, sample_batch: SampleBatch, key: chex.PRNGKey
    ) -> LossDict:
        # Call the original PPO loss

        base_loss = super().loss(agent_state, sample_batch, key)

        # Compute L2 norm of all policy parameters

        l2_penalty = sum(
            jnp.sum(jnp.square(p)) for p in jax.tree_util.tree_leaves(agent_state.params.policy_params)
        )
        # Add a new entry to the dict

        base_loss["l2_reg"] = l2_penalty
        return base_loss

To activate the regularization, add the weight to your configuration:

# in your experiment config

loss_weights:
  actor_loss: 1.0
  critic_loss: 0.5
  actor_entropy: 0.01
  approx_kl: 0.5
  l2_reg: 0.001   # <-- new entry

The existing evorl/workflows/rl_workflow.py already multiplies each key by config.loss_weights, so no workflow code changes are required.

Example: Critic Smoothness Regularizer for TD3

This example adds a gradient penalty to the TD3 critic to encourage smoothness.

# file: evorl/algorithms/custom_td3.py

import chex
import jax
import jax.numpy as jnp
from evorl.algorithms.td3 import TD3Agent
from evorl.types import LossDict, AgentState, SampleBatch

class CustomTD3Agent(TD3Agent):
    """TD3 agent with a critic smoothness regularizer."""

    def critic_loss(
        self, agent_state: AgentState, sample_batch: SampleBatch, key: chex.PRNGKey
    ) -> LossDict:
        # Original TD3 critic loss

        loss_dict = super().critic_loss(agent_state, sample_batch, key)

        # Smoothness: penalize large gradients of Q w.r.t. actions

        def q_fn(params, obs, act):
            return self.critic_network.apply(params, obs, act)

        grad_q = jax.jacrev(q_fn, argnums=2)(
            agent_state.params.critic_params, sample_batch.obs, sample_batch.actions
        )
        smoothness = jnp.mean(jnp.square(grad_q))
        loss_dict["smoothness"] = smoothness
        return loss_dict

Add the weight to your config:

# in your experiment config

loss_weights:
  critic_loss: 1.0
  actor_loss: 1.0
  smoothness: 0.01   # <-- new entry

The loss_fn inside TD3Workflow extracts loss_dict from critic_loss and actor_loss; the added key will be summed with its weight automatically.

Example: Using the Custom Agent

Instantiate your custom agent and pass it to the workflow:

from evorl.algorithms.custom_ppo import CustomPPOAgent
from evorl.workflows.rl_workflow import PPOWorkflow

# or

from evorl.algorithms.custom_td3 import CustomTD3Agent
from evorl.algorithms.td3 import TD3Workflow

# In your training script (e.g. scripts/train.py)

agent = CustomPPOAgent(
    continuous_action=True,
    policy_network=make_policy_network(...),
    value_network=make_v_network(...),
    obs_preprocessor=running_statistics.normalize if normalize_obs else None,
    # other PPO hyper-parameters …

)

workflow = PPOWorkflow(
    env=env,
    agent=agent,
    optimizer=optimizer,
    evaluator=evaluator,
    config=config,
)

Key Implementation Files

File	Role	Link
evorl/agent.py	Base Agent, AgentState, LossFn definitions	agent.py
evorl/types.py	LossDict, AgentState, PyTreeData utilities	types.py
evorl/algorithms/ppo.py	Reference PPO implementation with a loss method	ppo.py
evorl/algorithms/td3.py	Reference TD3 implementation with critic_loss & actor_loss	td3.py
evorl/workflows/rl_workflow.py	On‑policy workflow that aggregates losses using config.loss_weights	rl_workflow.py
evorl/distributed/gradients.py	agent_gradient_update – the optimizer wrapper that consumes the loss dict	gradients.py
evorl/utils/running_statistics.py	Optional observation normalizer used by many agents	running_statistics.py

Summary

• Extend the Agent class: Create a subclass of PPOAgent, TD3Agent, or another base agent and override the appropriate loss method (loss for on-policy, critic_loss/actor_loss for off-policy).
• Return a LossDict: Your method must return a dictionary mapping string keys to JAX arrays, following the LossDict type defined in evorl/types.py.
• Configure weights: Add your custom loss key to config.loss_weights so the workflow in evorl/workflows/rl_workflow.py can scale it before gradient computation.
• Automatic gradient updates: The agent_gradient_update function in evorl/distributed/gradients.py automatically differentiates all entries in the merged loss dictionary, requiring no manual gradient handling.

Frequently Asked Questions

What is the LossDict type in EvoRL?

LossDict is a type alias defined in evorl/types.py representing Mapping[str, chex.Array]. It is a dictionary where keys are loss term names (like "actor_loss" or "l2_reg") and values are JAX scalar arrays. The workflow expects this format so it can apply per-loss weights before summing into a single scalar for backpropagation.

How do I weight multiple custom losses differently?

Add each custom loss key to the loss_weights configuration dictionary with its own coefficient. For example, if you have l2_reg and smoothness losses, set config.loss_weights = {"actor_loss": 1.0, "l2_reg": 0.001, "smoothness": 0.01}. The workflow in evorl/workflows/rl_workflow.py multiplies each loss value by its corresponding weight before the optimizer step.

Can I modify existing loss methods instead of adding new ones?

Yes. You can override the base loss, critic_loss, or actor_loss methods in your subclass, call super() to obtain the original LossDict, and then modify or add entries before returning. This pattern is shown in the L2 regularization and smoothness examples, where the base loss is preserved and new terms are appended.

Where does the actual gradient computation happen?

The gradient computation occurs in agent_gradient_update within evorl/distributed/gradients.py. This function receives the merged loss scalar (after weights are applied), computes gradients with respect to the agent parameters using JAX's autodiff, and applies updates via the optimizer. As long as your custom loss is part of the LossDict returned by the agent, it is automatically included in this gradient update without requiring manual jax.grad calls.

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