Distributional Soft Actor-Critic (DSAC)

DSAC is a state-of-the-art off-policy reinforcement learning algorithm that combines the benefits of Soft Actor-Critic (SAC) with distributional RL using Implicit Quantile Networks (IQN). It was specifically designed for robust control in aerospace applications where uncertainty estimation and smooth control signals are critical.

Architecture Overview

DSAC extends the standard SAC framework with several key innovations:

1. Distributional Critics (IQN): Instead of predicting a single Q-value, DSAC learns the full return distribution using twin Implicit Quantile Networks
2. CAPS Regularization: Conditional Action Policy Smoothness ensures smooth, flight-safe control commands
3. Risk-Sensitive Control: Supports risk distortion functions (CVaR, CPW, Wang) for conservative or aggressive policies

Key Components

Twin IQN Critics

The core innovation of DSAC is the use of Implicit Quantile Networks (IQN) as critics. Each critic predicts a distribution of possible Q-values rather than a single point estimate:

Key features of the IQN critic:

• Cosine Embedding: Quantile levels τ ∈ (0,1) are embedded using cosine basis functions: φ(τ)ᵢ = cos(π · i · τ)
• Hadamard Product: State-action features are combined with quantile embeddings via element-wise multiplication
• Quantile Huber Loss: Robust loss function that handles the asymmetric nature of quantile regression

Training Loop

The DSAC training loop follows the SAC structure with modifications for distributional learning:

Training Steps:

1. Sample Mini-Batch: Draw (s, a, r, s', done) from replay buffer
2. Compute Target: Calculate distributional Bellman target with entropy bonus
3. Update Critics: Minimize quantile Huber loss for both Z₁ and Z₂
4. Freeze Critics: Temporarily disable critic gradients
5. Update Actor: Maximize expected Q-value with CAPS regularization
6. Unfreeze Critics: Re-enable critic gradients
7. Update Temperature: Adjust entropy coefficient α (if automatic)
8. Soft Update Targets: Polyak averaging for target networks

CAPS Regularization

CAPS (Conditional Action Policy Smoothness) is critical for aerospace applications:

• Spatial Smoothness: Penalizes policy sensitivity to small state perturbations

\(L_{spatial} = \lambda_s \cdot \frac{1}{B} \|\mu(s) - \mu(s + \epsilon)\|^2\)

• Temporal Smoothness: Encourages consistent actions over time

\(L_{temporal} = \lambda_t \cdot \frac{1}{B} \|a_t - a_{t+1}\|^2\)

Risk Distortion Functions

DSAC supports risk-sensitive control through distortion of quantile levels:

Function	Formula	Use Case
Neutral	τ	Standard expected value
CVaR	clamp(τ · ξ, 0, 1)	Conservative (worst-case)
CPW	τ^ξ / (τ^ξ + (1-τ)^ξ)^{1/ξ}	Probability weighting
Wang	Φ(Φ⁻¹(τ) + ξ)	Normal transform

Key Differences vs SAC

Feature	SAC	DSAC
Critic Output	Single Q-value	N quantile values
Loss Function	MSE	Quantile Huber
Uncertainty	None	Full distribution
Smoothness	None	CAPS regularization
Risk Sensitivity	None	Distortion functions

Quick Start

import numpy as np
import torch
from tensoraerospace.agent import DSAC
from tensoraerospace.envs.b747 import ImprovedB747Env

def step_reference(steps: int, deg: float = 5.0) -> np.ndarray:
    ref = np.zeros((1, steps), dtype=np.float32)
    ref[:, steps // 5 :] = np.deg2rad(deg)
    return ref

device = "cuda" if torch.cuda.is_available() else "cpu"
num_steps = 800

env = ImprovedB747Env(
    initial_state=np.array([0.0, 0.0, 0.0, 0.0], dtype=float),
    reference_signal=step_reference(num_steps, deg=5.0),
    number_time_steps=num_steps,
    dt=0.02,
    reward_mode="step_response",
)

agent = DSAC(
    env,
    batch_size=256,
    memory_capacity=500_000,
    learning_starts=10_000,
    updates_per_step=1,
    num_quantiles=32,
    embedding_dim=64,
    hidden_layers=[64, 64],
    huber_threshold=1.0,
    lr=4.4e-4,
    policy_lr=4.4e-4,
    gamma=0.99,
    tau=0.005,
    caps_lambda_smoothness=400.0,
    caps_lambda_temporal=400.0,
    caps_noise_std=0.05,
    device=device,
    log_every_updates=50,
    automatic_entropy_tuning=True,
)

# Training
agent.train(num_episodes=100, save_best=True, save_path="./runs")
agent.close()

Hyperparameters

Critical Parameters

Parameter	Default	Description
num_quantiles	8	Number of quantile samples (16-64 recommended)
embedding_dim	64	Cosine embedding dimension
huber_threshold	1.0	Huber loss threshold κ
batch_size	256	Mini-batch size
learning_starts	10,000	Steps before training starts

CAPS Parameters

Parameter	Default	Description
caps_lambda_smoothness	400.0	Spatial smoothness weight
caps_lambda_temporal	400.0	Temporal smoothness weight
caps_noise_std	0.05	Noise for spatial perturbation

Optimization Parameters

Parameter	Default	Description
lr	4.4e-4	Critic learning rate
policy_lr	4.4e-4	Actor learning rate (defaults to lr)
gamma	0.99	Discount factor
tau	0.005	Soft update coefficient
target_update_interval	1	Target network update frequency

Risk Control

Parameter	Default	Description
risk_distortion	"neutral"	Distortion function name
risk_measure	1.0	Distortion parameter ξ

Training Tips

• Keep num_quantiles between 16-64 for stable training
• Higher caps_lambda values produce smoother but potentially slower-converging policies
• Use risk_distortion="cvar" with risk_measure < 1.0 for conservative flight control
• Reduce updates_per_step if actions become overly smooth or training unstable

Vectorized Training

For environments with parallel simulation (e.g., GPU-accelerated):

agent.train_vector(
    total_steps=500_000,
    warmup_steps=10_000,
    log_every=2_000,
    reward_window=200,
    save_best=True,
    save_path="./runs",
)

Unified training interface

DSAC implements the shared unified train() API from BaseRLModel:

def train(
    self,
    num_episodes: int = 1,
    *,
    max_steps: Optional[int] = None,
    save_best: bool = False,
    save_path: Optional[str] = None,
    verbose: bool = True,
    **kwargs,
) -> dict

Algorithm-specific options accepted via **kwargs:

• save_best_with_gradients (bool): include optimizer gradients when saving the best model checkpoint.

Returns a metrics dictionary with episode_rewards, best_reward, updates and total_steps.

API Reference

DSAC(env, *, updates_per_step=1, batch_size=256, memory_capacity=500000, lr=0.00044, policy_lr=None, gamma=0.99, tau=0.005, alpha=0.2, policy_type='Gaussian', target_update_interval=1, automatic_entropy_tuning=True, target_entropy_scale=1.0, min_alpha=0.0, exploration_noise_std=0.0, max_grad_norm=None, reward_clip=None, hidden_size=64, num_quantiles=8, num_quantiles_exp=None, embedding_dim=64, hidden_layers=None, layer_norm=True, huber_threshold=1.0, learning_starts=10000, warmup_action_scale=1.0, caps_lambda_smoothness=400.0, caps_lambda_temporal=400.0, caps_noise_std=0.05, risk_distortion='neutral', risk_measure=1.0, device='cpu', verbose_histogram=False, seed=42, log_dir=None, log_every_updates=1, wandb_project=None, wandb_entity=None, wandb_run_name=None, wandb_tags=None, wandb_config=None)

Bases: BaseRLModel

Distributional SAC (dsac-flight port).

train(num_episodes=1, *, max_steps=None, save_best=False, save_path=None, verbose=True, **kwargs)

Train DSAC for num_episodes (unified interface).

Parameters:

Name	Type	Description	Default
num_episodes	int	Number of training episodes.	1
max_steps	Optional[int]	Optional per-episode step cap.	None
save_best	bool	If True, checkpoint the best-reward model.	False
save_path	Optional[str]	Destination for best-reward checkpoints.	None
verbose	bool	If True, display a tqdm progress bar.	True
**kwargs	Any	Algorithm-specific options: save_best_with_gradients (bool): save optimizer gradients alongside model weights.	{}

Returns:

Name	Type	Description
dict	dict	Training metrics (episode_rewards, best_reward,
	dict	updates, total_steps).

to_device(device)

Move all DSAC components (nets, optim states) to the target device.

eval()

Switch all DSAC networks to eval mode.

from_pretrained(repo_name, access_token=None, version=None, load_gradients=False) classmethod

Load DSAC checkpoint from local dir or Hugging Face Hub.

push_to_hub(repo_name, access_token=None, save_path=None, include_gradients=False)

Save model checkpoint and upload it to Hugging Face Hub.

ZNet(*, n_states, n_actions, n_hidden_layers, n_hidden_units, n_cos, device)

Bases: Module

Wrapper around IQN that outputs Z(s, a; taus) with shape (B, N).

generate_taus(*, batch_size, n_taus, device) staticmethod

Uniform taus in (0,1), shape (B, N).

IQN(*, n_inputs, n_outputs, embedding_size, n_hidden_layers, n_hidden_units, device)

Bases: Module

Implicit Quantile Network (dsac-flight style).

NormalPolicyNet(*, obs_dim, action_dim, n_hidden_layers, n_hidden_units)

Bases: Module

Outputs an Independent Normal distribution for continuous actions.

risk_distortions

Risk distortion functions copied from dsac-flight.

normal_cdf(tau, mean=0.0, std=1.0)

CDF of the normal distribution.

normal_inverse_cdf(tau, mean=0.0, std=1.0)

Inverse CDF of the normal distribution.

neutral(tau, _xi)

Neutral distortion returns the original quantiles.

cvar(tau, xi)

Conditional value at risk distortion (clamped to [0,1]).

cpw(tau, xi)

Cumulative probability weighting.

wang(tau, xi)

Wang transform.

Acknowledgements

The DSAC implementation in TensorAeroSpace was inspired by and partially based on the excellent work by Peter Seres on risk-sensitive distributional reinforcement learning for flight control:

• Repository: peter-seres/dsac-flight
• Description: Risk-sensitive Distributional Reinforcement Learning for Flight Control (MSc thesis project)

We gratefully acknowledge Peter Seres's contribution to the field of distributional RL for aerospace applications. His implementation of DSAC with IQN critics for the PH-LAB (Cessna Citation II) research aircraft provided valuable insights and architectural decisions that informed our implementation.

Citation

If you use the DSAC agent in your research, please consider citing both TensorAeroSpace and the original dsac-flight repository.