Skip to content

lit_sb3.algorithms.base

source module lit_sb3.algorithms.base

Abstract base class for RL algorithms using PyTorch Lightning.

Classes

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

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

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