> ## Documentation Index
> Fetch the complete documentation index at: https://docs.relayintelligence.org/llms.txt
> Use this file to discover all available pages before exploring further.

# Training loops

> Write the training loop yourself and run each step on Relay's GPU, in the style of Thinking Machines' Tinker.

A training client is a GPU with π0.5 loaded. Your Python program is the training loop: it decides the data, the schedule, the loss and when to evaluate, and each call runs one step of it on the GPU.

```python theme={null}
from relay.training import ServiceClient, types

service_client = ServiceClient()
training_client = service_client.create_lora_training_client(base_model="pi05", rank=32)

for step, batch in enumerate(batches):                        # lists of types.Datum
    fwd_bwd = training_client.forward_backward(batch, "flow_matching")
    optim = training_client.optim_step(types.AdamParams(learning_rate=1e-4))
    if step % 100 == 0:
        print(step, fwd_bwd.result().metrics["loss"])

sampling_client = training_client.save_weights_and_get_sampling_client("step-1000")
actions = sampling_client.sample(observation).result()
training_client.close()
```

## Sessions

`create_lora_training_client` and `create_full_training_client` start a session: a GPU that loads the model, which takes about a minute (longer the first time). They return once it's ready.

```python theme={null}
service_client.create_lora_training_client("pi05", rank=32, targets=None, alpha=None, action_horizon=10)
service_client.create_full_training_client("pi05", targets=None, action_horizon=10)
```

`targets` says where training goes; by default the action expert for LoRA, the whole model in full. See [Where training goes](https://docs.relayintelligence.org/targets.md). `action_horizon` is the length of the action chunks you train on.

A session holds its GPU until you call `close()` or leave a `with` block, or until 10 minutes pass without a connection. What you haven't saved is lost.

```python theme={null}
with service_client.create_lora_training_client("pi05") as training_client:
    ...
```

`training_client.info` has the model's sizes: `{"model": "pi05", "action_horizon": 10, "action_dim": 32, "trainable_params": 16026656}`.

## Data

A batch is a list of `types.Datum`: an observation in OpenPI's format, as [`relay.Policy`](https://docs.relayintelligence.org/inference) takes, and the action chunk to learn for it.

```python theme={null}
datum = types.Datum(
    observation={
        "image": {"base_0_rgb": front, "left_wrist_0_rgb": wrist},   # (H, W, 3) uint8
        "state": state,                                               # up to 32 values
        "prompt": "pick up the fork",
    },
    actions=chunk,                                                    # (action_horizon, up to 32)
    weight=1.0,                                                       # in the loss's mean
)
```

* State and actions are as the model sees them (normalized), and are padded with zeros to 32.
* Within a batch, each camera's frames are one size. A camera a datum leaves out is masked.
* Frames are sent losslessly by default. Pass `codec=relay.Codec.JPEG_Q95` to the create call to send less.

## Steps run in order, without waiting

`forward_backward`, `optim_step`, `sample`, `save_state` and `load_state` send their work and return a `Future` at once. The GPU runs calls in the order you make them, so the next step is sent while the GPU runs this one. Call `.result()` to wait for a call and get its value.

> [!WARNING]
>
> An error in a step, such as a malformed datum, raises only from that step's `.result()`. Check a future regularly, as the loop above does every 100 steps.

## The steps

```python theme={null}
fwd_bwd = training_client.forward_backward(batch, "flow_matching")
fwd_bwd.result().metrics               # {"loss": 0.21}
fwd_bwd.result().loss_fn_outputs       # [{"loss": 0.19}, {"loss": 0.24}, ...], per datum
```

`forward_backward` computes the flow-matching loss (the weighted mean squared error of the predicted velocity) and adds its gradients to those accumulated so far. Call it several times before an `optim_step` to train on a batch too big for one call.

```python theme={null}
optim = training_client.optim_step(types.AdamParams(learning_rate=1e-4))
optim.result().metrics                 # {"grad_norm": 0.46, "skipped": False}
```

`optim_step` takes an AdamW step with the accumulated gradients, then clears them. You pass `AdamParams` every step, so any schedule is a line of Python:

| `AdamParams` | Default |
| - | - |
| `learning_rate` | required |
| `beta1`, `beta2` | 0.9, 0.95 |
| `eps` | 1e-8 |
| `weight_decay` | 0.0 |
| `grad_clip_norm` | 1.0 |

A step whose gradient isn't finite is skipped (`"skipped": True`).

## Custom losses

`forward_backward_custom` computes your loss on your machine, in PyTorch:

```python theme={null}
import torch

def loss_fn(data, outputs):
    error = (outputs["velocity"] - outputs["target"]).square()       # (B, action_horizon, 32)
    weights = torch.linspace(2, 1, error.shape[1])                    # near-term actions count more
    loss = (error.mean(-1) * weights).mean()
    return loss, {"weighted": loss.item()}

fwd_bwd = training_client.forward_backward_custom(batch, loss_fn)
```

The GPU runs the forward pass and sends back `velocity` (which requires grad), `target`, `x_t`, `noise` and `time`. Your function returns the loss and any metrics; the gradient of the loss with respect to `velocity` goes back to the GPU, which backpropagates it through the model. It waits for the forward pass, so that calls stay in order, and needs `pip install "relay-intelligence[train]"`.

`forward(batch)` returns the same outputs without gradients.

## Sampling

```python theme={null}
actions = training_client.sample(observation).result()               # (action_horizon, 32), with the weights now
actions = training_client.sample([obs_a, obs_b], num_steps=10).result()   # (2, action_horizon, 32)
```

`save_weights_and_get_sampling_client(name)` saves the weights and returns a client that samples with them, however training goes on, for evaluating a checkpoint while you keep training:

```python theme={null}
sampling_client = training_client.save_weights_and_get_sampling_client("step-1000")
sampling_client.sample(observation).result()
sampling_client.path        # the saved weights, adapters merged in, to serve
```

## Saving and resuming

```python theme={null}
training_client.save_state("step-5000").result()     # the trainable weights and the optimizer

# later, in a new session with the same model and targets
training_client = service_client.create_lora_training_client("pi05", rank=32)
training_client.load_state("step-5000").result()
```

States are kept under your account. A name can't contain a slash or start with a dot.


This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.