lit_sb3.algorithms.sac

Classes

Actor — Actor network (policy) for SAC.
CnnPolicy — Policy class (with both actor and critic) for SAC.
MultiInputPolicy — Policy class (with both actor and critic) for SAC.
SAC — Soft Actor-Critic (SAC).
SACPolicy — Policy class (with both actor and critic) for SAC.

source class Actor(observation_space: spaces.Space, action_space: spaces.Box, net_arch: list[int], features_extractor: nn.Module, features_dim: int, activation_fn: type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, full_std: bool = True, use_expln: bool = False, clip_mean: float = 2.0, normalize_images: bool = True)

Bases : BasePolicy

Actor network (policy) for SAC.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Box — Action space
net_arch : list[int] — Network architecture
features_extractor : nn.Module — Network to extract features (a CNN when using images, a nn.Flatten() layer otherwise)
features_dim : int — Number of features
activation_fn : type[nn.Module] — Activation function
use_sde : bool — Whether to use State Dependent Exploration or not
log_std_init : float — Initial value for the log standard deviation
full_std : bool — Whether to use (n_features x n_actions) parameters for the std instead of only (n_features,) when using gSDE.
use_expln : bool — Use expln() function instead of exp() when using gSDE to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, exp() is usually enough.
clip_mean : float — Clip the mean output when using gSDE to avoid numerical instability.
normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)

Attributes

device : torch.device — Infer which device this policy lives on by inspecting its parameters.
squash_output : bool — (bool) Getter for squash_output.

Methods

get_std — Retrieve the standard deviation of the action distribution.
reset_noise — Sample new weights for the exploration matrix, when using gSDE.
get_action_dist_params — Get the parameters for the action distribution.
forward
action_log_prob

source method Actor.get_std() → torch.Tensor

Retrieve the standard deviation of the action distribution.

Only useful when using gSDE. It corresponds to torch.exp(log_std) in the normal case, but is slightly different when using expln function (cf StateDependentNoiseDistribution doc).

source method Actor.reset_noise(batch_size: int = 1) → None

Sample new weights for the exploration matrix, when using gSDE.

Parameters

batch_size : int
batch_size : int — The number of samples in the batch.

source method Actor.get_action_dist_params(obs: PyTorchObs) → tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]

Get the parameters for the action distribution.

Parameters

obs : PyTorchObs
obs : PyTorchObs — The current observation.

Returns

tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]] — Mean, standard deviation and optional keyword arguments.

source method Actor.forward(obs: PyTorchObs, deterministic: bool = False) → torch.Tensor

source method Actor.action_log_prob(obs: PyTorchObs) → tuple[torch.Tensor, torch.Tensor]

source class CnnPolicy(observation_space: spaces.Space, action_space: spaces.Box, lr_schedule: Schedule, net_arch: list[int] | dict[str, list[int]] | None = None, activation_fn: type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, use_expln: bool = False, clip_mean: float = 2.0, features_extractor_class: type[BaseFeaturesExtractor] = NatureCNN, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None, n_critics: int = 2, share_features_extractor: bool = False)

Bases : SACPolicy

Policy class (with both actor and critic) for SAC.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Box — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
net_arch : list[int] | dict[str, list[int]] | None — The specification of the policy and value networks.
activation_fn : type[nn.Module] — Activation function
use_sde : bool — Whether to use State Dependent Exploration or not
log_std_init : float — Initial value for the log standard deviation
use_expln : bool — Use expln() function instead of exp() when using gSDE to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, exp() is usually enough.
clip_mean : float — Clip the mean output when using gSDE to avoid numerical instability.
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)
optimizer_class : type[torch.optim.Optimizer] — The optimizer to use, torch.optim.Adam by default
optimizer_kwargs : dict[str, Any] | None — Additional keyword arguments, excluding the learning rate, to pass to the optimizer
n_critics : int — Number of critic networks to create.
share_features_extractor : bool — Whether to share or not the features extractor between the actor and the critic (this saves computation time)

Attributes

device : torch.device — Infer which device this policy lives on by inspecting its parameters.
squash_output : bool — (bool) Getter for squash_output.

source class MultiInputPolicy(observation_space: spaces.Space, action_space: spaces.Box, lr_schedule: Schedule, net_arch: list[int] | dict[str, list[int]] | None = None, activation_fn: type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, use_expln: bool = False, clip_mean: float = 2.0, features_extractor_class: type[BaseFeaturesExtractor] = CombinedExtractor, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None, n_critics: int = 2, share_features_extractor: bool = False)

Bases : SACPolicy

Policy class (with both actor and critic) for SAC.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Box — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
net_arch : list[int] | dict[str, list[int]] | None — The specification of the policy and value networks.
activation_fn : type[nn.Module] — Activation function
use_sde : bool — Whether to use State Dependent Exploration or not
log_std_init : float — Initial value for the log standard deviation
use_expln : bool — Use expln() function instead of exp() when using gSDE to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, exp() is usually enough.
clip_mean : float — Clip the mean output when using gSDE to avoid numerical instability.
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)
optimizer_class : type[torch.optim.Optimizer] — The optimizer to use, torch.optim.Adam by default
optimizer_kwargs : dict[str, Any] | None — Additional keyword arguments, excluding the learning rate, to pass to the optimizer
n_critics : int — Number of critic networks to create.
share_features_extractor : bool — Whether to share or not the features extractor between the actor and the critic (this saves computation time)

Attributes

device : torch.device — Infer which device this policy lives on by inspecting its parameters.
squash_output : bool — (bool) Getter for squash_output.

source class SAC(policy: str | type[SACPolicy], learning_rate: float | Schedule = 0.0003, buffer_size: int = 1000000, learning_starts: int = 100, batch_size: int = 256, tau: float = 0.005, gamma: float = 0.99, train_freq: int | tuple[int, str] = 1, gradient_steps: int = 1, action_noise: ActionNoise | None = None, replay_buffer_class: type[ReplayBuffer] | None = None, replay_buffer_kwargs: dict[str, Any] | None = None, optimize_memory_usage: bool = False, n_steps: int = 1, ent_coef: str | float = 'auto', target_update_interval: int = 1, target_entropy: str | float = 'auto', use_sde: bool = False, sde_sample_freq: int = -1, use_sde_at_warmup: bool = False, stats_window_size: int = 100, policy_kwargs: dict[str, Any] | None = None, verbose: int = 0, seed: int | None = None, device: torch.device | str = 'auto', _init_setup_model: bool = True)

Bases : OffPolicyAlgorithm

Soft Actor-Critic (SAC).

Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, This implementation borrows code from original implementation (https://github.com/haarnoja/sac) from OpenAI Spinning Up (https://github.com/openai/spinningup), from the softlearning repo (https://github.com/rail-berkeley/softlearning/) and from Stable Baselines (https://github.com/hill-a/stable-baselines) Paper: https://arxiv.org/abs/1801.01290 Introduction to SAC: https://spinningup.openai.com/en/latest/algorithms/sac.html.

Note: we use double q target and not value target as discussed in https://github.com/hill-a/stable-baselines/issues/270

Parameters

policy : str | type[SACPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — learning rate for adam optimizer, the same learning rate will be used for all networks (Q-Values, Actor and Value function) it can be a function of the current progress remaining (from 1 to 0)
buffer_size : int — size of the replay buffer
learning_starts : int — how many steps of the model to collect transitions for before learning starts
batch_size : int — Minibatch size for each gradient update
tau : float — the soft update coefficient ("Polyak update", between 0 and 1)
gamma : float — the discount factor
train_freq : int | tuple[int, str] — Update the model every train_freq steps. Alternatively pass a tuple of frequency and unit like (5, "step") or (2, "episode").
gradient_steps : int — How many gradient steps to do after each rollout (see train_freq) Set to -1 means to do as many gradient steps as steps done in the environment during the rollout.
action_noise : ActionNoise | None — the action noise type (None by default), this can help for hard exploration problem. Cf common.noise for the different action noise type.
replay_buffer_class : type[ReplayBuffer] | None — Replay buffer class to use (for instance HerReplayBuffer). If None, it will be automatically selected.
replay_buffer_kwargs : dict[str, Any] | None — Keyword arguments to pass to the replay buffer on creation.
optimize_memory_usage : bool — Enable a memory efficient variant of the replay buffer at a cost of more complexity. See https://github.com/DLR-RM/stable-baselines3/issues/37#issuecomment-637501195
n_steps : int — When n_step > 1, uses n-step return (with the NStepReplayBuffer) when updating the Q-value network.
ent_coef : str | float — Entropy regularization coefficient. (Equivalent to inverse of reward scale in the original SAC paper.) Controlling exploration/exploitation trade-off. Set it to 'auto' to learn it automatically (and 'auto_0.1' for using 0.1 as initial value)
target_update_interval : int — update the target network every target_network_update_freq gradient steps.
target_entropy : str | float — target entropy when learning ent_coef (ent_coef = 'auto')
use_sde : bool — Whether to use generalized State Dependent Exploration (gSDE) instead of action noise exploration (default: False)
sde_sample_freq : int — Sample a new noise matrix every n steps when using gSDE Default: -1 (only sample at the beginning of the rollout)
use_sde_at_warmup : bool — Whether to use gSDE instead of uniform sampling during the warm up phase (before learning starts)
stats_window_size : int — Window size for the rollout logging, specifying the number of episodes to average the reported success rate, mean episode length, and mean reward over
tensorboard_log — the log location for tensorboard (if None, no logging)
policy_kwargs : dict[str, Any] | None — additional arguments to be passed to the policy on creation. See :ref:sac_policies
verbose : int — Verbosity level: 0 for no output, 1 for info messages (such as device or wrappers used), 2 for debug messages
seed : int | None — Seed for the pseudo random generators
device : torch.device | str — Device (cpu, cuda, ...) on which the code should be run. Setting it to auto, the code will be run on the GPU if possible.
_init_setup_model : bool — Whether or not to build the network at the creation of the instance

Attributes

device : torch.device — Return current device (Lightning-aware).
hparams : Union[AttributeDict, MutableMapping] — meth:save_hyperparameters. It is mutable by the user. For the frozen set of initial hyperparameters, use :attr:hparams_initial.
hparams_initial : AttributeDict — meth:save_hyperparameters. These contents are read-only. Manual updates to the saved hyperparameters can instead be performed through :attr:hparams.
example_input_array : Optional[Union[Tensor, tuple, dict]] — meth:forward method. The return type is interpreted as follows:
current_epoch : int — The current epoch in the Trainer, or 0 if not attached.
global_step : int — Total training batches seen across all epochs.
global_rank : int — The index of the current process across all nodes and devices.
local_rank : int — The index of the current process within a single node.
on_gpu : bool — Returns True if this model is currently located on a GPU.
strict_loading : bool — Determines how Lightning loads this model using .load_state_dict(..., strict=model.strict_loading).
logger : Optional[Union[Logger, FabricLogger]] — Reference to the logger object in the Trainer.
loggers : Union[list[Logger], list[FabricLogger]] — Reference to the list of loggers in the Trainer.
device_mesh : Optional['DeviceMesh'] — Strategies like ModelParallelStrategy will create a device mesh that can be accessed in the :meth:~lightning.pytorch.core.hooks.ModelHooks.configure_model hook to parallelize the LightningModule.
env : VecEnv | None — Return the active training env from the trainer, or None.

Methods

configure_optimizers
training_step — Off-policy SAC step: collect env transitions then do gradient updates.
learn

source method SAC.configure_optimizers() → OptimizerLRScheduler

source method SAC.training_step(batch: Any, batch_idx: int) → None

Off-policy SAC step: collect env transitions then do gradient updates.

source method SAC.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 4, tb_log_name: str = 'SAC', reset_num_timesteps: bool = True, progress_bar: bool = False, default_root_dir: str | Path | None = None) → SelfSAC

source class SACPolicy(observation_space: spaces.Space, action_space: spaces.Box, lr_schedule: Schedule, net_arch: list[int] | dict[str, list[int]] | None = None, activation_fn: type[nn.Module] = nn.ReLU, use_sde: bool = False, log_std_init: float = -3, use_expln: bool = False, clip_mean: float = 2.0, features_extractor_class: type[BaseFeaturesExtractor] = FlattenExtractor, features_extractor_kwargs: dict[str, Any] | None = None, normalize_images: bool = True, optimizer_class: type[torch.optim.Optimizer] = torch.optim.Adam, optimizer_kwargs: dict[str, Any] | None = None, n_critics: int = 2, share_features_extractor: bool = False)

Bases : BasePolicy

Policy class (with both actor and critic) for SAC.

Parameters

observation_space : spaces.Space — Observation space
action_space : spaces.Box — Action space
lr_schedule : Schedule — Learning rate schedule (could be constant)
net_arch : list[int] | dict[str, list[int]] | None — The specification of the policy and value networks.
activation_fn : type[nn.Module] — Activation function
use_sde : bool — Whether to use State Dependent Exploration or not
log_std_init : float — Initial value for the log standard deviation
use_expln : bool — Use expln() function instead of exp() when using gSDE to ensure a positive standard deviation (cf paper). It allows to keep variance above zero and prevent it from growing too fast. In practice, exp() is usually enough.
clip_mean : float — Clip the mean output when using gSDE to avoid numerical instability.
features_extractor_class : type[BaseFeaturesExtractor] — Features extractor to use.
features_extractor_kwargs : dict[str, Any] | None — Keyword arguments to pass to the features extractor.
normalize_images : bool — Whether to normalize images or not, dividing by 255.0 (True by default)
optimizer_class : type[torch.optim.Optimizer] — The optimizer to use, torch.optim.Adam by default
optimizer_kwargs : dict[str, Any] | None — Additional keyword arguments, excluding the learning rate, to pass to the optimizer
n_critics : int — Number of critic networks to create.
share_features_extractor : bool — Whether to share or not the features extractor between the actor and the critic (this saves computation time)

Attributes

device : torch.device — Infer which device this policy lives on by inspecting its parameters.
squash_output : bool — (bool) Getter for squash_output.

Methods

reset_noise — Sample new weights for the exploration matrix, when using gSDE.
make_actor
make_critic
forward
set_training_mode — Put the policy in either training or evaluation mode.

source method SACPolicy.reset_noise(batch_size: int = 1) → None

Sample new weights for the exploration matrix, when using gSDE.

Parameters

batch_size : int
batch_size : int — The number of samples in the batch.

source method SACPolicy.make_actor(features_extractor: BaseFeaturesExtractor | None = None) → Actor

source method SACPolicy.make_critic(features_extractor: BaseFeaturesExtractor | None = None) → ContinuousCritic

source method SACPolicy.forward(obs: PyTorchObs, deterministic: bool = False) → torch.Tensor

source method SACPolicy.set_training_mode(mode: bool) → None

Put the policy in either training or evaluation mode.

This affects certain modules, such as batch normalisation and dropout.

Parameters

mode : bool — if true, set to training mode, else set to evaluation mode