Skip to content

lit_sb3.algorithms.trpo

source package lit_sb3.algorithms.trpo

Classes

  • TRPO — Trust Region Policy Optimization (TRPO).

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

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