lit_sb3.trainer

RLTrainer — a lightning.Trainer subclass for reinforcement learning.

Usage::

from lit_sb3 import PPO, RLEnvDataModule, RLTrainer

model = PPO("MlpPolicy")
dm = RLEnvDataModule(train_env="CartPole-v1")
trainer = RLTrainer(max_steps=100_000, enable_progress_bar=True)
trainer.fit(model, datamodule=dm)

The SB3-compatible model.learn(total_timesteps, env=...) API still works and simply delegates to RLTrainer.fit internally.

Classes

RLTrainer — A :class:lightning.Trainer subclass that drives reinforcement-learning runs.

Bases : L.Trainer

A :class:lightning.Trainer subclass that drives reinforcement-learning runs.

All standard lightning.Trainer keyword arguments are accepted and forwarded to the parent class, giving full access to Lightning features (mixed precision, custom loggers, DDP, plugins, …).

The RL-specific arguments control the training schedule and SB3 callbacks. accelerator, enable_model_summary, enable_checkpointing, and logger have RL-appropriate defaults but can be overridden.

Parameters

max_steps : int — Total environment steps to collect (RL timesteps budget). Lightning step-based stopping is disabled; this budget is enforced using RL env-step accounting instead.
rl_callback : MaybeCallback — SB3-style callback(s) fired at each RL step.
run_name : str — The run name used in loggers.

Attributes

device_ids : list[int] — List of device indexes per node.
num_devices : int — Number of devices the trainer uses per node.
model : Optional[torch.nn.Module] — The LightningModule, but possibly wrapped into DataParallel or DistributedDataParallel.
log_dir : Optional[str] — The directory for the current experiment. Use this to save images to, etc...
is_global_zero : bool — Whether this process is the global zero in multi-node training.
enable_validation : bool — Check if we should run validation during training.
default_root_dir : str — The default location to save artifacts of loggers, checkpoints etc.
early_stopping_callback : Optional[EarlyStopping] — class:~lightning.pytorch.callbacks.early_stopping.EarlyStopping callback in the Trainer.callbacks list, or None if it doesn't exist.
early_stopping_callbacks : list[EarlyStopping] — class:~lightning.pytorch.callbacks.early_stopping.EarlyStopping found in the Trainer.callbacks list.
checkpoint_callback : Optional[Checkpoint] — class:~lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint callback in the Trainer.callbacks list, or None if it doesn't exist.
checkpoint_callbacks : list[Checkpoint] — class:~lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint found in the Trainer.callbacks list.
progress_bar_callback : Optional[ProgressBar] — class:~lightning.pytorch.callbacks.progress.progress_bar.ProgressBar found in the Trainer.callbacks list, or None if one doesn't exist.
ckpt_path : Optional[_PATH] — meth:~lightning.pytorch.trainer.trainer.Trainer.fit, :meth:~lightning.pytorch.trainer.trainer.Trainer.validate, :meth:~lightning.pytorch.trainer.trainer.Trainer.test, or :meth:~lightning.pytorch.trainer.trainer.Trainer.predict.
sanity_checking : bool — Whether sanity checking is running.
received_sigterm : bool — Whether a signal.SIGTERM signal was received.
global_step : int — The number of optimizer steps taken (does not reset each epoch).
current_epoch : int — The current epoch, updated after the epoch end hooks are run.
is_last_batch : bool — Whether trainer is executing the last batch.
train_dataloader : Optional[TRAIN_DATALOADERS] — The training dataloader(s) used during trainer.fit().
val_dataloaders : Optional[EVAL_DATALOADERS] — The validation dataloader(s) used during trainer.fit() or trainer.validate().
test_dataloaders : Optional[EVAL_DATALOADERS] — The test dataloader(s) used during trainer.test().
predict_dataloaders : Optional[EVAL_DATALOADERS] — The prediction dataloader(s) used during trainer.predict().
num_training_batches : Union[int, float] — The number of training batches that will be used during trainer.fit().
num_sanity_val_batches : list[Union[int, float]] — The number of validation batches that will be used during the sanity-checking part of trainer.fit().
num_val_batches : list[Union[int, float]] — The number of validation batches that will be used during trainer.fit() or trainer.validate().
num_test_batches : list[Union[int, float]] — The number of test batches that will be used during trainer.test().
num_predict_batches : list[Union[int, float]] — The number of prediction batches that will be used during trainer.predict().
logger : Optional[Logger] — class:~lightning.pytorch.loggers.logger.Logger being used.
loggers : list[Logger] — class:~lightning.pytorch.loggers.logger.Logger used.
callback_metrics : _OUT_DICT — The metrics available to callbacks.
logged_metrics : _OUT_DICT — The metrics sent to the loggers.
progress_bar_metrics : _PBAR_DICT — The metrics sent to the progress bar.
estimated_stepping_batches : Union[int, float] — The estimated number of batches that will optimizer.step() during training.
rl_total_timesteps : int — Return the trainer-owned RL timestep budget.
rl_num_timesteps : int — Return the latest observed RL env-step count.
rl_budget_reached : bool — Return whether the trainer-owned RL timestep budget has been reached.
rl_progress : float — Return RL budget progress as a value in [0, 1].

Methods

fit — Train model for :attr:_rl_total_timesteps environment steps.
validate
test
predict

source property RLTrainer.rl_total_timesteps: int

Return the trainer-owned RL timestep budget.

source property RLTrainer.rl_num_timesteps: int

Return the latest observed RL env-step count.

source property RLTrainer.rl_budget_reached: bool

Return whether the trainer-owned RL timestep budget has been reached.

source property RLTrainer.rl_progress: float

Return RL budget progress as a value in [0, 1].

source method RLTrainer.fit(model: BaseAlgorithm, *, datamodule: Any = None, train_dataloaders: Any = None, val_dataloaders: Any = None, ckpt_path: str | Path | None = None, weights_only: bool | None = None) → BaseAlgorithm

Train model for :attr:_rl_total_timesteps environment steps.

The environment is provided via an :class:~lit_sb3.datamodule.RLEnvDataModule.

Parameters

model : BaseAlgorithm — The RL algorithm (a LightningModule subclass).
datamodule : Any — An :class:~lit_sb3.datamodule.RLEnvDataModule that provides the training (and optionally evaluation) environments.
train_dataloaders : Any — Unsupported for RLTrainer.
val_dataloaders : Any — Unsupported for RLTrainer.
ckpt_path : str | Path | None — Optional Lightning checkpoint path.
weights_only : bool | None — Optional Lightning checkpoint loading flag.

Returns

BaseAlgorithm — The trained model (same object, for chaining).

Raises

TypeError
ValueError

Raises

TypeError