Research map / Large language models / Agents, code and math

Reinforcement World Model Learning for LLM-based Agents

NeurIPS 2026 · Large language models › Agents, code and math · approach: Agent instruction-following benchmarks · 1 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes RWML, a self-supervised RL method that trains LLM agents to predict next states with an embedding-based sim-to-real reward, then fine-tunes with task-success RL on ALFWorld and τ²-Bench.

Problem

LLM-based agents struggle to anticipate action consequences and adapt to environment dynamics; existing world-model training relies on SFT with expert data or strong LLM synthetic data, which is costly and prioritizes token-level fidelity.

Approach

Train the policy LLM with GRPO to generate reasoning and a predicted next state, rewarded by a binary embedding cosine-similarity threshold against the real next state, using only interaction triplets with no expert data, strong-LLM labels, or task-success signals; subsample too-easy samples before training.

Key results

• RWML alone improves ALFWorld average from 13.0 to 32.6 (Qwen2.5-7B) without expert data or task-success reward.
• RWML + policy RL reaches 87.9 average on ALFWorld, outperforming GRPO task-success RL (81.0) by 6.9 points.
• RWML + policy RL reaches 43.7 average on tau2 Bench, outperforming GRPO (38.0) by 5.7 points and matching distillation (43.7).

Surprising findings

• LLM-as-a-judge reward is hackable and collapses ALFWorld performance (3.6 vs 32.6 for the embedding reward) while binarized judge is also much worse.
• World-model SFT collapses ALFWorld performance to 2.8 average, far below the base model's 13.0.

Limitations

• Gains on tau2 Bench depend on base model strength; Qwen2.5-7B struggles to transfer world knowledge to decisions.
• Reward relies on a threshold hyperparameter and embedding model choice; user responses are non-deterministic requiring different handling (ROUGE for tool outputs).
• Heuristic subsampling and hyperparameters chosen without extensive tuning.

Experimental setup

Datasets: ALFWorld tau2 Bench

Benchmarks: ALFWorld tau2 Bench MMLU-Redux IFEval MATH-500 GSM8k GPQA-Diamond LiveCodeBench

Metrics: success rate average accuracy forgetting degradation parameter change ratio

Baselines: ReAct prompting RFT GRPO task-success RL World Model SFT Implicit World Modeling ICRIT Distill expert SFT

Builds on: GRPO ReAct Implicit World Modeling Dyna Reinforced Finetuning

Closest papers by meaning

• Controlling Large Language Model with Latent Action (ICML 2025 · similarity 0.91)
• RLVR-World: Training World Models with Reinforcement Learning (NeurIPS 2025 · similarity 0.90)
• From Words to Actions: Unveiling the Theoretical Underpinnings of LLM-Driven Autonomous Systems (ICML 2024 · similarity 0.90)
• Agent World Model: Infinity Synthetic Environments for Agentic Reinforcement Learning (ICML 2026 · similarity 0.90)
• Do LLMs Build World Representations? Probing Through the Lens of State Abstraction (NeurIPS 2024 · similarity 0.90)
• Offline Training of Language Model Agents with Functions as Learnable Weights (ICML 2024 · similarity 0.89)