Build your own RL system

Use River's core sampling and training methods when your experiment needs a custom rollout schedule, advantage estimator, data mixture, or update rule. You choose what to sample, how to turn it into training data, and when to update the model. River executes the forward, backward, and optimizer operations.

This chapter assumes you understand the RL loop. It focuses on the boundary between your algorithm and the training service. The snippets operate inside an existing training session, with model referring to a live River model.

Choose the controls your algorithm needs

Choose a built-in loss and configure its parameters. You can implement new sampling, credit-assignment, and scheduling techniques around those losses. The API does not accept arbitrary Python loss functions, custom autograd code, or arbitrary optimizer implementations. A technique that requires a different server-side derivative needs an additional loss implementation.

If only the reward or environment changes, the RL library already supplies the surrounding loop. Moving to core primitives means your code also owns trajectory bookkeeping, normalization, policy provenance, and recovery.

Sample exactly the context you will train on

For a text-only experiment, render each conversation with the model's chat format and tokenize it once. Let prompt_ids be the resulting nonempty token lists, one per task. Pass those IDs directly to sampling:

groups = model.sample(
    prompt_token_ids=prompt_ids,
    num_samples=4,
    max_tokens=1024,
    temperature=1.0,
    top_p=1.0,
    top_k=-1,
    seed=0,
)

The result is indexed by prompt, then sample. Each Sample carries generated tokens, chosen-token logprobs, stop_reason, and policy_version. Chosen-token log probabilities are returned without setting logprobs=K; that option asks for alternative top-K tokens and adds work you usually do not need for an RL update.

These settings avoid temperature rescaling and top-p/top-k filtering in the initial experiment. When changing the sampling distribution, verify how the sampler reports log probabilities and how your objective uses them. Supplying old_logprobs alone does not make every sampling transformation an exact importance correction.

Keep the prompt IDs, generated IDs, original log probabilities, policy version, and generation settings with each rollout. Score readable text if appropriate, but construct training data from the recorded tokens. Decoding and retokenizing a response can change the sequence. A Sample with token_data_is_exact=False is unsuitable for an update that relies on exact rollout alignment.

Use stop_reason to apply your truncation policy before constructing a batch. Dropping incomplete attempts, assigning them zero reward, and training on their partial work produce different datasets. For group-relative estimators, decide whether excluded attempts still participate in the baseline.

Supply your own per-token signal

Your estimator produces an advantage for each generated token. A trajectory-level estimator can broadcast one scalar across the response; a process reward or credit-assignment method can assign different values to different spans. Preserve the grouping needed by your estimator even if results arrive out of order.

The training service consumes fields aligned to prediction positions. If a prompt has P tokens and a response has C tokens, the first response token is predicted at position P − 1. The last input position has no next-token target.

A zero advantage removes that position's direct RL loss contribution. It does not remove the token from the context seen by later predictions. The attention mask is not the loss mask.

This helper builds one text-only datum. token_advantages is your algorithm's vector for this sample, with one value per generated token:

def make_datum(prompt, sample, token_advantages):
    if not prompt:
        raise ValueError("A training prompt must contain at least one token")
    if not sample.token_data_is_exact:
        raise ValueError("Training requires exact sampled token data")
    if len(sample.tokens) != len(sample.logprobs):
        raise ValueError("Sampled tokens and log probabilities must align")
    if len(token_advantages) != len(sample.tokens):
        raise ValueError("Supply one advantage per generated token")

    ids = prompt + sample.tokens
    return {
        "input_ids": ids,
        "attention_mask": [1] * len(ids),
        "old_logprobs": [0.0] * (len(prompt) - 1) + sample.logprobs + [0.0],
        "advantages": [0.0] * (len(prompt) - 1) + token_advantages + [0.0],
    }

target_tokens is omitted: the server derives next-token targets from input_ids. For a scalar advantage a, pass [a] * len(sample.tokens). For selective credit, zero the generated positions you intend to exclude.

Now build batch by applying this helper to the samples your algorithm retains. Keep record metadata—task identity, policy version, reward components, and completion status—alongside the batch. Those records let you audit the estimator and reconstruct why a sample contributed to an update.

The concatenation above covers a single text response. In a multi-turn rollout, keep model-generated spans separate from tool and user messages. Train each span against the context that produced it. For images, preserve image chunks and their alignment as well as text tokens. The RL library's Trajectory represents these records; replacing the optimizer loop does not require replacing that representation.

Make batch weighting explicit

The core losses sum token contributions. Neither a request boundary nor an optimizer call implicitly turns that sum into the average your algorithm may intend. Decide the denominator across the whole logical update.

Let N be the number of response tokens included in your normalization, B the number of retained responses, and Cᵢ the included token count for response i:

Define which tokens count in these denominators. For example, removing zero-advantage groups and then dividing by the remaining token count changes update scale compared with keeping them in the denominator. A sequence with zero included tokens has no term in a per-sequence mean.

You can incorporate these factors into advantages before submission. For a single global factor, optim_step(gradient_scale=...) scales accumulated gradients before gradient clipping and Adam. It requires endpoint support for gradient_scale_v1; the client exposes that through model.get_server_capabilities().require("gradient_scale_v1").

Accumulate microbatches, then update once

Separate the logical training batch from its transport batches. The following function takes nonempty batch data from make_datum. normalization_tokens is the positive denominator selected by your algorithm; advantages have not already been divided by it.

def apply_update(model, batch, normalization_tokens, lr):
    if not batch or normalization_tokens <= 0:
        raise ValueError("Skip updates with no training contribution")
    model.get_server_capabilities().require("gradient_scale_v1")
    policy = model.get_policy_version()
    if policy is None:
        raise RuntimeError("This loop requires committed policy versions")

    microbatch_size = 8
    for start in range(0, len(batch), microbatch_size):
        model.forward_backward(
            batch[start : start + microbatch_size],
            loss_fn="cispo",
            eps_max=6.0,
            zero_out=(start == 0),
            expected_policy_id=policy.id,
        )

    return model.optim_step(
        lr=lr,
        beta1=0.9,
        beta2=0.95,
        eps=1e-8,
        weight_decay=0.0,
        grad_clip_norm=1.0,
        gradient_scale=1.0 / normalization_tokens,
        expected_policy_id=policy.id,
    )

zero_out=True clears existing gradients on the first microbatch. Later calls use False to accumulate. One optimizer step applies the aggregate. Using True on every call discards earlier contributions; taking an optimizer step per microbatch changes the weights between contributions and implements a different algorithm. The learning-rate schedule should advance on the logical updates your algorithm defines.

This example waits for each forward/backward result before proceeding. If one fails, it does not submit an optimizer step. Resolve the failed operation before restarting the logical batch, and reset gradients on that batch's first call.

model.train_step(...) is useful when you want a complete update in one call. It clears gradients and pipelines forward/backward with the optimizer operation. Use the separate methods when you need accumulation, inspection, or a decision between those operations. For requests within the normal upload limit, train_step submits the optimizer before the backward result is known; its failure behavior is therefore different from the loop above.

Change the algorithm at a specific boundary

Here are experiments you can implement without changing the training service:

• Curriculum or adaptive group size: change task selection and samples per task before sampling. Record the selection rule so evaluation still measures the intended task distribution.
• Reward shaping or a new baseline: change your estimator before creating token_advantages. Preserve reward components and raw scores so you can separate estimator effects from verifier changes.
• Process credit or span selection: change the per-token advantages and masks. Keep the original context and log probabilities for the included spans.
• Replay or multiple updates per rollout batch: reuse recorded trajectories while preserving their original old_logprobs. Re-read the current training policy before each update and enforce your policy-age rule. Replacing old log probabilities with current ones would change the estimator.
• SFT/RL mixtures: accumulate an SFT contribution with cross_entropy, then an RL contribution with zero_out=False, followed by one optimizer step. Set each objective's normalization and mixture coefficient explicitly through weights or advantages. Two separate optimizer steps are not the same as one step on the combined gradient.

model.forward(...) evaluates a supported loss without accumulating gradients. Use it for objective diagnostics on a fixed batch. Per-token log probabilities are available when the worker returns them; check the result before depending on them. A forward call is not an interface for sending a custom backward pass.

Schedule sampling and training independently

Start with one owner of the model's training sequence. submit_* methods let that owner submit work and collect results later without having several drivers race to update the same weights.

The independent sampling API accepts up to 128 samples per submission, including the expansion from num_samples. Its completion iterator is asynchronous:

async def collect_samples(model, prompt_ids):
    pending = model.submit_sampling_batch(
        prompt_token_ids=prompt_ids,
        num_samples=4,
        max_tokens=1024,
        temperature=1.0,
        return_prompt_token_ids=True,
    )
    async for result in pending.as_completed():
        if result.error is not None:
            raise result.error
        yield result.prompt_index, result.sample_index, result.sample

The caller supplies at most 32 prompts for this four-sample example. Feed results into your own group assembler using the original prompt and sample indexes; completion order is not group order. A ready sample is not necessarily a ready training group. Your estimator determines when the group has enough information. Do not turn infrastructure failures into zero task rewards by accident.

The iterator reports results independently; it does not define a training schedule. Cancelling a local wait does not cancel accepted server operations. Keep request IDs and resolve outstanding operations before replaying uncertain work or recovering a driver.

For microbatches submitted asynchronously, keep the first gradient reset ahead of every accumulation call and collect all required backward results before submitting an update that must depend on their success. The nonblocking submit_forward_backward rejects requests above its upload limit; use blocking forward_backward for the client's automatic splitting at datum boundaries.

Select and track the sampling policy

Low-level asynchronous sampling requires an explicit weight-version strategy:

These options belong to submit_sampling_batch and require server support. A latest snapshot may lag the committed training head. Publishing snapshots, bounding their age, and deciding when to wait or drop data are parts of your algorithm, not consequences of choosing a nonblocking call.

Use Sample.policy_version as the behavior-policy identity and OptimStepResult.policy_version as the committed result of an update. model.get_policy_version() reads the committed training policy. model.step advances at optimizer submission time, even if the operation later fails; it is not confirmation that the update committed.

Before training, check that rollout versions belong to the intended lineage and meet your age bound. expected_policy_id guards the weights used for a training operation; it does not enforce a staleness rule over your data. Record both the behavior policy and training policy for each contribution.

For mixed-policy trajectories or retained KV, track versions per generated span and distinguish sampling-policy age from cache age. The policy chapter details those choices. Training and sampling MoE models can also involve routing differences: the API exposes return_expert_routing, Sample.routing_datum_keys(required=True), compute_expert_flip_metric, and force_routing_replay for supported models. Use these controls when routing consistency is part of the experiment, and keep each capture tied to its sampled sequence.

Make the custom loop inspectable

For each logical update, retain the tasks and sampling settings, original rollout records, rewards, selected spans, normalization denominator, loss configuration, optimizer settings, and committed policy identity. This lets you compare algorithms on the same data and diagnose an unexpected gradient or change in reward.

Test data construction and update scheduling on a tiny recorded batch before scaling out. Verify that changing only the microbatch partition preserves the intended weighting, that prompt and tool tokens receive no direct RL loss, and that failed backward work cannot accidentally trigger a partial update.

A weight checkpoint alone does not save your custom scheduler, replay buffer, data position, or outstanding operations. Persist those alongside the run if you need recovery. Use Evaluation and recovery for the library-managed alternative, and compare saved policies on a fixed holdout before attributing gains to your new technique.

For a worked application of these controls, Distill a larger model uses teacher probabilities to construct token-level updates, including self-distillation and multi-teacher training.