Hierarchical, Hybrid, and Omnimodal Models › Worked Example: Comparing History Conditioning and Routing
World Models HandbookAll chapters Hierarchical, Hybrid, and Omnimodal Models › Worked Example: Comparing History Conditioning and RoutingOne section · Sections in this chapter6 / 9 ↓
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 now compare a flat baseline with two separate window-conditioned alternatives on a controlled 2D system with a hidden context. One compresses a history window into a continuous context latent; the other uses the window to route across experts. The first borrows a representation idea from hierarchy but recomputes its context every step, so it is not a true multirate temporal hierarchy. The system has two contexts that change how actions affect velocity:
- Context A is a damped point mass: each velocity component decays toward zero, and actions add acceleration.
- Context B adds a rotational coupling term in velocity space, as if the agent is inside a rotating frame.
Context is fixed per episode and its label is not supplied to the models. A flat model sees only (state, action). A history-conditioned model and an MoE router each receive a short state window that can provide more statistical evidence about the context. Finite state history without past actions does not guarantee exact identification, and a single post-divergence state can itself carry probabilistic context information.
Both contexts agree at the initial rest state. After velocity becomes nonzero, subsequent transitions reveal their different update laws, and a longer window may add evidence. Two experts provide the cardinality needed for a possible one-to-one context alignment, but training must still discover and validate that alignment. Multi-step rollout error is therefore the main outcome, while routing statistics diagnose how the MoE uses its capacity.
Generating the Data
We build a NumPy rollout that produces an episode given the hidden context. The system has four state dimensions (position and velocity in 2D) and two action dimensions (accelerations).
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
torch.manual_seed(7)
rng = np.random.default_rng(7)
DT = 0.05
EPISODE_HORIZON = 40
WINDOW = 8
LATENT_DIM = 2
N_EXPERTS = 2
HIDDEN = 64
TRAIN_EPOCHS = 25
LEARNING_RATE = 1e-3
BATCH_SIZE = 256
EVAL_HORIZON = 20
def rollout(context, horizon=EPISODE_HORIZON, actions=None):
if actions is None:
actions = 0.5 * rng.standard_normal((horizon, 2)).astype(np.float32)
else:
actions = np.asarray(actions, dtype=np.float32)
horizon = len(actions)
state = np.zeros(4, dtype=np.float32)
states = [state.copy()]
used_actions = []
for action in actions:
x, y, vx, vy = state
if context == 0:
dvx = action[0] - 0.2 * vx
dvy = action[1] - 0.2 * vy
else:
dvx = action[0] - 0.2 * vx + 1.2 * vy
dvy = action[1] - 0.2 * vy - 1.2 * vx
state = np.array(
[x + DT * vx, y + DT * vy, vx + DT * dvx, vy + DT * dvy],
dtype=np.float32,
)
states.append(state.copy())
used_actions.append(action)
return np.stack(states), np.stack(used_actions)
def make_dataset(n_episodes=400, horizon=EPISODE_HORIZON):
S, A, C = [], [], []
for i in range(n_episodes):
ctx = i % 2
s, a = rollout(ctx, horizon)
S.append(s)
A.append(a)
C.append(ctx)
return np.stack(S), np.stack(A), np.array(C)
S, A, C = make_dataset()Trajectories: 400, transition horizon: 40, state dim: 4 Contexts present: [0, 1]
The two contexts start from rest at the origin. They initially agree; once velocity is nonzero, an informative state-action transition can distinguish their noiseless update laws. Longer windows add statistical evidence. For the visualization, each context in a pair receives the exact same action sequence, so differences within a pair isolate the context dynamics.
trajectory_rng = np.random.default_rng(70)
paired_trajectories = {0: [], 1: []}
for _ in range(3):
paired_actions = 0.5 * trajectory_rng.standard_normal(
(EPISODE_HORIZON, 2)
).astype(np.float32)
for context in (0, 1):
states, _ = rollout(context, actions=paired_actions)
paired_trajectories[context].append(states)
Under the matched actions shown here, context B's rotational coupling curves the velocity response and reduces displacement on average across these three pairs. A state snapshot does not contain the context label, but after the dynamics have diverged its value may contain probabilistic evidence. A recent window supplies additional evidence about how velocity has evolved.
Preparing Windowed Transitions
For both window-conditioned models we build supervised examples of the form (window of recent states, current state, action, next state). The window offers evidence about which dynamics are in force, without supplying the context label or guaranteeing exact recovery.
def make_transitions(S, A, contexts, window=WINDOW):
"""Construct (window_states, s_t, a_t, s_{t+1}) tuples for every episode."""
N, T1, _ = S.shape
T = T1 - 1
wins, st, at, snext, ctxs = [], [], [], [], []
for i in range(N):
for t in range(window, T):
wins.append(S[i, t - window : t].reshape(-1))
st.append(S[i, t])
at.append(A[i, t])
snext.append(S[i, t + 1])
ctxs.append(contexts[i])
return (
torch.tensor(np.stack(wins), dtype=torch.float32),
torch.tensor(np.stack(st), dtype=torch.float32),
torch.tensor(np.stack(at), dtype=torch.float32),
torch.tensor(np.stack(snext), dtype=torch.float32),
torch.tensor(np.stack(ctxs), dtype=torch.long),
)
W_train, S_train, A_train, SN_train, C_train = make_transitions(
S, A, C, window=WINDOW
)Training transitions: 12800 Window flattened dim: 32
The output confirms the expected 32-feature dimension for an eight-state, four-dimensional window. The preceding slice and flattening operation supply the full contiguous recent-state window used for context inference.
Model Definitions
Three modules: a flat baseline that sees only (state, action), a history-conditioned model whose encoder compresses the window into a context latent, and a separate MoE with a learned router. The latter two receive more input and have more parameters than the flat baseline, so this is a diagnostic comparison, not a controlled estimate of an architectural advantage.
class FlatDynamics(nn.Module):
"""Baseline: predict next state from (state, action) only."""
def __init__(self, state_dim=4, action_dim=2, hidden=HIDDEN):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim + action_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, state_dim),
)
def forward(self, state, action):
delta = self.net(torch.cat([state, action], dim=-1))
return state + delta
class HistoryEncoder(nn.Module):
"""Compress a window of past states into a context latent."""
def __init__(
self, state_dim=4, window=WINDOW, latent_dim=LATENT_DIM, hidden=HIDDEN
):
super().__init__()
self.window = window
self.net = nn.Sequential(
nn.Linear(state_dim * window, hidden),
nn.ReLU(),
nn.Linear(hidden, latent_dim),
)
def forward(self, window_states):
return self.net(window_states)
class FastWithContext(nn.Module):
"""Predictor conditioned on the inferred context latent."""
def __init__(
self, state_dim=4, action_dim=2, latent_dim=LATENT_DIM, hidden=HIDDEN
):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim + action_dim + latent_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, state_dim),
)
def forward(self, state, action, latent):
delta = self.net(torch.cat([state, action, latent], dim=-1))
return state + delta
class HistoryConditionedDynamics(nn.Module):
"""History encoder + context-conditioned one-step dynamics."""
def __init__(
self,
state_dim=4,
action_dim=2,
window=WINDOW,
latent_dim=LATENT_DIM,
hidden=HIDDEN,
):
super().__init__()
self.history_encoder = HistoryEncoder(
state_dim, window, latent_dim, hidden
)
self.fast = FastWithContext(state_dim, action_dim, latent_dim, hidden)
def forward(self, window_states, state, action):
latent = self.history_encoder(window_states)
return self.fast(state, action, latent), latentTraining Loop
We share one training function across models. A lightweight forward adapter lets the flat, history-conditioned, and routed models use the same loop.
def train_model(
forward_fn,
model,
epochs=TRAIN_EPOCHS,
lr=LEARNING_RATE,
batch_size=BATCH_SIZE,
):
opt = torch.optim.Adam(model.parameters(), lr=lr)
dataset = TensorDataset(W_train, S_train, A_train, SN_train)
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
losses = []
for _ in range(epochs):
epoch_loss, n = 0.0, 0
for w, s, a, sn in loader:
pred = forward_fn(model, w, s, a)
loss = ((pred - sn) ** 2).mean()
opt.zero_grad()
loss.backward()
opt.step()
epoch_loss += loss.item() * w.shape[0]
n += w.shape[0]
losses.append(epoch_loss / n)
return losses
def flat_forward(model, window, state, action):
return model(state, action)
def history_forward(model, window, state, action):
pred, _ = model(window, state, action)
return predflat_model = FlatDynamics()
flat_losses = train_model(flat_forward, flat_model)
history_model = HistoryConditionedDynamics()
history_losses = train_model(history_forward, history_model)History-Conditioned vs. Flat Comparison
We compare training losses and multi-step rollout error on held-out test episodes.
def rollout_error(model, forward_fn, S, A, window=WINDOW, horizon=EVAL_HORIZON):
"""Average Euclidean position error over recursive rollout steps."""
errs = np.zeros(horizon)
counts = np.zeros(horizon)
for i in range(S.shape[0]):
s_true = S[i, : horizon + window + 1].copy()
s_pred = s_true[: window + 1].copy()
for t in range(window, window + horizon):
w = torch.tensor(
s_pred[t - window : t].reshape(1, -1), dtype=torch.float32
)
s = torch.tensor(s_pred[t].reshape(1, -1), dtype=torch.float32)
a = torch.tensor(A[i, t].reshape(1, -1), dtype=torch.float32)
with torch.no_grad():
s_pred_next = forward_fn(model, w, s, a).numpy()[0]
s_pred = np.concatenate([s_pred, s_pred_next[None]], axis=0)
errs[t - window] += np.linalg.norm(
s_pred[t + 1, :2] - s_true[t + 1, :2]
)
counts[t - window] += 1
return errs / counts
S_test, A_test, C_test = make_dataset(n_episodes=40, horizon=EPISODE_HORIZON)Wait, there is a subtlety: make_dataset uses the shared RNG, so S_test differs from S. That is intentional: we want a fresh test set. We now visualize both training curves and multi-step rollout error.
Both training losses are very small in this run. The history-conditioned model ends slightly lower, but it also receives a state window and has more parameters. No decisive difference appears in recursive rollouts.
flat_errs = rollout_error(flat_model, flat_forward, S_test, A_test)
history_errs = rollout_error(history_model, history_forward, S_test, A_test)
In this run the two models perform almost identically. The comparison neither isolates a benefit from history conditioning nor tests temporal abstraction: a controlled study would add a capacity-matched flat window baseline, match compute, and train across multiple seeds.
Adding a Mixture of Experts
Now let us compare the continuous context latent with a differentiable router and two experts. The experts see the same (state, action) inputs as the flat model, while the router uses the history window. The dense soft router blends both experts; it does not make a discrete mode decision, and specialization must be measured rather than assumed.
class MoEDynamics(nn.Module):
"""Mixture of experts with a window-conditioned router."""
def __init__(
self,
state_dim=4,
action_dim=2,
window=WINDOW,
n_experts=N_EXPERTS,
hidden=HIDDEN,
):
super().__init__()
self.n_experts = n_experts
self.experts = nn.ModuleList(
[
nn.Sequential(
nn.Linear(state_dim + action_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, state_dim),
)
for _ in range(n_experts)
]
)
self.router = nn.Sequential(
nn.Linear(state_dim * window, hidden),
nn.ReLU(),
nn.Linear(hidden, n_experts),
)
def forward(self, window_states, state, action):
x = torch.cat([state, action], dim=-1)
gates = torch.softmax(self.router(window_states), dim=-1)
deltas = torch.stack([e(x) for e in self.experts], dim=1) # (B, K, D)
delta = (gates.unsqueeze(-1) * deltas).sum(dim=1)
return state + delta, gates
def moe_forward(model, window, state, action):
pred, _ = model(window, state, action)
return pred
moe_model = MoEDynamics()
moe_losses = train_model(moe_forward, moe_model)moe_errs = rollout_error(moe_model, moe_forward, S_test, A_test)MoE final training loss: 0.0000 Flat final rollout error: 0.0261 History-conditioned final rollout error: 0.0267 MoE final rollout error: 0.0255
The MoE reaches a similarly small training loss, and its final position error is slightly below the flat baseline in this single seeded run. We next compute context-conditioned mean routing weights. This is one mode-alignment diagnostic, not proof of specialization: averages can hide sharp per-example routing or specialization along another axis.
def routing_by_context(model, S, A, C, window=WINDOW):
gates_by_context = {0: [], 1: []}
for i in range(S.shape[0]):
for t in range(window, S.shape[1] - 1):
w = torch.tensor(
S[i, t - window : t].reshape(1, -1), dtype=torch.float32
)
s = torch.tensor(S[i, t].reshape(1, -1), dtype=torch.float32)
a = torch.tensor(A[i, t].reshape(1, -1), dtype=torch.float32)
with torch.no_grad():
_, gates = model(w, s, a)
ctx = int(C[i])
gates_by_context[ctx].append(gates.numpy()[0])
arrays = {c: np.stack(values) for c, values in gates_by_context.items()}
means = {c: values.mean(axis=0) for c, values in arrays.items()}
mean_entropy = {
c: float((-values * np.log(values + 1e-12)).sum(axis=1).mean())
for c, values in arrays.items()
}
return means, mean_entropy
routing, routing_entropy = routing_by_context(moe_model, S_test, A_test, C_test)
context_routing_gap = 0.5 * np.abs(routing[0] - routing[1]).sum()
The context-conditioned means are nearly equal, so this diagnostic shows no clear one-to-one alignment between experts and hidden contexts. The conclusion is deliberately narrow. Per-example gate entropy is also useful, while expert-output differences, utilization, context probes, and expert ablations are needed before concluding that specialization is absent.
Context A mean gate entropy: 0.692 Context B mean gate entropy: 0.692 Total-variation gap between context mean gates: 0.004
Two observations close the example. First, neither the history-conditioned model nor the MoE establishes an advantage over the flat baseline in this single run. Second, the MoE's context-conditioned mean gates do not align with the known contexts. Mode alignment is useful when interpretable recovery is an explicit goal, but it should be paired with per-example routing statistics, probes, expert ablations, and task-level rollout performance. Low alignment alone does not prove that predictions average destructively across regimes.
Key Parameters
The key parameters for the worked history-conditioned and routed models are:
- WINDOW: Number of past states supplied to the history encoder and router. More history can add context evidence when it is relevant, but also increases input dimension and shortcut risk.
- LATENT_DIM: Dimensionality of the continuous context latent. The example does not claim that two dimensions are necessary or sufficient; that requires an ablation and a context probe.
- N_EXPERTS: Number of experts. Two experts permit a one-to-one alignment with two contexts, but the alignment still has to be learned and evaluated.
- HIDDEN: Width of the MLP hidden layers in the flat model, history encoder, context-conditioned predictor, router, and experts.
- LEARNING_RATE / BATCH_SIZE / TRAIN_EPOCHS: Optimizer step size, minibatch size, and training passes in the shared loop.
- DT: Forward-Euler simulation step. A smaller value generally reduces discretization error only when physical duration and the continuous action signal are held fixed while the step count changes accordingly.