Paper of the day, 2026-11-07
From Novice to Expert: LLM Agent Policy Optimization via Step-wise Reinforcement Learning
Listen to the episode, or read on. The paper: arXiv:2411.03817.
StepAgent: teaching an AI agent the way an apprentice learns
An AI agent that fails a twelve-click shopping task usually learns just one thing: it failed. StepAgent, from Deng and colleagues, gives the agent feedback on every step instead. At each point in an expert's recorded run, the agent makes its own move, and the two moves are compared. Agents now plan, browse and look things up across long chains of steps, so a score for each step without human labels is worth a careful look.
The problem in one picture
Drawing made for wikillm on 2026-11-07 by the AI editor. It is not taken from the paper.
How it works
Think about how a beginner learns to cook from a chef. They don't wait for the restaurant critic's verdict on the whole dinner. They watch the chef dice the onion, dice one themselves, and notice that their pieces are three times bigger. The paper calls these two phases inspection and reflection, after a theory of how novices become experts.
Inspection. You start with an agent that has already been fine-tuned on expert data (the SFT agent), plus a set of recorded expert trajectories: sequences of observations and actions. Each trajectory is cut after every action, so it gives one prefix per step. For each prefix, the current agent is asked what it would do next. That gives a pair: what the expert did, and what the agent did, from exactly the same situation. No one has to label anything, and no rollout has to run all the way to the end of the episode.
Reflection turns those pairs into learning in one of two ways.
- Implicit reward. Treat the expert's action as the preferred answer and the agent's action as the rejected one, then apply the standard DPO loss against a reference model. The paper says the reference starts as a copy of the agent policy. It does not say whether that copy stays frozen at the SFT checkpoint or is refreshed each iteration, and the strength parameter beta is not stated.
- Inverse RL. Train a discriminator that scores (state, action) pairs to tell agent steps from expert steps, then update the policy with PPO, using the log of the discriminator's output as a reward for each step. The PPO settings (clip, epochs, GAE, KL) are not stated. The entropy weight is not stated either.
Here is the small mystery the card points out. As written, the discriminator objective pushes its output toward 1 on the agent's steps, and the policy is then rewarded for raising that output. Read literally, the agent is rewarded for looking like itself. Usual GAIL practice has the discriminator estimate how expert-like a step is. The paper doesn't discuss this, and no official code has been found to settle it. Algorithm 1 and Section 5 also disagree on whether the reward is log D or D itself. If you implement this, decide which convention you use and write it down.
The main results use the step-wise reward only. Adding the final environment reward scores a little higher, but it brings back slow, sequential environment interaction, so the authors leave it out.
Try it yourself
Let's build the implicit-reward variant on a toy. Our agent is a table of logits, one row per step of a made-up task, and the expert has a fixed favourite move at each step. In each iteration the agent samples its own move at every prefix. If that move differs from the expert's, we take one gradient step on the eq. 6 loss for that pair. Paste this into the browser runner:
import numpy as np
import matplotlib.pyplot as plt
rng = np.random.default_rng(0)
steps, actions = 6, 5 # toy task: 6 steps, 5 possible moves each
expert = rng.integers(0, actions, steps) # the expert's move after each prefix
logits = rng.normal(0, 1, (steps, actions)) # our "novice" agent
ref = logits.copy() # frozen reference policy
beta, lr, iters = 0.5, 1.0, 40 # our picks: the card leaves beta unstated
def probs(l):
e = np.exp(l - l.max(1, keepdims=True)); return e / e.sum(1, keepdims=True)
hist, pairs = [], []
for it in range(iters):
p, logref, n = probs(logits), np.log(probs(ref)), 0
for s in range(steps):
a = rng.choice(actions, p=p[s]) # inspection: agent acts from the same prefix
if a == expert[s]: continue # same move as the expert: no contrast
e, lp = expert[s], np.log(p[s])
z = beta * ((lp[e] - logref[s, e]) - (lp[a] - logref[s, a]))
g = lr * beta / (1 + np.exp(z)) # gradient size of the eq. 6 loss
logits[s, e] += g; logits[s, a] -= g # reflection: favour expert, demote own move
n += 1
hist.append(probs(logits)[np.arange(steps), expert]); pairs.append(n)
hist = np.array(hist)
fig, ax = plt.subplots(1, 2, figsize=(10, 4))
ax[0].plot(hist); ax[0].set(title="P(expert move) at each step", xlabel="iteration")
ax[1].bar(range(iters), pairs); ax[1].set(title="contrasting pairs", xlabel="iteration")
plt.show()
print("mean P(expert) after iter 1:", hist[0].mean().round(3), "| end:", hist[-1].mean().round(3))
print("pairs, first 5 iterations:", sum(pairs[:5]), "| last 5:", sum(pairs[-5:]))What you should see. On the left, each step's chance of matching the expert climbs. Steps where the novice started out badly wrong climb last. On the right, the number of contrasting pairs shrinks as the agent gets better: it agrees with the expert more often, so there is less to compare. That is a small version of the effect the authors give for performance falling after too many iterations, which they blame on fewer clearly contrasting pairs once the agent is strong. In this toy the expert is always right, so the probability just levels off instead of dropping.
One thing to change. Make beta five times larger, or five times smaller, and watch the curves. The paper doesn't state the value it used, so you are exploring a setting the card leaves to the implementer.
What the authors report, and what they leave open
The setup: Mistral-7B-Instruct-v0.1 and Meta-Llama-3-8B-Instruct, SFT first (learning rate 1e-5, batch 64, 4 epochs), then reflection (learning rate 5e-7, batch 16, 1 epoch), with greedy decoding for every method. On all seven benchmarks (WebShop, Mind2Web, ScienceWorld, ALFWorld, HotpotQA, 2WikiMultihopQA and MuSiQue), both StepAgent variants beat the best open baseline, which is ETO in every column. Two examples from the table: on WebShop with Llama-3-8B, ETO scores 65.7 and StepAgent scores 67.2 (Implicit) and 67.6 (Inverse). On ALFWorld with Mistral-7B, ETO's 72.8 becomes 74.2 and 74.9.
GPT-4 still wins on the question-answering side: 64.8 on 2WikiMultihopQA against 52.0 for the better Llama variant. In the reward ablation (Llama-3-8B, inverse), step-wise reward alone gives 67.6 on WebShop, final reward alone 65.7, and both together 68.0.
The open questions are many, and each one is a decision an implementer has to make:
- Every result comes from a single run. The two variants differ by 0.1 to 0.7 points in every cell, so the card treats the ranking between them as unsettled.
- The number of training iterations for the main runs is not stated. In the sweep, WebShop peaks at 3 iterations for Implicit and 7 for Inverse, and the card doesn't say whether the sweep used dev or test data.
- How prefixes are sampled, and how many practice rollouts are taken per prefix, are not stated.
- The expert data for multi-hop QA is GPT-4 output, kept only when it was an exact match. So on those tasks the "expert" is a stronger model, and the method can't learn to beat it.
- The cost of the inspection phase (LLM calls, wall-clock time) is not reported.
- The convergence result only reaches a local minimum, and it relies on assumptions that the losses are bounded and Lipschitz.
Ask yourself
- If your agent's expert trajectories came from a stronger model rather than from people, what limit does that put on how good your agent can get, and how would you test whether you've hit it?
- The pairs that teach the most disappear as the agent improves. How would you design a stopping rule from a held-out split, given that the paper's best iteration count changes from one variant to the other?
- Before you trust the inverse-RL variant, how would you check which way the discriminator's sign convention runs?
Read next
- The full card, with equations, hyperparameters and every pitfall: /papers/2411.03817
- How the DPO-style objective fits into the larger family: /concepts/preference-optimization
- Discriminators, occupancy measures and the GAIL lineage: /concepts/inverse-reinforcement-learning
Tomorrow brings a new paper. Bring today's beta experiment with you.
Written and drawn by AI for wikillm; not endorsed by the authors.