Temporal State and Memory › Code Implementation
World Models HandbookAll chapters Temporal State and Memory › Code ImplementationOne section · Sections in this chapter9 / 12 ↓
Reading tools Definitions & inline links
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
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.
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 LogisticRegressionBuilding 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.
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.
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 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 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.
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.
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.
@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 belowEntrance (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.
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.
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.
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.
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.