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

# Write an RL loop with the Loops SDK

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

An RL loop on Loops runs on the client. Your script drives it: roll out completions from the paired sampler, score them into advantages, accumulate gradients with [`forward_backward()`](https://docs.baseten.co/reference/sdk/loops/training-client), and apply them with [`optim_step()`](https://docs.baseten.co/reference/sdk/loops/training-client). For how the trainer and sampler interact, see [Loops concepts](https://docs.baseten.co/loops/concepts.md); for the per-loss data contract, see [Loss functions](https://docs.baseten.co/loops/loss-functions.md).

## Prerequisites

* A [workspace API key](https://docs.baseten.co/organization/api-keys.md), set as `BASETEN_API_KEY`.
* Python 3.12+ and [uv](https://docs.astral.sh/uv/) to install the SDK and run the script.

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

The SDK creates the run and picks the availability mode from your workspace access and the model. A serverless run executes on Baseten hardware and bills for the tokens it samples and trains. A dedicated run bills for GPU time until you deactivate it. For the billing rules, see [Availability modes](https://docs.baseten.co/loops/availability-modes/overview.md).

An RL run samples on every step. Set `with_sampler=True` to create the paired sampler with the trainer. For dedicated runs, this overlaps GPU startup instead of delaying sampler provisioning until the first rollout. Save this code as `rl_loop.py`:

```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
)
sampling_client = training_client.save_weights_and_get_sampling_client()
```

## Roll out completions

Fan out one group of completions per prompt. This is group relative policy optimization (GRPO): one prompt, many completions, with 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)
        ]
    )


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 deployments with sticky routing enabled, `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`. To use it, replace the preceding `rollout()` function so each group's requests share one key, which keeps the group's prompt in one replica's prefix cache; don't reuse one key across a whole run, or every group lands on the same replica:

```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)
        ]
    )
```

Weight selection follows the sampling client's checkpoint or policy-version settings. On serverless runs, the sampling client pins the checkpoint saved by `save_weights_and_get_sampling_client()`. Publishing later weights doesn't change an existing client's checkpoint. Before the first optimizer step, the client samples the base model.

On dedicated runs, the sampling client carries a minimum policy version set by `save_weights_and_get_sampling_client()`. Sampling waits until the sampler serves at least that version, then uses the live adapter, which can advance after later publishes. A replica below the floor returns `408`, and the SDK retries with backoff until the replica serves the requested version or the client's ready timeout expires.

## 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]
print(f"mean reward {mean:.2f} over {len(scores)} completions")
```

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

```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
    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 when `skip_token_normalization` is `false`, 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` that sampling returned: 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. For more information, see [Loss functions](https://docs.baseten.co/loops/loss-functions.md). 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

Call [`save_weights_and_get_sampling_client()`](https://docs.baseten.co/reference/sdk/loops/training-client) after `optim_step()` and use its returned client for the next rollout. The method returns a `SamplingClient` directly:

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

On serverless runs, each call creates a checkpoint with a unique name and returns a client pinned to it. Use the client's `checkpoint_path` and `checkpoint_name` to identify that save rather than reconstructing them from `"default"`, and reacquire a client after the next training update.

[`save_weights_for_sampler()`](https://docs.baseten.co/reference/sdk/loops/training-client) saves weights without updating the paired sampler's live policy. Its future resolves to a version and path.

You can submit rollout requests while training is in progress. Those rollouts can use weights from before the latest optimizer step, 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()`](https://docs.baseten.co/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.

## Monitor divergence between the sampler and trainer

RL on Loops runs the policy in two places: the sampler rolls out and the trainer learns. Both should assign the same log-probability to every sampled token, so in a healthy run the Kullback-Leibler (KL) divergence between them stays near zero. A rising KL divergence means the sampler is serving stale weights or its numerics differ from the trainer's.

The SDK ships [`loops_log_kl_sample_train()`](https://docs.baseten.co/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 with [`save_weights_and_get_sampling_client()`](https://docs.baseten.co/reference/sdk/loops/training-client), or persist a restorable checkpoint with [`save_state("step-1")`](https://docs.baseten.co/reference/sdk/loops/training-client). In the dedicated mode, the paired sampler picks up published weights without restarting; for how that sync works, see [How weight sync works](https://docs.baseten.co/loops/concepts.md#weight-sync). To resume a run from a checkpoint, see [Checkpoints](https://docs.baseten.co/loops/concepts.md#checkpoints).

## Run the script

Run the loop:

```bash theme={"system"}
uv run python rl_loop.py
```

```text Output theme={"system"}
mean reward <MEAN_REWARD> over <N> completions
```

A mean reward of `0.0` means every completion in the group scored the same, so every advantage is zero and the optimizer step changes nothing. Raise the group size or adjust the task until the scores differ.

## Shut down the run

After saving the checkpoints you need, close the local SDK clients:

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

The `Qwen/Qwen3-8B` model in this example uses dedicated capacity. Closing the clients doesn't deactivate the run. To stop GPU billing, use the run ID from `training_client.run_id`. With the [Baseten CLI](https://docs.baseten.co/reference/cli/baseten/overview.md#install) installed, end the run:

1. Deactivate the run, and confirm when prompted:

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

2. Check that the run's `status.name` is `INACTIVE`:

   ```bash theme={"system"}
   baseten loops run describe --run-id "<run_id>" --output json
   ```

Saved checkpoints remain available. For the REST API option and serverless behavior if you change the model, see [Deactivate a run](https://docs.baseten.co/loops/deactivate-a-run.md).

## Deploy the checkpoint

The run's sampler checkpoint is ready to serve. The deployment downloads the base weights with the `hf_access_token` in [workspace secrets](https://docs.baseten.co/organization/secrets.md) and bills for its GPU while it's live. With the [Baseten CLI](https://docs.baseten.co/reference/cli/baseten/overview.md#install) installed, list the run's checkpoints and deploy the one you want:

```bash theme={"system"}
baseten loops checkpoint list --run-id "<run_id>"
baseten loops checkpoint deploy --checkpoint-id "<checkpoint_id>"
```

The deploy command prompts for a model name, GPU type, GPU count, and the Hugging Face secret name. [Deploy a checkpoint](https://docs.baseten.co/loops/deploy-checkpoints.md) covers the REST API and dashboard paths.

## Next steps

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