World Modeling for LLM Agent - Predicting or Acting in the Environment
Agentic reinforcement learning usually trains only on the agent’s own tokens: the commands it writes. The tool outputs it reads, such as query results, error messages and table histories, are masked out. The learning signal is one scalar reward per episode.
ECHO (Shrivastava et al., 2026) and related work (PaW) add a second signal. They put a cross-entropy loss on the observation tokens, so the model also learns to predict what the environment will say. They claim for
- faster learning
- better generalization
- less need for expert demonstrations
The intuition is that an agent that knows what a command will do should choose commands better.