# Custom losses

Write your training objective in PyTorch and run its loss and backward pass on
River's training backend. Custom losses let you experiment with objectives,
auxiliary terms, and weighting schemes for RL, distillation, and supervised
learning.

The client captures your function as a symbolic graph. River validates the
graph and executes the loss and its backward pass on the training workers,
alongside the model's logits.

For the design and application examples, read our blogpost: [Custom Losses at Scale](https://river.ai/custom-losses).

## Define, prepare, and train

Install the custom-loss extra. The examples below use `river-client` 0.17.1 or
later. Graph capture runs on the client's CPU.

```bash
pip install --upgrade 'river-client[custom-loss]>=0.17.1'
```

For an existing training `model`, such as the one in the
[SFT guide](/guides/sft/), define an objective and declare its input columns:

```python
import torch.nn.functional as F
import river_client as river


def masked_cross_entropy(logits, inputs):
    ce = F.cross_entropy(logits, inputs["targets"], reduction="none").float()
    mask = inputs["mask"]
    loss = (ce * mask).sum() / mask.sum().clamp_min(1)
    return loss, {"token_count": mask.sum()}


loss = model.prepare_loss(river.CustomLoss(
    masked_cross_entropy,
    inputs={"targets": "int64", "mask": "float32"},
))

data = [{
    "input_ids": [1, 2, 3, 4],  # Illustrative token IDs.
    "loss_data": {
        "targets": [2, 3, 4, 0],
        "mask": [0.0, 1.0, 1.0, 0.0],
    },
}]
result = model.forward_backward(data, loss_fn=loss, zero_out=True)
model.optim_step(lr=1e-4)  # Runs only after backward succeeds.
print(result.metrics)  # Includes custom/token_count.
```

The function receives logits shaped `[T, V]` for **one example**: `T` token
positions and `V` vocabulary entries. It returns a scalar tensor, optionally
paired with a dictionary of scalar metrics. Supply the declared columns in
`loss_data`; the schema specifies their names, dtypes, and shapes.

Keep the value returned by `prepare_loss`. Preparation captures the formula and
its constants and checks support on that model before training. Prepare again
when changing the formula, its constants, or the model/session. Put values that
change between examples in input columns.

Use `model.forward` to evaluate the prepared loss. Custom objectives return scalar
loss values and the metrics you declare, named `custom/<name>`.

## How River executes the objective

<figure class="docs-learning-figure">
<div class="docs-flow">
<div><strong>PyTorch function</strong><small>The client captures supported tensor operations on CPU.</small></div>
<span aria-hidden="true">→</span>
<div class="docs-model"><strong>Validated graph</strong><small>River checks operations, inputs, and resource requirements.</small></div>
<span aria-hidden="true">→</span>
<div><strong>Loss and backward</strong><small>Training workers execute the graph and propagate gradients into the model.</small></div>
</div>
<figcaption>Your Python code defines the objective; loss evaluation and training gradients stay on River's backend.</figcaption>
</figure>

A full sequence-by-vocabulary logits matrix can exceed GPU memory. River
processes token rows in smaller groups and combines their results before
finalizing each example's loss. It manages distributed execution and the loss
workspace.

The backend accepts validated graphs of supported tensor operations. Each
training request is checked for input and resource limits.

## Example: choose among candidate decisions

A classifier can reuse the model's language-model head by mapping each choice
to one token. At the decision position, normalize over those candidate tokens.

This example has one decision position per sequence. `target` is one-hot at that
position and zero elsewhere. `log_mask` is zero there and a large finite
negative value elsewhere; `weight` carries the example's weight at that position
and is zero elsewhere.

```python
import torch.nn.functional as F
import river_client as river


def decision_loss(logits, inputs):
    log_probs = F.log_softmax(logits, dim=-1)
    candidates = log_probs.gather(-1, inputs["candidates"]).float()
    normalizer = (candidates + inputs["log_mask"]).logsumexp((0, 1))
    target_log_prob = (inputs["target"] * candidates).sum()
    return (normalizer - target_log_prob) * inputs["weight"].sum()


K = 4  # Number of candidate decisions, each represented by a distinct token.
decision = model.prepare_loss(river.CustomLoss(
    decision_loss,
    inputs={
        "candidates": {"dtype": "int64", "shape": ("T", K)},
        "target": {"dtype": "float32", "shape": ("T", K)},
        "log_mask": {"dtype": "float32", "shape": ("T", 1)},
        "weight": "float32",
    },
))
```

Use valid candidate token IDs at every position, including masked positions.
Set `log_mask` to a finite value such as `-1e30` at non-decision positions. The
same objective can use a vision-capable model's image-conditioned logits.

## Use with AsyncTrainer

`AsyncTrainer` prepares the loss before starting rollouts. Use `river.CustomLoss`
for the tensor objective and `rl.CustomLoss` for its input adapter, sampling
requirements, and zero-advantage behavior. Keep reduction in the objective and
omit `normalize` and built-in loss options from the trainer configuration.

By default, the trainer supplies any declared subset of:

| Column | Dtype | Meaning |
| --- | --- | --- |
| `targets` | `int64` | Next-token targets aligned to prediction positions. |
| `old_logprobs` | `float32` | Log probabilities of the tokens when they were sampled. |
| `advantages` | `float32` | Values from the configured advantage estimator. |
| `mask` | `float32` | Eligible training positions, including the staleness policy. |

To form importance ratios, compute current-policy logprobs from
`logits / engine.temperature` inside the objective. Custom objectives require
unfiltered stochastic sampling: `top_p=1`, `top_k=-1` or `0`, and a positive
temperature.

For example, [SAPO](https://arxiv.org/abs/2511.20347) replaces hard clipping with
a smooth surrogate. With the single-turn `engine` and `rows` from the
[RL recipe](/guides/rl-sync/#training-recipe):

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


temperature = engine.temperature


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
    surrogate = (
        advantages.clamp_min(0) * (4.0 / tau_pos)
        / ((-tau_pos * (ratio - 1.0)).exp() + 1.0)
        + advantages.clamp_max(0) * (4.0 / tau_neg)
        / ((-tau_neg * (ratio - 1.0)).exp() + 1.0)
    )
    mask = inputs["mask"]
    return -(surrogate * mask).sum() / mask.sum().clamp_min(1)


loss = rl.CustomLoss(loss=river.CustomLoss(
    sapo,
    inputs={
        "targets": "int64",
        "old_logprobs": "float32",
        "advantages": "float32",
        "mask": "float32",
    },
))
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),
    loss=loss,
)
rl.run(trainer, rows, steps=100)
```

Each response contributes its own masked mean; the optimizer accumulates their
sum. The `tau` constants control the surrogate's smoothness.

### Request top-K logprobs

`old_logprobs` holds the sampled token's log probability at each position.
For alternative-token probabilities, set
`requirements=rl.LossRequirements(top_logprobs=K)` on `rl.CustomLoss` **before
starting rollouts**. The sampling distribution stays unchanged.

A client-side `input_adapter` receives the requested arrays in
`datum["sampling_metadata"]`:

| Field | Shape | Meaning |
| --- | --- | --- |
| `sampler_topk_ids` | `[T, K]` | Candidate IDs from the sampling policy. |
| `sampler_topk_logprobs` | `[T, K]` | Their log probabilities under the sampling policy. |
| `policy_age_steps` | `[T]` | Optimizer-step ages, when `policy_age=True` is also requested. |

These arrays are already aligned to prediction positions. Return the columns
your objective declares, and use `datum["loss_data"]["mask"]` for eligibility.
The following example uses retained sampling probabilities as soft targets:

```python
import numpy as np
import torch.nn.functional as F
import river_client as river
from river_client import rl

K = 8


class TopKInputs:
    def __call__(self, datum):
        metadata = datum["sampling_metadata"]
        ids = np.asarray(metadata["sampler_topk_ids"], dtype=np.int64)
        probabilities = np.exp(metadata["sampler_topk_logprobs"])
        return {
            "candidate_ids": np.maximum(ids, 0),  # Replace unused -1 IDs.
            "probabilities": np.where(ids >= 0, probabilities, 0).astype(np.float32),
            "mask": datum["loss_data"]["mask"],
        }

    def contract(self):
        return {"version": 1, "top_logprobs": K}


def topk_soft_targets(temperature):
    def objective(logits, inputs):
        logp = F.log_softmax(logits / temperature, dim=-1)
        selected = logp.gather(-1, inputs["candidate_ids"]).float()
        ce = -(selected * inputs["probabilities"]).sum(-1)
        mask = inputs["mask"]
        return (ce * mask).sum() / mask.sum().clamp_min(1)

    return rl.CustomLoss(
        loss=river.CustomLoss(objective, inputs={
            "candidate_ids": {"dtype": "int64", "shape": ("T", K)},
            "probabilities": {"dtype": "float32", "shape": ("T", K)},
            "mask": "float32",
        }),
        input_adapter=TopKInputs(),
        requirements=rl.LossRequirements(top_logprobs=K),
        train_zero_advantage_batches=True,
    )


# Use this loss in the AsyncTrainer construction above.
loss = topk_soft_targets(engine.temperature)
```

The sampler's original top-K probabilities serve as cross-entropy weights.

CISPO supports score centering directly through
`rl.CISPO(normalize="token", score_centering=rl.ScoreCentering(top_k=K))`,
which requests the sampling metadata automatically.

## Supported operations and validation

Custom losses use a supported subset of PyTorch:

| Computation | Supported forms |
| --- | --- |
| Vocabulary operations | `F.cross_entropy(..., reduction="none")` with default options; vocabulary-axis `log_softmax` followed by `gather` using declared `[T, K]` integer IDs. |
| Selected-token arithmetic | Arithmetic, clamps, `minimum`, `maximum`, square, `exp`, `expm1`, `log`, `sqrt`, `detach`, and float32 casts. |
| Per-example reductions | `sum`, `mean`, `amin`, `amax`, `logsumexp`, `var`, and `std`, with `keepdim=False` and no dtype override. Candidate-axis summation is also supported. |
| Scalar operations | Combine reduced scalars or apply supported nonlinear functions to them. |

Model-dependent rows stay independent until reduction. Keep reduced
model-dependent values in scalar computations, including detached values.
Input-only scalars, such as mask counts, can normalize token values. Use declared
columns for input data and Python numeric constants for fixed coefficients.
The graph must have the same structure across tensor values, sequence lengths,
devices, dtypes, and autograd states.

Preparation checks the graph. Submitted batches are checked for shape, input
size, and memory use. The loss workspace is one part of total training memory.
Use validation errors to identify unsupported expressions, fix input data, or
reduce a request's size. Contact River support for model or endpoint
availability errors.