Policy- und Weltmodellierungs-Co-Training für Sprachagenten

Zusammenfassung

Verstärkungslernen (RL) verbessert Agenten auf Basis großer Sprachmodelle (LLMs), indem es ihnen beibringt, welche Aktionen zu hohen Belohnungen führen, liefert jedoch nur wenig Aufschluss darüber, was diese Aktionen in der Umgebung bewirken. Weltmodellierung (WM) kann diese Lücke schließen, doch bestehende Ansätze erfordern oft separate Simulatoren, zusätzliche Trainingsphasen oder zusätzliche Rechenzeit während der Inferenz. Wir beobachten, dass On-Policy-RL-Rollouts bereits das benötigte Signal enthalten: Jeder Übergang paart eine Aktion mit der daraus resultierenden nächsten Beobachtung. Basierend auf dieser Beobachtung schlagen wir PaW vor, ein Co-Training-Framework für Policy und Weltmodellierung, das während des RL eine zusätzliche WM-Überwachung für dieselbe Policy hinzufügt, ohne das Inferenzparadigma zu ändern. Um die zusätzliche WM-Überwachung informativ und stabil zu gestalten, führt PaW drei Komponenten ein: aktionsentropiebasierte WM-Datenauswahl, rauschresistenter WM-Verlust und belohnungsadaptive Verlustgewichtung. Experimente mit drei Benchmarks für agentische Aufgaben zeigen konsistente Verbesserungen gegenüber starken RL-Baselines über verschiedene Modelle und RL-Algorithmen hinweg. Diese Ergebnisse deuten darauf hin, dass standardmäßige RL-Rollouts eine praktische Quelle für WM-Überwachung beim Training von Sprachagenten darstellen.

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.