---
title: "Training"
description: "Reference for HUD training: the TrainingClient API, supported training paths, loss functions, run results, and the hud models CLI for managing forks."
icon: "dumbbell"
mode: "wide"
---

Training advances a model's weights from the rewards a rollout already produces. A
[`TrainingClient`](#trainingclient) drives the managed path against one model slug; lower-level methods
expose the gradient-accumulation and custom-loss steps underneath it. For the end-to-end loop (forking a
model, rolling out a taskset in groups, feeding rewards back) see [training agents](/v6/guides/training-agents).

## Training paths

Each path accumulates gradients differently, then applies them with the same `optim_step`. All operate on
a [`TrainingClient`](#trainingclient) except the last.

| Path | How gradients accumulate | Promote weights |
| --- | --- | --- |
| Managed one-shot | `step(trajectories, learning_rate=...)` | folded into `step` |
| Split accumulate then update | `forward_backward(...)` once or repeated (or `num_substeps>1`) | `optim_step(...)` |
| Client-side custom loss | `forward_backward_custom(trajectories, loss_fn)` | `optim_step(...)` |
| Low-level | `forward(...)` then your loss then `backward(forward_id, weights)` | `optim_step(...)` |
| Supervised (SFT) | any path with `loss_fn="cross_entropy"` (`BuiltinLoss.CROSS_ENTROPY`) | as above |
| External RL stack | export each [`Run`](/v6/reference/types#run)'s `reward` and `trace_id`; no `TrainingClient` | owned by your stack |

Gradients accumulate across `forward_backward` / `forward_backward_custom` / `backward` calls until an
`optim_step` (or `step`, which calls both) applies them, checkpoints, and promotes the new weights.

## Objects

| Object | Role |
| --- | --- |
| Forked model slug | The gateway string you both sample and train; weights advance behind it in place. |
| [`create_agent`](/v6/reference/agents#create-agent) `completion_kwargs` | Marks rollouts for token capture (`extra_body={"return_token_ids": True}`); OpenAIChatConfig-only (see [Token capture](#token-capture)). |
| [`Job`](/v6/reference/types#job) / `group` | The receipt collecting a batch's graded [`Run`](/v6/reference/types#run)s; `group` sets rollouts per task, the GRPO group. |
| [`Run`](/v6/reference/types#run) | One graded rollout; carries its `reward` and either inline token samples or a `trace_id`. |
| `TrainingClient` | Drives managed training for one model slug. |
| `CheckpointResponse` | One node in the model's checkpoint tree. |

## TrainingClient

A `TrainingClient` drives managed training for one model: it accumulates gradients from rewarded
trajectories and advances the weights behind the model's gateway slug in place. Inputs are `Run`s (sent
inline) or `trace_id` strings (resolved server-side); the two can be mixed.

```python
TrainingClient(model, *, api_key=None, base_url=None, api_url=None)
```

| Argument | Default | Meaning |
|----------|---------|---------|
| `model` | - | Trainable model slug or id (the gateway string you also sample). |
| `api_key` | `settings.api_key` | HUD API key. |
| `base_url` | `settings.hud_rl_url` | Training (RL) service. |
| `api_url` | `settings.hud_api_url` | Catalog API (resolves the slug -> id once). |

### Methods

| Method | Returns | Purpose |
|--------|---------|---------|
| `forward_backward(trajectories, *, loss_fn, loss_fn_config=None, group_size=None, reward_scale=1.0, num_substeps=1)` | `ForwardBackwardResult` | Accumulate gradients with a built-in `loss_fn`. |
| `optim_step(*, learning_rate, beta1=0.9, beta2=0.95, eps=1e-8, weight_decay=0.0)` | `OptimStepResult` | Apply gradients, checkpoint, and promote the new weights. |
| `step(trajectories, *, learning_rate, loss_fn="importance_sampling", loss_fn_config=None, group_size=None, reward_scale=1.0, num_substeps=1, weight_decay=0.0)` | `OptimStepResult` | One `forward_backward` then one `optim_step`. |
| `forward_backward_custom(trajectories, loss_fn, *, group_size=None, reward_scale=1.0)` | `ForwardBackwardResult` | Accumulate gradients with a client-side loss (see [Custom losses](#custom-losses)). |
| `forward(trajectories, *, group_size=None, reward_scale=1.0)` | `ForwardResult` | Current-policy forward pass returning per-token tensors. |
| `backward(forward_id, weights, *, metrics=None)` | `ForwardBackwardResult` | Apply caller-computed per-token gradients to a forward pass. |
| `available_losses()` | `list[str]` | Built-in `loss_fn` names this model's provider supports. |
| `checkpoints()` | `list[CheckpointResponse]` | The checkpoint tree, each node with rewards, loss, counts, and a `metrics` blob. |
| `head()` | `CheckpointResponse \| None` | The active checkpoint (the weights the gateway serves), or `None` for base weights. |
| `set_head(checkpoint_id)` | `None` | Promote a checkpoint to head - roll back to, or branch from, it. |

Advantages are group-relative (GRPO-style) when `group_size` is set, normalized within contiguous
**groups** of that size; `None` treats the whole batch as one group. The `loss_fn` selects the
policy-gradient objective applied on top, defaulting to `importance_sampling` (not GRPO itself). The
batch must divide evenly into groups - `forward_backward` rejects a partial final group before spending
a forward pass. `num_substeps` splits the batch for gradient accumulation.

## Inputs

A training input is a recorded trajectory by id, or an inline one:

```python
TrainInput = str | TrajectoryPayload          # trace_id, or inline tokens + reward
```

The methods accept `str | Run | TrajectoryPayload`, mixed freely. A `Run` is converted automatically -
inline `TrajectoryPayload` when it carries token-level samples (local rollout), else its `trace_id`
(remote rollout). Build a `TrajectoryPayload` yourself when the tokens didn't come from a `Run` at all -
e.g. an opponent move sampled inside the environment during self-play, trained with its own reward.

| Type | Fields |
|------|--------|
| `TrajectorySample` | `prompt_token_ids`, `prompt_chunks`, `output_token_ids`, `output_logprobs` |
| `TrajectoryPayload` | `samples: list[TrajectorySample]`, `reward`, `trace_id=None` |

A sample's prompt is `prompt_token_ids` for a flat text prompt, or `prompt_chunks` (serialized text +
image chunks) for a multimodal one, where `prompt_token_ids` is empty. `output_logprobs` are per output
token under the sampling policy.

### Token capture

The token ids and logprobs a sample needs come from the gateway. `completion_kwargs` carries the flag,
and it lives only on `OpenAIChatConfig`, so token capture runs through `OpenAIChatAgent` (the agent
[`create_agent`](/v6/reference/agents#create-agent) selects for gateway models routed over Chat
Completions). Setting `extra_body={"return_token_ids": True}` is enough: that agent sets `logprobs=True`
automatically and writes a `Sample` onto each turn, which a `Run` converts to a `TrajectoryPayload` for
training.

```python
agent = create_agent("arith-rl", completion_kwargs={"extra_body": {"return_token_ids": True}})
```

Calling the gateway directly with your own OpenAI client (for a hand-built payload) takes both flags,
since nothing sets `logprobs` for you:

```python
resp = await client.chat.completions.create(
    model="arith-rl", messages=msgs, logprobs=True,
    extra_body={"return_token_ids": True},
)
choice = resp.choices[0]
sample = TrajectorySample(
    prompt_token_ids=choice.prompt_token_ids,   # gateway-added fields
    output_token_ids=choice.token_ids,
    output_logprobs=[t.logprob for t in choice.logprobs.content],
)
```

## Built-in losses

`loss_fn` is an open string validated against the model's provider; discover the set with
`await trainer.available_losses()`. `BuiltinLoss` lists the common Tinker names (each *is* a `str`):

| `BuiltinLoss` | Value | Use |
|---------------|-------|-----|
| `CROSS_ENTROPY` | `cross_entropy` | Supervised - imitate sampled tokens. |
| `IMPORTANCE_SAMPLING` | `importance_sampling` | On-policy PG, rollout-logprob ratio. |
| `PPO` | `ppo` | Clipped-surrogate PG. |
| `CISPO` | `cispo` | Clipped IS policy optimization. |
| `DRO` | `dro` | Direct reward optimization. |

`loss_fn_config` forwards provider-specific hyperparameters to the loss (e.g. `{"epsilon": 0.2}` for
the `ppo` clip). The supported keys are provider-defined and not every loss accepts config, so prefer
the defaults (`None`) unless a provider documents a key.

## Custom losses

`forward_backward_custom` runs the current-policy forward pass server-side, hands you per-token tensors,
runs your loss locally (torch autograd), and ships the per-token gradients back. Requires torch
(`pip install 'hud-python[train]'`).

```python
import torch
from hud.train import DatumTensors

def my_loss(data: list[DatumTensors], logprobs: list[torch.Tensor]):
    loss = logprobs[0].new_zeros(())
    for datum, policy_lp in zip(data, logprobs):
        ratio = torch.exp(policy_lp - torch.tensor(datum.sampling_logprobs))
        loss = loss - (ratio * datum.reward * torch.tensor(datum.mask)).sum()
    return loss, {"trained": float(len(data))}

await trainer.forward_backward_custom(batch, my_loss, group_size=8)   # accumulate gradients
await trainer.optim_step(learning_rate=1e-5)                          # apply, checkpoint, promote
```

Gradients accumulate on the model session until `optim_step` applies them, so a custom-loss step is not
complete until that call promotes the new weights.

`logprobs[i]` are the current policy for datum `i` as differentiable leaves. Everything else is constant
on the matching `DatumTensors`:

| `DatumTensors` | Meaning |
|----------------|---------|
| `logprobs` | Current-policy, per token (the differentiable leaf). |
| `sampling_logprobs` | Rollout policy, per token. |
| `mask` | `1.0` on action tokens, `0.0` on observation tokens. |
| `reward`, `traj_idx`, `group_idx` | Trajectory reward, source trajectory, GRPO group (or `None`). |

Under the hood `forward` returns a `ForwardResult` (`forward_id` + `data: list[DatumTensors]`);
`backward(forward_id, weights)` applies `weights[d][t] = -dC/dlogprobs`.

## Results

| Type | Fields |
|------|--------|
| `ForwardBackwardResult` | `metrics: dict[str, float]`, `num_datums` |
| `OptimStepResult` | `step`, `checkpoint_id`, `sampler_path`, `state_path`, `model` |
| `CheckpointResponse` | `id`, `name`, `checkpoint_name`, `is_active`, `prev_model_checkpoint_id`, `mean_reward`, `learning_rate`, `loss_fn`, `num_traces`, `num_datums`, `num_tokens`, `metrics`, `created_at` |

## `hud models` CLI

| Command | Purpose |
|---------|---------|
| `hud models list` | List gateway models. |
| `hud models fork <model> --name <slug>` | Fork a team-owned trainable model from an existing one. |
| `hud models checkpoints <model>` | List the checkpoint tree (active head marked). |
| `hud models head <model> [--set <checkpoint-id>]` | Show, or set (rollback/select), the active checkpoint. |

For the end-to-end loop see [training agents](/v6/guides/training-agents); for the
[`Job`](/v6/reference/types#job) and [`Run`](/v6/reference/types#run) the trainer consumes see
[types](/v6/reference/types), and for where rollouts run see [runtimes](/v6/reference/runtime).
