Aprendizado de Modelos de Mundo por Reforço para Agentes Baseados em LLM
Reinforcement World Model Learning for LLM-based Agents
February 5, 2026
Autores: Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu
cs.AI
Resumo
Os grandes modelos de linguagem (LLMs) demonstraram alto desempenho em tarefas centradas em linguagem. No entanto, em ambientes agentivos, os LLMs frequentemente lutam para antecipar consequências de ações e adaptar-se à dinâmica do ambiente, destacando a necessidade de capacidades de modelagem de mundo em agentes baseados em LLM. Propomos o Aprendizado de Modelo de Mundo por Reforço (RWML), um método auto supervisionado que aprende modelos de mundo condicionados por ação para agentes baseados em LLM em estados textuais, utilizando recompensas de diferença simulação-realidade. Nosso método alinha os próximos estados simulados produzidos pelo modelo com os próximos estados realizados observados a partir do ambiente, incentivando a consistência entre simulações internas do mundo e a dinâmica real do ambiente em um espaço de incorporação pré-treinado. Diferente da previsão de tokens do próximo estado, que prioriza a fidelidade a nível de token (ou seja, reproduzir a redação exata) em detrimento da equivalência semântica e pode levar ao colapso do modelo, nosso método fornece um sinal de treinamento mais robusto e é empiricamente menos suscetível a manipulação de recompensas do que o uso de LLM como juiz. Avaliamos nosso método no ALFWorld e no τ² Bench e observamos ganhos significativos em relação ao modelo base, apesar de ser totalmente auto supervisionado. Quando combinado com recompensas de sucesso na tarefa, nosso método supera o RL de recompensa direta de sucesso na tarefa em 6,9 e 5,7 pontos no ALFWorld e no τ² Bench, respectivamente, equiparando-se ao desempenho do treinamento com dados especializados.
English
Large language models (LLMs) have achieved strong performance in language-centric tasks. However, in agentic settings, LLMs often struggle to anticipate action consequences and adapt to environment dynamics, highlighting the need for world-modeling capabilities in LLM-based agents. We propose Reinforcement World Model Learning (RWML), a self-supervised method that learns action-conditioned world models for LLM-based agents on textual states using sim-to-real gap rewards. Our method aligns simulated next states produced by the model with realized next states observed from the environment, encouraging consistency between internal world simulations and actual environment dynamics in a pre-trained embedding space. Unlike next-state token prediction, which prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse, our method provides a more robust training signal and is empirically less susceptible to reward hacking than LLM-as-a-judge. We evaluate our method on ALFWorld and τ^2 Bench and observe significant gains over the base model, despite being entirely self-supervised. When combined with task-success rewards, our method outperforms direct task-success reward RL by 6.9 and 5.7 points on ALFWorld and τ^2 Bench respectively, while matching the performance of expert-data training.