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_ids 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:

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

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