[Algorithm] Expert Iteration and SFT

Author: vmoens

docs/source/reference/llms.rst

@@ -200,6 +200,7 @@ transforms).
 
     DataLoadingPrimer
     KLRewardTransform
+    RetrieveLogProb
     MCPToolTransform
     BrowserTransform
     PythonInterpreter
@@ -256,6 +257,9 @@ LLM post training require some appropriate versions of the losses implemented in
 GRPO
 ~~~~
 
+The :class:`~torchrl.objectives.llm.GRPOLoss` class is a thin wrapper around the :class:`~torchrl.objectives.PPOLoss` class
+that codes the LLM-specific functionnalities.
+
 .. currentmodule:: torchrl.objectives.llm
 
 .. autosummary::
@@ -265,3 +269,24 @@ GRPO
     GRPOLoss
     GRPOLossOutput
     MCAdvantage
+
+
+SFT
+~~~
+
+.. currentmodule:: torchrl.objectives.llm
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    SFTLoss
+    SFTLossOutput
+
+.. currentmodule:: torchrl.data.llm
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    TopKRewardSelector

sota-implementations/cql/online_config.yaml

@@ -44,7 +44,7 @@ optim:
   critic_lr: 3e-4
   weight_decay: 0.0
   batch_size: 256
-  optim_steps_per_batch: 200
+  optim_dialog_turns_per_batch: 200
 
 # Policy and model
 model:

sota-implementations/dreamer/config.yaml

@@ -26,7 +26,7 @@ optimization:
   value_lr: 8e-5
   kl_scale: 1.0
   free_nats: 3.0
-  optim_steps_per_batch: 80
+  optim_dialog_turns_per_batch: 80
   gamma: 0.99
   lmbda: 0.95
   imagination_horizon: 15

sota-implementations/dreamer/dreamer.py

@@ -137,7 +137,7 @@ def main(cfg: DictConfig):  # noqa: F821
         scaler3 = GradScaler()
 
     init_random_frames = cfg.collector.init_random_frames
-    optim_steps_per_batch = cfg.optimization.optim_steps_per_batch
+    optim_dialog_turns_per_batch = cfg.optimization.optim_dialog_turns_per_batch
     grad_clip = cfg.optimization.grad_clip
     eval_iter = cfg.logger.eval_iter
     eval_rollout_steps = cfg.logger.eval_rollout_steps
@@ -179,7 +179,7 @@ def compile_rssms(module):
             t_loss_actor = 0.0
             t_loss_critic = 0.0
             t_loss_model = 0.0
-            for _ in range(optim_steps_per_batch):
+            for _ in range(optim_dialog_turns_per_batch):
                 # sample from replay buffer
                 t_sample_init = time.time()
                 sampled_tensordict = replay_buffer.sample().reshape(-1, batch_length)