Object-Centric and Relational Worlds › Code Implementation
World Models HandbookAll chapters Object-Centric and Relational Worlds › Code ImplementationOne section · Sections in this chapter8 / 11 ↓
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 a small interacting-particle world, train an interaction network to predict its dynamics, measure how rollout error compounds, and test the central claim of the chapter: that a relational model trained on four-object scenes generalizes to larger graphs where a flat per-object model does not. We close with a minimal data-association tracker to make identity concrete. Throughout, our "slots" are the true object states, positions and velocities, so that questions about representation quality do not contaminate the questions about dynamics and identity we intend to isolate.
A Synthetic Interacting-Particle World
The world is a set of point masses in the plane connected by springs. Dynamics are semi-implicit Euler: accelerations come from spring forces along active edges plus global damping.
import numpy as np
rng = np.random.default_rng(11)
DT = 0.02 # integration step
SPRING_K = 6.0 # spring stiffness
DAMPING = 0.6 # global velocity damping
def spring_accel(pos, vel, senders, receivers, rest):
"""Acceleration from spring forces along directed edges plus damping."""
acc = np.zeros_like(pos)
diff = pos[senders] - pos[receivers] # (E, 2)
r = np.linalg.norm(diff, axis=1, keepdims=True) + 1e-9
f = SPRING_K * (r - rest[:, None]) * diff / r # force on receiver
np.add.at(acc, receivers, f)
acc -= DAMPING * vel
return acc
def simulate(pos0, vel0, senders, receivers, rest, steps):
pos, vel = np.array(pos0, float), np.array(vel0, float)
pos_traj, vel_traj = [pos.copy()], [vel.copy()]
for _ in range(steps):
acc = spring_accel(pos, vel, senders, receivers, rest)
vel = vel + DT * acc
pos = pos + DT * vel
pos_traj.append(pos.copy())
vel_traj.append(vel.copy())
return np.stack(pos_traj), np.stack(vel_traj)The directed-edge convention says that sender -> receiver carries the sender's effect on the receiver. The analytic simulator produces equal-and-opposite forces because its formula is odd under reversal. Duplicating edges lets a learned function represent both directions, but does not enforce antisymmetry. The tiny 1e-9 keeps the coincident-particle denominator nonzero.
Now the graph builders. Training scenes are jittered four-node squares; test scenes are random -nearest-neighbor graphs of various sizes, so evaluation involves both new geometries and new graph sizes. The mismatch between training and test topology is deliberate: it is the compositional generalization claim under test.
def edge_rest_lengths(pos0, edges, jitter=0.15):
base = np.array([np.linalg.norm(pos0[j] - pos0[i]) for i, j in edges])
noisy = base * (1 + rng.normal(scale=jitter, size=base.shape))
return np.clip(noisy, 0.2, None)
def make_square(jitter=0.05):
base = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]])
pos0 = base + rng.normal(scale=jitter, size=base.shape)
edges = [(0, 1), (1, 2), (2, 3), (3, 0)]
return pos0, edges, edge_rest_lengths(pos0, edges)
def make_random_graph(n, k=2):
pos0 = rng.normal(size=(n, 2))
edges = set()
for i in range(n):
d = np.linalg.norm(pos0 - pos0[i], axis=1)
for j in np.argsort(d)[1 : k + 1]:
edges.add((min(i, int(j)), max(i, int(j))))
edges = sorted(edges)
return pos0, edges, edge_rest_lengths(pos0, edges)
def directed_edges(edges, rest):
senders = np.array(
[i for i, j in edges] + [j for i, j in edges], dtype=np.int64
)
receivers = np.array(
[j for i, j in edges] + [i for i, j in edges], dtype=np.int64
)
return senders, receivers, np.concatenate([rest, rest])Rest lengths carry deliberate jitter. The true dynamics depend on current length minus rest length, so this edge attribute affects the prediction target. Even without the extra jitter, coordinate jitter would produce multiple rest lengths. The extra perturbation decorrelates rest length from initial geometry and makes ignoring the attribute costly; exact prediction across otherwise matched states requires that value to be available.
Generate the datasets: 256 training trajectories on squares and 16 held-out trajectories each for 4-, 6-, and 8-node random graphs, 40 steps each.
STEPS = 40
def build_dataset(n_graphs, graph_fn, steps=STEPS, vel_scale=0.6):
trajs, vels_all, meta = [], [], []
for _ in range(n_graphs):
pos0, edges, rest = graph_fn()
senders, receivers, rest_dir = directed_edges(edges, rest)
vel0 = rng.normal(scale=vel_scale, size=pos0.shape)
traj, vels = simulate(pos0, vel0, senders, receivers, rest_dir, steps)
trajs.append(traj)
vels_all.append(vels)
meta.append((senders, receivers, rest_dir))
return trajs, vels_all, meta
train_data = build_dataset(256, make_square)
test_data = {
n: build_dataset(16, lambda n=n: make_random_graph(n)) for n in (4, 6, 8)
}
n_train_transitions = 256 * STEPSTraining trajectories: 256 (nodes per graph: 4) Training transitions: 10240 Test graph sizes: [4, 6, 8] (16 trajectories each)
Every training trajectory is a 4-node graph; every test set changes geometry and the 6- and 8-node sets also change object count. The random 4-node set tests topology and geometry transfer, while the evaluated 6- and 8-node sets additionally test count transfer. The architecture accepts other compatible sizes, but this experiment establishes results only for 4, 6, and 8 nodes.
The next figure is the mental picture to keep for the rest of the section: a graph where each edge runs a small edge function and each node aggregates what arrives. When we later say the model "generalizes to 8 nodes," the mechanism behind that sentence is visible here: nothing in the picture depends on there being four nodes.
An Interaction Network in PyTorch
The model follows the recipe exactly. The edge function consumes relative position, relative velocity, and rest length; aggregation is a sum via index_add_; the node function maps own velocity plus aggregated messages to a predicted acceleration. Note that all edge features are relative, so translation invariance is built in rather than learned: shift the whole scene by any vector and every edge input is unchanged, hence so is every prediction.
import torch
import torch.nn as nn
torch.manual_seed(0)
class InteractionNet(nn.Module):
def __init__(self, edge_in=5, node_in=2, hidden=64):
super().__init__()
self.edge_mlp = nn.Sequential(
nn.Linear(edge_in, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
)
self.node_mlp = nn.Sequential(
nn.Linear(hidden + node_in, hidden),
nn.SiLU(),
nn.Linear(hidden, 2),
)
def forward(self, pos, vel, senders, receivers, rest):
edge_in = torch.cat(
[
pos[receivers] - pos[senders],
vel[receivers] - vel[senders],
rest[:, None],
],
dim=-1,
)
msgs = self.edge_mlp(edge_in)
agg = torch.zeros(pos.shape[0], msgs.shape[1]).index_add_(
0, receivers, msgs
)
return self.node_mlp(torch.cat([vel, agg], dim=-1))
class IndependentMLP(nn.Module):
"""Baseline: one shared MLP per node, no edges, no relations."""
def __init__(self, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(4, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, 2),
)
def forward(self, pos, vel, senders, receivers, rest):
return self.net(torch.cat([pos, vel], dim=-1))The comparison model is an independent per-object baseline, not a fixed-width flat scene model. It predicts from one node's absolute position and velocity and is itself applicable to any node count. Configuration-specific spring forces are not functions of those inputs alone. Under squared loss, the unrestricted Bayes predictor for this input is the conditional mean acceleration, with irreducible conditional variance from hidden relational context. The finite MLP's measured error can additionally include approximation and optimization error. Disjoint-graph concatenation supports variable node counts by offsetting indices; equal training size is not required.
def transition_items(traj, vels, senders, receivers, rest_dir):
items = []
for t in range(len(traj) - 1):
items.append(
dict(
pos=traj[t],
vel=vels[t],
acc=(vels[t + 1] - vels[t]) / DT,
senders=senders,
receivers=receivers,
rest=rest_dir,
)
)
return items
def batch_graphs(items):
pos = torch.tensor(
np.concatenate([it["pos"] for it in items]), dtype=torch.float32
)
vel = torch.tensor(
np.concatenate([it["vel"] for it in items]), dtype=torch.float32
)
acc = torch.tensor(
np.concatenate([it["acc"] for it in items]), dtype=torch.float32
)
rest = torch.tensor(
np.concatenate([it["rest"] for it in items]), dtype=torch.float32
)
senders, receivers, offset = [], [], 0
for it in items:
senders.append(it["senders"] + offset)
receivers.append(it["receivers"] + offset)
offset += it["pos"].shape[0]
senders = torch.tensor(np.concatenate(senders), dtype=torch.long)
receivers = torch.tensor(np.concatenate(receivers), dtype=torch.long)
return pos, vel, acc, senders, receivers, rest
def train_model(model, items, batch_indices, lr=3e-3):
opt = torch.optim.Adam(model.parameters(), lr=lr)
losses = []
for idx in batch_indices:
batch = batch_graphs([items[i] for i in idx])
pred = model(*batch[:2], *batch[3:])
loss = ((pred - batch[2]) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
losses.append(loss.item())
return lossesTrain both models on the square-graph transitions with the same precomputed minibatch sequence.
train_items = []
for traj, vels, (senders, receivers, rest_dir) in zip(*train_data):
train_items.extend(
transition_items(traj, vels, senders, receivers, rest_dir)
)
model_in = InteractionNet()
model_base = IndependentMLP()
TRAIN_STEPS = 500
BATCH_SIZE = 256
train_batch_indices = [
rng.integers(0, len(train_items), size=BATCH_SIZE)
for _ in range(TRAIN_STEPS)
]
losses_in = train_model(model_in, train_items, train_batch_indices)
losses_base = train_model(model_base, train_items, train_batch_indices)Interaction network: final loss 0.04309 Independent MLP: final loss 0.89749
In the seeded run below, the final acceleration loss is about 21 times lower for the interaction network. The representational advantage is architectural: an independent node lacks pair context. The measured finite-training gap still depends on capacity, optimization, minibatches, and training duration.
Training and Rollout Evaluation
Now the distinction that this book keeps returning to: one-step accuracy is cheap, closed-loop rollout is the real test. We roll each trained model out from the initial state of every test trajectory, feeding its own predictions back as inputs, and record position error per horizon.
def rollout_model(model, traj, vels, senders, receivers, rest, steps):
pos, vel = traj[0].copy(), vels[0].copy()
out = [pos.copy()]
t_s = torch.tensor(senders, dtype=torch.long)
t_r = torch.tensor(receivers, dtype=torch.long)
t_rest = torch.tensor(rest, dtype=torch.float32)
for _ in range(steps):
t_pos = torch.tensor(pos, dtype=torch.float32)
t_vel = torch.tensor(vel, dtype=torch.float32)
with torch.no_grad():
acc = model(t_pos, t_vel, t_s, t_r, t_rest).numpy()
vel = vel + DT * acc
pos = pos + DT * vel
out.append(pos.copy())
return np.stack(out)
def rollout_errors(model, data, steps=STEPS):
per_traj = []
for traj, vels, (senders, receivers, rest_dir) in zip(*data):
pred = rollout_model(
model, traj, vels, senders, receivers, rest_dir, steps
)
per_traj.append(((pred - traj) ** 2).mean(axis=(1, 2)))
return np.stack(per_traj).mean(axis=0)
errors_in = {n: rollout_errors(model_in, test_data[n]) for n in (4, 6, 8)}
errors_base = {n: rollout_errors(model_base, test_data[n]) for n in (4, 6, 8)}
example = test_data[6]
example_true = example[0][0]
example_pred = rollout_model(
model_in, example[0][0], example[1][0], *example[2][0], STEPS
)
bar_h = 20
ratio = {n: errors_base[n][bar_h] / errors_in[n][bar_h] for n in (4, 6, 8)}n=4: one-step MSE 8.41e-07, horizon-20 MSE 1.88e-02, baseline/interaction ratio at horizon 20: 3.5x n=6: one-step MSE 3.15e-07, horizon-20 MSE 7.76e-03, baseline/interaction ratio at horizon 20: 7.8x n=8: one-step MSE 2.12e-07, horizon-20 MSE 5.51e-03, baseline/interaction ratio at horizon 20: 10.6x
Two effects are visible. The interaction-network rollout error grows with horizon. Its horizon-20 errors at 4, 6, and 8 nodes remain in the same order of magnitude, while the independent-node baseline is worse at every evaluated size because it cannot recover configuration-specific pair forces. The baseline can still exploit own velocity and absolute position to learn a marginal predictor; without a component ablation, we do not attribute that partial fit specifically to damping. It is a reasonable current-state independent-node comparison, not a proof of the best possible trained per-object model.
For this seeded experiment, error is small at short horizons and grows during closed-loop rollout. Here, structure changes the observed rate rather than eliminating drift. The horizon-20 baseline-to-interaction ratios remain above one at all evaluated sizes. Both implemented models are size-flexible: the independent-node baseline shares a node MLP but lacks pair inputs, while the interaction network additionally shares an edge function and aggregates relational context. A separate fixed-dimensional joint-scene MLP would require redesign or retraining for a new count, but that is not the baseline plotted here.
We can also look at the trajectories themselves on a held-out 6-node graph.
The panels show both transfer and its limit. On this held-out 6-node graph, the rollout retains broad motion directions but several paths differ visibly in extent and curvature by 40 steps. The architecture provides separate places for relational messages and node updates, but this plot does not identify damping with one module or spring rhythm with the other. Shared size-independent functions are an architectural explanation for transfer across the evaluated counts, not proof that the learned modules recovered a unique physical decomposition.
Keeping Identity Through a Crossing
Finally, consider two noisy point tracks that cross while observation order is reshuffled each frame. We maintain a last-displacement constant-velocity prediction and solve a two-by-two Euclidean-cost assignment with the Hungarian algorithm. This is a SORT-like filtering-plus-assignment toy, not the SORT algorithm itself, which uses a Kalman state for detected boxes and an overlap-based association cost. Near the crossing, position-only current-frame association is highly uncertain under the stated equal-appearance noise model.
from scipy.optimize import linear_sum_assignment
T_TRACK = 60
track_rng = np.random.default_rng(0)
start = np.array([[-0.8, -0.3], [0.8, -0.3]])
end = np.array([[0.8, 0.3], [-0.8, 0.3]])
true_tracks = (
start[None]
+ (end - start)[None] * (np.arange(T_TRACK) / (T_TRACK - 1))[:, None, None]
)
obs = true_tracks + track_rng.normal(scale=0.04, size=true_tracks.shape)
## Shuffle observation order within each frame: no index privilege.
## true_id[t, o] records which particle generated the observation now at position o.
obs_shuffled = np.empty_like(obs)
true_id = np.zeros((T_TRACK, 2), dtype=int)
for t in range(T_TRACK):
perm = track_rng.permutation(2)
obs_shuffled[t] = obs[t][perm]
true_id[t] = perm # obs[t, i] was generated by particle i
track_pos = [obs_shuffled[0, 0].copy(), obs_shuffled[0, 1].copy()]
track_vel = [np.zeros(2), np.zeros(2)]
track_identity = [int(true_id[0, 0]), int(true_id[0, 1])]
track_hist = [[obs_shuffled[0, 0].copy()], [obs_shuffled[0, 1].copy()]]
switches = 0
for t in range(1, T_TRACK):
preds = np.array([track_pos[k] + track_vel[k] for k in range(2)])
cost = np.linalg.norm(obs_shuffled[t][:, None] - preds[None], axis=2)
rows, cols = linear_sum_assignment(cost) # rows: obs idx, cols: track idx
for o, k in zip(rows, cols):
new_pos = obs_shuffled[t, o]
track_vel[k] = new_pos - track_pos[k]
track_pos[k] = new_pos
track_hist[k].append(new_pos.copy())
assigned_identity = int(true_id[t, o])
if assigned_identity != track_identity[k]:
switches += 1
track_identity[k] = assigned_identityTrack-level identity changes over 60 frames: 4 Minimum true separation between the particles: 0.027 Minimum noisy-observation separation: 0.018
The counter initializes each track's generating identity at frame zero and increments whenever a later assignment changes it. Thus a nonzero value records a swap event rather than a change in shuffled observation index. In this seeded example the position-only predictor records a nonzero swap count near the crossing. That failure is the lesson: prediction adds context, but it does not make an ambiguous association identifiable.
Prediction helps away from the crossing by supplying where each track should be if its identity persists. Here it is insufficient at closest approach: noisy last-displacement estimates drive repeated swaps. Appearance features, calibrated Kalman uncertainty, or multi-frame smoothing could add evidence. Recurrent slot conditioning changes where continuity enters the system, but it can also swap unless identity consistency is checked explicitly.
Key Parameters
The key parameters for the interaction network and its training setup are:
- SPRING_K: Spring stiffness of the synthetic world. It sets the magnitude of the pairwise force the edge function must learn from relative positions and rest lengths.
- DAMPING: Global velocity damping applied in the node update. It is directly predictable from a node's own velocity, unlike configuration-specific spring forces.
- DT: Integration step for semi-implicit Euler. For a fixed physical horizon and within the method's stable regime, smaller values usually reduce discretization error but require more steps.
- STEPS: Number of rollout steps per trajectory. It fixes the horizon over which compounding rollout error is measured.
- hidden: Width of the hidden layers in the edge and node MLPs. Larger values give the message and update functions more capacity at a higher training cost.
- BATCH_SIZE: Number of transition items sampled per gradient step. Larger batches often reduce gradient variance at higher memory cost.
- TRAIN_STEPS: Number of optimizer updates applied to each model with the shared minibatch schedule.
- lr: Adam optimizer step size. Larger values can accelerate early progress but also increase instability or overshoot.