DRO¶

DRO (Richemond et al., 2024; Kimi Team, 2025) is an off-policy (and even offline) RL objective. Instead of clipping the probability ratio, it adds a quadratic penalty that keeps the learner close to the sampler:

[ \mathcal{L}{\text{DRO}}(\theta) = \mathbb{E}\right)^2\right] ]}\left[\log p_\theta(x) \cdot A(x) - \frac{1}{2}\beta \left(\log \frac{p_\theta(x)}{q(x)

Note that DRO expects a soft advantage formulation, which you compute client-side before building the Datum.

Equivalent code:

quadratic_term = (target_logprobs - sampling_logprobs) ** 2
dro_objective = target_logprobs * advantages - 0.5 * beta * quadratic_term
loss = -dro_objective.sum()

Inputs

target_tokens: array[(N,), int] — IDs sampled by (q).
logprobs: array[(N,), float] — sampling_logprobs.
advantages: array[(N,), float] — per-token advantage (soft formulation).

Outputs

logprobs — target_logprobs.

Diagnostics

loss:sum — total DRO loss.

Custom beta¶

Override the divergence penalty strength:

fb = training_client.forward_backward(
    data=data,
    loss_fn="dro",
    loss_fn_config={"beta": 0.05},
)
result = fb.result()