World Models and Internal Representations

Overview

For agents to act effectively in complex environments, they need to build internal representations of the world -- that is, "world models." World models enable agents to internally simulate the consequences of actions without trial-and-error in the real environment. From classical symbolic representations to LeCun's JEPA, and to the implicit world models of LLMs, this article explores the theory and practice of agent internal representations.

1. What Is a World Model?

A world model is an agent's internal representation of environmental dynamics, capable of predicting:

\[ \hat{s}_{t+1} = f_{\text{world}}(s_t, a_t) \]

Given the current state \(s_t\) and action \(a_t\), the world model predicts the next state \(\hat{s}_{t+1}\).

More generally, the world model can also predict rewards:

\[ (\hat{s}_{t+1}, \hat{r}_{t+1}) = f_{\text{world}}(s_t, a_t) \]

Core Functions of World Models

Function	Description	Example
Prediction	Predict consequences of actions	"If I push this cup, it will fall"
Planning	Search for optimal action sequences in internal simulation	"First open the door, then walk over, then pick it up"
Counterfactual reasoning	Reason about situations that did not occur	"If I hadn't closed the window, the rain would have come in"
Imagination	Generate possible scenarios	"If the weather is nice tomorrow, we could go to the park"
Learning	Generate training data through internal simulation	Dreamer learns policies in "dreams"

2. Representation Types for World Models

graph TD
    A[World Model Representations] --> B[Symbolic<br/>Representation]
    A --> C[Distributed<br/>Representation]
    A --> D[Hybrid<br/>Representation]

    B --> B1[Logical Formulae]
    B --> B2[Semantic Networks]
    B --> B3[STRIPS States]

    C --> C1[Neural Network Hidden States]
    C --> C2[Embedding Vectors]
    C --> C3[Latent Spaces]

    D --> D1[Symbolic + Vector]
    D --> D2[Structured Neural Representations]
    D --> D3[LLM Knowledge + External KB]

2.1 Symbolic Representation

Describes the world using logical formulae, predicates, or state variables:

\[ s = \{\text{on}(A, B), \text{on}(B, \text{table}), \text{clear}(A), \text{holding}(\text{nothing})\} \]

Pros: Precise, composable, interpretable

Cons: Difficult to represent continuous states, challenging to handle uncertainty, limited expressiveness

2.2 Distributed Representation (Neural Networks)

Represents world states using vectors or tensors in latent spaces:

\[ z_t = \text{Encoder}(o_t) \in \mathbb{R}^d \]

\[ \hat{z}_{t+1} = \text{Transition}(z_t, a_t) \]

\[ \hat{o}_{t+1} = \text{Decoder}(\hat{z}_{t+1}) \]

Representative Methods:

Method	Core Idea	Year
World Models (Ha & Schmidhuber)	VAE + MDN-RNN	2018
Dreamer (Hafner et al.)	RSSM latent dynamics model	2020
DreamerV3	Universal world model	2023
IRIS	Transformer world model	2023

2.3 LeCun's JEPA

The Joint Embedding Predictive Architecture (JEPA) proposed by Yann LeCun is a world model that predicts in representation space rather than pixel space:

\[ \hat{z}_y = f_{\theta}(z_x, z_c) \]

where:

\(z_x = \text{Encoder}_x(x)\): Input representation
\(z_c\): Representation of conditional information (e.g., action)
\(\hat{z}_y\): Predicted target representation
Training objective: minimize \(\|\hat{z}_y - z_y\|^2\)

Key Insight: Predicting in latent space is more efficient than predicting in pixel space, because the latent space filters out irrelevant details.

Cross-Reference

For the relationship between JEPA and embodied intelligence, see Representations and World Models. For a broader discussion of world models, see World Models.

3. LLM as a World Model

3.1 The Implicit World Model of LLMs

Have LLMs implicitly learned a world model through large-scale language pre-training?

Supporting Evidence:

Othello-GPT (Li et al., 2023): A Transformer trained on move sequences internally learned a board state representation
LLMs can correctly reason about physical causal relationships (e.g., "a cup pushed to the edge of a table will fall")
LLMs can perform mental simulation

Opposing Evidence:

LLMs perform poorly on tasks requiring precise spatial reasoning
LLMs "hallucinate" non-existent facts
LLM "world knowledge" may be merely statistical associations rather than causal understanding

3.2 Formalizing the LLM World Model

Viewing the LLM as a world model operating in "language description space":

\[ \hat{d}_{t+1} = \text{LLM}(d_t, a_t^{\text{lang}}) \]

where \(d_t\) is the world state described in natural language, and \(a_t^{\text{lang}}\) is the action described in natural language.

Differences from Traditional World Models:

Dimension	Traditional World Model	LLM World Model
State space	Numerical vectors / Logical formulae	Natural language descriptions
Transition function	Learned neural network	Pre-trained LLM
Training data	Environmental interaction trajectories	Internet text corpora
Precision	Numerically precise (within training domain)	Qualitatively accurate (but may hallucinate)
Generalization	Limited (in-domain)	Broad (cross-domain but shallow)

3.3 LLM-Based Mental Simulation

Agents can use LLMs for "mental simulation" to evaluate action plans:

Given task: Tidy the room

Mental Simulation 1: Tidy the desk first → Desk is clean → Space to sort files → High efficiency
Mental Simulation 2: Sweep the floor first → Floor is clean → But desk dust will fall again → Low efficiency
Conclusion: Should tidy the desk first, then sweep the floor

This "imagination" capability corresponds to forward search in planning:

\[ \text{score}(a) = \text{LLM\_evaluate}(\text{simulate}(s, a)) \]

4. World Models for Planning

4.1 Model-Based Planning Framework

graph TD
    S[Current State s_t] --> WM[World Model]
    A1[Action Candidate a1] --> WM
    A2[Action Candidate a2] --> WM
    A3[Action Candidate a3] --> WM
    WM --> S1[Predicted State s'_1]
    WM --> S2[Predicted State s'_2]
    WM --> S3[Predicted State s'_3]
    S1 --> EVAL[Evaluation Function]
    S2 --> EVAL
    S3 --> EVAL
    EVAL --> BEST[Select Best Action]

4.2 World Model Applications in LLM Agents

Application Scenario	Method	Description
Task planning	Mental simulation	Imagine the result of executing each subtask
Risk assessment	Counterfactual reasoning	"What happens if this operation fails?"
Plan comparison	Multi-path simulation	Simulate multiple plans and select the best
Safety checking	Consequence prediction	"Will this command delete important files?"
User modeling	Theory of Mind	Predict user reactions and needs

4.3 Insights from the Dreamer Paradigm

Dreamer (Hafner et al., 2020) trains policies in "dreams" (the world model), reducing the need for real environment interaction:

\[ \pi^* = \arg\max_{\pi} \mathbb{E}_{z_0 \sim q} \left[ \sum_{t=0}^{H} \gamma^t r(z_t, a_t) \right] \]

where state transitions occur entirely within the world model.

Implications for LLM Agents:

Complex operations can be rehearsed in the LLM's "imagination," executing only when confident
Similar to internal evaluation in Tree of Thoughts
Reduces unnecessary external tool calls (saving cost and time)

5. Evaluating Internal Representations

How to evaluate the quality of an agent's world model?

5.1 Prediction Accuracy

\[ \text{Error} = \mathbb{E}\left[\|s_{t+1} - \hat{s}_{t+1}\|^2\right] \]

5.2 Planning Support Capability

The value of a world model is ultimately reflected in planning quality:

\[ \text{Value} = R(\text{plan with world model}) - R(\text{plan without world model}) \]

5.3 Evaluating LLM World Models

Evaluation Dimension	Method	Example
Physical intuition	Physical reasoning benchmarks	"Does a ball rolling down a slope speed up or slow down?"
Causal reasoning	Counterfactual evaluation	"If it hadn't rained, would the grass still be wet?"
Spatial reasoning	Navigation tasks	"Walk two steps north, three steps east -- where am I now?"
Social cognition	Theory of Mind tests	"Sally doesn't know the ball was moved"
Temporal reasoning	Event ordering	"First boil water, then brew tea -- cannot reverse"

6. Future Directions

Causal world models: Moving from statistical associations to causal understanding
Multimodal world models: Integrating visual, linguistic, tactile, and other modalities
Hierarchical world models: State representations at different abstraction levels
Learnable world models: Continuously improving from interaction experience
Composable world models: Combining models from different domains

References

Ha, D. & Schmidhuber, J. (2018). World Models. arXiv:1803.10122.
Hafner, D. et al. (2020). Dream to Control: Learning Behaviors by Latent Imagination. ICLR 2020.
Hafner, D. et al. (2023). Mastering Diverse Domains through World Models. arXiv:2301.04104.
LeCun, Y. (2022). A Path Towards Autonomous Machine Intelligence. OpenReview.
Li, K. et al. (2023). Othello-GPT: Language Models Are Able to Infer World States. ICLR 2023.
Hao, S. et al. (2023). Reasoning with Language Model is Planning with World Model. EMNLP 2023.