diff --git a/cookbooks/rl-training/README.md b/cookbooks/rl-training/README.md index 5911985bb..dd6f0ce3c 100644 --- a/cookbooks/rl-training/README.md +++ b/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.