# Loss functions

## Built-in losses

Pass the loss name as `loss_fn` to `forward_backward`, with loss-specific
keyword arguments, for example:
`model.forward_backward(data, loss_fn="cispo", eps_max=8.0)`.

> **Cross-entropy and the four RL objectives below sum token losses.** To exclude
> a token, set its `weights` value to `0.0` for `cross_entropy`, or its `advantages`
> value to `0.0` for an RL loss. Mask prompt tokens this way to train only on the
> completion. `opsd` uses a different reduction, described below.

| `loss_fn` | Use | Required per-datum fields | Arguments (default) |
| --- | --- | --- | --- |
| `cross_entropy` | Supervised fine-tuning | `input_ids`, `weights` (`target_tokens` optional) | — |
| `importance_sampling` | Off-policy policy gradient | `input_ids`, `old_logprobs`, `advantages` | — |
| `ppo` | Clipped policy optimization | `input_ids`, `old_logprobs`, `advantages` | `clip_low` (0.2), `clip_high` (0.2) |
| `cispo` | Clipped importance sampling | `input_ids`, `old_logprobs`, `advantages` | `eps_max` (6.0) |
| `dro` | Direct reward optimization | `input_ids`, `old_logprobs`, `advantages` | `beta` (0.05) |
| `opsd` | Hint-conditioned self-distillation | Aligned student and teacher streams; see below | `ce_coef` (1.0), `kl_coef` (1.0) |

- **`cross_entropy`** trains the model to predict target tokens, scaled by
  per-token `weights`.
- **`importance_sampling`** scales each token's policy-gradient update by its
  advantage and the current-to-sampling policy probability ratio. It does not
  clip the ratio, so updates can be noisy when the policies diverge.
- **`ppo`** uses the smaller of the unclipped and clipped advantage-weighted
  ratios, with clipping bounds `[1 − clip_low, 1 + clip_high]`. This limits the
  incentive to move token probabilities too far in one update.
- **`cispo`** caps the importance ratio at `eps_max` and detaches it from the
  gradient, while differentiating through the current token log probability.
  It is used in ScaleRL-style training.
- **`dro`** adds a quadratic penalty, controlled by `beta`, for moving away from
  the sampler's log probabilities.

## Distillation objectives

For sampled reverse-KL feedback, use an RL loss with a detached per-token
advantage of `teacher_logprob - student_logprob`. Keep the student sampling
policy's log probabilities in `old_logprobs`. The
[distillation guide](https://docs.river.ai/guides/distillation/) explains the estimator and provides
examples for a larger teacher, a self-teacher, and multiple specialists.

`cross_entropy` also accepts two-dimensional `target_tokens` and `weights` per
datum, both shaped `[sequence_length, K]`. At each prediction position, supply
K candidate token IDs and their teacher probabilities as weights. This trains
against soft targets. Renormalizing over only the top-K tokens changes the
teacher distribution; it is an approximation, not full-vocabulary KL. The
returned per-position log probability is the weight-normalized mean over
candidates, not a vector of their individual log probabilities.

### Native self-distillation with opsd

`opsd` computes full-vocabulary **forward KL**, `KL(teacher || student)`, from a
hint-conditioned teacher using the same model weights. The teacher stream has
no gradient. It can also include supervised cross-entropy on separate student
positions. It accepts text-only data, not image chunks or an external teacher.

| Stream | Per-datum fields |
| --- | --- |
| Student | `input_ids`, `target_tokens`, `loss_mask`, `kl_mask`, `kl_ids` |
| Teacher | `teacher_input_ids`, `teacher_target_tokens`, `teacher_kl_mask`, `teacher_kl_ids` |
| Optional attention masks | `attention_mask`, `teacher_attention_mask` |

Each stream uses its own sequence length. Targets are next-token IDs aligned
with prediction positions. Binary `kl_mask` fields select paired positions;
the active `kl_ids` must be exactly `0..K-1` in each stream and pair identical
target tokens. All masks must be zero at the trailing no-target position.
Student `loss_mask` and `kl_mask` must be disjoint.

Unlike the token-summed losses above, `opsd` optimizes the **sum of per-datum
means**: each datum contributes `ce_coef × CE_sum / N_CE` plus
`kl_coef × KL_sum / N_KL`, with empty components contributing zero. A positive
`ce_coef` requires SFT tokens in every datum. Set `ce_coef=0.0` for KL-only
training; every datum must then contain KL tokens. Coefficients must be finite,
nonnegative, and at least one must be positive.

For a mean over B datums, apply `gradient_scale=1 / B` once at the optimizer step
on endpoints supporting `gradient_scale_v1`. Do not also divide by token count.
`opsd_objective_sum` reports the optimized scalar; `opsd_kl` and `ce_loss` report
aggregated token means for the components. See the
[aligned example](https://docs.river.ai/guides/distillation-opsd/).

## Custom losses

Install `river-client[custom-loss]` and define a scalar PyTorch objective.
Call `model.prepare_loss()` before collecting training data to capture the
graph and validate its operations and reductions:

```python
import torch.nn.functional as F
from river_client import CustomLoss


def weighted_ce(logits, inputs):
    nll = F.cross_entropy(logits, inputs["targets"], reduction="none")
    w = inputs["weights"]
    return (nll * w).sum() / w.sum().clamp_min(1.0)


loss = model.prepare_loss(CustomLoss(
    weighted_ce,
    inputs={"targets": "int64", "weights": "float32"},
))
data = [{
    "input_ids": [1, 2, 3, 4],  # Illustrative token IDs.
    "loss_data": {
        "targets": [2, 3, 4, 0],
        "weights": [0.0, 1.0, 1.0, 0.0],
    },
}]
result = model.forward_backward(data, loss_fn=loss)
model.optim_step(lr=1e-4)
```

Supply unpadded token IDs and the declared columns in `loss_data`. Targets align
with each position's logits; mask unwanted positions explicitly, including the
last position when no next-token target exists. This example normalizes by the
total weight across the request, including when execution is chunked.

### Custom losses with AsyncTrainer

`AsyncTrainer` prepares custom losses before rollouts and supplies declared
`targets`, `old_logprobs`, `advantages`, and `mask` inputs.

The example below uses `CustomLoss` to implement
[SAPO (Soft Adaptive Policy Optimization)](https://arxiv.org/abs/2511.20347),
which replaces hard clipping with a smooth sigmoid-shaped objective. It
averages within each response, then across responses, as in the paper. Use the
single-turn `engine` and `rows` from the [RL recipe](https://docs.river.ai/guides/rl-sync/#training-recipe).

```python
import torch.nn.functional as F
from river_client import CustomLoss, rl


async def train_custom(engine, rows):
    temperature = engine.temperature

    def custom_loss_inputs(datum):
        inputs = dict(datum["loss_data"])
        mask = inputs.pop("mask")
        length = max(sum(mask), 1)
        return {**inputs, "weights": [m / length for m in mask]}

    def sapo(logits, inputs):
        logprobs = -F.cross_entropy(
            logits / temperature, inputs["targets"], reduction="none"
        ).float()
        ratio = (logprobs - inputs["old_logprobs"]).exp()
        advantages = inputs["advantages"]
        tau_pos, tau_neg = 1.0, 1.05
        positive = advantages.clamp_min(0) * (4.0 / tau_pos)
        negative = advantages.clamp_max(0) * (4.0 / tau_neg)
        surrogate = (
            positive / ((-tau_pos * (ratio - 1.0)).exp() + 1.0)
            + negative / ((-tau_neg * (ratio - 1.0)).exp() + 1.0)
        )
        weights = inputs["weights"]
        return -(surrogate * weights).sum() / weights.sum().clamp_min(1.0)

    trainer = rl.AsyncTrainer(
        engine=engine,
        optimizer=rl.Adam(lr=1e-5),
        completion=rl.GroupCompletion(mode="wait"),
        groups_per_step=2,
        group_size=8,
        max_staleness=2,
        advantage=rl.GroupCentered(standardize=True),
        normalize="sum",
        loss=CustomLoss(sapo, inputs={
            "targets": "int64",
            "old_logprobs": "float32",
            "advantages": "float32",
            "weights": "float32",
        }),
        custom_loss_inputs=custom_loss_inputs,
    )
    async for step in trainer.run(rows, steps=100):
        print(step.n, step.metrics)
```

Call `await train_custom(engine, rows)` from your async application. The weights
mask non-response tokens and give each response equal weight. `tau_pos` and
`tau_neg` control SAPO's smoothness; they are separate from sampling temperature.

- Use `normalize="sum"`; handle normalization, temperature, and clipping in the
  objective. Mask every loss term.
- Custom losses wait for the full training batch by default, as in this example.
  Streaming backward is opt-in via `forward_backward_groups` or
  `forward_backward_batch` for supported additive objectives. Server microbatching
  and loss chunking preserve global normalization.
- All-zero-advantage batches skip training by default. Set
  `train_zero_advantage_batches=True` for advantage-independent terms such as CE.
- For extra inputs, pass a `custom_loss_inputs` callback returning the datum's
  declared `loss_data` columns. It runs on a worker thread, so synchronize access
  to shared mutable state.