lit_sb3.algorithms

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).
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).
BaseAlgorithm — Base class for RL algorithms, adapted to use PyTorch Lightning.
OffPolicyAlgorithm — Base class for off-policy algorithms (SAC, TD3, DQN, DDPG, …).
OnPolicyAlgorithm — Base class for on-policy algorithms (PPO, A2C, 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 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

source class BaseAlgorithm(policy: str | type[BasePolicy], learning_rate: float | Schedule, policy_kwargs: dict[str, Any] | None = None, stats_window_size: int = 100, verbose: int = 0, device: torch.device | str = 'auto', support_multi_env: bool = False, monitor_wrapper: bool = True, seed: int | None = None, use_sde: bool = False, sde_sample_freq: int = -1, supported_action_spaces: tuple[type[spaces.Space], ...] | None = None)

Bases : LightningModule

Base class for RL algorithms, adapted to use PyTorch Lightning.

Inherits from LightningModule so that device placement, mixed-precision, multi-GPU training, logging, and checkpointing are delegated to Lightning. The RL-specific training loop (environment interaction, rollout buffers, gradient updates) is implemented in OnPolicyAlgorithm and OffPolicyAlgorithm.

Original SB3 docstring follows:

Parameters

policy : str | type[BasePolicy] — The policy model to use (MlpPolicy, CnnPolicy, ...)
env — The environment to learn from (if registered in Gym, can be str. Can be None for loading trained models)
learning_rate : float | Schedule — learning rate for the optimizer, it can be a function of the current progress remaining (from 1 to 0)
policy_kwargs : dict[str, Any] | None — Additional arguments to be passed to the policy on creation
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
verbose : int — Verbosity level: 0 for no output, 1 for info messages (such as device or wrappers used), 2 for debug messages
device : torch.device | str — Compatibility-only constructor argument. Runtime device placement is resolved by RLTrainer during fit.
support_multi_env : bool — Whether the algorithm supports training with multiple environments (as in A2C)
monitor_wrapper : bool — When creating an environment, whether to wrap it or not in a Monitor wrapper.
seed : int | None — Seed for the pseudo random generators
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)
supported_action_spaces : tuple[type[spaces.Space], ...] | None — The action spaces supported by the algorithm.

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

create_trainer — Create an RLTrainer configured like the SB3-compatible learn shim.
on_fit_start — Called by Lightning before training starts. Initialise env observation.
set_logger — Setter for the custom Lightning logger used by RLTrainer.
log_metrics — Log a metrics dict via LightningModule self.log().
get_summary_writer — Return the active TensorBoard SummaryWriter if one is configured.
get_env — Returns the current environment (can be None if not defined).
get_vec_normalize_env — Return the VecNormalize wrapper of the training env.
set_env — Validate and attach the environment to the trainer.
learn — Return a trained model.
predict — Get the policy action from an observation (and optional hidden state).
set_random_seed — Set the seed of the pseudo-random generators.
set_parameters — Load parameters from a given zip-file or a nested dictionary containing parameters for.
load — Load the model from a zip-file.
get_parameters — Return the parameters of the agent. This includes parameters from different networks, e.g.
save — Save all the attributes of the object and the model parameters in a zip-file.
dump_logs — Write log data. (Implemented by OffPolicyAlgorithm and OnPolicyAlgorithm).

source property BaseAlgorithm.env: VecEnv | None

Return the active training env from the trainer, or None.

source property BaseAlgorithm.device: torch.device

Return current device (Lightning-aware).

source method BaseAlgorithm.create_trainer(max_steps: int, callback: MaybeCallback = None, run_name: str = 'run', log_every_n_steps: int | None = None, enable_progress_bar: bool | None = None, reset_num_timesteps: bool = True, default_root_dir: str | pathlib.Path | None = None) → L.Trainer

Create an RLTrainer configured like the SB3-compatible learn shim.

Parameters

max_steps : int — Total environment steps to collect.
callback : MaybeCallback — SB3-style callback(s) fired at each RL step.
run_name : str — The run name used in loggers.
log_every_n_steps : int | None — Logging frequency (if None, defaults to _default_log_every_n_steps).
enable_progress_bar : bool | None — Whether to display a progress bar using tqdm and rich.
default_root_dir : str | pathlib.Path | None — Root directory for logs and checkpoints.

Returns

L.Trainer — An RLTrainer configured with the same defaults used by learn().

source method BaseAlgorithm.on_fit_start() → None

Called by Lightning before training starts. Initialise env observation.

source method BaseAlgorithm.set_logger(logger: Any) → None

Setter for the custom Lightning logger used by RLTrainer.

.. warning::

        When passing a custom logger object, this overrides the trainer-managed
        default logger. Use ``RLTrainer(default_root_dir=...)`` to control the
        default TensorBoard save path.

source method BaseAlgorithm.log_metrics(metrics: dict[str, Any], step: int | None = None) → None

Log a metrics dict via LightningModule self.log().

source method BaseAlgorithm.get_summary_writer() → Any

Return the active TensorBoard SummaryWriter if one is configured.

source method BaseAlgorithm.get_env() → VecEnv | None

Returns the current environment (can be None if not defined).

Resolves through the trainer if available, otherwise returns None.

Returns

VecEnv | None — The current environment

source method BaseAlgorithm.get_vec_normalize_env() → VecNormalize | None

Return the VecNormalize wrapper of the training env.

if it exists.

Returns

VecNormalize | None — The VecNormalize env.

source method BaseAlgorithm.set_env(env: GymEnv, force_reset: bool = True) → None

Validate and attach the environment to the trainer.

The env is stored on the active RLTrainer, not on the model.

Parameters

env : GymEnv — The environment for learning a policy
force_reset : bool — Force call to reset() before training to avoid unexpected behavior. See issue https://github.com/DLR-RM/stable-baselines3/issues/597

Raises

RuntimeError

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

Return a trained model.

Delegates to :class:~lit_sb3.trainer.RLTrainer for backwards compatibility. Prefer constructing RLTrainer explicitly for full control over Lightning Trainer settings.

Parameters

total_timesteps : int — The total number of samples (env steps) to train on Note: it is a lower bound, see issue #1150 <https://github.com/DLR-RM/stable-baselines3/issues/1150>_
env : GymEnv | str | None — Training environment or env-id string. Required by the SB3-compatible learn() shim because environments are no longer stored on the model.
callback : MaybeCallback — callback(s) called at every step with state of the algorithm.
log_every_n_steps : int | None — How often to emit SB3 stats. For on-policy algos (PPO, A2C, …) this is the number of rollout cycles; for off-policy algos (SAC, TD3, …) it is the number of episodes. None uses the algorithm's default.
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)
progress_bar : bool — Display a progress bar using tqdm and rich.

Returns

SelfBaseAlgorithm — the trained model

source method BaseAlgorithm.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]

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.

Returns

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

source method BaseAlgorithm.set_random_seed(seed: int | None = None) → None

Set the seed of the pseudo-random generators.

(python, numpy, pytorch, gym, action_space).

Parameters

seed : int | None
seed : int | None — Random seed for reproducibility.

source method BaseAlgorithm.set_parameters(load_path_or_dict: str | TensorDict, exact_match: bool = True, device: torch.device | str = 'auto') → None

Load parameters from a given zip-file or a nested dictionary containing parameters for.

different modules (see get_parameters).

Parameters

load_path_or_iter — Location of the saved data (path or file-like, see save), or a nested dictionary containing nn.Module parameters used by the policy. The dictionary maps object names to a state-dictionary returned by torch.nn.Module.state_dict().
exact_match : bool — If True, the given parameters should include parameters for each module and each of their parameters, otherwise raises an Exception. If set to False, this can be used to update only specific parameters.
device : torch.device | str — Device on which the code should run.
load_path_or_dict : str | TensorDict — Path to model file or dict with parameters.

Raises

ValueError

source classmethod BaseAlgorithm.load(path: str | pathlib.Path | io.BufferedIOBase, env: GymEnv | None = None, device: torch.device | str = 'auto', custom_objects: dict[str, Any] | None = None, print_system_info: bool = False, force_reset: bool = True, **kwargs) → SelfBaseAlgorithm

Load the model from a zip-file.

Warning: load re-creates the model from scratch, it does not update it in-place! For an in-place load use set_parameters instead.

Parameters

path : str | pathlib.Path | io.BufferedIOBase — path to the file (or a file-like) where to load the agent from
env : GymEnv | None — the new environment to run the loaded model on (can be None if you only need prediction from a trained model) has priority over any saved environment
device : torch.device | str — Device on which the code should run.
custom_objects : dict[str, Any] | None — Dictionary of objects to replace upon loading. If a variable is present in this dictionary as a key, it will not be deserialized and the corresponding item will be used instead. Similar to custom_objects in keras.models.load_model. Useful when you have an object in file that can not be deserialized.
print_system_info : bool — Whether to print system info from the saved model and the current system info (useful to debug loading issues)
force_reset : bool — Force call to reset() before training to avoid unexpected behavior. See https://github.com/DLR-RM/stable-baselines3/issues/597
kwargs — extra arguments to change the model when loading

Returns

SelfBaseAlgorithm — new model instance with loaded parameters

Raises

ValueError
KeyError

source method BaseAlgorithm.get_parameters() → dict[str, dict]

Return the parameters of the agent. This includes parameters from different networks, e.g.

critics (value functions) and policies (pi functions).

Returns

dict[str, dict] — Mapping of from names of the objects to PyTorch state-dicts.

source method BaseAlgorithm.save(path: str | pathlib.Path | io.BufferedIOBase, exclude: Iterable[str] | None = None, include: Iterable[str] | None = None) → None

Save all the attributes of the object and the model parameters in a zip-file.

Parameters

path : str | pathlib.Path | io.BufferedIOBase — path to the file where the rl agent should be saved
exclude : Iterable[str] | None — name of parameters that should be excluded in addition to the default ones
include : Iterable[str] | None — name of parameters that might be excluded but should be included anyway

source method BaseAlgorithm.dump_logs() → None

Write log data. (Implemented by OffPolicyAlgorithm and OnPolicyAlgorithm).

Raises

NotImplementedError

source class OffPolicyAlgorithm(policy: str | type[BasePolicy], learning_rate: float | Schedule, 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, 'step'), 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, stats_window_size: int = 100, verbose: int = 0, device: torch.device | str = 'auto', support_multi_env: bool = False, monitor_wrapper: bool = True, seed: int | None = None, use_sde: bool = False, sde_sample_freq: int = -1, use_sde_at_warmup: bool = False, sde_support: bool = True, supported_action_spaces: tuple[type[spaces.Space], ...] | None = None)

Bases : BaseAlgorithm

Base class for off-policy algorithms (SAC, TD3, DQN, DDPG, …).

Off-policy training is driven step-by-step

each Lightning training_step call represents one iteration of "collect some env experience → do gradient updates from replay buffer".

Parameters

policy : str | type[BasePolicy] — Policy model
env — Training environment
learning_rate : float | Schedule — LR or LR schedule
buffer_size : int — Replay buffer capacity
learning_starts : int — Warm-up steps before gradient updates begin
batch_size : int — Mini-batch size for gradient updates
tau : float — Polyak soft-update coefficient
gamma : float — Discount factor
train_freq : int | tuple[int, str] — How often to collect env steps (per-step or per-episode)
gradient_steps : int — Gradient updates after each collection; -1 = auto
action_noise : ActionNoise | None — Exploration noise (for deterministic policies)
replay_buffer_class : type[ReplayBuffer] | None — Buffer class (auto-selected if None)
replay_buffer_kwargs : dict[str, Any] | None — Extra kwargs for the buffer
optimize_memory_usage : bool — Memory-efficient replay buffer variant
n_steps : int — N-step returns (>1 uses NStepReplayBuffer)
policy_kwargs : dict[str, Any] | None — Extra kwargs passed to the policy
stats_window_size : int — Episode-info averaging window
verbose : int — Verbosity level
device : torch.device | str — Torch device
support_multi_env : bool — Allow multiple parallel envs
monitor_wrapper : bool — Wrap env in Monitor
seed : int | None — RNG seed
use_sde : bool — Use gSDE exploration
sde_sample_freq : int — Noise re-sample frequency for gSDE
use_sde_at_warmup : bool — Use gSDE during warm-up phase
sde_support : bool — Whether the policy supports gSDE
supported_action_spaces : tuple[type[spaces.Space], ...] | None — Allowed action-space types

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 — Subclasses must override this to return their optimizers.
train_dataloader — Return an unbounded dataset that drives the training loop.
on_train_batch_end
training_step — One off-policy iteration: collect env transitions + gradient updates.
on_train_end
collect_rollouts
dump_logs
save_replay_buffer
load_replay_buffer
learn — Train for total_timesteps env steps.

source method OffPolicyAlgorithm.configure_optimizers() → OptimizerLRScheduler

Subclasses must override this to return their optimizers.

Raises

NotImplementedError

source method OffPolicyAlgorithm.train_dataloader() → DataLoader

Return an unbounded dataset that drives the training loop.

Each integer token yielded by InfiniteReplayBufferDataset triggers one call to training_step, which handles env interaction + gradient updates.

source method OffPolicyAlgorithm.on_train_batch_end(outputs: Any, batch: Any, batch_idx: int) → None

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

One off-policy iteration: collect env transitions + gradient updates.

Must be overridden by concrete algorithm classes.

Raises

NotImplementedError

source method OffPolicyAlgorithm.on_train_end() → None

source method OffPolicyAlgorithm.collect_rollouts(env: VecEnv, callback: BaseCallback, train_freq: TrainFreq, replay_buffer: ReplayBuffer, action_noise: ActionNoise | None = None, learning_starts: int = 0, log_interval: int | None = None) → RolloutReturn

source method OffPolicyAlgorithm.dump_logs() → None

source method OffPolicyAlgorithm.save_replay_buffer(path: str | pathlib.Path | io.BufferedIOBase) → None

source method OffPolicyAlgorithm.load_replay_buffer(path: str | pathlib.Path | io.BufferedIOBase, truncate_last_traj: bool = True) → None

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

Train for total_timesteps env steps.

Parameters

total_timesteps : int — Total environment steps to collect.
env : GymEnv | str | None — Training environment (gym env, env-id string, or VecEnv).
callback : MaybeCallback — SB3-style callback(s) fired at each RL step.
log_every_n_steps : int | None — Logging frequency.
tb_log_name : str — TensorBoard run name.
reset_num_timesteps : bool — Whether to reset the step counter.
progress_bar : bool — Display a progress bar.
default_root_dir : str | Path | None — Root directory for logs and checkpoints.

Returns

SelfOffPolicyAlgorithm — The trained model.

Raises

TypeError

source class OnPolicyAlgorithm(policy: str | type[ActorCriticPolicy], learning_rate: float | Schedule, n_steps: int, gamma: float, gae_lambda: float, ent_coef: float, vf_coef: float, max_grad_norm: float, use_sde: bool, sde_sample_freq: int, rollout_buffer_class: type[RolloutBuffer] | None = None, rollout_buffer_kwargs: dict[str, Any] | None = None, stats_window_size: int = 100, monitor_wrapper: bool = True, policy_kwargs: dict[str, Any] | None = None, verbose: int = 0, seed: int | None = None, device: torch.device | str = 'auto', _init_setup_model: bool = True, supported_action_spaces: tuple[type[spaces.Space], ...] | None = None)

Bases : BaseAlgorithm

Base class for on-policy algorithms (PPO, A2C, TRPO, …).

Subclasses must override training_step to compute and return the loss for a single mini-batch of rollout data.

Parameters

policy : str | type[ActorCriticPolicy] — Policy model (MlpPolicy, CnnPolicy, …)
env — Training environment
learning_rate : float | Schedule — LR or LR schedule (progress 1→0)
n_steps : int — Rollout steps per env per update
gamma : float — Discount factor
gae_lambda : float — GAE lambda
ent_coef : float — Entropy coefficient
vf_coef : float — Value-function coefficient
max_grad_norm : float — Gradient clipping threshold
use_sde : bool — Use gSDE exploration
sde_sample_freq : int — Noise re-sample frequency for gSDE
rollout_buffer_class : type[RolloutBuffer] | None — Buffer class (auto-selected if None)
rollout_buffer_kwargs : dict[str, Any] | None — Extra kwargs for the buffer constructor
stats_window_size : int — Episode-info averaging window
monitor_wrapper : bool — Wrap env in Monitor
policy_kwargs : dict[str, Any] | None — Extra kwargs passed to the policy
verbose : int — Verbosity level
seed : int | None — RNG seed
device : torch.device | str — Torch device
_init_setup_model : bool — Call _setup_model() in __init__
supported_action_spaces : tuple[type[spaces.Space], ...] | None — Allowed action-space types

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 for Lightning's automatic optimization.
train_dataloader — Wrap the rollout buffer in a DataLoader.
on_train_epoch_start — Collect a fresh rollout (unless skipped) and update progress/LR before each epoch.
training_step — Compute and return the loss for one mini-batch of rollout data.
on_train_end — Signal the RL callback that training has finished.
collect_rollouts — Collect experiences using the current policy and fill rollout_buffer.
learn — Train for total_timesteps environment steps.
dump_logs — Write log data (kept for SB3 API compatibility).

source method OnPolicyAlgorithm.configure_optimizers() → OptimizerLRScheduler

Return the policy optimizer for Lightning's automatic optimization.

source method OnPolicyAlgorithm.train_dataloader() → DataLoader

Wrap the rollout buffer in a DataLoader.

RolloutDataset holds a reference to the buffer, so the freshly collected rollout data (stored in on_train_epoch_start) is automatically served to each training_step call.

source method OnPolicyAlgorithm.on_train_epoch_start() → None

Collect a fresh rollout (unless skipped) and update progress/LR before each epoch.

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

Compute and return the loss for one mini-batch of rollout data.

Must be overridden by concrete algorithm classes (PPO, A2C, …).

Raises

NotImplementedError

source method OnPolicyAlgorithm.on_train_end() → None

Signal the RL callback that training has finished.

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

Collect experiences using the current policy and fill rollout_buffer.

Parameters

env : VecEnv — Training environment
callback : BaseCallback — Callback called at each step
rollout_buffer : RolloutBuffer — Buffer to fill
n_rollout_steps : int — Number of steps per env

Returns

bool — True if completed, False if callback terminated early

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

Train for total_timesteps environment steps.

Parameters

total_timesteps : int — Total environment steps to collect.
env : GymEnv | str | None — Training environment (gym env, env-id string, or VecEnv).
callback : MaybeCallback — SB3-style callback(s) fired at each RL step.
log_every_n_steps : int | None — Logging frequency.
tb_log_name : str — TensorBoard run name.
reset_num_timesteps : bool — Whether to reset the step counter.
progress_bar : bool — Display a progress bar.
default_root_dir : str | Path | None — Root directory for logs and checkpoints.

Returns

SelfOnPolicyAlgorithm — The trained model.

Raises

TypeError

source method OnPolicyAlgorithm.dump_logs(iteration: int = 0) → None

Write log data (kept for SB3 API compatibility).