Textual Belief States for World Models: Identifiable Representation Learning Under Strict Mediation

Abstract: 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.

• The paper introduces strict mediation to enforce latent state sufficiency, ensuring that textual representations effectively summarize interaction history.
• It proposes a novel factorized GRPO training method with discrete, variable-length textual latent states that prevent history bypass in predictions.
• Empirical evaluations on TextWorld and ScienceWorld demonstrate that strict mediation minimizes error propagation and boosts both state-level F1 and rollout stability.

Enforcing Identifiable Textual Belief States in World Models with Strict Mediation

Motivation and Problem Formulation

The challenge of constructing world models in partially observable text environments necessitates latent representations that summarize interaction history. Standard sequence modeling with LLMs in such contexts is hampered by the identifiability problem: model architectures frequently employ history bypass, allowing next-step predictors to access the full observation history and effectively sidestepping the latent state. As a result, predictive performance is decoupled from the informativeness of the latent state, undermining the utility of representation learning for model-based control and planning.

This paper rigorously formalizes identifiability in the context of text-based world models and proves that strict mediation—where prediction and dynamics depend exclusively on the latent state and the action—makes representation sufficiency empirically testable. This property is structurally violated in leaky or history-conditioned architectures, in which optimal prediction is compatible with trivial or uninformative latent states.

Proposed Method: Textual Latent States and Factorized GRPO

To enforce strict mediation in text world models, the authors propose:

• Textual latent states: Discrete, variable-length, natural-language state representations compatible with autoregressive LLMs, aligned with observation and action modalities, directly inspectable, and adaptable in representational capacity.
• Strict mediation architecture: The model comprises three stages—encoding the history to a textual latent state, performing action-conditioned state transitions, and decoding the next textual observation from the resulting state. No access to the original history is permitted during transition and prediction.
• Training via factorized GRPO (Generalized Reinforced Prompt Optimization): Since textual latent states are non-differentiable—obviating standard variational ELBO training—the approach relies on treating each stochastic decoding stage as a separate RL policy and optimizes rolled-out tree-structured world trajectories using a fact-level reward. The factorized design ensures transition and decoding never access the observation history, closing the leaky bypass.

Theoretical Analysis

Three core propositions are proved:

1. Predictive sufficiency under strict mediation: Observed predictive accuracy in a strictly mediated model certifies that the learned latent state is a sufficient statistic for history with respect to prediction.
2. Non-identifiability in leaky architectures: Perfect next-step accuracy can be achieved with entirely uninformative latent states, given the decoder's access to the full history.
3. Training-inference consistency and error propagation: Only strict mediation ensures training and rollout conditions are matched, yielding bounded error propagation under contractive state dynamics. In contrast, leaky/historied rollouts compound errors geometrically due to the recursive inclusion of synthetic or incorrect predictions.

Empirical Evaluation

Evaluation is conducted on two canonical text-based partially observed environments, TextWorld and ScienceWorld. The primary metrics are:

• Observation-level F1: Fact-level comparison of predicted vs. ground-truth observations.
• State-level F1: Direct evaluation of the learned textual latent states against ground-truth minimal sufficient fact sets needed to render observations.
• Compression Ratio: The efficiency of the latent representation in summarizing the history.
• Long-horizon rollout stability: Degradation in predictive performance with increasing rollout horizon.

Key Empirical Findings

• StrictState-Text achieves equal or superior one-step predictive accuracy compared to stateless (history-conditioned) and leaky architectures, confirming that strict mediation does not relegate predictive capacity.
• State F1 for strict mediation architecture outperforms leaky by up to 57% on ScienceWorld and 25% on TextWorld at high complexity, verifying that strict mediation forces the model to encode richer, more informative hidden states.
• Rollout performance gap between strict and leaky models increases with horizon (up to 98% relative improvement at horizon 9), empirically confirming the theoretical prediction concerning error propagation.
• Textual states can outperform or match continuous embeddings under strict mediation, due to their adaptive length and semantic compositionality.

Practical Implications and Future Directions

Enforcement of strict mediation in text world models makes representation quality a measurable quantity, enables interpretability via direct inspection of textual latent states, and closes the discrepancy between offline training and sequential rollout. This has profound implications for model-based RL, hierarchical planning, and sequential decision-making in text environments (including simulation, interactive fiction, and autonomous task completion via LLMs).

However, the current approach is limited to offline, instrumented, partially observable, and relatively synthetic domains. The auxiliary rewards for structuring and compressing latent state representations rely on fact extraction modules; generalization to less structured, more ambiguous, or real-world natural language settings will require alternative approaches (e.g., learned similarity metrics, LLM-based semantic evaluation). Likewise, the use of a single-shot encoder in strict mediation makes recovery from informational omissions impossible without more complex incrementally-refined or hierarchical encoders.

Areas for further research include:

• Extending strict mediation to recurrent or hierarchical state encoding schemes.
• Connecting established representation-level gains to downstream performance in long-horizon planning and counterfactual reasoning.
• Adapting the reward design for domains without structured fact extraction, possibly using self-supervised or LLM-discriminated objectives.

Conclusion

This work establishes that strict latent state mediation, rigorously enforced via a tree-structured, factorized RL objective, is both achievable and beneficial in text-based world models. The resulting identifiable, interpretable textual belief states form a direct and empirically verifiable bridge between predictive performance and representation quality—addressing a core deficiency of existing LLM-based simulators. These findings are foundational for both representation learning theory in partially observed sequential settings and for the practical development of robust, interpretable, and reliable AI agents operating in language-driven environments.

For further technical details and comprehensive empirical results, see "Textual Belief States for World Models: Identifiable Representation Learning Under Strict Mediation" (2606.27681).