Training Decision Models on Agent-Visited States

Why People Wanted It

A decision model can answer a fixed multiple-choice question well and still be unprepared for the state an agent actually reaches. The agent’s choices change what happens next; an early mistake can lead to a later state that never appeared in a clean demonstration. The real question is not just how to teach the model the right answer in familiar states, but how to train it on the situations its host agent will actually create—including the messy ones after an error.

The Core Move

Train the decision model on states recorded from the harness, include states after deliberately replayed mistakes, then let it make choices in the agent loop and learn from task outcomes. In the supervised stage, a target can mix a gold answer with a teacher’s probability distribution:

$y(a\mid s)=\alpha\,\mathbf{1}[a=a^\star]+(1-\alpha)\,p_{\text{JEV}}(a\mid s)$

When the teacher ranks the gold answer first, the target combines a hard label with the teacher’s softer preferences; otherwise it uses the gold label alone.

Quiz

Why include decision states that follow an agent’s mistakes in training?

Because mistake states can be labeled without knowing which options were available.

Because every error state must be assigned the same action.

Because it removes the need to evaluate the model in an agent loop.

Because earlier choices alter later states, including states that successful demonstrations may omit.

$ check

Why Visited States Include the Agent’s Own Consequences

Now the version for readers who want more of the machinery: the training distribution is built from harness states rather than only from a static question bank. The harness was run on training tasks and each decision state was recorded; a scripted policy also repeated mistakes seen during development so later states could follow those errors. This matters because an agent’s earlier choice changes its history and therefore changes the next decision prompt. Training only on successful demonstrations leaves that part of the deployed state distribution underrepresented.

In the reported supervised collection, 587 training tasks produced 15,694 decision states, and 30.9% followed a mistake. These figures describe this training setup, not a general recipe for how many post-error examples every agent needs. The central constraint is structural: labels must correspond to the state actually presented, including the history created by prior decisions. Why should the model imitate a teacher’s uncertainty as well as the answer key?

Why Supervision Mixes a Gold Answer with Teacher Probabilities

Here is the part that gets technical: when JEV ranked the correct option first, the target blended a one-hot gold label with JEV’s probabilities, using $\alpha=0.3$. In words, gold supplies a firm preference for the correct choice, while the teacher distribution preserves some of its relative preferences among options. If JEV did not rank the gold answer first, the target reverted to the one-hot gold label alone.

The model was fitted with cross-entropy plus a squared-error Brier term, which penalizes mismatches between predicted option probabilities and the target probabilities. That makes the objective care about the probability distribution, not only whether the top option is right. It is a useful distinction for a model whose outputs may later be used as confidence estimates, although this module is about training data and not the separate mechanism that decides when to hand a choice to the model. What does supervised fitting miss once the model starts making choices itself?

Why the Agent Loop Adds Consequence-Based Training

Now the deeper training step: the supervised model makes decisions on training tasks while a fixed agent LLM handles the rest, so its choices affect the episode that supplies its reward. The reward combines task success, a reduction for LLM calls, and a penalty for an irreversible wrong call:

$R_i=S_i\left(1-0.5\min\left(1,\frac{N_i}{10}\right)\right)-0.2E_i$

In words, success earns a reward reduced as LLM calls accumulate up to the stated cap, while an irreversible wrong call incurs a penalty. Groups of episodes provide relative advantages, and each decision in an episode shares its episode-level advantage in the GRPO objective; recurring states also receive a step-level term, with a KL penalty keeping the policy near the supervised model. This connects a local choice to the outcome of the trajectory it helped produce. It also makes training depend on the fixed executor and the reward design: what counts as a good decision if the task reward or call cost changes?

Minimal Causal Code

This plain loop makes the key causal link visible: choose from the current state, then use the resulting state and outcome to update the decision model. It is not the paper’s full GRPO implementation.

import torch

for task in training_tasks:
    state = harness.reset(task)
    episode = []
    done = False

    while not done:
        options = harness.options(state)
        probs = decision_model(state, options)
        action = sample_option(probs)

        next_state, outcome, done, info = harness.step(action)
        episode.append((state, options, action, outcome, info))
        state = next_state

    reward = task_reward(episode)
    loss = policy_loss(decision_model, episode, reward)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

What You Gain, and What You Lose

Training on harness-visited states aligns the examples with the prompts the model encounters in control, and replaying mistakes exposes it to states that successful demonstrations tend to omit. Supervised labels give a direct signal; agent-loop training can additionally reward consequences such as task success and fewer LLM calls.

The tradeoff is that collecting and training on trajectories costs task runs, and the learned behavior is shaped by the harness, fixed executor, and chosen reward. A model can learn the wrong incentives if the reward misprices calls or fails to represent important costs. Visited-state training also does not guarantee robustness to states the training agent never reaches. The point is not that trajectories erase distribution shift; it is that the training distribution can include consequences of the model’s own decisions rather than pretending every decision arrives in isolation.