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

> Learn to add custom loss functions in EvoRL for gradient-based updates. Implement new loss methods in your Agent subclass and register weights in your config for seamless integration.

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

---

**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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.

```python
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:

```yaml
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`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) and [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.

```python

# 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:

```yaml

# 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`](https://github.com/emi-group/evorl/blob/main/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.

```python

# 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:

```yaml

# 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:

```python
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`](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) | Base `Agent`, `AgentState`, `LossFn` definitions | [agent.py](https://github.com/emi-group/evorl/blob/main/evorl/agent.py) |
| [`evorl/types.py`](https://github.com/emi-group/evorl/blob/main/evorl/types.py) | `LossDict`, `AgentState`, `PyTreeData` utilities | [types.py](https://github.com/emi-group/evorl/blob/main/evorl/types.py) |
| [`evorl/algorithms/ppo.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) | Reference PPO implementation with a `loss` method | [ppo.py](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/ppo.py) |
| [`evorl/algorithms/td3.py`](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) | Reference TD3 implementation with `critic_loss` & `actor_loss` | [td3.py](https://github.com/emi-group/evorl/blob/main/evorl/algorithms/td3.py) |
| [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) | On‑policy workflow that aggregates losses using `config.loss_weights` | [rl_workflow.py](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) |
| [`evorl/distributed/gradients.py`](https://github.com/emi-group/evorl/blob/main/evorl/distributed/gradients.py) | `agent_gradient_update` – the optimizer wrapper that consumes the loss dict | [gradients.py](https://github.com/emi-group/evorl/blob/main/evorl/distributed/gradients.py) |
| [`evorl/utils/running_statistics.py`](https://github.com/emi-group/evorl/blob/main/evorl/utils/running_statistics.py) | Optional observation normalizer used by many agents | [running_statistics.py](https://github.com/emi-group/evorl/blob/main/evorl/utils/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`](https://github.com/emi-group/evorl/blob/main/evorl/types.py).
- **Configure weights**: Add your custom loss key to `config.loss_weights` so the workflow in [`evorl/workflows/rl_workflow.py`](https://github.com/emi-group/evorl/blob/main/evorl/workflows/rl_workflow.py) can scale it before gradient computation.
- **Automatic gradient updates**: The `agent_gradient_update` function in [`evorl/distributed/gradients.py`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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`](https://github.com/emi-group/evorl/blob/main/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.