Write an RL loop with the Loops SDK

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(), and apply them with optim_step(). This page walks through one full pass. See Loops concepts for how the trainer and sampler interact underneath, and Loss functions for the per-loss data contract.

Prerequisites

A workspace API key, set as BASETEN_API_KEY.
Python 3.12+ and 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:

uv init loops-rl
cd loops-rl
uv add baseten-loops

Provision the run

The SDK creates the run and runs it in the serverless or dedicated mode according to your workspace access and model configuration. 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:

rl_loop.py

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

See Availability modes for how a run’s mode affects billing.

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:

rl_loop.py

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 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. 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 selection follows the sampling client’s checkpoint or policy-version settings:

rl_loop.py

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

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.

rl_loop.py

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

rl_loop.py

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:

rl_loop.py

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. 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() after optim_step() and use its returned client for the next rollout. The method returns a SamplingClient directly:

rl_loop.py

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". Keep using the returned client for that rollout group, and reacquire a client after the next training update. On dedicated runs, the returned client sets a minimum policy version rather than pinning an exact checkpoint. save_weights_for_sampler() saves weights without updating the paired sampler’s live policy. Its future resolves to a version and path. An existing serverless sampling client keeps its checkpoint, and an existing sampling client in the dedicated mode keeps its original minimum version. 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() 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 KL 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 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(), 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(), or persist a restorable checkpoint with save_state("step-1"), which takes a name for the checkpoint. A serverless sampling client selects the saved checkpoint; in the dedicated mode, the paired sampler picks up published weights without restarting. See How weight sync works. To resume a run from a checkpoint, see Checkpoints.

Shut down the run

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

rl_loop.py

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 installed, end the run:

Deactivate the run, and confirm when prompted:

baseten loops run deactivate --run-id "<run_id>"

Check that the run’s status.name is INACTIVE:

baseten loops run describe --run-id "<run_id>" --output json

Saved checkpoints remain available. See Deactivate a run for the REST API option and serverless behavior if you change the model.

Next steps

Loss functions for the full per-loss data contract and loss_fn_config defaults.
TrainingClient reference for forward_backward(), forward(), and optim_step() signatures.
SamplingClient reference for sample(), compute_logprobs(), and affinity_key.

Examples

Post-training

Models

Engines

Custom Docker servers

Custom Python models

Chains

Baseten for Model Labs

Write an RL loop with the Loops SDK

Prerequisites

Install

Provision the run

Roll out completions

Score rewards into advantages

Step the trainer

Publish weights for the sampler

Evaluate without training

Monitor KL divergence between the sampler and trainer

Checkpoint mid-run

Shut down the run

Next steps

​Prerequisites

​Install

​Provision the run

​Roll out completions

​Score rewards into advantages

​Step the trainer

​Publish weights for the sampler

​Evaluate without training

​Monitor KL divergence between the sampler and trainer

​Checkpoint mid-run

​Shut down the run

​Next steps

Prerequisites

Install

Provision the run

Roll out completions

Score rewards into advantages

Step the trainer

Publish weights for the sampler

Evaluate without training

Monitor KL divergence between the sampler and trainer

Checkpoint mid-run

Shut down the run

Next steps