Training loops

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.

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.

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

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 takes, and the action chunk to learn for it.

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.

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

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.

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:

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

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:

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

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.

Get started

Inference

Training

Reference

Training loops

Sessions

Data

Steps run in order, without waiting

The steps

Custom losses

Sampling

Saving and resuming