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

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.get_weights() gathers every DTensor leaf via
  full_tensor().
• 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.

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

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