dew.rl

Reinforcement learning as array math, one module per function.

advantage turns rewards into advantages, surrogate turns log-probabilities and advantages into a scalar loss. Neither knows about a model, an objective, the trainer or a batch dict, so an objective composes them and this package is testable on fixed tensors.

The import arrow points one way. dew.rl may read dew; nothing under dew outside dew.rl and dew.objectives.rl may read dew.rl. That one-way arrow keeps the split into a separate distribution a directory move (docs/design/plan.md, section 5.1). tests/test_rl_imports.py walks the tree and fails when either half of that stops being true.

advantage and surrogate port Apache-2.0 code from Tunix and verl and carry their notice; the rest of Dew is MIT.

Name	Summary
behavior_importance_weights	Detached token TIS weights from recorded raw and behavior policies.
clipped_surrogate	Token-mean reduction of the dual-clipped policy terms.
clipped_value_loss_terms	verl d040717 compute_value_loss, before its token-mask reduction.
gae	Generalized advantage estimation over [B, T] rewards and values.
group_advantage	Group-relative advantage of [B] rewards, group completions per prompt.
k3_kl	Schulman’s k3 estimator of KL(pi || pi_ref), per token.
masked_mean	Mean of x over the positions mask keeps.
masked_whiten	x centred and scaled by its masked mean and unbiased deviation.
preference_logsigmoid	Pair-mean reduction of the DPO sigmoid loss.
rloo_advantage	Each completion against the mean of the rest of its group.
sequence_log_ratio	GSPO’s sequence-level log ratio, carrying a per-token gradient.
token_log_ratio	Per-token log pi(a) - log pi_old(a), clamped to +-20.
token_mean	Sum over the unmasked positions of x, divided by how many there are.

behavior_importance_weights

function source

def behavior_importance_weights(
    old_log_probs: jax.Array,
    behavior_log_probs: jax.Array,
    mask: jax.Array,
    cap: float,
) -> jax.Array

Detached token TIS weights from recorded raw and behavior policies.

Port of verl compute_rollout_correction_weights(token), revision d040717b21af2e23e8e789a3e354cff2394ae2de: exponentiate the log ratio clamped to +-20, mask padding, then cap the weight. No batch normalization or rejection sampling is implied. Token TIS and filtered sampling do not recover an unbiased full-trajectory raw-policy expectation.

clipped_surrogate

function source

def clipped_surrogate(
    log_ratio: jax.Array,
    advantages: jax.Array,
    mask: jax.Array,
    epsilon_low: float = 0.2,
    epsilon_high: float = 0.2,
    dual_clip: float | None = 3.0,
) -> tuple[jax.Array, dict[str, jax.Array]]

Token-mean reduction of the dual-clipped policy terms.

clipped_value_loss_terms

function source

def clipped_value_loss_terms(
    predicted: jax.Array,
    returns: jax.Array,
    old_values: jax.Array,
    clip: float = 0.2,
) -> jax.Array

verl d040717 compute_value_loss, before its token-mask reduction.

The larger squared error of the live prediction and the prediction clipped around recorded values is multiplied by one half. Targets and recorded values are detached rollout data.

gae

function source

def gae(
    token_rewards: jax.Array,
    values: jax.Array,
    mask: jax.Array,
    gamma: float,
    lam: float,
) -> tuple[jax.Array, jax.Array]

Generalized advantage estimation over [B, T] rewards and values.

delta_t = r_t + gamma * V(s_{t+1}) - V(s_t) and A_t = delta_t + gamma * lam * A_{t+1}, backwards from the last step (arXiv:1506.02438). A masked step contributes nothing and passes the running advantage and the next value through unchanged, so a padded tail leaves the real steps before it undiscounted. Positions outside the mask hold the neighbouring step’s carry, as in both references, for the loss to mask again.

Returns the whitened advantages and the unwhitened returns, in that order. returns = A + V happens before the whitening in both references.

The recursion runs in float32 whatever the caller’s dtype, the way Tunix’s GRPO loss casts its log-probabilities. The whitening subtracts two nearby numbers, and bf16 rounds the difference away.

group_advantage

function source

def group_advantage(
    rewards: jax.Array,
    group: int,
    normalise_by_std: bool = True,
    eps: float = GROUP_EPS,
) -> jax.Array

Group-relative advantage of [B] rewards, group completions per prompt.

The rollout expands each prompt into group rows next to each other, so a group is a reshape. verl groups by a uid column, which allows ragged groups. Dew’s rollout cannot produce those.

normalise_by_std=False is Dr.GRPO (arXiv:2503.20783), which subtracts the group mean without scaling by the deviation. verl spells the same switch norm_adv_by_std_in_grpo.

k3_kl

function source

def k3_kl(log_probs: jax.Array, ref_log_probs: jax.Array) -> jax.Array

Schulman’s k3 estimator of KL(pi || pi_ref), per token.

exp(d) - d - 1 for d = log pi_ref - log pi, which is non-negative, unbiased and lower variance than -d (http://joschu.net/blog/kl-approx.html). verl’s kl_penalty_forward("k3") clamps d to +-20 before the exponential and the estimate to +-10 after it, and this follows verl. Without the second clamp one drifted token contributes exp(20) to the penalty and dominates the step.

Aggregate it the way the policy loss is aggregated, token_mean(kl, mask), and add beta times that.

masked_mean

function source

def masked_mean(x: jax.Array, mask: jax.Array, axis=None) -> jax.Array

Mean of x over the positions mask keeps.

Outside the mask the values are replaced through jnp.where, verl’s masked_sum form. A nan in a padded position survives a multiply by a zero mask and reaches the loss, and a padded position is where an uninitialised value sits.

masked_whiten

function source

def masked_whiten(x: jax.Array, mask: jax.Array) -> jax.Array

x centred and scaled by its masked mean and unbiased deviation.

Both references whiten GAE advantages this way, Bessel correction included, and both leave the positions outside the mask in the output for the loss to mask again. With one unmasked position the correction divides by zero. verl raises there; this runs under jit and cannot.

preference_logsigmoid

function source

def preference_logsigmoid(
    policy_chosen: jax.Array,
    policy_rejected: jax.Array,
    ref_chosen: jax.Array,
    ref_rejected: jax.Array,
    mask_chosen: jax.Array,
    mask_rejected: jax.Array,
    beta: float,
) -> jax.Array

Pair-mean reduction of the DPO sigmoid loss.

rloo_advantage

function source

def rloo_advantage(rewards: jax.Array, group: int) -> jax.Array

Each completion against the mean of the rest of its group.

r_i - mean(r_j, j != i), which is group / (group - 1) times the centred reward. Tunix writes the first form and verl the second (arXiv:2402.14740). This follows Tunix.

sequence_log_ratio

function source

def sequence_log_ratio(
    log_probs: jax.Array,
    old_log_probs: jax.Array,
    mask: jax.Array,
    segments: jax.Array | None = None,
) -> jax.Array

GSPO’s sequence-level log ratio, carrying a per-token gradient.

The sequence ratio is the geometric mean of the token ratios, so its log is the masked mean of the token log ratios (arXiv:2507.18071, equation 6). Written as logp - sg(logp) + sg(mean) the value of every token in a sequence is that one mean, while the derivative with respect to each token’s log-probability is that token’s own. Dropping either stop-gradient leaves the value untouched and changes every gradient. The test pins the gradients for that reason.

A sequence is a row, or with segments one packed chain of it (segment_mean), which is what verl’s per-row pooling sees when each chain is its own row.

Tunix clamps the token log ratios to +-20 before pooling them. This pools the raw difference, as verl’s compute_policy_loss_gspo does. The clamp at 10 on the result bounds what is exponentiated either way.

token_log_ratio

function source

def token_log_ratio(log_probs: jax.Array, old_log_probs: jax.Array) -> jax.Array

Per-token log pi(a) - log pi_old(a), clamped to +-20.

verl calls this negative_approx_kl and negates it for its ppo_kl metric, which is -token_mean(token_log_ratio(...), mask).

token_mean

function source

def token_mean(x: jax.Array, mask: jax.Array) -> jax.Array

Sum over the unmasked positions of x, divided by how many there are.

verl’s agg_loss(loss_agg_mode="token-mean") and Tunix’s aggregate_loss("token-mean"). The denominator is the exact token count, without masked_mean’s 1e-8. A loss that is 1e-8 off scales the gradient by the same factor, so both references keep the two reductions separate. A batch with no unmasked token divides by zero and surfaces as a nan.