COMAP: Co-Evolving World Models and Agent Policies for LLM Agents

2026-06-01Artificial Intelligence

Artificial IntelligenceComputation and Language
AI summary

The authors propose COMAP, a system where an AI's understanding of the world (world model) and its decision-making (agent policy) improve together by learning from their own actions and outcomes continuously. Unlike previous models that stay the same after training, COMAP updates its world model based on the agent's recent experiences, helping it make better predictions and decisions over time. This approach works well in tasks like planning, browsing the web, and using tools, outperforming other methods. The authors also show that repeatedly updating both parts leads to more accurate world predictions and smarter long-term choices.

language agentsworld modelson-policy learningself-distillationagent policiesclosed-loop interactionembodied task planningfuture-aware reflectionlong-horizon decision-makingQwen3-4B
Authors
Youwei Liu, Jian Wang, Hanlin Wang, Wenjie Li
Abstract
Equipping language agents with world models enables them to anticipate environment dynamics and evaluate candidate actions before execution. However, existing textual world models are typically fixed after training, preventing them from adapting to the on-policy state-action distributions induced by an evolving agent. Meanwhile, agent-improvement methods often rely on external rewards or verifiers, limiting their applicability in realistic interactive environments. In this paper, we propose COMAP, a novel framework that co-evolves textual world models and agent policies through closed-loop interaction. At each decision step, the world model predicts future state feedback for candidate actions, and the agent performs future-aware reflection by estimating the reliability of this feedback and refining its action accordingly. The resulting on-policy trajectories are then used to update the world model via self-distillation, allowing it to better match the agent's evolving interaction distribution. Across embodied task planning, Web navigation, and tool-use benchmarks, COMAP consistently outperforms competitive baselines, e.g., +16.75% relative improvement with Qwen3-4B. Further analyses show that the co-evolutionary loop improves the world model's prediction accuracy over time and leads to more effective long-horizon decision-making. Our code is available at: https://github.com/loyiv/CoMAP.