General agents need world models

Thank you for your detailed and thoughtful review and helpful comments. We hope to address your main concerns about the core claims of our paper, which we believe stem from a misunderstanding of theorem 1, and have implemented your suggestions for improving the paper.

What do we actually show? Your review notes that we prove a bound on the agent's ability to achieve zero-shot generalization and from this conclude that to do zero-shot generalization, agents must learn world models. We agree that if this was our result, it would be highly questionable.

What we actually do is assume the agent satisfies this regret bound (rather than prove it does). I.e. we assume the agent has some minimal degree of competence at zero-shot learning. We then formally prove (as verified by reviewer Xhe1) that for any agent satisfying this assumption, a world model is encoded in the agent's policy. We derive an (unsupervised) algorithm for recovering this world model from the policy (similar to [4]), and prove an error bound on the accuracy of the world model recovered, which depends on the agent's regret.

We agree that this could be made clearer in the paper, and have done so by specifying the algorithm that recovers the world model outside of the proof of Theorem 1 (see Algorithm 1 in response to reviewer aY71) and re-writing the results section to improve clarity and discussion of our results.

Is this obvious? The question of whether or not AI systems have or need world models is hotly debated (see for example [1]), and the subject of significant empirical research [2,3]. A similar result to ours [4] was recently proven for domain generalization (rather than zero-shot learning) and received an award at ICLR 2024.

While we agree it feels intuitive that agents should have a world model, there are many ways agents can can reach goal states without one; e.g. through numerous heuristics (schemas, similarity-based reasoning, ...), online learning, etc. Humans can switch between using model-based reasoning or heuristics to generalize to new tasks, depending on the situation. And many biological agents that are thought to be purely model-free (`stimulus-response' agents [5]). It is unclear if general AI systems like LLMs have world models, or if they can generalize purely via heuristics. Before now, there was no formal result showing that world models are necessary for generalization. And indeed we show that for myopic tasks, where the agent is optimizing for immediate outcomes, world models are not necessary.

We tackle this question by tying world models to a key capability---zero-shot learning. This has consequences for how we design agents (model-based v.s. model free) and reveals fundamental limitations on agent capabilities. Further, we show that this world model can be extracted, which has consequences for safety and interpretability.

How have we improved the paper based on your feedback

1. We agree with your point that the paper lacks insights into theorem 1 and the procedure for extracting world models. We now give an explicit algorithm for recovering world models (see response to reviewer aY71), and have extended the discussion of theorem 1 and other results to 2 pages and reduced repetition.

2. We have included experiments demonstrating that our algorithm can be applied to real world agents (see response to reviewer aY71)

3. We have introduced a new theorem which proves that learning a world model is not necessary for myopic agents, which optimize for immediate outcomes to their actions (depth-1 goals). This relates to your question on if the need for a world model is obvious or not. Theorem and proof can be seen here https://imgur.com/a/QzrXt0W

Reviewer questions

1. Why is it a generative model?. The world model we recover can be used to simulate environment trajectories. This is opposed to the purely state-based world models often studied (e.g. in [2]).
2. Three temporal operators... Typo corrected. We were referring to the trivial (Now) operator, but have removed this.
3. Should eq. 2 rhs be $\pi^*$ We are taking the max over $\pi$, which is equivalent.
4. What is $p$ in theorem 1. Typo corrected.
5. Comments on Brookes et. al. 2024, have removed.
6. Line 206 rh column, I don't see what this prob ratio refers to?. This is the ratio of the estimated transition probability to the true value, i.e. the relative error. It is given by dividing the inequality in theorem 1 by $P_{ss'}(a)$

[1] https://x.com/ylecun/status/1667947166764023808

[2] Li, et al. "Emergent world representations: Exploring a sequence model trained on a synthetic task." ICLR (2023) oral

[3] Gurnee et. al. "Language Models Represent Space and Time." ICLR (2024)

[4] Richens et. al. "Robust agents learn causal world models". ICLR (2024) oral

[5] Tomasello. The evolution of agency: Behavioral organization from lizards to humans. MIT Press, 2022.

Thank you for you helpful comments. As noted by Reviewer Xhe1, our paper does propose a new method for eliciting world models from agents. However, this was quite unclear in the submitted draft, and we have included an explicit algorithm (below) in the manuscript to clarify this.

Following you recommendation we have made the following changes to the paper;

1. Explicit algorithms for recovering world models from agents (below)
2. New experiments, validating these algorithms on real agents (details below).
3. Discussion of real-world tasks where our results can be applied (for example [1] recently developed goal conditioned agents that can generalize zero-shot to arbitrary linear temporal logic goals).

[1] Jackermeier et al. "DeepLTL: Learning to Efficiently Satisfy Complex LTL Specifications for Multi-Task RL." ICLR 2025 (oral)

New experiments

Motivation: can our algorithm for recovering an agent's world model (below) be applied to real-world agents that perhaps maximally violate our theoretical assumptions? Namely, the strict regret bound we assume in Theorem 1.

Experimental setup

Our experiment involves extracting a world model from a model-based language agent in 120 randomly generated cMDP environments using our algorithm (below). We show the agent strongly violates our assumptions, but nonetheless this algorithm can still recover the agent's world model.

1. 120 randomly generated environments described by cMPs, with between 5 and 40 states and 3 and 20 actions.
2. Our goal-conditioned agent is an LLM (Gemini Flash 2.0), with an explicit, private world model.
3. We then attempt to learn this private world model using Algorithm 1 (above) given only the agent's policy

A figure of our results can be viewed at the following URL: https://imgur.com/a/1gNe15c

We also note that our paper is a theory paper, which as pointed out by reviewer Xhe1 significantly extends important recent theory work. We hope it can be judged on these merits, without requiring experiments that extend on the current state of the art empirical work (for example [1] trained LTL conditioned agents and was an oral at ICLR 2025, and used environments simpler than Atari).

Algorithm 1: Estimate Transition Probability $\hat{P}_{ss'}(a)$ from Policy $\pi$

Input:

• Goal-conditioned policy $\pi(a_t | h_t; \psi)$
• Choice of state $s$, action $a$, outcome $s'$
• Precision parameter $n \in \mathbb{N}$ (related to maximum goal depth $2n+1$)
• An alternative action $b \neq a$

Function: EstimateTransitionProbability($\pi, s, a, s', n, b$)

1. Initialize $k^* \gets n$
2. For $k = 1$ to $n$:
   • Define base LTL components:
     • $\varphi_0 \gets [A_0=a]$ (Take action $a$)
     • $\varphi'_0 \gets [A_0=b]$ (Take action $b$)
     • $\varphi_1 \gets ◇ [A=a, S = s]$ (Transitions eventually to state $s$ and takes action $a$)
     • $\varphi_2 \gets ○ [S=s']$ (Transition Next to state $s'$)
     • $\varphi'_2 \gets ○ [S\neq s']$ (Transition Next to any state other than $s'$)
   • Define composite goal:
     • $\psi_0 \gets \langle\varphi_1, \varphi_2'\rangle$ (Sequential goal labelled Fail)
     • $\psi_1 \gets \langle\varphi_1, \varphi_2\rangle$ (Sequential goal labelled Success)
     • $\psi_a(k,n) \gets \bigvee_{\text{sequences with } r \le k \text{ successes}} \langle \varphi_0, (\psi_0 \text{ or } \psi_1)_{\times n} \rangle$
     • $\psi_a(k,n) \gets \bigvee_{\text{sequences with } r > k \text{ successes}} \langle \varphi_0', (\psi_0 \text{ or } \psi_1)_{\times n} \rangle$
   • $\psi_{a,b}(k,n) \gets \psi_a(k,n) \vee \psi_b(k,n)$
   • $a_0 \gets \pi(a_0 | s_0; \psi_{a,b}(k,n))$ (Query the policy for the first action)
   • If $a_0 = a$:
     • $k^* \gets k$
     • break (Found smallest $k$ s.t. where agent prefers goal involving $\le k$ successes)
3. Estimate $\hat{P}_{ss'}(a) \gets (k^*-1/2)/n$
4. Return $\hat{P}_{ss'}(a)$

Thank you for thoroughly reviewing our paper and for your helpful comments. In particular, the following comment on the need to clarify the size of the `universal’ set of goals, compared to the set of finite horizon trajectories.

Reviewer: It not obvious that the set of 'universal' goal-directed tasks is small compared to the set of finite trajectories. I would be happy to raise my score if the authors provide insights on this aspect.

You correctly point out that we assume the agent can generalize to any composite goal of depth n (denoted $\mathbf{\Psi}_n$). What we failed to make clear in the paper is that the proof of Theorem 1 actually only requires that the agent can generalize to only a very small subset of $\mathbf{\Psi}_n$. This is what led us to comment on the existence of small `universal’ goal sets, which are sufficient for learning a world model and hence generalizing to more complex goals.

Precisely, the set of goals we require the agent to generalize to is $\\{ \psi_{a, b}(k, n) \\}_{k=1}^{k=n}$

where $\psi_{a, b}(k, n)$ are described at the start of the proofs of Lemma 5 and Theorem 1. The cardinality of this set is n, which is much smaller than $\mathbf{\Psi}_n$, and in general much smaller than the number of finite time trajectories, which scale as $\mathcal{O}(|S|^T |A|^{T-1})$ for horizon $T$. We can interpret this as growing exponentially with the goal depth n, as satisfying our sequential goals of depth n requires trajectories at least of length $T \geq n/2$ for these goals.

We have made this all clear in the paper, and thank the reviewer for pointing it out.

Typo corrected on p.5. second sentence.