lit_sb3.callbacks

source package lit_sb3.callbacks

Classes

• BaseCallback — Base class for callback.

• CallbackList — Class for chaining callbacks.

• CheckpointCallback — Callback for saving a model every save_freq calls.

• ConvertCallback — Convert functional callback (old-style) to object.

• EvalCallback — Callback for evaluating an agent.

• EventCallback — Base class for triggering callback on event.

• EveryNTimesteps — Trigger a callback every n_steps timesteps.

• LogEveryNTimesteps — Log data every n_steps timesteps.

• ProgressBarCallback — Display a progress bar when training SB3 agent.

• RawStatisticsCallback — Callback for logging raw per-episode returns and lengths to TensorBoard.

• SaveVecNormalizeCallback — Callback for saving a :class:~lit_sb3.vec_env.VecNormalize wrapper periodically.

• StopTrainingOnMaxEpisodes — Stop the training once a maximum number of episodes are played.

• StopTrainingOnNoModelImprovement — Stop the training early if there is no new best model (new best mean reward) after more than N consecutive.

• StopTrainingOnRewardThreshold — Stop the training once a threshold in episodic reward.

• SummaryWriterCallback — Resolve and expose the active TensorBoard SummaryWriter.

• TrialEvalCallback — Callback used for evaluating and reporting an Optuna trial.

source class BaseCallback(verbose: int = 0)

Bases : L.Callback, ABC

Base class for callback.

Inherits from lightning.Callback so instances can be passed directly to a lightning.Trainer via its callbacks argument. When training runs through :meth:BaseAlgorithm.learn the SB3 lifecycle methods (on_training_start / on_training_end) are called manually before Trainer.fit; the Lightning shims below detect that and skip to avoid double-calls. When a callback is passed directly to a Trainer without going through learn, the shims invoke the SB3 lifecycle instead.

Parameters

• verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

Attributes

• state_key : str — Identifier for the state of the callback.

Methods

• init_callback — Initialize the callback by saving references to the.

• on_training_start

• on_rollout_start

• on_step — This method will be called by the model after each call to env.step().

• on_training_end

• on_rollout_end

• update_locals — Update the references to the local variables.

• update_child_locals — Update the references to the local variables on sub callbacks.

• on_fit_start — Lightning hook: called at the very beginning of Trainer.fit.

• on_fit_end — Lightning hook: called at the very end of Trainer.fit.

source property BaseCallback.training_env: VecEnv

source property BaseCallback.summary_writer: Any

source method BaseCallback.init_callback(model: BaseAlgorithm) → None

Initialize the callback by saving references to the.

RL model and the training environment for convenience.

source method BaseCallback.on_training_start(locals_: dict[str, Any], globals_: dict[str, Any]) → None

source method BaseCallback.on_rollout_start() → None

source method BaseCallback.on_step() → bool

This method will be called by the model after each call to env.step().

For child callback (of an EventCallback), this will be called when the event is triggered.

Returns

• bool — If the callback returns False, training is aborted early.

source method BaseCallback.on_training_end() → None

source method BaseCallback.on_rollout_end() → None

source method BaseCallback.update_locals(locals_: dict[str, Any]) → None

Update the references to the local variables.

Parameters

• locals_ : dict[str, Any] — the local variables during rollout collection

source method BaseCallback.update_child_locals(locals_: dict[str, Any]) → None

Update the references to the local variables on sub callbacks.

Parameters

• locals_ : dict[str, Any] — the local variables during rollout collection

source method BaseCallback.on_fit_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

Lightning hook: called at the very beginning of Trainer.fit.

When invoked through :meth:BaseAlgorithm.learn, on_training_start has already been called manually, so this is a no-op. When a callback is passed directly to Trainer(callbacks=[cb]) this initialises the SB3 side of the callback instead.

source method BaseCallback.on_fit_end(trainer: L.Trainer, pl_module: L.LightningModule) → None

Lightning hook: called at the very end of Trainer.fit.

When invoked through :meth:BaseAlgorithm.learn, on_training_end has already been called by LightningModule.on_train_end, so this is a no-op. When a callback is passed directly to Trainer this tears down the SB3 side instead.

source class CallbackList(callbacks: list[BaseCallback])

Bases : BaseCallback

Class for chaining callbacks.

Parameters

• callbacks : list[BaseCallback] — A list of callbacks that will be called sequentially.

Attributes

• state_key : str — Identifier for the state of the callback.

Methods

• update_child_locals — Update the references to the local variables.

source method CallbackList.update_child_locals(locals_: dict[str, Any]) → None

Update the references to the local variables.

Parameters

• locals_ : dict[str, Any] — the local variables during rollout collection

source class CheckpointCallback(save_freq: int, save_path: str | Path, name_prefix: str = 'rl_model', save_replay_buffer: bool = False, save_vecnormalize: bool = False, verbose: int = 0)

Bases : BaseCallback

Callback for saving a model every save_freq calls.

to env.step(). By default, it only saves model checkpoints, you need to pass save_replay_buffer=True, and save_vecnormalize=True to also save replay buffer checkpoints and normalization statistics checkpoints.

.. warning::

When using multiple environments, each call to env.step() will effectively correspond to n_envs steps. To account for that, you can use save_freq = max(save_freq // n_envs, 1)

Parameters

• save_freq : int — Save checkpoints every save_freq call of the callback.

• save_path : str | Path — Path to the folder where the model will be saved.

• name_prefix : str — Common prefix to the saved models

• save_replay_buffer : bool — Save the model replay buffer

• save_vecnormalize : bool — Save the VecNormalize statistics

• verbose : int — Verbosity level: 0 for no output, 2 for indicating when saving model checkpoint

Attributes

• state_key : str — Identifier for the state of the callback.

source class ConvertCallback(callback: Callable[[dict[str, Any], dict[str, Any]], bool] | None, verbose: int = 0)

Bases : BaseCallback

Convert functional callback (old-style) to object.

Parameters

• callback : Callable[[dict[str, Any], dict[str, Any]], bool] | None

• verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

Attributes

• state_key : str — Identifier for the state of the callback.

source class EvalCallback(eval_env: gym.Env | VecEnv, callback_on_new_best: BaseCallback | None = None, callback_after_eval: BaseCallback | None = None, n_eval_episodes: int = 5, eval_freq: int = 10000, log_path: str | Path | None = None, best_model_save_path: str | Path | None = None, deterministic: bool = True, render: bool = False, verbose: int = 1, warn: bool = True)

Bases : EventCallback

Callback for evaluating an agent.

.. warning::

When using multiple environments, each call to env.step() will effectively correspond to n_envs steps. To account for that, you can use eval_freq = max(eval_freq // n_envs, 1)

Parameters

• eval_env : gym.Env | VecEnv — The environment used for initialization

• callback_on_new_best : BaseCallback | None — Callback to trigger when there is a new best model according to the mean_reward

• callback_after_eval : BaseCallback | None — Callback to trigger after every evaluation

• n_eval_episodes : int — The number of episodes to test the agent

• eval_freq : int — Evaluate the agent every eval_freq call of the callback.

• log_path : str | Path | None — Path to a folder where the evaluations (evaluations.npz) will be saved. It will be updated at each evaluation.

• best_model_save_path : str | Path | None — Path to a folder where the best model according to performance on the eval env will be saved.

• deterministic : bool — Whether the evaluation should use a stochastic or deterministic actions.

• render : bool — Whether to render or not the environment during evaluation

• verbose : int — Verbosity level: 0 for no output, 1 for indicating information about evaluation results

• warn : bool — Passed to evaluate_policy (warns if eval_env has not been wrapped with a Monitor wrapper)

Attributes

• state_key : str — Identifier for the state of the callback.

Methods

• on_train_start

• update_child_locals — Update the references to the local variables.

source method EvalCallback.on_train_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

source method EvalCallback.update_child_locals(locals_: dict[str, Any]) → None

Update the references to the local variables.

Parameters

• locals_ : dict[str, Any] — the local variables during rollout collection

source class EventCallback(callback: BaseCallback | None = None, verbose: int = 0)

Bases : BaseCallback

Base class for triggering callback on event.

Parameters

• callback : BaseCallback | None — Callback that will be called when an event is triggered.

• verbose : int — Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages

Attributes

• state_key : str — Identifier for the state of the callback.

Methods

• init_callback

• update_child_locals — Update the references to the local variables.

source method EventCallback.init_callback(model: BaseAlgorithm) → None

source method EventCallback.update_child_locals(locals_: dict[str, Any]) → None

Update the references to the local variables.

Parameters

• locals_ : dict[str, Any] — the local variables during rollout collection

source class EveryNTimesteps(n_steps: int, callback: BaseCallback)

Bases : EventCallback

Trigger a callback every n_steps timesteps.

Parameters

• n_steps : int — Number of timesteps between two trigger.

• callback : BaseCallback — Callback that will be called when the event is triggered.

Attributes

• state_key : str — Identifier for the state of the callback.

source class LogEveryNTimesteps(n_steps: int)

Bases : EveryNTimesteps

Log data every n_steps timesteps.

Parameters

• n_steps : int — Number of timesteps between two trigger.

Attributes

• state_key : str — Identifier for the state of the callback.

source class ProgressBarCallback()

Bases : BaseCallback

Display a progress bar when training SB3 agent.

using tqdm and rich packages.

Attributes

• state_key : str — Identifier for the state of the callback.

source class RawStatisticsCallback(verbose: int = 0)

Bases : BaseCallback

Callback for logging raw per-episode returns and lengths to TensorBoard.

Requires TensorBoard logging to be active. The callback bypasses Lightning's default log-throttling so that every episode is recorded.

Attributes

• state_key : str — Identifier for the state of the callback.

source class SaveVecNormalizeCallback(save_freq: int, save_path: str | Path, name_prefix: str | None = None, verbose: int = 0)

Bases : BaseCallback

Callback for saving a :class:~lit_sb3.vec_env.VecNormalize wrapper periodically.

Parameters

• save_freq : int — Save every save_freq steps.

• save_path : str | Path — Folder where the VecNormalize pickle is saved.

• name_prefix : str | None — If set, save timestamped files instead of overwriting a single vecnormalize.pkl.

• verbose : int — Verbosity level.

Attributes

• state_key : str — Identifier for the state of the callback.

source class StopTrainingOnMaxEpisodes(max_episodes: int, verbose: int = 0)

Bases : BaseCallback

Stop the training once a maximum number of episodes are played.

For multiple environments presumes that, the desired behavior is that the agent trains on each env for max_episodes and in total for max_episodes * n_envs episodes.

Parameters

• max_episodes : int — Maximum number of episodes to stop training.

• verbose : int — Verbosity level: 0 for no output, 1 for indicating information about when training ended by reaching max_episodes

Attributes

• state_key : str — Identifier for the state of the callback.

source class StopTrainingOnNoModelImprovement(max_no_improvement_evals: int, min_evals: int = 0, verbose: int = 0)

Bases : BaseCallback

Stop the training early if there is no new best model (new best mean reward) after more than N consecutive.

evaluations.

It is possible to define a minimum number of evaluations before start to count evaluations without improvement.

It must be used with the EvalCallback.

Parameters

• max_no_improvement_evals : int — Maximum number of consecutive evaluations without a new best model.

• min_evals : int — Number of evaluations before start to count evaluations without improvements.

• verbose : int — Verbosity level: 0 for no output, 1 for indicating when training ended because no new best model

Attributes

• state_key : str — Identifier for the state of the callback.

source class StopTrainingOnRewardThreshold(reward_threshold: float, verbose: int = 0)

Bases : BaseCallback

Stop the training once a threshold in episodic reward.

has been reached (i.e. when the model is good enough).

It must be used with the EvalCallback.

Parameters

• reward_threshold : float — Minimum expected reward per episode to stop training.

• verbose : int — Verbosity level: 0 for no output, 1 for indicating when training ended because episodic reward threshold reached

Attributes

• state_key : str — Identifier for the state of the callback.

source class SummaryWriterCallback()

Bases : L.Callback

Resolve and expose the active TensorBoard SummaryWriter.

Attributes

• state_key : str — Identifier for the state of the callback.

Methods

• on_train_start

• on_validate_start

• on_test_start

• on_predict_start

• get_writer

source method SummaryWriterCallback.on_train_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

source method SummaryWriterCallback.on_validate_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

source method SummaryWriterCallback.on_test_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

source method SummaryWriterCallback.on_predict_start(trainer: L.Trainer, pl_module: L.LightningModule) → None

source classmethod SummaryWriterCallback.get_writer(trainer: L.Trainer) → Any

source class TrialEvalCallback(eval_env: VecEnv, trial: Any, n_eval_episodes: int = 5, eval_freq: int = 10000, deterministic: bool = True, verbose: int = 0, best_model_save_path: str | Path | None = None, log_path: str | Path | None = None)

Bases : EvalCallback

Callback used for evaluating and reporting an Optuna trial.

Evaluates the model every eval_freq steps and reports the mean reward to the Optuna trial. If the trial should be pruned, training is stopped early.

Parameters

• eval_env : VecEnv — Environment used for evaluation.

• trial : Any — Optuna trial instance.

• n_eval_episodes : int — Number of episodes per evaluation.

• eval_freq : int — Evaluate every eval_freq steps.

• deterministic : bool — Use deterministic actions during evaluation.

• verbose : int — Verbosity level.

• best_model_save_path : str | Path | None — Where to save the best model.

• log_path : str | Path | None — Where to save evaluation logs.

Attributes

• state_key : str — Identifier for the state of the callback.