[Feature] Add Learner primitive (LocalLearner, FSDP2Learner)

[theap06]

Summary

Introduces torchrl.trainers.Learner: a backend-agnostic entry point for taking
one optimization step on a tensordict batch with a given LossModule.
LocalLearner is the single-process reference implementation; FSDP2Learner
shards the same model with fully_shard and reuses the training step unchanged.

It plays the same role for training that Collector plays for data collection
and LLMWrapperBase plays for generation/scoring: a fixed, TensorDict-native
contract with interchangeable backends, so algorithm code does not need to know
whether the update runs on one device, under sharded training, or on a remote
training process.

Design

• Learner.update() is concrete, in the base class, and touches only
  self.model / self.optimizer / self.clip_grad_norm /
  self.grad_accum_steps: zero_grad -> forward -> sum the loss module's
  "loss"-prefixed output keys -> backward -> optional grad-norm clip ->
  optimizer step. This is what lets FSDP2Learner reuse the exact same step as
  LocalLearner -- sharded training only changes model construction and weight
  gathering, not the step itself.
• get_weights() is the one place sharding is not transparent: it must return
  plain tensors (for WeightSyncScheme.send, which already accepts a
  TensorDictBase), even when the learner's parameters are sharded.
• FSDP2Learner does not decide sharding granularity or device mesh -- it
  accepts a model the caller has already wrapped with fully_shard, exactly as
  bare FSDP2 usage works. Keeping that decision in caller code avoids
  FSDP2Learner becoming a second place those choices are made.

Two contracts worth reading before using this

The optimizer defines what is trained, not model. model is the
weight-sync source and the gradient-sync handle; the parameters that get clipped
and stepped are the ones in optimizer.param_groups. Many TorchRL losses hold
their trainable parameters on the loss module as TensorDictParams, and the
losses that expand their networks (SACLoss, REDQLoss, ...) hold copies of
the modules you passed in:

loss_module = SACLoss(actor, qvalue)
learner = LocalLearner(actor, Adam(loss_module.parameters()))  # correct
learner = LocalLearner(actor, Adam(actor.parameters()))        # critics never train

The second line is silent -- the critics are differentiated on every step and
never updated -- so update() now checks before its first optimizer step and
raises if any parameter received a gradient that no param group covers. Grad-norm
clipping likewise covers optimizer.param_groups, not model.parameters().

A Learner owns exactly one optimizer, so algorithms that deliberately use
several (per-network learning rates, a separate entropy-temperature optimizer)
are not expressible as a single Learner today; use one per optimizer.

Checkpointing is checkpoint() / load_checkpoint(), not state_dict(). A
bare Optimizer is not an nn.Module, so nn.Module.state_dict() structurally
cannot carry its state and a resume would reset Adam's moments. Overriding
state_dict to return it anyway would break the nn.Module contract -- a parent
module calls child.state_dict(destination=...) and discards the return value,
so nesting a Learner inside any other module would silently drop all of its
state. Separate names keep both contracts intact.

Gradients are not checkpointed, so a checkpoint taken mid-accumulation-window
resets the accumulation counter to 0 with a warning rather than resuming at a
non-zero step with empty gradients (which would step after fewer micro-batches
than grad_accum_steps and under-scale that update).

Notes

• update() raises on a non-scalar summed loss, which is what a loss built with
  reduction="none" produces -- previously this surfaced as a torch-internal
  "grad can be implicitly created only for scalar outputs" error.
• The "loss" prefix cannot be tightened to "loss_": "loss" with no
  underscore is a real out_key (DQNLoss, GAILLoss, OnlineDTLoss).
• With FSDP2Learner.get_weights()'s default cpu_offload=True, the full
  weights are returned on rank 0 only and every other rank gets an empty
  tensordict. That is deliberate (gathering the full model onto every rank does
  not scale), but it means an unconditional
  scheme.send(learner.get_weights()) sends nothing from non-zero ranks. Pass
  cpu_offload=False to gather everywhere.
• With gradient accumulation the output key set differs between accumulation
  calls and step calls (grad_norm only appears on the latter).
• No *Config companion yet. These classes have no Trainer wiring in this PR,
  so a Hydra config would have nothing to instantiate against; the configs land
  with the trainer integration (follow-up 2 below).
• These are new, unreleased classes, so the checkpoint() rename carries no
  deprecation shim.

Planned follow-ups (not in this PR)

1. Multi-rank FSDP2Learner verification on real multi-GPU hardware -- the one
   thing I could not test here. The single-rank gloo/CPU tests exercise the
   fully_shard/DTensor code paths but not actual cross-rank communication.
2. Refactor an existing recipe's hand-rolled training loop (e.g. the
   reward-model recipe, once [Feature] Add RewardModelLoss objective for RLHF reward-model training #3922 lands) onto LocalLearner, as the first real
   consumer, with the *Config companions.
3. A RemoteLearner design writeup scoping one external backend (TorchTitan is
   the most PyTorch-native of the candidates and the most plausible first
   integration target) before any implementation.

Tested

test/test_trainer.py: 100 passed, 19 skipped (26 of those are the
Learner/FSDP2Learner tests). Single-rank gloo/CPU only; multi-GPU is
follow-up 1.

c75e16b fixes a checkpoint bug that only surfaced once the suite actually ran:
get_state_dict(full_state_dict=True, cpu_offload=True) does not guarantee
fresh tensors — cpu_offload only copies when the shards are off-CPU, so on a
CPU or single-rank mesh the "gathered" state aliases the live shards and training
on after a checkpoint silently mutates it. FSDP2Learner.checkpoint now clones,
like the base implementation.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3926

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

Reviewed the rebased head 3614a44, with particular attention to the latest optimizer-coverage/checkpoint contract changes and the final FSDP2 checkpoint-cloning fix. The base and FSDP2 checkpoint paths now clone tensor state before returning it, accumulation restores safely reset partial windows, and the optimizer-coverage check runs before the first step. I ran the Learner/FSDP2-focused test selection locally: 27 passed, including the FSDP2 checkpoint round trip, accumulation/reference comparison, and plain-tensor weight gather. I did not find a blocking issue in the current code. The meaningful residual risk is unchanged from the PR description: these checks use world_size=1 gloo/CPU, so actual multi-rank collective behavior and rank-0-only checkpoint/weight gathering still need real distributed coverage.