Loss Functions

PyTRIO provides built-in loss functions for supervised learning and reinforcement learning.

You can select a loss function by passing a string to forward_backward:

future = training_client.forward_backward(
    data,
    loss_fn="cross_entropy", 
)
result = future.result()

Built-In Loss Functions

PyTRIO currently supports the following built-in loss functions:

cross_entropy

For supervised learning, PyTRIO implements the standard cross-entropy loss, also known as negative log-likelihood. This loss optimizes the policy $p_\theta$ to maximize the log probability of token $x$:

$L(\theta) = -\mathbb{E}_x[\log p_\theta(x)]$

Here weights are 0 or 1 and are usually generated by renderer.build_supervised_example(), which returns (model_input, weights) and is used to specify the target assistant turns that should be trained.

The implementation is:

# Apply weights and compute elementwise loss
elementwise_loss = -target_logprobs * weights
# Apply sum reduction to get the total loss
loss = elementwise_loss.sum()  # scalar

The cross_entropy loss requires the following fields in Datum.loss_fn_inputs:

• target_tokens: array[(N,), int] | array[(N, K), int]: target token IDs
• weights: array[(N,), float] | array[(N, K), float]: token-level loss weights, usually from the renderer

Output:

• logprobs: array[(N,), float] | array[(N, K), float]: log probabilities of the requested target tokens

Metric:

• loss_sum: aggregated loss returned by the SDK, a scalar

importance_sampling

For reinforcement learning, PyTRIO implements a common variant of the policy-gradient objective for practical cases where the learning policy $p$ and the sampling policy $q$ differ, for example because of non-determinism that makes the data off-policy.

The issue is that the objective:

$L(\theta) = \mathbb{E}_{x \sim p_\theta}[A(x)]$

can become biased when $x \sim q$ from the sampler is not exactly the same as the desired $x \sim p_\theta$ from the learner. To correct this bias, PyTRIO uses the improved importance-sampling objective:

$L_{\text{IS}}(\theta) = \mathbb{E}_{x \sim q}\left[\frac{p_\theta(x)}{q(x)} A(x)\right]$

This objective gives the correct expected reward. In the formula:

• $\log p_\theta(x)$ (target_logprobs) comes from the learner and is computed during the forward pass in forward_backward.
• $\log q(x)$ (sampling_logprobs) comes from the sampler and is recorded during sampling as the correction term.

The implementation is:

# Compute probability ratio
prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
# Compute importance-weighted loss
loss = -(prob_ratio * advantages).sum()

The importance_sampling loss requires the following fields in Datum.loss_fn_inputs:

• target_tokens: array[(N,), int]: target token IDs from sampler $q$
• logprobs: array[(N,), float]: token-level sampling_logprobs
• advantages: array[(N,), float]: RL advantage values. Positive values reinforce the token, while negative values suppress it.

Output:

• logprobs: array[(N,), float]: token-level target_logprobs

Metric:

• loss_sum: aggregated loss returned by the SDK, a scalar

ppo

PPO (Schulman et al., 2017) addresses issues in standard policy-gradient methods by introducing a clipped objective function. It restricts policy updates to the neighborhood of the sampling distribution, preventing excessive policy drift when multiple gradient updates are applied to the same rollout distribution.

The objective prevents overly large policy updates by clipping the importance ratio $\frac{p_\theta(x)}{q(x)}$, where $p_\theta$ is the learner policy and $q$ is the sampling policy. Note that PPO clipping and loss computation are both performed independently at the token level.

The PPO clipped objective is:

$L_{\text{CLIP}}(\theta) = -\mathbb{E}_{x \sim q}\left[\text{clip}\!\left(\frac{p_\theta(x)}{q(x)},\, 1 - \epsilon_{\text{low}},\, 1 + \epsilon_{\text{high}}\right) \cdot A(x)\right]$

The final PPO loss combines the clipped and unclipped objectives:

$L_{\text{PPO}}(\theta) = -\mathbb{E}_{x \sim q}\left[\min\!\left(\frac{p_\theta(x)}{q(x)} \cdot A(x),\; \text{clip}\!\left(\frac{p_\theta(x)}{q(x)},\, 1 - \epsilon_{\text{low}},\, 1 + \epsilon_{\text{high}}\right) \cdot A(x)\right)\right]$

Here $\epsilon_{\text{low}}$ and $\epsilon_{\text{high}}$ are hyperparameters, currently fixed to 0.2 in PyTRIO.

The implementation is:

# Compute probability ratio
prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
# Apply clipping
clipped_ratio = torch.clamp(prob_ratio, clip_low_threshold, clip_high_threshold)
# Compute both objectives
unclipped_objective = prob_ratio * advantages
clipped_objective = clipped_ratio * advantages
# Take minimum (most conservative)
ppo_objective = torch.min(unclipped_objective, clipped_objective)
# PPO loss is negative of objective
loss = -ppo_objective.sum()

The ppo loss requires the following fields in Datum.loss_fn_inputs:

• target_tokens: array[(N,), int]: target token IDs from sampler $q$
• logprobs: array[(N,), float]: token-level sampling_logprobs
• advantages: array[(N,), float]: RL advantage values

Output:

• logprobs: array[(N,), float]: token-level target_logprobs

Metric:

• loss_sum: aggregated loss returned by the SDK, a scalar

You can also customize the clipping thresholds with loss_fn_config:

fwd_bwd_future = await training_client.forward_backward_async(
    data=data,
    loss_fn="ppo",
    loss_fn_config={"clip_low_threshold": 0.9, "clip_high_threshold": 1.1}
)
fwd_bwd_result = await fwd_bwd_future

cispo

CISPO (Clipped Importance Sampling Policy Optimization) is a policy-gradient method. Like PPO, it uses the importance ratio $\frac{p_\theta(x)}{q(x)}$. The difference is that PPO clips the objective directly, while CISPO clips the importance ratio and uses the clipped ratio as the coefficient on target_logprobs.

The CISPO objective is:

$L_{\text{CISPO}}(\theta) = \mathbb{E}_{x \sim q}\left[\operatorname{sg}\left(\text{clip}\left(\frac{p_\theta(x)}{q(x)}, 1-\epsilon_{\text{low}}, 1+\epsilon_{\text{high}}\right)\right) \cdot \log p_\theta(x) \cdot A(x)\right]$

Here $\operatorname{sg}$ means stop-gradient. The clipped ratio is detached, so it only controls the gradient coefficient and does not backpropagate through the ratio itself.

The implementation is:

# Compute probability ratio
prob_ratio = torch.exp(target_logprobs - sampling_logprobs)
# Apply clipping
clipped_ratio = torch.clamp(prob_ratio, clip_low_threshold, clip_high_threshold)
# Compute CISPO objective (detach the clipped ratio)
cispo_objective = clipped_ratio.detach() * target_logprobs * advantages
# CISPO loss is negative of objective
loss = -cispo_objective.sum()

The cispo loss requires the following fields in Datum.loss_fn_inputs:

• target_tokens: array[(N,), int]: target token IDs from sampler $q$
• logprobs: array[(N,), float]: token-level sampling_logprobs
• advantages: array[(N,), float]: RL advantage values

Output:

• logprobs: array[(N,), float]: token-level target_logprobs

Metric:

• loss_sum: aggregated loss returned by the SDK, a scalar

CISPO uses one-sided clipping by default: no lower bound and an upper bound only, with clip_low_threshold=0.0 and clip_high_threshold=4.0. You can also set these explicitly with loss_fn_config:

fwd_bwd_future = await training_client.forward_backward_async(
    data=data,
    loss_fn="cispo",
    loss_fn_config={"clip_low_threshold": 0.0, "clip_high_threshold": 4.0}
)

In async or off-policy training, the sampling policy $q$ may lag behind the current learner policy $p_\theta$. A high lower bound can keep stale tokens weighted strongly even after the current policy has moved away from them, weakening the attenuation that importance sampling is meant to provide. Leaving the lower bound disabled is usually more robust.

dro

DRO (Direct Reward Optimization) is a general off-policy and even offline reinforcement learning method. It adds a quadratic penalty to the reward-weighted logprob term to constrain how far the learner policy $p_\theta$ moves from the sampling policy $q$.

The DRO objective is:

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

Here $\beta$ controls the strength of the quadratic penalty. Note that DRO uses a softer formulation of advantage estimation, and those advantages need to be constructed on the client side before being passed in.

The implementation is:

# Compute quadratic penalty term
quadratic_term = (target_logprobs - sampling_logprobs) ** 2
# Compute DRO objective
dro_objective = target_logprobs * advantages - 0.5 * beta * quadratic_term
# DRO loss is negative of objective
loss = -dro_objective.sum()

The dro loss requires the following fields in Datum.loss_fn_inputs:

• target_tokens: array[(N,), int]: target token IDs from sampler $q$
• logprobs: array[(N,), float]: token-level sampling_logprobs
• advantages: array[(N,), float]: RL advantage values

Output:

• logprobs: array[(N,), float]: token-level target_logprobs

Metric:

• loss_sum: aggregated loss returned by the SDK, a scalar

You can customize $\beta$ with loss_fn_config:

fwd_bwd_future = await training_client.forward_backward_async(
    data=data,
    loss_fn="dro",
    loss_fn_config={"beta": 0.05}
)