lit_sb3.algorithms.base

[docs] module lit_sb3.algorithms.base

"""Abstract base class for RL algorithms using PyTorch Lightning."""

from __future__ import annotations

from abc import abstractmethod
from collections import deque
from collections.abc import Iterable
import io
import pathlib
import time
from typing import Any, ClassVar, Optional, TypeVar, Union
import warnings

import gymnasium as gym
from gymnasium import spaces
import lightning as L  # noqa: N812
from lightning import LightningModule
import numpy as np
import torch

from lit_sb3.callbacks.base import BaseCallback, CallbackList
from lit_sb3.callbacks.convert import ConvertCallback
from lit_sb3.callbacks.progress import ProgressBarCallback
from lit_sb3.core.noise import ActionNoise
from lit_sb3.core.preprocessing import check_for_nested_spaces, is_image_space, is_image_space_channels_first
from lit_sb3.core.types import GymEnv, MaybeCallback, Schedule, TensorDict
from lit_sb3.policies import BasePolicy
from lit_sb3.utils.device import get_device
from lit_sb3.utils.env import is_wrapped, maybe_make_env
from lit_sb3.utils.save import load_from_zip_file, recursive_getattr, recursive_setattr, save_to_zip_file
from lit_sb3.utils.schedule import FloatSchedule
from lit_sb3.utils.spaces import check_for_correct_spaces
from lit_sb3.utils.system import get_system_info
from lit_sb3.utils.training import set_random_seed, update_learning_rate
from lit_sb3.vec_env import (
    DummyVecEnv,
    VecEnv,
    VecNormalize,
    VecTransposeImage,
    is_vecenv_wrapped,
    unwrap_vec_normalize,
)
from lit_sb3.vec_env.patch_gym import _convert_space, _patch_env
from lit_sb3.wrappers.monitor import Monitor

SelfBaseAlgorithm = TypeVar("SelfBaseAlgorithm", bound="BaseAlgorithm")


class BaseAlgorithm(LightningModule):
[docs]
    """Base class for RL algorithms, adapted to use PyTorch Lightning.

    Inherits from ``LightningModule`` so that device placement, mixed-precision,
    multi-GPU training, logging, and checkpointing are delegated to Lightning.
    The RL-specific training loop (environment interaction, rollout buffers,
    gradient updates) is implemented in ``OnPolicyAlgorithm`` and
    ``OffPolicyAlgorithm``.

    Original SB3 docstring follows:

    Args:
        policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
        env: The environment to learn from
            (if registered in Gym, can be str. Can be None for loading trained models)
        learning_rate: learning rate for the optimizer,
            it can be a function of the current progress remaining (from 1 to 0)
        policy_kwargs: Additional arguments to be passed to the policy on creation
        stats_window_size: Window size for the rollout logging, specifying the number of episodes to average
            the reported success rate, mean episode length, and mean reward over
        verbose: Verbosity level: 0 for no output, 1 for info messages (such as device or wrappers used), 2 for
            debug messages
        device: Compatibility-only constructor argument.
            Runtime device placement is resolved by RLTrainer during fit.
        support_multi_env: Whether the algorithm supports training
            with multiple environments (as in A2C)
        monitor_wrapper: When creating an environment, whether to wrap it
            or not in a Monitor wrapper.
        seed: Seed for the pseudo random generators
        use_sde: Whether to use generalized State Dependent Exploration (gSDE)
            instead of action noise exploration (default: False)
        sde_sample_freq: Sample a new noise matrix every n steps when using gSDE
            Default: -1 (only sample at the beginning of the rollout)
        supported_action_spaces: The action spaces supported by the algorithm.
    """

    # Policy aliases (see _get_policy_from_name())
    policy_aliases: ClassVar[dict[str, type[BasePolicy]]] = {}
    policy: BasePolicy
    observation_space: spaces.Space
    action_space: spaces.Space
    n_envs: int
    lr_schedule: Schedule

    def __init__(
        self,
        policy: str | type[BasePolicy],
        learning_rate: float | Schedule,
        policy_kwargs: dict[str, Any] | None = None,
        stats_window_size: int = 100,
        verbose: int = 0,
        device: torch.device | str = "auto",
        support_multi_env: bool = False,
        monitor_wrapper: bool = True,
        seed: int | None = None,
        use_sde: bool = False,
        sde_sample_freq: int = -1,
        supported_action_spaces: tuple[type[spaces.Space], ...] | None = None,
    ) -> None:
        super().__init__()

        if isinstance(policy, str):
            self.policy_class = self._get_policy_from_name(policy)
        else:
            self.policy_class = policy
        self.policy = None  # type: ignore[assignment]
        self.verbose = verbose
        self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs
        del device

        self.num_timesteps = 0
        # Used for updating schedules
        self._total_timesteps = 0
        # Used for computing fps, it is updated at each call of learn()
        self._num_timesteps_at_start = 0
        self.seed = seed
        self.action_noise: ActionNoise | None = None
        self.start_time = 0.0
        self.learning_rate = learning_rate
        self._last_obs = None  # type: Optional[Union[np.ndarray, dict[str, np.ndarray]]]
        self._last_episode_starts = None  # type: Optional[np.ndarray]
        # When using VecNormalize:
        self._last_original_obs = None  # type: Optional[Union[np.ndarray, dict[str, np.ndarray]]]
        self._episode_num = 0
        # Used for gSDE only
        self.use_sde = use_sde
        self.sde_sample_freq = sde_sample_freq
        # Track the training progress remaining (from 1 to 0)
        # this is used to update the learning rate
        self._current_progress_remaining = 1.0
        # Buffers for logging
        self._stats_window_size = stats_window_size
        self.ep_info_buffer = None  # type: Optional[deque]
        self.ep_success_buffer = None  # type: Optional[deque]
        # For logging (and TD3 delayed updates)
        self._n_updates = 0  # type: int
        # Whether the user passed a custom logger or not
        self._custom_logger = False
        self._trainer_logger: Any = None
        self._vec_normalize_env: VecNormalize | None = None
        self._monitor_wrapper = monitor_wrapper
        self._support_multi_env = support_multi_env
        self._supported_action_spaces = supported_action_spaces
        self._policy_name = policy if isinstance(policy, str) else None

    def _init_env_metadata(
        self,
        env: VecEnv,
        supported_action_spaces: tuple[type[spaces.Space], ...] | None = None,
        support_multi_env: bool = False,
    ) -> None:
        """Extract and validate env metadata without storing the env."""
        self.observation_space = env.observation_space
        self.action_space = env.action_space
        self.n_envs = env.num_envs
        self._vec_normalize_env = unwrap_vec_normalize(env)

        if supported_action_spaces is not None:
            assert isinstance(self.action_space, supported_action_spaces), (
                f"The algorithm only supports {supported_action_spaces} as action spaces "
                f"but {self.action_space} was provided"
            )

        if not support_multi_env and self.n_envs > 1:
            raise ValueError(
                "Error: the model does not support multiple envs; it requires a single vectorized environment."
            )

        # Catch common mistake: using MlpPolicy/CnnPolicy instead of MultiInputPolicy
        if self._policy_name in ["MlpPolicy", "CnnPolicy"] and isinstance(self.observation_space, spaces.Dict):
            raise ValueError(
                f"You must use `MultiInputPolicy` when working with dict observation space, not {self._policy_name}"
            )

        if self.use_sde and not isinstance(self.action_space, spaces.Box):
            raise ValueError("generalized State-Dependent Exploration (gSDE) can only be used with continuous actions.")

        if isinstance(self.action_space, spaces.Box):
            assert np.all(np.isfinite(np.array([self.action_space.low, self.action_space.high]))), (
                "Continuous action space must have a finite lower and upper bound"
            )

    @property
    def env(self) -> VecEnv | None:
[docs]
        """Return the active training env from the trainer, or None."""
        # Try the RLTrainer reference (set before _setup_learn)
        rl_trainer = getattr(self, "_rl_trainer_ref", None)
        if rl_trainer is not None and hasattr(rl_trainer, "_rl_env"):
            return rl_trainer._rl_env
        # Try Lightning's trainer (set after super().fit())
        try:
            trainer = self.trainer
            if trainer is not None and hasattr(trainer, "_rl_env"):
                return trainer._rl_env  # type: ignore[union-attr]
        except RuntimeError:
            pass
        return None

    @property
    def device(self) -> torch.device:
[docs]
        """Return current device (Lightning-aware)."""
        rl_trainer = getattr(self, "_rl_trainer_ref", None)
        if rl_trainer is not None:
            strategy = getattr(rl_trainer, "strategy", None)
            root_device = getattr(strategy, "root_device", None)
            if isinstance(root_device, torch.device):
                return root_device

        setup_device_override = getattr(self, "_setup_device_override", None)
        if isinstance(setup_device_override, torch.device):
            return setup_device_override

        try:
            # After trainer attaches, use Lightning's device tracking
            return super().device
        except Exception:
            return torch.device("cpu")

    def create_trainer(
[docs]
        self,
        max_steps: int,
        callback: MaybeCallback = None,
        run_name: str = "run",
        log_every_n_steps: int | None = None,
        enable_progress_bar: bool | None = None,
        reset_num_timesteps: bool = True,
        default_root_dir: str | pathlib.Path | None = None,
    ) -> L.Trainer:
        """Create an RLTrainer configured like the SB3-compatible learn shim.

        Args:
            max_steps: Total environment steps to collect.
            callback: SB3-style callback(s) fired at each RL step.
            run_name: The run name used in loggers.
            log_every_n_steps: Logging frequency (if None, defaults to ``_default_log_every_n_steps``).
            enable_progress_bar: Whether to display a progress bar using tqdm and rich.
            default_root_dir: Root directory for logs and checkpoints.

        Returns:
            An RLTrainer configured with the same defaults used by learn().
        """
        from lit_sb3.trainer import RLTrainer

        trainer_kwargs: dict[str, Any] = dict(
            max_steps=max_steps,
            rl_callback=callback,
            run_name=run_name,
            log_every_n_steps=log_every_n_steps,
            enable_progress_bar=enable_progress_bar,
        )
        if default_root_dir is not None:
            trainer_kwargs["default_root_dir"] = default_root_dir
        trainer = RLTrainer(**trainer_kwargs)
        trainer._rl_reset_num_timesteps = reset_num_timesteps
        return trainer

    def on_fit_start(self) -> None:
[docs]
        """Called by Lightning before training starts. Initialise env observation."""
        if self.ep_info_buffer is None:
            self.ep_info_buffer = deque(maxlen=self._stats_window_size)
            self.ep_success_buffer = deque(maxlen=self._stats_window_size)
        env = self.env
        if self._last_obs is None and env is not None:
            self._last_obs = env.reset()
            self._last_episode_starts = np.ones((env.num_envs,), dtype=bool)
            if self._vec_normalize_env is not None:
                self._last_original_obs = self._vec_normalize_env.get_original_obs()
        self.start_time = time.time_ns()
        self.set_random_seed(self.seed)

    @staticmethod
    def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> VecEnv:
        """Wrap environment with the appropriate wrappers if needed.

        For instance, to have a vectorized environment
        or to re-order the image channels.

        Args:
            env:
            verbose: Verbosity level: 0 for no output, 1 for indicating wrappers used
            monitor_wrapper: Whether to wrap the env in a ``Monitor`` when possible.

            env: The environment.

        Returns:
            The wrapped environment.
        """
        if not isinstance(env, VecEnv):
            # Patch to support gym 0.21/0.26 and gymnasium
            env = _patch_env(env)
            if not is_wrapped(env, Monitor) and monitor_wrapper:
                if verbose >= 1:
                    print("Wrapping the env with a `Monitor` wrapper")
                env = Monitor(env)
            if verbose >= 1:
                print("Wrapping the env in a DummyVecEnv.")
            env = DummyVecEnv([lambda: env])

        # Make sure that dict-spaces are not nested (not supported)
        check_for_nested_spaces(env.observation_space)

        if not is_vecenv_wrapped(env, VecTransposeImage):
            wrap_with_vectranspose = False
            if isinstance(env.observation_space, spaces.Dict):
                # If even one of the keys is a image-space in need of transpose, apply transpose
                # If the image spaces are not consistent (for instance one is channel first,
                # the other channel last), VecTransposeImage will throw an error
                for space in env.observation_space.spaces.values():
                    wrap_with_vectranspose = wrap_with_vectranspose or (
                        is_image_space(space) and not is_image_space_channels_first(space)  # type: ignore[arg-type]
                    )
            else:
                wrap_with_vectranspose = is_image_space(env.observation_space) and not is_image_space_channels_first(
                    env.observation_space  # type: ignore[arg-type]
                )

            if wrap_with_vectranspose:
                if verbose >= 1:
                    print("Wrapping the env in a VecTransposeImage.")
                env = VecTransposeImage(env)

        return env

    @abstractmethod
    def _setup_model(self) -> None:
        """Create networks, buffer and optimizers."""

    def set_logger(self, logger: Any) -> None:
[docs]
        """Setter for the custom Lightning logger used by ``RLTrainer``.

        .. warning::

                    When passing a custom logger object, this overrides the trainer-managed
                    default logger. Use ``RLTrainer(default_root_dir=...)`` to control the
                    default TensorBoard save path.
        """
        self._trainer_logger = logger
        self._custom_logger = True
[docs]
    def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:  # noqa: ARG002
        """Log a metrics dict via LightningModule ``self.log()``."""
        for key, value in metrics.items():
            if isinstance(value, torch.Tensor):
                value = value.detach().cpu().item() if value.numel() == 1 else value.detach().cpu()
            elif isinstance(value, np.generic):
                value = value.item()
            self.log(key, value, on_step=True, on_epoch=False, prog_bar=False, logger=True)

    def get_summary_writer(self) -> Any:
[docs]
        """Return the active TensorBoard SummaryWriter if one is configured."""
        trainer = getattr(self, "_rl_trainer_ref", None)
        if trainer is None:
            try:
                trainer = self.trainer
            except RuntimeError:
                trainer = None
        if trainer is None:
            return None

        from lit_sb3.callbacks.summary_writer import SummaryWriterCallback

        return SummaryWriterCallback.get_writer(trainer)

    def _setup_lr_schedule(self) -> None:
        """Transform to callable if needed."""
        self.lr_schedule = FloatSchedule(self.learning_rate)

    def _update_current_progress_remaining(self, num_timesteps: int, total_timesteps: int) -> None:
        """Compute current progress remaining (starts from 1 and ends to 0).

        Args:
            num_timesteps: current number of timesteps
            total_timesteps:
            total_timesteps: Total number of timesteps.
        """
        self._current_progress_remaining = 1.0 - float(num_timesteps) / float(total_timesteps)

    def _update_learning_rate(self, optimizers: list[torch.optim.Optimizer] | torch.optim.Optimizer) -> None:
        """Update optimizer LR according to current LR schedule."""
        current_lr = self.lr_schedule(self._current_progress_remaining)
        # Log to Lightning logger (on_step=False is required when called from epoch hooks)
        self.log(
            "train/learning_rate",
            current_lr,
            on_step=False,
            on_epoch=True,
            prog_bar=False,
            logger=True,
        )
        if not isinstance(optimizers, list):
            optimizers = [optimizers]
        for optimizer in optimizers:
            update_learning_rate(optimizer, current_lr)  # ty:ignore[invalid-argument-type]

    def _excluded_save_params(self) -> list[str]:
        """Returns the names of the parameters that should be excluded from being.

        saved by pickling. E.g. replay buffers are skipped by default
        as they take up a lot of space. PyTorch variables should be excluded
        with this so they can be stored with ``torch.save``.

        Returns:
            List of parameters that should be excluded from being saved with pickle.
        """
        return [
            "policy",
            "device",
            "replay_buffer",
            "rollout_buffer",
            "_vec_normalize_env",
            "_custom_logger",
            "_trainer_logger",
            "_rl_trainer_ref",
        ]

    def _get_policy_from_name(self, policy_name: str) -> type[BasePolicy]:
        """Get a policy class from its name representation.

        The goal here is to standardize policy naming, e.g.
        all algorithms can call upon "MlpPolicy" or "CnnPolicy",
        and they receive respective policies that work for them.

        Args:
            policy_name: Alias of the policy

        Returns:
            A policy class (type)
        """
        if policy_name in self.policy_aliases:
            return self.policy_aliases[policy_name]
        else:
            raise ValueError(f"Policy {policy_name} unknown")

    def _get_torch_save_params(self) -> tuple[list[str], list[str]]:
        """Get the name of the torch variables that will be saved with.

        PyTorch ``torch.save``, ``torch.load`` and ``state_dicts`` instead of the default
        pickling strategy. This is to handle device placement correctly.

        Names can point to specific variables under classes, e.g.
        "policy.optimizer" would point to ``optimizer`` object of ``self.policy``
        if this object.

        Returns:
            List of Torch variables whose state dicts to save (e.g. torch.nn.Modules),
            and list of other Torch variables to store with ``torch.save``.
        """
        state_dicts = ["policy"]

        return state_dicts, []

    def _init_callback(
        self,
        callback: MaybeCallback,
    ) -> BaseCallback:
        """Args.

        callback: Callback(s) called at every step with state of the algorithm.
            :param progress_bar: Display a progress bar using tqdm and rich.
            :return: A hybrid callback calling `callback` and performing evaluation.
        """
        # Convert a list of callbacks into a callback
        if isinstance(callback, list):
            callback = CallbackList(callback)  # ty:ignore[invalid-argument-type]

        # Convert functional callback to object
        if not isinstance(callback, BaseCallback):
            callback = ConvertCallback(callback)

        callback.init_callback(self)
        return callback

    def _setup_learn(
        self,
        max_steps: int,
        callback: MaybeCallback = None,
        run_name: str = "run",
    ) -> tuple[int, BaseCallback]:
        """Initialize different variables needed for training.

        Args:
            max_steps: The total number of samples (env steps) to train on
            callback: Callback(s) called at every step with state of the algorithm.
            run_name: The run name used in loggers.
            progress_bar: Display a progress bar using tqdm and rich.

        Returns:
            Total timesteps and callback(s)
        """
        self.start_time = time.time_ns()

        if self.ep_info_buffer is None:
            # Initialize buffers if they don't exist, or reinitialize if resetting counters
            self.ep_info_buffer = deque(maxlen=self._stats_window_size)
            self.ep_success_buffer = deque(maxlen=self._stats_window_size)

        if self.action_noise is not None:
            self.action_noise.reset()

        if True:
            self.num_timesteps = 0
            self._episode_num = 0
        else:
            # Make sure training timesteps are ahead of the internal counter
            max_steps += self.num_timesteps
        self._total_timesteps = max_steps
        self._num_timesteps_at_start = self.num_timesteps

        # Avoid resetting the environment when calling ``.learn()`` consecutive times
        if True or self._last_obs is None:
            env = self.env
            assert env is not None
            self._last_obs = env.reset()
            self._last_episode_starts = np.ones((env.num_envs,), dtype=bool)
            # Retrieve unnormalized observation for saving into the buffer
            if self._vec_normalize_env is not None:
                self._last_original_obs = self._vec_normalize_env.get_original_obs()

        # Create eval callback if needed
        callback = self._init_callback(callback)

        return max_steps, callback

    def _update_info_buffer(self, infos: list[dict[str, Any]], dones: np.ndarray | None = None) -> None:
        """Retrieve reward, episode length, episode success and update the buffer.

        if using Monitor wrapper or a GoalEnv.

        Args:
            infos: List of additional information about the transition.
            dones: Termination signals
        """
        assert self.ep_info_buffer is not None
        assert self.ep_success_buffer is not None

        if dones is None:
            dones = np.array([False] * len(infos))
        for idx, info in enumerate(infos):
            maybe_ep_info = info.get("episode")
            maybe_is_success = info.get("is_success")
            if maybe_ep_info is not None:
                self.ep_info_buffer.extend([maybe_ep_info])
            if maybe_is_success is not None and dones[idx]:
                self.ep_success_buffer.append(maybe_is_success)

    def get_env(self) -> VecEnv | None:
[docs]
        """Returns the current environment (can be None if not defined).

        Resolves through the trainer if available, otherwise returns None.

        Returns:
            The current environment
        """
        return self.env

    def _get_rl_trainer(self) -> Any | None:
        """Return the active RLTrainer reference if one is attached."""
        rl_trainer = getattr(self, "_rl_trainer_ref", None)
        if rl_trainer is not None and hasattr(rl_trainer, "_rl_total_timesteps"):
            return rl_trainer

        try:
            trainer = self.trainer
            if trainer is not None and hasattr(trainer, "_rl_total_timesteps"):
                return trainer
        except RuntimeError:
            pass
        return None

    def _sync_rl_trainer_timesteps(self) -> None:
        """Mirror the latest env-step count into the attached RLTrainer."""
        rl_trainer = self._get_rl_trainer()
        if rl_trainer is not None and hasattr(rl_trainer, "_update_rl_timesteps"):
            rl_trainer._update_rl_timesteps(self.num_timesteps)

    def _has_reached_rl_budget(self) -> bool:
        """Return whether the trainer-owned RL timestep budget has been reached."""
        self._sync_rl_trainer_timesteps()
        rl_trainer = self._get_rl_trainer()
        if rl_trainer is not None and hasattr(rl_trainer, "rl_budget_reached"):
            return bool(rl_trainer.rl_budget_reached)
        return self.num_timesteps >= getattr(self, "_total_timesteps", self.num_timesteps + 1)

    def get_vec_normalize_env(self) -> VecNormalize | None:
[docs]
        """Return the ``VecNormalize`` wrapper of the training env.

        if it exists.

        Returns:
            The ``VecNormalize`` env.
        """
        return self._vec_normalize_env

    def set_env(self, env: GymEnv, force_reset: bool = True) -> None:
[docs]
        """Validate and attach the environment to the trainer.

        The env is stored on the active RLTrainer, not on the model.

        Args:
            env: The environment for learning a policy
            force_reset: Force call to ``reset()`` before training
                to avoid unexpected behavior.
                See issue https://github.com/DLR-RM/stable-baselines3/issues/597
        """
        env = self._wrap_env(env, self.verbose)
        assert env.num_envs == self.n_envs, (
            "The number of environments to be set is different from the number of environments in the model: "
            f"({env.num_envs} != {self.n_envs}), whereas `set_env` requires them to be the same. To load a model with "
            f"a different number of environments, you must use `{self.__class__.__name__}.load(path, env)` instead"
        )
        check_for_correct_spaces(env, self.observation_space, self.action_space)
        self._vec_normalize_env = unwrap_vec_normalize(env)

        if force_reset:
            self._last_obs = None

        self.n_envs = env.num_envs
        # Store on trainer if available
        rl_trainer = getattr(self, "_rl_trainer_ref", None)
        if rl_trainer is not None and hasattr(rl_trainer, "_rl_env"):
            rl_trainer._rl_env = env
            return
        try:
            trainer = self.trainer
            if trainer is not None and hasattr(trainer, "_rl_env"):
                trainer._rl_env = env  # type: ignore[union-attr]
                return
        except RuntimeError:
            pass
        raise RuntimeError("Cannot set env: no RLTrainer is attached. Pass env via RLEnvDataModule.")

    # ------------------------------------------------------------------
    # RLTrainer protocol — implemented by OnPolicyAlgorithm / OffPolicyAlgorithm
    # ------------------------------------------------------------------

    @property
    @abstractmethod
    def _default_log_every_n_steps(self) -> int:
        """Default ``log_every_n_steps`` used when the caller doesn't supply one."""

    @abstractmethod
    def _prepare_fit(
        self,
        callback: BaseCallback,
        log_every_n_steps: int,
    ) -> None:
        """Inject per-algorithm state required by Lightning hooks.

        Called by :class:`~lit_sb3.trainer.RLTrainer` (and the ``learn()`` shim)
        after ``_setup_learn`` and before ``L.Trainer.fit``.

        Args:
            callback: Already-initialised callback returned by ``_setup_learn``.
            log_every_n_steps: Resolved logging frequency (caller's value or
                ``_default_log_every_n_steps``).
        """

    @abstractmethod
    def _rl_max_epochs(self, total_timesteps: int) -> int:
        """Return the ``max_epochs`` value for ``L.Trainer``.

        On-policy: ``ceil(total_timesteps / (n_steps * n_envs))``.
        Off-policy: ``1`` (single epoch driven by an unbounded replay driver and
        trainer-owned RL budget checks).

        Args:
            total_timesteps: Effective total timesteps.
        """

    def _pre_fit_hook(self, callback: BaseCallback) -> None:
        """Algorithm-specific work executed after ``on_training_start`` but before ``L.Trainer.fit``.

        Default: no-op.  On-policy algorithms override this to pre-collect the
        first rollout so Lightning's ``setup_data()`` prefetch finds a populated
        buffer.

        Args:
            callback: The active training callback.
        """

    # ------------------------------------------------------------------

    @abstractmethod
    def learn(
[docs]
        self: SelfBaseAlgorithm,
        total_timesteps: int,
        env: GymEnv | str | None = None,
        callback: MaybeCallback = None,
        log_every_n_steps: int | None = None,
        tb_log_name: str = "run",
        reset_num_timesteps: bool = True,
        progress_bar: bool = False,
    ) -> SelfBaseAlgorithm:
        """Return a trained model.

        Delegates to :class:`~lit_sb3.trainer.RLTrainer` for backwards
        compatibility.  Prefer constructing ``RLTrainer`` explicitly for full
        control over Lightning Trainer settings.

        Args:
            total_timesteps: The total number of samples (env steps) to train on
                Note: it is a lower bound, see `issue #1150 <https://github.com/DLR-RM/stable-baselines3/issues/1150>`_
            env: Training environment or env-id string. Required by the SB3-compatible
                ``learn()`` shim because environments are no longer stored on the model.
            callback: callback(s) called at every step with state of the algorithm.
            log_every_n_steps: How often to emit SB3 stats. For on-policy algos
                (PPO, A2C, …) this is the number of rollout cycles; for off-policy
                algos (SAC, TD3, …) it is the number of episodes.  ``None`` uses the
                algorithm's default.
            tb_log_name: the name of the run for TensorBoard logging
            reset_num_timesteps: whether or not to reset the current timestep number (used in logging)
            progress_bar: Display a progress bar using tqdm and rich.

        Returns:
            the trained model
        """

    def predict(
[docs]
        self,
        observation: np.ndarray | dict[str, np.ndarray],
        state: tuple[np.ndarray, ...] | None = None,
        episode_start: np.ndarray | None = None,
        deterministic: bool = False,
    ) -> 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).

        Args:
            observation: the input observation
            state: The last hidden states (can be None, used in recurrent policies)
            episode_start: 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: Whether or not to return deterministic actions.

        Returns:
            the model's action and the next hidden state
            (used in recurrent policies)
        """
        return self.policy.predict(observation, state, episode_start, deterministic)

    def set_random_seed(self, seed: int | None = None) -> None:
[docs]
        """Set the seed of the pseudo-random generators.

        (python, numpy, pytorch, gym, action_space).

        Args:
            seed:
            seed: Random seed for reproducibility.
        """
        if seed is None:
            return
        set_random_seed(seed, using_cuda=self.device.type == torch.device("cuda").type)
        self.action_space.seed(seed)
        env = self.env
        if env is not None:
            env.seed(seed)

    def set_parameters(
[docs]
        self,
        load_path_or_dict: str | TensorDict,
        exact_match: bool = True,
        device: torch.device | str = "auto",
    ) -> None:
        """Load parameters from a given zip-file or a nested dictionary containing parameters for.

        different modules (see ``get_parameters``).

        Args:
            load_path_or_iter: Location of the saved data (path or file-like, see ``save``), or a nested
                dictionary containing nn.Module parameters used by the policy. The dictionary maps
                object names to a state-dictionary returned by ``torch.nn.Module.state_dict()``.
            exact_match: If True, the given parameters should include parameters for each
                module and each of their parameters, otherwise raises an Exception. If set to False, this
                can be used to update only specific parameters.
            device: Device on which the code should run.
            load_path_or_dict: Path to model file or dict with parameters.
        """
        params = {}
        if isinstance(load_path_or_dict, dict):
            params = load_path_or_dict
        else:
            _, params, _ = load_from_zip_file(load_path_or_dict, device=device, load_data=False)

        # Keep track which objects were updated.
        # `_get_torch_save_params` returns [params, other_pytorch_variables].
        # We are only interested in former here.
        objects_needing_update = set(self._get_torch_save_params()[0])
        updated_objects = set()

        for name in params:
            attr = None
            try:
                attr = recursive_getattr(self, name)
            except Exception as e:
                # What errors recursive_getattr could throw? KeyError, but
                # possible something else too (e.g. if key is an int?).
                # Catch anything for now.
                raise ValueError(f"Key {name} is an invalid object name.") from e

            if isinstance(attr, torch.optim.Optimizer):
                # Optimizers do not support "strict" keyword...
                # Seems like they will just replace the whole
                # optimizer state with the given one.
                # On top of this, optimizer state-dict
                # seems to change (e.g. first ``optim.step()``),
                # which makes comparing state dictionary keys
                # invalid (there is also a nesting of dictionaries
                # with lists with dictionaries with ...), adding to the
                # mess.
                #
                # TL;DR: We might not be able to reliably say
                # if given state-dict is missing keys.
                #
                # Solution: Just load the state-dict as is, and trust
                # the user has provided a sensible state dictionary.
                attr.load_state_dict(params[name])
            else:
                # Assume attr is torch.nn.Module
                attr.load_state_dict(params[name], strict=exact_match)
            updated_objects.add(name)

        if exact_match and updated_objects != objects_needing_update:
            raise ValueError(
                "Names of parameters do not match agents' parameters: "
                f"expected {objects_needing_update}, got {updated_objects}"
            )

    @classmethod
    def load(  # noqa: C901
[docs]
        cls: type[SelfBaseAlgorithm],
        path: str | pathlib.Path | io.BufferedIOBase,
        env: GymEnv | None = None,
        device: torch.device | str = "auto",
        custom_objects: dict[str, Any] | None = None,
        print_system_info: bool = False,
        force_reset: bool = True,
        **kwargs,
    ) -> SelfBaseAlgorithm:
        """Load the model from a zip-file.

        Warning: ``load`` re-creates the model from scratch, it does not update it in-place!
        For an in-place load use ``set_parameters`` instead.

        Args:
            path: path to the file (or a file-like) where to
                load the agent from
            env: the new environment to run the loaded model on
                (can be None if you only need prediction from a trained model) has priority over any saved environment
            device: Device on which the code should run.
            custom_objects: Dictionary of objects to replace
                upon loading. If a variable is present in this dictionary as a
                key, it will not be deserialized and the corresponding item
                will be used instead. Similar to custom_objects in
                ``keras.models.load_model``. Useful when you have an object in
                file that can not be deserialized.
            print_system_info: Whether to print system info from the saved model
                and the current system info (useful to debug loading issues)
            force_reset: Force call to ``reset()`` before training
                to avoid unexpected behavior.
                See https://github.com/DLR-RM/stable-baselines3/issues/597
            kwargs: extra arguments to change the model when loading

        Returns:
            new model instance with loaded parameters
        """
        resolved_device = get_device(device)

        if print_system_info:
            print("== CURRENT SYSTEM INFO ==")
            get_system_info()

        data, params, pytorch_variables = load_from_zip_file(
            path,
            device=resolved_device,
            custom_objects=custom_objects,
            print_system_info=print_system_info,
        )

        assert data is not None, "No data found in the saved file"
        assert params is not None, "No params found in the saved file"

        # Remove stored device information and replace with ours
        if "policy_kwargs" in data:
            if "device" in data["policy_kwargs"]:
                del data["policy_kwargs"]["device"]
            # backward compatibility, convert to new format
            saved_net_arch = data["policy_kwargs"].get("net_arch")
            if saved_net_arch and isinstance(saved_net_arch, list) and isinstance(saved_net_arch[0], dict):
                data["policy_kwargs"]["net_arch"] = saved_net_arch[0]

        if "policy_kwargs" in kwargs and kwargs["policy_kwargs"] != data["policy_kwargs"]:
            raise ValueError(
                f"The specified policy kwargs do not equal the stored policy kwargs."
                f"Stored kwargs: {data['policy_kwargs']}, specified kwargs: {kwargs['policy_kwargs']}"
            )

        if "observation_space" not in data or "action_space" not in data:
            raise KeyError("The observation_space and action_space were not given, can't verify new environments")

        # Gym -> Gymnasium space conversion
        for key in {"observation_space", "action_space"}:
            data[key] = _convert_space(data[key])

        if env is not None:
            # Wrap first if needed
            env = cls._wrap_env(env, data["verbose"])
            # Check if given env is valid
            check_for_correct_spaces(env, data["observation_space"], data["action_space"])
            # Discard `_last_obs`, this will force the env to reset before training
            # See issue https://github.com/DLR-RM/stable-baselines3/issues/597
            if force_reset and data is not None:
                data["_last_obs"] = None
            # `n_envs` must be updated. See issue https://github.com/DLR-RM/stable-baselines3/issues/1018
            if data is not None:
                data["n_envs"] = env.num_envs
        else:
            # Use stored env, if one exists. If not, continue as is (can be used for predict)
            if "_init_env" in data:
                env = data.pop("_init_env", None)
            elif "env" in data:
                env = data.pop("env", None)

        model = cls(
            policy=data["policy_class"],
            device=resolved_device,
            _init_setup_model=False,  # type: ignore[call-arg]
        )  # ty:ignore[missing-argument]

        # load parameters
        model.__dict__.update(data)
        model.__dict__.update(kwargs)

        # If an env was provided or stored, extract metadata
        if env is not None:
            if not isinstance(env, VecEnv):
                env = cls._wrap_env(env, data.get("verbose", 0))
            model._init_env_metadata(env, model._supported_action_spaces, model._support_multi_env)

        model.__dict__["_setup_device_override"] = resolved_device
        model._setup_model()
        delattr(model, "_setup_device_override")

        try:
            # put state_dicts back in place
            model.set_parameters(params, exact_match=True, device=resolved_device)
        except RuntimeError as e:
            # Patch to load policies saved using SB3 < 1.7.0
            # the error is probably due to old policy being loaded
            # See https://github.com/DLR-RM/stable-baselines3/issues/1233
            if "pi_features_extractor" in str(e) and "Missing key(s) in state_dict" in str(e):
                model.set_parameters(params, exact_match=False, device=resolved_device)
                warnings.warn(
                    "You are probably loading a A2C/PPO model saved with SB3 < 1.7.0, "
                    "we deactivated exact_match so you can save the model "
                    "again to avoid issues in the future "
                    "(see https://github.com/DLR-RM/stable-baselines3/issues/1233 for more info). "
                    f"Original error: {e} \n"
                    "Note: the model should still work fine, this only a warning.",
                    stacklevel=2,
                )
            else:
                raise
        except ValueError as e:
            # Patch to load DQN policies saved using SB3 < 2.4.0
            # The target network params are no longer in the optimizer
            # See https://github.com/DLR-RM/stable-baselines3/pull/1963
            saved_optim_params = params["policy.optimizer"]["param_groups"][0]["params"]  # type: ignore[index]
            n_params_saved = len(saved_optim_params)
            n_params = len(model.policy.optimizer.param_groups[0]["params"])
            if n_params_saved == 2 * n_params:
                # Truncate to include only online network params
                params["policy.optimizer"]["param_groups"][0]["params"] = saved_optim_params[:n_params]  # type: ignore[index]

                model.set_parameters(params, exact_match=True, device=resolved_device)
                warnings.warn(
                    "You are probably loading a DQN model saved with SB3 < 2.4.0, "
                    "we truncated the optimizer state so you can save the model "
                    "again to avoid issues in the future "
                    "(see https://github.com/DLR-RM/stable-baselines3/pull/1963 for more info). "
                    f"Original error: {e} \n"
                    "Note: the model should still work fine, this only a warning.",
                    stacklevel=2,
                )
            else:
                raise

        # put other pytorch variables back in place
        if pytorch_variables is not None:
            for name in pytorch_variables:
                # Skip if PyTorch variable was not defined (to ensure backward compatibility).
                # This happens when using SAC/TQC.
                # SAC has an entropy coefficient which can be fixed or optimized.
                # If it is optimized, an additional PyTorch variable `log_ent_coef` is defined,
                # otherwise it is initialized to `None`.
                if pytorch_variables[name] is None:
                    continue
                # Set the data attribute directly to avoid issue when using optimizers
                # See https://github.com/DLR-RM/stable-baselines3/issues/391
                recursive_setattr(model, f"{name}.data", pytorch_variables[name].data)

        # Sample gSDE exploration matrix, so it uses the right device
        # see issue #44
        if model.use_sde:
            model.policy.reset_noise()  # type: ignore[operator]
        return model

    def get_parameters(self) -> dict[str, dict]:
[docs]
        """Return the parameters of the agent. This includes parameters from different networks, e.g.

        critics (value functions) and policies (pi functions).

        Returns:
            Mapping of from names of the objects to PyTorch state-dicts.
        """
        state_dicts_names, _ = self._get_torch_save_params()
        params = {}
        for name in state_dicts_names:
            attr = recursive_getattr(self, name)
            # Retrieve state dict
            params[name] = attr.state_dict()
        return params

    def save(
[docs]
        self,
        path: str | pathlib.Path | io.BufferedIOBase,
        exclude: Iterable[str] | None = None,
        include: Iterable[str] | None = None,
    ) -> None:
        """Save all the attributes of the object and the model parameters in a zip-file.

        Args:
            path: path to the file where the rl agent should be saved
            exclude: name of parameters that should be excluded in addition to the default ones
            include: name of parameters that might be excluded but should be included anyway
        """
        # Copy parameter list so we don't mutate the original dict
        data = self.__dict__.copy()

        # Exclude is union of specified parameters (if any) and standard exclusions
        if exclude is None:
            exclude = []
        exclude = set(exclude).union(self._excluded_save_params())

        # Do not exclude params if they are specifically included
        if include is not None:
            exclude = exclude.difference(include)

        state_dicts_names, torch_variable_names = self._get_torch_save_params()
        all_pytorch_variables = state_dicts_names + torch_variable_names
        for torch_var in all_pytorch_variables:
            # We need to get only the name of the top most module as we'll remove that
            var_name = torch_var.split(".")[0]
            # Any params that are in the save vars must not be saved by data
            exclude.add(var_name)

        # Remove parameter entries of parameters which are to be excluded
        for param_name in exclude:
            data.pop(param_name, None)

        # Build dict of torch variables
        pytorch_variables = None
        if torch_variable_names is not None:
            pytorch_variables = {}
            for name in torch_variable_names:
                attr = recursive_getattr(self, name)
                pytorch_variables[name] = attr

        # Build dict of state_dicts
        params_to_save = self.get_parameters()

        save_to_zip_file(path, data=data, params=params_to_save, pytorch_variables=pytorch_variables)

    def dump_logs(self) -> None:
[docs]
        """Write log data. (Implemented by OffPolicyAlgorithm and OnPolicyAlgorithm)."""
        raise NotImplementedError()

    def _dump_logs(self, *args) -> None:
        warnings.warn(
            "algo._dump_logs() is deprecated in favor of algo.dump_logs(). It will be removed in SB3 v2.7.0",
            stacklevel=2,
        )
        self.dump_logs(*args)