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

# RL on Loops

> Build a reinforcement learning loop on Loops: roll out completions from the paired sampler, score them into advantages, and step the trainer.

RL on Loops is a client-side loop. Your script drives it: roll out completions from the paired sampler, score them into advantages, accumulate gradients with [`forward_backward()`](/reference/sdk/loops/training-client), and apply them with [`optim_step()`](/reference/sdk/loops/training-client). This page walks through one full pass. See [Loops concepts](/loops/concepts) for how the trainer and sampler interact underneath, and [Loss functions](/loops/loss-functions) for the per-loss data contract.

## Prerequisites

* **API key**: A [workspace API key](/organization/api-keys) with org access to Loops. Follow [Set up your API key and SDK](/quickstart#set-up-your-api-key-and-sdk) and export the key as `BASETEN_API_KEY`.
* **Python 3.12+ and [uv](https://docs.astral.sh/uv/)**: This guide uses `uv` to install the Loops client and run the RL script.
* **Baseten CLI**: Install the [Baseten CLI](/reference/cli/baseten/overview) for the `baseten loops` commands in [Shut down the run](#shut-down-the-run).

<Note>
  Loops is in early access. To enable it for your workspace, [fill out the signup form](https://www.baseten.co/talk-to-us/loops-signup/).
</Note>

## Install

Install `baseten-loops` into a uv project. Create one first if you don't have it:

```bash theme={"system"}
uv init loops-rl
cd loops-rl
uv add baseten-loops
```

## Provision the run

An RL run samples on every step, so boot the sampler alongside the trainer with `with_sampler=True` instead of paying its cold start on the first rollout:

```python rl_loop.py theme={"system"}
from baseten.loops import ServiceClient

training_client = ServiceClient().create_lora_training_client(
    base_model="Qwen/Qwen3-8B",
    with_sampler=True,
    name="rl-example",  # without a name, the run shows up as the base model ID
)
# Publish the current weights and get a client that follows them.
sampling_client = training_client.save_weights_and_get_sampling_client()
```

See [Overlap trainer and sampler startup](/loops/concepts#overlap-trainer-and-sampler-startup). The trainer and sampler bill for their GPUs while they run, so [shut the run down](#shut-down-the-run) when you finish.

## Roll out completions

Fan out one group of completions per prompt. A GRPO-style group means one prompt, many completions, advantages computed relative to the group:

```python rl_loop.py theme={"system"}
import asyncio

from baseten.loops import ModelInput, SamplingParams

tokenizer = training_client.get_tokenizer()
prompt_ids = tokenizer.encode("What is 17 times 24?", add_special_tokens=False)
group_size = 4

prompt = ModelInput.from_ints(prompt_ids)

async def rollout():
    return await asyncio.gather(*[
        sampling_client.sample_async(
            prompt,
            sampling_params=SamplingParams(temperature=1.0, top_p=0.95),
        )
        for _ in range(group_size)
    ])

# asyncio.run() drives the fan-out from a plain script.
responses = asyncio.run(rollout())
```

Each response carries `sequences`, each with the sampled `tokens`, per-token `logprobs`, and `stop_reason`. Pass `include_prompt_logprobs=True` when you need the prompt's per-position logprobs too.

On sticky-enabled deployments, `affinity_key` pins every request that shares a key to one sampler replica. It is a keyword argument on `sample()` and `sample_async()`, not a field on `SamplingParams`. Set it to one key per rollout group so the group's shared prompt stays in one replica's prefix cache; don't reuse one key across a whole run, or every group lands on the same replica. Weight consistency across the group comes from the version floor, not from `affinity_key`:

```python rl_loop.py theme={"system"}
async def rollout():
    return await asyncio.gather(*[
        sampling_client.sample_async(
            prompt,
            sampling_params=SamplingParams(temperature=1.0, top_p=0.95),
            affinity_key="rollout-group-0",
        )
        for _ in range(group_size)
    ])
```

The sampling client carries the policy-version floor set by `save_weights_and_get_sampling_client()`: `sample()` blocks until the sampler serves at least the version you published, then serves the live adapter, so rollouts come from the policy you trained. When a replica is below that floor, it returns `408`, and the SDK retries with backoff until the replica serves the requested version (up to the client's ready timeout, one hour by default).

## Score rewards into advantages

Score each completion, then normalize within its group: subtract the group mean, and divide by the group's standard deviation for unit-variance advantages. When every completion in a group earns the same reward the standard deviation is zero, so guard the division; those advantages are all `0.0`, which makes the group a no-op for the policy losses rather than an error.

```python rl_loop.py theme={"system"}
sequences = [response.sequences[0] for response in responses]
scores = [1.0 if "408" in tokenizer.decode(seq.tokens) else 0.0 for seq in sequences]

mean = sum(scores) / len(scores)
std = (sum((score - mean) ** 2 for score in scores) / len(scores)) ** 0.5
advantages = [(score - mean) / std if std else 0.0 for score in scores]
```

Each completion becomes one [`Datum`](/reference/sdk/loops/types) carrying the full sequence and two per-position tensors: the rollout logprobs (the behavior policy) and the advantages, tiled across the completion's token positions. Prompt positions carry `0.0` in both tensors.

```python rl_loop.py theme={"system"}
from baseten.loops import Datum, TensorData

prompt_len = len(prompt_ids)

def build_datum(sequence, advantage):
    tokens = prompt_ids + sequence.tokens  # prompt + completion
    L = len(tokens)
    return Datum(
        model_input=ModelInput.from_ints(tokens),
        loss_fn_inputs={
            "logprobs": TensorData(
                data=[0.0] * prompt_len + sequence.logprobs, dtype="float32", shape=[L]
            ),
            "advantages": TensorData(
                data=[0.0] * prompt_len + [advantage] * len(sequence.tokens),
                dtype="float32",
                shape=[L],
            ),
        },
    )

datums = [
    build_datum(sequence, advantage)
    for sequence, advantage in zip(sequences, advantages)
]
```

Advantages are per-position and must match the sequence length in every `Datum` in the batch. They also serve as the active-token mask for the policy losses: a position with advantage `0.0` contributes nothing to the loss or to the denominator, so a completion whose group-normalized advantage lands at zero drops out of the update.

When the rollout samples at a temperature other than `1.0`, set it on `SamplingParams`. There is no matching parameter on the trainer, so the sampler and the trainer can score the same token under different distributions; the divergence metrics in the next section show the drift.

## Step the trainer

Accumulate gradients, then apply them:

```python rl_loop.py theme={"system"}
from baseten.loops import AdamParams

fb = training_client.forward_backward(
    data=datums,
    loss_fn="ppo",  # or "cispo", "importance_sampling", "dppo", "dro"
).result(timeout=600.0)

training_client.optim_step(AdamParams(learning_rate=1e-4)).result(timeout=600.0)
```

The ratio-based losses need the per-position `logprobs` from above: the importance ratio is `exp(new_logprobs - logprobs)`, so the behavior policy each trajectory was sampled from has to travel with the datum. The losses differ in how they clip or penalize that ratio; see [Loss functions](/loops/loss-functions). For supervised fine-tuning, use `cross_entropy` with `target_tokens` instead, which reads neither `logprobs` nor `advantages`.

`forward_backward()` adds to the trainer's gradient buffer, so you can call it several times between `optim_step()` calls to accumulate gradients across microbatches or rollout batches. Pass `top_k_logprobs` (up to `20`) to get the top token log-probabilities per active position back in the result, which is useful for debugging the ratio.

Invalid `loss_fn_config` values fail fast: the server validates the config before running the pass and returns the validation error verbatim.

## Publish weights for the sampler

Rollouts only use a guaranteed-new policy after you publish the weights and reacquire a pinned sampling client. Use [`save_weights_and_get_sampling_client()`](/reference/sdk/loops/training-client) after `optim_step()` when the next rollout must wait for the new version:

```python rl_loop.py theme={"system"}
sampling_client = training_client.save_weights_and_get_sampling_client("default")
```

[`save_weights_for_sampler()`](/reference/sdk/loops/training-client) is the publish-only variant. Its future resolves to the version and path it published, but an existing sampling client keeps its original minimum version.

The sampler and the trainer are separate servers, so the next step's rollouts can run while the current step trains. Rollouts started before an `optim_step()` land are one step behind the trainer's weights, which is why the ratio-based losses take the rollout logprobs: the importance ratio corrects for the gap between the policy that sampled and the policy being trained.

## Evaluate without training

[`forward()`](/reference/sdk/loops/training-client) runs the same loss math with no backward pass and leaves the gradient buffer untouched, so it is safe to interleave with gradient accumulation. Use it to evaluate a held-out batch mid-run without disturbing the gradients you have accumulated.

## Watch sampler and trainer agreement

RL on Loops runs the policy in two places: the sampler rolls out and the trainer learns. They should agree on the log-probability of every sampled token. A healthy run keeps their divergence near zero. Rising values mean the two sides disagree, from stale weights on the sampler or a numerics mismatch between training and inference.

The SDK ships [`loops_log_kl_sample_train()`](/reference/sdk/loops/helpers), which logs a train/sample divergence panel to Weights & Biases: the `kl_sample_train_v1_loops`, `kl_sample_train_v2_loops`, and `kl_sample_train_k3` estimators (Tinker's first- and second-order estimators plus Schulman's k3), importance-ratio and effective-sample-size statistics, per-token difference histograms, and a worst-divergent-tokens table.

## Checkpoint mid-run

Publish the current weights and get a sampling client that serves them with [`save_weights_and_get_sampling_client()`](/reference/sdk/loops/training-client), or persist a restorable checkpoint with [`save_state("step-1")`](/reference/sdk/loops/training-client), which takes a name for the checkpoint. The paired sampler picks up published weights without restarting; see [How weight sync works](/loops/concepts#how-weight-sync-works). To resume a run from a checkpoint, see [Checkpoints](/loops/concepts#checkpoints).

## Shut down the run

The trainer and sampler bill for their GPUs until you deactivate the run, and they keep running after your script exits. Close the client when you finish training, then check what's live and tear it down:

```bash theme={"system"}
baseten loops run list
baseten loops run deactivate --run-id <run_id> --yes
```

[`baseten loops run list`](/reference/cli/baseten/loops-run) lists your active runs, one row per run, with the run ID, base model, and status of each. Pass the run ID to [`deactivate`](/reference/cli/baseten/loops-run), which tears down both the trainer and its paired sampler. Your checkpoints survive the shutdown.

## Next steps

* [Loss functions](/loops/loss-functions) for the full per-loss data contract and `loss_fn_config` defaults.
* [`TrainingClient` reference](/reference/sdk/loops/training-client) for `forward_backward()`, `forward()`, and `optim_step()` signatures.
* [`SamplingClient` reference](/reference/sdk/loops/sampling-client) for `sample()`, `compute_logprobs()`, and `affinity_key`.