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

# Router replay

> Replay the sampler's expert choices in the trainer, so RL on mixture-of-experts models trains each token through the experts that generated it.

Router replay, also called [rollout routing replay (R3)](https://arxiv.org/abs/2510.11370), makes the trainer send each sampled token through the same experts the sampler chose for it when you train a mixture-of-experts (MoE) model. To use it, record your samples and attach each training token's source to its [`Datum`](https://docs.baseten.co/reference/sdk/loops/types).

## Record samples and attach provenance

To record routes and attach them to a training example, create the sampling client with `record_stability_info=True`, sample with `include_routed_experts=True`, and describe the datum's tokens with `with_provenance()`:

```python theme={"system"}
from baseten.loops import Datum, PromptProvenanceSpan, SampledProvenanceSpan

sampling_client = training_client.save_weights_and_get_sampling_client(
    record_stability_info=True,
)
response = await sampling_client.sample_async(
    prompt=prompt, num_samples=1, sampling_params=params,
    include_routed_experts=True,
)
seq = response.sequences[0]

p, s = prompt.length, len(seq.tokens)
datum = Datum(model_input=..., loss_fn_inputs=...).with_provenance(
    # where each input token came from
    model_input=[
        PromptProvenanceSpan(sequence_id=seq.sequence_id, length=p),
        SampledProvenanceSpan(sequence_id=seq.sequence_id, length=s - 1),
    ],
    # where each target token came from (shifted by one)
    loss_fn_inputs=[
        PromptProvenanceSpan(sequence_id=seq.sequence_id, offset=1, length=p - 1),
        SampledProvenanceSpan(sequence_id=seq.sequence_id, length=s),
    ],
)
```

* `record_stability_info=True` records every sample from the client. Each returned [`SampledSequence`](https://docs.baseten.co/reference/sdk/loops/types) carries a `sequence_id` that refers to its record. The record and its routes stay on Baseten, and the datum carries only the references. Create the client from the training client, as in the snippet, so the run's own sampler records the routes that its trainer replays.
* `include_routed_experts=True` captures the expert routes with each record. It requires a client created with `record_stability_info=True`.
* `with_provenance()` returns a copy of the datum with two span lists. The `model_input` spans describe the input tokens in order, and the `loss_fn_inputs` spans describe `target_tokens` in order. Each span names a sample's `sequence_id`, whether its tokens come from the sample's prompt (`PromptProvenanceSpan`) or its completion (`SampledProvenanceSpan`), the position of its first token in that part (`offset`, `0` by default), and the number of tokens (`length`).

With the [Tinker compatibility package](https://docs.baseten.co/loops/tinker-compatibility.md), import `PromptProvenanceSpan` and `SampledProvenanceSpan` from `tinker.types`.

## Requirements and limitations

* Router replay works in both [availability modes](https://docs.baseten.co/loops/availability-modes/overview.md) on MoE models whose sampler captures routes. On other models, a `sample()` call with `include_routed_experts=True` fails.
* Recording takes tokenized text prompts. A recording client rejects prompts that contain images.
* `dpo` doesn't support router replay. A `forward_backward()` call with `loss_fn="dpo"` fails when its datums carry provenance.

## Routing mismatch between the sampler and trainer

Each MoE layer has a router that sends every token to a few of the layer's experts. The sampler and the trainer both run the router, but with different kernels and numerics, so for some tokens they choose different experts. The trainer then computes those tokens' log-probabilities and gradients through experts that didn't generate them.

These routing differences widen the [KL divergence between the sampler and trainer](https://docs.baseten.co/loops/rl-run-health.md#sampler-and-trainer-kl-divergence) even when both run the same weights. The RL losses read the gap as a change in the policy, and in MoE models the mismatch can grow until training becomes unstable. R3 removes this source of mismatch: the sampler records the experts it chose for each token, and the trainer reuses them in its forward and backward passes. The trainer still computes the router's gating weights from its own logits, so the router keeps training.

## Next steps

* [Check whether an RL run is learning](https://docs.baseten.co/loops/rl-run-health.md#sampler-and-trainer-kl-divergence): Monitor the KL divergence between the sampler and trainer.
* [Write an RL loop with the Loops SDK](https://docs.baseten.co/examples/rl-sdk-loop.md): Sample, score, and train in your own loop.
* [SamplingClient](https://docs.baseten.co/reference/sdk/loops/sampling-client.md): The `sample()` options for route capture.