Training

Reference for HUD training: the TrainingClient API, supported training paths, loss functions, run results, and the hud models CLI for managing forks.

Training advances a model’s weights from the rewards a rollout already produces. A 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.

Training paths

Each path accumulates gradients differently, then applies them with the same optim_step. All operate on a 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’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 completion_kwargs	Marks rollouts for token capture (extra_body={"return_token_ids": True}); OpenAIChatConfig-only (see Token capture).
Job / group	The receipt collecting a batch’s graded Runs; group sets rollouts per task, the GRPO group.
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 Runs (sent inline) or trace_id strings (resolved server-side); the two can be mixed.

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).
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:

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 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.

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:

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[train]').

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; for the Job and Run the trainer consumes see types, and for where rollouts run see runtimes.