From Novice to Expert: LLM Agent Policy Optimization via Step-wise Reinforcement Learning

In plain words

A language-model agent that works through a task in several steps, such as browsing a shop website or answering a question by looking things up, usually learns from a single score at the end: the task succeeded or it did not. That score says little about which of the many steps went wrong, so learning is slow.

This paper's idea is to borrow a teaching pattern from how people become experts. A novice watches what an expert does at each point, tries the same step themselves, and notices where they differ. StepAgent does this with recorded expert trajectories. At each step it compares the expert's move with the agent's move from the same situation, and turns that difference into a per-step learning signal. No human has to label the steps.

The paper offers two ways to turn those comparisons into updates: one that directly prefers the expert's move over the agent's, and one that first learns a scoring function that separates expert moves from agent moves. The authors report that both improve agents on web, science-simulator and multi-hop question-answering tasks.

TL;DR

StepAgent is a reinforcement-learning recipe for LLM agents that uses step-wise rewards built from expert data, not from human annotation. For each expert step, the agent is asked to act from the same prefix, and the pair is used as a preference (implicit reward) or fed to a discriminator (inverse RL). The authors report gains over SFT, PPO, DPO, RFT, SPIN, NAT and ETO on WebShop, Mind2Web, ScienceWorld, ALFWorld and three multi-hop QA sets, with Mistral-7B and Llama-3-8B as backbones §7.1. Results are single runs reported by the authors.

Problem

Agent tasks are multi-step, and the usual reward is a single scalar for the whole trajectory §1. Step-wise supervision would help, but per-step value labels are rarely available in agent datasets, and annotating them is costly §1. Estimating step values with Monte Carlo tree search needs the agent to interact with the environment until a final outcome is reached, which cannot be parallelised §1. Imitating the complete expert trajectory with only a final reward has two costs: each rollout requires sequential environment interaction, and the agent must absorb every expert action at once §4.1.

The method assumes an SFT base agent already exists, and that expert trajectories (observation, action sequences) are available for the task [§3.2, §4]. It also assumes the agent can be sampled from at any prefix of an expert trajectory §4.1, eq. 5.

Method

Problem setup. The agent is a policy over actions given the history of observations and actions. An episode ends with a reward in [0, 1] for the whole trajectory §3.1. SFT maximises the likelihood of expert action tokens only, with observation tokens masked §3.2, eq. 2, eq. 3.

Inspection §4.1. Each expert trajectory is cut after every action, giving one prefix per step, and each prefix is treated as a short learning target §4.1, eq. 4. For each target, the current agent policy generates its own action from the same prefix §4.1, eq. 5. The pair (expert action, agent action) is the step-wise comparison that drives reflection.

Reflection §4.2. Two strategies consume the comparisons.

1. Implicit-reward RL §4.2.1. The expert action is the preferred response and the agent's action is the dispreferred one, at the same prefix. The loss is the DPO objective, with a frozen reference model that the paper says is initialised from the agent policy §4.2.1, eq. 6:

   L_\text{implicit}(\pi_\theta, \pi_\text{ref}) = -\mathbb{E}\Big[\log\sigma\Big(\beta\log\tfrac{\pi_\theta(\hat{a}_i|\hat{s}_i)}{\pi_\text{ref}(\hat{a}_i|\hat{s}_i)}-\beta\log\tfrac{\pi_\theta(a^{\theta}_i|\hat{s}_i)}{\pi_\text{ref}(a^{\theta}_i|\hat{s}_i)}\Big)\Big] eq. 6

2. Inverse RL §4.2.2. A discriminator D(s, a) in (0, 1) is trained to tell agent steps from expert steps. The paper gives the occupancy-measure and Jensen-Shannon formulation that motivates this §4.2.2, eq. 7, eq. 8, eq. 9. The discriminator objective is written as

   \mathbb{E}_{\pi_\theta}[\log D(s,a)] + \mathbb{E}_{\pi_e}[\log(1-D(s,a))] - \lambda H(\pi_\theta) eq. 10

   The policy is then updated with PPO, using the step-wise reward log D(s, a) alg. 1. The complete loop is given in Algorithm 1.

Algorithm 1 (inverse-RL variant) alg. 1. Shown in pseudocode form below. Each iteration runs inspection on sampled expert steps, then reflection with a discriminator update and one PPO step.

# StepAgent with inverse RL  [alg. 1]
# Input:  expert trajectories T, SFT base policy pi_theta0
# Output: final policy pi_theta

def stepagent_irl(expert_trajectories, base_policy, num_iterations):
    # Start the policy from the SFT base agent  [alg. 1]
    policy = base_policy
    discriminator = init_discriminator()
    for iteration in range(num_iterations):
        # Inspection: sample expert prefixes and let the agent act on each
        # Sampling rate is not stated in the paper  [alg. 1] [eq. 4] [eq. 5]
        expert_steps = sample_step_prefixes(expert_trajectories)
        agent_steps = [policy.act(prefix) for prefix in expert_steps]
        # Reflection: update D so agent steps score high, expert steps low
        # Alg. 1 calls this a loss; eq. 10 is written as a maximisation
        # [alg. 1] [eq. 10]
        discriminator.update(agent_steps, expert_steps)
        # Step-wise reward from the discriminator  [alg. 1]
        rewards = [log(discriminator(s, a)) for (s, a) in agent_steps]
        # One PPO update on the step-wise rewards  [alg. 1]
        policy = policy.ppo_update(agent_steps, rewards)
    return policy

Reward combination. The main results use step-wise reward only. Adding the final environment reward, the paper says, gives better numbers but requires environment interaction that cannot be parallelised, so the authors leave it out of the main results §7.2, tab. 3.

Theory §5. The paper assumes the losses in eq. 6 and eq. 10 are bounded and Lipschitz continuous §5, assumption. Under that assumption, and because the updates are gradient steps, the policy converges to a local minimum §5. Proposition 1 states that the agent's occupancy measure can approximate the expert's after several iterations, and a second proposition argues that optimising either loss is equivalent to reducing a KL or JS divergence §5.

Hyperparameters

Dataset sizes are in tab. 1: WebShop 1,938 train and 200 test; Mind2Web 1,009 train and 912 test; ScienceWorld 1,483 train, 194 dev, 241 test; ALFWorld 3,321 train, 140 dev, 134 test; HotpotQA, 2WikiMultihopQA and MuSiQue are listed with their full sizes tab. 1. The multi-hop experiments use the subset in §6.1, not these full sizes.

Results

Author-reported. Setup: Mistral-7B-Instruct-v0.1 and Meta-Llama-3-8B-Instruct, fine-tuned with the SFT and reflection settings above, evaluated on the test sets in tab. 1. Metrics are the environment reward for WebShop and ScienceWorld, success rate for ALFWorld, macro element accuracy for Mind2Web, and exact match for the three multi-hop QA sets §6.3. Best baseline is the strongest of SFT, PPO, DPO, RFT, SPIN, NAT and ETO for each column, from tab. 2. The table shows the two StepAgent variants, Implicit and Inverse.

Closed models in the same table: GPT-3.5 (gpt-3.5-turbo-1106) and GPT-4 (gpt-4-0125-preview) with the base prompt §6.2. GPT-4 scores higher than StepAgent-Mistral on HotpotQA (39.4 against 31.0), 2WikiMultihopQA (64.8 against 46.8) and MuSiQue (28.2 against 27.7). It also scores higher than StepAgent-Llama-3-8B on HotpotQA (39.4 against 38.1) and 2WikiMultihopQA (64.8 against 52.0). On the other columns both StepAgent backbones score higher than GPT-4 tab. 2.

Ablation on reward type, Llama-3-8B with inverse RL tab. 3: step-wise reward alone gives 67.6 on WebShop, 64.1 on ScienceWorld and 34.1 on MuSiQue; final reward alone gives 65.7, 62.5 and 32.3; both together give 68.0, 64.8 and 34.8 tab. 3. The main results use the step-wise reward alone §7.2.

Model size fig. 3. The paper compares Mistral-7B with Mistral-Nemo-Instruct-2407 (called "Mistral 13B") across WebShop, ScienceWorld, HotpotQA and 2WikiMultihopQA and reports the larger model is better on those tasks §7.3, fig. 3. The numbers are shown only in the figure, so they are not transcribed here.

Paper vs code

No official code. The arXiv abstract page lists no code link and the LaTeX source contains no repository URL, only dataset links that are commented out §6.1. A GitHub repository search for "StepAgent" returned four unrelated projects, none by the authors. A search with "StepAgent" restricted to the RUC-NLPIR organisation returned nothing; a search for the method description terms and a search for "Zhirui Deng" also returned nothing relevant.

Pitfalls

• Expert data for multi-hop QA is not human. GPT-4 generated the expert trajectories, filtered to exact match = 1 §6.4. The "expert" is therefore a stronger model's output on the training split, and the method has no way to exceed that expert on these tasks.
• Sampling of prefixes is underspecified. Equation 4 includes every prefix of every expert trajectory, but Algorithm 1 says "each sampled step-wise expert trajectory" eq. 4 alg. 1. The sampling rate and the number of practice rollouts per prefix in the main runs are not stated. The text reports gains up to three practice rollouts and none beyond §7.4, fig. 4.
• Training length is unstated for main runs. The sweep shows the best iteration count differs by method on WebShop (3 for implicit, 7 for inverse), and performance degrades past the peak, which the authors attribute to fewer clearly contrasting positive and negative pairs once the agent is stronger §7.4, fig. 4. Choose the count by a held-out split, not the test set.
• Selection split for the sweeps is not stated. §7.4 says the sweep used WebShop and HotpotQA but does not say whether it used dev or test data. WebShop has no dev split in tab. 1.
• Sign of the discriminator reward looks inverted relative to GAIL. Eq. 10 maximises E_pi[log D] + E_pi_e[log(1 - D)], which pushes D toward 1 on agent pairs. The policy then maximises log D alg. 1, so it is rewarded for steps that look like its own. GAIL-style training usually uses D as the probability that a pair came from the expert. The paper does not discuss this. Check it against the code if it is released; the card cannot resolve it.
• Three inconsistent uses of the discriminator output. Algorithm 1 uses log D(s, a) as the reward alg. 1. Section 5 Step 2 uses D(s, a) itself as the reward §5. Algorithm 1 calls eq. 10 a "loss" to train on, while eq. 10 is written as a supremum alg. 1 eq. 10.
• Implicit-RL reference model is ambiguous. The paper says it is "initialized with the agent policy" §4.2.1 but does not say whether it is frozen at the SFT checkpoint or refreshed each iteration. In the DPO loss of eq. 6 this choice changes the target.
• Several constants are missing. Beta (eq. 6), lambda (eq. 8), PPO clip, KL and GAE settings, and PPO epochs per batch are not stated alg. 1. An implementer will have to choose them and report that choice.
• Convergence is weaker than the abstract suggests. The abstract says the agent's action distribution converges toward the expert's. The analysis assumes bounded and Lipschitz losses §5, assumption, proves convergence only to a local minimum under gradient descent §5, and gives no rate or bound on the approximation in Proposition 1 §5.
• Small gaps are within run-to-run noise. The Implicit and Inverse variants differ by 0.1 to 0.7 points in every cell tab. 2. The paper reports one run per setting and no variance, so the ranking between the two variants is not established.
• Text and table disagree on one margin. The text says StepAgent beats ETO by 2.9 points on HotpotQA §7.1. That matches Llama-3-8B (38.1 against 35.2) tab. 2. The Mistral margin is 2.8 (31.0 against 28.2).
• Step-wise reward alone is a deliberate trade-off. The combined step and final reward is best in tab. 3 (68.0 on WebShop), but the paper leaves it out because final rewards need sequential environment interaction §7.2.
• Environment cost is not reported. The number of LLM calls for inspection, the wall-clock time of the reflection stage, and the cost per practice rollout are not given. The only compute statement is the GPU type and count §6.4.
• "Mistral 13B" is a different model. The footnote names Mistral-Nemo-Instruct-2407 §7.3. Do not assume a Mistral 13B checkpoint.
• Closed-model baselines are dated. The GPT model names in §6.2 are pinned snapshots, so later API models are not covered by these numbers.

Lineage

• Builds on: Benner's novice-to-expert theory, the conceptual basis §1, §4. Its technical components are the DPO loss §4.2.1, and the GAIL-style occupancy-measure and discriminator formulation with PPO §4.2.2.
• Compared against: SFT, PPO, DPO, RFT, SPIN, NAT and ETO §6.2. Of these, ETO and SPIN are the strongest baselines in tab. 2.
• Superseded by: not known from content already in this wiki.