ModuRL Guide

This guide is for Rust developers who want to build reinforcement learning programs with ModuRL and Candle. It starts with PPO on CartPole, then explains the library types behind that example, stochastic and deterministic actor-critic algorithms, and the value-based DQN and DDQN paths.

The guide assumes basic Rust and Cargo knowledge. It does not teach reinforcement learning or neural networks from first principles.

ModuRL is early, and API stability is not guaranteed. Start with Getting Started to run and assemble a PPO CartPole program. For off-policy stochastic actor training, read Soft Actor-Critic. For deterministic continuous-control training, read Deterministic Actor-Critic Training before choosing DDPG or TD3. For discrete-action value-based training, read Value-Based Training.

The README provides project status and the shortest repository-based commands. Rustdoc provides the precise contracts for public traits, structs, and builders.

Getting Started

This page builds a small PPO training program in a new Cargo binary crate.

What You Will Build

The program trains PPO on several CartPole environments at once. It uses an MLP to produce policy logits, a probabilistic policy model to sample actions, and another MLP to estimate state values.

Prerequisites

Install a Rust toolchain that supports edition 2024.

Create a Small Program

Create a new binary crate:

cargo new modurl-hello
cd modurl-hello

Start by adding the libraries used by this example:

[package]
name = "modurl-hello"
version = "0.1.0"
edition = "2024"

[dependencies]
modurl = "0.1"
modurl_gym = "0.1"
candle-core = "0.11.0"
candle-nn = "0.11.0"

The snippets below are consecutive pieces of src/main.rs. Add each one in order; the file compiles after the final Train section.

Imports

First, bring the Candle types, optimizer, ModuRL traits, and CartPole environment into scope:

use candle_core::{Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use modurl::prelude::*;
use modurl_gym::classic_control::cartpole::CartPoleV1;

modurl::prelude::* brings the common ModuRL traits and training types into scope. Agent and MultiGym are traits that make .learn(), .observation_space(), and .action_space() available.

A Few Terms

Before we build the program, it is worth separating two similar names; agent and actor.

In ModuRL, an Agent is the object that can act in an environment and learn from it. For this example, the agent is a PPOAgent. It contains the PPO rollout logic, loss calculations, optimizers, policy distribution, and neural networks.

The actor network is the source of the policy’s action scores. It produces logits, which are raw scores for each action. ProbabilisticPolicyModel<CategoricalDistribution> interprets those scores as a categorical policy: it can sample an action from them and later compute the log probability of that action. PPO needs both operations during training.

You will also see MLP below. MLP builds a dense feed-forward neural network: linear layers with an activation between them. We use one MLP for the actor network and one MLP for the critic network.

Choose a Device

Start main by choosing where tensors and models will live:

fn main() {
    let device = Device::Cpu;

The CPU device is the simplest first run. See Run on CUDA or Metal when the CPU version works.

Create Environments

PPO learns from batches of environment steps, so create several CartPole environments and wrap them as one vectorized environment:

    let envs = (0..4)
        .map(|_| CartPoleV1::builder().device(&device).build().unwrap())
        .collect::<Vec<_>>();
    let mut env = VectorizedGymWrapper::from(envs);

Each inner CartPoleV1 is a single environment. VectorizedGymWrapper stacks them so PPO can collect multiple transitions per step.

Read the Spaces

Ask the environment for its observation and action spaces before building the networks:

    let observation_space = env.observation_space();
    let action_space = env.action_space();

The observation space determines the input size for both networks. The action space determines how many action logits the actor network must produce.

Build the Actor Network

The actor network maps observations to action logits:

    let actor_var_map = VarMap::new();
    let actor_vb = VarBuilder::from_varmap(&actor_var_map, candle_core::DType::F32, &device);
    let actor_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(action_space.shape()[0])
        .vb(actor_vb)
        .activation(Tensor::tanh)
        .hidden_layer_sizes(vec![64, 64])
        .name("actor".to_string())
        .build()
        .expect("failed to build actor network");

The MLP here has two hidden dense layers of width 64. Its output size is the number of action logits. For CartPole, that means two logits: one for pushing left and one for pushing right. Later, CategoricalDistribution will interpret those logits as a discrete policy.

Build the Critic Network

The critic maps observations to one value estimate:

    let critic_var_map = VarMap::new();
    let critic_vb = VarBuilder::from_varmap(&critic_var_map, candle_core::DType::F32, &device);
    let critic_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(1)
        .vb(critic_vb)
        .activation(Tensor::tanh)
        .hidden_layer_sizes(vec![64, 64])
        .name("critic".to_string())
        .build()
        .expect("failed to build critic network");

This is another MLP with the same hidden-layer shape. PPO uses the critic to estimate how good each observed state is, so the critic network has the same input size as the actor network but only one output.

Create Optimizers

Give the actor network and critic network separate AdamW optimizers:

    let optimizer_config = ParamsAdamW {
        lr: 3e-4,
        ..Default::default()
    };

    let actor_optimizer = AdamW::new(actor_var_map.all_vars(), optimizer_config.clone())
        .expect("failed to build actor optimizer");
    let critic_optimizer = AdamW::new(critic_var_map.all_vars(), optimizer_config)
        .expect("failed to build critic optimizer");

Each optimizer receives the variables from the matching network’s VarMap.

Assemble PPO

Wrap the actor network in ProbabilisticPolicyModel so PPOAgent can sample actions and evaluate their log probabilities:

    let policy =
        ProbabilisticPolicyModel::<CategoricalDistribution>::new(actor_network);

    let network_info = PPONetworkInfo::Separate(
        SeparatePPONetwork::builder()
            .actor_network(policy)
            .critic_network(critic_network)
            .actor_optimizer(actor_optimizer)
            .critic_optimizer(critic_optimizer)
            .build(),
    );

policy is the probabilistic policy PPO trains. In this configuration, the policy model and value model are independent neural networks with independent optimizers.

Train

Finally, build the PPOAgent and run learning:

    let mut agent = PPOAgent::builder()
        .action_space(action_space)
        .network_info(network_info)
        .batch_size(1024)
        .mini_batch_size(64)
        .clip_range(ConstantSchedule::new(0.2))
        .training_horizon(10_000)
        .device(device)
        .build().unwrap();

    agent.learn(&mut env, 10_000).expect("PPO learning failed");
    println!("Training complete.");
}

batch_size controls how many transitions PPO collects before an update. mini_batch_size controls how those transitions are split during optimization. The clip range is the PPO policy-update bound. training_horizon defines the total number of environment transitions over which parameter schedules progress. This way, the schedules can span multiple learn calls.

Run the program with:

cargo run

When the training loop finishes, the program prints Training complete.. The default run collects 10,000 environment steps, so it does more than compile the program but is smaller than a longer experiment.

Complete File

After applying the pieces above, src/main.rs should look like this:

use candle_core::{Device, Tensor};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use modurl::prelude::*;
use modurl_gym::classic_control::cartpole::CartPoleV1;

fn main() {
    let device = Device::Cpu;

    let envs = (0..4)
        .map(|_| CartPoleV1::builder().device(&device).build().unwrap())
        .collect::<Vec<_>>();
    let mut env = VectorizedGymWrapper::from(envs);

    let observation_space = env.observation_space();
    let action_space = env.action_space();

    let actor_var_map = VarMap::new();
    let actor_vb = VarBuilder::from_varmap(&actor_var_map, candle_core::DType::F32, &device);
    let actor_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(action_space.shape()[0])
        .vb(actor_vb)
        .activation(Tensor::tanh)
        .hidden_layer_sizes(vec![64, 64])
        .name("actor".to_string())
        .build()
        .expect("failed to build actor network");

    let critic_var_map = VarMap::new();
    let critic_vb = VarBuilder::from_varmap(&critic_var_map, candle_core::DType::F32, &device);
    let critic_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(1)
        .vb(critic_vb)
        .activation(Tensor::tanh)
        .hidden_layer_sizes(vec![64, 64])
        .name("critic".to_string())
        .build()
        .expect("failed to build critic network");

    let optimizer_config = ParamsAdamW {
        lr: 3e-4,
        ..Default::default()
    };

    let actor_optimizer = AdamW::new(actor_var_map.all_vars(), optimizer_config.clone())
        .expect("failed to build actor optimizer");
    let critic_optimizer = AdamW::new(critic_var_map.all_vars(), optimizer_config)
        .expect("failed to build critic optimizer");

    let policy =
        ProbabilisticPolicyModel::<CategoricalDistribution>::new(actor_network);

    let network_info = PPONetworkInfo::Separate(
        SeparatePPONetwork::builder()
            .actor_network(policy)
            .critic_network(critic_network)
            .actor_optimizer(actor_optimizer)
            .critic_optimizer(critic_optimizer)
            .build(),
    );

    let mut agent = PPOAgent::builder()
        .action_space(action_space)
        .network_info(network_info)
        .batch_size(1024)
        .mini_batch_size(64)
        .clip_range(ConstantSchedule::new(0.2))
        .training_horizon(10_000)
        .device(device)
        .build().unwrap();

    agent.learn(&mut env, 10_000).expect("PPO learning failed");
    println!("Training complete.");
}

Where to Go Next

You now have a PPO training program with vectorized CartPole environments, a stochastic categorical policy, and a value model. Read Understand a PPO Training Run to learn what the example’s metrics show, or Use Vectorized Environments to work directly with batched environments.

Core Concepts

This chapter explains the names you will see in ModuRL examples. The goal is not to explain all of reinforcement learning at once. The goal is to make the types and builder arguments in the code feel less mysterious.

Agent

An Agent is the top-level object that interacts with an environment.

In ModuRL, an agent can do two things:

• act from observations
• learn from a vectorized environment

ModuRL provides agent types named after the algorithms they implement. PPOAgent implements PPO, SACAgent implements Soft Actor-Critic, DDPGAgent implements DDPG, TD3Agent implements TD3, DQNAgent implements DQN, and DDQNAgent implements Double DQN.

An agent coordinates the pieces needed for training: model modules, optimizers, schedules, buffers, distributions, and update logic.

Environments and Spaces

An environment is the world the agent interacts with.

A single environment implements the Gym interface. It can reset, accept an action, and return the next observation and reward.

Training usually needs more than one transition at a time, so ModuRL also uses vectorized environments. A MultiGym behaves like a batch of environments. When the agent sends one batch of actions, the vectorized environment steps each inner environment and returns a batch of results.

VectorizedGymWrapper is the simple wrapper used in the getting-started example. It turns several single environments into one vectorized environment. With the multithreading feature enabled, MultithreadedVectorizedGymWrapper can step the inner environments on worker threads. Applications can also implement the public MultiGym trait when they need a different way to manage or step a batch of environments.

Vectorized environments auto-reset each inner environment when it returns done or truncated. The states field in MultiGymStepInfo contains the next state to continue training from. If an inner environment ended, that next state is already the reset state for the next episode.

A Space describes what observations or actions look like.

For observations, the space tells you the input shape your networks need. For actions, the space tells you what kind of actions the environment accepts.

CartPole has a discrete action space. That is why the PPO example uses CategoricalDistribution: the policy chooses from a fixed set of actions.

Spaces also help convert policy outputs into environment actions. In PPO, the policy produces action information as tensors. The action space defines how those tensors map to valid actions for the environment.

Models and Policies

Models are neural network modules.

MLP builds a dense feed-forward network. Dense means each linear layer connects all of its inputs to all of its outputs. In the getting-started example, both networks are MLPs with two hidden layers.

MLP does not know about PPO by itself. It maps input tensors to output tensors. The agent and policy model decide what those outputs mean.

A policy is the rule used to choose actions. A stochastic policy samples from a distribution. PPO wraps its actor module in ProbabilisticPolicyModel<D>, where D defines how to interpret the actor’s output. D can be a distribution supplied by ModuRL or a user-defined type that implements the public Distribution trait. This gives PPO the operations it needs:

• sample an action representation
• compute the log probability of an action
• compute the policy entropy

For DQN and DDQN, an MLP is a Q-network. Its output has one value for each discrete action. The agent selects an action from those values with epsilon-greedy exploration instead of a probability distribution.

For DDPG and TD3, an actor MLP returns a continuous action directly. These algorithms add Gaussian noise while collecting experience, but the underlying policy is deterministic. Their scalar state-action critics receive an observation and action together and return one Q value.

Read Models, Policies, and Distributions for how tensors move between actors, distributions, and action spaces. Read Value-Based Training for discrete-action Q-networks, or Deterministic Actor-Critic Training for continuous actors and state-action critics.

Tensors and Devices

ModuRL uses Candle tensors.

A tensor has a shape, data type, and device. The shape must match what the model or environment expects. The device says where the tensor lives, such as CPU, CUDA, or Metal.

Most examples start with:

let device = Device::Cpu;

The important rule is that tensors and models used together should live on the same device. If a model is on the CPU, the observations passed to it should also be on the CPU.

Training Configuration

An optimizer updates model parameters.

In PPO with separate networks, the two model modules can have separate optimizers. The optimizer receives variables from the matching VarMap.

A schedule is a value that changes during training. For example, ConstantSchedule keeps the PPO clip range fixed. LinearSchedule can move a value from one number to another over training progress.

In the CartPole Example

The CartPole PPO example configures a probabilistic policy that wraps a model which produces policy scores. It also configures a separate model that estimates values. PPOAgent uses the policy and value model while it collects experience and updates their parameters.

These are roles in the PPO configuration, not separate library architecture types. The public types in the example are MLP, ProbabilisticPolicyModel<CategoricalDistribution>, and PPOAgent.

Read Understand a PPO Training Run to inspect the results of the example, or Use Vectorized Environments to work with batched environment steps directly.

Models, Policies, and Distributions

A module returns a tensor. An environment expects an action. The agent decides what the tensor means and how it becomes that action.

Two possible paths are:

direct: model -> action representation -> action space -> environment
probabilistic:
    probabilistic policy -> action representation -> action space -> environment

The algorithm determines whether the action path uses a distribution. A model can return an action representation that an action space understands directly. DDPG and TD3 use this direct path for their deterministic actors, then add Gaussian noise inside the agent during collection. An agent can also insert other selection logic. DQN and DDQN, for example, choose from Q-values with epsilon-greedy selection.

On the probabilistic path, the policy uses a model and a distribution together. The distribution is part of the policy rather than a standalone stage.

Models Produce Tensors

A model transforms an input tensor into an output tensor. MLP supplies dense layers and activation functions, but it does not assign meaning to its output.

Let B be the batch size and I the number of input features. A dense MLP receives [B, I]. The features may be observations, the output of an earlier module, or other values prepared by an agent. Other architectures may use different input shapes.

The component that receives a model’s output determines its shape and meaning. A critic, for example, returns [B, 1] state-value estimates. A Q-network returns action values. A probabilistic actor returns parameters for a distribution.

Probabilistic Policies Use a Distribution

ProbabilisticPolicyModel<D> owns the wrapped module and implements ProbabilisticPolicy. Its sample and mode operations pass the module output to D::from_outputs, then call the corresponding distribution operation. Its log_prob_and_entropy operation calls D::dist_eval when an algorithm needs those values.

Distribution is a public trait. ModuRL currently supplies CategoricalDistribution and GaussianDistribution, but applications can add their own implementations:

let policy =
    ProbabilisticPolicyModel::<MyDistribution>::new(actor);

A custom implementation provides from_outputs, sample, mode, dist_eval, and an associated Error type. It must document its model-output layout and returned tensor shapes. Its action representation must also match the chosen Space. The Rust type system does not check these tensor shapes.

Spaces Produce Environment Actions

Space::tensor_from_neurons converts an action representation into the tensor passed to the environment.

Discrete selects the index of the largest component. BoxSpace clamps each component to its lower and upper bounds. PPO retains the original sample for log-probability calculations while sending the converted action to the environment.

Built-In Distributions

Categorical Distribution

Let C be the number of discrete choices. CategoricalDistribution expects one logit for each choice:

The logits are unnormalized scores. Sampling adds an independent random perturbation called Gumbel noise to each score. Taking the largest perturbed score samples choices according to the logits. Discrete then selects that score’s index.

CartPole has two actions, so the getting-started actor returns two logits per observation.

Gaussian Distribution

Let A be the number of components in a one-dimensional continuous action space with shape [A]. GaussianDistribution expects two values for each component. Along dimension 1, all A means come first, followed by all A log standard deviations:

[mean_0, ..., mean_(A-1), log_std_0, ..., log_std_(A-1)]

GaussianDistribution applies exp to the log standard deviations before sampling. Neither half of the model output contains log probabilities. Each action component uses an independent Gaussian, and dist_eval sums the component log probabilities and entropies into one value per batch row.

The actor can return the complete [B, 2 * A] tensor itself. It can also combine state-dependent means with a separate trainable log_std, as the MuJoCo PPO example does.

Read Getting Started for a categorical policy, Soft Actor-Critic for policies that expose exact or sampled action expectations, Deterministic Actor-Critic Training for direct continuous actors, Value-Based Training for action selection without a distribution, and crates/examples/examples/ppo_mujoco.rs for a Gaussian policy.

PPO

PPOAgent is a ModuRL agent that implements Proximal Policy Optimization (PPO). It collects transitions from a MultiGym, then updates a stochastic policy and a value model from that experience.

The getting-started program uses the separate-network PPO configuration. It gives the policy model and the value model their own MLP, VarMap, and Adam optimizer. PPONetworkInfo::Shared is also available when a configuration needs one shared model followed by separate policy and value heads.

Start with Getting Started for the complete CartPole program. Then read Understand a PPO Training Run before changing its configuration.

The continuous-control example supports ant, half-cheetah, hopper, and walker2d. Normally enable one environment feature; rendering may be added to that feature to open a viewer. Cargo features are additive, so builds that enable several environment features use the priority ant, half-cheetah, hopper, then walker2d. PPO writes TensorBoard events below runs/ppo_mujoco/ and also displays terminal graphs.

PPO can use any compatible Distribution implementation. ModuRL currently supplies categorical and Gaussian distributions, and applications can define their own. Models, Policies, and Distributions explains the extension point, the built-in tensor layouts, and how sampled representations become environment actions.

For value-based agents in discrete action spaces, read Value-Based Training.

Understand a PPO Training Run

The Getting Started program confirms that training completed, but it does not record metrics. Add a PPOLogger when you need to compare runs or diagnose a configuration. PPOLogger::log receives update metrics, while PPOLogger::log_collection receives the newest environment rewards and any episodes completed during that vectorized step.

Each PPOLogEntry contains tensors for one optimization minibatch. Aggregate them before graphing or comparing them across a run. The graphs are useful for comparing or diagnosing runs.

Start With Episode Performance

CartPole gives a reward of one for each step that keeps the pole balanced. The example graphs completed episode returns and lengths, so longer episodes are the clearest sign that the policy is improving.

PPOCollectionLogEntry::completed_episodes contains one PPOEpisodeLogEntry for each environment that terminated or truncated on that vectorized step. Each episode entry reports its environment index, return, length, ending flags, and collection timestep. Partial episodes carry across PPO rollout boundaries and repeated learn calls.

PPOCollectionLogEntry::infos contains the typed step metadata from each vectorized environment in the same environment order as collection_rewards. Wrappers can add logging-only values there without changing the reward used for training. For example, RecordRawRewardGym records rewards before outer reward wrappers normalize or clip them.

Read the Other Metrics Together

The logger exposes the following values through PPOLogEntry:

• actor_loss: the policy objective used for the update
• critic_loss: the value-model objective used for the update
• entropy: how spread out the policy’s action choices are
• kl_divergence: how far the updated policy moved from the policy that collected the rollout
• explained_variance: how well the value model accounts for return variation
• rewards: rewards from the collected transitions

Loss values do not have one universally good target. They are most useful when compared with episode length and with earlier runs using the same configuration.

Entropy commonly falls as the policy becomes more certain. A rapid collapse can mean the policy stopped exploring too early. KL divergence shows the size of the policy change. Large, erratic changes suggest that the update settings may be too aggressive.

Explained variance shows whether the value model predicts returns better than a constant average prediction: values nearer one are better, but early updates can be noisy.

Change One Setting at a Time

The CartPole example sets these PPO builder values:

.batch_size(2048)
.mini_batch_size(64)
.ent_coef(0.005)
.clip_range(ConstantSchedule::new(0.2))
.num_epochs(10)
.training_horizon(120_000)

Use one changed value per experiment. Keep the device, environment count, model shape, and remaining PPO settings fixed while comparing the graph. That makes a change in episode length or KL divergence easier to attribute to its cause.

Next, read Use Vectorized Environments to see how PPO receives batches of transitions.

A2C

A2CAgent implements synchronous Advantage Actor-Critic by configuring an inner PPO agent. The wrapper disables policy and value clipping, performs one optimization epoch, and uses the complete rollout as one minibatch. These algorithm-defining settings cannot be changed through the A2C builder.

The remaining defaults follow Stable-Baselines3 A2C: gamma is 0.99, gae_lambda is 1.0, advantages and returns are not normalized, vf_coef is 0.5, ent_coef is 0.0, and gradient clipping is 0.5. The application still supplies the optimizer as part of its network configuration.

Use SharedA2CNetwork when the actor and critic share a model, or SeparateA2CNetwork when they have independent models and optimizers. Wrap the result with A2CNetworkInfo::shared or A2CNetworkInfo::separate before building the agent. A2C also provides typed errors, log entries, and the A2CLogger trait. Pass an A2C logger directly to the agent’s logging_info builder field.

batch_size is the total number of transitions in one rollout, across every vectorized environment. Choose a value divisible by the environment count so the agent performs exactly one full-batch update.

The MuJoCo example uses a separate Gaussian actor and critic. Run it with one of the supported environment features:

cargo run -p examples --example a2c_mujoco --features ant
cargo run -p examples --example a2c_mujoco --features half-cheetah
cargo run -p examples --example a2c_mujoco --features hopper
cargo run -p examples --example a2c_mujoco --features walker2d

Choose one command for the environment you want to train. Read PPO for the underlying on-policy training model.

Add rendering alongside the environment feature to open a viewer, for example --features ant,rendering. A2C writes TensorBoard events below runs/a2c_mujoco/ and also displays terminal graphs. If several environment features are enabled by an additive feature build, the example uses ant, half-cheetah, hopper, then walker2d as its selection priority.

Soft Actor-Critic

Soft Actor-Critic (SAC) trains one stochastic actor and two Q-value critics from replayed transitions. The algorithm traditionally known as SAC almost always uses exactly two critics. ModuRL accepts any nonempty critic ensemble, but use two unless you are intentionally studying another aggregation scheme.

SAC is primarily a continuous-control algorithm. This chapter follows the continuous MuJoCo example and explains the actor, critic, entropy, replay, and device choices you need to adapt it. You should already be familiar with MultiGym, Candle VarMaps, and probabilistic policies from Core Concepts.

Run Continuous SAC

The example supports Ant, HalfCheetah, Hopper, and Walker2d:

cargo run --release -p examples --example sac_mujoco --features ant
cargo run --release -p examples --example sac_mujoco --features half-cheetah
cargo run --release -p examples --example sac_mujoco --features hopper
cargo run --release -p examples --example sac_mujoco --features walker2d

Normally enable one environment feature; add rendering to it to open a viewer. Additive builds with several environment features select ant, half-cheetah, hopper, then walker2d. The program selects CUDA when it is available and otherwise uses the CPU. It trains for one million collected transitions, then displays terminal graphs for episode performance and SAC diagnostics.

The complete source is crates/examples/examples/sac_mujoco.rs. It contains four parts:

1. A Gaussian parameter module and squashed distribution for bounded actions.
2. Two independently optimized online/target critic pairs.
3. Automatic entropy tuning.
4. An SACAgent that owns replay and runs collection and optimization.

The agent stores every collected transition in replay, including transitions from the initial random-action phase. Once collection_timestep reaches training_start, each new transition triggers one replay optimization step. This gives the implementation an update-to-data ratio of one. A vectorized step with N environments stores N transitions and, after the threshold, triggers N updates.

Build the Continuous Policy

GaussianModule returns distribution parameters rather than an environment action. ProbabilisticPolicyModel combines the module with a distribution to form the actor’s policy. For an action with A components, GaussianDistribution expects A means followed by A log standard deviations.

The module uses state-dependent means and one trainable log standard deviation per action component:

struct GaussianModule {
    mean: MLP,
    log_std: Tensor,
}

impl Module for GaussianModule {
    fn forward(&self, observations: &Tensor) -> candle_core::Result<Tensor> {
        let mean = self.mean.forward(observations)?;
        let log_std = self
            .log_std
            .broadcast_as(mean.shape())?
            .clamp(-20.0, 2.0)?;
        Tensor::cat(&[mean, log_std], 1)
    }
}

Build the distribution with the environment’s action shape:

let action_shape = action_space.shape();
let action_size = action_shape.iter().product();

let actor_vars = VarMap::new();
let actor_vb = VarBuilder::from_varmap(
    &actor_vars,
    DType::F32,
    &device,
);
let mean = sac_mlp(
    actor_vb.pp("mean"),
    observation_size,
    action_size,
    0.01,
    "mlp",
)?;
let log_std = actor_vb.get_with_hints(
    (1, action_size),
    "log_std",
    Init::Const(0.0),
)?;

let distribution = TransformedDistribution::new(
    GaussianDistribution::new(action_shape)?,
    TanhTransform,
);
let policy = ProbabilisticPolicyModel::with_distribution(
    GaussianModule { mean, log_std },
    distribution,
);

TanhTransform keeps policy candidates between -1 and 1, matching the included MuJoCo environments’ action bounds. The transform also includes its Jacobian correction in action log probabilities.

Build Two Critics

Each SACCritic contains an online network, a target network, the corresponding parameter maps, and an optimizer for the online parameters. The constructor copies the online parameters into the target parameter map. Both maps must contain the same parameter names.

A continuous SAC critic receives a state and an action and returns one Q-value. ScalarStateActionCritic concatenates each state/action pair before calling the wrapped module:

let critic = SACCritic::builder()
    .online_network(ScalarStateActionCritic::new(
        online_network,
    ))
    .target_network(ScalarStateActionCritic::new(
        target_network,
    ))
    .online_vars(&online_vars)
    .target_vars(&mut target_vars)
    .optimizer(critic_optimizer)
    .build()?;

Create this pair twice with separate online parameters, target parameters, and optimizers. The agent aggregates both critics when it builds targets and updates the actor.

The default aggregation mode is SACCriticAggregationMode::Min. This is the traditional SAC choice: using the lower estimate helps reduce optimistic Q-value errors.

ModuRL also exposes Mean, Median, and Max for experiments with different critic ensembles:

Configure Entropy

The entropy coefficient, usually written as alpha, balances expected return against policy entropy. Larger values put more weight on keeping the policy’s action distribution stochastic. Use automatic entropy tuning for the standard SAC workflow:

let log_alpha = Var::from_vec(vec![0.0f32], (), &device)?;
let alpha_optimizer = AdamW::new(
    vec![log_alpha.clone()],
    optimizer_parameters.clone(),
)?;
let entropy = SACEntropyConfiguration::automatic(log_alpha, alpha_optimizer);

automatic asks the policy for its default target entropy. The built-in Gaussian policy uses the negative action-component count. An affine distribution transform adjusts that default for its scale.

Pass a ParameterSchedule when the desired entropy should change during the run:

let target_entropy = LinearSchedule::new(-2.0, -1.0);
let entropy = SACEntropyConfiguration::automatic_with_target_schedule(
    log_alpha,
    alpha_optimizer,
    target_entropy,
);

LinearSchedule moves the target entropy from -2.0 to -1.0. Choose values appropriate for the environment’s action dimension rather than copying these illustrative values.

Use a fixed coefficient when you need a known constant or want to reproduce a configuration without an entropy optimizer:

let entropy = SACEntropyConfiguration::<AdamW>::fixed(0.2);

The optimizer type annotation is required because the fixed variant does not contain an optimizer value. A fixed coefficient must be finite and non-negative.

Assemble the Agent

The continuous example assembles the pieces as follows:

let mut agent = SACAgent::builder()
    .policy(policy)
    .actor_optimizer(actor_optimizer)
    .critics(vec![critic_1, critic_2])
    .entropy_configuration(entropy)
    .action_space(action_space)
    .observation_space(observation_space)
    .aggregation_mode(SACCriticAggregationMode::Min)
    .training_horizon(total_timesteps)
    .replay_storage_config(ReplayStorageConfig::new(
        ReplayDeviceStrategy::OneDevice(device),
    ))
    .build()?;

agent.learn(&mut env, total_timesteps)?;

training_horizon defines how many collected transitions parameter schedules take to reach their final values. Schedule progress continues across learn calls and stops at the end of the horizon.

Important defaults and constraints are:

samples controls the Monte Carlo estimate used by continuous distributions. Larger values cost more critic evaluations but use more candidates to estimate the actor and target expectations.

Q-value clipping is an optional experimental stabilization:

.q_value_clip(0.5)

The value is a positive bound on a critic’s update around its target-network estimate. Its useful scale depends on environment rewards, so it remains opt-in.

Implement a Custom Critic Adapter

Use ScalarStateActionCritic unless your critic has a different input or output layout. A custom SACCriticNetwork must follow these contracts:

actor_values must preserve gradients through differentiable candidate actions while excluding gradients to critic parameters.

Discrete SAC Is an API Variant

SACAgent can also work with categorical policies and DiscreteVectorHeadCritic. Experimental discrete-SAC configurations can select mean critic aggregation with .aggregation_mode(SACCriticAggregationMode::Mean) and add an entropy-drift penalty with .entropy_change_penalty(0.5). It is not the canonical example because traditional SAC is a continuous-control algorithm.

Read Understand an SAC Training Run to add a logger and interpret its metrics. Read Run on CUDA or Metal when replay storage should remain on the CPU while optimization runs on an accelerator.

Understand an SAC Training Run

The SAC examples confirm that the actor, critics, replay buffer, and entropy configuration can train together. Add a SACLogger when you need to compare runs or diagnose one of those components.

SACLogger::log_update receives metrics from replay optimization. SACLogger::log_collection receives rewards and completed episodes from the current environment step. Keep these streams separate: an update describes a sampled replay batch, while collection describes the policy’s newest behavior.

Add a Logger

This logger prints occasional update values and every completed episode:

use modurl::prelude::*;

struct ConsoleLogger;

impl<I> SACLogger<I> for ConsoleLogger {
    fn log_update(&mut self, entry: &SACLogEntry) {
        if entry.update_index % 1_000 != 0 {
            return;
        }

        let actor_loss = entry
            .actor_loss
            .to_scalar::<f32>()
            .expect("actor loss must be scalar");
        let alpha = entry
            .alpha
            .to_scalar::<f32>()
            .expect("alpha must be scalar");

        println!(
            "step={} update={} actor_loss={actor_loss:.4} alpha={alpha:.4}",
            entry.collection_timestep,
            entry.update_index,
        );
    }

    fn log_collection(&mut self, entry: &SACCollectionLogEntry<I>) {
        for episode in &entry.completed_episodes {
            println!(
                "step={} env={} return={} length={} terminated={} truncated={}",
                episode.collection_timestep,
                episode.environment_index,
                episode.episode_return,
                episode.episode_length,
                episode.terminated,
                episode.truncated,
            );
        }
    }
}

Pass a mutable reference while building the agent:

let mut logger = ConsoleLogger;

let mut agent = SACAgent::builder()
    // Keep the remaining SAC configuration unchanged.
    .logger(&mut logger)
    .build()?;

agent.learn(&mut env, total_timesteps)?;

The repository’s sac_mujoco example uses SACGrapher to aggregate these callbacks into terminal graphs.

Start With Episode Performance

SACCollectionLogEntry::completed_episodes contains one SACEpisodeLogEntry for each environment that ended during the latest vectorized step. Each entry reports:

• environment_index: which inner environment completed
• episode_return: the sum of that episode’s rewards
• episode_length: the number of environment steps
• terminated: whether the environment reached a terminal state
• truncated: whether the episode was truncated
• collection_timestep: the total collected-transition count at completion

Episode return is the clearest measure of task performance. Compare it across several episodes rather than treating one episode as a trend. Episode length is useful when termination timing is meaningful for the environment.

SACCollectionLogEntry also contains:

• collection_rewards: one reward per inner environment from the newest step
• infos: the corresponding typed environment metadata
• collection_timestep: the total number of collected transitions
• replay_len: the number of entries currently held in replay

log_collection runs once per vectorized environment step. With N inner environments, one callback normally advances collection_timestep by N.

Read Update Metrics as a Group

After training_start, SAC performs an optimization step for each collected transition. Each call to log_update describes one sampled replay batch. update_index is the zero-based optimization count, while collection_timestep identifies the transition that triggered the update.

Let B be replay batch size and K the number of policy candidates. For a categorical policy, K is the action count. For a sampled continuous policy, K is the configured samples value.

Losses do not have a universal target value. Reward scale, model architecture, and entropy configuration all change their magnitude. Compare them with episode return and with earlier runs using the same environment and reward handling.

Critic loss measures disagreement with the soft Bellman targets. A persistent increase alongside unstable Q-values can indicate an aggressive learning rate or reward scale. A small critic loss alone does not prove that the policy is good.

Actor loss combines expected Q-values and the entropy term. It can be negative, and its raw value is not a task score. Use it to spot abrupt changes rather than to rank policies.

Alpha controls the strength of the entropy term. A larger alpha places more weight on uncertain actions. With automatic tuning, compare alpha, policy entropy, and target entropy together.

The expected policy entropy for one update is:

-mean(sum(policy_weights * policy_log_probabilities, candidate_axis))

Entropy commonly falls when a policy becomes more certain. Whether that is healthy depends on the target entropy and episode performance. A rapid collapse with poor returns suggests that exploration ended too soon.

entropy_change_loss appears only when the stabilization configuration enables the replay-to-current entropy penalty. alpha_loss and target_entropy appear only with automatic entropy tuning.

Compare Collection and Replay Values

collection_rewards show the newest environment behavior. replay_rewards come from a randomly sampled historical batch. They need not move together at each callback.

The same distinction applies to time:

• collection_timestep counts transitions gathered from environments.
• update_index counts replay optimization steps.
• replay_len grows until it reaches replay capacity.

Graph completed episode returns against collection_timestep. Graph update losses, alpha, expected policy entropy, and mean Bellman targets at a lower frequency if logging every optimization step is too expensive.

Return to Soft Actor-Critic to change entropy, critic aggregation, or stabilization. Read Run on CUDA or Metal to separate replay storage from optimization.

Deterministic Actor-Critic Training

Use DDPG or TD3 when the environment has a continuous action space and you want a deterministic policy trained from replay. Both algorithms use an actor that returns one action directly and a critic that estimates the value of a state-action pair.

This chapter explains the pieces shared by DDPGAgent and TD3Agent. Read DDPG for the smaller, single-critic configuration. Read TD3 when you want twin critics, target-policy smoothing, and delayed actor updates.

Choose DDPG or TD3

DDPG and TD3 share the same collection and replay loop. TD3 changes how the agent calculates targets and when it updates the actor:

DDPG is the direct deterministic actor-critic baseline. TD3 adds safeguards against overestimated Q values and an actor that exploits narrow errors in the critic. Start with TD3 when those safeguards are appropriate. Choose DDPG when you need the single-critic algorithm or a simpler baseline for comparison.

ModuRL also lets TD3Agent use any nonempty critic ensemble and explicit aggregation modes. Those settings are experimental variants, not canonical TD3.

Follow the Shared Training Loop

Both agents collect and train in the same order:

1. Before training_start, sample actions uniformly from the BoxSpace.
2. After training_start, run the online actor and add Gaussian exploration noise.
3. Clamp the resulting actions to the action-space bounds and store each transition in replay.
4. On collection timesteps selected by update_frequency, sample a replay batch and update every online critic.
5. Update the online actor at the algorithm’s actor-update interval.
6. Polyak-update the target actors and critics whenever the actor is updated.

Collection timesteps count transitions, not calls to the vectorized environment. A step with N environments adds N transitions. Each transition whose global index is a multiple of update_frequency causes one replay update after the warm-up threshold.

training_horizon records the intended number of collected transitions and tracks progress across calls to learn. DDPG and TD3 do not currently expose scheduled hyperparameters, but the shared progress counter still determines the global collection timestep used by training and logging.

Match the Actor Contract

The online and target actors are ordinary Candle Modules. For observations shaped [batch, ...observation_shape], each actor must return values shaped [batch, ...action_shape].

The supplied BoxSpace clamps those values to its bounds. The MuJoCo examples use tanh as the actor’s output activation and bounds of -1.0..=1.0:

let actor = MLP::builder()
    .input_size(observation_size)
    .output_size(action_size)
    .vb(actor_vb)
    .hidden_layer_sizes(vec![64, 64])
    .activation(Tensor::relu)
    .output_activation(Tensor::tanh)
    .name("actor".to_owned())
    .build()?;

Build the online and target actors with separate VarMaps and identical parameter names. Agent construction copies the online parameters into the target map. Construction returns DeterministicActorCriticError::ActorParameterMapMismatch if the names differ.

The exploration-noise standard deviation is measured in the actor’s output units. ModuRL adds that noise before BoxSpace clamps the action. Use act_deterministic when evaluating a trained policy without exploration noise.

Match the Critic Contract

DeterministicCritic is the deterministic-agent name for SACCritic. Each critic owns:

• an online state-action network
• a target state-action network
• separate online and target VarMaps
• an optimizer for the online parameters

A scalar critic receives a state and action and returns one Q value. Wrap a module whose input width is observation_size + action_size with ScalarStateActionCritic:

let critic = DeterministicCritic::builder()
    .online_network(ScalarStateActionCritic::new(
        online_network,
    ))
    .target_network(ScalarStateActionCritic::new(
        target_network,
    ))
    .online_vars(&online_vars)
    .target_vars(&mut target_vars)
    .optimizer(critic_optimizer)
    .build()?;

As with the actors, the online and target critic parameter names must match. The constructor initializes the target parameters from the online parameters. Each TD3 critic needs its own networks, parameter maps, and optimizer.

Keep Replay and Models on the Intended Devices

Wrap ReplayDeviceStrategy::OneDevice(device) in ReplayStorageConfig to store replay and run optimization on the same device. A configuration using ReplayDeviceStrategy::Hybrid can keep detached replay transitions on one device, such as the CPU, and move sampled batches to the optimization device.

The strategy does not move models or environments. Build the environment, actors, critics, target networks, and optimizers on the optimization device. Read Run on CUDA or Metal before splitting replay storage from optimization.

The DDPG chapter uses these pieces to build the single-critic baseline. The TD3 chapter shows how to use two critics, target-policy noise, and delayed actor updates.

DDPG

Deep Deterministic Policy Gradient (DDPG) trains one deterministic actor and one Q-value critic from replayed transitions. This page shows how the complete MuJoCo example connects those networks to DDPGAgent.

You should already understand the actor, critic, target-network, and replay contracts in Deterministic Actor-Critic Training.

Run the DDPG Example

The example supports Ant, HalfCheetah, Hopper, and Walker2d:

cargo run --release -p examples --example ddpg_mujoco --features ant
cargo run --release -p examples --example ddpg_mujoco --features half-cheetah
cargo run --release -p examples --example ddpg_mujoco --features hopper
cargo run --release -p examples --example ddpg_mujoco --features walker2d

Normally enable one environment feature; add rendering to it to open a viewer. Additive builds with several environment features select ant, half-cheetah, hopper, then walker2d. The program selects CUDA when it is available and otherwise uses the CPU. It trains for one million collected transitions, then displays terminal graphs for optimization and episode metrics.

The complete source is crates/examples/examples/ddpg_mujoco.rs. It contains four parts:

1. Separate online and target actor networks.
2. One online/target critic pair.
3. A bounded continuous action space and replay-device strategy.
4. A DDPGAgent that owns collection and optimization.

actor and critic are private helper functions defined in that example, not functions provided by ModuRL. actor builds an actor MLP. critic builds the online and target critic networks and packages them as a DeterministicCritic.

Build the Online and Target Actors

Both actor networks must have the same architecture and parameter names. Only the online actor has an optimizer:

let online_actor_variables = VarMap::new();
let mut target_actor_variables = VarMap::new();

let online_actor = actor(
    &online_actor_variables,
    observation_size,
    action_size,
    &device,
)?;
let target_actor = actor(
    &target_actor_variables,
    observation_size,
    action_size,
    &device,
)?;
let actor_optimizer = AdamW::new(
    online_actor_variables.all_vars(),
    actor_optimizer_parameters,
)?;

DDPGAgent copies the online actor parameters to the target actor during construction. Later, it moves each target parameter toward the corresponding online parameter by tau after every replay update.

Build One Critic

DDPG requires exactly one DeterministicCritic. Its online network learns from replay. Its target network supplies the next-state Q estimate:

let critic = DeterministicCritic::builder()
    .online_network(ScalarStateActionCritic::new(
        online_critic,
    ))
    .target_network(ScalarStateActionCritic::new(
        target_critic,
    ))
    .online_vars(&online_critic_variables)
    .target_vars(&mut target_critic_variables)
    .optimizer(critic_optimizer)
    .build()?;

The actor learns to maximize the online critic’s mean Q estimate. Equivalently, the optimizer minimizes the negative mean Q value.

Assemble the Agent

The example passes all owned modules and borrowed parameter maps to the builder:

let mut agent = DDPGAgent::builder()
    .online_actor(online_actor)
    .target_actor(target_actor)
    .online_actor_vars(&online_actor_variables)
    .target_actor_vars(&mut target_actor_variables)
    .actor_optimizer(actor_optimizer)
    .critic(critic)
    .action_space(action_space)
    .observation_space(observation_space)
    .replay_storage_config(ReplayStorageConfig::new(
        ReplayDeviceStrategy::OneDevice(device),
    ))
    .gamma(0.99)
    .tau(0.005)
    .exploration_noise(0.1)
    .replay_capacity(1_000_000)
    .batch_size(256)
    .training_start(10_000)
    .training_horizon(TOTAL_TIMESTEPS)
    .logger(&mut grapher)
    .build()?;

agent.learn(&mut env, TOTAL_TIMESTEPS)?;

The action_space must be a BoxSpace. Its shape must match the actor output, and its bounds must match the environment’s accepted actions. The included MuJoCo environments use one-dimensional action vectors bounded by -1.0..=1.0; their vector length depends on the selected environment.

Important defaults and constraints are:

The example overrides training_start to collect 10,000 random transitions. Those transitions remain in replay. After warm-up, the agent adds exploration noise to online-actor actions and updates the actor, critic, and target networks at every selected replay update.

Evaluate Without Exploration Noise

Agent::act includes Gaussian exploration noise. Use act_deterministic for evaluation:

let actions = agent.act_deterministic(&observations)?;

The observations must be shaped [batch, ...observation_shape]. The returned actions are shaped [batch, ...action_shape] and are clamped to the configured BoxSpace.

Read Understand a Deterministic Actor-Critic Training Run to add a logger and interpret replay-update metrics. Read TD3 to add twin critics, target-policy smoothing, and delayed actor updates.

TD3

Twin Delayed Deep Deterministic Policy Gradient (TD3) keeps DDPG’s deterministic actor and replay loop, then adds three safeguards:

1. Two critics reduce optimistic target estimates.
2. Noise on target actions smooths the critic target.
3. Delayed actor and target-network updates give the critics more update steps.

This page starts from the network and replay contracts in Deterministic Actor-Critic Training and focuses on the TD3-specific choices.

Run the TD3 Example

The example supports Ant, HalfCheetah, Hopper, and Walker2d:

cargo run --release -p examples --example td3_mujoco --features ant
cargo run --release -p examples --example td3_mujoco --features half-cheetah
cargo run --release -p examples --example td3_mujoco --features hopper
cargo run --release -p examples --example td3_mujoco --features walker2d

Normally enable one environment feature; add rendering to it to open a viewer. Additive builds with several environment features select ant, half-cheetah, hopper, then walker2d. The program selects CUDA when it is available and otherwise uses the CPU. It trains for one million collected transitions, then displays terminal graphs for optimization and episode metrics.

The complete source is crates/examples/examples/td3_mujoco.rs. Its actor construction is the same as the DDPG example. The difference is the critic ensemble and the TD3 builder settings.

actor and critic are private helper functions defined in that example, not functions provided by ModuRL. actor builds an actor MLP. critic builds one pair of online and target critic networks and packages them as a DeterministicCritic.

Build Two Independent Critics

Canonical TD3 uses two critics. Build each critic with separate online and target networks, parameter maps, and optimizers:

let critic_1 = critic(
    &online_critic_variables_1,
    &mut target_critic_variables_1,
    observation_size,
    action_size,
    &optimizer_parameters,
    &device,
)?;
let critic_2 = critic(
    &online_critic_variables_2,
    &mut target_critic_variables_2,
    observation_size,
    action_size,
    &optimizer_parameters,
    &device,
)?;

Do not share a VarMap or optimizer between the critics. Their independent errors are what make the smaller of the two target estimates useful.

Configure the Three TD3 Safeguards

Pass both critics and the TD3-specific update settings:

let mut agent = TD3Agent::builder()
    .online_actor(online_actor)
    .target_actor(target_actor)
    .online_actor_vars(&online_actor_variables)
    .target_actor_vars(&mut target_actor_variables)
    .actor_optimizer(actor_optimizer)
    .critics(vec![critic_1, critic_2])
    .action_space(action_space)
    .observation_space(observation_space)
    .replay_storage_config(ReplayStorageConfig::new(
        ReplayDeviceStrategy::OneDevice(device),
    ))
    .gamma(0.99)
    .tau(0.005)
    .exploration_noise(0.1)
    .target_policy_noise(0.2)
    .target_noise_clip(0.5)
    .actor_update_interval(2)
    .replay_capacity(1_000_000)
    .batch_size(256)
    .training_start(10_000)
    .training_horizon(TOTAL_TIMESTEPS)
    .logger(&mut grapher)
    .build()?;

agent.learn(&mut env, TOTAL_TIMESTEPS)?;

target_policy_noise is the standard deviation of Gaussian noise added to the target actor’s output. target_noise_clip limits that noise component by component. The BoxSpace then clamps the noisy target action to the environment bounds.

By default, target_aggregation_mode is SACCriticAggregationMode::Min. The smaller target Q estimate becomes the bootstrap value shared by both critic losses.

actor_update_interval defaults to 2. The agent updates every critic on each replay optimization, but updates the actor and all target networks only on every second optimization. On skipped actor updates, logger fields such as actor_loss are None.

Important defaults and constraints are:

Collection exploration noise and target-policy noise solve different problems. exploration_noise changes actions sent to the environment after warm-up. target_policy_noise changes only next actions used to calculate replay targets.

Keep Canonical Actor and Target Aggregation

When actor_aggregation_mode is omitted, the actor maximizes the first online critic’s Q estimate. This is canonical TD3 behavior.

ModuRL supports nonempty ensembles of other sizes and explicit aggregation for experiments:

.target_aggregation_mode(SACCriticAggregationMode::Mean)
.actor_aggregation_mode(SACCriticAggregationMode::Median)

Min, Mean, Median, and Max use the same elementwise aggregation implemented for SAC critic ensembles. Setting actor_aggregation_mode makes the actor optimize the selected aggregate instead of the first critic.

These choices define algorithm variants. Omit them when reproducing canonical TD3 with two critics.

Evaluate the Deterministic Actor

As with DDPG, Agent::act adds collection exploration noise. Evaluate the online actor with:

let actions = agent.act_deterministic(&observations)?;

Read Understand a Deterministic Actor-Critic Training Run to distinguish critic updates from delayed actor updates in logs.

Understand a Deterministic Actor-Critic Training Run

The DDPG and TD3 examples confirm that their actors, critics, target networks, and replay buffers train together. Add a DDPGLogger or TD3Logger when you need to compare runs or diagnose one of those components.

Each logger receives two streams. log receives one entry per replay optimization. log_collection receives rewards and completed episodes from one vectorized environment step. An update describes a sampled replay batch; collection describes the policy’s newest behavior.

Add a Logger

Both logger traits use the shared deterministic actor-critic entry types. One logger can therefore support both algorithms:

use modurl::prelude::*;

struct ConsoleLogger;

fn log_update(entry: &DeterministicActorCriticLogEntry) {
    if entry.update_index % 1_000 != 0 {
        return;
    }

    let critic_loss = entry.critic_losses[0]
        .to_scalar::<f32>()
        .expect("critic loss must be scalar");
    let actor_loss = entry.actor_loss
        .as_ref()
        .map(|loss| loss.to_scalar::<f32>().expect("actor loss must be scalar"));

    println!(
        "step={} update={} critic_loss={critic_loss:.4} \
         actor_updated={} actor_loss={actor_loss:?}",
        entry.collection_timestep,
        entry.update_index,
        entry.actor_updated,
    );
}

fn log_collection<I>(entry: &DeterministicActorCriticCollectionLogEntry<I>) {
    for episode in &entry.completed_episodes {
        println!(
            "step={} env={} return={} length={} terminated={} truncated={}",
            episode.collection_timestep,
            episode.environment_index,
            episode.episode_return,
            episode.episode_length,
            episode.terminated,
            episode.truncated,
        );
    }
}

impl<I> DDPGLogger<I> for ConsoleLogger {
    fn log(&mut self, entry: &DeterministicActorCriticLogEntry) {
        log_update(entry);
    }

    fn log_collection(
        &mut self,
        entry: &DeterministicActorCriticCollectionLogEntry<I>,
    ) {
        log_collection(entry);
    }
}

impl<I> TD3Logger<I> for ConsoleLogger {
    fn log(&mut self, entry: &DeterministicActorCriticLogEntry) {
        log_update(entry);
    }

    fn log_collection(
        &mut self,
        entry: &DeterministicActorCriticCollectionLogEntry<I>,
    ) {
        log_collection(entry);
    }
}

Pass a mutable reference while building either agent:

let mut logger = ConsoleLogger;

let mut agent = TD3Agent::builder()
    // Keep the remaining TD3 configuration unchanged.
    .logger(&mut logger)
    .build()?;

The agent borrows the logger. Drop the agent before reading or displaying values held by the concrete logger, as the terminal graph examples do.

Read Replay-Update Metrics

DeterministicActorCriticLogEntry exposes these values:

DDPG sets actor_updated on every replay update. TD3 sets it according to actor_update_interval. When it is false, actor_loss, policy_q_values, and policy_actions are all None; critic metrics remain present.

Compare critic_q_values with bellman_targets when a critic loss changes unexpectedly. Compare policy_actions with replay_actions to distinguish the current actor from behavior stored earlier in replay.

In canonical TD3, policy_q_values come from the first online critic. If actor_aggregation_mode is configured, they contain the aggregate used by the actor objective.

Read Collection Metrics

DeterministicActorCriticCollectionLogEntry describes the newest vectorized environment step:

Each completed episode records its environment index, return, length, ending condition, and global collection timestep. Partial episodes carry across vectorized steps until the environment terminates or truncates.

During the initial random-action phase, collection entries arrive but replay update entries do not. After training_start, update entries occur only at timesteps selected by update_frequency.

You can now separate current policy behavior from replay optimization and, for TD3, delayed actor updates from critic-only updates. The repository’s DeterministicActorCriticGrapher applies the same split to terminal plots in both MuJoCo examples.

Value-Based Training

Use DQN or DDQN when your environment has a discrete action space and you want the agent to learn an action value for every possible action.

This chapter explains the pieces shared by DQNAgent and DDQNAgent. Read DQN for a complete CartPole program. Read Double DQN when you want DDQN’s target calculation instead.

Choose DQN or DDQN

Both agents train an online Q-network and periodically copy its parameters to a target Q-network. They differ only in how they form the next-state training target:

Start with DQN when you need the standard DQN update. Choose DDQN when you want to reduce overestimation of action values. In DQN, the target network both chooses the largest next-state value and uses that value in the training target. Taking a maximum tends to favor values that are accidentally too high. DDQN uses the online network to choose the action and the target network to evaluate it, which reduces that optimistic bias. The builder fields and training loop are otherwise the same.

Terms

A Q-network maps one observation to one value per discrete action. Its output width must equal the number of actions. An online Q-network is the network the optimizer updates. A target Q-network has the same architecture and variable names, but the agent only refreshes its parameters by copying the online network at a fixed interval.

Epsilon-greedy exploration chooses a random valid action with probability epsilon and otherwise chooses the online network’s highest-valued action. epsilon_schedule controls epsilon over the configured training horizon.

Use a Dueling Q-Network

DuelingMLP separates its final representation into a scalar state value and one advantage per action. It mean-centers the advantages and combines the two streams into the same [batch, action_count] output expected from any Q-network:

#![allow(unused)]
fn main() {
let online_q_network = DuelingMLP::builder()
    .input_size(observation_space.shape()[0])
    .output_size(2)
    .vb(online_vb)
    .hidden_layer_sizes(vec![64, 64])
    .value_hidden_layer_sizes(vec![64])
    .advantage_hidden_layer_sizes(vec![64])
    .build()?;
}

hidden_layer_sizes configures the shared trunk. The value and advantage hidden-layer fields configure the two independent streams after that trunk. Leave either stream’s list empty when you want its output head to connect directly to the shared features.

Pass identically configured online and target DuelingMLP instances to either DQNAgent or DDQNAgent. Dueling changes the network architecture; DDQN independently changes the next-state target calculation, so the two techniques can be used together.

Where to Go Next

Build the complete DQN CartPole program. It shows the two Q-networks, replay configuration, and training call in one place. Then use the small, documented change in Double DQN to change its target calculation.

DQN

This page builds a DQN agent for CartPole. It uses one vectorized CartPole environment, two identically shaped Q-networks, an epsilon schedule, and an experience replay buffer. You need Rust, Cargo, and the dependencies from Getting Started.

The program trains for 500,000 environment transitions, so the first complete run takes longer than the PPO quick-start. Lower training_horizon and the argument to learn together when you only want to check that the program runs.

The Q-Networks

Create a Q-network for the online parameters and an identically shaped one for the target parameters. CartPole has four observation values and two discrete actions, so the network reads the environment’s observation shape and produces two Q-values.

DQNAgent needs a Discrete action space. CartPoleV1 exposes the action space through the general Space trait, so this example supplies the known CartPole action count with Discrete::new(2).

Complete Program

Place this program in src/main.rs:

use candle_core::{DType, Device};
use candle_nn::{AdamW, Optimizer, ParamsAdamW, VarBuilder, VarMap};
use modurl::prelude::*;
use modurl_gym::classic_control::cartpole::CartPoleV1;

fn main() {
    let device = Device::Cpu;
    let envs = vec![CartPoleV1::builder().device(&device).build().unwrap()];
    let mut env = VectorizedGymWrapper::from(envs);
    let observation_space = env.observation_space();

    let online_var_map = VarMap::new();
    let online_q_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(2)
        .vb(VarBuilder::from_varmap(&online_var_map, DType::F32, &device))
        .hidden_layer_sizes(vec![64, 64])
        .build()
        .expect("failed to build the online Q-network");

    let mut target_var_map = VarMap::new();
    let target_q_network = MLP::builder()
        .input_size(observation_space.shape()[0])
        .output_size(2)
        .vb(VarBuilder::from_varmap(&target_var_map, DType::F32, &device))
        .hidden_layer_sizes(vec![64, 64])
        .build()
        .expect("failed to build the target Q-network");

    let optimizer = AdamW::new(
        online_var_map.all_vars(),
        ParamsAdamW {
            lr: 2.5e-4,
            ..Default::default()
        },
    )
    .expect("failed to build the optimizer");

    let mut agent = DQNAgent::builder()
        .action_space(Discrete::new(2))
        .observation_space(observation_space)
        .online_q_network(online_q_network)
        .target_q_network(target_q_network)
        .online_vars(&online_var_map)
        .target_vars(&mut target_var_map)
        .optimizer(optimizer)
        .replay_capacity(10_000)
        .batch_size(128)
        .training_start(10_000)
        .update_frequency(10)
        .target_update_interval(500)
        .training_horizon(500_000)
        .epsilon_schedule(|progress: f64| {
            let exploration_progress = (progress / 0.5).min(1.0);
            1.0 + (0.05 - 1.0) * exploration_progress
        })
        .replay_storage_config(ReplayStorageConfig::new(
            ReplayDeviceStrategy::OneDevice(device.clone()),
        ))
        .build()
        .expect("DQN configuration should be valid");

    agent.learn(&mut env, 500_000).expect("DQN learning failed");
    println!("Training complete.");
}

The online network is the only network passed to AdamW; the target network is not optimized directly. At construction, the agent copies the online parameters to the target network. It repeats that copy every 500 transitions.

The epsilon schedule decreases exploration from 1.0 to 0.05 during the first half of the 500,000-transition horizon. The agent collects 10,000 transitions before its first update, then samples 128 replay entries every 10 transitions.

Run the program with:

cargo run

When training finishes, it prints Training complete.. Read Value-Based Training for the DQN and DDQN distinction, or Double DQN to use the alternative target calculation. To record and interpret training metrics, read Understand a Q-Learning Training Run.

Graph a Training Run

Run the repository example when you want terminal graphs instead of the minimal program above:

cargo run --example dqn_cartpole

The example prints each completed CartPole episode’s collection step, length, and return during training. After its 500,000-transition run, it plots DQN loss, exploration epsilon, mean selected Q-values, episode returns, and episode lengths. The update metrics come from replay batches; the episode graphs come from the current collection stream. Read Understand a Q-Learning Training Run for the distinction.

Double DQN

Double DQN, or DDQN, uses the same setup, builder fields, and training loop as DQN. Its difference is the next-state target: the online Q-network chooses the next action, and the target Q-network evaluates that selected action. This separation avoids using the same values for both decisions.

Start with the complete DQN CartPole program. Keep its environment, two Q-networks, two VarMaps, optimizer, replay configuration, epsilon schedule, and learn call unchanged. Then replace the agent construction with this version:

let mut agent = DDQNAgent::builder()
    .action_space(Discrete::new(2))
    .observation_space(observation_space)
    .online_q_network(online_q_network)
    .target_q_network(target_q_network)
    .online_vars(&online_var_map)
    .target_vars(&mut target_var_map)
    .optimizer(optimizer)
    .replay_capacity(10_000)
    .batch_size(128)
    .training_start(10_000)
    .update_frequency(10)
    .target_update_interval(500)
    .training_horizon(500_000)
    .epsilon_schedule(|progress: f64| {
        let exploration_progress = (progress / 0.5).min(1.0);
        1.0 + (0.05 - 1.0) * exploration_progress
    })
    .replay_storage_config(ReplayStorageConfig::new(
        ReplayDeviceStrategy::OneDevice(device.clone()),
    ))
    .build()
    .expect("DDQN configuration should be valid");

agent.learn(&mut env, 500_000).expect("DDQN learning failed");

DDQNAgent has the same configuration validation as DQNAgent: the replay capacity must be at least the batch size; the replay capacity, batch size, update frequency, target-update interval, and training horizon must be nonzero; and gamma and epsilon values must stay in the inclusive 0.0..=1.0 range.

This is a change in the agent’s learning target, not a change in the Q-network shape or device setup. Use Value-Based Training for the shared configuration rules. To record and interpret training metrics, read Understand a Q-Learning Training Run.

Understand a Q-Learning Training Run

The DQN and DDQN programs confirm that training completed, but they do not record metrics. Add a logger when you need to compare runs or diagnose a configuration. Each logger receives an update entry when the agent optimizes from replay and a collection entry after every vectorized environment step.

The repository’s dqn_cartpole example uses both entry types. It prints completed CartPole episodes during training and draws loss, exploration, Q-value, return, and length graphs when training ends. Run it with:

cargo run --example dqn_cartpole

Implement the logger trait that matches the agent. Both traits receive the same entry types:

struct TrainingLogger;

impl DQNLogger for TrainingLogger {
    fn log(&mut self, entry: &QLogEntry) {
        println!("update {}: loss = {:?}", entry.update_index, entry.loss);
    }

    fn log_collection(&mut self, entry: &QCollectionLogEntry) {
        for episode in &entry.completed_episodes {
            println!(
                "environment {}: return = {}, length = {}",
                episode.environment_index,
                episode.episode_return,
                episode.episode_length,
            );
        }
    }
}

let mut logger = TrainingLogger;
let mut agent = DQNAgent::builder()
    // The remaining DQN configuration is unchanged.
    .logger(&mut logger)
    .build()
    .expect("DQN configuration should be valid");

For DDQN, implement DDQNLogger and pass the logger to DDQNAgent::builder() in the same way. log_collection has a default no-op implementation, so a logger can ignore collection metrics when it only needs update metrics. The agent borrows the logger, so keep it alive for as long as the agent is in use.

Read Update Metrics

QLogEntry describes one optimization update from a sampled replay batch. The agent does not emit it for every environment transition. It starts after replay warm-up and then follows the update frequency. The entry reaches the logger before the optimizer applies that update.

The logger exposes these values through QLogEntry:

• loss: the mean-squared error between the online network’s value for each sampled action and its Q-learning target
• epsilon: the exploration probability used while collecting experience
• learning_rate: the optimizer’s current learning rate
• q_values: the online network’s value for the selected action in each sampled replay transition
• replay_rewards: the one-step rewards in that sampled replay batch
• update_index: the zero-based index of this optimization update
• collection_timestep: the total number of environment transitions collected when the agent formed this update

q_values and replay_rewards are tensors for a replay batch, not single summary numbers. Aggregate them, such as with a mean, before graphing or comparing runs. A replay batch can contain transitions from old episodes and from earlier versions of the policy.

Read Collection Metrics

QCollectionLogEntry describes the newest environment interaction, not a sample from replay. The agent emits it after every call to MultiGym::step, including during replay warm-up.

• collection_rewards: one reward for each inner environment from the latest vectorized step
• epsilon: the exploration probability that selected the actions for that step
• collection_timestep: the total number of transitions collected after that step
• completed_episodes: episodes that ended in this vectorized step

completed_episodes is empty when every inner environment continues its current episode. It contains one QEpisodeLogEntry for each inner environment that ended in the latest vectorized step, so one collection entry can report several completed episodes.

Each QEpisodeLogEntry contains the following summary for one finished episode:

• environment_index: which inner environment produced the episode
• episode_return: the sum of every reward collected since that environment’s last reset
• episode_length: the number of environment steps collected since that reset
• terminated and truncated: the environment’s ending flags
• collection_timestep: the total number of collected transitions when that environment finished

The agent resets its running return and length for that environment after recording the entry. An entry never represents an in-progress episode, and it never combines rewards from several episodes. Episode returns and lengths are current collection metrics, so they are not mixed with older replay data.

Read the Metrics Together

Loss has no universal target. Compare it with earlier runs that use the same environment, reward scale, model shape, and replay settings. A noisy loss is normal because every update samples a different replay batch. A persistently growing loss or non-finite values is a signal to inspect the configuration.

The Q-values are estimates, not a direct success score. Their magnitude depends on the reward scale, gamma, and the remaining rewards the agent expects. Use them to spot abrupt changes or divergence, and use episode returns and lengths to decide whether the policy is improving.

Epsilon should follow the schedule you configured. In the DQN example it falls from 1.0 to 0.05; if it remains high, the agent continues to choose random actions often. If it falls too quickly, the collection stream may contain too little exploration when the agent begins learning.

Change One Setting at a Time

Keep the environment, model shape, replay capacity, and epsilon schedule fixed while comparing a change to the learning rate, target-update interval, or update frequency. That makes a change in loss or Q-values easier to attribute to one configuration decision.

Return to DQN or Double DQN to change the training program.

Environments

An environment defines the interaction loop an agent learns from: reset to an initial observation, apply an action, and return the next observation and reward. ModuRL represents one environment with Gym and a batch of environments with MultiGym.

The getting-started example uses CartPoleV1 for individual environments and VectorizedGymWrapper to train from several of them at once.

modurl_gym includes these Gymnasium-compatible environments:

BipedalWalkerV3 implements the standard environment with uneven grass terrain; the hardcore obstacle variant is not included.

These structs expose Gymnasium’s unwrapped dynamics. Registry time limits are applied explicitly with TimeLimitGym: use 500 steps for AcrobotV1, 200 for PendulumV1, and the registry horizon appropriate to the other environment. Keeping the limit in a wrapper makes truncation visible and composable.

modurl_mujoco provides the Gymnasium v5 AntV5, HalfCheetahV5, HopperV5, HumanoidV5, and Walker2dV5 environments. See that crate’s README for model, installation, metadata, and parity details.

Read Use Vectorized Environments before writing manual training or evaluation loops. Read Build a Custom Gym Environment when you need a new environment type.

Environment Wrappers

Wrappers change an environment’s observations, rewards, metadata, or episode boundaries while preserving the Gym interface. When wrappers are nested, the innermost wrapper processes a transition first.

Core Wrappers

These wrappers are available from modurl::wrappers and work with any compatible Gym.

Wrapper order determines which values a wrapper sees. For example, placing RecordEpisodeStatisticsGym inside ClipRewardGym records the underlying return while the agent receives clipped rewards.

Atari Wrappers

These wrappers are available from modurl_ale::wrappers and implement the standard Atari preprocessing steps.

EpisodicLifeGym requires metadata that implements AtariLives. AtariInfo and EpisodeStatisticsInfo<I> provide that implementation when their inner metadata exposes Atari lives.

Batching adapters such as VectorizedGymWrapper and StackedMultiGym are covered in Use Vectorized Environments.

Use Vectorized Environments

MultiGym steps several environments with one batch of actions. PPO uses this interface so one rollout step can collect several transitions.

VectorizedGymWrapper turns a Vec<G> of ordinary Gym values into a vectorized environment:

let envs = (0..4)
    .map(|_| CartPoleV1::builder().device(&device).build().unwrap())
    .collect::<Vec<_>>();
let mut env = VectorizedGymWrapper::from(envs);

With the multithreading feature enabled, MultithreadedVectorizedGymWrapper runs each inner Gym on a persistent worker thread. Pass constructors so every inner environment is created on the thread that owns it, together with representative observation and action spaces:

let constructors = (0..4)
    .map(|_| {
        let device = device.clone();
        move || CartPoleV1::builder().device(&device).build().unwrap()
    })
    .collect();
let mut env = MultithreadedVectorizedGymWrapper::new(
    constructors,
    observation_space,
    action_space,
);

Each inner environment remains one unit of work. The wrapper preserves the batching and auto-reset behavior of VectorizedGymWrapper. If an inner environment returns an error, reset the complete batch before stepping it again.

Reset Once, Then Step

Before a manual loop, call reset once to receive one initial observation per inner environment. Pass that batch to Agent::act, then pass the returned batch of actions to MultiGym::step.

let mut states = env.reset()?;

loop {
    let actions = agent.act(&states)?;
    let step = env.step(actions)?;
    states = step.states;
}

states has one next observation for every inner environment, so it is ready for the next call to act.

Understand Auto-Reset

When an inner environment returns done or truncated, ModuRL resets that one environment immediately. The states field then contains the first observation of its next episode. This lets the next batched step continue without a special reset branch.

The terminal observation is still available. terminal_states contains an entry for each inner environment: Some(state) when that environment ended and None when it continued.

If code needs the true next state for each transition, call transition_next_states:

let step = env.step(actions)?;
let transition_next_states = step.transition_next_states()?;
let next_states_for_the_loop = step.states;

transition_next_states uses a terminal state where one exists and the normal next state otherwise. The second value, step.states, remains the right input for the following action-selection step.

PPOAgent::learn handles this distinction while it collects experience. You only need it when you write a loop that consumes transitions yourself.

Use One Shared World for Several Players

A custom MultiGym can use batch rows for coupled players instead of independent simulations. The environment must consume every player’s action before it advances the shared world and must keep termination and reset behavior consistent across those rows.

For example, a custom two-player game can expose players as the leading tensor dimension:

let mut env = CoupledGame::new()?;
let states = env.reset()?; // [players, ...observation_shape]
let actions = agent.act(&states)?; // [players, ...action_shape]
let step = env.step(actions)?; // advances the shared game once

For a shared-policy agent, acting on the whole observation batch applies the same policy to every player and provides self-play without creating duplicate physics simulations. A coupled implementation should end and reset every player row together whenever the shared episode ends.

Stack Several Batched Environments

StackedMultiGym combines several homogeneous MultiGym values into one flat batch. For example, four two-player games become eight batch rows:

let games = (0..4)
    .map(|seed| {
        let mut game = CoupledGame::new()?;
        game.seed(seed);
        Ok(game)
    })
    .collect::<Result<Vec<_>, GameError>>()?;
let mut env = StackedMultiGym::new(games)?;

let states = env.reset()?; // [8, ...observation_shape]
let actions = agent.act(&states)?; // [8, ...action_shape]
let step = env.step(actions)?; // steps each shared game once

Rows are ordered first by inner gym and then by that gym’s own row order. group_offsets() maps the flattened rows back to their inner gyms. All inner gyms must expose the same observation and action shapes, and each inner gym keeps responsibility for its own auto-reset behavior.

With the multithreading feature enabled, MultithreadedStackedMultiGym runs each inner MultiGym on a persistent worker thread. Pass constructors so every inner gym is created on the thread that owns it, together with representative observation and action spaces:

let constructors = (0..4)
    .map(|seed| move || {
        let mut game = CoupledGame::new().unwrap();
        game.seed(seed);
        game
    })
    .collect();
let mut env = MultithreadedStackedMultiGym::new(
    constructors,
    observation_space,
    action_space,
)?;

The whole inner gym remains one unit of work, so coupled player rows are never split across threads. If an inner gym returns an error, reset the complete stack before stepping it again.

Next, read Build a Custom Gym Environment to provide your own single-environment implementation.

Build a Custom Gym Environment

Implement Gym for one environment. Then place one or more instances in a VectorizedGymWrapper when an agent needs batched interaction.

This small environment has one floating-point observation. Action 0 moves its state left and action 1 moves it right. An episode ends when the state reaches either bound.

Define the Environment

The following code belongs in src/counter_env.rs:

use candle_core::{Device, Tensor};
use modurl::prelude::*;

pub struct CounterEnv {
    state: i32,
    device: Device,
}

impl CounterEnv {
    pub fn new(device: Device) -> Self {
        Self { state: 0, device }
    }

    fn observation(&self) -> candle_core::Result<Tensor> {
        Tensor::from_vec(vec![self.state as f32], (1,), &self.device)
    }
}

impl Gym for CounterEnv {
    type Error = candle_core::Error;
    type SpaceError = candle_core::Error;

    fn reset(&mut self) -> Result<ResetInfo, Self::Error> {
        self.state = 0;
        Ok(ResetInfo {
            state: self.observation()?,
            info: (),
        })
    }

    fn step(&mut self, action: Tensor) -> Result<StepInfo, Self::Error> {
        match action.to_vec0::<u32>()? {
            0 => self.state -= 1,
            1 => self.state += 1,
            _ => panic!("action is outside the action space"),
        }

        let done = self.state.abs() >= 4;
        Ok(StepInfo {
            state: self.observation()?,
            reward: 1.0,
            done,
            truncated: false,
            info: (),
        })
    }

    fn observation_space(&self) -> Box<dyn Space<Error = Self::SpaceError>> {
        Box::new(BoxSpace::new_with_universal_bounds(
            vec![1],
            -4.0,
            4.0,
            &self.device,
        ))
    }

    fn action_space(&self) -> Box<dyn Space<Error = Self::SpaceError>> {
        Box::new(Discrete::new(2))
    }
}

reset returns the initial observation and step consumes one action and returns the observation that follows it, its reward, and the episode flags. The default Gym information type is (), so ordinary environments use ResetInfo and StepInfo without an explicit type parameter. Environments with additional typed metadata can instead implement Gym<MyInfo>.

In src/main.rs, declare the module and bring the environment into scope:

mod counter_env;

use counter_env::CounterEnv;

The Space values are part of the contract. The observation space must match the tensors returned by reset and step. The action space must accept the actions that step understands.

Vectorize the Environment

Build several instances, then wrap them exactly as in the CartPole example:

let envs = (0..4)
    .map(|_| CounterEnv::new(device.clone()))
    .collect::<Vec<_>>();
let env = VectorizedGymWrapper::from(envs);

VectorizedGymWrapper handles the batched action split and auto-reset behavior. The individual environment only needs to implement the single-environment Gym contract.

How-to Guides

These pages solve focused tasks after you understand a CartPole training example. They assume that you already know the surrounding concepts and link back to the tutorial when a complete program is more useful.

Start with Run on CUDA or Metal after the CPU version of the program works.

Run on CUDA or Metal

First run the CPU version of your program. Then enable one Candle backend feature and construct a device for that backend.

CUDA

Enable CUDA on your direct candle-core dependency in Cargo.toml:

candle-core = { version = "0.11", features = ["cuda"] }

In a program, replace Device::Cpu with:

let device = Device::new_cuda(0)?;

0 selects the first CUDA device. The CUDA runtime and a Candle build with CUDA support must be available on the machine.

Metal

Enable Metal on your direct candle-core dependency in Cargo.toml:

candle-core = { version = "0.11", features = ["metal"] }

In a program, replace Device::Cpu with:

let device = Device::new_metal(0)?;

0 selects the first Metal device. Metal builds require a supported Apple platform.

Keep Values on One Device

Pass the same device to the environment builder and to VarBuilder. That places environment observations and model parameters on the same backend.

let env = CartPoleV1::builder().device(&device).build().unwrap();
let vb = VarBuilder::from_varmap(&var_map, candle_core::DType::F32, &device);

If the selected device is unavailable, Candle returns an error when the program constructs it. Fix the backend installation or return to Device::Cpu.

Split Replay Storage From Optimization

Replay-based agents can keep a large replay buffer on the CPU while running models and optimization on an accelerator:

let optimization_device = Device::new_cuda(0)?;
let storage_device = Device::Cpu;

let env = CartPoleV1::builder()
    .device(&optimization_device)
    .build()
    .unwrap();

let actor_vb = VarBuilder::from_varmap(
    &actor_vars,
    candle_core::DType::F32,
    &optimization_device,
);

let device_strategy = ReplayDeviceStrategy::Hybrid {
    optimization_device: optimization_device.clone(),
    storage_device,
};
let replay_storage_config = ReplayStorageConfig::new(device_strategy);

let mut agent = SACAgent::builder()
    // Build the actor, critics, target critics, optimizers, and entropy
    // variable on optimization_device.
    .replay_storage_config(replay_storage_config)
    // Keep the remaining SAC configuration unchanged.
    .build()?;

ReplayStorageConfig controls replay observation representation and uses its ReplayDeviceStrategy to move replay entries and sampled batches. It does not move an environment or model parameters for you.

Set up everything you create for SAC on optimization_device: the environment, actor, critics, target critics, optimizers, and automatic entropy variable. For DDPG or TD3, this includes both the online and target actors as well as every critic pair. You do not create any of these components on storage_device.

Internally, the agent transfers detached transitions to storage_device when it adds them to replay. It transfers sampled replay batches back to optimization_device before each update.

This strategy trades transfer time for accelerator memory. Start with ReplayDeviceStrategy::OneDevice and measure the run. Switch to Hybrid when replay memory is the limiting resource and the transfer cost is acceptable.