Critical-State RL: Finding the Call Worth Training

Why People Wanted It

A multi-turn tool interaction can contain several model calls, yet only some calls determine whether the task succeeds. Training every generated token treats a refusal, a useful tool choice, and a downstream recovery as if they were equally responsible for the reward. The real question is not simply which trajectory earned reward, but which occurrence has a learnable action signal that is relevant to the task.

A mixed batch of successes and failures does not answer that. Outcomes may vary because the continuation is noisy even when the decision at the candidate call never changes. Critical-State RL therefore separates action-dependent variation from continuation noise, checks that the local label is tied to benchmark return and has room to improve, and then updates only the selected call.

The Core Move

Find the occurrence whose action-conditioned local reward is both useful and variable, then apply the learning loss to that call alone. At a fixed prefix, the variance of the local label splits into action signal and continuation noise:

$\operatorname{Var}(z\mid x)=\operatorname{Var}_{a}(\mathbb{E}[z\mid x,a])+\mathbb{E}_{a}[\operatorname{Var}(z\mid x,a)].$

The first term asks whether changing the action changes its expected label; the second asks how much the label still fluctuates when the action is held fixed.

Quiz

A candidate group has mixed rewards, but the candidate action is identical in every rollout. What does Critical-State RL need to test before calling the group trainable?

Whether the batch has at least one success and one failure.

Whether the longest trajectory has the highest return.

Whether all surrounding calls receive policy gradients.

Whether action-conditioned label means vary, separately from within-action continuation variation.

$ check

Why a Mixed Reward Group Can Be Misleading

Now the version for readers who want more of the machinery: a nonzero reward variance across completions does not establish that the candidate action carries a trainable signal. Let $x$ be the prefix, $a$ the action at the candidate occurrence, and $z$ its occurrence-local label. The law of total variance gives the split above: variation between action-conditioned means is distinct from variation among continuations for the same action.

A concrete counterexample is the no-think Gemma missing-function pilot: sampled groups had mixed rewards, but every surviving rollout made the same clean refusal with no tool call. Reward varied from $0$ to $1$ with within-group standard deviation $0.477$, so a mixed-outcome filter would retain a group whose observed decision behavior did not distinguish the rewards. The diagnostic asks a sharper question: if the action were fixed, would the reward still vary? If so, raw group variance may mostly be continuation noise. That leaves the question of how to estimate the two sources without confusing them.

Why Fixed Prefixes Separate Signal from Noise

Here is the part that gets technical: nested sampling keeps the prefix fixed while comparing actions, then keeps each action fixed while resampling reward-only continuations. Sample a policy prefix $x$; draw candidate actions at that prefix; for each action, draw repeated continuations and observe the local label. The mean label for an action estimates its conditional mean, while variation among its repeated continuations estimates continuation noise.

The action-dependent component is the variance across those action means, with a finite-sampling correction for estimation noise. Repeating this across sampled prefixes and averaging yields a candidate-level estimate. Pooling raw labels across model-generated prefixes would mix upstream-prefix variation with action variation; using ground-truth teacher-forced prefixes can instead hide recovery-turn trainability when the candidate action distribution is saturated. After candidates pass the action-sufficiency and headroom screens, the diagnostic ranks them by estimated action-dependent label variance under a common label. What makes the selected call worth improving rather than merely variable?

Why Selection Needs Both Headroom and Local Training

Now the exact qualification conditions: a phase is critical only when its occurrence-local label mediates the action-to-benchmark-return relationship (action-sufficiency), some action’s label mean exceeds the reference-policy mean (headroom), and action-conditioned label means vary under the base policy (trainability). Task structure supplies the candidate phase and label; reference-policy and base-policy rollouts measure headroom and trainability.

The diagnosis is not itself the update. Once an occurrence is selected, optimize the policy with its local label and apply the RL loss only to that call’s generated tokens. Its prefix supplies context; a downstream continuation can supply the consequence and reward, but receives no gradient. This makes the selected occurrence—not the whole trajectory—the training unit. The remaining design question is what useful context or consequence can be supplied without moving the gradient boundary.

Minimal Causal Code

This sketch isolates the training boundary. In practice, diagnosis needs nested samples and the task-specific label; the loop below makes the local update visible rather than fast.

# One selected occurrence: prefix is context; continuation supplies reward only.
for batch in training_batches:
    prefix = batch.prefix
    action_tokens = policy.generate_call(prefix)
    continuation = environment.run_without_grad(action_tokens)
    label = local_label(action_tokens, continuation)

    logp = policy.log_prob(prefix, action_tokens)
    loss = -(label * logp).mean()

    optimizer.zero_grad()
    loss.backward()  # gradients flow through the selected call only
    optimizer.step()

What You Gain, and What You Lose

You gain a principled way to distinguish action-dependent reward signal from a merely mixed batch, and a clear boundary for local credit assignment. The selection rule also makes headroom explicit instead of treating every variable reward as an opportunity.

You pay for nested sampling: prefixes, candidate actions, and repeated continuations must be evaluated before training, and finite-sample estimates can misrank candidates. The label must also be designed so that it is action-sufficient for benchmark return; a convenient proxy can grade a local behavior without capturing the full task chain. Finally, local updates deliberately leave surrounding calls without gradient. That can focus credit, but it also means the method does not directly train those calls from this selected occurrence's reward.

Quiz

After an occurrence is selected, where does the RL loss apply?

To every token in the full trajectory, including all continuations.

To the prefix and continuation, but not to the selected call.

Only to the final benchmark score, with no call tokens involved.

Only to the selected call’s generated tokens.

$ check