Skip to content

lit_sb3.algorithms.on_policy

source module lit_sb3.algorithms.on_policy

On-policy algorithm base class — Lightning-integrated.

One Lightning *epoch* corresponds to one rollout collection cycle

  1. on_train_epoch_start → collect_rollouts() fills the RolloutBuffer.
  2. train_dataloader → wraps the buffer in a DataLoader (via RolloutDataset).
  3. training_step → per mini-batch gradient update (implemented by subclasses).
  4. configure_optimizers → returns the policy optimizer.

learn() translates total_timesteps into Lightning max_epochs and calls Trainer.fit(self), so the public SB3 API is fully preserved.

Classes

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).