Textual Belief States for World Models: Identifiable Representation Learning Under Strict Mediation
For researchers building world models for text-based environments, this work provides a principled solution to ensure identifiable and interpretable latent representations, which is crucial for reliable planning and reasoning.
The paper addresses the problem of unidentifiable latent representations in LLM-based world models due to history bypass. By enforcing strict latent state mediation with discrete textual states and a novel training method (fGRPO), they achieve up to 57% gains in representation quality and 98% improvements in rollout performance on TextWorld and ScienceWorld benchmarks.
World models in partially observed environments rely on latent representations that summarize interaction history, but in many modern LLM-based architectures predictive performance fails to reflect representation quality due to history bypass, rendering the latent state unidentifiable. Strict latent state mediation, requiring predictions to depend only on the latent state and action, is a classical principle that resolves this, but enforcing it in text-based settings is an open challenge: textual latent states are discrete and non-differentiable, precluding variational training, and expressive LLM decoders readily ignore the bottleneck. We show how to make strict mediation work in the text domain. We formalize why it is necessary, showing that strict mediation makes representation quality empirically testable while history-leaky architectures break this connection. We then introduce textual latent states, which are discrete, interpretable, and variable-length, and factorized GRPO (fGRPO), a tree-structured reinforcement learning method that enforces strict mediation during training. Experiments on TextWorld and ScienceWorld show preserved one-step prediction accuracy alongside up to 57\% gains in representation quality and 98\% improvements in rollout performance, increasing with task complexity and horizon.