How to Implement Off-Policy Algorithms (SAC, TD3, DDPG) in EvoRL: A Complete Guide

Implement SAC, TD3, and DDPG in EvoRL by creating Agent classes that define network initialization, action computation, and loss functions, then wrapping them with OffPolicyWorkflowTemplate for replay buffer management and training loops.

EvoRL is a JAX-based framework that unifies evolutionary and reinforcement learning algorithms. To implement off-policy algorithms such as SAC, TD3, and DDPG in EvoRL, you follow a consistent pattern: define an Agent class that implements the algorithm-specific logic, then use the generic OffPolicyWorkflowTemplate to handle environment interaction, replay buffer storage, and distributed training across multiple devices.

Core Architecture: Agent and Workflow Separation

EvoRL separates algorithm logic from training infrastructure. Every off-policy algorithm requires two components:

1. Agent Class – Implements init, compute_actions, evaluate_actions, actor_loss, and critic_loss (plus alpha_loss for SAC). This class holds the policy and value networks, observation normalizers, and optimizer states.
2. Workflow Class – Inherits from OffPolicyWorkflowTemplate in evorl/algorithms/offpolicy_utils.py. It constructs the environment, replay buffer, evaluator, and links them to the agent.

This separation allows you to implement a new off-policy algorithm by only writing the agent logic, reusing the entire training scaffolding.

SAC, TD3, and DDPG Implementation Details

Each algorithm resides in its own module under evorl/algorithms/, following identical structural patterns but implementing distinct loss functions and action sampling strategies.

SAC Agent Implementation (evorl/algorithms/sac.py)

The SACAgent class implements Soft Actor-Critic with automatic entropy tuning. Key features include:

• Network Parameters: Uses SACNetworkParams (a PyTreeDict) containing actor_params, critic_params, and target_critic_params.
• Action Computation: compute_actions samples from a tanh-normal distribution using get_tanh_norm_dist, applying the reparameterization trick for differentiability.
• Loss Functions:
  • critic_loss: Computes mean-squared error against soft Q-targets, using the minimum of two Q-values for stability.
  • actor_loss: Maximizes the expected Q-value minus the entropy bonus (alpha * log_prob).
  • alpha_loss: Adjusts the temperature parameter to match the target entropy.

TD3 Agent Implementation (evorl/algorithms/td3.py)

The TD3Agent implements Twin Delayed Deep Deterministic Policy Gradient with target policy smoothing:

• Deterministic Policy: Unlike SAC, TD3 outputs deterministic actions via the actor network. Exploration is handled externally by adding Gaussian noise (exploration_epsilon) in compute_actions.
• Twin Critics: Maintains two Q-networks (critic_params and target_critic_params for both) and uses the minimum Q-value for target computation to reduce overestimation bias.
• Target Policy Smoothing: Adds clipped noise to target actions during critic updates (policy_noise and clip_policy_noise parameters).
• Delayed Updates: The actor_loss is only computed every d updates (configured in the workflow), using either the first critic or the minimum of both critics based on critics_in_actor_loss setting.

DDPG Agent Implementation (evorl/algorithms/ddpg.py)

The DDPGAgent provides the foundational deterministic policy gradient algorithm:

• Single Critic: Unlike TD3, DDPG uses a single Q-network (critic_params and target_critic_params), making it simpler but more prone to overestimation.
• Deterministic Actions: Similar to TD3, actions are deterministic with optional exploration noise.
• Simple Target Updates: Uses standard soft target updates (soft_target_update) without the additional smoothing or twin critic logic found in TD3.

Network Builders and Initialization

Each algorithm provides a factory function to construct the necessary networks:

• make_mlp_sac_agent in evorl/algorithms/sac.py
• make_mlp_td3_agent in evorl/algorithms/td3.py
• make_mlp_ddpg_agent in evorl/algorithms/ddpg.py

These functions:

1. Create policy networks using make_policy_network from evorl/networks/
2. Create Q-networks using make_q_network from the same module
3. Optionally initialize observation normalizers via running_statistics.init_state from evorl/utils/running_statistics.py
4. Return the configured Agent instance ready for workflow integration

The networks are pure Flax modules defined in evorl/networks/linear.py and evorl/networks/layer_norm.py, supporting configurable hidden layer sizes and layer normalization.

The Off-Policy Training Loop

The OffPolicyWorkflowTemplate in evorl/algorithms/offpolicy_utils.py provides the generic scaffolding that powers SAC, TD3, and DDPG:

Replay Buffer Management

• Warm-up: Performs initial random rollouts to populate the replay buffer before training begins
• Storage: Uses flatten_rollout_trajectory from evorl/utils/rl_toolkits.py to process and store transitions
• Sampling: Draws minibatches for gradient updates

Training Step Execution

The step method (identical across all three algorithms):

1. Rollout: Collects trajectories using the current policy
2. Storage: Cleans and flattens trajectories into the replay buffer
3. Updates: Samples batches and runs agent_gradient_update, which handles:
   • JAX pmap for multi-device training
   • Optimizer state management
   • Parameter replacement via PyTree operations
4. Target Updates: Applies soft_target_update from evorl/utils/rl_toolkits.py to slowly update target network parameters

Distributed Training

The template automatically handles multi-device configuration, rescaling batch sizes and rollout lengths across JAX devices while maintaining consistent training logic.

Running Training with Hydra Configuration

EvoRL uses Hydra for configuration management. To train an off-policy algorithm:

# configs/sac.yaml

workflow_cls: evorl.algorithms.sac.SACWorkflow
env:
  env_name: HalfCheetah-v4
  env_type: brax
agent_network:
  num_critics: 2
  critic_hidden_layer_sizes: [256, 256]
  actor_hidden_layer_sizes: [256, 256]
optimizer:
  lr: 3e-4
  grad_clip_norm: 0.5
replay_buffer_capacity: 1000000
rollout_length: 256
num_updates_per_iter: 1
total_timesteps: 1000000

Launch training:

python -m evorl.scripts.train --config-name sac

The scripts/train.py entry point handles workflow instantiation from the config and launches the training loop.

Extending the Framework: Custom Off-Policy Agents

To implement a variant off-policy algorithm, extend the existing agent classes and register a new workflow:

# evorl/algorithms/my_td3.py

import jax.numpy as jnp
from evorl.algorithms.td3 import TD3Agent

class RegularizedTD3Agent(TD3Agent):
    """TD3 with L2 regularization on actor outputs."""
    
    def actor_loss(self, agent_state, sample_batch, key):
        loss_dict = super().actor_loss(agent_state, sample_batch, key)
        
        # Add penalty on raw action magnitudes

        raw_actions = self.actor_network.apply(
            agent_state.params.actor_params, 
            sample_batch.obs
        )
        l2_penalty = jnp.mean(jnp.square(raw_actions))
        
        return loss_dict.replace(
            actor_loss=loss_dict.actor_loss + 1e-4 * l2_penalty
        )

Then create the workflow:

# evorl/algorithms/my_workflow.py

from evorl.algorithms.offpolicy_utils import OffPolicyWorkflowTemplate
from .my_td3 import RegularizedTD3Agent

class RegularizedTD3Workflow(OffPolicyWorkflowTemplate):
    @classmethod
    def name(cls):
        return "RegularizedTD3"
    
    @classmethod
    def _build_from_config(cls, config):
        # Reuse TD3 workflow construction logic

        # ... (environment setup, network creation)

        agent = RegularizedTD3Agent(
            critic_network=critic_network,
            actor_network=actor_network,
            # ... other parameters from config

        )
        # ... (optimizer, evaluator, replay buffer setup)

        return cls(env, agent, optimizer, evaluator, replay_buffer, config)

Run the custom workflow:

python -m evorl.scripts.train workflow_cls=evorl.algorithms.my_workflow.RegularizedTD3Workflow

Summary

• EvoRL implements SAC, TD3, and DDPG through a unified architecture separating algorithm logic (Agent) from training infrastructure (Workflow).
• Agent classes (SACAgent, TD3Agent, DDPGAgent) in evorl/algorithms/ define network initialization, action computation, and loss functions specific to each algorithm.
• OffPolicyWorkflowTemplate in evorl/algorithms/offpolicy_utils.py provides the generic training loop, replay buffer management, multi-device sharding, and soft target updates.
• Network factories (make_mlp_sac_agent, etc.) construct Flax-based policy and Q-networks using modules from evorl/networks/.
• Hydra configuration in configs/ and the entry point scripts/train.py enable launching experiments without code changes.
• Extension pattern allows custom algorithms by inheriting from existing agents and workflows, requiring only the implementation of modified loss functions or action selection logic.

Frequently Asked Questions

What is the difference between SAC, TD3, and DDPG implementations in EvoRL?

SAC (evorl/algorithms/sac.py) uses a stochastic policy with tanh-normal distributions, twin critics with minimum value selection, and automatic entropy tuning via an alpha_loss function. TD3 (evorl/algorithms/td3.py) employs a deterministic policy with target policy smoothing, delayed actor updates, and clipped noise addition to target actions. DDPG (evorl/algorithms/ddpg.py) uses a simpler deterministic policy with a single critic network and standard soft target updates without the twin critic or smoothing mechanisms found in TD3.

How does EvoRL handle the replay buffer for off-policy algorithms?

EvoRL manages replay buffers through the OffPolicyWorkflowTemplate in evorl/algorithms/offpolicy_utils.py. The template automatically handles warm-up phases with random rollouts, trajectory flattening via flatten_rollout_trajectory from evorl/utils/rl_toolkits.py, and minibatch sampling during the training step. The workflow rescales buffer capacity and batch sizes automatically when running on multiple JAX devices, ensuring consistent sampling across distributed training runs.

Can I mix evolutionary algorithms with SAC or TD3 in EvoRL?

Yes, EvoRL is designed to unify evolutionary and reinforcement learning methods. While SAC, TD3, and DDPG follow the OffPolicyWorkflowTemplate, evolutionary algorithms typically use different workflow templates. You can combine these approaches by implementing custom workflows that alternate between evolutionary steps (using population-based methods) and off-policy RL updates, or by creating hybrid agents that inherit from both off-policy agents and evolutionary operators available in the evorl/algorithms/ directory.

Where are the network architectures defined for off-policy agents?

Network architectures are defined in evorl/networks/ using Flax modules. The policy networks use make_policy_network and Q-networks use make_q_network, both found in the networks subpackage (e.g., evorl/networks/linear.py and evorl/networks/layer_norm.py). Each algorithm provides factory functions (make_mlp_sac_agent, make_mlp_td3_agent, make_mlp_ddpg_agent) that instantiate these networks with algorithm-specific configurations, such as the number of critics (twin for SAC and TD3, single for DDPG) and hidden layer sizes.

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