[Algorithm] Add RLHF reward-model training recipe

.github/unittest/linux_sota/scripts/test_sota.py

@@ -24,6 +24,19 @@
   logger.eval_iter=2 \
   logger.eval_episodes=4 \
   checkpoint.save_iter=2
+""",
+    "reward_model_training": """python sota-implementations/reward_model_training/reward_model.py \
+  model.name= \
+  data.dataset_name= \
+  data.synthetic_size=32 \
+  data.batch_size=8 \
+  data.max_length=32 \
+  optim.max_iters=3 \
+  logger.eval_iter=2 \
+  logger.eval_iters=1 \
+  logger.log_interval=1 \
+  logger.backend= \
+  checkpoint.save_iter=2
 """,
     "diffusion_bc": """python sota-implementations/diffusion_bc/diffusion_bc.py \
   optim.gradient_steps=55 \

docs/source/reference/llms.rst

@@ -544,6 +544,18 @@ SFT
     SFTLoss
     SFTLossOutput
 
+Reward Model Training
+~~~~~~~~~~~~~~~~~~~~~
+
+.. currentmodule:: torchrl.objectives.llm
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    RewardModelLoss
+    RewardModelLossOutput
+
 .. currentmodule:: torchrl.data.llm
 
 .. autosummary::

docs/source/reference/llms_objectives.rst

@@ -29,3 +29,13 @@ SFT
 
     SFTLoss
     SFTLossOutput
+
+Reward Model Training
+---------------------
+
+.. autosummary::
+    :toctree: generated/
+    :template: rl_template.rst
+
+    RewardModelLoss
+    RewardModelLossOutput

sota-check/README.md

@@ -26,7 +26,7 @@ export MUJOCO_GL=egl
 conda create -n rl-sota-bench python=3.10 -y 
 conda install anaconda::libglu -y
 pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu121
-pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit hydra-core transformers
+pip3 install "gymnasium[atari,mujoco]" vmas tqdm wandb pygame "moviepy<2.0.0" imageio submitit hydra-core transformers datasets
 
 cd /path/to/tensordict
 python setup.py develop

sota-check/run_reward_model_training.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+#SBATCH --job-name=reward_model_training
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=8
+#SBATCH --gres=gpu:1
+#SBATCH --output=slurm_logs/reward_model_training_%j.txt
+#SBATCH --error=slurm_errors/reward_model_training_%j.txt
+
+current_commit=$(git rev-parse --short HEAD)
+project_name="torchrl-example-check-$current_commit"
+group_name="reward_model_training"
+export PYTHONPATH=$(dirname $(dirname $PWD))
+python $PYTHONPATH/sota-implementations/reward_model_training/reward_model.py \
+  model.name=gpt2 \
+  model.device=cuda:0 \
+  logger.backend=wandb \
+  logger.project_name="$project_name" \
+  logger.group_name="$group_name"
+
+# Capture the exit status of the Python command
+exit_status=$?
+# Write the exit status to a file
+if [ $exit_status -eq 0 ]; then
+  echo "${group_name}_${SLURM_JOB_ID}=success" >> report.log
+else
+  echo "${group_name}_${SLURM_JOB_ID}=error" >> report.log
+fi

sota-check/submitit-release-check.sh

@@ -64,6 +64,7 @@ scripts=(
     run_ppo_atari.sh
     run_ppo_mujoco.sh
     run_rnd_mujoco.sh
+    run_reward_model_training.sh
     run_sac.sh
     run_td3.sh
     run_td3bc.sh

sota-implementations/reward_model_training/README.md

@@ -0,0 +1,98 @@
+# RLHF reward-model training
+
+Trains a scalar **reward model** from pairwise human-preference data using the
+Bradley-Terry objective
+[`RewardModelLoss`](../../torchrl/objectives/llm/reward.py). Given a prompt and two
+responses -- a `chosen` one preferred by an annotator and a `rejected` one -- the model
+learns to assign a higher score to the chosen response:
+
+```
+loss = -log σ(r_θ(x, y_chosen) - r_θ(x, y_rejected))
+```
+
+This is the reward-modelling stage that precedes policy optimization (e.g.
+[GRPO](../grpo/)) in RLHF pipelines.
+
+This recipe is a minimal **single-node reference implementation**, not a large-scale
+training stack. See [Scaling and integration](#scaling-and-integration) for how it is
+meant to fit into serious training.
+
+## Install
+
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+Train on the default summarization-comparisons dataset with a small Qwen backbone:
+
+```bash
+python reward_model.py model.name=Qwen/Qwen2.5-0.5B
+```
+
+The backbone is **model-agnostic**: `model.name` accepts any Hugging Face model id
+usable with `AutoModelForSequenceClassification` (a single-output head is attached
+automatically), e.g. `facebook/opt-125m` or a larger instruct model.
+
+Use a different preference dataset (must expose `chosen`/`rejected`, and optionally
+`prompt` fields):
+
+```bash
+python reward_model.py model.name=Qwen/Qwen2.5-0.5B data.dataset_name=Anthropic/hh-rlhf \
+  data.split_train=train data.split_val=test
+```
+
+Add the score-centering regularizer or freeze part of the backbone:
+
+```bash
+python reward_model.py loss.center_coeff=0.01 optim.freeze_frac=0.7
+```
+
+## Configuration
+
+Key fields in [`config.yaml`](config.yaml):
+
+| Group | Field | Description |
+|-------|-------|-------------|
+| `model` | `name` | HF backbone id; empty -> tiny from-scratch model (CI). |
+| `model` | `device` | Training device; `null` auto-selects CUDA/MPS/CPU. |
+| `data` | `dataset_name` | HF preference dataset; empty -> synthetic data (CI). |
+| `data` | `max_length`, `batch_size`, `max_samples` | Tokenization / sampling. |
+| `optim` | `max_iters`, `lr`, `freeze_frac`, `clip_grad` | Optimization. |
+| `loss` | `reduction`, `center_coeff` | Bradley-Terry loss options. |
+| `logger` | `backend`, `eval_iter`, `eval_iters` | Logging / evaluation. |
+
+## Hermetic smoke run
+
+Leaving `model.name` and `data.dataset_name` empty builds a tiny from-scratch model and
+a synthetic preference dataset, so the recipe runs end-to-end with no download and
+without the `datasets` dependency (this is what CI exercises):
+
+```bash
+python reward_model.py model.name= data.dataset_name= \
+  optim.max_iters=3 logger.eval_iter=2 logger.eval_iters=1 logger.backend=
+```
+
+## Scaling and integration
+
+This recipe is a minimal single-node baseline: a readable, single-process training loop
+that exercises the loss and data format end to end. It is **not** a large-scale
+reward-model training stack, and is not meant to become one.
+
+TorchRL owns the **contract**, not the parallelism. What this recipe provides:
+
+- a canonical TensorDict **preference-data format** (prompt / chosen / rejected),
+- a canonical **reward-model loss**
+  ([`RewardModelLoss`](../../torchrl/objectives/llm/reward.py), Bradley-Terry),
+- a small **reference trainer** (this recipe),
+- and, in the future, **adapters/scorers** so an externally trained reward model plugs
+  back into TorchRL GRPO/PPO/SFT pipelines.
+
+For serious distributed reward-model training (parallelism, sharded checkpointing,
+fault tolerance, packed/streamed datasets, mixed precision, activation checkpointing,
+multi-node orchestration), use a dedicated backend such as TRL/Accelerate/FSDP,
+TorchTitan/FSDP2, Megatron/NeMo, or OpenRLHF, while **preserving the TorchRL data and
+scoring contract** so the resulting model plugs straight back into TorchRL's RLHF
+pipelines. Interfacing with those backends, rather than reimplementing them, is the
+intended direction for scaling this and the rest of the TorchRL LLM stack.

sota-implementations/reward_model_training/config.yaml

@@ -0,0 +1,51 @@
+seed: 42
+
+# reward-model backbone
+model:
+  # Any HF model id usable with AutoModelForSequenceClassification (e.g.
+  # Qwen/Qwen2.5-0.5B, facebook/opt-125m). Leave empty to build a tiny
+  # from-scratch model (used by CI).
+  name: Qwen/Qwen2.5-0.5B
+  # Training device. null = auto-select CUDA/MPS/CPU via get_available_device().
+  device: null
+
+# preference data
+data:
+  # HF preference dataset with prompt/chosen/rejected fields. Leave empty to use a
+  # hermetic synthetic dataset (no download, no `datasets` dependency).
+  dataset_name: CarperAI/openai_summarize_comparisons
+  split_train: train
+  split_val: valid1
+  max_length: 550
+  batch_size: 16
+  max_samples: null  # cap the number of pairs loaded (null = all)
+  synthetic_size: 64  # number of synthetic pairs when dataset_name is empty
+
+# optim
+optim:
+  max_iters: 10000
+  lr: 1.0e-5
+  weight_decay: 0.0
+  clip_grad: 1.0
+  freeze_frac: 0.0  # fraction of backbone layers to freeze (0 = train all)
+
+# loss
+loss:
+  reduction: mean
+  center_coeff: null  # set a float to add the score-centering regularizer
+
+# logging / eval
+logger:
+  backend: wandb
+  project_name: reward_model_training
+  group_name: null
+  exp_name: RewardModel_${model.name}
+  mode: online
+  eval_iter: 500
+  eval_iters: 20
+  log_interval: 10
+
+# checkpointing
+checkpoint:
+  save_dir: ./reward_model_ckpt
+  save_iter: 1000

sota-implementations/reward_model_training/requirements.txt

@@ -0,0 +1,5 @@
+datasets
+hydra-core
+tqdm
+transformers
+wandb