Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Games and Learned Game Engines › A Worked Example: Learning a Grid-World Engine

Michael BrenndoerferJuly 25, 20268 min read

World Models HandbookAll chapters

Part 10 of 12Application Domains

Chapter 55Games and Learned Game Engines › A Worked Example: Learning a Grid-World Engine

Reading section 4 of 10 · A Worked Example: Learning a Grid-World Engine · 8 minSection 4 of 10 · 8 min

Section 4 of 10 · 8 min read

Games and Learned Game Engines › A Worked Example: Learning a Grid-World EngineOne section · Sections in this chapter4 / 10 ↓

Sections in this chapter10

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 now build a $5 \times 5$ grid-world engine with a movable agent and fixed goal. Its two-channel observation marks the agent and goal, and its after-action reward is negative Manhattan distance to the goal. A small network predicts the next observation and reward. We will test how its predictions change under recursion, then use it to score candidate action sequences while executing and evaluating the resulting choices in the reference environment.

The observation contains 50 binary entries: two 5 × 5 channels, one for the agent and one for the goal. That makes mistakes easy to inspect. The experiment is a teaching example, not a timing or scaling result for pixel-based engines.

The Environment and Its DataLink Copied

We import the scientific stack and set NumPy and PyTorch seeds so this notebook run can be reproduced in a compatible environment. A seed does not remove sampling uncertainty, implementation differences, or training-data effects; later curves describe one seeded run.

In[3]:

Code

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)
np.random.seed(0)

Actions are the four cardinal moves, clipped at the boundary. The reward is the negative Manhattan distance from the agent to the goal after the action; reaching the goal gives zero, the maximum reward. An action that leaves distance unchanged has unchanged reward, not an improvement signal. Episodes end at the goal or after twenty steps. The budget bounds trajectory length but does not make state coverage uniform.

In[4]:

Code

class GridWorld:
    """A tiny 5x5 grid world with a movable agent and a random goal."""

    ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # up, down, left, right

    def __init__(self, size=5, max_steps=20, seed=0):
        self.size = size
        self.max_steps = max_steps
        self.rng = np.random.default_rng(seed)
        self.agent = None
        self.goal = None
        self.t = 0

    def reset(self):
        self.agent = self.rng.integers(0, self.size, size=2)
        self.goal = self.rng.integers(0, self.size, size=2)
        while np.array_equal(self.goal, self.agent):
            self.goal = self.rng.integers(0, self.size, size=2)
        self.t = 0
        return self._obs()

    def _obs(self):
        obs = np.zeros((2, self.size, self.size), dtype=np.float32)
        obs[0, self.agent[0], self.agent[1]] = 1.0
        obs[1, self.goal[0], self.goal[1]] = 1.0
        return obs

    def step(self, a):
        dy, dx = self.ACTIONS[a]
        self.agent = np.array(
            [
                int(np.clip(self.agent[0] + dy, 0, self.size - 1)),
                int(np.clip(self.agent[1] + dx, 0, self.size - 1)),
            ]
        )
        self.t += 1
        dist = float(np.abs(self.agent - self.goal).sum())
        reward = -dist
        done = np.array_equal(self.agent, self.goal) or self.t >= self.max_steps
        return self._obs(), reward, done

To generate data, we interact with the environment using uniformly sampled actions. This is a simple behavior policy, not passive observation: the collector chooses actions and receives their consequences. Nor does uniform action sampling imply uniform state coverage; boundaries, early termination, and the reset distribution shape which transitions appear. Interaction Data and Passive Observation explains why the behavior policy and coverage matter.

In[5]:

Code

def collect_random_transitions(n_episodes=2000, seed=1):
    env = GridWorld(seed=seed)
    obs_list, act_list, next_obs_list, rew_list = [], [], [], []
    for _ in range(n_episodes):
        obs = env.reset()
        done = False
        while not done:
            a = int(env.rng.integers(0, 4))
            next_obs, r, done = env.step(a)
            obs_list.append(obs)
            act_list.append(a)
            next_obs_list.append(next_obs)
            rew_list.append(r)
            obs = next_obs
    return (
        np.array(obs_list, dtype=np.float32),
        np.array(act_list, dtype=np.int64),
        np.array(next_obs_list, dtype=np.float32),
        np.array(rew_list, dtype=np.float32),
    )


obs_arr, act_arr, next_obs_arr, rew_arr = collect_random_transitions()

Out[6]:

Console

Collected 32158 transitions
Obs shape: (32158, 2, 5, 5)
Mean reward: -3.488
Reward range: -8.0 to -0.0

Each transition is $(o_t,a_t,o_{t+1},r_t)$ , where $r_t$ is the reward after action $a_t$ and $o_t \in \{0,1\}^{2\times5\times5}$ . The model maps a flattened observation and one-hot action to a predicted next observation and scalar reward. Unlike Dreamer, it has no recurrent latent state. It is deliberately small enough to inspect, and it can still expose the same gap between one-step prediction and recursive use.

The ModelLink Copied

The architecture is minimal. A two-hidden-layer MLP takes the flattened observation and one-hot action; separate linear heads predict the next observation and reward. The heads allow distinct losses or weights, although this example simply sums two mean-squared errors. Other world models may use different state representations, objectives, and rollout tests; this MLP is not a template for all of them.

In[7]:

Code

OBS_DIM = 2 * 5 * 5  # two channels (agent, goal), each a 5x5 grid
N_ACTIONS = 4


class WorldModel(nn.Module):
    def __init__(self, obs_dim=OBS_DIM, n_actions=N_ACTIONS, hidden=128):
        super().__init__()
        self.n_actions = n_actions
        self.trunk = nn.Sequential(
            nn.Linear(obs_dim + n_actions, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        self.next_obs_head = nn.Linear(hidden, obs_dim)
        self.reward_head = nn.Linear(hidden, 1)

    def forward(self, obs_flat, a):
        a_oh = F.one_hot(a, self.n_actions).float()
        h = self.trunk(torch.cat([obs_flat, a_oh], dim=-1))
        return self.next_obs_head(h), self.reward_head(h).squeeze(-1)

In[8]:

Code

W_model = WorldModel()
optimizer = torch.optim.Adam(W_model.parameters(), lr=1e-3)

obs_t = torch.from_numpy(obs_arr.reshape(len(obs_arr), -1))
act_t = torch.from_numpy(act_arr)
nobs_t = torch.from_numpy(next_obs_arr.reshape(len(next_obs_arr), -1))
rew_t = torch.from_numpy(rew_arr)

n_val = 5000
train_obs, train_act, train_nobs, train_rew = (
    obs_t[:-n_val],
    act_t[:-n_val],
    nobs_t[:-n_val],
    rew_t[:-n_val],
)
val_obs, val_act, val_nobs, val_rew = (
    obs_t[-n_val:],
    act_t[-n_val:],
    nobs_t[-n_val:],
    rew_t[-n_val:],
)

Training sums mean-squared error for next-observation entries and reward. Equal coefficients are a simple demonstration choice, not evidence that the two error scales are comparable or appropriately balanced. Their values should be inspected separately when that balance matters; Stochasticity and Uncertainty discusses predictive uncertainty rather than prescribing this weighting.

In[9]:

Code

loss_history = []
for epoch in range(15):
    W_model.train()
    perm = torch.randperm(len(train_obs))
    epoch_loss = 0.0
    for i in range(0, len(train_obs), 256):
        idx = perm[i : i + 256]
        pred_no, pred_r = W_model(train_obs[idx], train_act[idx])
        target_no = train_nobs[idx] + 0.02 * torch.randn_like(train_nobs[idx])
        target_r = train_rew[idx] + 0.10 * torch.randn_like(train_rew[idx])
        loss = F.mse_loss(pred_no, target_no) + F.mse_loss(pred_r, target_r)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item() * len(idx)
    loss_history.append(epoch_loss / len(train_obs))

Out[10]:

Console

Final training loss: 0.0344

Out[11]:

Visualization

Line plot of world-model MSE loss dropping over fifteen epochs. — Training joint MSE for next-observation and reward prediction over 15 epochs. The loss falls during optimization, but this training curve does not establish accurate held-out transitions or reliable recursive rollouts. Small random perturbations were added to both training targets.

The training objective falls, but it cannot by itself tell us how well the model predicts held-out transitions. The environment is deterministic, while the training targets include small random perturbations. The next tests separate fitting this objective from one-step prediction and recursive rollout fidelity.

One-Step Accuracy Is Not EnoughLink Copied

We start with 5,000 validation transitions set aside from the end of the collected array. This is a transition-level split: the boundary need not be episode-disjoint, and the small grid permits repeated states. We compute the mean Euclidean (L2) distance between flattened predicted and true next observations:

\frac{1}{N} \sum_{i=1}^{N} \bigl\| \hat{o}_{t+1}^{(i)} - o_{t+1}^{(i)} \bigr\|_2

over that held-out set of $N$ transitions. Here:

$N$ : the number of held-out transitions evaluated
$\hat{o}_{t+1}^{(i)}$ : the model's predicted next observation for the $i$ -th held-out transition, flattened to a vector
$o_{t+1}^{(i)}$ : the true next observation for the $i$ -th held-out transition, flattened to a vector
$\| \cdot \|_2$ : the Euclidean (L2) norm, so each example contributes the overall pixel-space distance between prediction and truth

In[12]:

Code

with torch.no_grad():
    pred_next, pred_r = W_model(val_obs, val_act)

per_example_err = ((pred_next - val_nobs) ** 2).sum(dim=-1).sqrt().numpy()
reward_err = (pred_r - val_rew).abs().numpy()

Out[13]:

Console

Mean 1-step L2 error: 0.9972
Mean reward absolute error: 0.0492

Out[14]:

Visualization

Histogram of held-out one-step L2 prediction errors, clustered around one rather than zero. — Distribution of one-step L2 distances between predicted and true next observations on held-out transitions. Errors cluster around one in this run rather than near zero, so the histogram does not support an almost-exact prediction claim. A one-step distribution also cannot establish recursive rollout fidelity.

The errors cluster around an L2 distance of one on the flattened two-channel grid, not around zero. That is a meaningful one-step mismatch in this small environment: the clustered error should not be mistaken for near-exact predictions. One-step validation remains useful, but it does not reveal how errors change when predictions become inputs to later predictions. We therefore measure recursive rollout error separately.