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.