Ning Lu, Baijiong Lin, Shengcai Liu, Jiahao Wu, Haoze Lv, Yanbin Wei, Lingting Zhu, Shengju Qian, Xin Wang, Ying-Cong Chen, Qi Wang, Ke Tang
6 min
Abstract
Reinforcement learning (RL) improves large language model (LLM) agents by teaching them which actions lead to high rewards, but provides little supervision on what those actions do to the environment. World modeling (WM) can fill this gap, yet existing approaches often require separate simulators, extra training stages, or additional inference-time computation. We observe that on-policy RL rollouts already contain the needed signal: each transition pairs an action with its resulting next observation. Based on this observation, we propose PaW, a Policy and World modeling co-training framework that adds auxiliary WM supervision to the same policy during RL, without changing the inference paradigm. To make auxiliary WM supervision informative and stable, PaW introduces three components: action-entropy-based WM data selection, noise-tolerant WM loss, and reward-adaptive loss balancing. Experiments on three agentic task benchmarks show consistent improvements over strong RL baselines across models and RL algorithms. These results suggest that standard RL rollouts are a practical source of WM supervision for language-agent training.
Alex: So if the model tries to learn from that noise, it starts treating randomness as a rule?
Sam: Precisely — that's called overfitting. To address this, PaW uses a specific filtering technique. Rather than trying to predict every detail of the next observation perfectly, it limits how much the model is penalized for getting unpredictable or noisy parts wrong. Think of a chef learning to cook. They learn to ignore the random sizzle and clatter of a busy kitchen and focus only on what actually matters — the effect of heat, or how much salt changes the flavor.
Alex: And crucially, this filtering only happens during training. The agent doesn't need to do any of that extra work when it's actually out in the world.
Sam: Right. The world modeling is a training tool, not a runtime cost. The agent just performs better because it has a more accurate internal model of how things work.
Alex: But are all training experiences equally worth learning from?
Sam: Not at all — and that's where the second key idea comes in. The framework uses what the researchers call "Action-Entropy Data Selection." Here's the intuition: when the agent is very confident about what to do next, that moment isn't very informative. It already knows the answer. But when the agent is genuinely uncertain — when several options seem equally plausible — that's a moment where it has the most to learn.
Alex: So it selects for the moments of genuine uncertainty, and uses those to train the world model?
Sam: Exactly. It prioritizes the transitions where the agent was most unsure, because those are the situations where understanding the environment's response matters most. It's like studying for an exam by focusing on the questions you got wrong, rather than drilling the ones you already know cold.
Alex: So the two ideas work together — filter out the noise, and then prioritize the most informative moments. The agent ends up with a much cleaner, more reliable picture of how its world operates.
Sam: That's the core logic. It's a more disciplined approach to learning from experience.
Alex: Does this actually translate into better performance on real tasks?
Sam: The results suggest it does. The researchers tested PaW on tasks like ALFWorld — which involves completing household chores described in text — and WebShop, which simulates navigating an online store to find specific products. In both cases, PaW consistently improved the success rates of the underlying agents.
Alex: And does it hold up as the tasks get more complex?
Sam: It appears to scale reasonably well. The researchers observed meaningful gains across different model sizes, and the approach also helped in search-augmented question answering — tasks where the agent has to navigate multiple steps to locate and synthesize information.
Alex: So it's not just about getting the right answer. It's about the agent actually learning the *process* of how to get there.
Sam: That's the key shift. By forcing the agent to predict the environment's response at each step, you move it from being a pattern-matcher to something that understands the consequences of its actions. That understanding is what makes it more robust when conditions change.
Alex: What are the limitations? No framework is complete — where does this one fall short?
Sam: The authors are candid about this. The current version uses one-step predictions — it looks at the immediate next state, but doesn't explicitly model long-term consequences. If the real payoff requires a long chain of correct moves, the agent might still struggle, because small early errors can compound as the sequence grows longer.
Alex: It's like learning chess by only thinking one move ahead. You'll understand the board, but you might miss a trap set ten moves away.
Sam: That's a fair comparison. The authors also note that the training data currently isn't deduplicated — meaning if the agent repeatedly practices the same simple moves, the world model gets biased toward those common patterns and may stop learning from genuinely new situations.
Alex: So future work might look at multi-step modeling and more diverse training data — making the agent's internal map of the world both deeper and broader.
Sam: That's the direction the authors point toward. But even with those limitations, the core contribution is meaningful: you don't always need to build a complex, separate simulator to help an agent learn. Sometimes the data you already have is enough — if you know how to use it properly.
Alex: It's a clean, efficient approach to a long-standing challenge in AI training. Thanks for walking us through it, Sam.
Sam: It was a good one to unpack. The elegance is in the simplicity — reuse what you have, filter out the noise, and focus on what's genuinely informative.
Alex: And that's a principle that applies well beyond AI. Thanks for listening to ResearchPod.