torchrl.trainers package#
The trainer package provides utilities to write reusable training scripts. The core idea is to use a trainer that implements a nested loop, where the outer loop runs the data collection steps and the inner loop the optimization steps.
Key Features#
Modular hook system: Customize training at 18 different points in the loop
Checkpointing support: Pass a
torchrl.checkpoint.Checkpointfor the unified manifest format. Legacytorch,torchsnapshot, andmemmapbackends remain readable during the migration window.Algorithm trainers: High-level trainers for PPO, A2C, REINFORCE, SAC, offline-to-online SAC, DQN, DDPG, IQL, CQL, and TD3 with Hydra configuration
Builder helpers: Utilities for constructing collectors, losses, and replay buffers
Quick Example#
from torchrl.trainers import Trainer
from torchrl.trainers import UpdateWeights, LogScalar
# Create trainer
trainer = Trainer(
collector=collector,
total_frames=1000000,
loss_module=loss,
optimizer=optimizer,
)
# Register hooks
UpdateWeights(collector, 10).register(trainer)
LogScalar("reward").register(trainer)
# Train
trainer.train()