# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations
import abc
import contextlib
import itertools
import json
import math
import pathlib
import signal
import sys
import time
import warnings
import weakref
from collections import defaultdict, OrderedDict
from collections.abc import Callable, Collection, Mapping, Sequence
from copy import deepcopy
from textwrap import indent
from typing import Any, Literal
import numpy as np
import torch
from packaging import version
from tensordict import NestedKey, NonTensorData, pad, TensorDict, TensorDictBase
from tensordict._tensorcollection import TensorCollection
from tensordict.nn import TensorDictModule
from tensordict.utils import expand_right
from torch import nn, optim
from torchrl._utils import (
_CKPT_BACKEND,
implement_for,
KeyDependentDefaultDict,
logger as torchrl_logger,
rl_warnings,
timeit,
VERBOSE,
)
from torchrl.checkpoint import (
Checkpoint,
CheckpointRotation,
GlobalRNGState,
resolve_checkpoint_path,
StopOnSignal,
)
from torchrl.collectors import BaseCollector, Evaluator
from torchrl.collectors.utils import split_trajectories
from torchrl.data.replay_buffers import (
PrioritizedSampler,
TensorDictPrioritizedReplayBuffer,
TensorDictReplayBuffer,
)
from torchrl.data.utils import DEVICE_TYPING
from torchrl.envs.common import EnvBase
from torchrl.envs.utils import ExplorationType, set_exploration_type
from torchrl.objectives.common import LossModule
from torchrl.objectives.utils import TargetNetUpdater
from torchrl.record.loggers import Logger
_TORCH_GRAD_SCALER_HAS_DEVICE = version.parse(torch.__version__).release >= (2, 3)
if _TORCH_GRAD_SCALER_HAS_DEVICE:
from torch.amp import GradScaler
else:
from torch.cuda.amp import GradScaler
try:
from tqdm import tqdm
_has_tqdm = True
except ImportError:
_has_tqdm = False
try:
from torchsnapshot import Snapshot, StateDict
_has_ts = True
except ImportError:
_has_ts = False
REPLAY_BUFFER_CLASS = {
"prioritized": TensorDictPrioritizedReplayBuffer,
"circular": TensorDictReplayBuffer,
}
# Mapping of metric names to logger methods - controls how different metrics are logged
LOGGER_METHODS = {
"grad_norm": "log_scalar",
"loss": "log_scalar",
}
# Format strings for different data types in progress bar display
TYPE_DESCR = {float: "4.4f", int: ""}
REWARD_KEY = ("next", "reward")
_OPTIM_STEPS_UNSET = object()
@implement_for("torch", "2.3")
def _make_grad_scaler(device_type: str, enabled: bool) -> GradScaler:
return GradScaler(device_type, enabled=enabled)
@implement_for("torch", None, "2.3")
def _make_grad_scaler(device_type: str, enabled: bool) -> GradScaler: # noqa: F811
return GradScaler(enabled=enabled)
# On Windows, a memory-mapped checkpoint keeps the file locked for as long as
# the loaded tensors are alive, so the checkpoint could neither be deleted nor
# safely re-saved. Only default to ``mmap=True`` on other platforms.
_MMAP_CKPT_DEFAULT = sys.platform != "win32"
@implement_for("torch", "2.4")
def _torch_load_defaults() -> dict[str, Any]:
return {"weights_only": True, "mmap": _MMAP_CKPT_DEFAULT}
@implement_for("torch", None, "2.4")
def _torch_load_defaults() -> dict[str, Any]: # noqa: F811
# The weights-only unpickler of torch < 2.4 does not allow
# ``torch.device``, which TensorDict state-dicts carry as their
# ``__device`` metadata entry.
return {"weights_only": False, "mmap": _MMAP_CKPT_DEFAULT}
class _TrainerCheckpointState:
"""State-dict view over Trainer progress counters."""
def __init__(self, trainer: Trainer) -> None:
self.trainer = trainer
def state_dict(self) -> dict[str, Any]:
return dict(self.trainer._get_state())
def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
self.trainer.collected_frames = state_dict["collected_frames"]
self.trainer._last_log = state_dict["_last_log"]
self.trainer._last_save = state_dict["_last_save"]
self.trainer._optim_count = state_dict["_optim_count"]
class _ExecutionCheckpointState:
"""Semantic state view over a private learner execution backend."""
def __init__(self, trainer: Trainer) -> None:
self.trainer = trainer
def state_dict(self) -> dict[str, Any]:
backend = self.trainer._execution_backend
if backend is None or not backend.is_alive():
raise RuntimeError(
"Remote learner checkpointing requires a running learner backend."
)
return {
"backend": backend.state_dict(),
"controller": self.trainer._execution_controller_state(),
}
def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
backend = self.trainer._execution_backend
if backend is None:
raise RuntimeError("Remote learner execution backend is missing.")
if not backend.is_alive():
backend.start()
backend.load_state_dict(state_dict["backend"])
self.trainer._load_execution_controller_state(state_dict.get("controller", {}))
self.trainer._published_model_version = -1
def _state_dict_to_td(sd: dict) -> TensorDict:
"""Convert a state dict to a :class:`~tensordict.TensorDict`.
Tensor values are stored directly; everything else is wrapped in
:class:`~tensordict.NonTensorData` so that :meth:`~tensordict.TensorDict.dumps`
can persist the whole state without pickle dependencies.
"""
return TensorDict(
{
k: v if isinstance(v, torch.Tensor) else NonTensorData(v)
for k, v in sd.items()
},
[],
)
def _td_to_state_dict(td: TensorDict) -> dict:
"""Inverse of :func:`_state_dict_to_td`.
Unwraps :class:`~tensordict.NonTensorData` back to plain Python values and
leaves tensors (including :class:`~tensordict.MemoryMappedTensor`) as-is.
"""
return {k: v.data if isinstance(v, NonTensorData) else v for k, v in td.items()}
[docs]
class TrainerHookBase:
"""An abstract hooking class for torchrl Trainer class."""
@abc.abstractmethod
def state_dict(self) -> dict[str, Any]:
raise NotImplementedError
@abc.abstractmethod
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
raise NotImplementedError
[docs]
@abc.abstractmethod
def register(self, trainer: Trainer, name: str):
"""Registers the hook in the trainer at a default location.
Args:
trainer (Trainer): the trainer where the hook must be registered.
name (str): the name of the hook.
.. note::
To register the hook at another location than the default, use
:meth:`~torchrl.trainers.Trainer.register_op`.
"""
raise NotImplementedError
class OptimizationStepper(TrainerHookBase):
"""Performs a single optimization step in a Trainer.
The optimization stepper encapsulates the logic executed for each ``sub_batch``
during training. This is useful for algorithms that require multiple optimizers,
delayed updates (e.g. TD3, where critics are updated every step while the actor
and target networks are updated less frequently), or multiple backward passes
within one training iteration.
The :class:`~torchrl.trainers.Trainer` calls :meth:`step` inside its optimization
loop and handles post-optimization hooks (e.g. target network updates, priority
updates, schedulers) and logging around this call.
Subclasses should return a :class:`~tensordict.TensorDictBase` of detached scalar
values suitable for logging.
"""
_trainer: Trainer
def step(self, trainer: Trainer, sub_batch: TensorDictBase) -> TensorDictBase:
"""Perform one optimization step on a ``sub_batch``.
Args:
trainer (Trainer): The trainer executing the optimization loop.
sub_batch (TensorDictBase): Batch used for this optimization step.
Returns:
A TensorDict containing detached scalar metrics for logging.
"""
raise NotImplementedError
def _step(self, context: Any, sub_batch: TensorDictBase) -> TensorDictBase:
"""Run this stepper against a private optimization context."""
return self.step(context, sub_batch)
def state_dict(self) -> dict[str, Any]:
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
return
def register(self, trainer: Trainer, name: str = "optimization_stepper") -> None:
"""Register the stepper with a Trainer for checkpointing."""
# Register as a module so it is included in Trainer checkpoints.
# This is not a hook stage (i.e., it is not registered via ``register_op``).
trainer.register_module(name, self)
self._trainer = trainer
class DefaultOptimizationStepper(OptimizationStepper):
"""Default optimization step implementation.
This stepper computes losses via ``trainer.loss_module(sub_batch)`` and applies a
single optimizer update with ``trainer.optimizer`` (including gradient clipping
when configured).
Optionally, a subset of loss entries can be selected via ``loss_components``.
In that case, only the selected keys contribute to the backward pass.
"""
def __init__(self, loss_components: Sequence[str] | None = None) -> None:
if loss_components is not None and not loss_components:
raise ValueError(
"loss_components list cannot be empty. "
"Set to None to act on all components of the loss."
)
self.loss_components = (
set(loss_components) if loss_components is not None else None
)
@staticmethod
def _compute_and_clip_grad_norm(
optimizer: optim.Optimizer,
clip_grad_norm: bool,
clip_norm: float | None,
) -> float:
params = []
for param_group in optimizer.param_groups:
params += param_group["params"]
if clip_grad_norm and clip_norm is not None:
gn = nn.utils.clip_grad_norm_(params, clip_norm)
else:
gn = sum([p.grad.pow(2).sum() for p in params if p.grad is not None]).sqrt()
if clip_norm is not None:
nn.utils.clip_grad_value_(params, clip_norm)
return float(gn)
def step(self, trainer: Trainer, sub_batch: TensorDictBase) -> TensorDictBase:
losses_td = trainer.compute_loss(sub_batch)
if trainer.optimizer is None:
raise RuntimeError(
"DefaultOptimizationStepper requires an optimizer. "
"Pass `optimizer=` to Trainer or use a custom "
"OptimizationStepper that owns its optimizer(s)."
)
if self.loss_components is not None:
items = [
item for key, item in losses_td.items() if key in self.loss_components
]
else:
items = [item for key, item in losses_td.items() if key.startswith("loss")]
loss = sum(items)
loss.backward()
gn = self._compute_and_clip_grad_norm(
trainer.optimizer,
trainer.clip_grad_norm,
trainer.clip_norm,
)
losses_td["grad_norm"] = torch.tensor(gn)
trainer.optimizer.step()
trainer.optimizer.zero_grad()
return losses_td
[docs]
class MixedPrecisionOptimizationStepper(OptimizationStepper):
"""Optimization step with mixed precision and gradient accumulation.
This stepper wraps each forward/backward pass in ``torch.amp.autocast``
and optionally scales gradients with ``torch.amp.GradScaler`` (for fp16).
It also implements *gradient accumulation*: gradients are accumulated for
``gradient_accumulation_steps`` micro-batches before the optimizer is
stepped and zeroed.
It can be used with any :class:`~torchrl.trainers.Trainer`; LLM trainers
such as :class:`~torchrl.trainers.algorithms.GRPOTrainer` construct it by
default.
Args:
optimizer (optim.Optimizer): The optimizer to use.
Keyword Args:
mixed_precision (bool, optional): Whether to enable mixed-precision
training. Default: ``False``.
autocast_dtype (torch.dtype, optional): The dtype to use inside
``autocast``. Default: ``torch.bfloat16``.
gradient_accumulation_steps (int, optional): Number of micro-batches
over which gradients are accumulated before a step. Default: ``1``.
clip_norm (float, optional): Maximum gradient norm for clipping.
Default: ``1.0``.
device_type (str, optional): Device type passed to ``autocast`` and
``GradScaler`` (e.g. ``"cuda"`` or ``"cpu"``). Defaults to the
device type of the optimizer's first parameter.
.. note::
``GradScaler`` is only enabled when ``mixed_precision=True`` *and*
``autocast_dtype=torch.float16``. With bfloat16 (the recommended
dtype for modern GPUs) the scaler is a no-op and is not created.
"""
def __init__(
self,
optimizer: optim.Optimizer,
*,
mixed_precision: bool = False,
autocast_dtype: torch.dtype = torch.bfloat16,
gradient_accumulation_steps: int = 1,
clip_norm: float | None = 1.0,
device_type: str | None = None,
) -> None:
if gradient_accumulation_steps < 1:
raise ValueError("gradient_accumulation_steps must be >= 1")
self.optimizer = optimizer
self.mixed_precision = mixed_precision
self.autocast_dtype = autocast_dtype
self.gradient_accumulation_steps = gradient_accumulation_steps
self.clip_norm = clip_norm
if device_type is None:
params = [p for group in optimizer.param_groups for p in group["params"]]
device_type = params[0].device.type if params else "cpu"
self.device_type = device_type
# GradScaler is only useful for fp16; bf16 doesn't need it.
self._use_scaler = mixed_precision and (autocast_dtype == torch.float16)
self.scaler = _make_grad_scaler(self.device_type, self._use_scaler)
# Internal micro-batch counter (reset after every optimizer step).
self._micro_step: int = 0
self._optimizer_step_count: int = 0
self._skipped_nonfinite_steps: int = 0
@property
def optimizer_step_count(self) -> int:
"""Number of completed optimizer steps.
Discounts gradient-accumulation micro-steps and steps skipped by the
GradScaler on overflow or by the non-finite guards. Read by hooks that
act on an optimizer-step cadence (e.g.
:class:`~torchrl.trainers.UpdateWeights` with
``interval_unit="optim_steps"``).
"""
return self._optimizer_step_count
# ------------------------------------------------------------------
# Checkpointing (optimizer + scaler state)
# ------------------------------------------------------------------
def state_dict(self) -> dict[str, Any]:
if self._micro_step % self.gradient_accumulation_steps != 0:
raise RuntimeError(
f"Cannot save stepper state mid-accumulation. (micro_step={self._micro_step}, "
f"accumulation_steps={self.gradient_accumulation_steps}). "
"Adjust your save_interval to align with the gradient accumulation window."
)
sd: dict[str, Any] = {
"optimizer": self.optimizer.state_dict(),
"micro_step": self._micro_step,
"optimizer_step_count": self._optimizer_step_count,
"skipped_nonfinite_steps": self._skipped_nonfinite_steps,
}
if self._use_scaler:
sd["scaler"] = self.scaler.state_dict()
return sd
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
self.optimizer.load_state_dict(state_dict["optimizer"])
self._micro_step = state_dict.get("micro_step", 0)
self._optimizer_step_count = state_dict.get("optimizer_step_count", 0)
self._skipped_nonfinite_steps = state_dict.get("skipped_nonfinite_steps", 0)
if self._use_scaler and "scaler" in state_dict:
self.scaler.load_state_dict(state_dict["scaler"])
# ------------------------------------------------------------------
# Core step
# ------------------------------------------------------------------
[docs]
def step(self, trainer: Trainer, sub_batch: TensorDictBase) -> TensorDictBase:
"""Perform one forward pass and scaled backward pass.
The optimizer is only stepped and zeroed every
``gradient_accumulation_steps`` calls.
Args:
trainer (Trainer): The owning :class:`~torchrl.trainers.Trainer`.
sub_batch (TensorDictBase): Mini-batch used for this step.
Returns:
A :class:`~tensordict.TensorDict` with scalar metrics (losses,
grad_norm) suitable for logging.
"""
# ---- forward pass (optionally under autocast) ----
with torch.amp.autocast(
self.device_type,
enabled=self.mixed_precision,
dtype=self.autocast_dtype,
):
losses_td = trainer.compute_loss(sub_batch)
# Sum all loss_* keys and normalise by accumulation steps.
loss_items = [v for k, v in losses_td.items() if k.startswith("loss")]
if not loss_items:
raise RuntimeError(
"The loss module returned no 'loss_*' keys. "
"Make sure your loss module prefixes scalar outputs with 'loss'."
)
loss = sum(loss_items) / self.gradient_accumulation_steps
# ---- non-finite loss guard ----
# A single non-finite loss would poison the gradients accumulated so
# far, so the whole accumulation window is dropped and restarted.
if not torch.isfinite(loss.detach()).all():
self._skipped_nonfinite_steps += 1
torchrl_logger.warning(
f"Skipping optimization step because the loss is non-finite: "
f"{loss.detach()}. Skipped non-finite steps: "
f"{self._skipped_nonfinite_steps}."
)
self.optimizer.zero_grad(set_to_none=True)
self._micro_step = 0
return self._reduce_metrics(losses_td)
# ---- backward pass ----
if self._use_scaler:
self.scaler.scale(loss).backward()
else:
loss.backward()
self._micro_step += 1
# ---- optimizer step every `gradient_accumulation_steps` micro-batches ----
grad_norm = None
if self._micro_step % self.gradient_accumulation_steps == 0:
if self._use_scaler:
self.scaler.unscale_(self.optimizer)
grad_norm = 0.0
if self.clip_norm is not None:
params = [
p for group in self.optimizer.param_groups for p in group["params"]
]
grad_norm = float(nn.utils.clip_grad_norm_(params, self.clip_norm))
if self._use_scaler:
scale_before = self.scaler.get_scale()
self.scaler.step(self.optimizer)
self.scaler.update()
# If scale dropped, an overflow occurred and optimizer.step() was skipped.
if self.scaler.get_scale() >= scale_before:
self._optimizer_step_count += 1
elif self.clip_norm is not None and not math.isfinite(grad_norm):
# Without a GradScaler (e.g. bf16 or full precision) nothing
# skips overflowing updates, so guard the step explicitly.
self._skipped_nonfinite_steps += 1
torchrl_logger.warning(
f"Skipping optimizer step because the gradient norm is "
f"non-finite: {grad_norm}. Skipped non-finite steps: "
f"{self._skipped_nonfinite_steps}."
)
grad_norm = 0.0
else:
self.optimizer.step()
self._optimizer_step_count += 1
self.optimizer.zero_grad(set_to_none=True)
metrics = self._reduce_metrics(losses_td)
if grad_norm is not None:
metrics["grad_norm"] = torch.tensor(grad_norm)
return metrics
@staticmethod
def _reduce_metrics(losses_td: TensorDictBase) -> TensorDictBase:
"""Detach the loss output and reduce non-scalar entries to their mean.
Steppers must return scalar metrics suitable for logging; loss modules
such as :class:`~torchrl.objectives.llm.GRPOLoss` also emit per-token
diagnostics (KL divergences) that are reduced here.
"""
metrics = {}
for key, value in losses_td.detach().items():
metrics[key] = value.float().mean() if value.numel() > 1 else value
return TensorDict(metrics, [])
[docs]
class Trainer:
"""A generic Trainer class.
A trainer is responsible for collecting data and training the model.
To keep the class as versatile as possible, Trainer does not construct any
of its specific operations: they all must be hooked at specific points in
the training loop.
To build a Trainer, one needs an iterable data source (a :obj:`collector`), a
loss module and an optimizer.
Args:
collector (Sequence[TensorDictBase]): An iterable returning batches of
data in a TensorDict form of shape [batch x time steps].
total_frames (int): Total number of frames to be collected during
training.
loss_module (LossModule): A module that reads TensorDict batches
(possibly sampled from a replay buffer) and return a loss
TensorDict where every key points to a different loss component.
optimizer (optim.Optimizer): An optimizer that trains the parameters
of the model.
logger (Logger, optional): a Logger that will handle the logging.
optim_steps_per_batch (int, optional): number of optimization steps
per collection of data. An trainer works as follows: a main loop
collects batches of data (epoch loop), and a sub-loop (training
loop) performs model updates in between two collections of data.
If `None`, the trainer will use the number of workers as the number of optimization steps.
clip_grad_norm (bool, optional): If True, the gradients will be clipped
based on the total norm of the model parameters. If False,
all the partial derivatives will be clamped to
(-clip_norm, clip_norm). Default is ``True``.
clip_norm (Number, optional): value to be used for clipping gradients.
Default is None (no clip norm).
progress_bar (bool, optional): If True, a progress bar will be
displayed using tqdm. If tqdm is not installed, this option
won't have any effect. Default is ``True``
seed (int, optional): Seed to be used for the collector, pytorch and
numpy. Default is ``None``.
save_trainer_interval (int, optional): How often the trainer should be
saved to disk, in frame count. Default is 10000.
log_interval (int, optional): How often the values should be logged,
in frame count. Default is 10000.
save_trainer_file (path, optional): path where to save the trainer.
Default is None (no saving)
checkpoint (Checkpoint, optional): unified checkpoint object used for
scheduled saves and restores. The trainer registers any missing
standard components on this object, including the process-global
RNG state under ``"rng"``, which :meth:`load_from_file` restores
after every other component. When omitted, the legacy
``CKPT_BACKEND`` path is retained during the compatibility window.
checkpoint_rotation (CheckpointRotation, optional): retention policy used
for scheduled unified checkpoints. Requires ``checkpoint`` and cannot
be combined with ``save_trainer_file``.
checkpoint_metadata (Callable, optional): function called with the trainer
before each rotated save. Its returned mapping is added to the
checkpoint manifest metadata.
async_collection (bool, optional): Whether to collect data asynchronously.
This will only work if the replay buffer is registered within the data collector.
If using this, the UTD ratio (Update to Data) will be logged under the key "utd_ratio".
Default is False.
log_timings (bool, optional): If True, automatically register a LogTiming hook to log
timing information for all hooks to the logger (e.g., wandb, tensorboard).
Timing metrics will be logged with prefix "time/" (e.g., "time/hook/UpdateWeights").
Default is False.
auto_log_optim_steps (bool, optional): If True, automatically log ``optim_steps`` and the
keys of the averaged loss TensorDict at the end of every optimization loop, in addition
to anything ``post_optim_complete_log`` hooks return. Set to False to fully delegate
this logging to user-registered hooks. Default is True.
replay_buffer (optional): Replay owner used by a remote learner backend.
target_net_updater (TargetNetUpdater, optional): Target updater serialized
with the learner object graph.
batch_size (int, optional): Global learner batch size. Defaults to the
replay buffer batch size.
learner_backend (str): Optimization placement, ``"local"`` or ``"ray"``.
Defaults to ``"local"``.
learner_backend_options (dict, optional): Backend-specific options.
learner_poll_interval (float): Remote replay polling interval. Defaults
to ``0.05`` seconds.
"""
@classmethod
def __new__(cls, *args, **kwargs):
# Training state trackers (used for logging and checkpointing)
cls._optim_count: int = 0 # Total number of optimization steps completed
cls._collected_frames: int = 0 # Total number of frames collected (deprecated)
cls._last_log: dict[
str, Any
] = {} # Tracks when each metric was last logged (for log_interval control)
cls._last_save: int = (
0 # Tracks when trainer was last saved (for save_interval control)
)
cls.collected_frames = 0 # Total number of frames collected (current)
cls._app_state = None # Application state for checkpointing
return super().__new__(cls)
def __init__(
self,
*,
collector: BaseCollector,
total_frames: int,
frame_skip: int,
optim_steps_per_batch: int,
loss_module: LossModule | Callable[[TensorDictBase], TensorDictBase],
optimizer: optim.Optimizer | None = None,
optimization_stepper: OptimizationStepper | None = None,
replay_buffer: Any | None = None,
target_net_updater: TargetNetUpdater | None = None,
batch_size: int | None = None,
learner_backend: Literal["local", "ray"] = "local",
learner_backend_options: dict[str, Any] | None = None,
learner_poll_interval: float = 0.05,
logger: Logger | None = None,
clip_grad_norm: bool = True,
clip_norm: float | None = None,
progress_bar: bool = True,
seed: int | None = None,
save_trainer_interval: int = 10000,
log_interval: int = 10000,
save_trainer_file: str | pathlib.Path | None = None,
checkpoint: Checkpoint | None = None,
checkpoint_rotation: CheckpointRotation | None = None,
checkpoint_metadata: Callable[[Trainer], Mapping[str, Any]] | None = None,
num_epochs: int = 1,
async_collection: bool = False,
log_timings: bool = False,
auto_log_optim_steps: bool = True,
) -> None:
# objects
self.frame_skip = frame_skip
self.collector = collector
self.loss_module = loss_module
self.optimizer = optimizer
self.replay_buffer = replay_buffer
self.target_net_updater = target_net_updater
self.batch_size = batch_size
self.learner_backend = learner_backend
self.learner_backend_options = dict(learner_backend_options or {})
if learner_backend not in ("local", "ray"):
raise ValueError("learner_backend must be 'local' or 'ray'.")
if learner_backend == "ray":
if isinstance(optim_steps_per_batch, bool) or not isinstance(
optim_steps_per_batch, int
):
raise TypeError(
"optim_steps_per_batch must be an integer with "
"learner_backend='ray'."
)
if optim_steps_per_batch <= 0:
raise ValueError(
"optim_steps_per_batch must be positive with "
"learner_backend='ray'."
)
if learner_poll_interval <= 0:
raise ValueError("learner_poll_interval must be positive.")
self.learner_poll_interval = float(learner_poll_interval)
self._execution_backend = None
self._published_model_version = -1
self.logger = logger
self.async_collection = async_collection
# Logging frequency control - how often to log each metric (in frames)
self._log_interval = log_interval
# seeding
self.seed = seed
if seed is not None:
self.set_seed()
# constants
self.optim_steps_per_batch = optim_steps_per_batch
self.total_frames = total_frames
self.num_epochs = num_epochs
self.clip_grad_norm = clip_grad_norm
self.clip_norm = clip_norm
if progress_bar and not _has_tqdm:
warnings.warn(
"tqdm library not found. "
"Consider installing tqdm to use the Trainer progress bar."
)
self.progress_bar = progress_bar and _has_tqdm
self.save_trainer_interval = save_trainer_interval
self.save_trainer_file = save_trainer_file
self.checkpoint = checkpoint
if checkpoint_rotation is not None and checkpoint is None:
raise ValueError("checkpoint_rotation requires checkpoint.")
if checkpoint_rotation is not None and save_trainer_file is not None:
raise ValueError(
"checkpoint_rotation and save_trainer_file cannot both be set."
)
if checkpoint_metadata is not None and checkpoint_rotation is None:
raise ValueError("checkpoint_metadata requires checkpoint_rotation.")
if checkpoint_metadata is not None and not callable(checkpoint_metadata):
raise TypeError("checkpoint_metadata must be callable.")
self.checkpoint_rotation = checkpoint_rotation
self.checkpoint_metadata = checkpoint_metadata
self._checkpoint_state = _TrainerCheckpointState(self)
self._execution_checkpoint_state = _ExecutionCheckpointState(self)
self._checkpoint_skip_warnings: set[str] = set()
self.auto_log_optim_steps = auto_log_optim_steps
self._log_dict = defaultdict(list)
self._stop_training = False
self._stop_reason = None
# Hook collections for different stages of the training loop
self._batch_process_ops = (
[]
) # Process collected batches (e.g., reward normalization)
self._post_steps_ops = [] # After optimization steps (e.g., weight updates)
# Logging hook collections - different points in training loop where logging can occur
self._post_steps_log_ops = (
[]
) # After optimization steps (e.g., validation rewards)
self._pre_steps_log_ops = (
[]
) # Before optimization steps (e.g., rewards, frame counts)
self._post_optim_log_ops = (
[]
) # After each optimization step (e.g., gradient norms)
self._pre_epoch_log_ops = (
[]
) # Before each epoch logging (e.g., epoch-specific metrics)
self._post_epoch_log_ops = (
[]
) # After each epoch logging (e.g., epoch completion metrics)
self._post_optim_complete_log_ops = (
[]
) # After all optimization steps for a batch (e.g., logging average_losses)
# Regular hook collections for non-logging operations
self._pre_epoch_ops = (
[]
) # Before each epoch (e.g., epoch setup, cache clearing)
self._post_epoch_ops = (
[]
) # After each epoch (e.g., epoch cleanup, weight syncing)
# Optimization-related hook collections
self._pre_optim_ops = [] # Before optimization steps (e.g., cache clearing)
self._post_loss_ops = (
[]
) # After loss computation, operates on batch (e.g., priority updates)
self._process_loss_ops = (
[]
) # Transform loss values before optimizer (e.g., scaling, clipping)
self._optimizer_ops = [] # During optimization (e.g., gradient clipping)
self._process_optim_batch_ops = (
[]
) # Process batches for optimization (e.g., subsampling)
self._post_optim_ops = [] # After optimization (e.g., weight syncing)
self._setup_ops = [] # Before training starts (e.g., warmups, lazy init)
self._shutdown_ops = [] # At training end (e.g., final eval, publish)
self._modules = {}
self.optimization_stepper = optimization_stepper
if self.optimization_stepper is not None and self.learner_backend == "local":
self.optimization_stepper.register(self, name="optimization_stepper")
if (
self.optimizer is not None
and self.optimization_stepper is None
and self.learner_backend == "local"
):
# Only auto-create the OptimizerHook when no stepper is
# provided. When a stepper is present it may access
# trainer.optimizer directly, so creating the hook would leave a
# dead hook in _optimizer_ops that never fires.
with warnings.catch_warnings():
warnings.simplefilter("ignore", FutureWarning)
optimizer_hook = OptimizerHook(self.optimizer)
optimizer_hook.register(self)
if log_timings:
log_timing_hook = LogTiming(prefix="time", percall=True, erase=False)
log_timing_hook.register(self)
if self.learner_backend == "ray":
if replay_buffer is None:
raise ValueError(
"replay_buffer is required with learner_backend='ray'."
)
if not isinstance(loss_module, LossModule):
raise TypeError(
"learner_backend='ray' requires an ordinary LossModule object."
)
global_batch_size = batch_size
if global_batch_size is None:
global_batch_size = getattr(replay_buffer, "batch_size", None)
if global_batch_size is None:
raise ValueError(
"batch_size is required when replay_buffer.batch_size is unset."
)
# Kept lazy to break the intentional Trainer/stepper import cycle.
from torchrl.trainers._ray_execution import _RayTrainerExecution
prepare_weight_sync = getattr(collector, "_learner_weight_sync", None)
if prepare_weight_sync is None:
raise TypeError(
"learner_backend='ray' requires a collector that exposes "
"TorchRL WeightSyncScheme receivers."
)
self._execution_backend = _RayTrainerExecution(
loss_module=loss_module,
optimizer=optimizer,
optimization_stepper=optimization_stepper,
target_net_updater=target_net_updater,
replay_buffer=replay_buffer,
global_batch_size=global_batch_size,
options=self.learner_backend_options,
seed=seed,
clip_grad_norm=clip_grad_norm,
clip_norm=clip_norm,
update_replay_priority=not getattr(
optimization_stepper, "updates_replay_priority", False
),
weight_sync_factory=prepare_weight_sync,
)
if self.checkpoint is not None:
self._sync_checkpoint_components()
[docs]
def compute_loss(
self, sub_batch: TensorDictBase, method: str | None = None
) -> TensorDictBase | tuple[Any, ...]:
"""Evaluate the configured loss through the active execution boundary."""
if method is None:
return self.loss_module(sub_batch)
return getattr(self.loss_module, method)(sub_batch)
def register_module(self, module_name: str, module: Any) -> None:
if module_name in self._modules:
raise RuntimeError(
f"{module_name} is already registered, choose a different name."
)
self._modules[module_name] = module
if self.checkpoint is not None:
self._sync_checkpoint_components()
@staticmethod
def _is_checkpointable(component: Any) -> bool:
return (
callable(getattr(component, "dump", None))
and callable(getattr(component, "load", None))
) or (
callable(getattr(component, "state_dict", None))
and callable(getattr(component, "load_state_dict", None))
)
def _checkpoint_policy(self) -> Any | None:
for name in ("actor_network", "value_network", "local_value_network"):
policy = getattr(self.loss_module, name, None)
if policy is not None:
return policy
return None
def _checkpoint_optimizer(self) -> Any | None:
"""Return the optimizer owned by the Trainer or a legacy hook."""
optimizer = self.optimizer
if self._is_checkpointable(optimizer):
return optimizer
optimizer_hook = self._modules.get("optimizer")
hook_optimizer = getattr(optimizer_hook, "optimizer", optimizer_hook)
if self._is_checkpointable(hook_optimizer):
return hook_optimizer
return None
def _sync_checkpoint_components(
self, checkpoint: Checkpoint | None = None
) -> Checkpoint:
if checkpoint is None:
checkpoint = self.checkpoint
if checkpoint is None:
checkpoint = Checkpoint()
def register(name: str, component: Any) -> None:
if name in checkpoint or component is None:
return
if self._is_checkpointable(component):
checkpoint.register(name, component)
elif name not in self._checkpoint_skip_warnings:
torchrl_logger.warning(
"Skipping non-checkpointable Trainer component %r (%s).",
name,
type(component).__name__,
)
self._checkpoint_skip_warnings.add(name)
if self.learner_backend == "ray":
# These adapters refer to already-live services, so Checkpoint's
# deterministic name-based load order is safe. Driver copies of
# loss/optimizer state are not authoritative in this mode and are
# intentionally omitted. Policy weights are republished after the
# complete checkpoint has loaded.
register("replay_buffer", self.replay_buffer)
register("collector", self.collector)
register("trainer_state", self._checkpoint_state)
register("logger", self.logger)
register("exploration", getattr(self, "greedy_module", None))
register("exploration", getattr(self, "exploration_module", None))
for name, module in self._modules.items():
register(f"trainer_module.{name}", module)
register("rng", GlobalRNGState())
register("learner_execution", self._execution_checkpoint_state)
return checkpoint
register("policy", self._checkpoint_policy())
register("loss_module", self.loss_module)
register("optimizer", self._checkpoint_optimizer())
register("collector", self.collector)
replay_buffer = getattr(self, "replay_buffer", None)
if not self._is_checkpointable(replay_buffer):
replay_buffer_hook = self._modules.get("replay_buffer")
hook_replay_buffer = getattr(
replay_buffer_hook, "replay_buffer", replay_buffer_hook
)
replay_buffer = (
hook_replay_buffer
if self._is_checkpointable(hook_replay_buffer)
else replay_buffer_hook
)
register("replay_buffer", replay_buffer)
register("logger", self.logger)
register("exploration", getattr(self, "greedy_module", None))
register("exploration", getattr(self, "exploration_module", None))
register("target_updater", getattr(self, "target_net_updater", None))
register("trainer_state", self._checkpoint_state)
for name, module in self._modules.items():
if name in ("optimizer", "replay_buffer") and name in checkpoint:
continue
register(f"trainer_module.{name}", module)
register("rng", GlobalRNGState())
return checkpoint
def _wrap_hook_with_timing(
self, op: Callable, hook_name: str | None = None
) -> Callable:
"""Wrap a hook with timing measurement.
Args:
op: The hook/operation to wrap
hook_name: Optional name for the hook. If not provided, will be inferred from op.
Returns:
A wrapped version of the hook that measures execution time.
"""
if hook_name is None:
hook_name = getattr(
op,
"__name__",
op.__class__.__name__ if hasattr(op, "__class__") else "unknown_hook",
)
def timed_hook(*args, **kwargs):
with timeit(f"hook/{hook_name}"):
return op(*args, **kwargs)
# Preserve original attributes for debugging
timed_hook.__wrapped__ = op
timed_hook.__name__ = hook_name
return timed_hook
def _get_state(self):
if _CKPT_BACKEND == "torchsnapshot":
state = StateDict(
collected_frames=self.collected_frames,
_last_log=self._last_log,
_last_save=self._last_save,
_optim_count=self._optim_count,
)
else:
state = OrderedDict(
collected_frames=self.collected_frames,
_last_log=self._last_log,
_last_save=self._last_save,
_optim_count=self._optim_count,
)
return state
@property
def app_state(self):
optimizer = self._checkpoint_optimizer()
self._app_state = {
"state": StateDict(**self._get_state()),
"collector": self.collector,
"loss_module": self.loss_module,
**({"optimizer": optimizer} if optimizer is not None else {}),
**{k: item for k, item in self._modules.items() if k != "optimizer"},
}
return self._app_state
def state_dict(self) -> dict:
state = self._get_state()
state_dict: OrderedDict[str, Any] = OrderedDict(
collector=self.collector.state_dict(),
loss_module=self.loss_module.state_dict(),
state=state,
)
optimizer = self._checkpoint_optimizer()
if optimizer is not None:
state_dict["optimizer"] = optimizer.state_dict()
for key, item in self._modules.items():
# The standard optimizer component is emitted above, while a
# legacy OptimizerHook may also be registered under this name.
if key == "optimizer":
continue
state_dict[key] = item.state_dict()
return state_dict
def load_state_dict(self, state_dict: dict) -> None:
model_state_dict = state_dict["loss_module"]
collector_state_dict = state_dict["collector"]
self.loss_module.load_state_dict(model_state_dict)
self.collector.load_state_dict(collector_state_dict)
optimizer = self._checkpoint_optimizer()
optimizer_state_dict = state_dict.get("optimizer")
if optimizer is not None and optimizer_state_dict:
optimizer.load_state_dict(optimizer_state_dict)
for key, item in self._modules.items():
if key == "optimizer":
continue
item.load_state_dict(state_dict[key])
self.collected_frames = state_dict["state"]["collected_frames"]
self._last_log = state_dict["state"]["_last_log"]
self._last_save = state_dict["state"]["_last_save"]
self._optim_count = state_dict["state"]["_optim_count"]
[docs]
def request_stop(self, reason: str | None = None) -> None:
"""Signal that training should stop at the next loop boundary."""
self._stop_training = True
self._stop_reason = reason
[docs]
@contextlib.contextmanager
def stop_on_signal(
self, signals: Collection[int] = (signal.SIGINT, signal.SIGTERM)
):
"""Stop training cleanly when the process receives a termination signal.
Wrap :meth:`train` in this context. The first signal calls
:meth:`request_stop`, so the loop finishes the current batch, writes a
final checkpoint when a save destination is configured, shuts the
collector down and returns. A second signal raises
:class:`KeyboardInterrupt`. Previous handlers are restored on exit.
Args:
signals (Collection[int], optional): signal numbers to handle.
Defaults to ``SIGINT`` and ``SIGTERM``.
Examples:
>>> with trainer.stop_on_signal(): # doctest: +SKIP
... trainer.train()
"""
with StopOnSignal(
signals,
on_request=lambda name: self.request_stop(f"received {name}"),
) as stop:
yield stop
def _save_trainer(self) -> None:
if self.checkpoint is not None:
checkpoint = self._sync_checkpoint_components()
if self.checkpoint_rotation is not None:
self.checkpoint_rotation.save(
checkpoint,
step=self.collected_frames,
metadata=self._checkpoint_manifest_metadata(),
)
else:
checkpoint.save(self.save_trainer_file)
return
warnings.warn(
"The default Trainer checkpoint format will change from the legacy "
"CKPT_BACKEND format to torchrl.checkpoint in v0.15. Pass "
"checkpoint=Checkpoint(...) to opt in now.",
FutureWarning,
stacklevel=2,
)
warnings.warn(
"The legacy CKPT_BACKEND trainer checkpoint path is deprecated and "
"will be removed in v0.16. Use checkpoint=Checkpoint(...).",
DeprecationWarning,
stacklevel=2,
)
if _CKPT_BACKEND == "torchsnapshot":
if not _has_ts:
raise ImportError(
"torchsnapshot not found. Consider installing torchsnapshot or "
"using the torch checkpointing backend (`CKPT_BACKEND=torch`)"
)
Snapshot.take(app_state=self.app_state, path=self.save_trainer_file)
elif _CKPT_BACKEND == "torch":
# Write to a temporary file and atomically swap it in: an
# interrupted save cannot destroy the previous checkpoint, and
# tensors still memory-mapped from a previous
# ``load_from_file(mmap=True)`` keep reading from the old inode
# instead of from a truncated file.
file = pathlib.Path(self.save_trainer_file)
tmp_file = file.with_name(file.name + ".tmp")
try:
torch.save(self.state_dict(), tmp_file)
tmp_file.replace(file)
except BaseException:
tmp_file.unlink(missing_ok=True)
raise
elif _CKPT_BACKEND == "memmap":
state = self.state_dict()
path = pathlib.Path(self.save_trainer_file)
path.mkdir(parents=True, exist_ok=True)
# Persist all module state dicts using TensorDict memmap.
# Non-tensor values (scalars, bools, nested dicts) are wrapped in
# NonTensorData automatically by _state_dict_to_td, so no pickle
# dependency is needed.
for key in dict.fromkeys(
("loss_module", "collector", "optimizer", *self._modules)
):
if key not in state:
continue
sd = state[key]
if sd:
_state_dict_to_td(sd).dumps(str(path / key))
# Persist non-tensor training counters as JSON.
with open(path / "state.json", "w") as f:
json.dump(dict(state["state"]), f)
else:
raise NotImplementedError(
f"CKPT_BACKEND should be one of {_CKPT_BACKEND.backends}, got {_CKPT_BACKEND}."
)
def _has_checkpoint_destination(self) -> bool:
return (
self.checkpoint_rotation is not None or self.save_trainer_file is not None
)
def _checkpoint_manifest_metadata(self) -> dict[str, Any]:
metadata = {}
if self.checkpoint_metadata is not None:
custom_metadata = self.checkpoint_metadata(self)
if not isinstance(custom_metadata, Mapping):
raise TypeError("checkpoint_metadata must return a mapping.")
metadata.update(custom_metadata)
metadata.update(
collected_frames=self.collected_frames,
optim_steps=self._optim_count,
)
return metadata
def _save_interval_elapsed(self) -> bool:
return (self.collected_frames - self._last_save) > self.save_trainer_interval
def _save_due(self, force_save: bool = False) -> bool:
"""Whether a destination is configured and a save is due now."""
return self._has_checkpoint_destination() and (
force_save or self._save_interval_elapsed()
)
def save_trainer(self, force_save: bool = False) -> None:
if not self._has_checkpoint_destination():
return
if self._save_interval_elapsed():
self._last_save = self.collected_frames
elif not force_save:
return
self._save_trainer()
def _save_trainer_at_boundary(self, *, force_save: bool = False) -> None:
"""Save at a training-loop boundary, pausing free-running collection."""
if self._save_due(force_save):
with self._collection_paused():
self.save_trainer(force_save=force_save)
@contextlib.contextmanager
def _collection_paused(self):
"""Pause an asynchronous collector while a checkpoint is written."""
pause = (
getattr(self.collector, "pause", None) if self.async_collection else None
)
with contextlib.ExitStack() as stack:
if pause is not None:
try:
stack.enter_context(pause())
except NotImplementedError:
if "collector.pause" not in self._checkpoint_skip_warnings:
torchrl_logger.warning(
"%s does not implement pause(); asynchronous checkpoints "
"are written while collection continues.",
type(self.collector).__name__,
)
self._checkpoint_skip_warnings.add("collector.pause")
yield
def _resolve_checkpoint_path(self, file: str | pathlib.Path) -> str | pathlib.Path:
"""Map a rotation directory to its newest checkpoint; pass other inputs through."""
if not isinstance(file, (str, pathlib.PurePath)):
return file
path = pathlib.Path(file).expanduser()
if not path.is_dir() or (path / "state.json").exists():
# A file, an archive, or a legacy memmap trainer directory.
return file
return resolve_checkpoint_path(path)
[docs]
def load_from_file(self, file: str | pathlib.Path, **kwargs) -> Trainer:
"""Loads a file and its state-dict in the trainer.
Keyword arguments are passed to the :func:`~torch.load` function for
legacy torch checkpoints and unified components explicitly saved with
the torch state-dict payload format. Unified checkpoints additionally
accept ``strict`` to control missing or incompatible components.
Arguments are ignored when ``CKPT_BACKEND=memmap``.
.. note::
Unified state-dict components use TensorDict storage by default and
do not invoke the pickle loader. For explicit torch payloads and
``CKPT_BACKEND=torch`` checkpoints, ``weights_only=True`` is the
default for safer deserialization. Pass ``weights_only=False``
explicitly only if the state dict contains custom objects. On
torch < 2.4 the default is ``weights_only=False`` because the
weights-only unpickler of those versions cannot deserialize the
``torch.device`` instances contained in TensorDict state-dicts.
.. note::
Explicit torch payloads and ``CKPT_BACKEND=torch`` checkpoints use
``mmap=True`` by default. Pass ``mmap=False`` for legacy pre-zipfile
``torch.save`` files or file-like objects. On Windows the default
is ``mmap=False`` because a mapped checkpoint keeps the file locked,
preventing deletion or re-save.
.. note::
Unified checkpoint tensors are mapped to CPU by default. Pass an
explicit ``map_location`` to select another device mapping.
.. note::
After restoring an independently registered policy component, the
trainer synchronizes the collector once so local policy copies and
remote workers observe the restored learner weights.
.. note::
``file`` may also be a :class:`~torchrl.checkpoint.CheckpointRotation`
directory, in which case its newest checkpoint is restored.
"""
file = self._resolve_checkpoint_path(file)
if Checkpoint.is_checkpoint(file):
checkpoint = self.checkpoint
if checkpoint is None:
checkpoint = Checkpoint()
map_location = kwargs.pop("map_location", "cpu")
strict = kwargs.pop("strict", None)
for key, value in _torch_load_defaults().items():
kwargs.setdefault(key, value)
checkpoint = self._sync_checkpoint_components(checkpoint)
load_kwargs = {
"map_location": map_location,
"tensor_load_kwargs": kwargs,
"strict": strict,
}
registered = set(checkpoint.components)
load_rng = (
"rng" in registered and "rng" in Checkpoint.manifest(file)["components"]
)
registered.discard("rng")
if self.learner_backend == "ray":
# Service owners must be restored before learner actors create
# rank-aware clients. The learner state is deliberately last.
ordered = [
name
for name in ("replay_buffer", "collector")
if name in registered
]
ordered.extend(
sorted(
registered.difference(
{"replay_buffer", "collector", "learner_execution"}
)
)
)
if "learner_execution" in registered:
ordered.append("learner_execution")
loaded = set()
for name in ordered:
result = checkpoint.load(file, components=[name], **load_kwargs)
loaded.update(result.loaded)
else:
result = checkpoint.load(file, components=registered, **load_kwargs)
loaded = set(result.loaded)
if "learner_execution" in loaded:
self._publish_execution_weights(force=True)
elif "policy" in loaded:
# The collector payload is restored before the independently
# registered policy component. Synchronize once more after all
# components have loaded so local copies, remote workers, and
# weight-transport caches observe the restored learner policy.
policy = checkpoint.components["policy"]
if policy is self._checkpoint_policy():
self.collector.update_policy_weights_()
else:
self.collector.update_policy_weights_(policy)
if load_rng:
result = checkpoint.load(file, components=["rng"], **load_kwargs)
loaded.update(result.loaded)
elif _CKPT_BACKEND == "torchsnapshot":
snapshot = Snapshot(path=file)
snapshot.restore(app_state=self.app_state)
elif _CKPT_BACKEND == "torch":
for key, value in _torch_load_defaults().items():
kwargs.setdefault(key, value)
if isinstance(file, pathlib.Path):
# Older torch versions require a string path when mmap is set.
file = str(file)
loaded_dict: OrderedDict = torch.load(file, **kwargs)
self.load_state_dict(loaded_dict)
elif _CKPT_BACKEND == "memmap":
path = pathlib.Path(file)
state: dict = {}
for key in dict.fromkeys(
("loss_module", "collector", "optimizer", *self._modules)
):
key_path = path / key
if key_path.exists():
state[key] = _td_to_state_dict(
TensorDict.load_memmap(str(key_path))
)
else:
state[key] = {}
with open(path / "state.json") as f:
state["state"] = json.load(f)
self.load_state_dict(state)
return self
def set_seed(self):
seed = self.collector.set_seed(self.seed, static_seed=False)
torch.manual_seed(seed)
np.random.seed(seed)
@property
def collector(self) -> BaseCollector:
return self._collector
@collector.setter
def collector(self, collector: BaseCollector) -> None:
self._collector = collector
def register_op(
self,
dest: Literal[
"batch_process",
"pre_optim_steps",
"process_optim_batch",
"post_loss",
"process_loss",
"optimizer",
"post_steps",
"post_optim",
"pre_steps_log",
"post_steps_log",
"post_optim_log",
"pre_epoch_log",
"post_epoch_log",
"post_optim_complete_log",
"pre_epoch",
"post_epoch",
"setup",
"shutdown",
],
op: Callable,
**kwargs,
) -> None:
if self.learner_backend == "ray" and dest not in {
"batch_process",
"post_steps",
"pre_steps_log",
"post_steps_log",
"setup",
"shutdown",
}:
raise RuntimeError(
f"The {dest!r} hook stage executes inside the local optimization "
"loop and is unavailable with learner_backend='ray'. Move this "
"behavior into an OptimizationStepper or a learner-owned component."
)
# Wrap hook with timing for performance monitoring
# Get hook name from registered modules if available
hook_name = None
for name, module in self._modules.items():
if module is op or (callable(module) and module.__call__ is op):
hook_name = name
break
timed_op = self._wrap_hook_with_timing(op, hook_name)
if dest == "batch_process":
_check_input_output_typehint(
op, input=TensorDictBase, output=TensorDictBase
)
self._batch_process_ops.append((timed_op, kwargs))
elif dest == "pre_optim_steps":
_check_input_output_typehint(op, input=None, output=None)
self._pre_optim_ops.append((timed_op, kwargs))
elif dest == "process_optim_batch":
_check_input_output_typehint(
op, input=TensorDictBase, output=TensorDictBase
)
self._process_optim_batch_ops.append((timed_op, kwargs))
elif dest == "post_loss":
warnings.warn(
"The 'post_loss' hook stage will be replaced by OptimizationStepper "
"in a future release. Use 'post_optim' for post-optimization hooks "
"(e.g. priority updates), or provide a custom OptimizationStepper.",
FutureWarning,
stacklevel=2,
)
_check_input_output_typehint(
op, input=TensorDictBase, output=TensorDictBase
)
self._post_loss_ops.append((timed_op, kwargs))
elif dest == "process_loss":
warnings.warn(
"The 'process_loss' hook stage will be replaced by OptimizationStepper "
"in a future release. Move loss-transformation logic into a custom "
"OptimizationStepper.",
FutureWarning,
stacklevel=2,
)
_check_input_output_typehint(
op, input=TensorDictBase, output=TensorDictBase
)
self._process_loss_ops.append((timed_op, kwargs))
elif dest == "optimizer":
warnings.warn(
"The 'optimizer' hook stage will be replaced by OptimizationStepper "
"in a future release. Use DefaultOptimizationStepper for equivalent "
"behaviour.",
FutureWarning,
stacklevel=2,
)
_check_input_output_typehint(
op, input=[TensorDictBase, bool, float, int], output=TensorDictBase
)
self._optimizer_ops.append((timed_op, kwargs))
elif dest == "post_steps":
_check_input_output_typehint(op, input=None, output=None)
self._post_steps_ops.append((timed_op, kwargs))
elif dest == "post_optim":
_check_input_output_typehint(op, input=None, output=None)
self._post_optim_ops.append((timed_op, kwargs))
elif dest == "pre_steps_log":
_check_input_output_typehint(
op, input=TensorDictBase, output=tuple[str, float]
)
self._pre_steps_log_ops.append((timed_op, kwargs))
elif dest == "post_steps_log":
_check_input_output_typehint(
op, input=TensorDictBase, output=tuple[str, float]
)
self._post_steps_log_ops.append((timed_op, kwargs))
elif dest == "post_optim_log":
_check_input_output_typehint(
op, input=TensorDictBase, output=tuple[str, float]
)
self._post_optim_log_ops.append((timed_op, kwargs))
elif dest == "pre_epoch_log":
_check_input_output_typehint(
op, input=TensorDictBase, output=tuple[str, float]
)
self._pre_epoch_log_ops.append((timed_op, kwargs))
elif dest == "post_epoch_log":
_check_input_output_typehint(
op, input=TensorDictBase, output=tuple[str, float]
)
self._post_epoch_log_ops.append((timed_op, kwargs))
elif dest == "post_optim_complete_log":
_check_input_output_typehint(
op, input=[int, TensorDictBase | None], output=tuple[str, float]
)
self._post_optim_complete_log_ops.append((timed_op, kwargs))
elif dest == "pre_epoch":
_check_input_output_typehint(op, input=None, output=None)
self._pre_epoch_ops.append((timed_op, kwargs))
elif dest == "post_epoch":
_check_input_output_typehint(op, input=None, output=None)
self._post_epoch_ops.append((timed_op, kwargs))
elif dest == "setup":
_check_input_output_typehint(op, input=None, output=None)
self._setup_ops.append((timed_op, kwargs))
elif dest == "shutdown":
_check_input_output_typehint(op, input=None, output=None)
self._shutdown_ops.append((timed_op, kwargs))
else:
raise RuntimeError(
f"The hook collection {dest} is not recognised. Choose from:"
f"(batch_process, pre_optim_steps, process_optim_batch, post_loss, "
f"process_loss, optimizer, post_steps, post_optim, pre_steps_log, "
f"post_steps_log, post_optim_log, pre_epoch_log, post_epoch_log, "
f"post_optim_complete_log, setup, shutdown, pre_epoch, post_epoch)"
)
register_hook = register_op
# Process batch
def _process_batch_hook(self, batch: TensorDictBase) -> TensorDictBase:
for op, kwargs in self._batch_process_ops:
out = op(batch, **kwargs)
if isinstance(out, TensorDictBase):
batch = out
return batch
def _post_steps_hook(self) -> None:
for op, kwargs in self._post_steps_ops:
op(**kwargs)
def _post_optim_log(self, batch: TensorDictBase) -> None:
"""Execute logging hooks that run AFTER EACH optimization step.
These hooks log metrics that are computed after each individual optimization step,
such as gradient norms, individual loss components, or step-specific metrics.
Called after each optimization step within the optimization loop.
"""
for op, kwargs in self._post_optim_log_ops:
result = op(batch, **kwargs)
if result is not None:
self._log(**result)
def _pre_optim_hook(self):
for op, kwargs in self._pre_optim_ops:
op(**kwargs)
def _process_optim_batch_hook(self, batch):
for op, kwargs in self._process_optim_batch_ops:
out = op(batch, **kwargs)
if isinstance(out, TensorDictBase):
batch = out
return batch
def _post_loss_hook(self, batch):
for op, kwargs in self._post_loss_ops:
out = op(batch, **kwargs)
if isinstance(out, TensorDictBase):
batch = out
return batch
def _process_loss_hook(
self, sub_batch: TensorDictBase, losses_td: TensorDictBase
) -> TensorDictBase:
"""Apply registered loss transformation hooks before optimization.
Unlike ``post_loss`` hooks which operate on the batch (e.g., for priority updates),
``process_loss`` hooks transform the loss TensorDict itself. These hooks receive
both the sub_batch and the losses, and should return the modified losses.
Use cases include loss scaling, clipping, or applying importance weights.
Args:
sub_batch: The batch of data used to compute the losses.
losses_td: The TensorDict containing loss components from the loss module.
Returns:
The (possibly modified) losses TensorDict.
"""
for op, kwargs in self._process_loss_ops:
out = op(sub_batch, losses_td, **kwargs)
if isinstance(out, TensorDictBase):
losses_td = out
return losses_td
def _optimizer_hook(self, batch):
for i, (op, kwargs) in enumerate(self._optimizer_ops):
out = op(batch, self.clip_grad_norm, self.clip_norm, i, **kwargs)
if isinstance(out, TensorDictBase):
batch = out
return batch.detach()
def _post_optim_hook(self):
for op, kwargs in self._post_optim_ops:
op(**kwargs)
def _pre_epoch_log_hook(self, batch: TensorDictBase) -> None:
"""Execute logging hooks that run BEFORE each epoch of optimization.
These hooks log metrics that should be computed before starting a new epoch
of optimization steps. Called once per epoch within the optimization loop.
"""
for op, kwargs in self._pre_epoch_log_ops:
result = op(batch, **kwargs)
if result is not None:
self._log(**result)
def _pre_epoch_hook(self, batch: TensorDictBase, **kwargs) -> None:
"""Execute regular hooks that run BEFORE each epoch of optimization.
These hooks perform non-logging operations before starting a new epoch
of optimization steps. Called once per epoch within the optimization loop.
"""
for op, kwargs in self._pre_epoch_ops:
batch = op(batch, **kwargs)
return batch
def _post_epoch_log_hook(self, batch: TensorDictBase) -> None:
"""Execute logging hooks that run AFTER each epoch of optimization.
These hooks log metrics that should be computed after completing an epoch
of optimization steps. Called once per epoch within the optimization loop.
"""
for op, kwargs in self._post_epoch_log_ops:
result = op(batch, **kwargs)
if result is not None:
self._log(**result)
def _post_optim_complete_log_hook(
self, optim_steps: int, average_losses: TensorDictBase | None
) -> None:
"""Execute logging hooks that run AFTER all steps in the optimization loop.
These hooks log metrics that use the total step count and averaged loss TensorDict.
Called once per optimization loop.
"""
for op, kwargs in self._post_optim_complete_log_ops:
result = op(optim_steps, average_losses, **kwargs)
if result is not None:
self._log(**result)
if self.auto_log_optim_steps:
if average_losses is not None:
self._log(optim_steps=optim_steps, **average_losses)
else:
self._log(optim_steps=optim_steps)
def _post_epoch_hook(self) -> None:
"""Execute regular hooks that run AFTER each epoch of optimization.
These hooks perform non-logging operations after completing an epoch
of optimization steps. Called once per epoch within the optimization loop.
"""
for op, kwargs in self._post_epoch_ops:
op(**kwargs)
def _pre_steps_log_hook(self, batch: TensorDictBase) -> None:
"""Execute logging hooks that run BEFORE optimization steps.
These hooks typically log metrics from the collected batch data,
such as rewards, frame counts, or other batch-level statistics.
Called once per batch collection, before any optimization occurs.
"""
for op, kwargs in self._pre_steps_log_ops:
result = op(batch, **kwargs)
if result is not None:
self._log(**result)
def _post_steps_log_hook(self, batch: TensorDictBase) -> None:
"""Execute logging hooks that run AFTER optimization steps.
These hooks typically log metrics that depend on the optimization results,
such as validation rewards, evaluation metrics, or post-training statistics.
Called once per batch collection, after all optimization steps are complete.
"""
for op, kwargs in self._post_steps_log_ops:
result = op(batch, **kwargs)
if result is not None:
self._log(**result)
def _setup_hook(self) -> None:
for op, kwargs in self._setup_ops:
op(**kwargs)
def _shutdown_hook(self) -> None:
for op, kwargs in self._shutdown_ops:
op(**kwargs)
def train(self):
if self.learner_backend == "ray":
return self._train_with_execution_backend()
if self.progress_bar:
self._pbar = tqdm(total=self.total_frames, initial=self.collected_frames)
self._pbar_str = {}
setup_complete = False
try:
if self.async_collection:
self.collector.start()
while self.collector.getattr_rb("write_count") == 0:
time.sleep(0.1)
# Create async iterator that monitors write_count progress
iterator = self._async_iterator()
else:
iterator = self.collector
self._setup_hook()
setup_complete = True
for batch in iterator:
if not self.async_collection and batch is not None:
batch = self._process_batch_hook(batch)
current_frames = (
batch.get(("collector", "mask"), torch.tensor(batch.numel()))
.sum()
.item()
* self.frame_skip
)
self.collected_frames += current_frames
else:
# Batch is None: either async collection, or a synchronous
# collector that writes directly to the replay buffer (e.g.
# LLM collectors created with a replay_buffer). Frames are
# tracked via the buffer write count in both cases.
batch = None
cf = self.collected_frames
if self.replay_buffer is not None:
self.collected_frames = self._replay_write_count()
else:
self.collected_frames = self.collector.getattr_rb("write_count")
current_frames = self.collected_frames - cf
# LOGGING POINT 1: Pre-optimization logging (e.g., rewards, frame counts)
self._pre_steps_log_hook(batch)
if self.collected_frames >= self.collector.init_random_frames:
self.optim_steps(batch)
self._post_steps_hook()
# LOGGING POINT 2: Post-optimization logging (e.g., validation rewards, evaluation metrics)
self._post_steps_log_hook(batch)
if self._stop_training:
if self._stop_reason and VERBOSE:
torchrl_logger.info(
f"Trainer stopping early: {self._stop_reason}"
)
self._save_trainer_at_boundary(force_save=True)
break
if self.progress_bar:
self._pbar.update(current_frames)
self._pbar_description()
if self.collected_frames >= self.total_frames:
self._save_trainer_at_boundary(force_save=True)
break
self._save_trainer_at_boundary()
finally:
try:
if setup_complete:
self._shutdown_hook()
finally:
self.collector.shutdown()
def _train_with_execution_backend(self) -> None:
"""Run collection while the private backend owns optimization state."""
backend = self._execution_backend
if backend is None:
raise RuntimeError("The configured learner execution backend is missing.")
if self.progress_bar:
self._pbar = tqdm(total=self.total_frames, initial=self.collected_frames)
self._pbar_str = {}
setup_complete = False
try:
if hasattr(self.replay_buffer, "start"):
self.replay_buffer.start()
backend.start()
self._publish_execution_weights(force=True)
self._setup_hook()
setup_complete = True
previous_write_count = self._replay_write_count()
if self.async_collection:
self.collector.start()
iterator = itertools.repeat(None)
else:
iterator = iter(self.collector)
for batch in iterator:
previous_frames = self.collected_frames
if batch is not None:
batch = self._process_batch_hook(batch)
current_frames = int(
batch.get(("collector", "mask"), torch.tensor(batch.numel()))
.sum()
.item()
)
replay_batch = batch
if ("collector", "mask") in replay_batch.keys(True):
replay_batch = replay_batch[
replay_batch.get(("collector", "mask"))
]
else:
replay_batch = replay_batch.reshape(-1)
self.replay_buffer.extend(replay_batch)
self.collected_frames += current_frames * self.frame_skip
self._pre_steps_log_hook(batch)
else:
write_count = self._replay_write_count()
current_frames = max(0, write_count - previous_write_count)
previous_write_count = write_count
self.collected_frames += current_frames * self.frame_skip
if (
self.collected_frames
>= getattr(self.collector, "init_random_frames", 0)
and len(self.replay_buffer) >= backend.global_batch_size
and current_frames > 0
):
num_steps = self.optim_steps_per_batch * self.num_epochs
receipt = backend.step(num_steps)
self._optim_count += receipt.optim_steps
metrics = receipt.metrics.flatten_keys(".").to_dict()
if self.auto_log_optim_steps:
metrics["optim_steps"] = self._optim_count
self._log(**metrics)
self._publish_execution_weights()
self._post_steps_hook()
if batch is not None:
self._post_steps_log_hook(batch)
if self.progress_bar and self.collected_frames > previous_frames:
self._pbar.update(self.collected_frames - previous_frames)
self._pbar_description()
if self._stop_training or self.collected_frames >= self.total_frames:
self._save_execution_checkpoint(
force_save=True, resume_collection=False
)
break
self._save_execution_checkpoint()
if self.async_collection and current_frames == 0:
time.sleep(self.learner_poll_interval)
finally:
if setup_complete:
self._shutdown_hook()
try:
self.collector.shutdown()
finally:
backend.shutdown()
def _save_execution_checkpoint(
self, *, force_save: bool = False, resume_collection: bool = True
) -> None:
if not self._save_due(force_save):
return
if self.async_collection:
pause = getattr(self.collector, "pause", None)
if not callable(pause):
raise RuntimeError(
"Asynchronous remote checkpointing requires a collector "
"with a pause() boundary."
)
with pause(resume=resume_collection):
self.save_trainer(force_save=force_save)
else:
self.save_trainer(force_save=force_save)
def _publish_execution_weights(self, *, force: bool = False) -> None:
backend = self._execution_backend
model_version = backend.model_version
if not force and model_version <= self._published_model_version:
return
model_weights_key, auxiliary_weights = self._execution_weight_publication()
published_version = backend.publish_weights(
expected_version=model_version,
model_weights_key=model_weights_key,
auxiliary_weights=auxiliary_weights,
)
if published_version != model_version:
raise RuntimeError(
f"Published model version {published_version} does not match "
f"learner version {model_version}."
)
self._published_model_version = published_version
def _execution_weight_publication(
self,
) -> tuple[NestedKey | None, TensorDictBase | None]:
return None, None
@staticmethod
def _compose_execution_weight_publication(
auxiliary_module: nn.Module | None,
) -> tuple[NestedKey | None, TensorDictBase | None]:
"""Compose learner policy weights with a controller-owned module."""
if auxiliary_module is None:
return None, None
auxiliary_weights = TensorDict(
{"module": TensorDict({"1": TensorDict.from_module(auxiliary_module)})}
)
return ("module", "0"), auxiliary_weights
def _execution_controller_state(self) -> dict[str, Any]:
return {
"published_model_version": self._published_model_version,
"learner_generation": self._execution_backend.generation,
}
def _load_execution_controller_state(self, state_dict: Mapping[str, Any]) -> None:
# A restored backend always starts a fresh generation and republishes
# its semantic model version before collection resumes.
del state_dict
self._published_model_version = -1
def _replay_write_count(self) -> int:
write_count = self.replay_buffer.write_count
if callable(write_count):
write_count = write_count()
return int(write_count)
def _async_iterator(self):
"""Create an iterator for async collection that monitors replay buffer write_count.
This iterator yields None batches and terminates when total_frames is reached
based on the replay buffer's write_count rather than using a fixed range.
This ensures the training loop properly consumes the entire collector output.
"""
while True:
current_write_count = self.collector.getattr_rb("write_count")
# Check if we've reached the target frames
if current_write_count >= self.total_frames:
break
else:
yield None
def __del__(self):
try:
self.collector.shutdown()
except Exception:
pass
def shutdown(self):
if VERBOSE:
torchrl_logger.info("shutting down collector")
self.collector.shutdown()
[docs]
def optim_steps(
self,
batch: TensorDictBase,
*,
optim_steps_per_batch: int | None | object = _OPTIM_STEPS_UNSET,
num_epochs: int | object = _OPTIM_STEPS_UNSET,
) -> None:
"""Run the configured optimization loop for one collected batch.
Keyword overrides are applied only to this call and do not change the
trainer configuration. They are useful for algorithms that need a
one-time optimization schedule while retaining the standard Trainer
hooks and logging behavior.
"""
average_losses = None
self._pre_optim_hook()
if optim_steps_per_batch is _OPTIM_STEPS_UNSET:
optim_steps_per_batch = self.optim_steps_per_batch
if num_epochs is _OPTIM_STEPS_UNSET:
num_epochs = self.num_epochs
j = -1
for _ in range(num_epochs):
# LOGGING POINT 3: Pre-epoch logging (e.g., epoch-specific metrics)
self._pre_epoch_log_hook(batch)
# Regular pre-epoch operations (e.g., epoch setup)
batch_processed = self._pre_epoch_hook(batch)
if optim_steps_per_batch is None:
prog = itertools.count()
else:
prog = range(optim_steps_per_batch)
for j in prog:
self._optim_count += 1
try:
sub_batch = self._process_optim_batch_hook(batch_processed)
except StopIteration:
break
if sub_batch is None:
break
if self.optimization_stepper is not None:
losses_detached = self.optimization_stepper.step(self, sub_batch)
self._post_optim_hook()
else:
losses_td = self.loss_module(sub_batch)
self._post_loss_hook(sub_batch)
losses_td = self._process_loss_hook(sub_batch, losses_td)
losses_detached = self._optimizer_hook(losses_td)
self._post_optim_hook()
del losses_td
# LOGGING POINT 4: Post-optimization step logging (e.g., gradient norms, step-specific metrics)
self._post_optim_log(sub_batch)
if average_losses is None:
average_losses: TensorDictBase = losses_detached
else:
for key, item in losses_detached.items():
val = average_losses.get(key)
average_losses.set(key, val * j / (j + 1) + item / (j + 1))
del sub_batch, losses_detached
# LOGGING POINT 5: Post-epoch logging (e.g., epoch completion metrics)
self._post_epoch_log_hook(batch)
# Regular post-epoch operations (e.g., epoch cleanup)
self._post_epoch_hook()
if j >= 0:
# LOGGING POINT 6: After all optimization for this batch (e.g., logging average_losses)
self._post_optim_complete_log_hook(self._optim_count, average_losses)
def _log(self, log_pbar=False, **kwargs) -> None:
"""Main logging method that handles both logger output and progress bar updates.
This method is called from various hooks throughout the training loop to log metrics.
It maintains a history of logged values and controls logging frequency based on log_interval.
Args:
log_pbar: If True, the value will also be displayed in the progress bar
**kwargs: Key-value pairs to log, where key is the metric name and value is the metric value
"""
collected_frames = self.collected_frames
for key, item in kwargs.items():
# Store all values in history regardless of logging frequency
self._log_dict[key].append(item)
# Check if enough frames have passed since last logging for this key
if (collected_frames - self._last_log.get(key, 0)) > self._log_interval:
self._last_log[key] = collected_frames
_log = True
else:
_log = False
# Determine logging method (defaults to "log_scalar")
method = LOGGER_METHODS.get(key, "log_scalar")
# Log to external logger (e.g., tensorboard, wandb) if conditions are met
if _log and self.logger is not None:
if (
method == "log_scalar"
and isinstance(item, torch.Tensor)
and item.ndim > 0
):
continue
getattr(self.logger, method)(key, item, step=collected_frames)
# Update progress bar if requested and method is scalar
if method == "log_scalar" and self.progress_bar and log_pbar:
if isinstance(item, torch.Tensor):
item = item.item()
self._pbar_str[key] = item
def _pbar_description(self) -> None:
"""Update the progress bar description with current metric values.
This method formats and displays the current values of metrics that have
been marked for progress bar display (log_pbar=True) in the logging hooks.
"""
if self.progress_bar:
self._pbar.set_description(
", ".join(
[
f"{key}: {self._pbar_str[key]:{TYPE_DESCR.get(type(self._pbar_str[key]), '4.4f')}}"
for key in sorted(self._pbar_str.keys())
]
)
)
def __repr__(self) -> str:
loss_str = indent(f"loss={self.loss_module}", 4 * " ")
collector_str = indent(f"collector={self.collector}", 4 * " ")
optimizer_str = indent(f"optimizer={self.optimizer}", 4 * " ")
logger = indent(f"logger={self.logger}", 4 * " ")
string = "\n".join(
[
loss_str,
collector_str,
optimizer_str,
logger,
]
)
string = f"Trainer(\n{string})"
return string
def _get_list_state_dict(hook_list):
out = []
for item, kwargs in hook_list:
if hasattr(item, "state_dict"):
out.append((item.state_dict(), kwargs))
else:
out.append((None, kwargs))
return out
def _load_list_state_dict(list_state_dict, hook_list):
for i, ((state_dict_item, kwargs), (item, _)) in enumerate(
zip(list_state_dict, hook_list)
):
if state_dict_item is not None:
item.load_state_dict(state_dict_item)
hook_list[i] = (item, kwargs)
[docs]
class SelectKeys(TrainerHookBase):
"""Selects keys in a TensorDict batch.
Args:
keys (iterable of strings): keys to be selected in the tensordict.
Examples:
>>> trainer = make_trainer()
>>> key1 = "first key"
>>> key2 = "second key"
>>> td = TensorDict(
... {
... key1: torch.randn(3),
... key2: torch.randn(3),
... },
... [],
... )
>>> trainer.register_op("batch_process", SelectKeys([key1]))
>>> td_out = trainer._process_batch_hook(td)
>>> assert key1 in td_out.keys()
>>> assert key2 not in td_out.keys()
"""
def __init__(self, keys: Sequence[str]):
if isinstance(keys, str):
raise RuntimeError(
"Expected keys to be an iterable of str, got str instead"
)
self.keys = keys
def __call__(self, batch: TensorDictBase) -> TensorDictBase:
return batch.select(*self.keys)
def state_dict(self) -> dict[str, Any]:
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
pass
[docs]
def register(self, trainer, name="select_keys") -> None:
trainer.register_op("batch_process", self)
trainer.register_module(name, self)
[docs]
class ReplayBufferTrainer(TrainerHookBase):
"""Replay buffer hook provider.
Args:
replay_buffer (TensorDictReplayBuffer): replay buffer to be used.
batch_size (int, optional): batch size when sampling data from the
latest collection or from the replay buffer. If none is provided,
the replay buffer batch-size will be used (preferred option for
unchanged batch-sizes).
memmap (bool, optional): if ``True``, a memmap tensordict is created.
Default is ``False``.
device (device, optional): device where the samples must be placed.
Default to ``None``.
flatten_tensordicts (bool, optional): if ``True``, the tensordicts will be
flattened (or equivalently masked with the valid mask obtained from
the collector) before being passed to the replay buffer. Otherwise,
no transform will be achieved other than padding (see :obj:`max_dims` arg below).
Defaults to ``False``.
max_dims (sequence of int, optional): if :obj:`flatten_tensordicts` is set to False,
this will be a list of the length of the batch_size of the provided
tensordicts that represent the maximum size of each. If provided,
this list of sizes will be used to pad the tensordict and make their shape
match before they are passed to the replay buffer. If there is no
maximum value, a -1 value should be provided.
iterate (bool, optional): if ``True``, the replay buffer will be iterated over
in a loop. Defaults to ``False`` (call to :meth:`~torchrl.data.ReplayBuffer.sample` will be used).
Examples:
>>> rb_trainer = ReplayBufferTrainer(replay_buffer=replay_buffer, batch_size=N)
>>> trainer.register_op("batch_process", rb_trainer.extend)
>>> trainer.register_op("process_optim_batch", rb_trainer.sample)
>>> trainer.register_op("post_loss", rb_trainer.update_priority)
"""
def __init__(
self,
replay_buffer: TensorDictReplayBuffer,
batch_size: int | None = None,
memmap: bool = False,
device: DEVICE_TYPING | None = None,
flatten_tensordicts: bool = False,
max_dims: Sequence[int] | None = None,
iterate: bool = False,
) -> None:
self.replay_buffer = replay_buffer
if hasattr(replay_buffer, "update_tensordict_priority"):
self._update_priority = self.replay_buffer.update_tensordict_priority
else:
if isinstance(replay_buffer.sampler, PrioritizedSampler):
raise ValueError(
"Prioritized sampler not supported for replay buffer trainer if not within a TensorDictReplayBuffer"
)
self._update_priority = None
self.batch_size = batch_size
self.memmap = memmap
self.device = device
self.flatten_tensordicts = flatten_tensordicts
self.max_dims = max_dims
self.iterate = iterate
if iterate:
self.replay_buffer_iter = iter(self.replay_buffer)
def extend(self, batch: TensorDictBase) -> TensorDictBase:
if self.flatten_tensordicts:
if ("collector", "mask") in batch.keys(True):
batch = batch[batch.get(("collector", "mask"))]
else:
if "truncated" in batch["next"]:
batch["next", "truncated"][..., -1] = True
batch = batch.reshape(-1)
else:
if self.max_dims is not None:
pads = []
for d in range(batch.ndimension()):
pad_value = (
0
if self.max_dims[d] == -1
else self.max_dims[d] - batch.batch_size[d]
)
pads += [0, pad_value]
batch = pad(batch, pads)
batch = batch.cpu()
self.replay_buffer.extend(batch)
return batch
def sample(self, batch: TensorDictBase) -> TensorDictBase:
if self.iterate:
try:
sample = next(self.replay_buffer_iter)
except StopIteration:
# reset the replay buffer
self.replay_buffer_iter = iter(self.replay_buffer)
raise
else:
sample = self.replay_buffer.sample(batch_size=self.batch_size)
return sample.to(self.device) if self.device is not None else sample
def update_priority(self, batch: TensorDictBase) -> None:
if self._update_priority is not None:
self._update_priority(batch)
def state_dict(self) -> dict[str, Any]:
return {
"replay_buffer": self.replay_buffer.state_dict(),
}
def load_state_dict(self, state_dict) -> None:
self.replay_buffer.load_state_dict(state_dict["replay_buffer"])
[docs]
def register(self, trainer: Trainer, name: str = "replay_buffer"):
trainer.register_op("batch_process", self.extend)
trainer.register_op("process_optim_batch", self.sample)
trainer.register_op("post_loss", self.update_priority)
trainer.register_module(name, self)
[docs]
class OptimizerHook(TrainerHookBase):
"""Add an optimizer for one or more loss components.
.. deprecated::
``OptimizerHook`` will be replaced by
:class:`~torchrl.trainers.DefaultOptimizationStepper` in a future release.
Args:
optimizer (optim.Optimizer): An optimizer to apply to the loss_components.
loss_components (Sequence[str], optional): The keys in the loss TensorDict
for which the optimizer should be appled to the respective values.
If omitted, the optimizer is applied to all components with the
names starting with `loss_`.
Examples:
>>> optimizer_hook = OptimizerHook(optimizer, ["loss_actor"])
>>> trainer.register_op("optimizer", optimizer_hook)
"""
def __init__(
self,
optimizer: optim.Optimizer,
loss_components: Sequence[str] | None = None,
):
warnings.warn(
"OptimizerHook will be replaced by DefaultOptimizationStepper "
"in a future release.",
FutureWarning,
stacklevel=2,
)
if loss_components is not None and not loss_components:
raise ValueError(
"loss_components list cannot be empty. "
"Set to None to act on all components of the loss."
)
self.optimizer = optimizer
self.loss_components = loss_components
if self.loss_components is not None:
self.loss_components = set(self.loss_components)
def _grad_clip(self, clip_grad_norm: bool, clip_norm: float) -> float:
params = []
for param_group in self.optimizer.param_groups:
params += param_group["params"]
if clip_grad_norm and clip_norm is not None:
gn = nn.utils.clip_grad_norm_(params, clip_norm)
else:
gn = sum([p.grad.pow(2).sum() for p in params if p.grad is not None]).sqrt()
if clip_norm is not None:
nn.utils.clip_grad_value_(params, clip_norm)
return float(gn)
def __call__(
self,
losses_td: TensorDictBase,
clip_grad_norm: bool,
clip_norm: float,
index: int,
) -> TensorDictBase:
loss_components = (
[item for key, item in losses_td.items() if key in self.loss_components]
if self.loss_components is not None
else [item for key, item in losses_td.items() if key.startswith("loss")]
)
loss = sum(loss_components)
loss.backward()
grad_norm = self._grad_clip(clip_grad_norm, clip_norm)
losses_td[f"grad_norm_{index}"] = torch.tensor(grad_norm)
self.optimizer.step()
self.optimizer.zero_grad()
return losses_td
def state_dict(self) -> dict[str, Any]:
state_dict = getattr(self.optimizer, "state_dict", None)
return state_dict() if callable(state_dict) else {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
load_state_dict = getattr(self.optimizer, "load_state_dict", None)
if state_dict and callable(load_state_dict):
load_state_dict(state_dict)
[docs]
def register(self, trainer, name="optimizer") -> None:
trainer.register_op("optimizer", self)
trainer.register_module(name, self)
[docs]
class ClearCudaCache(TrainerHookBase):
"""Clears cuda cache at a given interval.
Examples:
>>> clear_cuda = ClearCudaCache(100)
>>> trainer.register_op("pre_optim_steps", clear_cuda)
"""
def __init__(self, interval: int):
self.interval = interval
self.count = 0
def __call__(self, *args, **kwargs):
self.count += 1
if self.count % self.interval == 0:
torch.cuda.empty_cache()
def state_dict(self) -> dict[str, Any]:
return {"count": self.count}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
self.count = state_dict["count"]
[docs]
def register(self, trainer: Trainer, name: str = "clear_cuda_cache"):
trainer.register_module(name, self)
trainer.register_op("pre_optim_steps", self)
class LogTiming(TrainerHookBase):
"""Hook to log timing information collected by timeit context managers.
This hook extracts timing data from the global timeit registry and logs it
to the trainer's logger (e.g., wandb, tensorboard). It's useful for profiling
different parts of the training loop.
Args:
prefix (str, optional): Prefix to add to timing metric names.
Default is "time".
percall (bool, optional): If True, log average time per call.
If False, log total time. Default is True.
erase (bool, optional): If True, reset timing data after each log.
Default is False.
Examples:
>>> # Log timing data after each optimization step
>>> log_timing = LogTiming(prefix="time", percall=True)
>>> trainer.register_op("post_optim_log", log_timing)
>>> # Log timing data after each batch collection
>>> log_timing = LogTiming(prefix="time", erase=True)
>>> trainer.register_op("post_steps_log", log_timing)
Note:
This hook works with timing data collected using the `timeit` context manager.
For example, hooks registered with `register_op` are automatically wrapped
with timing measurement.
"""
def __init__(
self,
prefix: str = "time",
percall: bool = True,
erase: bool = False,
):
self.prefix = prefix
self.percall = percall
self.erase = erase
def __call__(self, batch: TensorDictBase | None = None) -> dict:
"""Extract timing data and return as a dict for logging.
Args:
batch: The batch (unused, but required by hook signature)
Returns:
Dictionary of timing metrics with the format {metric_name: value}
"""
timing_dict = timeit.todict(percall=self.percall, prefix=self.prefix)
if self.erase:
timeit.erase()
return timing_dict
def state_dict(self) -> dict[str, Any]:
"""Return state dict for checkpointing."""
return {
"prefix": self.prefix,
"percall": self.percall,
"erase": self.erase,
}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
"""Load state dict from checkpoint."""
self.prefix = state_dict.get("prefix", "time")
self.percall = state_dict.get("percall", True)
self.erase = state_dict.get("erase", False)
def register(self, trainer: Trainer, name: str | None = None):
if name is None:
name = "log_timing"
trainer.register_module(name, self)
trainer.register_op("post_steps_log", self)
[docs]
class LogScalar(TrainerHookBase):
"""Generic scalar logger hook for any tensor values in the batch.
This hook can log any scalar values from the collected batch data, including
rewards, action norms, done states, and any other metrics. It automatically
handles masking and computes both mean and standard deviation.
Args:
key (NestedKey): the key where to find the value in the input batch.
Can be a string for simple keys or a tuple for nested keys.
Default is `torchrl.trainers.trainers.REWARD_KEY` (= `("next", "reward")`).
logname (str, optional): name of the metric to be logged. If None, will use
the key as the log name. Default is None.
log_pbar (bool, optional): if ``True``, the value will be logged on
the progression bar. Default is ``False``.
include_std (bool, optional): if ``True``, also log the standard deviation
of the values. Default is ``True``.
reduction (str, optional): reduction method to apply. Can be "mean", "sum",
"min", "max". Default is "mean".
Examples:
>>> # Log training rewards
>>> log_reward = LogScalar(("next", "reward"), "r_training", log_pbar=True)
>>> trainer.register_op("pre_steps_log", log_reward)
>>> # Log action norms
>>> log_action_norm = LogScalar("action", "action_norm", include_std=True)
>>> trainer.register_op("pre_steps_log", log_action_norm)
>>> # Log done states (as percentage)
>>> log_done = LogScalar(("next", "done"), "done_percentage", reduction="mean")
>>> trainer.register_op("pre_steps_log", log_done)
"""
def __init__(
self,
key: NestedKey = REWARD_KEY,
logname: str | None = None,
log_pbar: bool = False,
include_std: bool = True,
reduction: str = "mean",
):
self.key = key
self.logname = logname if logname is not None else str(key)
self.log_pbar = log_pbar
self.include_std = include_std
self.reduction = reduction
# Validate reduction method
if reduction not in ["mean", "sum", "min", "max"]:
raise ValueError(
f"reduction must be one of ['mean', 'sum', 'min', 'max'], got {reduction}"
)
def _apply_reduction(self, tensor: torch.Tensor) -> torch.Tensor:
"""Apply the specified reduction to the tensor."""
if self.reduction == "mean":
return tensor.float().mean()
elif self.reduction == "sum":
return tensor.sum()
elif self.reduction == "min":
return tensor.min()
elif self.reduction == "max":
return tensor.max()
else:
raise ValueError(f"Unknown reduction: {self.reduction}")
def __call__(self, batch: TensorDictBase) -> dict:
# Get the tensor from the batch
tensor = batch.get(self.key)
# Apply mask if available
if ("collector", "mask") in batch.keys(True):
mask = batch.get(("collector", "mask"))
tensor = tensor[mask]
# Compute the main statistic
main_value = self._apply_reduction(tensor).item()
# Prepare the result dictionary
result = {
self.logname: main_value,
"log_pbar": self.log_pbar,
}
# Add standard deviation if requested
if self.include_std and tensor.numel() > 1:
std_value = tensor.float().std().item()
result[f"{self.logname}_std"] = std_value
return result
def state_dict(self) -> dict[str, Any]:
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
pass
[docs]
def register(self, trainer: Trainer, name: str | None = None):
if name is None:
name = f"log_{self.logname}"
trainer.register_op("pre_steps_log", self)
trainer.register_module(name, self)
[docs]
class RewardNormalizer(TrainerHookBase):
"""Reward normalizer hook.
Args:
decay (:obj:`float`, optional): exponential moving average decay parameter.
Default is 0.999
scale (:obj:`float`, optional): the scale used to multiply the reward once
normalized. Defaults to 1.0.
eps (:obj:`float`, optional): the epsilon jitter used to prevent numerical
underflow. Defaults to ``torch.finfo(DEFAULT_DTYPE).eps``
where ``DEFAULT_DTYPE=torch.get_default_dtype()``.
reward_key (str or tuple, optional): the key where to find the reward
in the input batch. Defaults to ``("next", "reward")``
Examples:
>>> reward_normalizer = RewardNormalizer()
>>> trainer.register_op("batch_process", reward_normalizer.update_reward_stats)
>>> trainer.register_op("process_optim_batch", reward_normalizer.normalize_reward)
"""
def __init__(
self,
decay: float = 0.999,
scale: float = 1.0,
eps: float | None = None,
log_pbar: bool = False,
reward_key=None,
):
self._normalize_has_been_called = False
self._update_has_been_called = False
self._reward_stats = OrderedDict()
self._reward_stats["decay"] = decay
self.scale = scale
if eps is None:
eps = torch.finfo(torch.get_default_dtype()).eps
self.eps = eps
if reward_key is None:
reward_key = REWARD_KEY
self.reward_key = reward_key
@torch.no_grad()
def update_reward_stats(self, batch: TensorDictBase) -> None:
reward = batch.get(self.reward_key)
if ("collector", "mask") in batch.keys(True):
reward = reward[batch.get(("collector", "mask"))]
if self._update_has_been_called and not self._normalize_has_been_called:
# We'd like to check that rewards are normalized. Problem is that the trainer can collect data without calling steps...
# raise RuntimeError(
# "There have been two consecutive calls to update_reward_stats without a call to normalize_reward. "
# "Check that normalize_reward has been registered in the trainer."
# )
pass
decay = self._reward_stats.get("decay", 0.999)
sum = self._reward_stats["sum"] = (
decay * self._reward_stats.get("sum", 0.0) + reward.sum()
)
ssq = self._reward_stats["ssq"] = (
decay * self._reward_stats.get("ssq", 0.0) + reward.pow(2).sum()
)
count = self._reward_stats["count"] = (
decay * self._reward_stats.get("count", 0.0) + reward.numel()
)
self._reward_stats["mean"] = sum / count
if count > 1:
var = self._reward_stats["var"] = (ssq - sum.pow(2) / count) / (count - 1)
else:
var = self._reward_stats["var"] = torch.zeros_like(sum)
self._reward_stats["std"] = var.clamp_min(self.eps).sqrt()
self._update_has_been_called = True
def normalize_reward(self, tensordict: TensorDictBase) -> TensorDictBase:
tensordict = tensordict.to_tensordict() # make sure it is not a SubTensorDict
reward = tensordict.get(self.reward_key)
if reward.device is not None:
reward = reward - self._reward_stats["mean"].to(reward.device)
reward = reward / self._reward_stats["std"].to(reward.device)
else:
reward = reward - self._reward_stats["mean"]
reward = reward / self._reward_stats["std"]
tensordict.set(self.reward_key, reward * self.scale)
self._normalize_has_been_called = True
return tensordict
def state_dict(self) -> dict[str, Any]:
return {
"_reward_stats": deepcopy(self._reward_stats),
"scale": self.scale,
"_normalize_has_been_called": self._normalize_has_been_called,
"_update_has_been_called": self._update_has_been_called,
}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
# deepcopy to decouple the normalizer stats from the caller's tensors
# (which may e.g. be mmap-backed by a checkpoint file)
for key, value in state_dict.items():
setattr(self, key, deepcopy(value))
[docs]
def register(self, trainer: Trainer, name: str = "reward_normalizer"):
trainer.register_op("batch_process", self.update_reward_stats)
trainer.register_op("process_optim_batch", self.normalize_reward)
trainer.register_module(name, self)
def mask_batch(batch: TensorDictBase) -> TensorDictBase:
"""Batch masking hook.
If a tensordict contained padded trajectories but only single events are
needed, this hook can be used to select the valid events from the original
tensordict.
Args:
batch:
Examples:
>>> trainer = mocking_trainer()
>>> trainer.register_op("batch_process", mask_batch)
"""
if ("collector", "mask") in batch.keys(True):
mask = batch.get(("collector", "mask"))
return batch[mask]
return batch
[docs]
class BatchSubSampler(TrainerHookBase):
"""Data subsampler for online RL sota-implementations.
This class subsamples a part of a whole batch of data just collected from the
environment.
Args:
batch_size (int): sub-batch size to collect. The provided batch size
must be equal to the total number of items in the output tensordict,
which will have size [batch_size // sub_traj_len, sub_traj_len].
sub_traj_len (int, optional): length of the trajectories that
sub-samples must have in online settings. Default is -1 (i.e.
takes the full length of the trajectory)
min_sub_traj_len (int, optional): minimum value of :obj:`sub_traj_len`, in
case some elements of the batch contain few steps.
Default is -1 (i.e. no minimum value)
Examples:
>>> td = TensorDict(
... {
... key1: torch.stack([torch.arange(0, 10), torch.arange(10, 20)], 0),
... key2: torch.stack([torch.arange(0, 10), torch.arange(10, 20)], 0),
... },
... [2, 10],
... )
>>> trainer.register_op(
... "process_optim_batch",
... BatchSubSampler(batch_size=batch_size, sub_traj_len=sub_traj_len),
... )
>>> td_out = trainer._process_optim_batch_hook(td)
>>> assert td_out.shape == torch.Size([batch_size // sub_traj_len, sub_traj_len])
"""
def __init__(
self, batch_size: int, sub_traj_len: int = 0, min_sub_traj_len: int = 0
) -> None:
self.batch_size = batch_size
self.sub_traj_len = sub_traj_len
self.min_sub_traj_len = min_sub_traj_len
def __call__(self, batch: TensorDictBase) -> TensorDictBase:
"""Sub-sampled part of a batch randomly.
If the batch has one dimension, a random subsample of length
self.bach_size will be returned. If the batch has two or more
dimensions, the last batch dimension represents time. All leading
batch dimensions represent independent trajectories. The resulting
subsample contains consecutive samples across time.
"""
if batch.ndimension() == 1:
return batch[torch.randperm(batch.shape[0])[: self.batch_size]]
batch = batch.reshape(-1, batch.shape[-1])
sub_traj_len = self.sub_traj_len if self.sub_traj_len > 0 else batch.shape[1]
if ("collector", "mask") in batch.keys(True):
# if a valid mask is present, it's important to sample only
# valid steps
traj_len = batch.get(("collector", "mask")).sum(-1)
sub_traj_len = max(
self.min_sub_traj_len,
min(sub_traj_len, traj_len.min().int().item()),
)
else:
traj_len = (
torch.ones(batch.shape[0], device=batch.device, dtype=torch.bool)
* batch.shape[1]
)
len_mask = traj_len >= sub_traj_len
valid_trajectories = torch.arange(batch.shape[0], device=batch.device)[len_mask]
batch_size = self.batch_size // sub_traj_len
if batch_size == 0:
raise RuntimeError(
"Resulting batch size is zero. The batch size given to "
"BatchSubSampler must be equal to the total number of elements "
"that will result in a batch provided to the loss function."
)
traj_idx = valid_trajectories[
torch.randint(
valid_trajectories.numel(), (batch_size,), device=batch.device
)
]
if sub_traj_len < batch.shape[1]:
_traj_len = traj_len[traj_idx]
seq_idx = (
torch.rand_like(_traj_len, dtype=torch.float)
* (_traj_len - sub_traj_len)
).int()
seq_idx = seq_idx.unsqueeze(-1).expand(-1, sub_traj_len)
elif sub_traj_len == batch.shape[1]:
seq_idx = torch.zeros(
batch_size, sub_traj_len, device=batch.device, dtype=torch.long
)
else:
raise ValueError(
f"sub_traj_len={sub_traj_len} is not allowed. Accepted values "
f"are in the range [1, {batch.shape[1]}]."
)
seq_idx = seq_idx + torch.arange(sub_traj_len, device=seq_idx.device)
td = batch[traj_idx].clone()
td = td.apply(
lambda t: t.gather(
dim=1,
index=expand_right(seq_idx, (batch_size, sub_traj_len, *t.shape[2:])),
),
batch_size=(batch_size, sub_traj_len),
)
if ("collector", "mask") in batch.keys(True) and not td.get(
("collector", "mask")
).all():
raise RuntimeError("Sampled invalid steps")
return td
def state_dict(self) -> dict[str, Any]:
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
pass
[docs]
def register(self, trainer: Trainer, name: str = "batch_subsampler"):
trainer.register_op(
"process_optim_batch",
self,
)
trainer.register_module(name, self)
[docs]
class LogValidationReward(TrainerHookBase):
"""Recorder hook for :class:`~torchrl.trainers.Trainer`.
Args:
record_interval (int): total number of optimization steps
between two calls to the recorder for testing.
record_frames (int): number of frames to be recorded during
testing.
frame_skip (int): frame_skip used in the environment. It is
important to let the trainer know the number of frames skipped at
each iteration, otherwise the frame count can be underestimated.
For logging, this parameter is important to normalize the reward.
Finally, to compare different runs with different frame_skip,
one must normalize the frame count and rewards. Defaults to ``1``.
policy_exploration (ProbabilisticTDModule): a policy
instance used for
(1) updating the exploration noise schedule;
(2) testing the policy on the recorder.
Given that this instance is supposed to both explore and render
the performance of the policy, it should be possible to turn off
the explorative behavior by calling the
`set_exploration_type(ExplorationType.DETERMINISTIC)` context manager.
environment (EnvBase): An environment instance to be used
for testing.
exploration_type (ExplorationType, optional): exploration mode to use for the
policy. By default, no exploration is used and the value used is
``ExplorationType.DETERMINISTIC``. Set to ``ExplorationType.RANDOM`` to enable exploration
log_keys (sequence of str or tuples or str, optional): keys to read in the tensordict
for logging. Defaults to ``[("next", "reward")]``.
out_keys (Dict[str, str], optional): a dictionary mapping the ``log_keys``
to their name in the logs. Defaults to ``{("next", "reward"): "r_evaluation"}``.
suffix (str, optional): suffix of the video to be recorded.
log_pbar (bool, optional): if ``True``, the reward value will be logged on
the progression bar. Default is `False`.
"""
ENV_DEPREC = (
"the environment should be passed under the 'environment' key"
" and not the 'recorder' key."
)
def __init__(
self,
*,
record_interval: int,
record_frames: int,
frame_skip: int = 1,
policy_exploration: TensorDictModule,
environment: EnvBase = None,
exploration_type: ExplorationType = ExplorationType.RANDOM,
log_keys: list[str | tuple[str]] | None = None,
out_keys: dict[str | tuple[str], str] | None = None,
suffix: str | None = None,
log_pbar: bool = False,
recorder: EnvBase = None,
) -> None:
if environment is None and recorder is not None:
warnings.warn(self.ENV_DEPREC)
environment = recorder
elif environment is not None and recorder is not None:
raise ValueError("environment and recorder conflict.")
self.policy_exploration = policy_exploration
self.environment = environment
self.record_frames = record_frames
self.frame_skip = frame_skip
self._count = 0
self.record_interval = record_interval
self.exploration_type = exploration_type
if log_keys is None:
log_keys = [("next", "reward")]
if out_keys is None:
out_keys = KeyDependentDefaultDict(lambda x: x)
out_keys[("next", "reward")] = "r_evaluation"
self.log_keys = log_keys
self.out_keys = out_keys
self.suffix = suffix
self.log_pbar = log_pbar
@torch.inference_mode()
def __call__(self, batch: TensorDictBase) -> dict:
out = None
if self._count % self.record_interval == 0:
with set_exploration_type(self.exploration_type):
if isinstance(self.policy_exploration, torch.nn.Module):
self.policy_exploration.eval()
self.environment.eval()
td_record = self.environment.rollout(
policy=self.policy_exploration,
max_steps=self.record_frames,
auto_reset=True,
auto_cast_to_device=True,
break_when_any_done=False,
).clone()
td_record = split_trajectories(td_record)
if isinstance(self.policy_exploration, torch.nn.Module):
self.policy_exploration.train()
self.environment.train()
self.environment.transform.dump(suffix=self.suffix)
out = {}
for key in self.log_keys:
value = td_record.get(key).float()
if key == ("next", "reward"):
mask = td_record["mask"]
mean_value = value[mask].mean() / self.frame_skip
total_value = value.sum(dim=td_record.ndim - 1).mean()
out[self.out_keys[key]] = mean_value
out["total_" + self.out_keys[key]] = total_value
continue
out[self.out_keys[key]] = value
out["log_pbar"] = self.log_pbar
self._count += 1
self.environment.close()
return out
def state_dict(self) -> dict:
return {
"_count": self._count,
"recorder_state_dict": self.environment.state_dict(),
}
def load_state_dict(self, state_dict: dict) -> None:
self._count = state_dict["_count"]
self.environment.load_state_dict(state_dict["recorder_state_dict"])
[docs]
def register(self, trainer: Trainer, name: str = "recorder"):
trainer.register_module(name, self)
trainer.register_op(
"post_steps_log",
self,
)
def _resolve_module(trainer: Trainer, path: str):
"""Resolve a module from a trainer using a string path.
Args:
trainer (Trainer): The trainer instance to resolve from.
path (str): A dot-separated path to the module (e.g., "loss_module.actor_network").
Returns:
The resolved module.
Raises:
AttributeError: If the path cannot be resolved.
Examples:
>>> module = _resolve_module(trainer, "loss_module.actor_network")
>>> module = _resolve_module(trainer, "collector.policy")
"""
obj = trainer
for attr in path.split("."):
obj = getattr(obj, attr)
return obj
[docs]
class EvaluatorHook(TrainerHookBase):
"""Schedule asynchronous evaluation from a :class:`~torchrl.trainers.Trainer`.
The hook snapshots the training policy when an evaluation is triggered, polls
completed results after each collected batch, and logs them under the
``evaluation/`` namespace. If several evaluation intervals elapse while an
evaluation is running, they are coalesced into one evaluation with the latest
policy weights when the evaluator becomes available.
Args:
evaluator (Evaluator): Evaluator service used to run rollouts.
Keyword Args:
every_frames (int): Number of collected frames between evaluations.
policy (str or Callable, optional): Dot-separated path resolved from the
trainer, or a callable receiving the trainer and returning an
:class:`~torch.nn.Module` or :class:`~tensordict.TensorDictBase`.
Defaults to ``"loss_module.actor_network"``.
run_at_start (bool, optional): Whether to evaluate the initial policy.
Defaults to ``False``.
run_at_end (bool, optional): Whether to request a final evaluation with
the latest policy weights. A final evaluation is skipped when the
latest completed evaluation already used the same frame count.
Defaults to ``True``.
wait_at_end (bool, optional): Whether shutdown waits for pending and final
evaluations so their metrics are logged. When ``False``, outstanding
work is handed to :meth:`Evaluator.shutdown` and may be cancelled by
the evaluator backend. Defaults to ``True``.
wait_at_end_timeout (float or None, optional): Maximum seconds to wait for
a pending evaluation during shutdown. ``None`` waits without a time
limit. Defaults to ``60.0``.
Examples:
>>> from torchrl.collectors import Evaluator
>>> from torchrl.trainers import EvaluatorHook
>>> evaluator = Evaluator(make_eval_env, eval_policy, max_steps=1_000) # doctest: +SKIP
>>> EvaluatorHook(evaluator, every_frames=10_000).register(trainer) # doctest: +SKIP
.. note::
Checkpoints contain only the next due frame and the last completed
evaluation frame. In-flight evaluator work is intentionally not
serialized and is discarded when resuming from a checkpoint.
"""
def __init__(
self,
evaluator: Evaluator,
*,
every_frames: int,
policy: str
| Callable[[Trainer], nn.Module | TensorDictBase] = (
"loss_module.actor_network"
),
run_at_start: bool = False,
run_at_end: bool = True,
wait_at_end: bool = True,
wait_at_end_timeout: float | None = 60.0,
):
if isinstance(every_frames, bool) or not isinstance(every_frames, int):
raise TypeError("every_frames must be an integer.")
if every_frames <= 0:
raise ValueError("every_frames must be positive.")
if not isinstance(policy, str) and not callable(policy):
raise TypeError("policy must be a string path or a callable.")
self.evaluator = evaluator
self.every_frames = every_frames
self.policy = policy
self.run_at_start = run_at_start
self.run_at_end = run_at_end
self.wait_at_end = wait_at_end
self.wait_at_end_timeout = wait_at_end_timeout
self._next_due_frame = 0 if run_at_start else every_frames
self._last_completed_frame: int | None = None
self._last_triggered_frame: int | None = None
self._trainer: Trainer | None = None
def _policy_weights(self) -> TensorDictBase:
source = self.policy
if isinstance(source, str):
source = _resolve_module(self._trainer, source)
elif not isinstance(source, nn.Module):
source = source(self._trainer)
if isinstance(source, nn.Module):
return Evaluator.extract_weights(source)
if isinstance(source, TensorDictBase):
return source.detach().clone().cpu()
raise TypeError(
"EvaluatorHook policy must resolve to an nn.Module or TensorDictBase, "
f"got {type(source).__name__}."
)
def _trigger(self, step: int) -> bool:
accepted = self.evaluator.trigger_eval(self._policy_weights(), step=step)
if accepted:
self._last_triggered_frame = step
return accepted
@staticmethod
def _evaluation_name(name: str) -> str:
_, separator, suffix = name.partition("/")
return f"evaluation/{suffix if separator else name}"
def _log_result(self, result: Mapping[str, Any]) -> None:
if not result:
return
normalized = {
self._evaluation_name(name): value for name, value in result.items()
}
step_value = normalized.get("evaluation/step", self._last_triggered_frame)
if isinstance(step_value, torch.Tensor):
step_value = step_value.item()
if step_value is None:
step_value = self._trainer.collected_frames
step = int(step_value)
self._last_completed_frame = step
if step >= self._next_due_frame:
intervals = (step - self._next_due_frame) // self.every_frames + 1
self._next_due_frame += intervals * self.every_frames
scalar_metrics = {}
for name, value in normalized.items():
if name.endswith("/video"):
if self._trainer.logger is not None:
self._trainer.logger.log_video(name, value, step=step)
continue
self._trainer._log_dict[name].append(value)
self._trainer._last_log[name] = step
scalar_metrics[name] = value
if scalar_metrics and self._trainer.logger is not None:
self._trainer.logger.log_metrics(scalar_metrics, step=step)
def _poll(self) -> None:
while True:
result = self.evaluator.poll()
if result is None:
return
self._log_result(result)
def _schedule(self) -> None:
self._poll()
frames = int(self._trainer.collected_frames)
if frames < self._next_due_frame or self.evaluator.pending:
return
self._trigger(frames)
def _setup(self) -> None:
try:
self._schedule()
except BaseException:
self.evaluator.shutdown()
raise
def _drain(self) -> None:
self._poll()
while self.evaluator.pending:
result = self.evaluator.wait(timeout=self.wait_at_end_timeout)
if result is not None:
self._log_result(result)
elif self.evaluator.pending:
raise RuntimeError(
"Evaluator.wait() returned without completing pending work."
)
self._poll()
def _shutdown(self) -> None:
try:
self._poll()
if self.wait_at_end:
self._drain()
frames = int(self._trainer.collected_frames)
if self.run_at_end and self._last_completed_frame != frames:
if not self.evaluator.pending and self._trigger(frames):
if self.wait_at_end:
self._drain()
finally:
self.evaluator.shutdown()
def state_dict(self) -> dict[str, Any]:
return {
"next_due_frame": self._next_due_frame,
"last_completed_frame": self._last_completed_frame,
}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
self._next_due_frame = int(state_dict["next_due_frame"])
last_completed_frame = state_dict.get("last_completed_frame")
self._last_completed_frame = (
None if last_completed_frame is None else int(last_completed_frame)
)
self._last_triggered_frame = None
[docs]
def register(self, trainer: Trainer, name: str = "evaluator_hook") -> None:
self._trainer = trainer
trainer.register_module(name, self)
trainer.register_op("setup", self._setup)
trainer.register_op("post_steps", self._schedule)
trainer.register_op("shutdown", self._shutdown)
[docs]
class UpdateWeights(TrainerHookBase):
"""A collector weights update hook class.
This hook must be used whenever the collector policy weights sit on a
different device than the policy weights being trained by the Trainer.
In that case, those weights must be synced across devices at regular
intervals. If the devices match, this will result in a no-op.
Args:
collector (BaseCollector, optional): A data collector where the policy
weights must be synced. Not required when a ``sender`` is given.
update_weights_interval (int, optional): Interval where the sync must
take place, counted in units of ``interval_unit``. Default: ``1``.
policy_weights_getter (Callable, optional): A callable that returns the policy
weights to sync. Used for backward compatibility. If both this and
weight_update_map are provided, weight_update_map takes precedence.
weight_update_map (dict[str, str], optional): A mapping from destination paths
(keys in collector's weight_sync_schemes) to source paths on the trainer.
Example: ``{"policy": "loss_module.actor_network", "replay_buffer.transforms[0]": "loss_module.critic_network"}``.
trainer (Trainer, optional): The trainer instance, required when using
weight_update_map to resolve source paths, or when
``interval_unit="optim_steps"`` (to read the optimizer step count).
Keyword Args:
sender (optional): A weight-sync sender object exposing an
``update_weights()`` method (e.g. the sender returned by a
:class:`~torchrl.weight_update.weight_sync_schemes.WeightSyncScheme`'s
``create_sender()``). When provided, weights are pushed through the
sender instead of ``collector.update_policy_weights_()``. This is
used by LLM trainers whose inference engine (vLLM, SGLang) is fed
by a standalone sender.
interval_unit (str, optional): Unit of ``update_weights_interval``:
``"batches"`` (default) counts collected batches and registers the
hook at the ``post_steps`` stage; ``"optim_steps"`` counts
optimizer steps and registers the hook at the ``post_optim`` stage,
enabling weight pushes in the middle of an optimization loop.
Examples:
>>> # Legacy usage with policy_weights_getter
>>> update_weights = UpdateWeights(
... trainer.collector, T,
... policy_weights_getter=lambda: TensorDict.from_module(policy)
... )
>>> trainer.register_op("post_steps", update_weights)
>>> # New usage with weight_update_map
>>> update_weights = UpdateWeights(
... trainer.collector, T,
... weight_update_map={
... "policy": "loss_module.actor_network",
... "replay_buffer.transforms[0]": "loss_module.critic_network"
... },
... trainer=trainer
... )
>>> trainer.register_op("post_steps", update_weights)
>>> # Sender-based usage with optimizer-step cadence (LLM trainers)
>>> update_weights = UpdateWeights(
... update_weights_interval=10,
... trainer=trainer,
... sender=sender,
... interval_unit="optim_steps",
... )
>>> update_weights.register(trainer)
"""
def __init__(
self,
collector: BaseCollector | None = None,
update_weights_interval: int = 1,
policy_weights_getter: Callable[[Any], Any] | None = None,
weight_update_map: dict[str, str] | None = None,
trainer: Trainer | None = None,
*,
sender: Any | None = None,
interval_unit: Literal["batches", "optim_steps"] = "batches",
):
self.collector = collector
self.update_weights_interval = update_weights_interval
self.counter = 0
self.policy_weights_getter = policy_weights_getter
self.weight_update_map = weight_update_map
self.trainer = trainer
self.sender = sender
self.interval_unit = interval_unit
self._last_update_count = 0
# Validate inputs
if update_weights_interval < 1:
raise ValueError("update_weights_interval must be >= 1")
if interval_unit not in ("batches", "optim_steps"):
raise ValueError(
f"interval_unit must be 'batches' or 'optim_steps', got {interval_unit!r}"
)
if sender is None and collector is None:
raise ValueError("either a collector or a sender must be provided")
if sender is not None and (
policy_weights_getter is not None or weight_update_map is not None
):
raise ValueError(
"sender is mutually exclusive with policy_weights_getter and "
"weight_update_map"
)
if weight_update_map is not None and trainer is None:
raise ValueError("trainer must be provided when using weight_update_map")
if interval_unit == "optim_steps" and trainer is None:
raise ValueError(
"trainer must be provided when interval_unit='optim_steps'"
)
def _optimizer_step_count(self) -> int:
"""Read the optimizer step count from the trainer.
Prefers the stepper's ``optimizer_step_count`` (which discounts
gradient-accumulation micro-steps and skipped steps) and falls back to
the trainer's raw optimization-loop counter.
"""
stepper = getattr(self.trainer, "optimization_stepper", None)
count = getattr(stepper, "optimizer_step_count", None)
if count is None:
count = self.trainer._optim_count
return count
def __call__(self):
if self.interval_unit == "optim_steps":
count = self._optimizer_step_count()
if count - self._last_update_count < self.update_weights_interval:
return
self._last_update_count = count
else:
self.counter += 1
if self.counter % self.update_weights_interval != 0:
return
if self.sender is not None:
self.sender.update_weights()
# New approach: use weight_update_map if provided
elif self.weight_update_map is not None:
self._update_with_map()
# Legacy approach: use policy_weights_getter
else:
weights = (
self.policy_weights_getter()
if self.policy_weights_getter is not None
else None
)
if weights is not None:
self.collector.update_policy_weights_(weights)
else:
self.collector.update_policy_weights_()
def _update_with_map(self):
"""Update weights using the weight_update_map."""
from torchrl.weight_update.weight_sync_schemes import WeightStrategy
weights_dict = {}
for destination, source_path in self.weight_update_map.items():
# Resolve the source module from the trainer
source_module = _resolve_module(self.trainer, source_path)
# Get the scheme for this destination to know the extraction strategy
if (
hasattr(self.collector, "_weight_sync_schemes")
and self.collector._weight_sync_schemes
and destination in self.collector._weight_sync_schemes
):
scheme = self.collector._weight_sync_schemes[destination]
strategy = WeightStrategy(extract_as=scheme.strategy_str)
weights = strategy.extract_weights(source_module)
else:
# Fallback: use TensorDict extraction if no scheme found
weights = TensorDict.from_module(source_module)
weights_dict[destination] = weights
# Send all weights atomically
self.collector.update_policy_weights_(weights_dict=weights_dict)
[docs]
def register(self, trainer: Trainer, name: str = "update_weights"):
if self.trainer is None:
self.trainer = trainer
trainer.register_module(name, self)
stage = "post_optim" if self.interval_unit == "optim_steps" else "post_steps"
trainer.register_op(
stage,
self,
)
def state_dict(self) -> dict:
return {
"counter": self.counter,
"last_update_count": self._last_update_count,
}
def load_state_dict(self, state_dict) -> None:
self.counter = state_dict.get("counter", 0)
self._last_update_count = state_dict.get("last_update_count", 0)
[docs]
class CountFramesLog(TrainerHookBase):
"""A frame counter hook.
Args:
frame_skip (int): frame skip of the environment. This argument is
important to keep track of the total number of frames, not the
apparent one.
log_pbar (bool, optional): if ``True``, the reward value will be logged on
the progression bar. Default is `False`.
Examples:
>>> count_frames = CountFramesLog(frame_skip=frame_skip)
>>> trainer.register_op("pre_steps_log", count_frames)
"""
@classmethod
def __new__(cls, *args, **kwargs):
cls.frame_count = 0
return super().__new__(cls)
def __init__(self, frame_skip: int, log_pbar: bool = False):
self.frame_skip = frame_skip
self.log_pbar = log_pbar
def __call__(self, batch: TensorDictBase) -> dict:
if ("collector", "mask") in batch.keys(True):
current_frames = (
batch.get(("collector", "mask")).sum().item() * self.frame_skip
)
else:
current_frames = batch.numel() * self.frame_skip
self.frame_count += current_frames
return {"n_frames": self.frame_count, "log_pbar": self.log_pbar}
[docs]
def register(self, trainer: Trainer, name: str = "count_frames_log"):
trainer.register_module(name, self)
trainer.register_op(
"pre_steps_log",
self,
)
def state_dict(self) -> dict:
return {"frame_count": self.frame_count}
def load_state_dict(self, state_dict) -> None:
self.frame_count = state_dict["frame_count"]
def _check_input_output_typehint(
func: Callable, input: type | list[type], output: type
):
# Placeholder for a function that checks the types input / output against expectations
return
def flatten_dict(d):
"""Flattens a dictionary with sub-dictionaries accessed through point-separated (:obj:`"var1.var2"`) fields."""
out = {}
for key, item in d.items():
if isinstance(item, dict):
item = flatten_dict(item)
for _key, _item in item.items():
out[".".join([key, _key])] = _item
else:
out[key] = item
return out
[docs]
class TargetNetUpdaterHook(TrainerHookBase):
"""A hook for target parameters update.
Examples:
>>> # define a loss module
>>> loss_module = SACLoss(actor_network, qvalue_network)
>>> # define a target network updater
>>> target_net_updater = SoftUpdate(loss_module)
>>> # define a target network updater hook
>>> target_net_updater_hook = TargetNetUpdaterHook(target_net_updater)
>>> # register the target network updater hook
>>> trainer.register_op("post_optim", target_net_updater_hook)
"""
def __init__(self, target_params_updater: TargetNetUpdater):
if not isinstance(target_params_updater, TargetNetUpdater):
raise ValueError(
f"Expected a target network updater, got {type(target_params_updater)=}"
)
self.target_params_updater = target_params_updater
def __call__(self, tensordict: TensorCollection | None = None):
self.target_params_updater.step()
return tensordict
[docs]
def register(self, trainer: Trainer, name: str):
trainer.register_op("post_steps", self)
[docs]
class ValueEstimatorHook(TrainerHookBase):
"""A hook that computes value estimates over a collected batch.
Wraps a value estimator module such as :class:`~torchrl.objectives.value.GAE`
and applies it to the whole collected batch at the ``pre_epoch`` stage, so
that advantage and value-target entries are available when the loss module
consumes sub-batches during optimization.
In async-collection mode the training loop has no batch to hand to the
``pre_epoch`` stage (``batch`` is ``None``); the hook then passes the batch
through untouched, and value estimates are expected to be computed
elsewhere (e.g. by a replay-buffer transform).
Args:
value_estimator (Callable[[TensorDictBase], TensorDictBase]): the value
estimator to apply to the collected batch, e.g. an instance of
:class:`~torchrl.objectives.value.GAE`.
Examples:
>>> gae = GAE(gamma=0.99, lmbda=0.95, value_network=critic, average_gae=True)
>>> value_estimator_hook = ValueEstimatorHook(gae)
>>> value_estimator_hook.register(trainer)
"""
def __init__(
self, value_estimator: Callable[[TensorDictBase], TensorDictBase]
) -> None:
self.value_estimator = value_estimator
def __call__(self, batch: TensorDictBase | None) -> TensorDictBase | None:
if batch is None:
return batch
return self.value_estimator(batch)
def state_dict(self) -> dict[str, Any]:
if hasattr(self.value_estimator, "state_dict"):
return {"value_estimator": self.value_estimator.state_dict()}
return {}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
if "value_estimator" in state_dict and hasattr(
self.value_estimator, "load_state_dict"
):
self.value_estimator.load_state_dict(state_dict["value_estimator"])
[docs]
def register(self, trainer: Trainer, name: str = "value_estimator"):
trainer.register_op("pre_epoch", self)
trainer.register_module(name, self)
[docs]
class LRSchedulerHook(TrainerHookBase):
"""A hook that steps a learning-rate scheduler during training.
Args:
scheduler (torch.optim.lr_scheduler.LRScheduler): the scheduler to step.
interval (Literal["batch", "optim"], optional): ``"batch"`` to step the
scheduler once per collected batch, or ``"optim"`` to step it after
every optimization step. With ``"optim"``, the number of scheduler
steps per collected batch scales with ``num_epochs`` and the number
of sub-batches per batch. Defaults to ``"batch"``.
Once registered with a trainer, the hook only steps the scheduler when at
least one optimization step has run since its last call, so the learning
rate is not decayed during warmup phases (e.g. while
``collector.init_random_frames`` has not been reached).
Examples:
>>> scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100)
>>> lr_scheduler_hook = LRSchedulerHook(scheduler)
>>> lr_scheduler_hook.register(trainer)
"""
def __init__(
self,
scheduler: torch.optim.lr_scheduler.LRScheduler,
interval: Literal["batch", "optim"] = "batch",
) -> None:
if interval not in ("batch", "optim"):
raise ValueError(f"interval must be 'batch' or 'optim', got {interval}")
self.scheduler = scheduler
self.interval = interval
self._trainer_ref: weakref.ReferenceType[Trainer] | None = None
self._last_optim_count: int = 0
def __call__(self, batch: TensorDictBase | None = None) -> TensorDictBase | None:
trainer = self._trainer_ref() if self._trainer_ref is not None else None
if trainer is not None:
optim_count = trainer._optim_count
if optim_count == self._last_optim_count:
# no optimization step has run since the last call (e.g. during
# the init_random_frames warmup): keep the learning rate as is
return batch
self._last_optim_count = optim_count
self.scheduler.step()
return batch
def state_dict(self) -> dict[str, Any]:
return {
"scheduler": self.scheduler.state_dict(),
"last_optim_count": self._last_optim_count,
}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
self.scheduler.load_state_dict(state_dict["scheduler"])
self._last_optim_count = state_dict.get("last_optim_count", 0)
[docs]
def register(self, trainer: Trainer, name: str = "lr_scheduler"):
self._trainer_ref = weakref.ref(trainer)
dest = "post_optim" if self.interval == "optim" else "post_steps"
trainer.register_op(dest, self)
trainer.register_module(name, self)
[docs]
class UTDRHook(TrainerHookBase):
"""Hook for logging Update-to-Data (UTD) ratio during async collection.
The UTD ratio measures how many optimization steps are performed per
collected data sample, providing insight into training efficiency during
asynchronous data collection. This metric is particularly useful for
off-policy algorithms where data collection and training happen concurrently.
The UTD ratio is calculated as: (batch_size * update_count) / write_count
where:
- batch_size: Size of batches sampled from replay buffer
- update_count: Total number of optimization steps performed
- write_count: Total number of samples written to replay buffer
Args:
trainer (Trainer): The trainer instance to monitor for UTD calculation.
Must have async_collection=True for meaningful results.
Note:
This hook is only meaningful when async_collection is enabled, as it
relies on the replay buffer's write_count to track data collection progress.
"""
def __init__(self, trainer: Trainer):
self.trainer = trainer
def __call__(self, batch: TensorDictBase | None = None) -> dict:
if (
hasattr(self.trainer, "replay_buffer")
and self.trainer.replay_buffer is not None
):
write_count = self.trainer.replay_buffer.write_count
batch_size = self.trainer.replay_buffer.batch_size
else:
write_count = self.trainer.collector.getattr_rb("write_count")
batch_size = self.trainer.collector.getattr_rb("batch_size")
if not write_count:
return {}
if batch_size is None and rl_warnings():
warnings.warn("Batch size is not set. Using 1.")
batch_size = 1
update_count = self.trainer._optim_count
utd_ratio = batch_size * update_count / write_count
return {
"utd_ratio": utd_ratio,
"write_count": write_count,
"update_count": update_count,
"log_pbar": False,
}
[docs]
def register(self, trainer: Trainer, name: str = "utdr_hook"):
"""Register the UTD ratio hook with the trainer.
Args:
trainer (Trainer): The trainer to register with.
name (str): Name to use when registering the hook module.
"""
trainer.register_op("pre_steps_log", self)
trainer.register_module(name, self)
[docs]
def state_dict(self) -> dict[str, Any]:
"""Return state dictionary for checkpointing."""
return {}
[docs]
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
"""Load state from dictionary."""
[docs]
class EarlyStopping(TrainerHookBase):
"""Early stopping hook for :class:`~torchrl.trainers.Trainer`.
This hook monitors a scalar metric and stops training when that metric
does not improve according to a configured criterion.
By default, the hook monitors ``"r_evaluation"``.
Args:
monitor (NestedKey, optional): Metric name to monitor.
Defaults to ``"r_evaluation"``.
mode (Literal["min", "max"], optional): One of ``"min"`` or ``"max"``.
In ``"max"`` mode, larger metric values are considered better.
Defaults to ``"max"``.
min_delta (float, optional): Minimum absolute improvement required to
qualify as better. Defaults to ``0.0``.
patience (int, optional): Maximum number of non-improving frames
allowed before stopping. Defaults to ``100_000``.
wait_for (int, optional): Number of initial frames to ignore before
checking the stopping criterion. Defaults to ``1_000_000``.
check_finite (bool, optional): If ``True``, non-finite metric values
(NaN or inf) trigger early stopping. Defaults to ``True``.
Examples:
>>> LogScalar(("next", "reward"), "r_training").register(trainer)
>>> EarlyStopping(monitor="r_training", patience=10_000).register(trainer)
"""
def __init__(
self,
*,
monitor: NestedKey = "r_evaluation",
mode: Literal["min", "max"] = "max",
min_delta: float = 0.0,
patience: int = 100_000,
wait_for: int = 1_000_000,
check_finite: bool = True,
) -> None:
if mode not in {"min", "max"}:
raise ValueError(f"mode must be either 'min' or 'max', got {mode}.")
if patience < 0:
raise ValueError(f"patience must be >= 0, got {patience}.")
if wait_for < 0:
raise ValueError(f"wait_for must be >= 0, got {wait_for}.")
self.monitor = monitor
self.mode = mode
self.min_delta = float(min_delta)
self.patience = int(patience)
self.wait_for = int(wait_for)
self.check_finite = check_finite
self.best_score: float | None = None
self.stop_reason: str | None = None
self._trainer: Trainer | None = None
self._last_improvement_frame: int | None = None
def _resolve_metric(self, trainer: Trainer) -> float:
metric_values = trainer._log_dict.get(self.monitor, None)
if not metric_values:
raise RuntimeError(
"EarlyStopping could not find monitored metric "
f"'{self.monitor}' in trainer._log_dict."
)
metric = metric_values[-1]
if isinstance(metric, torch.Tensor):
if metric.numel() != 1:
raise RuntimeError(
"EarlyStopping expects scalar metrics, "
f"got shape {tuple(metric.shape)} for '{self.monitor}'."
)
metric = float(metric.item())
else:
metric = float(metric)
return metric
def _is_improvement(self, score: float, best_score: float) -> bool:
if self.mode == "max":
return score > (best_score + self.min_delta)
return score < (best_score - self.min_delta)
def _stop(self, trainer: Trainer, reason: str) -> None:
self.stop_reason = reason
trainer.request_stop(reason)
def __call__(self, batch: TensorDictBase | None = None) -> None:
if self._trainer is not None and self._trainer._stop_training:
return
if self._trainer is None:
raise RuntimeError("EarlyStopping is not attached to a trainer.")
trainer = self._trainer
score = self._resolve_metric(trainer)
current_frame = int(trainer.collected_frames)
if current_frame < self.wait_for:
return
if self.check_finite and not math.isfinite(score):
self._stop(
trainer,
f"Monitored metric '{self.monitor}' became non-finite ({score}).",
)
return
if self.best_score is None:
self.best_score = score
self._last_improvement_frame = current_frame
return
if self._is_improvement(score, self.best_score):
self.best_score = score
self._last_improvement_frame = current_frame
else:
if current_frame - self._last_improvement_frame >= self.patience:
self._stop(
trainer,
f"Monitored metric '{self.monitor}' did not improve for "
f"{current_frame - self._last_improvement_frame} frames.",
)
def state_dict(self) -> dict[str, Any]:
return {
"best_score": self.best_score,
"stop_reason": self.stop_reason,
"_last_improvement_frame": self._last_improvement_frame,
}
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
self.best_score = state_dict.get("best_score", None)
self.stop_reason = state_dict.get("stop_reason", None)
self._last_improvement_frame = state_dict.get("_last_improvement_frame", None)
[docs]
def register(self, trainer: Trainer, name: str = "early_stopping") -> None:
self._trainer = trainer
trainer.register_op("post_steps_log", self)
trainer.register_module(name, self)
