# StableReinforce

> **Alpha**
>
> StableReinforce is in alpha, and its API may change.

StableReinforce (`stable_reinforce`) is a REINFORCE variant that stays stable even when the sampling policy differs from the policy being trained, such as in async RL or when training and sampling use different numerics. In our internal post-training, this loss function has led to extremely stable training, even under highly off-policy async RL and low precision numerics. The objective uses minimal additional information captured at sampling time, so you need to sample with `record_stability_info=True` and annotate your training tokens with corresponding `sequence_id`s from sampling.

## Sampling with stability info

Create the sampling client with `record_stability_info=True`. Every sample from that client is recorded, and each returned sequence carries a `sequence_id` that training requests use to name it:

```python
sampling_client = training_client.save_weights_and_get_sampling_client(
    record_stability_info=True
)
response = await sampling_client.sample_async(
    prompt=prompt,
    sampling_params=params,
    num_samples=4,
)
sequence = response.sequences[0]
sequence.sequence_id  # names this sequence in a training request
```

`ServiceClient.create_sampling_client` and `TrainingClient.create_sampling_client` take the same argument.

Recorded sequences belong to the session that sampled them, and can be trained on for one day after sampling. Each training request that uses a sequence restarts that day.

## Training

StableReinforce is one of the features that need [provenance](https://tinker-docs.thinkingmachines.ai/tinker/datum/#provenance): every datum must identify the sampled sequence behind each of its tokens. Beyond the general provenance checks, it requires:

- Both span lists on every datum.
- At least one `SampledProvenanceSpan` in the `loss_fn_inputs` spans, each covering its whole sequence (`offset=0` and the sequence's full length). Spans on the model input may cover part of a sequence.
- Every token with a nonzero advantage under a `SampledProvenanceSpan` whose sequence was sampled with `record_stability_info=True`.

With datums built as in the [provenance example](https://tinker-docs.thinkingmachines.ai/tinker/datum/#provenance):

```python
future = await training_client.forward_backward_async(
    data, loss_fn="stable_reinforce", loss_fn_config={"kl_coeff": 0.01}
)
result = await future.result_async()
```

**Input tensors:**

- `target_tokens: array[(N,), int]` — Target token IDs (from the sampler (q))
- `logprobs: array[(N,), float]` — `sampling_logprobs` for the tokens
- `advantages: array[(N,), float]` — Advantage values for RL

**Output tensors:**

- `logprobs: array[(N,), float]` — `target_logprobs` for the tokens

**Output diagnostics:**

- `loss:sum` (scalar) — Summed loss, for monitoring

## KL penalty

The loss includes a penalty on the divergence between the sampler (q) and the learner (p\_\\theta). It is estimated with the k1 estimator on each sampled token (x_t), the tokens under a `SampledProvenanceSpan`:

\[ \\mathcal{L}_{\\text{KL}}(\\theta) = \\beta \\sum_t \\left(\\log q(x_t) - \\log p_\\theta(x_t)\\right) \]

This is a sample estimate of (\\mathrm{KL}(q ,|, p\_\\theta)) scaled by (\\beta). Its gradient, (-\\beta \\sum_t \\nabla\_\\theta \\log p\_\\theta(x_t)), raises the learner's probability of every sampled token regardless of its advantage, keeping the learner close to the sampler.

Set (\\beta) through `loss_fn_config`:

- `kl_coeff` — (\\beta), a finite float ≥ 0 (default `0.01`); `0` turns the penalty off