lit_sb3

Classes

A2C — Advantage Actor Critic (A2C).
ARS — Augmented Random Search: https://arxiv.org/abs/1803.07055.
CrossQ — Implementation of Batch Normalization in Deep Reinforcement Learning (CrossQ).
DDPG — Deep Deterministic Policy Gradient (DDPG).
DQN — Deep Q-Network (DQN).
HerReplayBuffer — Hindsight Experience Replay (HER) buffer.
MaskablePPO — Proximal Policy Optimization algorithm (PPO) (clip version) with Invalid Action Masking.
PPO — Proximal Policy Optimization algorithm (PPO) (clip version).
QRDQN — Quantile Regression Deep Q-Network (QR-DQN).
RecurrentPPO — Proximal Policy Optimization algorithm (PPO) (clip version) with support for recurrent policies (LSTM).
RLEnvDataModule — A :class:~lightning.LightningDataModule that holds environments for each RL stage.
RLTrainer — A :class:lightning.Trainer subclass that drives reinforcement-learning runs.
SAC — Soft Actor-Critic (SAC).
TD3 — Twin Delayed DDPG (TD3).
TQC — Controlling Overestimation Bias with Truncated Mixture of Continuous.
TRPO — Trust Region Policy Optimization (TRPO).

source class A2C(policy: str | type[ActorCriticPolicy], learning_rate: float | Schedule = 0.0007, n_steps: int = 5, gamma: float = 0.99, gae_lambda: float = 1.0, ent_coef: float = 0.0, vf_coef: float = 0.5, max_grad_norm: float = 0.5, rms_prop_eps: float = 1e-05, use_rms_prop: bool = True, use_sde: bool = False, sde_sample_freq: int = -1, rollout_buffer_class: type[RolloutBuffer] | None = None, rollout_buffer_kwargs: dict[str, Any] | None = None, normalize_advantage: 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 : OnPolicyAlgorithm

Advantage Actor Critic (A2C).

Paper: https://arxiv.org/abs/1602.01783 Code: This implementation borrows code from https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail and and Stable Baselines (https://github.com/hill-a/stable-baselines)

Introduction to A2C: https://hackernoon.com/intuitive-rl-intro-to-advantage-actor-critic-a2c-4ff545978752

Parameters

policy : str | type[ActorCriticPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — The learning rate, it can be a function of the current progress remaining (from 1 to 0)
n_steps : int — The number of steps to run for each environment per update (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
gamma : float — Discount factor
gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator. Equivalent to classic advantage when set to 1.
ent_coef : float — Entropy coefficient for the loss calculation
vf_coef : float — Value function coefficient for the loss calculation
max_grad_norm : float — The maximum value for the gradient clipping
rms_prop_eps : float — RMSProp epsilon. It stabilizes square root computation in denominator of RMSProp update
use_rms_prop : bool — Whether to use RMSprop (default) or Adam as optimizer
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)
rollout_buffer_class : type[RolloutBuffer] | None — Rollout buffer class to use. If None, it will be automatically selected.
rollout_buffer_kwargs : dict[str, Any] | None — Keyword arguments to pass to the rollout buffer on creation.
normalize_advantage : bool — Whether to normalize or not the advantage
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:a2c_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.
automatic_optimization : bool — If set to False you are responsible for calling .backward(), .step(), .zero_grad().
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
on_train_epoch_start
training_step — A2C loss on the full rollout buffer (single gradient step).
on_train_epoch_end
learn

source method A2C.configure_optimizers() → OptimizerLRScheduler

source method A2C.on_train_epoch_start() → None

source method A2C.training_step(batch: Any, batch_idx: int) → torch.Tensor

A2C loss on the full rollout buffer (single gradient step).

source method A2C.on_train_epoch_end() → None

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

source class ARS(policy: str | type[ARSPolicy], n_delta: int = 8, n_top: int | None = None, learning_rate: float | Schedule = 0.02, delta_std: float | Schedule = 0.05, zero_policy: bool = True, alive_bonus_offset: float = 0, n_eval_episodes: int = 1, policy_kwargs: dict[str, Any] | None = None, stats_window_size: int = 100, seed: int | None = None, verbose: int = 0, device: torch.device | str = 'cpu', _init_setup_model: bool = True)

Bases : BaseAlgorithm

Augmented Random Search: https://arxiv.org/abs/1803.07055.

Original implementation: https://github.com/modestyachts/ARS C++/Cuda Implementation: https://github.com/google-research/tiny-differentiable-simulator/ 150 LOC Numpy Implementation: https://github.com/alexis-jacq/numpy_ARS/blob/master/asr.py

Parameters

policy : str | type[ARSPolicy] — The policy to train, can be an instance of ARSPolicy, or a string from ["LinearPolicy", "MlpPolicy"]
env — The environment to train on, may be a string if registered with gym
n_delta : int — How many random perturbations of the policy to try at each update step.
n_top : int | None — How many of the top delta to use in each update step. Default is n_delta
learning_rate : float | Schedule — Float or schedule for the step size
delta_std : float | Schedule — Float or schedule for the exploration noise
zero_policy : bool — Boolean determining if the passed policy should have it's weights zeroed before training.
alive_bonus_offset : float — Constant added to the reward at each step, used to cancel out alive bonuses.
n_eval_episodes : int — Number of episodes to evaluate each candidate.
policy_kwargs : dict[str, Any] | None — Keyword arguments to pass to the policy on creation. See :ref:ars_policies
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
seed : int | None — Random seed for the training
verbose : int — Verbosity level: 0 no output, 1 info, 2 debug
device : torch.device | str — Torch device to use for training, defaults to "cpu"
_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.
automatic_optimization : bool — If set to False you are responsible for calling .backward(), .step(), .zero_grad().
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

evaluate_candidates — Evaluate each candidate.
dump_logs — Dump information to the logger.
learn — Return a trained model.
set_parameters

source method ARS.evaluate_candidates(candidate_weights: torch.Tensor, callback: BaseCallback, async_eval: AsyncEval | None) → torch.Tensor

Evaluate each candidate.

Parameters

candidate_weights : torch.Tensor — The candidate weights to be evaluated.
callback : BaseCallback — Callback that will be called at each step (or after evaluation in the multiprocess version)
async_eval : AsyncEval | None — The object for asynchronous evaluation of candidates.

Returns

torch.Tensor — The episodic return for each candidate.

source method ARS.dump_logs() → None

Dump information to the logger.

source method ARS.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 1, tb_log_name: str = 'ARS', reset_num_timesteps: bool = True, async_eval: AsyncEval | None = None, progress_bar: bool = False, default_root_dir: str | Path | None = None) → SelfARS

Return a trained model.

Parameters

total_timesteps : int — The total number of samples (env steps) to train on
env : GymEnv | str | None — Training environment (gym env, env-id string, or VecEnv).
callback : MaybeCallback — callback(s) called at every step with state of the algorithm.
log_every_n_steps : int | None — The number of timesteps before logging.
tb_log_name : str — the name of the run for TensorBoard logging
reset_num_timesteps : bool — whether or not to reset the current timestep number (used in logging)
async_eval : AsyncEval | None — The object for asynchronous evaluation of candidates.
progress_bar : bool — Display a progress bar using tqdm and rich.
default_root_dir : str | Path | None — Unused for ARS; accepted for API compatibility.

Returns

SelfARS — the trained model

Raises

TypeError

source method ARS.set_parameters(load_path_or_dict: str | dict[str, dict], exact_match: bool = True, device: torch.device | str = 'auto') → None

source class CrossQ(policy: str | type[CrossQPolicy], learning_rate: float | Schedule = 0.001, buffer_size: int = 1000000, learning_starts: int = 100, batch_size: int = 256, 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_entropy: str | float = 'auto', policy_delay: int = 3, 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

Implementation of Batch Normalization in Deep Reinforcement Learning (CrossQ).

Paper: https://openreview.net/forum?id=PczQtTsTIX Reference implementation: https://github.com/araffin/sbx.

Parameters

policy : str | type[CrossQPolicy] — The policy model to use (MlpPolicy)
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
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_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:crossq_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 CrossQ step: collect env transitions then do gradient updates.
learn

source method CrossQ.configure_optimizers() → OptimizerLRScheduler

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

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

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

source class DDPG(policy: str | type[TD3Policy], learning_rate: float | Schedule = 0.001, 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, 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 : TD3

Deep Deterministic Policy Gradient (DDPG).

Deterministic Policy Gradient: http://proceedings.mlr.press/v32/silver14.pdf DDPG Paper: https://arxiv.org/abs/1509.02971 Introduction to DDPG: https://spinningup.openai.com/en/latest/algorithms/ddpg.html

Note: we treat DDPG as a special case of its successor TD3.

Parameters

policy : str | type[TD3Policy] — 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.
policy_kwargs : dict[str, Any] | None — additional arguments to be passed to the policy on creation. See :ref:ddpg_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

learn

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

source class DQN(policy: str | type[DQNPolicy], learning_rate: float | Schedule = 0.0001, buffer_size: int = 1000000, learning_starts: int = 100, batch_size: int = 32, tau: float = 1.0, gamma: float = 0.99, train_freq: int | tuple[int, str] = 4, gradient_steps: int = 1, replay_buffer_class: type[ReplayBuffer] | None = None, replay_buffer_kwargs: dict[str, Any] | None = None, optimize_memory_usage: bool = False, n_steps: int = 1, target_update_interval: int = 10000, exploration_fraction: float = 0.1, exploration_initial_eps: float = 1.0, exploration_final_eps: float = 0.05, max_grad_norm: float = 10, 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

Deep Q-Network (DQN).

Paper: https://arxiv.org/abs/1312.5602, https://www.nature.com/articles/nature14236 Default hyperparameters are taken from the Nature paper, except for the optimizer and learning rate that were taken from Stable Baselines defaults.

Parameters

policy : str | type[DQNPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — The learning rate, 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) default 1 for hard update
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.
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.
target_update_interval : int — update the target network every target_update_interval environment steps.
exploration_fraction : float — fraction of entire training period over which the exploration rate is reduced
exploration_initial_eps : float — initial value of random action probability
exploration_final_eps : float — final value of random action probability
max_grad_norm : float — The maximum value for the gradient clipping
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:dqn_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 DQN step: collect env transitions then do gradient updates.
predict — Overrides the base_class predict function to include epsilon-greedy exploration.
learn

source method DQN.configure_optimizers() → OptimizerLRScheduler

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

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

source method DQN.predict(observation: np.ndarray | dict[str, np.ndarray], state: tuple[np.ndarray, ...] | None = None, episode_start: np.ndarray | None = None, deterministic: bool = False) → tuple[np.ndarray, tuple[np.ndarray, ...] | None]

Overrides the base_class predict function to include epsilon-greedy exploration.

Parameters

observation : np.ndarray | dict[str, np.ndarray] — the input observation
state : tuple[np.ndarray, ...] | None — The last states (can be None, used in recurrent policies)
episode_start : np.ndarray | None — The last masks (can be None, used in recurrent policies)
deterministic : bool — Whether or not to return deterministic actions.

Returns

tuple[np.ndarray, tuple[np.ndarray, ...] | None] — the model's action and the next state (used in recurrent policies)

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

source class HerReplayBuffer(buffer_size: int, observation_space: spaces.Dict, action_space: spaces.Space, env: VecEnv, device: torch.device | str = 'auto', n_envs: int = 1, optimize_memory_usage: bool = False, handle_timeout_termination: bool = True, n_sampled_goal: int = 4, goal_selection_strategy: GoalSelectionStrategy | str = 'future', copy_info_dict: bool = False)

Bases : DictReplayBuffer

Hindsight Experience Replay (HER) buffer.

Paper: https://arxiv.org/abs/1707.01495

Replay buffer for sampling HER (Hindsight Experience Replay) transitions.

.. note::

Compared to other implementations, the future goal sampling strategy is inclusive: the current transition can be used when re-sampling.

Parameters

buffer_size : int — Max number of element in the buffer
observation_space : spaces.Dict — Observation space
action_space : spaces.Space — Action space
env : VecEnv — The training environment
device : torch.device | str — PyTorch device
n_envs : int — Number of parallel environments
optimize_memory_usage : bool — Enable a memory efficient variant Disabled for now (see https://github.com/DLR-RM/stable-baselines3/pull/243#discussion_r531535702)
handle_timeout_termination : bool — Handle timeout termination (due to timelimit) separately and treat the task as infinite horizon task. https://github.com/DLR-RM/stable-baselines3/issues/284
n_sampled_goal : int — Number of virtual transitions to create per real transition, by sampling new goals.
goal_selection_strategy : GoalSelectionStrategy | str — Strategy for sampling goals for replay. One of ['episode', 'final', 'future']
copy_info_dict : bool — Whether to copy the info dictionary and pass it to compute_reward() method. Please note that the copy may cause a slowdown. False by default.

Methods

set_env — Sets the environment.
add
sample — Sample elements from the replay buffer.
truncate_last_trajectory — If called, we assume that the last trajectory in the replay buffer was finished.

source method HerReplayBuffer.set_env(env: VecEnv) → None

Sets the environment.

Parameters

env : VecEnv
env : VecEnv — The environment.

Raises

ValueError

source method HerReplayBuffer.add(obs: dict[str, np.ndarray], next_obs: dict[str, np.ndarray], action: np.ndarray, reward: np.ndarray, done: np.ndarray, infos: list[dict[str, Any]]) → None

source method HerReplayBuffer.sample(batch_size: int, env: VecNormalize | None = None) → DictReplayBufferSamples

Sample elements from the replay buffer.

Parameters

batch_size : int — Number of element to sample
env : VecNormalize | None — Associated VecEnv to normalize the observations/rewards when sampling

Returns

DictReplayBufferSamples — Samples

Raises

RuntimeError

source method HerReplayBuffer.truncate_last_trajectory() → None

If called, we assume that the last trajectory in the replay buffer was finished.

(and truncate it). If not called, we assume that we continue the same trajectory (same episode).

source class MaskablePPO(policy: str | type[MaskableActorCriticPolicy], learning_rate: float | Schedule = 0.0003, n_steps: int = 2048, batch_size: int | None = 64, n_epochs: int = 10, gamma: float = 0.99, gae_lambda: float = 0.95, clip_range: float | Schedule = 0.2, clip_range_vf: None | float | Schedule = None, normalize_advantage: bool = True, ent_coef: float = 0.0, vf_coef: float = 0.5, max_grad_norm: float = 0.5, rollout_buffer_class: type[RolloutBuffer] | None = None, rollout_buffer_kwargs: dict[str, Any] | None = None, target_kl: float | None = None, 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 : OnPolicyAlgorithm

Proximal Policy Optimization algorithm (PPO) (clip version) with Invalid Action Masking.

Based on the original Stable Baselines 3 implementation.

Introduction to PPO: https://spinningup.openai.com/en/latest/algorithms/ppo.html Background on Invalid Action Masking: https://arxiv.org/abs/2006.14171

Parameters

policy : str | type[MaskableActorCriticPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — The learning rate, it can be a function of the current progress remaining (from 1 to 0)
n_steps : int — The number of steps to run for each environment per update (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
batch_size : int | None — Minibatch size
n_epochs : int — Number of epoch when optimizing the surrogate loss
gamma : float — Discount factor
gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator
clip_range : float | Schedule — Clipping parameter, it can be a function of the current progress remaining (from 1 to 0).
clip_range_vf : None | float | Schedule — Clipping parameter for the value function, it can be a function of the current progress remaining (from 1 to 0). This is a parameter specific to the OpenAI implementation. If None is passed (default), no clipping will be done on the value function. IMPORTANT: this clipping depends on the reward scaling.
normalize_advantage : bool — Whether to normalize or not the advantage
ent_coef : float — Entropy coefficient for the loss calculation
vf_coef : float — Value function coefficient for the loss calculation
max_grad_norm : float — The maximum value for the gradient clipping
target_kl : float | None — Limit the KL divergence between updates, because the clipping is not enough to prevent large update see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213) By default, there is no limit on the kl div.
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:ppo_mask_policies
verbose : int — the verbosity level: 0 no output, 1 info, 2 debug
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.
automatic_optimization : bool — If set to False you are responsible for calling .backward(), .step(), .zero_grad().
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

collect_rollouts — Collect experiences using the current policy and fill a RolloutBuffer.
predict — Get the policy action from an observation (and optional hidden state).
on_train_epoch_start — Collect rollouts with masking, then set policy to train mode.
training_step — Compute MaskablePPO clipped-surrogate loss for one mini-batch.
on_train_epoch_end — Log epoch-level stats.
learn

source method MaskablePPO.collect_rollouts(env: VecEnv, callback: BaseCallback, rollout_buffer: RolloutBuffer, n_rollout_steps: int, use_masking: bool = True) → bool

Collect experiences using the current policy and fill a RolloutBuffer.

The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning.

This method is largely identical to the implementation found in the parent class.

Parameters

env : VecEnv — The training environment
callback : BaseCallback — Callback that will be called at each step (and at the beginning and end of the rollout)
rollout_buffer : RolloutBuffer — Buffer to fill with rollouts
n_steps — Number of experiences to collect per environment
use_masking : bool — Whether or not to use invalid action masks during training
n_rollout_steps : int — Number of steps to rollout.

Returns

bool — True if function returned with at least n_rollout_steps collected, False if callback terminated rollout prematurely.

Raises

ValueError

source method MaskablePPO.predict(observation: np.ndarray | dict[str, np.ndarray], state: tuple[np.ndarray, ...] | None = None, episode_start: np.ndarray | None = None, deterministic: bool = False, action_masks: np.ndarray | None = None) → tuple[np.ndarray, tuple[np.ndarray, ...] | None]

Get the policy action from an observation (and optional hidden state).

Includes sugar-coating to handle different observations (e.g. normalizing images).

Parameters

observation : np.ndarray | dict[str, np.ndarray] — the input observation
state : tuple[np.ndarray, ...] | None — The last hidden states (can be None, used in recurrent policies)
episode_start : np.ndarray | None — The last masks (can be None, used in recurrent policies) this correspond to beginning of episodes, where the hidden states of the RNN must be reset.
deterministic : bool — Whether or not to return deterministic actions.
action_masks : np.ndarray | None — Boolean masks for invalid actions (True = valid).

Returns

tuple[np.ndarray, tuple[np.ndarray, ...] | None] — the model's action and the next hidden state (used in recurrent policies)

source method MaskablePPO.on_train_epoch_start() → None

Collect rollouts with masking, then set policy to train mode.

source method MaskablePPO.training_step(batch: Any, batch_idx: int) → torch.Tensor

Compute MaskablePPO clipped-surrogate loss for one mini-batch.

source method MaskablePPO.on_train_epoch_end() → None

Log epoch-level stats.

source method MaskablePPO.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 1, tb_log_name: str = 'MaskablePPO', reset_num_timesteps: bool = True, use_masking: bool = True, progress_bar: bool = False) → SelfMaskablePPO

source class PPO(policy: str | type[ActorCriticPolicy], learning_rate: float | Schedule = 0.0003, n_steps: int = 2048, batch_size: int = 64, n_epochs: int = 10, gamma: float = 0.99, gae_lambda: float = 0.95, clip_range: float | Schedule = 0.2, clip_range_vf: None | float | Schedule = None, normalize_advantage: bool = True, ent_coef: float = 0.0, vf_coef: float = 0.5, max_grad_norm: float = 0.5, use_sde: bool = False, sde_sample_freq: int = -1, rollout_buffer_class: type[RolloutBuffer] | None = None, rollout_buffer_kwargs: dict[str, Any] | None = None, target_kl: float | None = None, 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 : OnPolicyAlgorithm

Proximal Policy Optimization algorithm (PPO) (clip version).

Paper: https://arxiv.org/abs/1707.06347 Code: This implementation borrows code from OpenAI Spinning Up (https://github.com/openai/spinningup/) https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail and Stable Baselines (PPO2 from https://github.com/hill-a/stable-baselines)

Introduction to PPO: https://spinningup.openai.com/en/latest/algorithms/ppo.html

Parameters

policy : str | type[ActorCriticPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — The learning rate, it can be a function of the current progress remaining (from 1 to 0)
n_steps : int — The number of steps to run for each environment per update (i.e. rollout buffer size is n_steps * n_envs where n_envs is number of environment copies running in parallel) NOTE: n_steps * n_envs must be greater than 1 (because of the advantage normalization) See https://github.com/pytorch/pytorch/issues/29372
batch_size : int — Minibatch size
n_epochs : int — Number of epoch when optimizing the surrogate loss
gamma : float — Discount factor
gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator
clip_range : float | Schedule — Clipping parameter, it can be a function of the current progress remaining (from 1 to 0).
clip_range_vf : None | float | Schedule — Clipping parameter for the value function, it can be a function of the current progress remaining (from 1 to 0). This is a parameter specific to the OpenAI implementation. If None is passed (default), no clipping will be done on the value function. IMPORTANT: this clipping depends on the reward scaling.
normalize_advantage : bool — Whether to normalize or not the advantage
ent_coef : float — Entropy coefficient for the loss calculation
vf_coef : float — Value function coefficient for the loss calculation
max_grad_norm : float — The maximum value for the gradient clipping
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)
rollout_buffer_class : type[RolloutBuffer] | None — Rollout buffer class to use. If None, it will be automatically selected.
rollout_buffer_kwargs : dict[str, Any] | None — Keyword arguments to pass to the rollout buffer on creation
target_kl : float | None — Limit the KL divergence between updates, because the clipping is not enough to prevent large update see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213) By default, there is no limit on the kl div.
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:ppo_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.
automatic_optimization : bool — If set to False you are responsible for calling .backward(), .step(), .zero_grad().
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 — Return the policy optimizer.
on_train_epoch_start — Run base rollout collection hook then set policy to train mode.
training_step — Compute PPO clipped-surrogate loss for one mini-batch of rollout data.
on_train_epoch_end — Log epoch-level PPO stats.
learn

source method PPO.configure_optimizers() → OptimizerLRScheduler

Return the policy optimizer.

source method PPO.on_train_epoch_start() → None

Run base rollout collection hook then set policy to train mode.

source method PPO.training_step(batch: Any, batch_idx: int) → torch.Tensor

Compute PPO clipped-surrogate loss for one mini-batch of rollout data.

source method PPO.on_train_epoch_end() → None

Log epoch-level PPO stats.

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

source class QRDQN(policy: str | type[QRDQNPolicy], learning_rate: float | Schedule = 5e-05, buffer_size: int = 1000000, learning_starts: int = 100, batch_size: int = 32, tau: float = 1.0, gamma: float = 0.99, train_freq: int | tuple[int, str] = 4, gradient_steps: int = 1, replay_buffer_class: type[ReplayBuffer] | None = None, replay_buffer_kwargs: dict[str, Any] | None = None, optimize_memory_usage: bool = False, n_steps: int = 1, target_update_interval: int = 10000, exploration_fraction: float = 0.005, exploration_initial_eps: float = 1.0, exploration_final_eps: float = 0.01, max_grad_norm: float | None = None, 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

Quantile Regression Deep Q-Network (QR-DQN).

Paper: https://arxiv.org/abs/1710.10044.

Adapted from sb3-contrib to use PyTorch Lightning.

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 QRDQN step: collect env transitions then do gradient updates.
predict — Epsilon-greedy action selection for QRDQN.
learn

source method QRDQN.configure_optimizers() → OptimizerLRScheduler

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

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

source method QRDQN.predict(observation: np.ndarray | dict[str, np.ndarray], state: tuple[np.ndarray, ...] | None = None, episode_start: np.ndarray | None = None, deterministic: bool = False) → tuple[np.ndarray, tuple[np.ndarray, ...] | None]

Epsilon-greedy action selection for QRDQN.

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

source class RecurrentPPO(policy: str | type[RecurrentActorCriticPolicy], learning_rate: float | Schedule = 0.0003, n_steps: int = 128, batch_size: int | None = 128, n_epochs: int = 10, gamma: float = 0.99, gae_lambda: float = 0.95, clip_range: float | Schedule = 0.2, clip_range_vf: None | float | Schedule = None, normalize_advantage: bool = True, ent_coef: float = 0.0, vf_coef: float = 0.5, max_grad_norm: float = 0.5, use_sde: bool = False, sde_sample_freq: int = -1, target_kl: float | None = None, 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 : OnPolicyAlgorithm

Proximal Policy Optimization algorithm (PPO) (clip version) with support for recurrent policies (LSTM).

Based on the original Stable Baselines 3 implementation.

Introduction to PPO: https://spinningup.openai.com/en/latest/algorithms/ppo.html

Parameters

policy : str | type[RecurrentActorCriticPolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str)
learning_rate : float | Schedule — The learning rate, it can be a function of the current progress remaining (from 1 to 0)
n_steps : int — The number of steps to run for each environment per update (i.e. batch size is n_steps * n_env where n_env is number of environment copies running in parallel)
batch_size : int | None — Minibatch size
n_epochs : int — Number of epoch when optimizing the surrogate loss
gamma : float — Discount factor
gae_lambda : float — Factor for trade-off of bias vs variance for Generalized Advantage Estimator
clip_range : float | Schedule — Clipping parameter, it can be a function of the current progress remaining (from 1 to 0).
clip_range_vf : None | float | Schedule — Clipping parameter for the value function, it can be a function of the current progress remaining (from 1 to 0). This is a parameter specific to the OpenAI implementation. If None is passed (default), no clipping will be done on the value function. IMPORTANT: this clipping depends on the reward scaling.
normalize_advantage : bool — Whether to normalize or not the advantage
ent_coef : float — Entropy coefficient for the loss calculation
vf_coef : float — Value function coefficient for the loss calculation
max_grad_norm : float — The maximum value for the gradient clipping
target_kl : float | None — Limit the KL divergence between updates, because the clipping is not enough to prevent large update see issue #213 (cf https://github.com/hill-a/stable-baselines/issues/213) By default, there is no limit on the kl div.
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:ppo_recurrent_policies
verbose : int — the verbosity level: 0 no output, 1 info, 2 debug
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.
automatic_optimization : bool — If set to False you are responsible for calling .backward(), .step(), .zero_grad().
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

collect_rollouts — Collect experiences using the current policy and fill a RolloutBuffer.
train — Update policy using the currently gathered rollout buffer.
learn

source method RecurrentPPO.collect_rollouts(env: VecEnv, callback: BaseCallback, rollout_buffer: RolloutBuffer, n_rollout_steps: int) → bool

Collect experiences using the current policy and fill a RolloutBuffer.

The term rollout here refers to the model-free notion and should not be used with the concept of rollout used in model-based RL or planning.

Parameters

env : VecEnv — The training environment
callback : BaseCallback — Callback that will be called at each step (and at the beginning and end of the rollout)
rollout_buffer : RolloutBuffer — Buffer to fill with rollouts
n_steps — Number of experiences to collect per environment
n_rollout_steps : int — Number of steps to rollout.

Returns

bool — True if function returned with at least n_rollout_steps collected, False if callback terminated rollout prematurely.

source method RecurrentPPO.train() → None

Update policy using the currently gathered rollout buffer.

source method RecurrentPPO.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 1, tb_log_name: str = 'RecurrentPPO', reset_num_timesteps: bool = True, progress_bar: bool = False) → SelfRecurrentPPO

source class RLEnvDataModule(train_env: GymEnv | str | None = None, val_env: GymEnv | str | None = None, test_env: GymEnv | str | None = None, predict_env: GymEnv | str | None = None, monitor_wrapper: bool = True, verbose: int = 0)

Bases : L.LightningDataModule

A :class:~lightning.LightningDataModule that holds environments for each RL stage.

Instead of :meth:train_dataloader, :meth:val_dataloader, etc., this module exposes :attr:train_env, :attr:val_env, :attr:test_env, and :attr:predict_env. The :class:~lit_sb3.trainer.RLTrainer extracts the appropriate environment from the datamodule at each stage.

Parameters

train_env : GymEnv | str | None — Environment (or env-id string) used for training.
val_env : GymEnv | str | None — Environment (or env-id string) used for validation / evaluation.
test_env : GymEnv | str | None — Environment (or env-id string) used for testing / evaluation.
predict_env : GymEnv | str | None — Environment (or env-id string) used for prediction.
monitor_wrapper : bool — Whether to wrap raw Gym envs in a Monitor wrapper.
verbose : int — Verbosity level forwarded to env wrapping utilities.

Attributes

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.

Methods

setup — Resolve and wrap environments for the requested stage.
teardown — Release resolved environments for the requested stage.
on_exception — Ensure vectorized envs are closed when Lightning aborts the run.
get_env — Return the environment configured for stage.
train_dataloader — RLEnvDataModule exposes envs, not training dataloaders.
val_dataloader — RLEnvDataModule exposes envs, not validation dataloaders.
test_dataloader — RLEnvDataModule exposes envs, not test dataloaders.
predict_dataloader — RLEnvDataModule exposes envs, not predict dataloaders.
state_dict — Persist serializable datamodule config and env specs, never live env instances.
load_state_dict — Restore serializable env specs and datamodule config from checkpoint state.
load_from_checkpoint — Load the datamodule and then restore env specs through :meth:load_state_dict.
close — Close all resolved environments and free resources.

source method RLEnvDataModule.setup(stage: str | None = None) → None

Resolve and wrap environments for the requested stage.

Called automatically by Lightning before each stage, or can be called manually.

source method RLEnvDataModule.teardown(stage: str | None = None) → None

Release resolved environments for the requested stage.

source method RLEnvDataModule.on_exception(exception: BaseException) → None

Ensure vectorized envs are closed when Lightning aborts the run.

source method RLEnvDataModule.get_env(stage: Literal['fit', 'validate', 'test', 'predict'], *, resolve: bool = True) → VecEnv | None

Return the environment configured for stage.

Parameters

stage : Literal['fit', 'validate', 'test', 'predict'] — Lightning stage name.
resolve : bool — Whether to call :meth:setup before reading the stage env.

source method RLEnvDataModule.train_dataloader() → None

RLEnvDataModule exposes envs, not training dataloaders.

Raises

TypeError

source method RLEnvDataModule.val_dataloader() → None

RLEnvDataModule exposes envs, not validation dataloaders.

Raises

TypeError

source method RLEnvDataModule.test_dataloader() → None

RLEnvDataModule exposes envs, not test dataloaders.

Raises

TypeError

source method RLEnvDataModule.predict_dataloader() → None

RLEnvDataModule exposes envs, not predict dataloaders.

Raises

TypeError

source method RLEnvDataModule.state_dict() → dict[str, Any]

Persist serializable datamodule config and env specs, never live env instances.

source method RLEnvDataModule.load_state_dict(state_dict: dict[str, Any]) → None

Restore serializable env specs and datamodule config from checkpoint state.

source classmethod RLEnvDataModule.load_from_checkpoint(checkpoint_path: Any, map_location: Any = None, hparams_file: Any = None, weights_only: bool | None = None, **kwargs: Any) → RLEnvDataModule

Load the datamodule and then restore env specs through :meth:load_state_dict.

source method RLEnvDataModule.close() → None

Close all resolved environments and free resources.

Bases : L.Trainer

A :class:lightning.Trainer subclass that drives reinforcement-learning runs.

All standard lightning.Trainer keyword arguments are accepted and forwarded to the parent class, giving full access to Lightning features (mixed precision, custom loggers, DDP, plugins, …).

The RL-specific arguments control the training schedule and SB3 callbacks. accelerator, enable_model_summary, enable_checkpointing, and logger have RL-appropriate defaults but can be overridden.

Parameters

max_steps : int — Total environment steps to collect (RL timesteps budget). Lightning step-based stopping is disabled; this budget is enforced using RL env-step accounting instead.
rl_callback : MaybeCallback — SB3-style callback(s) fired at each RL step.
run_name : str — The run name used in loggers.

Attributes

device_ids : list[int] — List of device indexes per node.
num_devices : int — Number of devices the trainer uses per node.
model : Optional[torch.nn.Module] — The LightningModule, but possibly wrapped into DataParallel or DistributedDataParallel.
log_dir : Optional[str] — The directory for the current experiment. Use this to save images to, etc...
is_global_zero : bool — Whether this process is the global zero in multi-node training.
enable_validation : bool — Check if we should run validation during training.
default_root_dir : str — The default location to save artifacts of loggers, checkpoints etc.
early_stopping_callback : Optional[EarlyStopping] — class:~lightning.pytorch.callbacks.early_stopping.EarlyStopping callback in the Trainer.callbacks list, or None if it doesn't exist.
early_stopping_callbacks : list[EarlyStopping] — class:~lightning.pytorch.callbacks.early_stopping.EarlyStopping found in the Trainer.callbacks list.
checkpoint_callback : Optional[Checkpoint] — class:~lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint callback in the Trainer.callbacks list, or None if it doesn't exist.
checkpoint_callbacks : list[Checkpoint] — class:~lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint found in the Trainer.callbacks list.
progress_bar_callback : Optional[ProgressBar] — class:~lightning.pytorch.callbacks.progress.progress_bar.ProgressBar found in the Trainer.callbacks list, or None if one doesn't exist.
ckpt_path : Optional[_PATH] — meth:~lightning.pytorch.trainer.trainer.Trainer.fit, :meth:~lightning.pytorch.trainer.trainer.Trainer.validate, :meth:~lightning.pytorch.trainer.trainer.Trainer.test, or :meth:~lightning.pytorch.trainer.trainer.Trainer.predict.
sanity_checking : bool — Whether sanity checking is running.
received_sigterm : bool — Whether a signal.SIGTERM signal was received.
global_step : int — The number of optimizer steps taken (does not reset each epoch).
current_epoch : int — The current epoch, updated after the epoch end hooks are run.
is_last_batch : bool — Whether trainer is executing the last batch.
train_dataloader : Optional[TRAIN_DATALOADERS] — The training dataloader(s) used during trainer.fit().
val_dataloaders : Optional[EVAL_DATALOADERS] — The validation dataloader(s) used during trainer.fit() or trainer.validate().
test_dataloaders : Optional[EVAL_DATALOADERS] — The test dataloader(s) used during trainer.test().
predict_dataloaders : Optional[EVAL_DATALOADERS] — The prediction dataloader(s) used during trainer.predict().
num_training_batches : Union[int, float] — The number of training batches that will be used during trainer.fit().
num_sanity_val_batches : list[Union[int, float]] — The number of validation batches that will be used during the sanity-checking part of trainer.fit().
num_val_batches : list[Union[int, float]] — The number of validation batches that will be used during trainer.fit() or trainer.validate().
num_test_batches : list[Union[int, float]] — The number of test batches that will be used during trainer.test().
num_predict_batches : list[Union[int, float]] — The number of prediction batches that will be used during trainer.predict().
logger : Optional[Logger] — class:~lightning.pytorch.loggers.logger.Logger being used.
loggers : list[Logger] — class:~lightning.pytorch.loggers.logger.Logger used.
callback_metrics : _OUT_DICT — The metrics available to callbacks.
logged_metrics : _OUT_DICT — The metrics sent to the loggers.
progress_bar_metrics : _PBAR_DICT — The metrics sent to the progress bar.
estimated_stepping_batches : Union[int, float] — The estimated number of batches that will optimizer.step() during training.
rl_total_timesteps : int — Return the trainer-owned RL timestep budget.
rl_num_timesteps : int — Return the latest observed RL env-step count.
rl_budget_reached : bool — Return whether the trainer-owned RL timestep budget has been reached.
rl_progress : float — Return RL budget progress as a value in [0, 1].

Methods

fit — Train model for :attr:_rl_total_timesteps environment steps.
validate
test
predict

source property RLTrainer.rl_total_timesteps: int

Return the trainer-owned RL timestep budget.

source property RLTrainer.rl_num_timesteps: int

Return the latest observed RL env-step count.

source property RLTrainer.rl_budget_reached: bool

Return whether the trainer-owned RL timestep budget has been reached.

source property RLTrainer.rl_progress: float

Return RL budget progress as a value in [0, 1].

source method RLTrainer.fit(model: BaseAlgorithm, *, datamodule: Any = None, train_dataloaders: Any = None, val_dataloaders: Any = None, ckpt_path: str | Path | None = None, weights_only: bool | None = None) → BaseAlgorithm

Train model for :attr:_rl_total_timesteps environment steps.

The environment is provided via an :class:~lit_sb3.datamodule.RLEnvDataModule.

Parameters

model : BaseAlgorithm — The RL algorithm (a LightningModule subclass).
datamodule : Any — An :class:~lit_sb3.datamodule.RLEnvDataModule that provides the training (and optionally evaluation) environments.
train_dataloaders : Any — Unsupported for RLTrainer.
val_dataloaders : Any — Unsupported for RLTrainer.
ckpt_path : str | Path | None — Optional Lightning checkpoint path.
weights_only : bool | None — Optional Lightning checkpoint loading flag.

Returns

BaseAlgorithm — The trained model (same object, for chaining).

Raises

TypeError
ValueError

Raises

TypeError

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 TD3(policy: str | type[TD3Policy], learning_rate: float | Schedule = 0.001, 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, policy_delay: int = 2, target_policy_noise: float = 0.2, target_noise_clip: float = 0.5, 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

Twin Delayed DDPG (TD3).

Addressing Function Approximation Error in Actor-Critic Methods.

Original implementation: https://github.com/sfujim/TD3 Paper: https://arxiv.org/abs/1802.09477 Introduction to TD3: https://spinningup.openai.com/en/latest/algorithms/td3.html

Parameters

policy : str | type[TD3Policy] — 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.
policy_delay : int — Policy and target networks will only be updated once every policy_delay steps per training steps. The Q values will be updated policy_delay more often (update every training step).
target_policy_noise : float — Standard deviation of Gaussian noise added to target policy (smoothing noise)
target_noise_clip : float — Limit for absolute value of target policy smoothing noise.
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:td3_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 TD3 step: collect env transitions then do gradient updates.
learn

source method TD3.configure_optimizers() → OptimizerLRScheduler

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

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

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

source class TQC(policy: str | type[TQCPolicy], learning_rate: float = 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', top_quantiles_to_drop_per_net: int = 2, 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

Controlling Overestimation Bias with Truncated Mixture of Continuous.

Distributional Quantile Critics (TQC).

Paper: https://arxiv.org/abs/2005.04269 This implementation adapts TQC from sb3-contrib to use PyTorch Lightning.

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 TQC step: collect env transitions then do gradient updates.
learn

source method TQC.configure_optimizers() → OptimizerLRScheduler

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

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

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

source class TRPO(policy: str | type[ActorCriticPolicy], learning_rate: float | Schedule = 0.001, n_steps: int = 2048, batch_size: int = 128, gamma: float = 0.99, cg_max_steps: int = 15, cg_damping: float = 0.1, line_search_shrinking_factor: float = 0.8, line_search_max_iter: int = 10, n_critic_updates: int = 10, gae_lambda: float = 0.95, use_sde: bool = False, sde_sample_freq: int = -1, rollout_buffer_class: type[RolloutBuffer] | None = None, rollout_buffer_kwargs: dict[str, Any] | None = None, normalize_advantage: bool = True, target_kl: float = 0.01, sub_sampling_factor: int = 1, 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 : OnPolicyAlgorithm

Trust Region Policy Optimization (TRPO).

Paper: https://arxiv.org/abs/1502.05477 Adapted from sb3-contrib to use PyTorch Lightning.

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
on_train_epoch_start — Configure TRPO to process full buffer in one pass, then collect rollout.
training_step — TRPO update: CG actor update + critic gradient descent.
hessian_vector_product — Compute the matrix-vector product with the Fisher information matrix.
learn

source method TRPO.configure_optimizers() → OptimizerLRScheduler

source method TRPO.on_train_epoch_start() → None

Configure TRPO to process full buffer in one pass, then collect rollout.

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

TRPO update: CG actor update + critic gradient descent.

source method TRPO.hessian_vector_product(params: list[nn.Parameter], grad_kl: torch.Tensor, vector: torch.Tensor, retain_graph: bool = True) → torch.Tensor

Compute the matrix-vector product with the Fisher information matrix.

source method TRPO.learn(total_timesteps: int, env: GymEnv | str | None = None, callback: MaybeCallback = None, log_every_n_steps: int | None = 1, tb_log_name: str = 'TRPO', reset_num_timesteps: bool = True, progress_bar: bool = False) → SelfTRPO