Reinforcement Learning for Calibrated Decisions

Why People Wanted It

A same-speaker detector can select the correct answer while attaching an unhelpful probability to it. Thresholding the probability into “same” or “different” hides how tentative or confident the prediction was.

The real question is therefore not just whether training rewards the right answer. It is whether the training signal evaluates the probability distribution behind that answer.

Reinforcement Learning for Calibrated Decisions (RLCD) adds an optional stage starting from the supervised fine-tuning checkpoint. Each audio-conditioned question becomes a one-step decision problem: sample answers, evaluate them against the known outcome, and update the model. There is no continuing environment or future reward to predict.

The Core Move

Use a group of sampled answers to estimate probability-sensitive feedback, then update each sampled answer’s log-probability relative to a reference that excludes that draw.

For question $i$ and sampled draw $r$, the central comparison is

$$A_{ir}=G_{ir}-b_{ir}.$$

Here, $G_{ir}$ is the influence return—the feedback assigned to that draw by the group-reward construction. The baseline $b_{ir}$ is a comparison value built without that draw. Their difference, $A_{ir}$, is the advantage used in the policy-gradient update.

The important distinction is where this feedback comes from. It is constructed to target probability-scoring objectives, not merely to assign a reward whenever a sampled answer matches the label.

Quiz

What distinguishes RLCD's sampled group reward from simply rewarding every correct sampled answer?

It replaces the observed outcome with the most frequent sampled answer.

It rewards higher confidence regardless of which answer is correct.

It accounts for both outcome agreement and the sampled decision distribution.

$ check

Why Sampling Is an Exploration Policy

Here is the part that gets technical: the distribution used to collect answers and the distribution being trained play different roles.

Let $x_i$ be the audio-conditioned question and $y_i$ its observed answer. An action $a$ is one valid answer—not an audio frame or an intermediate reasoning step. The model’s decision policy is

$$p_{\theta,i}(a)=\pi_\theta(a\mid x_i).$$

This is the probability that parameters $\theta$ assign to answer $a$ in context $x_i$. The supported decision types determine what counts as an action:

Decision type	Available actions	Example
Choice	Supplied candidate options	Selecting an instrument category
Noul	False or true	Whether two recordings contain the same speaker
Score	Ordered rating levels	Selecting an audio-quality level

Choice and Score use categorical distributions over their valid actions. Noul uses the two-class distribution $(1-p_i,p_i)$, where $p_i$ is the probability of true. The decision state is extracted at a causal position before the answer. The observed outcome evaluates the sampled decisions; it is not supplied as the answer in that decision state.

For each question, RLCD draws $M=32$ independent actions with replacement from an exploratory behavior policy. All draws use the distribution obtained from a single forward pass. “With replacement” matters: two draws can select the same answer and still count as distinct samples.

Importance weights correct for sampling from the exploratory behavior policy rather than directly from the decision policy being optimized. Without this correction, the frequency of an answer in the group would reflect the sampling distribution without properly accounting for the target decision distribution.

This separates two jobs: exploration supplies the samples; the corrected feedback determines the learning signal. What, then, makes that feedback about probability quality?

Why the Group Reward Includes Distribution Shape

Now the version for readers who want more of the machinery: a probability score evaluates the entire predicted distribution against an observed outcome, and the group reward estimates that kind of objective from samples.

For a categorical decision, write the observed outcome as a one-hot target:

$$t_{ik}=\mathbf{1}[k=y_i].$$

The target is one for the observed answer and zero for the other answers. The Brier discrepancy is the sum of squared probability errors:

$$B_i=\sum_{k=1}^{K_i}\bigl(p_{\theta,i}(k)-t_{ik}\bigr)^2.$$

Here, $K_i$ is the number of valid outcomes. Unlike checking only the highest-probability answer, this expression evaluates every probability entry.

For Noul, encode false as zero and true as one. The two entries are complementary, so the discrepancy reduces to

$$B_i=(p_i-y_i)^2+\bigl((1-p_i)-(1-y_i)\bigr)^2 =2(p_i-y_i)^2.$$

Both entries contribute the same squared error. The factor of two comes from scoring the full two-class distribution, not from a second binary decision.

For ordered Score outputs, the objective also includes a normalized ranked-probability discrepancy $D_i$. This evaluates probability allocation along the ordered levels, rather than treating the levels as unrelated categories. Its weight is

$$\lambda_i=0.25\,\mathbf{1}[\tau_i=\mathrm{Score}],$$

where $\tau_i$ identifies the decision type. Thus the combined target is $B_i+\lambda_iD_i$: Brier error for every decision type, plus an ordinal contribution only for Score.

The sampled group reward has two complementary parts. An outcome-agreement term compares sampled answers with the observed answer. A pairwise term compares sampled answers with one another and characterizes the concentration of their distribution. Together, these connect the sampled feedback to the probability-scoring target.

Neither term should be mistaken for the whole mechanism. Counting correct draws alone misses the distribution-sensitive construction; rewarding concentration alone could favor confident mistakes. Nor does the group vote replace the observed label. The remaining question is how to assign this group feedback to individual draws.

Why the Baseline Leaves the Sampled Action Out

Here is the deeper optimization detail: excluding a draw from its baseline prevents that comparison value from depending on the very action whose log-probability is being differentiated.

For each draw $r$, RLCD constructs the influence return $G_{ir}$ and subtracts $b_{ir}$. The baseline excludes draw $r$ from both the outcome-agreement and pairwise terms. Exclusion is by sample index, not by answer category: if other draws select the same answer, those other draws remain available.

Conditioned on the question, outcome, and fixed model parameters used for sampling, the other independent draws do not depend on which action draw $r$ selected. This is the relevant independence. The baseline still uses the observed outcome and the remaining samples.

During differentiation, the feedback side is held fixed. Gradients flow through sampled-action log-probabilities; sampled actions, rewards, influence returns, baselines, and importance weights all receive stop-gradient treatment. In other words, training differentiates how likely the sampled answer is, not the arithmetic that produced its feedback.

With the specified reward, importance correction, and leave-one-out construction, the expected surrogate gradient equals the gradient of

$$B_i+\lambda_iD_i.$$

This is a statement about the estimator in expectation. It does not mean every sampled update improves the objective, and it does not apply to an arbitrary replacement reward or baseline.

Minimal Causal Code

This PyTorch loop makes the parameter update and gradient boundary visible. It is the final update interface, not a complete implementation of the group-reward estimator. rlcd_feedback must supply the exploratory sampling, importance correction, influence returns, and leave-one-out baselines described above.

Its coefficients are detached, signed coefficients for a loss minimized by gradient descent. They must include the appropriate advantage, importance factors, sign, and normalization; they are not simply correct/incorrect rewards.

import torch


def rlcd_step(model, optimizer, contexts, labels, decision_types,
              rlcd_feedback):
    optimizer.zero_grad()

    # Read the decision state before the answer, then obtain valid-action logits.
    logits = model.decision_logits(contexts)
    probs = logits.softmax(dim=-1)

    # One model forward pass; feedback construction cannot backpropagate.
    with torch.no_grad():
        actions, coefficients = rlcd_feedback(
            probs=probs,
            labels=labels,
            decision_types=decision_types,
            num_samples=32,
        )

    loss = logits.new_zeros(())
    for i in range(logits.shape[0]):
        log_probs = logits[i].log_softmax(dim=-1)
        for r in range(actions.shape[1]):
            a = actions[i, r]
            coefficient = coefficients[i, r].detach()
            # Only this sampled-action log-probability carries gradients.
            loss = loss + coefficient * log_probs[a]

    loss.backward()
    optimizer.step()  # The model parameters are the state changed here.
    return loss.detach()

The interface assumes logits have already been restricted or masked to valid actions. There is no recurrent environment-state update: this is a one-step decision problem. The visible update changes model parameters after the sampled feedback has been fixed.

Backpropagating through feedback construction would violate the stated gradient boundary. Conversely, detaching arbitrary feedback does not make it RLCD: the estimator inside rlcd_feedback is essential.

What You Gain, and What You Lose

You gain an explicit probability-quality target during an optional stage after supervised training. Brier error evaluates probabilities rather than only selected labels, and the Score-specific contribution respects ordered ratings. The sampled construction provides a policy-gradient route to those objectives.

You pay for an additional training stage, sampled-action processing, importance correction, and group-feedback computation. Sharing a forward pass across the draws does not remove those additional computations.

You also take on implementation constraints. The exploration correction, draw-level exclusion, and stop-gradient boundaries must agree with the reward construction. A generic loop that rewards correct answers is not an interchangeable substitute.

Finally, an objective aimed at probability quality is not a guarantee of calibrated predictions on every downstream task. Correctness, probability error, and useful confidence remain distinct things to evaluate.

Quiz

Why exclude draw r from both terms of its leave-one-out baseline?

To make the baseline independent of the sampled action being differentiated.

To prevent the observed label from influencing training.

To remove the need for importance correction.

$ check