Skip to content

lit_sb3.datamodule

source module lit_sb3.datamodule

RLEnvDataModule — a LightningDataModule that provides environments instead of dataloaders.

In supervised learning, a LightningDataModule manages datasets and dataloaders for each stage (fit, validate, test, predict). In reinforcement learning the "data" is the environment: the agent generates its own training data by interacting with the environment.

RLEnvDataModule follows the same stage-based pattern but exposes environments rather than dataloaders::

dm = RLEnvDataModule(
    train_env="CartPole-v1",
    val_env="CartPole-v1",
)
trainer = RLTrainer(max_timesteps=100_000)
trainer.fit(model, datamodule=dm)
trainer.validate(model, datamodule=dm)

Classes

  • RLEnvDataModule — A :class:~lightning.LightningDataModule that holds environments for each RL stage.

source class RLEnvDataModule(train_env: GymEnv | str | None = None, val_env: GymEnv | str | None = None, test_env: GymEnv | str | None = None, predict_env: GymEnv | str | None = None, monitor_wrapper: bool = True, verbose: int = 0)

Bases : L.LightningDataModule

A :class:~lightning.LightningDataModule that holds environments for each RL stage.

Instead of :meth:train_dataloader, :meth:val_dataloader, etc., this module exposes :attr:train_env, :attr:val_env, :attr:test_env, and :attr:predict_env. The :class:~lit_sb3.trainer.RLTrainer extracts the appropriate environment from the datamodule at each stage.

Parameters

  • train_env : GymEnv | str | None — Environment (or env-id string) used for training.

  • val_env : GymEnv | str | None — Environment (or env-id string) used for validation / evaluation.

  • test_env : GymEnv | str | None — Environment (or env-id string) used for testing / evaluation.

  • predict_env : GymEnv | str | None — Environment (or env-id string) used for prediction.

  • monitor_wrapper : bool — Whether to wrap raw Gym envs in a Monitor wrapper.

  • verbose : int — Verbosity level forwarded to env wrapping utilities.

Attributes

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

Methods

  • setup — Resolve and wrap environments for the requested stage.

  • teardown — Release resolved environments for the requested stage.

  • on_exception — Ensure vectorized envs are closed when Lightning aborts the run.

  • get_env — Return the environment configured for stage.

  • train_dataloader — RLEnvDataModule exposes envs, not training dataloaders.

  • val_dataloader — RLEnvDataModule exposes envs, not validation dataloaders.

  • test_dataloader — RLEnvDataModule exposes envs, not test dataloaders.

  • predict_dataloader — RLEnvDataModule exposes envs, not predict dataloaders.

  • state_dict — Persist serializable datamodule config and env specs, never live env instances.

  • load_state_dict — Restore serializable env specs and datamodule config from checkpoint state.

  • load_from_checkpoint — Load the datamodule and then restore env specs through :meth:load_state_dict.

  • close — Close all resolved environments and free resources.

source method RLEnvDataModule.setup(stage: str | None = None) → None

Resolve and wrap environments for the requested stage.

Called automatically by Lightning before each stage, or can be called manually.

source method RLEnvDataModule.teardown(stage: str | None = None) → None

Release resolved environments for the requested stage.

source method RLEnvDataModule.on_exception(exception: BaseException) → None

Ensure vectorized envs are closed when Lightning aborts the run.

source method RLEnvDataModule.get_env(stage: Literal['fit', 'validate', 'test', 'predict'], *, resolve: bool = True) → VecEnv | None

Return the environment configured for stage.

Parameters

  • stage : Literal['fit', 'validate', 'test', 'predict'] — Lightning stage name.

  • resolve : bool — Whether to call :meth:setup before reading the stage env.

source method RLEnvDataModule.train_dataloader() → None

RLEnvDataModule exposes envs, not training dataloaders.

Raises

  • TypeError

source method RLEnvDataModule.val_dataloader() → None

RLEnvDataModule exposes envs, not validation dataloaders.

Raises

  • TypeError

source method RLEnvDataModule.test_dataloader() → None

RLEnvDataModule exposes envs, not test dataloaders.

Raises

  • TypeError

source method RLEnvDataModule.predict_dataloader() → None

RLEnvDataModule exposes envs, not predict dataloaders.

Raises

  • TypeError

source method RLEnvDataModule.state_dict() → dict[str, Any]

Persist serializable datamodule config and env specs, never live env instances.

source method RLEnvDataModule.load_state_dict(state_dict: dict[str, Any]) → None

Restore serializable env specs and datamodule config from checkpoint state.

source classmethod RLEnvDataModule.load_from_checkpoint(checkpoint_path: Any, map_location: Any = None, hparams_file: Any = None, weights_only: bool | None = None, **kwargs: Any) → RLEnvDataModule

Load the datamodule and then restore env specs through :meth:load_state_dict.

source method RLEnvDataModule.close() → None

Close all resolved environments and free resources.