Soft Actor‑Critic (SAC)¶

SAC is an off-policy actor-critic with entropy maximization: it learns a stochastic policy while increasing expected reward and entropy (exploration). Our implementation employs twin Q-networks, a target critic, Gaussian/deterministic policy options, a replay buffer, soft updates, and optional automatic entropy tuning.

Components¶

Twin Q-networks: QNetwork(state, action) -> (Q1, Q2) plus target critic_target
Policy: GaussianPolicy (default) or DeterministicPolicy (no entropy term)
Replay buffer: ReplayMemory for mini-batch sampling
Soft target update: soft_update(target, source, tau)
Automatic entropy tuning: optimizes alpha toward \(H_{\text{target}} = -\dim(\mathcal{A})\)

Theory (as implemented)¶

Soft Q target (double Q + entropy):

\[ \begin{aligned} & a' \sim \pi_\theta(\cdot|s')\ ,\ \log \pi_\theta(a'|s'), \\ & Q_{\text{targ}}(s,a) = r + \gamma\, \big( \min(Q_1(s',a'), Q_2(s',a')) - \alpha\, \log \pi_\theta(a'|s') \big) \end{aligned} \]

Critic update (MSE to target): \(\mathcal{L}_{Q_i} = \mathbb{E}[(Q_i(s,a) - Q_{\text{targ}})^2]\)
Policy update (reparameterization):

\[ \mathcal{L}_\pi = \mathbb{E}_{s\sim \mathcal{D},\ \epsilon\sim\mathcal{N}}\big[ \alpha\, \log \pi_\theta(f_\theta(\epsilon; s) | s) - Q_{\min}(s, f_\theta(\epsilon; s)) \big] \]

Automatic \(\alpha\) tuning (optional):

\[ \mathcal{L}_\alpha = -\,\mathbb{E}_{a\sim\pi}\big[\log \alpha\, (\log \pi_\theta(a|s) + H_{\text{target}})\big]\ ,\quad \alpha \leftarrow e^{\log \alpha} \]

Quick start¶

import gymnasium as gym
from tensoraerospace.agent.sac.sac import SAC

env = gym.make('LinearLongitudinalF16-v0', number_time_steps=2000)
agent = SAC(env,
            updates_per_step=1,
            batch_size=64,
            memory_capacity=100000,
            lr=3e-4,
            gamma=0.99,
            tau=0.005,
            alpha=0.2,
            policy_type='Gaussian',
            target_update_interval=1,
            automatic_entropy_tuning=True,
            hidden_size=256,
            device='cpu')

agent.train(num_episodes=100)
agent.save('./runs')

Tip

For continuous action spaces keep GaussianPolicy with automatic_entropy_tuning=True to stabilize exploration.

Unified training interface¶

All TensorAeroSpace RL agents share a common train() signature defined on BaseRLModel:

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

For SAC the algorithm-specific options accepted via **kwargs are:

save_best_with_gradients (bool): include optimizer gradients in best-model checkpoints.

Example:

stats = agent.train(
    num_episodes=100,
    max_steps=500,
    save_best=True,
    save_path='./runs/sac_best',
)
print(stats['best_reward'], len(stats['episode_rewards']))

Practical tips¶

Increase batch_size and memory_capacity for steadier gradients
Choose tau around 0.005–0.02 for soft target updates
With a deterministic policy set alpha=0 and disable auto tuning
When using DeterministicPolicy with action_space=None, note that action_scale and action_bias are now torch.Tensor values (not Python floats)

Gymnasium 5-tuple API

This implementation uses the modern Gymnasium 5-tuple step API internally:

next_state, reward, terminated, truncated, info = env.step(action)
done = terminated or truncated

If you are migrating from older code that used the 4-tuple API (next_state, reward, done, info = env.step(action)), ensure your environment is compatible with Gymnasium and returns the 5-tuple.

API reference¶

`SAC(env, updates_per_step=1, batch_size=32, memory_capacity=10000000, lr=0.0003, policy_lr=0.0003, gamma=0.99, tau=0.005, alpha=0.2, policy_type='Gaussian', target_update_interval=1, automatic_entropy_tuning=False, hidden_size=256, 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

Soft Actor-Critic (SAC) algorithm for reinforcement learning.

Parameters:

Name	Type	Description	Default
`env`	`Any`	Environment (Gym API compatible).	required
`updates_per_step`	`int`	Updates per interaction step.	`1`
`batch_size`	`int`	Mini-batch size.	`32`
`memory_capacity`	`int`	Replay buffer capacity.	`10000000`
`lr`	`float`	Learning rate.	`0.0003`
`gamma`	`float`	Discount coefficient.	`0.99`
`tau`	`float`	Soft update coefficient for target network.	`0.005`
`alpha`	`float`	Entropy coefficient (for policy).	`0.2`
`policy_type`	`str`	Policy type ("Gaussian" or "Deterministic").	`'Gaussian'`
`target_update_interval`	`int`	Target network update interval.	`1`
`automatic_entropy_tuning`	`bool`	Automatic entropy tuning.	`False`
`hidden_size`	`int`	Hidden layer size of networks.	`256`
`device`	`str \| device`	Device for computations.	`'cpu'`
`verbose_histogram`	`bool`	Histogram logging in TensorBoard.	`False`
`seed`	`int`	Random number generator seed.	`42`

Attributes:

Name	Type	Description
`critic`		Critic network.
`critic_optim`		Optimizer for updating critic weights.
`critic_target`		Target critic network.
`policy`	`Union[GaussianPolicy, DeterministicPolicy]`	Agent policy.
`policy_optim`		Optimizer for updating policy weights.

Initialize SAC agent, networks, replay buffer, and optimizers.

`select_action(state, evaluate=False)` ¶

Select action based on current state.

Parameters:

Name	Type	Description	Default
`state`	`ndarray`	Current state of the agent.	required
`evaluate`	`bool`	Evaluation mode flag.	`False`

Returns:

Name	Type	Description
`action`	`ndarray`	Selected action.

`select_action_batch(states, *, evaluate=False, return_tensor=False)` ¶

Select actions for a batch of states.

This is the recommended API for vectorized environments.

Parameters:

Name	Type	Description	Default
`states`	`Union[ndarray, Tensor]`	Batch of states of shape (N, obs_dim). Can be numpy or torch.	required
`evaluate`	`bool`	If True, use deterministic evaluation action.	`False`
`return_tensor`	`bool`	If True, return a torch Tensor on agent device, else numpy.	`False`

Returns:

Type	Description
`Union[ndarray, Tensor]`	Actions with shape (N, act_dim).

`update_parameters(memory, batch_size, updates)` ¶

Update network parameters based on a mini-batch from memory.

Parameters:

Name	Type	Description	Default
`memory`	`ReplayMemory`	Memory for storing transitions.	required
`batch_size`	`int`	Mini-batch size.	required
`updates`	`int`	Number of updates.	required

Returns:

Name	Type	Description
`qf1_loss`	`float`	Loss value for the first Q-network.
`qf2_loss`	`float`	Loss value for the second Q-network.
`policy_loss`	`float`	Loss value for the policy.
`alpha_loss`	`float`	Loss value for the alpha coefficient.
`alpha_tlogs`	`float`	Value of the alpha coefficient.

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

Train SAC for the given number of 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` lets the environment decide when an episode terminates.	`None`
`save_best`	`bool`	If True, save a checkpoint whenever a new best episode reward is reached.	`False`
`save_path`	`Optional[str]`	Destination for best-reward checkpoints.	`None`
`verbose`	`bool`	If True, show a tqdm progress bar.	`True`
`**kwargs`	`Any`	Algorithm-specific options. Recognized keys: `save_best_with_gradients` (bool): include optimizer gradients when saving the best model.	`{}`

Returns:

Name	Type	Description
`dict`	`dict`	Training metrics with keys `episode_rewards`,
	`dict`	`best_reward` and `updates`.

`train_vector(*, total_steps, warmup_steps=10000, log_every=2000, reward_window=200, save_best=False, save_path=None, save_best_with_gradients=False)` ¶

Train SAC on a vectorized (batched) environment.

Expected env API (Gymnasium-like, batched): reset() -> (obs[N, obs_dim], info) step(action[N, act_dim]) -> (obs, reward[N], terminated[N], truncated[N], info)

Notes

If env has auto_reset=True, it may reset done envs internally. We still use terminated/truncated from the current step for episode accounting.
For replay bootstrap we use terminated only (not truncated), consistent with train().

`close()` ¶

Flush and close TensorBoard writer.

`get_param_env()` ¶

Return serializable configuration of environment and policy.

`save(path=None, save_gradients=False)` ¶

Save PyTorch model to the specified directory.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path \| None`	Save path. If None, creates a folder with current date and time in the working directory.	`None`
`save_gradients`	`bool`	Save optimizer states for continuing training (Adam moments, etc.).	`False`

Returns:

Type	Description
`None`	None

`__load(path, load_gradients=False)` `classmethod` ¶

Load a SAC agent from checkpoint folder.

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

Load pretrained model from local directory or Hugging Face.

Parameters:

Name	Type	Description	Default
`repo_name`	`str`	Path to local folder with weights or repository name in format "namespace/repo_name" on Hugging Face.	required
`access_token`	`Optional[str]`	Access token for private HF repository.	`None`
`version`	`Optional[str]`	Revision/branch/tag of HF repository.	`None`
`load_gradients`	`bool`	Load optimizer states for continuing training.	`False`

Returns:

Name	Type	Description
`SAC`	`SAC`	Initialized agent.

`ReplayMemory(capacity, seed)` ¶

Replay buffer for off-policy reinforcement learning algorithms.

Parameters:

Name	Type	Description	Default
`capacity`	`int`	Maximum number of transitions to store.	required
`seed`	`int`	Random seed used for sampling.	required

Attributes:

Name	Type	Description
`capacity`	`int`	Maximum capacity.
`buffer`	`List`	Stored transitions.
`position`	`int`	Current write position.

Initialize replay buffer with fixed capacity and seed.

Parameters:

Name	Type	Description	Default
`capacity`	`int`	Maximum number of transitions.	required
`seed`	`int`	Random seed for sampling.	required

`push(state, action, reward, next_state, done)` ¶

Add a transition to the replay buffer.

Parameters:

Name	Type	Description	Default
`state`	`ndarray`	Current state.	required
`action`	`ndarray`	Action taken.	required
`reward`	`Union[float, ndarray]`	Reward received.	required
`next_state`	`ndarray`	Next state.	required
`done`	`Union[bool, float]`	Episode termination flag (bool) or mask (0.0/1.0).	required

`sample(batch_size)` ¶

Sample a batch of transitions.

Parameters:

Name	Type	Description	Default
`batch_size`	`int`	Batch size.	required

Returns:

Type	Description
`ndarray`	Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
`ndarray`	`(state, action, reward, next_state, done)` batch.

`len()` ¶

Return the current number of stored transitions.

Returns:

Name	Type	Description
`int`	`int`	Buffer size.

`save_buffer(env_name, suffix='', save_path=None)` ¶

Save the replay buffer to disk.

Parameters:

Name	Type	Description	Default
`env_name`	`str`	Environment name used in the default path.	required
`suffix`	`str`	Optional filename suffix.	`''`
`save_path`	`str`	Optional explicit save path. If None, a default path under `checkpoints/` is used.	`None`

`load_buffer(save_path)` ¶

Load a replay buffer from disk.

Parameters:

Name	Type	Description	Default
`save_path`	`str`	Path to a previously saved buffer.	required

`ValueNetwork(num_inputs, hidden_dim)` ¶

Bases: Module

Neural network for value function approximation.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	Number of input features.	required
`hidden_dim`	`int`	Hidden layers dimension.	required

Attributes:

Name	Type	Description
`linear1`	`Linear`	First linear layer.
`linear2`	`Linear`	Second linear layer.
`linear3`	`Linear`	Third linear layer.

Initialize ValueNetwork.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	Number of input features.	required
`hidden_dim`	`int`	Hidden layers dimension.	required

`forward(state)` ¶

Forward pass of neural network.

Parameters:

Name	Type	Description	Default
`state`	`Tensor`	Input state tensor.	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Output value tensor.

`QNetwork(num_inputs, num_actions, hidden_dim)` ¶

Bases: Module

Neural network for Q function evaluation.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	Number of input features.	required
`num_actions`	`int`	Number of actions.	required
`hidden_dim`	`int`	Hidden layers dimension.	required

Attributes:

Name	Type	Description
`linear1`	`Linear`	First linear layer for Q1.
`linear2`	`Linear`	Second linear layer for Q1.
`linear3`	`Linear`	Third linear layer for Q1.
`linear4`	`Linear`	First linear layer for Q2.
`linear5`	`Linear`	Second linear layer for Q2.
`linear6`	`Linear`	Third linear layer for Q2.

Initialize twin Q-network architecture.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	State dimension.	required
`num_actions`	`int`	Action dimension.	required
`hidden_dim`	`int`	Hidden layer width.	required

`forward(state, action)` ¶

Forward pass for Q-value estimation.

Parameters:

Name	Type	Description	Default
`state`	`Tensor`	State tensor.	required
`action`	`Tensor`	Action tensor.	required

Returns:

Type	Description
`Tuple[Tensor, Tensor]`	Tuple[torch.Tensor, torch.Tensor]: `(Q1, Q2)` tensors.

`GaussianPolicy(num_inputs, num_actions, hidden_dim, action_space=None)` ¶

Bases: Module

Gaussian policy used by SAC.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	Number of input features.	required
`num_actions`	`int`	Number of actions.	required
`hidden_dim`	`int`	Hidden layer dimension.	required
`action_space`	`Optional[Space]`	Action space. Defaults to None.	`None`

Attributes:

Name	Type	Description
`linear1`	`Linear`	First linear layer.
`linear2`	`Linear`	Second linear layer.
`mean_linear`	`Linear`	Linear layer for mean value.
`log_std_linear`	`Linear`	Linear layer for log standard deviation.
`action_scale`	`Tensor`	Action scale.
`action_bias`	`Tensor`	Action bias.

Initialize Gaussian policy network and action scaling.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	State dimension.	required
`num_actions`	`int`	Action dimension.	required
`hidden_dim`	`int`	Hidden layer width.	required
`action_space`	`Space \| None`	Optional gym space to derive scaling/bias.	`None`

`forward(state)` ¶

Compute mean and log standard deviation for a batch of states.

`sample(state)` ¶

Sample an action using the reparameterization trick.

Returns:

Type	Description
`Tuple[Tensor, Tensor, Tensor]`	Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: `(action, log_prob, mean)`.

`DeterministicPolicy(num_inputs, num_actions, hidden_dim, action_space=None)` ¶

Bases: Module

Deterministic policy used by SAC (e.g., for evaluation).

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	Number of input features.	required
`num_actions`	`int`	Number of actions.	required
`hidden_dim`	`int`	Hidden layer dimension.	required
`action_space`	`Optional[Space]`	Action space used to scale/bias actions.	`None`

Attributes:

Name	Type	Description
`linear1`	`Linear`	First linear layer.
`linear2`	`Linear`	Second linear layer.
`mean`	`Linear`	Linear layer producing the mean action.
`noise`	`Tensor`	Noise tensor for exploration.
`action_scale`	`Tensor \| float`	Action scaling factor.
`action_bias`	`Tensor \| float`	Action bias.

Initialize deterministic policy network for SAC evaluation.

Parameters:

Name	Type	Description	Default
`num_inputs`	`int`	State dimension.	required
`num_actions`	`int`	Action dimension.	required
`hidden_dim`	`int`	Hidden layer width.	required
`action_space`	`Space \| None`	Optional gym space to derive scaling/bias.	`None`

`forward(state)` ¶

Compute the deterministic action mean for a batch of states.

`sample(state)` ¶

Sample an action by adding bounded Gaussian noise to the mean.

Returns:

Type	Description
`Tuple[Tensor, Tensor, Tensor]`	Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: `(action, log_prob, mean)`.

Soft Actor‑Critic (SAC)¶

Components¶

Theory (as implemented)¶

Quick start¶

Unified training interface¶

Practical tips¶

API reference¶

select_action(state, evaluate=False) ¶

select_action_batch(states, *, evaluate=False, return_tensor=False) ¶

update_parameters(memory, batch_size, updates) ¶

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

train_vector(*, total_steps, warmup_steps=10000, log_every=2000, reward_window=200, save_best=False, save_path=None, save_best_with_gradients=False) ¶

close() ¶

get_param_env() ¶

save(path=None, save_gradients=False) ¶

__load(path, load_gradients=False) classmethod ¶

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

ReplayMemory(capacity, seed) ¶

push(state, action, reward, next_state, done) ¶

sample(batch_size) ¶

__len__() ¶

save_buffer(env_name, suffix='', save_path=None) ¶

load_buffer(save_path) ¶

ValueNetwork(num_inputs, hidden_dim) ¶

forward(state) ¶

QNetwork(num_inputs, num_actions, hidden_dim) ¶

forward(state, action) ¶

GaussianPolicy(num_inputs, num_actions, hidden_dim, action_space=None) ¶

forward(state) ¶

sample(state) ¶

DeterministicPolicy(num_inputs, num_actions, hidden_dim, action_space=None) ¶

forward(state) ¶

sample(state) ¶

`select_action(state, evaluate=False)` ¶

`select_action_batch(states, *, evaluate=False, return_tensor=False)` ¶

`update_parameters(memory, batch_size, updates)` ¶

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

`train_vector(*, total_steps, warmup_steps=10000, log_every=2000, reward_window=200, save_best=False, save_path=None, save_best_with_gradients=False)` ¶

`close()` ¶

`get_param_env()` ¶

`save(path=None, save_gradients=False)` ¶

`__load(path, load_gradients=False)` `classmethod` ¶

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

`ReplayMemory(capacity, seed)` ¶

`push(state, action, reward, next_state, done)` ¶

`sample(batch_size)` ¶

`len()` ¶

`save_buffer(env_name, suffix='', save_path=None)` ¶

`load_buffer(save_path)` ¶

`ValueNetwork(num_inputs, hidden_dim)` ¶

`forward(state)` ¶

`QNetwork(num_inputs, num_actions, hidden_dim)` ¶

`forward(state, action)` ¶

`GaussianPolicy(num_inputs, num_actions, hidden_dim, action_space=None)` ¶

`forward(state)` ¶

`sample(state)` ¶

`DeterministicPolicy(num_inputs, num_actions, hidden_dim, action_space=None)` ¶

`forward(state)` ¶

`sample(state)` ¶