docs(cookbooks): formalize the RL training guide

cookbooks/rl-training/README.md

@@ -1,121 +1,135 @@
 # RL Training
 
-On-policy reinforcement learning with the HUD SDK: roll out a taskset with the
-current weights, train on the resulting trajectories, and let the updated weights
-serve the next rollout — all under one model string.
-
-`hud.TrainingClient` targets one **trainable gateway model**. Training advances
-the weights behind that string in place (the HUD training service checkpoints and
-promotes them), so the *same* `model` you sample with is the one you train, and
-each `optim_step` closes the on-policy loop.
-
-| File | What it does |
-|------|--------------|
-| `env.py` | A tiny verifiable env: ask for `a + b`, reward 1.0 if correct (quickstart fallback) |
-| `common.py` | Resolves the rollout source: a deployed taskset on remote boxes, or the local env |
-| `simple_train.py` | The loop with a built-in server-side loss (`importance_sampling`) |
-| `ppo_custom_loss.py` | The loop with a client-side custom loss (GLM-5.2 double-sided IS) |
+This cookbook shows how to train a HUD gateway model on trajectories from a
+HUD taskset. Each iteration collects grouped rollouts, grades them in the
+environment, and updates the model behind the same gateway identifier.
+Subsequent rollouts therefore use the latest weights.
 
-## Run
+Two implementations are included. `simple_train.py` uses a built-in
+server-side loss, while `ppo_custom_loss.py` defines the policy-gradient
+loss in PyTorch and sends the resulting per-token gradients to the training
+service.
+
+| File | Purpose |
+|------|---------|
+| `env.py` | Local arithmetic environment used by the example |
+| `common.py` | Selects a deployed taskset or the local environment |
+| `simple_train.py` | On-policy training with a built-in loss |
+| `ppo_custom_loss.py` | On-policy training with a custom double-sided importance-sampling loss |
 
-Needs `HUD_API_KEY` (from your environment or `.env`). List the gateway models
-on your account, pick a trainable one (the **Trainable** column marks them), and
-set it as the `MODEL` constant at the top of `simple_train.py` /
-`ppo_custom_loss.py`:
+## Setup
+
+Set `HUD_API_KEY` in the environment or in a local `.env` file. Then list
+the gateway models available to the account:
 
 ```bash
-hud models list          # Name | Model (API) | ID | Provider | Agent | Trainable
+hud models list
 ```
 
-**Train on a deployed taskset (the real flow).** You've built a taskset and
-pushed it (`hud deploy` + `hud sync`); now train on it. Set the `TASKSET`
-constant in `common.py` to its name/id and rollouts run on **remote HUD
-boxes** — nothing local:
+Choose a model marked **Trainable** and assign its identifier to `MODEL` in
+the training script.
+
+## Run
+
+The default configuration uses the arithmetic taskset in `env.py` and runs
+it through `LocalRuntime`:
 
 ```bash
 uv run simple_train.py --steps 10
-uv run ppo_custom_loss.py --steps 10
 ```
 
-**Quickstart (self-contained).** Leave `TASKSET` empty and a tiny local
-arithmetic taskset runs against the bundled `env.py`:
+To train on a deployed taskset, set `TASKSET` in `common.py` to the taskset
+name or id. `load_taskset_and_runtime()` will load it with
+`Taskset.from_api(...)` and execute rollouts through `HUDRuntime`. The
+training command remains the same:
 
 ```bash
 uv run simple_train.py --steps 10
 ```
 
-The swap is `common.py`'s `load_taskset_and_runtime()` — `Taskset.from_api(name)`
-+ `HUDRuntime()` for the deployed case, `Taskset(...)` + `LocalRuntime("env.py")`
-for the local one. **The training code is identical either way.**
+The custom-loss example uses the same task and rollout configuration:
 
-## The loop
+```bash
+uv run ppo_custom_loss.py --steps 10
+```
 
-Both scripts are the same five lines — the only difference is the training call:
+Both scripts accept `--group`, `--learning-rate`, and `--max-concurrent`.
+The defaults are a group size of 8, a learning rate of `1e-5`, and at most
+8 concurrent rollouts.
+
+## Training flow
+
+A `Job` spans the training session and accumulates its runs. Each iteration
+selects the runs added by the latest rollout and trains on that batch:
 
 ```python
-taskset, runtime = load_taskset_and_runtime()  # deployed+remote, or local
-session = await Job.start("rl", group=8)  # one job spans the session
-for step in range(steps):
-    start = len(session.runs)
-    await taskset.run(agent, runtime=runtime, job=session)  # roll out current weights
-    batch = session.runs[start:]  # this step's runs
-    await trainer.step(batch, learning_rate=1e-5, group_size=8)  # train + promote
+batch_start = len(session.runs)
+await taskset.run(agent, runtime=runtime, job=session)
+batch = session.runs[batch_start:]
+
+await trainer.step(batch, learning_rate=1e-5, group_size=8)
 ```
 
-The loop only ever touches `job.runs`, so where the rollouts executed — a remote
-leased box or your laptop — is irrelevant to training. Passing the `Run` is
-enough either way:
-
-- **Remote (`HUDRuntime`)** runs fold back only reward + `trace_id`; their full
-  token-level trajectory lives on the platform (collected server-side during the
-  rollout). The client sends the `trace_id` and the training service resolves the
-  trajectory + reward from it.
-- **Local (`LocalRuntime`)** runs carry the token-level `Sample` on each agent
-  turn in `run.trace`, so the client sends the trajectory inline (works even with
-  telemetry off).
-
-You can also pass `trace_id` strings directly, and mix them with `Run`s.
-
-## Two loss tiers
-
-**Built-in (`simple_train.py`).** `trainer.step(...)` = one `forward_backward`
-with a server-side loss, then one `optim_step`. The client stays dependency-light
-(no torch). `loss_fn` mirrors Tinker's native set — `cross_entropy` (supervised),
-`importance_sampling`, `ppo`, `cispo`, `dro`; the policy-gradient ones compute
-advantages from rewards server-side (GRPO over each `group_size` chunk).
-
-**Custom (`ppo_custom_loss.py`).** `trainer.forward_backward_custom(batch, loss_fn)`
-splits the step so *you* write the loss:
-
-1. `forward` (service) runs the current-policy pass and returns per-token tensors
-   (`DatumTensors`: current-policy logprobs π_θ, rollout logprobs q, action mask,
-   reward, group index).
-2. your `loss_fn` builds a differentiable loss over the π_θ logprobs (torch, here).
-3. `backward` (service) applies the resulting per-token gradients.
-
-This mirrors Tinker's `forward_backward_custom` and its `weights = -dC/dlogprobs`
-convention, split across the service boundary. Build the loss out of the
-**provided** logprob tensors (don't re-wrap from `.data`) or gradients won't flow.
-
-## What this supports (and what it doesn't)
-
-The custom path expresses token-level methods whose only moving part is the
-advantage / loss math over per-token tensors:
-
-- **GLM-5.2 direct double-sided IS** (the worked example): reuse rollout logprobs
-  as the behavior proxy, ratio `r = exp(logπ_θ − logπ_rollout)`, hard-mask tokens
-  outside `[1 − ε_l, 1 + ε_h]`, token-level normalization.
-- **Compaction** is free: a rollout is a variable-length list of variable-length
-  turns, and training has no constraint on how many turns a trajectory has or
-  their relative lengths — every turn's `Sample` is a trainable unit.
-- Critic-free credit assignment (TEMPO-style tree-TD, MemPO per-segment,
-  broadcast-advantage + token-level loss) is all advantage math you can write in
-  `loss_fn`.
-
-The one thing the Tinker backend cannot do natively is **train a value network**
-(its loss API is over logprobs, not a value head). GLM-5.2's critic exists only to
-produce token-level advantages, and advantages are an input — so for true
-critic-PPO you host a decoupled critic in the training service (**Option A**:
-value model + GAE, fed as the `advantages` input; deps beyond `tinker` such as a
-small value model are expected there) rather than on Tinker. The examples here use
-a critic-free group baseline as the stand-in.
+`trainer.step(...)` performs `forward_backward` followed by `optim_step`.
+The optimizer step checkpoints and promotes the updated weights behind the
+gateway model, so the next call to `taskset.run(...)` samples the new policy.
+
+The trajectory representation depends on the runtime:
+
+| Runtime | Data sent to training |
+|---------|-----------------------|
+| `HUDRuntime` | Reward and `trace_id`; the training service resolves the token-level trajectory stored by the platform |
+| `LocalRuntime` | Reward and the token-level `Sample` recorded on each agent turn in `run.trace` |
+
+`TrainingClient` also accepts `trace_id` strings directly. A training batch
+may contain `Run` objects, trace ids, or both.
+
+## Built-in losses
+
+`simple_train.py` calls `forward_backward` with
+`loss_fn="importance_sampling"`. The available server-side losses are:
+
+| Loss | Use |
+|------|-----|
+| `cross_entropy` | Supervised training |
+| `importance_sampling` | Group-relative policy-gradient training |
+| `ppo` | PPO objective |
+| `cispo` | CISPO objective |
+| `dro` | Distributionally robust objective |
+
+For the policy-gradient losses, the service computes group-relative
+advantages from rewards using each consecutive `group_size` set of
+trajectories. The built-in path does not require PyTorch on the client.
+
+## Custom loss
+
+`ppo_custom_loss.py` uses `forward_backward_custom` to implement GLM-5.2
+direct double-sided importance sampling. The computation is split across
+the client and training service:
+
+1. The service runs the current-policy forward pass and returns
+   `DatumTensors`, including policy logprobs, rollout logprobs, action masks,
+   rewards, and group indices.
+2. The client computes a differentiable PyTorch loss from the policy
+   logprobs.
+3. The service applies the per-token gradients, after which `optim_step`
+   updates and promotes the model.
+
+The example uses the rollout logprobs as the behavior-policy proxy, computes
+`r = exp(logπ_θ - logπ_rollout)`, masks tokens outside
+`[1 - ε_l, 1 + ε_h]`, and normalizes by the number of trained tokens.
+The loss must use the policy logprob tensors returned by the service.
+Constructing new tensors from `.data` disconnects the computation graph.
+
+### Scope
+
+The custom API supports objectives defined from per-token logprobs, masks,
+rewards, group membership, and externally supplied advantages. It supports
+variable-length, multi-turn trajectories because each turn's `Sample` is an
+independent training datum. Critic-free methods such as grouped baselines,
+tree-based credit assignment, and per-segment advantages can be expressed
+in the client loss.
+
+The backend does not train a value head. A critic-based PPO implementation
+must run the value model separately and pass its token-level advantages into
+the policy loss. The included example uses a group-mean baseline instead.