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

2026-06-25

AI Summary

This work addresses the issue of non-identifiable representations in large language model–based world models operating in partially observable environments, where historical information often bypasses latent variables. To resolve this, the study introduces the strict mediation principle—adapted from causal inference—into the textual domain for the first time, enforcing that predictions depend solely on discrete textual latent states and actions, thereby guaranteeing identifiability. By integrating tree-structured reinforcement learning with a factorized GRPO (fGRPO) algorithm, the method learns interpretable, variable-length textual belief states in TextWorld and ScienceWorld. The approach maintains high one-step prediction accuracy while improving representation quality by up to 57% and rollout performance by as much as 98%, with gains amplifying significantly as task complexity and planning horizon increase, offering a theoretically grounded, testable framework for representation quality in textual world models.

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.

Problem

• world models
• latent state identifiability
• strict mediation
• textual belief states
• history bypass

Innovation

• strict mediation
• textual latent states
• identifiable representation
• fGRPO
• world models