Technical Summary: Textual Belief States for World Models

Problem Statement

In partially observed environments, world models must maintain a latent representation that summarizes interaction history to enable prediction, planning, and reasoning. A critical challenge in modern Large Language Model (LLM)-based architectures is the identifiability problem. In many existing designs, the latent state coexists with other information sources (such as the full observation history) that can substitute for it. Consequently, a model can achieve high predictive accuracy by "bypassing" the latent state and relying directly on the history. This renders the latent state unidentifiable; predictive performance no longer constrains or reflects the quality of the learned representation.

While classical world models enforce strict latent state mediation—requiring all predictions to depend only on the latent state and the current action—applying this principle to text-based settings is difficult. Textual latent states are discrete and non-differentiable, precluding standard variational training (e.g., ELBO). Furthermore, expressive LLM decoders trained with such objectives often ignore the bottleneck, leading to posterior collapse. The paper addresses the challenge of enforcing strict mediation in the text domain to ensure that representation quality is empirically testable and that the latent state acts as a sufficient statistic for the history.

Methodology

1. Textual Latent States

The authors propose representing latent states as discrete, interpretable, variable-length text strings (specifically JSON predicate-argument lists) rather than continuous vectors. This approach offers four advantages:

• Alignment: It matches the LLM backbone, avoiding lossy modality conversions.
• Adaptive Bottleneck: Capacity scales with content rather than fixed dimensions.
• Interpretability: States are directly inspectable and can be evaluated against ground-truth facts.
• Modality Match: Observations are semantic text, so textual states preserve structural integrity.

2. Model Architecture

The system utilizes a single pretrained autoregressive language model with shared parameters $\theta$ operating in three distinct, prompt-defined stages:

1. Encoding: $\hat{s}_t \sim p_\theta(\cdot | h_t)$, generating a latent state from the history $h_t$.
2. Transition (Dynamics): $\hat{s}_{t+1} \sim p_\theta(\cdot | \hat{s}_t, a_t)$, predicting the next state based only on the current latent state and action.
3. Decoding: $\hat{o}_{t+1} \sim p_\theta(\cdot | \hat{s}_{t+1}, a_t)$, generating the next observation from the new latent state.

Crucially, once $\hat{s}_t$ is generated, the transition and decoding stages have no access to the original history $h_t$. This enforces strict mediation.

3. Training via Factorized GRPO (fGRPO)

Since textual states are non-differentiable, gradient-based optimization through the latent space is impossible. The authors treat each generation stage as a stochastic policy and optimize using Reinforcement Learning (RL).

• Limitation of Standard GRPO: Standard Group Relative Policy Optimization (GRPO) operates on a single autoregressive continuation, which would allow the model to access history through the shared prompt context, violating strict mediation.
• fGRPO Solution: The authors introduce Factorized GRPO, a tree-structured extension. For a given history, the model samples $A$ candidate latent states, then $B$ candidate next states per encoding, and $C$ candidate observations per transition. This forms a tree where rewards are computed at the leaves (comparing predicted vs. ground-truth observations) and propagated upward. This structure ensures that the transition and decoding steps are conditioned strictly on the sampled latent states, not the original history.

4. Reward Design

The reward signal combines predictive accuracy with structural regularization to encourage compact, sufficient states:

• Observation Accuracy ($R_{obs}$): Fact-level F1 score between extracted facts from the predicted and ground-truth observations.
• Format Validity ($R_{fmt}$): A binary check ensuring the latent state is a valid JSON structure.
• Auxiliary Rewards:
  • Support ($R_{sup}$): Encourages consistency between the predicted observation and the underlying state facts.
  • Deduplication ($R_{dedup}$): Penalizes semantically redundant facts.
  • Ablation ($R_{abl}$): Measures the fraction of facts whose removal degrades prediction, encouraging the retention of only necessary information.

Key Contributions

1. Formalization of Strict Mediation: The paper proves that under strict mediation, optimal predictive accuracy implies the latent state is a sufficient statistic of the history (Proposition 1). Conversely, in "leaky" architectures where history bypass is allowed, optimal accuracy can be achieved with an uninformative latent state (Proposition 2).
2. Theoretical Guarantees on Rollout: Strict mediation structurally aligns training and inference conditions. The authors show that this prevents the geometric error growth seen in history-based or leaky rollouts, enabling bounded error propagation under contraction assumptions (Proposition 3).
3. fGRPO Algorithm: The introduction of factorized GRPO, which enables training discrete, non-differentiable textual latent states under strict mediation constraints without variational objectives.
4. Empirical Validation: Experiments on TextWorld and ScienceWorld demonstrate that strict mediation yields representations that are both predictive and identifiable.

Experimental Results

The authors evaluated their approach against three baselines:

• Stateless: Predicts directly from history (upper bound on accuracy, no explicit representation).
• LeakyState: Uses a textual latent state but allows history bypass during prediction.
• StrictState-Emb: Uses strict mediation but with continuous embeddings (variational training).

Key Findings:

• Predictive Accuracy: StrictState-Text matches the one-step prediction accuracy (obs.F1) of LeakyState and Stateless baselines, confirming that strict mediation does not sacrifice performance.
• Representation Quality: Under matched prediction accuracy, StrictState-Text achieves significantly higher state-level F1 (up to 57% gains in ScienceWorld and 30% in TextWorld for complex tasks) compared to LeakyState. This confirms that leaky models learn uninformative states that rely on history bypass.
• Rollout Stability: Strict mediation significantly improves multi-step rollout performance. Gains in state-level F1 increase with the prediction horizon, reaching 98% improvements at horizon 9 on TextWorld. This validates the theoretical claim that strict mediation mitigates cumulative error propagation.
• Ablation: Removing auxiliary rewards (deduplication, ablation, support) degrades long-horizon performance, indicating these components are critical for learning compact, sufficient states.

Significance and Claims

The paper claims that strict latent state mediation is a necessary condition for making representation quality empirically testable in world models. Without it, predictive accuracy is a poor proxy for representation quality because the model can "cheat" by accessing history directly.

By successfully enforcing strict mediation in the text domain using fGRPO, the authors demonstrate that:

1. Identifiability is achievable: It is possible to learn discrete, interpretable latent states where the quality of the representation is directly diagnosable from predictive performance.
2. Textual states are effective: Variable-length textual states provide a more flexible and interpretable bottleneck than fixed-dimensional embeddings for discrete, compositional facts.
3. Stability is improved: Strict mediation aligns training and inference, preventing the error accumulation that plagues long-horizon rollouts in leaky or stateless architectures.

The authors note that their evaluation focuses on representation quality rather than downstream control or planning, and the experiments are limited to synthetic, instrumented environments (TextWorld and ScienceWorld). They acknowledge structural limitations, such as the burden of compression falling entirely on a single-shot encoder and the reliance on structured fact extraction for auxiliary rewards. However, the work establishes a methodological foundation for building reliable, interpretable world models in text-based settings.