Why Do LLM Agents Fail in Exploring New Environments? A World-Modeling Perspective

AI summary

Overview

Research area: Reinforcement learning for large language model (LLM) agents, with a focus on world modeling, exploration, and generalization to unfamiliar environments.

Technical level: Advanced. The paper assumes familiarity with reinforcement learning (PPO, advantage estimation, KL/clipping), LLM finetuning (SFT), and agentic scaffolding (ReAct-style think-then-act loops), plus metrics such as Pass@1 and Pass@k.

Scope in one sentence: The paper diagnoses a failure mode called "exploration collapse" — Pass@k dropping during RL in environments unfamiliar to the policy — and proposes SPA, an explore-then-exploit recipe that grounds the agent in state estimation and transition prediction before RL.

What This Paper Is About

LLM agents often fail to improve when reinforcement learning is applied in new environments whose states the pretrained model has never seen. The authors trace this to weak grounding in environment states and dynamics, and test whether explicitly teaching the agent to estimate the current state and predict the next state — before optimizing for reward — restores the ability to explore across many plausible solution paths.

Key Contributions

1. Identification of "exploration collapse." The paper characterizes a systematic divergence in unseen environments: Pass@1 (success of the highest-probability trajectory) edges up during RL while Pass@k (whether at least one of k sampled trajectories succeeds) drops markedly, indicating increasingly brittle exploration. Environments closer to the pretraining distribution (ALFWorld, WebShop) show no such decline.

2. A grounding diagnosis. The collapse is traced to weak grounding in environment states and dynamics, quantified by state-description perplexity. Sokoban (163.9), FrozenLake (187.1), and Sudoku (15.5) show high state PPL relative to their number of possible grid-cell values (7, 6, and 5 respectively), while ALFWorld (6.0) and WebShop (11.7) are much lower — a proxy for "representational familiarity" of state strings to the LM.

3. SPA (Self-exPerience Agent). An explore-then-exploit framework combining (i) state estimation, which concatenates raw states with an abstract representation giving coordinates of key entities (player, boxes, goals), and (ii) transition modeling, trained via self-experience SFT where the model's own interaction trajectories are supervised on current-state and next-state prediction. World modeling is confined to the SFT cold-start stage rather than injected as an inference-time planner or a reward signal.

4. Controlled ablations isolating four ingredients: grounded state representations, explicit transition modeling, self-experience trajectories from a sufficiently strong exploration policy, and adequate coverage of transition data.

Main Findings

• Pass@k collapse in unfamiliar environments. In Sokoban, FrozenLake, and Sudoku, standard RL causes Pass@k to consistently deteriorate over training while Pass@1 improves only marginally. In ALFWorld and WebShop, both metrics improve together and Pass@k "fastly goes to 100 percent."

• Large main-table gains. On Qwen2.5-1.5B-Instruct, SPA raises Sokoban Pass@1 from 25.6 to 59.8 and FrozenLake Pass@1 from 22.1 to 70.9, with Pass@8 of 69.5 (Sokoban) and 75.0 (FrozenLake). The paper states this lets sub-3B models surpass a 20B baseline on these tasks.

• Consistent ordering across models. For every model evaluated, SPA outperforms State Estimation RL, which in turn outperforms vanilla RL. Examples (Pass@1/Pass@8, in units of 10⁻²): Qwen2.5-0.5B-Instruct Sokoban 36.7/45.3 vs. vanilla 16.9/35.9; Qwen2.5-3B Sokoban 49.7/58.2 vs. vanilla 31.4/35.5; LLaMA3.2-1B-Instruct Sokoban 53.0/68.0 vs. vanilla 21.2/39.8; LLaMA3.2-1B-Instruct Sudoku 81.3/100 vs. vanilla 0.1/1.2.

• SPA beats the reward-shaping baseline (VAGEN) under matched settings. In their reimplementation, VAGEN yields little improvement (e.g., 44.5/50.0 on Sokoban for Qwen2.5-1.5B vs. SPA's 59.8/69.5). The authors tentatively attribute this to SPA confining world-model learning to SFT and leaving the RL objective unperturbed, avoiding multi-objective interference.

• Transition modeling is the key ingredient. Masking the SFT loss on current and next states (replacing observation and prediction spans with [MASKED] and setting their loss to 0) yields no improvement for downstream PPO.

• Ground-truth supervision is necessary. Training SFT on the model's own self-belief states instead of ground-truth states degrades RL: 39.2 Pass@1 / 44.2 Pass@8, versus SPA's 59.8/69.5 and State Estimation RL's 52.7/53.9 on the same Sokoban setup.

• World modeling helps even without state estimation. SPA without state estimation reaches 36.3 Pass@1 / 40.6 Pass@8, versus vanilla RL 25.6/34.0.

• The prediction slot is not redundant with planning. Dropping next-state prediction while keeping free-form planning lowers Pass@1/Pass@8 from 36.3/40.6 to 25.6/34.0. On-policy prediction accuracy rises from 31.9% to 46.6% during RL, indicating the policy refines rather than ignores its predictions.

• Ground-truth coordinates matter. Replacing all ground-truth (x, y) coordinates for player/box/goal with i.i.d. random coordinates from the same grid causes training to collapse: Pass@1 and Pass@8 stagnate at low levels with fluctuation, whereas with correct coordinates both metrics improve steadily.

• A strong exploration policy is needed for data collection. Replacing the SFT data generator with a uniform random-action generator (sampling from {Up, Down, Left, Right} with a fixed reasoning token sequence) produces substantially worse downstream RL: 1-epoch RandSFT gives 0.1 Pass@1 / 0.4 Pass@8, and 5-epoch RandSFT gives 20.2/50.0, versus SPA's 1-epoch 29.2/52.7 and 5-epoch 59.8/69.5.

• Format filtering helps a weak model. On LLaMA-3.2-1B-Instruct in Sokoban, removing misaligned SFT samples (enforcing minimal structural validity of <observation>–<prediction> reasoning followed by <answer>) improves training stability and both Pass@k and Pass@1.

• Longer world-modeling SFT improves downstream RL. Going from 1 to 5 SFT epochs, Pass@1 rises from 0.29 to 0.60 and Pass@k from 0.53 to 0.70. Average actions per episode drop from 8.47 to 6.44, action effectiveness rebounds from an early dip at Epoch 2's 0.60 to 0.77 by Epoch 5, and response length climbs from 65 tokens at Epoch 1 then stabilizes in the 160 to 175 range for Epochs 2–5 — so gains are not caused by verbosity.

• Explore-then-exploit dynamics. Pass@k continues to increase in the early phase of training (steps 0–400), remaining high even at k=1024, and only decreases in the later phase (steps 400–1000). The authors note this differs from single-turn RLVR, where base models achieve higher Pass@k at large k.

• Useful exploration, not just more exploration. Grounded SPA runs visit fewer total states than vanilla runs, yet cover far more success-relevant states with higher action diversity, per trajectory metrics in Appendix D.

• Easy-to-hard transfer works; cross-game transfer does not. Training the world model with SFT on FrozenLake 4×4 and then running RL on 6×6 substantially outperforms non-world-modeling baselines, though this is a transfer analysis rather than a same-budget comparison. Cross-complexity Sokoban (trained 6×6 one box, tested 10×10 two boxes) gives 0.9 Pass@1 / 3.1 Pass@k versus a baseline trained directly on the complex task at 0.1 / 0.8. Cross-game Sokoban→FrozenLake gives 15.9 Pass@1 / 49.2 Pass@k versus a FrozenLake baseline at 17.2 / 49.2, showing cross-game generalization is difficult.

• ALFWorld check. SPA also delivers consistent Pass@1 and Pass@k gains on ALFWorld with Qwen2.5-0.5B-Instruct, chosen because performance saturates quickly for the 1.5B model; the authors describe this as initial evidence and leave broader real-world validation to future work.

Methodology in Plain English

The authors first measure how unfamiliar each environment is by computing the perplexity of serialized state descriptions relative to the number of possible values each grid cell can take — a proxy for how surprising the state string is to the language model. They then track Pass@1 and Pass@k throughout RL training to see whether exploration broadens or narrows.

To fix the problem, they split world modeling into two pieces. First, state estimation: raw symbolic states (for example, a 6×6 Sokoban board serialized as text) are augmented with an abstract representation giving coordinates of the player, boxes, and goals, which lowers state perplexity. Second, transition modeling: the base model is prompted to freely interact with the environment, producing trajectories in which the reasoning trace must explicitly write the current state inside <observation> tags and the predicted next state inside <prediction> tags before emitting actions inside <answer> tags. The model's self-belief states are then replaced with the environment's ground-truth states, so SFT supervises a masked token-level cross-entropy loss only over the <think>...</think> and <answer>...</answer> spans, with observation tokens masked out.

The resulting SFT checkpoint initializes standard RL with PPO over a masked objective that applies only to answer tokens, using batch size 32 and eight rollouts per environment at each update step over 1000 training steps. The codebase is built on RAGEN, and the comparison baseline is VAGEN reimplemented in the same state-estimation setting. The critical design choice is that world modeling lives in the SFT cold start, not in the RL reward and not as an inference-time planner — the authors emphasize the learned predictive distribution is tied to the collecting policy, so they call it a grounded initialization rather than an action-conditioned simulator.

Ablations isolate each ingredient: masking the state/prediction supervision, substituting self-belief states for ground truth, dropping state estimation, randomizing coordinates, swapping the data-collecting policy for random actions, filtering format-noncompliant samples, and varying SFT epochs.

Why This Matters

Impact on research. The paper reframes a common failure of agentic RL — agents that do not improve in new environments — as a grounding problem rather than a pure optimization problem, and shows that a cheap SFT cold start can restore Pass@k scaling. It extends the known Pass@1/Pass@k divergence from single-turn RLVR to multi-turn, stateful environments, and offers a policy-centric alternative to reward-shaped world modeling. The claim is deliberately narrow: the internalized world model is an initialization, not a planner, and the authors state they do not claim the divergence itself is new.

Real-world applications:

• Computer-use and GUI agents that must operate on unfamiliar software on first contact, where serialized screen states are far from anything in pretraining.
• Tool-use and API-calling agents dropped into new tool ecosystems whose state and transition semantics must be learned from interaction.
• Web navigation and search agents facing site layouts outside their pretraining distribution.
• Embodied control and game-playing systems where agents must learn sparse-reward dynamics with limited rollout budget.

Industry relevance. The recipe is simple, uses the model's own trajectories without external knowledge or larger teacher models, and is reported to let sub-3B models surpass a 20B baseline on the evaluated tasks — attractive for cost-sensitive deployment. It also requires access to privileged simulator state during training, which the authors acknowledge as a scope limitation rather than a property of the method.

Future Directions

• Scaling and richer tasks. The main results are on grid-style interactive benchmarks (Sokoban, FrozenLake, Sudoku) with small- to mid-scale open-weight models, and ALFWorld provides only a single-model check. Whether SPA transfers to larger models, web navigation, tool use, and embodied control remains unverified.

• Reducing reliance on privileged state. SPA supervises SFT with ground-truth states and coordinates exposed by the simulator, and ablations indicate this supervision is necessary. Environments without accessible privileged state would need it estimated, for instance via a learned perception module — a case not studied here.

• A formal mechanism account. The authors restore exploration evidence through Pass@k recovery plus post-hoc trajectory metrics, but leave a fully formalized account of the exploration–exploitation tradeoff (in terms of policy entropy or regret) and a finer-grained causal account of why transition-prediction SFT restores exploration to future work.

• Stability and better baselines. Training remains unstable in stochastic settings, errors in the instruction/data-generation stage degrade the supervised transition data, and the VAGEN comparison is a reimplementation under matched settings that may not reflect its best configuration. Next steps suggested include uncertainty-aware transition modeling, richer modalities, and investigating SPA training stability.

Target Audience

Researchers and engineers working on RL for LLM agents, especially those training multi-turn, tool-using, or computer-use agents that must operate outside the pretraining distribution. It is also useful for practitioners interested in world models for language agents, and for readers tracking the Pass@1/Pass@k exploration literature from single-turn RLVR. Readers need prior familiarity with PPO-style RL and modern LLM finetuning pipelines to follow the training setup and ablations in detail.

Authors’ abstract

Large Language Models (LLMs) as agents often fail to improve in new environments. We identify and characterize a failure mode we call exploration collapse: under reinforcement learning (RL) in environments whose states are unfamiliar to the policy, Pass@k, the probability that at least one of k sampled trajectories succeeds, drops markedly over training even as Pass@1 edges up, revealing increasingly brittle exploration; environments closer to the pretraining distribution show no such decline. We trace this collapse to weak grounding in environment states and dynamics, and study a simple remedy: explicitly teaching the agent to estimate the current state and predict its transitions before optimizing for reward. We instantiate it as SPA, an explore-then-exploit recipe that cold-starts the policy with a Self-Experience supervised finetuning (SFT) stage, collecting the model's own interaction trajectories and supervising state and next-state prediction, and then runs standard RL. The resulting world model serves as a grounded initialization for RL rather than an inference-time planner. Across unseen environments, SPA consistently and substantially improves over vanilla RL: for example, it raises the Sokoban success rate from 25.6% to 59.8% on Qwen2.5-1.5B-Instruct, letting sub-3B models surpass a 20B baseline on these tasks. Controlled studies indicate that the gains track four factors: grounded state representations, explicit transition modeling, self-experience trajectories from a sufficiently strong exploration policy, and adequate coverage of transition data.