Temporal State and Memory

Michael BrenndoerferJune 16, 202656 min read

Part of World Models Handbook

Explains how world models maintain state under partial observability using recurrent memory, gated LSTMs and GRUs, attention, episodic storage.

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Temporal State and Memory

So far in Part III we have built up a picture of what a world model should represent: observations and sensors in Part III, Chapter 1, sufficient state in Chapter 2, and latent variables in Chapter 3. Each of those chapters answered a question about the content of the model's state at a single moment in time. This chapter addresses a question that is easy to state and surprisingly hard to answer in practice: what should the model's state be when the world's state cannot be recovered from a single observation? This is not an edge case. Many environments an agent cares about, from a robot navigating a building to a language model reading a long document to a game agent with fog of war, hide information from the current observation and store it, instead, in the past.

You already know the general formal answer from Part II: optimal prediction and control under partial observability may require the history of observations and actions, or a task-sufficient state distilled from that history. A Bayesian belief is one such state when the environment model is known. But formal answers do not build themselves. Saying "condition on the history" is a specification, not an implementation. Someone must decide how the history is stored, how far back it reaches, how quickly it can be updated, and what gets forgotten. Those decisions can determine whether an agent can solve the task, and they involve real tradeoffs in computation, memory capacity, and learnability. Those decisions are the subject of this chapter. We will work through recurrent memory, the distinction between short-term and long-term state, episodic and external memory, and the diagnostic failure mode that motivates all of it: state aliasing in non-Markov environments.

Along the way, keep one distinction from Part I, Chapter 2 firmly in view. A world model makes predictions; an agent uses those predictions to act. Memory lives in both places, and the requirements differ. A predictor may need to remember a cue seen a thousand steps ago to forecast an outcome; a controller needs that same memory available now, at decision time, in a form it can act on. This difference in requirements is subtle but consequential. A memory that preserves information faithfully but buries it in a high-dimensional vector that no policy can read from cheaply, or one that is fast to read but discards exactly the facts that distinguish good actions from bad ones, will pass some evaluations and fail others. A memory representation that serves one role but not the other is not yet a world model in the sense of Part I, Chapter 3. Keep asking, as we go: does this memory mechanism serve prediction, decision, or both?

Why Stateless Models Break: A Corridor in the Dark

Imagine an agent entering a long, unlit corridor. At the entrance, it glimpses a sign: an arrow pointing left or right. The sign indicates which branch at the far end of the corridor hides a reward. Then the agent walks forward for fifty steps through a featureless corridor: every observation looks identical. At the end, it must turn left or right.

This scenario is deliberately minimal, and its minimality is what makes it diagnostic. Nothing about it is exotic: no adversarial structure, no complicated dynamics, no partial rewards. The only complication is that one piece of decision-critical information appears once and then vanishes from the observation channel. That is enough to break every stateless model.

During the walk, the agent's current observation carries no information about the correct choice. Yet the correct choice is fully determined by the agent's history, specifically by the sign at step zero. We can make the failure of stateless models precise by splitting the prediction problem into two parts that a stateless model handles very differently:

  • Next-observation prediction is easy without memory. In the corridor, the next observation has the same wall mean plus independent sensor noise. A Markov predictor can recover that mean and incur only the irreducible noise error. A dynamics-only evaluation can therefore report low one-step error, and short rollouts may not expose the missing branch cue.
  • The realized branch cannot be identified without memory. The current wall observation supports only the prior over branches, which is 50/50 in this example. Expected rewards and values are still well-defined once a future policy and information set are specified, but they average over an unresolved branch rather than identify the episode's correct terminal choice. No amount of processing of the wall observation alone can recover the missing episode-specific cue.

This asymmetry is the central lesson. Partial observability does not always show up as prediction error in the near future. It shows up as irreducible uncertainty about distant outcomes and about which decisions are correct. A model evaluated only on one-step dynamics can be confidently, catastrophically wrong about everything that matters for control. This is one more instance of the recurring theme of this handbook: open-loop prediction quality is not the same as closed-loop decision usefulness. The corridor is the cleanest possible demonstration of that gap. In less stylized systems, the same issue can appear either as an abrupt loss of a cue or as gradual decay across many partially informative observations.

State aliasing

Two underlying situations are aliased when they produce the same observation, or the same observation distribution, but demand different behavior or imply different futures. In the noisy corridor, the two branch states have statistically indistinguishable wall observations but opposite correct terminal choices.

The rest of the chapter is about building state representations that un-alias such situations, and about the engineering tradeoffs involved in doing so at scale. The strategy is always the same in outline: carry forward some trace of the past, and make sure the trace survives until it is needed. The differences among approaches lie in what gets carried, for how long, at what cost, and how reliably the carrying can be learned.

History as State

The natural response to partial observability is to condition on the observation history itself. This section examines why that fallback is always safe but never practical, and what a usable replacement must preserve.

Where the Markov Assumption Fails

In Part II, Chapter 1, you met the Markov assumption: the current state is a sufficient statistic for the future. In a Markov decision process, p(st+1∣st,at)p(s_{t+1} \mid s_t, a_t) captures everything, and no amount of additional history can improve a prediction conditioned on sts_t. Partially observable settings break this at the level of observations. The environment's underlying state sts_t may still be Markovian; the problem is that the agent sees only an observation ot∼p(ot∣st)o_t \sim p(o_t \mid s_t), and the map from states to observations is many-to-one. Many distinct underlying states collapse onto the same observation, and once the collapse happens, the agent cannot tell which of them it is in. Note carefully that the difficulty is not that the world is unpredictable; it is that the agent's view of the world is lossy.

Given a full history τt=(o0,a0,o1,a1,…,ot)\tau_t = (o_0, a_0, o_1, a_1, \ldots, o_t), two facts hold:

  1. The history is always enough. Conditioning on τt\tau_t is never worse than conditioning on oto_t, because the history contains the current observation. Any quantity of interest, such as p(ot+1∣τt)p(o_{t+1} \mid \tau_t) or the value V(τt)V(\tau_t), is well-defined. In a precise sense, history is the universal fallback: when in doubt about whether some information matters for the future, conditioning on history is always safe.
  2. The history is usually too much. It grows without bound, most of it is redundant, and downstream models (policies, planners, predictors) cannot consume an ever-growing sequence efficiently. An agent that reprocesses its entire history at every timestep pays a cost that grows linearly with episode length at best, and it spends most of that compute re-deriving conclusions it has already reached.

So the practical design problem is to compute a useful state btb_t from history under a finite memory budget. When an exact finite sufficient statistic exists, we would like to recover it. Otherwise, the goal is an approximation that preserves the distinctions relevant to the predictions and decisions we care about. A corridor where only one bit matters is easy; a world where thousands of features of the past can each matter is not. In Part II, Chapter 2 you saw the classical answer for a finite-state POMDP with known transition and observation models: the belief state, a probability distribution over underlying states, updated recursively by Bayesian filtering. That belief is a fixed-dimensional sufficient statistic for optimal control, though it need not be the smallest possible representation. General continuous-state beliefs need not have an exact finite-dimensional parameterization. Here we tackle the learned case: the models are unknown, and a useful history summary must be learned from data. We give up the guarantees of exact Bayesian filtering in exchange for applicability where we cannot write down the transition and observation models.

Memory, defined operationally

For this handbook, memory is any mechanism that turns history τt\tau_t into a state or store that makes past information available to later predictions and decisions. A recurrent hidden vector is fixed-size; a sliding window or attention cache is bounded by its context; an episodic database may grow with experience. Nothing in this definition requires the mechanism to look biological. The definition is deliberately behavioral: judge memory by the predictions and decisions it supports, while keeping its storage bound, write policy, and retrieval mechanism explicit.

What Must Survive the Compression

It is worth asking what information a memory must preserve. Three kinds, with different lifetimes, and the differences among them drive most of the architectural choices in this chapter:

  • Cue information. One-time signals (the corridor sign) whose consequences persist. The memory must retain them, unchanged, for as long as they matter. Failure here is silent and total: a forgotten cue is indistinguishable from a cue that never existed, and no amount of clever downstream processing can recover it.
  • Integration information. Quantities built up over many steps: accumulated velocity, running averages, elapsed time since an event, total charge drained from a battery. The memory must aggregate. No single observation contains these quantities; they exist only as summaries, and the memory mechanism is the only place they can live.
  • Recent-context information. Short-lived but high-bandwidth: the last few frames of a video, the recent trajectory of a joint. The memory must carry this faithfully but can let it decay. What matters is fidelity over the last handful of steps, not persistence.

These three requirements pull in opposite directions. A single scalar hidden state cannot serve all three well. Cue retention wants constancy over time; integration wants accumulation; recent context wants high capacity and fast turnover. A mechanism optimized for one, for example a heavily damped accumulator, is actively bad at another, for example a buffer that must reflect the latest frame before it settles. This tension is why memory architectures end up with multiple compartments, which is the subject of the next two sections. It also connects back to Part III, Chapter 2: the question "which features of history are sufficient?" is the temporal version of the sufficient-state question, and the answer is a set of dynamical roles, not a single snapshot.

Recurrent Memory

Recurrent networks are a standard learned answer to the compression problem: a fixed-size state updated one step at a time. This section covers how recurrence works, why plain recurrence can struggle over long horizons, and how gated architectures provide a more direct retention path.

Recurrence as Learned Compression

The workhorse solution is a recurrent state: a fixed-size vector hth_t updated step by step,

ht=fθ(ht−1, ot, at−1),h_t = f_\theta(h_{t-1},\, o_t,\, a_{t-1}),

where:

  • hth_t: the recurrent state at time tt, a fixed-size vector summarizing the history
  • ht−1h_{t-1}: the previous recurrent state, carrying everything absorbed so far
  • oto_t: the current observation
  • at−1a_{t-1}: the action taken at the previous step
  • fθf_\theta: a neural network, with parameters θ\theta, that decides how to merge the new observation and the last action into the running summary

Everything the model "knows" about the past lives in hth_t; the history itself is discarded after being absorbed. This is the compression decision in its purest form: one fixed-size vector stands in for an arbitrarily long past, and the network fθf_\theta is the sole arbiter of what makes it into the summary. Prediction then proceeds from the recurrent state:

o^t+1=gϕ(ht,at),r^t+1=qϕ(ht,at),\hat{o}_{t+1} = g_\phi(h_t, a_t), \qquad \hat{r}_{t+1} = q_\phi(h_t, a_t),

where:

  • o^t+1\hat{o}_{t+1}: the predicted next observation (the hat marks a prediction, as opposed to the observed ot+1o_{t+1})
  • r^t+1\hat{r}_{t+1}: the predicted next reward
  • gϕg_\phi: the observation head, a network with parameters ϕ\phi mapping the recurrent state and a candidate action to the next-observation prediction
  • qϕq_\phi: the reward head, mapping the same inputs to the predicted reward
  • hth_t: the recurrent state defined above
  • ata_t: the action whose consequences are being predicted

Notice that the prediction heads condition on both the recurrent state and the proposed action. They therefore produce action-conditional predictions, and a planner can compare candidate actions by running the heads. Interpreting those predictions as causal consequences additionally requires the data and model to identify the effects of intervening on the action. A policy or planner can operate on the same hth_t. This is exactly the structure of the recurrent state-space models we will build in detail in Part V, Chapter 1; here we focus on the memory inside the recurrence, not the full world-model architecture.

Conceptually, a recurrent network is doing online compression: at every step it decides, implicitly, which parts of the new observation to merge into the running summary and which parts of the old summary to discard. The compression is shaped by the training objective and data. When trained for the role, hth_t can approximate a belief-like filter, a running integral, or a short-term buffer. The catch is that nothing in the architecture guarantees retention over the task horizon. Gated architectures provide a more direct retention path, but optimization must still learn when to use it.

The Long-Range Gradient Problem

To see why plain recurrence struggles, consider what training demands. If the correct action at step TT depends on a cue at step 00, then the loss at step TT must send gradients back through every intermediate hidden state. Learning to preserve the cue means learning weight values for which the cue's influence survives that chain, and whether such weights are discoverable by gradient descent depends on the geometry of the chain itself. For a simple recurrent update ht=tanh⁡(Wht−1+Uxt)h_t = \tanh(W h_{t-1} + U x_t), the derivative of the final state with respect to the initial state is a product of Jacobians:

Jt=∂ht∂ht−1,∂hT∂h0=JTJT−1⋯J1.J_t = \frac{\partial h_t}{\partial h_{t-1}}, \qquad \frac{\partial h_T}{\partial h_0} = J_T J_{T-1} \cdots J_1.

where:

  • h0h_0: the hidden state at the step where the cue arrives
  • hTh_T: the hidden state at the step where the loss is computed
  • Jt=∂ht/∂ht−1J_t = \partial h_t / \partial h_{t-1}: the Jacobian of recurrent step tt, the matrix of partial derivatives describing how each component of ht−1h_{t-1} influences each component of hth_t
  • JTJT−1⋯J1J_T J_{T-1} \cdots J_1: the ordered product of these TT matrices, with the latest step on the left, chaining the cue's influence from step 00 to step TT
  • xtx_t: the input at step tt (for a world model, typically the current observation together with the last action)
  • WW: the recurrent weight matrix, multiplying the previous hidden state ht−1h_{t-1}
  • UU: the input weight matrix, multiplying xtx_t

For matrix-valued, time-varying Jacobians, growth depends on the product's singular directions as well as the norm of each factor. A useful bound is

∥JTJT−1⋯J1∥≤∏t=1T∥Jt∥.\left\lVert J_T J_{T-1} \cdots J_1 \right\rVert \leq \prod_{t=1}^{T} \left\lVert J_t \right\rVert.

Thus uniformly contractive factors are sufficient for decay, while factors with norm above one make growth possible but do not force it: directions can rotate, contract, or cancel. The scalar constant-gain model below isolates the basic effect. If every step multiplies a scalar signal by ρ\rho, its magnitude after TT steps is ∣ρ∣T|\rho|^T: it decays when ∣ρ∣<1|\rho| < 1, stays constant when ∣ρ∣=1|\rho| = 1, and grows when ∣ρ∣>1|\rho| > 1. Over 50 steps, 0.950≈0.005150.9^{50} \approx 0.00515, about 0.515 percent of the initial magnitude.

We can make the scalar approximation concrete with a simple computation. The plot shows decay (ρ=0.9\rho = 0.9), constancy (ρ=1.0\rho = 1.0), and growth (ρ=1.1\rho = 1.1) over 50 steps. It illustrates repeated scalar gain, not the exact behavior of every recurrent Jacobian product.

In[3]:
Code
import numpy as np

horizons = np.arange(1, 51)
rho_values = [0.9, 1.0, 1.1]
grad_norms = {rho: rho ** horizons.astype(float) for rho in rho_values}
Out[4]:
Visualization
Log-scale line chart of gradient magnitude versus horizon for rho values 0.9, 1.0, and 1.1.
Illustrative scalar gradient magnitude rho^T across a recurrent chain. With rho = 0.9 the signal falls below one percent within fifty steps; rho = 1.0 remains constant; and rho = 1.1 grows. Matrix-valued recurrent Jacobians need not follow one fixed scalar gain.

In this scalar illustration, a ten-percent change around one produces orders-of-magnitude differences over a corridor-length span. Gated additive updates are a widely used way to give optimization a controllable near-identity path. They make long-range retention easier to learn, but they are not necessary: other recurrent parameterizations and explicit memory mechanisms can also preserve long dependencies.

Without such a path, the network may need to learn retention through the same long Jacobian chain that attenuates its training signal.

Gating: Additive Memory Channels

Long short-term memory (LSTM) networks and gated recurrent units (GRUs) attack the problem by changing the shape of the update. Instead of always overwriting the state through a squashing nonlinearity, they introduce an additive channel guarded by learned gates. The key insight is that an additive update, one that adds new content to old content rather than replacing it, provides a gradient path whose gain can be held near one, so information can flow across many steps without exponential decay. For an LSTM, the core equations are:

it=σ(Wixt+Uiht−1+bi)(input gate)ft=σ(Wfxt+Ufht−1+bf)(forget gate)c~t=tanh⁡(Wcxt+Ucht−1+bc)(candidate)ct=ft⊙ct−1+it⊙c~t(cell state)ωt=σ(Woxt+Uoht−1+bo)(output gate)ht=ωt⊙tanh⁡(ct)\begin{aligned} i_t &= \sigma(W_i x_t + U_i h_{t-1} + b_i) & &\text{(input gate)}\\ f_t &= \sigma(W_f x_t + U_f h_{t-1} + b_f) & &\text{(forget gate)}\\ \tilde{c}_t &= \tanh(W_c x_t + U_c h_{t-1} + b_c) & &\text{(candidate)}\\ c_t &= f_t \odot c_{t-1} + i_t \odot \tilde{c}_t & &\text{(cell state)}\\ \omega_t &= \sigma(W_o x_t + U_o h_{t-1} + b_o) & &\text{(output gate)}\\ h_t &= \omega_t \odot \tanh(c_t) & & \end{aligned}

where:

  • xtx_t: the input at step tt, for a world model typically the current observation together with the last action
  • ht−1h_{t-1}: the previous hidden state, the network's working output
  • ctc_t: the cell state, the protected long-term memory carried across steps
  • c~t\tilde{c}_t: the candidate value, a proposal for new information to write into the cell state
  • iti_t: the input gate, controlling how much of the candidate is written into the cell state
  • ftf_t: the forget gate, controlling how much of the previous cell state survives
  • ωt\omega_t: the output gate, controlling how much of the cell state is exposed as the hidden state (written ωt\omega_t rather than the conventional oto_t to avoid clashing with the observation symbol used elsewhere in this chapter)
  • hth_t: the new hidden state, the gated readout of the cell state
  • σ\sigma: the logistic sigmoid, which maps any real number to (0,1)(0, 1) so each gate acts as a soft switch
  • ⊙\odot: elementwise multiplication
  • Wi,Wf,Wc,WoW_i, W_f, W_c, W_o and Ui,Uf,Uc,UoU_i, U_f, U_c, U_o: weight matrices acting on the input and on the previous hidden state; bi,bf,bc,bob_i, b_f, b_c, b_o: their bias vectors

Look at the cell-state update ct=ft⊙ct−1+it⊙c~tc_t = f_t \odot c_{t-1} + i_t \odot \tilde{c}_t. It is a residual update: when the forget gate ftf_t is near 11 and the input gate iti_t near 00, information flows through the direct cell-state branch nearly unchanged. Across several such steps, the direct gain is the product of the forget gates, so quantitative retention still depends on how close each gate remains to one. In the corridor, the desired behavior is to write the cue at step 0, preserve it during the featureless walk, and expose it at the decision point. Because the gates depend on the input and hidden state, gradient-based training can learn a write-and-retain policy, provided the objective supplies a useful signal.

The GRU is a streamlined variant with a single update gate ztz_t and a reset gate rtr_t:

zt=σ(Wzxt+Uzht−1)(update gate)rt=σ(Wrxt+Urht−1)(reset gate)h~t=tanh⁡(Wxt+U(rt⊙ht−1))(candidate)ht=(1−zt) ht−1+zt h~t(new state)\begin{aligned} z_t &= \sigma(W_z x_t + U_z h_{t-1}) & &\text{(update gate)}\\ r_t &= \sigma(W_r x_t + U_r h_{t-1}) & &\text{(reset gate)}\\ \tilde{h}_t &= \tanh(W x_t + U (r_t \odot h_{t-1})) & &\text{(candidate)}\\ h_t &= (1 - z_t)\, h_{t-1} + z_t\, \tilde{h}_t & &\text{(new state)} \end{aligned}

where:

  • ztz_t: the update gate, blending how much of the old state is kept versus replaced by the candidate
  • rtr_t: the reset gate, controlling how much of the old state the candidate can consult
  • h~t\tilde{h}_t: the candidate hidden state, the proposed new content
  • Wz,Wr,WW_z, W_r, W and Uz,Ur,UU_z, U_r, U: weight matrices acting on the input and on the previous state (biases omitted for brevity)
  • σ\sigma and ⊙\odot: the logistic sigmoid and elementwise multiplication, as in the LSTM equations

Notice that the GRU hidden state itself now has the residual form: the update ht=(1−zt)ht−1+zth~th_t = (1 - z_t) h_{t-1} + z_t \tilde{h}_t is an interpolation between keeping and replacing, so the direct "keep" branch has gain 1−zt1 - z_t. The GRU merges the LSTM's separate cell and hidden states into one, using three affine gate/candidate groups where the displayed LSTM uses four. Which architecture performs better is task-dependent. The LSTM's output gate controls how much cell content is exposed as hth_t, but it does not isolate the cell completely: hth_t feeds the next step's gates and candidate, so noisy exposed state can still influence ct+1c_{t+1} indirectly.

Fast and slow state

It is often useful to give the model two coupled recurrent states: a fast state updated every step (carrying recent context) and a slow state updated rarely or with strong damping (carrying cues and integrals). This is the architectural answer to the tension we identified in "What Must Survive the Compression": rather than forcing one vector to serve constancy, accumulation, and turnover simultaneously, give each dynamical role its own compartment and couple them. Hierarchical recurrent models make this split explicit, updating the slow state every kk steps or on event boundaries, which also shortens the gradient path each timescale must span. The same idea reappears in modern state-space models with multi-timescale dynamics, covered in Part V, Chapter 2.

Training Practicalities That Shape Memory

Two training choices quietly influence how good the learned memory is. They are easy to overlook because they do not appear in the architecture diagram:

  • Truncated backpropagation through time (TBPTT). Unrolled gradients are cut at a window of length KK, so a loss cannot assign credit directly through more than KK recurrent transitions. This is not a hard bound on forward memory: hidden state can be carried across windows, and shorter-range or auxiliary losses can teach a retention rule that generalizes farther. But a terminal-only 50-step dependency receives no direct cue-to-loss gradient when K=20K=20, which can make it difficult to learn. Treat KK as a credit-assignment limit and test retention at the horizon the task needs.
  • Scheduled or curriculum length. Training first on shorter episodes and then extending the horizon can help optimization. One plausible explanation is that shorter examples first establish useful update behavior, which longer examples then refine. That mechanism is a hypothesis unless an ablation separates it from data ordering, regularization, and other optimization effects.

A related trap: if observations occasionally contain the needed cue deterministically and observably, the model can learn memory. If the cue is buried in noise, or appears only under rare action sequences, the data distribution itself may never teach the gates to open. Memory learning is a joint property of architecture, gradient budget, and data distribution. When a memory model fails, before concluding that the architecture is inadequate, check whether the training signal demanded retention, and check whether the truncation window ever spanned the dependence in question.

Short-Term Versus Long-Term State

The word memory covers two jobs with different requirements: filtering noisy observations into a current estimate, and retaining older information that still matters. This section separates the two jobs and surveys designs that give each its own mechanism.

Two Problems, Not One

The phrase "memory" hides a useful distinction. Short-term state answers: what is happening right now, given that observations are noisy, partial, or delayed? This is a filtering problem, and the recurrent state at short horizons is a learned filter, the amortized cousin of the Bayesian filters of Part II, Chapter 2. Like those filters, its job is to fuse a stream of noisy observations into a current estimate, and like them, it should be judged on the accuracy of that estimate at the present moment. Long-term state answers: what happened in the past that still matters? This is a retention problem, and it has different engineering requirements: persistence, capacity, and addressing. A filter is judged on responsiveness and noise rejection; a retention mechanism is judged on whether it holds content without corruption over long spans. Mechanisms optimized for one are frequently poor at the other.

Conflating the two causes predictable failures, and it is worth spelling out the mechanism behind each:

  • A per-step filtering objective can dominate a sparse delayed-retention objective when its aggregate gradient is larger. Loss scales, sampling, architecture, and optimization determine whether this happens, so compare the objectives rather than assuming a GRU will always favor one.
  • Conversely, a model with only slow accumulators cannot track fast-moving objects or contacts. If the entire state is heavily damped, the freshest information is smeared across many steps, and the model reacts to where the world was, not where it is.

Practical designs therefore separate timescales deliberately, giving each compartment the write policy and update rate its content requires:

  • Context windows. A raw buffer can preserve its stored tensors exactly for the last LL steps, then evicts them at a hard cutoff. An attention context keeps projected token representations separately addressable within that span, but encoding and learned retrieval can still lose information.
  • Recurrent summaries. A GRU/LSTM state compresses absorbed history into a fixed-size state and forgets according to its learned update. Known running statistics often admit compact recurrent updates; arbitrary historical detail generally does not.
  • Persistent slots. Dedicated variables for cue-like quantities that must not decay at all: task identity, goal, accumulated score, "which branch did the sign point to." These are written explicitly and rarely, and the rest of the model reads them as stable context.

The last category anticipates episodic and external memory, and the distinction matters for a reason that will keep recurring in this handbook: different memories have different write policies. Short-term state is written automatically and constantly, every step, no decisions required. Cue state should be written only when something new is detected, and a model that rewrites it every step will eventually overwrite it with the featureless corridor observations. Getting the write policy right is not an implementation detail; it is the difference between memory that preserves what matters and memory that grinds it away.

Attention over History: Memory as a Readable Buffer

Recurrent state is memory under write-only pressure: information enters, is compressed, and can only influence the future through whatever survives in the summary vector. If the compression discards something, there is no second chance, because the raw history is gone. Transformers, introduced to world modeling through sequence models over trajectories, take the opposite approach: keep the history and read from it directly. Attention computes

Attention⁡(Q,K,V)=softmax⁡ ⁣(QK⊤d)V.\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left(\frac{Q K^\top}{\sqrt{d}}\right) V.

where:

  • QQ: the query matrix, projections of the current token
  • KK: the key matrix, projections of each stored past token
  • VV: the value matrix, the content carried by each stored token
  • QK⊤QK^\top: the matrix of dot products between every query and every key, scoring how relevant each stored token is to the current one
  • dd: the key dimension; the factor d\sqrt{d} keeps the dot-product scores from growing too large before the softmax
  • softmax⁡(⋅)\operatorname{softmax}(\cdot): converts each row of scores into positive weights that sum to one

The queries, keys, and values are projections of stored token representations (and the current input). Reading the formula operationally: at each step, the model compares its query with every retained key, converts those scores into weights, and retrieves a weighted mixture of the values. A cue that remains encoded in context is directly addressable at step 50; it does not have to traverse 50 recurrent updates. Presence does not guarantee recall, however. The learned projections may discard the cue or the attention weights may fail to select it.

The result is a tradeoff: neither mechanism strictly improves on the other. The figure isolates one part of it: cumulative attention-score work over an autoregressive rollout. With a key-value cache, the new query at step tt attends to tt stored keys, so one step costs O(t)O(t) and the cumulative score work through step TT is O(T2)O(T^2). A recurrent update costs O(1)O(1) per step with fixed state width, or O(T)O(T) cumulatively. Architecture width, layers, and projection costs are omitted, so the curves show scaling rather than wall-clock ratios.

In[5]:
Code
import numpy as np

rollout_len = np.arange(8, 513, 8)
attention_work = rollout_len.astype(float) ** 2
recurrent_work = rollout_len.astype(float)
Out[6]:
Visualization
Log-scale line chart of normalized cumulative work versus rollout length, showing a quadratic attention curve and a linear recurrent curve.
Illustrative cumulative score/update work through an autoregressive rollout of length T, with architecture-dependent constants omitted. Cached attention is linear per new step and quadratic over the rollout; a fixed-width recurrent update is constant per step and linear over the rollout. KV-cache storage, not shown, grows linearly with T.

The curves separate with rollout length, but they do not establish a hardware-independent crossover or an orders-of-magnitude runtime claim. Long-context planning repeats this work for every simulated step and candidate future. Three costs need separate attention:

  • Compute. Cached attention grows linearly for each new step and quadratically over a full rollout; a fixed-width recurrent update is constant per step and linear over the rollout.
  • Storage and cutoffs. A standard key-value cache grows linearly with retained tokens. Once a representation leaves a finite context, it is unavailable unless another mechanism summarizes or stores it.
  • Temporal position. A model that reasons about elapsed time needs a representation that preserves the relevant timing distinctions, whether through positional features, recurrence, or another clock-like state.

One useful hybrid pairs a recurrent summary with bounded attention over recent latents. Attention serves high-fidelity short-term context while recurrence carries a compact rolling summary. We will examine transformer-based and state-space sequence architectures in Part V, Chapter 2. For now, hold onto the design question the two families answer differently: should old information be compressed into a summary, or kept and addressed? Neither answer is universally right; the choice depends on which parts of the past matter and at what fidelity.

Episodic and External Memory

Recurrent state and attention windows store information automatically. Some tasks call for a store the agent writes to selectively and reads from by content. This section covers two representative instantiations of that third family of memory.

Content-Addressed Memory

Both recurrent state and context windows update memory automatically from the interaction stream. Recurrence compresses absorbed history into a fixed state, while attention retains every token representation that fits in its finite context. Some tasks instead benefit from a store that the agent writes to selectively and reads by content. Exact soft lookup over NN slots costs O(Nd)O(Nd) for slot width dd, as the denominator below makes visible; approximate indexes can reduce lookup cost at the price of retrieval error and index maintenance. Neural Turing machines and differentiable neural computers formalized differentiable external memory as a matrix Mt∈RN×dM_t \in \mathbb{R}^{N \times d} manipulated by read and write heads:

mt=∑i=1Nwt,i Mt[i]wt,i=exp⁡ ⁣(kt⊤Mt[i])∑j=1Nexp⁡ ⁣(kt⊤Mt[j])\begin{aligned} m_t &= \sum_{i=1}^{N} w_{t,i}\, M_t[i] \\ w_{t,i} &= \frac{\exp\!\big(k_t^\top M_t[i]\big)}{\sum_{j=1}^{N} \exp\!\big(k_t^\top M_t[j]\big)} \end{aligned}

where:

  • MtM_t: the memory matrix at time tt, with NN slots
  • Mt[i]M_t[i]: the contents of slot ii
  • ktk_t: a query vector produced from the current state
  • wt,iw_{t,i}: the attention weight on slot ii at time tt, a softmax over the similarity scores kt⊤Mt[j]k_t^\top M_t[j] between the query and every slot
  • mtm_t: the retrieved vector, the weighted average of the slots most similar to the query

The model retrieves the stored entries most similar to what it currently needs. The mechanism is content-addressed in the same sense that a lookup table keyed by description is: the address is the content itself, or rather the model's learned encoding of it. Writes are more delicate than reads, because a bad write destroys information irreversibly: the model must decide where to write (allocation, often by usage weights so old entries are not clobbered) and what to write (erase and add vectors, as in the LSTM gating philosophy). The gating philosophy reappears here deliberately: writes are soft, gated modifications rather than hard overwrites, so that gradient descent can learn the write policy.

Soft addressing and allocation can become diffuse or saturated during optimization, so differentiable external-memory systems require careful training and evaluation. The same design pattern, selective writes and content-based reads, also provides a useful way to understand episodic control and retrieval-augmented models, even when they replace differentiable memory heads with concrete data structures and explicit retrieval logic.

Episodic Memory for Control

Episodic control keeps a database of experienced transitions with their realized returns and answers value queries by lookup rather than by generalization:

Q^(s,a)≈weighted average of stored returns from the k nearest neighbors of (s,a).\hat{Q}(s, a) \approx \text{weighted average of stored returns from the } k \text{ nearest neighbors of } (s, a).

where:

  • Q^(s,a)\hat{Q}(s, a): the lookup-based value estimate for the state-action pair (s,a)(s, a)
  • kk: the number of nearest neighbors retrieved from the database of stored transitions
  • "stored returns": the returns observed after visiting each stored neighbor, not values generalized by a network

The appeal is immediate reuse: after a high return is stored, the next matching query can consult it without waiting for parameter updates. This differs most sharply from one-step temporal-difference learning, which backs a target through one sampled transition at a time. Multi-step returns, eligibility traces, replay, and model-based planning can propagate credit farther, so episodic lookup is not the only alternative. From our current perspective, the deeper point is that episodic control is memory in the service of decisions: it stores evidence for "what worked here before?" The store need not support world prediction, which is a useful reminder that decision usefulness and predictive fidelity are separate criteria.

Retrieval-Augmented World Models

The same lookup principle applies to prediction. One concrete retrieval-augmented design fetches the kk stored trajectories most similar to a query latent state and conditions its prediction on them. Other systems may retrieve transitions, demonstrations, compressed memories, or documents. Instead of relying solely on fixed parameters to encode every dynamical pattern, the model consults stored experience. This has three attractive properties:

  • Non-parametric long-term memory. The store can grow with experience rather than squeezing every pattern into fixed parameters. An entry that remains retained is not overwritten by later parameter updates, but availability still depends on capacity, eviction, durability, key stability, and retrieval quality. A retention policy may itself use frequency or importance.
  • Interpretable predictions. A prediction can be accompanied by the retrieved neighbors that support it, which helps debugging and trust. When a rollout looks wrong, you can inspect the trajectories the model used as evidence and ask whether they were appropriate analogies, a level of transparency that opaque parameterized predictions cannot offer.
  • Locality assumptions made explicit. The model predicts "the future will resemble the futures of similar past situations." Parametric models also encode generalization assumptions, though not necessarily a nearest-neighbor locality assumption; retrieval makes its similarity rule directly inspectable. When similarity in latent space is a good proxy for dynamical similarity, this works well; when it is not, retrieval confidently retrieves the wrong analogies.

The failure mode deserves emphasis because it is a version of model exploitation that we will meet repeatedly in Part VII. An erroneous or outdated entry can steer the agent toward a repeated action. If the write or update rule then adds duplicates, increases the entry's weight, or suppresses contrary evidence, the loop becomes self-reinforcing; a read-only store does not become more entrenched merely because it was queried. Memory therefore needs write discipline and a policy for invalidating experience after the world changes.

Non-Markov Environments and State Aliasing

With the mechanisms in hand, we can now state the failure they exist to prevent. This section defines state aliasing precisely, shows how it manifests, and gives practical diagnostics for detecting inadequate memory.

Precisely What Goes Wrong

Let us now state the aliasing problem with care, because it is the standard diagnostic for memory failures. Assume the environment has an underlying Markov state sts_t, and the agent observes oto_t through a many-to-one channel. Consider two situations with identical observations:

Aliased situations

Situations (τt,τt′)(\tau_t, \tau'_t) are aliased for a target when they induce the same current observation, ot=ot′o_t = o'_t, but the history-conditioned target differs. For example, for some action aa,

p(ot+1,rt+1∣τt,a)≠p(ot+1,rt+1∣τt′,a).p(o_{t+1}, r_{t+1} \mid \tau_t, a) \neq p(o_{t+1}, r_{t+1} \mid \tau'_t, a).

They may instead require disjoint optimal actions. An observation-only model must assign both histories the same prediction or action distribution because it receives the same oto_t.

The formal definition rewards a close reading. The two histories agree on the observation-conditional and disagree on the history-conditional. That gap is exactly the information the memory is responsible for supplying, and it is also exactly what a stateless model cannot represent, since a stateless model's predictions are a function of oto_t alone and therefore identical across the two situations.

Aliasing has different consequences depending on what differs across the aliased situations, and the two cases fail in instructively different ways:

  • Different futures, same optimal action. The agent may still behave correctly, but a point predictor trained with a mean-seeking loss can average mutually exclusive futures into an infeasible outcome, such as a door halfway between open and closed. Planning on that average can distort values when rewards or dynamics are nonlinear, although affine values preserve expectations. A calibrated distribution can instead represent both unresolved modes.
  • Disjoint optimal actions. No observation-only reactive policy can select a different unique optimum in both situations. This creates an irreducible performance ceiling for that policy class when both situations occur. Better planning over the same aliased state cannot remove it; a history-dependent representation can.

Whether a given pair of situations is aliased depends on the policy and prediction task as well as the environment. Two corridor positions with identical walls but different door positions are only behaviorally aliased if the agent will eventually act differently there. If every policy under consideration treats the two doors identically, the distinction never matters, and spending memory capacity on it is waste. This is the temporal shadow of the bisimulation-style equivalences you met when we discussed state abstraction in Part III, Chapter 2: two histories are equivalent when they cannot be told apart by any future interaction, and a good memory is precisely a map from histories to a state space where only non-equivalent histories land in different places. The practical corollary is that memory requirements are task-relative: a shorter or cheaper memory may be perfectly sufficient for one controller and fatally inadequate for another operating in the same environment.

Diagnosing Memory Failures

How do you know, in practice, that a model's memory is inadequate? The corridor suggests several complementary checks. None is conclusive alone:

  • The one-step / long-horizon gap. Compare one-step error with multi-step rollouts or distant outcomes. A large gap is a screening signal, not a measurement of missing information: compounding transition error, rollout distribution shift, stochasticity, and decoder instability can produce the same pattern. Follow it with controlled memory tests.
  • Conditional uncertainty. A point estimate at the mean of mutually exclusive futures indicates mean-seeking prediction under unresolved uncertainty. A calibrated multimodal distribution can be the correct output of a stateless model that has forgotten which branch occurred. Evidence for disambiguating memory is history-conditioned concentration on the appropriate mode, not multimodality by itself.
  • Probe the state directly. Train a specified probe on one set of episodes and evaluate it on held-out episodes to predict the hidden cue from btb_t. High held-out accuracy shows that the cue is decodable by that probe class; it does not prove every downstream head can use it. Plot accuracy against time since the cue to estimate a task- and threshold-dependent memory horizon.

A concrete decay curve makes the diagnostic vivid. The figure below defines two synthetic curves by modeling probe accuracy as a cue signal with per-step retention rate gg: accuracy falls as 0.5+0.5 gΔt0.5 + 0.5 \, g^{\Delta t}, where Δt\Delta t is steps since the cue. A protected curve has gg near one; a leaking curve has gg well below one. The horizontal band is a chosen near-chance threshold for this illustration, not a statistical confidence interval.

In[7]:
Code
import numpy as np

dt = np.arange(0, 51)
retention_protected = (
    0.999  # gate stays open: cue signal survives nearly intact
)
retention_leaking = 0.94  # gates leak: cue signal decays each step

acc_protected = 0.5 + 0.5 * retention_protected ** dt.astype(float)
acc_leaking = 0.5 + 0.5 * retention_leaking ** dt.astype(float)
Out[8]:
Visualization
Two synthetic decay curves of probe accuracy versus steps since cue, one staying high and one entering a shaded near-chance band.
Illustrative probe-accuracy curves defined by 0.5 + 0.5 g^Delta-t. For g = 0.94, the curve enters the chosen 0.5-0.6 near-chance band at about 27 discrete steps; for g = 0.999 it remains high through 50 steps. The threshold is a diagnostic convention, not a universal horizon.

With the chosen 0.6 threshold, the leaking curve enters the shaded band at step 27. That crossing is an illustrative operational horizon, not a property of the retention rate alone; change the probe, threshold, or task and the number changes. Aggregate metrics can still hide this temporal pattern.

This held-out probe test is inexpensive and useful alongside behavioral evaluation; we will return to it in Part XI.

Behavioral ablation. Compare a memory-enabled policy with a reactive policy while matching observations, policy capacity, data, optimization, and compute as closely as possible. The performance gap estimates the benefit of the entire memory-enabled setup under those controls, not memory in isolation, and an empirical reactive baseline is not generally the best possible reactive policy.

Predictive State: Memory Without Reconstruction

Before leaving the formal territory, one more classical viewpoint is worth having in your pocket. Predictive state representations (PSRs) define state through predictions of tests: action-observation sequences such as "if I push left, then the light will be green." Two histories with the same sufficient vector of test predictions occupy the same predictive state, regardless of their hidden causes. In a linear PSR, each action-observation pair (a,o)(a,o) selects a matrix MaoM_{ao} and vector maom_{ao}. If q(h)=p(Q∣h)q(h)=p(Q\mid h) is the row vector of core-test predictions after history hh and the predicted observation probability satisfies q(h)mao>0q(h)m_{ao} > 0, the normalized update is

q(hao)=q(h)Maoq(h)mao.q(hao) = \frac{q(h) M_{ao}}{q(h) m_{ao}}.

The numerator is linear in q(h)q(h), while division by the predicted observation probability makes the normalized update linear-fractional rather than linear.

This viewpoint also clarifies modern recurrent predictors. A model trained only on observable prediction targets shares the predictive-state motivation: its hidden vector is useful insofar as it supports those targets, without reconstructing a privileged hidden cause. It is not automatically a classical PSR, because one-step prediction need not produce a sufficient set of controlled multi-step test predictions or the update above. The warning still applies: a representation trained only for selected predictions may omit everything the loss leaves out, which is why controlled probes remain worthwhile.

Worked Example: The Dark Corridor

Let us make the corridor concrete and quantitative. We will define a minimal POMDP, then compare two world-model candidates:

  • a Markov predictor, an MLP that sees only the current observation and action, and
  • a memory-based predictor, a GRU that sees the full history through its recurrent state.

Both are trained on the same objective: predict, at every timestep, the eventual outcome (left or right branch). This is the decision-relevant prediction that the corridor makes hard. Training on it at every timestep, rather than only at the end, matters: it gives the memory model a dense signal for retaining the cue, a point we will return to in the takeaways. Next-observation prediction, remember, is easy for both models here, which is precisely why this task is a good trap for dynamics-only evaluation.

The environment, in detail:

  • Underlying state: position along the corridor, plus the hidden branch side.
  • Observation: at position 00, a two-dimensional one-hot cue revealing the branch side. At positions 11 through 1010, a constant "wall" observation; every reading also carries small independent sensor noise. The observation channel destroys the cue information one step after the entrance.
  • Action features: forward during the corridor and a shared terminal decision marker at the end. The code predicts which branch should be chosen; it does not encode the executed left/right choice as an input action.
  • Target: the branch side, known from the start of each episode, which the models must predict from the history so far.

We expect the following, and the code will confirm it:

  1. The Markov predictor is accurate at step 00 (the cue is visible) and fluctuates around chance afterward. In the balanced population, the optimal current-input-only accuracy from step 1 onward is 50%50\%; a finite held-out sample need not equal exactly 0.5 at every timestep.
  2. The GRU maintains near-perfect accuracy throughout, because its recurrent state preserves the cue.
  3. A held-out linear probe recovers the branch from the GRU state, while a two-dimensional PCA view shows separated mid-corridor clusters and trajectories.

Code Implementation

We will build the corridor, train both models, and inspect the learned representation. The example is small enough for CPU execution; exact runtime depends on hardware, software versions, thread settings, and system load.

Now the imports we need for modeling and analysis.

In[9]:
Code
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression

Building the Dark Corridor

Each episode is a sequence of observations and actions of length 11: one cue step plus ten corridor steps. The observation is 3-dimensional (cue_left, cue_right, wall) and the action is a 2-dimensional one-hot vector. We generate training and test sets with different random seeds so the test episodes are fully held out.

In[10]:
Code
CORRIDOR_LEN = 10
T = CORRIDOR_LEN + 1  # one cue step plus corridor steps


def sample_episode(rng, corridor_len=CORRIDOR_LEN):
    steps = corridor_len + 1
    side = int(rng.integers(2))  # 0 = left branch, 1 = right branch
    obs = np.zeros((steps, 3), dtype=np.float32)
    act = np.zeros((steps, 2), dtype=np.float32)
    # Step 0: the entrance cue reveals the branch side
    obs[0, side] = 1.0
    # Steps 1..corridor_len: featureless walls (up to sensor noise)
    for t in range(1, corridor_len + 1):
        obs[t, 2] = 1.0
        act[t - 1, 0] = 1.0  # keep moving forward
    # Shared terminal-decision marker. It does not encode left versus right,
    # and it is noninformative about the target even though it is an input.
    act[corridor_len, 1] = 1.0
    # Sensor noise: independent of the branch, so it gives a stateless model
    # nothing to exploit, while keeping the memory task realistic.
    obs += rng.normal(0.0, 0.1, size=obs.shape).astype(np.float32)
    return obs, act, np.float32(side)


def make_dataset(n_episodes, seed, corridor_len=CORRIDOR_LEN):
    rng = np.random.default_rng(seed)
    obs_l, act_l, tgt_l = [], [], []
    for _ in range(n_episodes):
        o, a, s = sample_episode(rng, corridor_len=corridor_len)
        obs_l.append(o)
        act_l.append(a)
        tgt_l.append(s)
    obs = torch.tensor(np.stack(obs_l))
    act = torch.tensor(np.stack(act_l))
    steps = corridor_len + 1
    tgt = torch.tensor(np.stack(tgt_l)).unsqueeze(1).repeat(1, steps).float()
    return torch.cat([obs, act], dim=-1), tgt  # inputs: (obs, action)


x_train, y_train = make_dataset(2000, seed=1)
x_test, y_test = make_dataset(400, seed=2)

A quick look at one episode makes the aliasing tangible: from step 1 onward the observations are statistically indistinguishable regardless of branch, since the only variation is sensor noise independent of the side.

Out[11]:
Console
Input tensor shape (episodes, timesteps, features): (2000, 11, 5)
Branch balance in test set: 0.485

Observations at steps 0, 1, and 2 of one episode:
  t=0: obs=[ 1.08  0.03 -0.13]
  t=1: obs=[0.09 0.04 0.95]
  t=2: obs=[0.06 0.04 1.03]

As promised, step 0 carries the branch information and every later step is a featureless wall observation whose sensor noise is independent of the branch. The two appended action features mark the corridor phase (forward or terminal decision); they are fixed across branch labels and reveal nothing about the target. This toy input therefore uses a same-timestep phase/action marker rather than the at−1a_{t-1} convention in the earlier generic recurrence. Any model limited to the current input is guessing from step 1 on.

Two Predictor Candidates

The Markov baseline is an MLP applied independently to each timestep's input. Because it processes each timestep in isolation, it has no mechanism whatsoever for carrying information across steps; whatever it predicts at step tt is a fixed function of that step's input alone. The memory model is a GRU over the sequence, with a shared linear head reading out a branch prediction at every step. Both predict from the same inputs. The functional distinction relevant to aliasing is that the GRU carries prior inputs forward whereas the MLP does not, but the models also differ in architecture, parameter count, and compute. The comparison is therefore diagnostic, not a controlled ablation that isolates history access.

In[12]:
Code
class MarkovMLP(nn.Module):
    def __init__(self, in_dim=5, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(in_dim, hidden), nn.ReLU(), nn.Linear(hidden, 1)
        )

    def forward(self, x):  # x: (B, T, D)
        return self.net(x).squeeze(-1)  # (B, T) per-step logits


class MemoryGRU(nn.Module):
    def __init__(self, in_dim=5, hidden=32):
        super().__init__()
        self.gru = nn.GRU(in_dim, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)

    def forward(self, x):
        h, _ = self.gru(x)  # h: (B, T, hidden) recurrent states
        return self.head(h).squeeze(-1)

Note the deliberately modest GRU size: a hidden dimension of 32 has ample representational capacity for the one-bit target, though capacity alone does not ensure learning. Both models use binary cross-entropy on per-step branch prediction and receive the same 40 epochs and minibatch updates. They do not have equal parameter or compute budgets: the GRU is larger and sequential.

In[13]:
Code
def train_model(model, x, y, epochs=40, batch=128, lr=1e-3):
    opt = torch.optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.BCEWithLogitsLoss()
    n = x.shape[0]
    for epoch in range(epochs):
        perm = torch.randperm(n)
        for i in range(0, n, batch):
            idx = perm[i : i + batch]
            logits = model(x[idx])
            loss = loss_fn(logits, y[idx])
            opt.zero_grad()
            loss.backward()
            opt.step()
    return model


torch.manual_seed(0)
markov = train_model(MarkovMLP(), x_train, y_train)
memory = train_model(MemoryGRU(), x_train, y_train)

Evaluating Prediction by Timestep

The interesting quantity is accuracy as a function of steps since the cue. This is the evaluation the corridor was designed to demand: aggregate accuracy over the whole episode would average together the informative entrance and the aliased corridor, hiding exactly the failure we care about. We evaluate both models per timestep on the held-out episodes.

In[14]:
Code
@torch.no_grad()
def accuracy_per_step(model, x, y):
    probs = torch.sigmoid(model(x))
    preds = (probs > 0.5).float()
    return (preds == y).float().mean(dim=0).numpy()  # (T,) accuracy per step


acc_markov = accuracy_per_step(markov, x_test, y_test)
acc_memory = accuracy_per_step(memory, x_test, y_test)

mid = (
    CORRIDOR_LEN // 2
)  # mid-corridor timestep, used in the report and plots below
Out[15]:
Console
Entrance (t=0)   : Markov 1.000 | GRU 1.000
Mid-corridor (t=5)  : Markov 0.495 | GRU 1.000
Corridor end (t=10) : Markov 0.477 | GRU 1.000

The pattern is exactly the trap we described. The Markov model performs well at the entrance, where the cue is visible, and fluctuates around the chance level of 0.5 afterward. Its mid-corridor input is statistically independent of the branch. The GRU holds near-perfect branch accuracy across the corridor because branch information survives in its recurrent state. The claim about easy one-step dynamics prediction is analytic here: the next wall observation has a fixed mean plus independent noise. This code does not train or evaluate a separate dynamics head.

Inspecting the Learned Memory

Accuracy tells us the branch head works. We next use two different diagnostics. A logistic probe is trained on the first half of the held-out episodes' hidden states and evaluated on the second half. Separately, PCA gives a qualitative two-dimensional view of mid-corridor states and mean trajectories. PCA is not itself a probe, and its geometry does not identify which gate caused the representation.

In[16]:
Code
with torch.no_grad():
    h_all, _ = memory.gru(x_test)  # (N, T, hidden)
h_all = h_all.numpy()
sides = y_test[:, 0].numpy().astype(int)

pca = PCA(n_components=2).fit(h_all.reshape(-1, h_all.shape[-1]))
coords = pca.transform(h_all.reshape(-1, h_all.shape[-1])).reshape(
    len(sides), T, 2
)

mid_coords = coords[:, mid, :]  # hidden states at mid-corridor
traj_left = coords[sides == 0].mean(axis=0)  # (T, 2) mean trajectory
traj_right = coords[sides == 1].mean(axis=0)

probe_split = len(sides) // 2
probe_accuracy = []
for t in range(T):
    probe = LogisticRegression(random_state=0, max_iter=1000)
    probe.fit(h_all[:probe_split, t, :], sides[:probe_split])
    probe_accuracy.append(
        probe.score(h_all[probe_split:, t, :], sides[probe_split:])
    )

print(
    "Held-out cue-probe accuracy by timestep:",
    np.round(probe_accuracy, 3),
)

First, the per-timestep accuracy comparison.

Out[17]:
Visualization
Line chart: Markov accuracy drops to 0.5 after the cue; GRU accuracy stays near 1.0 across the corridor.
Branch-prediction accuracy per timestep on held-out dark-corridor episodes. The Markov model is accurate only while the cue is visible (t=0) and collapses to chance afterward; the GRU maintains near-perfect accuracy by carrying the cue in its recurrent state.

The vertical drop of the Markov curve between steps 0 and 1 is the signature of aliasing in this constructed process: one step of the observation channel erases the only decision-relevant information. In other systems the corresponding pattern may be abrupt, gradual, or obscured by other errors, so the clean version is a diagnostic reference rather than a universal curve shape.

Next, the mid-corridor hidden states. If the GRU's memory has truly un-aliased the two corridors, the two branches should form separated clusters even though the inputs at this timestep carry no information about the branch.

Out[18]:
Visualization
Scatter plot of mid-corridor hidden states in 2D PCA, two well-separated clusters colored by branch side.
PCA projection of GRU hidden states at mid-corridor, where the observation distribution is identical across branch labels even though individual noise realizations differ. The hidden states form two separated clusters by true branch, consistent with retained cue information.

The two branches form separated clusters in this projection. The explicit logistic probe above checks the stronger claim on held-out episodes; in this run it recovers the branch at every timestep. The PCA remains a qualitative view, while the held-out probe supplies the decodability measurement.

Finally, the trajectories. The two branches start apart at the cue and stay apart, which is exactly the "protected cue" behavior we hoped the gated recurrence would learn.

Out[19]:
Visualization
Two diverging trajectories in 2D PCA space over time, one per branch, staying separated after the cue.
Mean GRU hidden-state trajectories (in PCA space) for the two branch sides over the episode. The trajectories separate at the cue step and remain separated through the corridor, showing persistent cue retention rather than decay.

The separation established at the cue step remains visible across the corridor. This is consistent with the GRU's gated additive hidden-state update; a GRU has no separate LSTM cell state, and the PCA trajectory does not establish the mechanism causally. If the trajectories and held-out probe accuracy both converged toward chance with delay, that would be evidence of leaking cue information.

What This Example Teaches

Three takeaways generalize well beyond this toy corridor:

  • The objective supplies retention pressure. Every post-cue branch loss rewards a state that preserves the entrance cue. A next-wall-only objective would not distinguish the two branches in this constructed process. This explains the available learning signal; it is not an ablation proving a unique causal mechanism.
  • Use complementary diagnostics. Per-timestep task accuracy checks behavior across delay, a held-out probe tests cue decodability, and PCA trajectories give a qualitative view of representation dynamics. None alone identifies every memory failure.
  • Keep the reactive baseline. In richer environments a current observation may be partially informative, so a reactive model can be intermittently rather than uniformly wrong. Aggregate metrics may hide rare temporal failures; report performance by delay and compare matched systems with and without history access.

Limitations and Impact

Recurrent and attention-based sequence models are major tools for learned world models, including latent-imagination agents in Part VIII and large video models in Part IX. Learned temporal state is one enabling ingredient among representation learning, optimization, planning algorithms, data, and compute. Earlier systems also handled partial observability through finite histories, frame stacking, finite-state controllers, sampling over histories, and engineered estimators. Modern learned memory reduces some of that engineering burden; it does not replace the other ingredients or every earlier approach.

But the limitations are real, and knowing them will save you weeks of debugging:

  • Memory horizons are learned, not guaranteed. TBPTT length KK limits direct credit assignment across more than KK recurrent steps; it does not impose a hard forward-memory cutoff. Longer retention may emerge from transferable local dynamics, dense losses, auxiliary objectives, or hierarchy, but terminal-only dependencies beyond KK are difficult to teach directly. Test memory at the horizon the task requires.
  • The bottleneck can lose exactly what matters. Expected training loss weights common examples heavily unless sampling or loss design compensates, so rare decision-critical features can be squeezed out. A planner may then act as if a forgotten hazard were impossible. A calibrated predictive distribution over unresolved alternatives can expose that risk, but only if the model represents the uncertainty rather than collapsing it to a point estimate.
  • Long contexts shift the problem rather than solve it. Attention avoids repeatedly compressing all retained tokens into one recurrent vector, but token projections, learned retrieval, finite windows, and cumulative compute still introduce loss and cost. Retrieval adds capacity and inspectability while introducing similarity, retention, staleness, and write-policy failures.
  • Evaluation remains unsettled. Short-horizon metrics can hide aliasing, while rollout gaps are not specific to memory. Horizon sweeps, held-out cue probes, controlled history ablations, and task return answer different questions; use them together. We return to this evaluation problem in Part XI.

There is also a deeper conceptual limitation worth sitting with. This chapter emphasizes an individual's interaction history. Abstraction and composition address a complementary source of knowledge: structure shared across situations. Episodic retrieval can transfer from similar episodes, and compositional structure can also help identify the current instance, so the division is not exclusive. Their emphasis differs: memory supports instance-specific continuity, while compositional representations support structural reuse. The next two chapters develop that complementary direction through object-centric and relational representations, then self-models and affordances.

Summary

This chapter turned the abstract fact of partial observability into a concrete design discipline. The through-line was a single question, what should the state be when a single observation is not enough, answered first in principle (condition on history, then compress) and then in architecture (gated recurrence, attention, external stores):

  • History is the general fallback under partial observability. A practical memory aims to construct a task-sufficient state when one exists, or a useful fixed-budget approximation otherwise. Learned recurrent states can act as filters or belief-like summaries, but they do not automatically represent normalized Bayesian beliefs.
  • Recurrence can struggle over long horizons. Long-range derivatives are products of Jacobians and can vanish or explode; the scalar constant-gain example illustrates exponential behavior, while matrix products depend on time-varying directions and spectra. LSTM and GRU updates provide direct gated retention paths. TBPTT limits direct credit-assignment distance, not forward memory as a hard theorem.
  • Short-term and long-term memory are different problems. Filtering recent context, integrating slow quantities, and retaining one-time cues have different requirements and different write policies, and good architectures give them separate compartments rather than forcing one vector to do all three jobs.
  • Memory can be compressive or addressable. Recurrent summaries trade detail for compactness. Attention and external stores keep representations separately addressable subject to encoding, capacity, eviction, write, and retrieval losses. Compute, averaging, cutoff, similarity, and staleness are characteristic but overlapping risks rather than one exclusive failure per family.
  • Aliasing is the diagnostic lens. Identical observations paired with different target conditionals or disjoint optimal actions mark where observation-only models break. The corridor executed per-timestep evaluation, a held-out latent probe, and PCA visualizations. Controlled reactive-policy baselines are a further recommended behavioral ablation; this example did not execute one.

You now have the vocabulary and the failure-mode literacy to read any recurrent world-model architecture critically: ask what its compartments are, what write policy each uses, what horizon its training spanned, and how you would detect an aliasing failure if one occurred. The next chapter asks a complementary question: instead of compressing history into an opaque vector, can we factor the world into objects and relations, so that both memory and dynamics become structured and compositional?

Quiz

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about temporal state and memory in world models.

Temporal State and Memory

Question 1 of 80 of 8 completed
In the dark corridor example, why does a stateless Markov model report low one-step prediction error yet fail at the task?

Comments

No comments yet. Be the first to share your thoughts!

Reference

Citation details

Cite or share this article.

BIBTEXAcademic
@misc{brenndoerfer2026temporalstate, author = {Michael Brenndoerfer}, title = {Temporal State and Memory}, year = {2026}, url = {https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models}, organization = {mbrenndoerfer.com}, note = {Accessed: 2026-10-15} }
APAAcademic
Michael Brenndoerfer (2026). Temporal State and Memory. Retrieved from https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models
MLAAcademic
Michael Brenndoerfer. "Temporal State and Memory." 2026. Web. October 15, 2026. <https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models>.
CHICAGOAcademic
Michael Brenndoerfer. "Temporal State and Memory." Accessed October 15, 2026. https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models.
HARVARDAcademic
Michael Brenndoerfer (2026) 'Temporal State and Memory'. Available at: https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models (Accessed: October 15, 2026).
SimpleBasic
Michael Brenndoerfer (2026). Temporal State and Memory. https://mbrenndoerfer.com/writing/temporal-state-and-memory-recurrent-world-models

About the author

Continue with the full handbook

This chapter is part of World Models Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.

Explore World Models Handbook
Newsletter

Stay up to date

Get articles, book updates, and news delivered to your inbox.

No spam, unsubscribe anytime.

or

Join the community

Sign in to remove popups, track your reading progress, and join the discussion.