Skip to content

lit_sb3.algorithms.off_policy

source module lit_sb3.algorithms.off_policy

Off-policy algorithm base class — Lightning-integrated.

Architecture

  • automatic_optimization = False — all gradient calls are explicit.
  • train_dataloader returns an InfiniteReplayBufferDataset that keeps Lightning advancing batches until the trainer-owned RL budget stops the run.
  • training_step (overridden by subclasses) collects env transitions and performs manual gradient updates.
  • learn() wraps Lightning Trainer.fit while keeping the SB3 API.

Classes

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

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