Skip to content

lit_sb3.algorithms.ppo_mask

source package lit_sb3.algorithms.ppo_mask

Classes

  • MaskablePPO — Proximal Policy Optimization algorithm (PPO) (clip version) with Invalid Action Masking.

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