Co-entrenamiento de Política y Modelado del Mundo para Agentes de Lenguaje

Resumen

El aprendizaje por refuerzo (AR) mejora los agentes basados en modelos de lenguaje grande (MLG) al enseñarles qué acciones conducen a altas recompensas, pero proporciona poca supervisión sobre qué efectos tienen dichas acciones en el entorno. El modelado del mundo (MM) puede llenar este vacío, aunque los enfoques existentes suelen requerir simuladores separados, etapas de entrenamiento adicionales o cómputo extra en tiempo de inferencia. Observamos que las ejecuciones de AR dentro de política ya contienen la señal necesaria: cada transición empareja una acción con la observación resultante posterior. Basándonos en esta observación, proponemos PaW, un marco de coentrenamiento de política y modelado del mundo que añade supervisión auxiliar de MM a la misma política durante el AR, sin alterar el paradigma de inferencia. Para que la supervisión auxiliar de MM sea informativa y estable, PaW introduce tres componentes: selección de datos de MM basada en la entropía de la acción, pérdida de MM tolerante al ruido y balanceo de pérdida adaptativo a la recompensa. Los experimentos en tres conjuntos de tareas de agente muestran mejoras consistentes frente a bases sólidas de AR en distintos modelos y algoritmos de AR. Estos resultados sugieren que las ejecuciones estándar de AR constituyen una fuente práctica de supervisión de MM para el entrenamiento de agentes lingüísticos.

English

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.