Vision-Language-Action and World-Action Models › Code Implementation: A Miniature World-Action Model
World Models HandbookAll chapters Vision-Language-Action and World-Action Models › Code Implementation: A Miniature World-Action ModelOne section · Sections in this chapter7 / 10 ↓
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 complete, tiny example on a language-conditioned reaching task. The task is deliberately simple, because the point is to observe the mechanisms, not to reproduce a benchmark: a language token selects one of four goals, and a world-action model must both pick actions and predict the consequences of actions. Along the way we will see action quantization, mode averaging, compounding prediction error, and the gap between one-step error and closed-loop success. Every mechanism we isolate here has a scaled-up counterpart in the real systems discussed above.
Simulating language-conditioned demonstrations
We start with a two-dimensional point mass with velocity damping, which is a linear system we can reason about exactly. The state is , and the action is a velocity command.
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nnDT = 0.1
DAMP = 0.9
START = np.array([-1.0, 0.0])
GOALS = np.array([[1.0, -0.6], [1.0, -0.2], [1.0, 0.2], [1.0, 0.6]])
COLORS = ["red", "blue", "green", "yellow"]
VOCAB = {"reach": 0, "red": 1, "blue": 2, "green": 3, "yellow": 4}
SUCCESS_RADIUS = 0.15
def step(state, action):
"""Double integrator with velocity damping: s' = f(s, a)."""
velocity = DAMP * state[2:] + (1.0 - DAMP) * action
position = state[:2] + DT * velocity
return np.concatenate([position, velocity])
def expert_action(state, goal, kp=3.0, kv=3.0, v_max=1.0):
"""Cascaded controller: clip a desired velocity, then track it."""
desired = np.clip(kp * (goal - state[:2]), -v_max, v_max)
return np.clip(kv * (desired - state[2:]), -1.0, 1.0)The dynamics are deterministic and linear, so this network has enough capacity to represent the transition rule. A learned fit can still have optimization and finite-data error, while free-running use adds the distinct problem of feeding model-predicted states back into the input. The simple system lets us separate those effects more clearly than a contact-rich robot would.
Next we collect demonstrations with feedback. At each realized state, the controller recomputes its command before execution noise is added. The recorded state-action-next-state triples therefore follow the noisy trajectory the controller sees, rather than replaying commands from a clean trajectory open-loop.
def rollout(action_fn, steps, start=None):
state = (
np.concatenate([START, np.zeros(2)])
if start is None
else np.asarray(start, dtype=float).copy()
)
states = [state.copy()]
for _ in range(steps):
state = step(state, np.clip(action_fn(state), -1.0, 1.0))
states.append(state.copy())
return np.array(states)
def instruction_tokens(color):
return np.array([VOCAB["reach"], VOCAB[color]], dtype=np.int64)rng = np.random.default_rng(0)
N_EPISODES = 400
STEPS_PER_EPISODE = 60
ACTION_NOISE = 0.35
state_log, token_log, action_log, next_log, episode_log = [], [], [], [], []
for episode in range(N_EPISODES):
color = COLORS[episode % len(COLORS)]
goal = GOALS[COLORS.index(color)]
state = np.concatenate([START, np.zeros(2)])
for t in range(STEPS_PER_EPISODE):
commanded = expert_action(state, goal)
executed = np.clip(
commanded + rng.normal(0.0, ACTION_NOISE, size=2), -1.0, 1.0
)
nxt = step(state, executed)
state_log.append(state.copy())
token_log.append(instruction_tokens(color))
action_log.append(executed)
next_log.append(nxt)
episode_log.append(episode)
state = nxt
states = np.array(state_log)
tokens = np.array(token_log)
actions = np.array(action_log)
next_states = np.array(next_log)
episodes = np.array(episode_log)
DEMO_EPISODES_PER_COLOR = 5
demo_examples = {}
for index, color in enumerate(COLORS):
paths = []
for episode in range(index, N_EPISODES, len(COLORS)):
if len(paths) == DEMO_EPISODES_PER_COLOR:
break
block = states[
episode * STEPS_PER_EPISODE : (episode + 1) * STEPS_PER_EPISODE
]
paths.append(block[:, :2])
demo_examples[color] = paths
final_positions = next_states.reshape(N_EPISODES, STEPS_PER_EPISODE, 4)[
:, -1, :2
]
final_goal_distance = np.linalg.norm(
final_positions - GOALS[np.arange(N_EPISODES) % len(GOALS)], axis=1
)
The demonstration set has 24,000 transitions. The fraction within the goal tolerance uses the final state reached after noisy, feedback-corrected commands. Every episode starts from the same state, so the recorded paths cover only part of the state space.
The final positions come from the recorded next state after each episode's last noisy action. The following summary reports distance to the instructed goal for those realized trajectories.
transitions: 24000 state dim: 4 action dim: 2 instruction vocabulary size: 5 expert episodes settling within 0.15: 100.0% mean expert settling distance: 0.0173
The four instruction colors map to four separated goals, and the spread within each color shows the effect of action noise and feedback correction. The trajectories still cover a narrow region around the demonstrated tasks, not the full state space. The policy has no training examples for most states outside those paths.
The action-space choice: regression versus discrete bins
Before training anything, let us quantify the cost of discretizing actions and the cost of regressing them.
BIN_COUNTS = np.array([4, 8, 16, 32, 64, 128, 256, 512, 1024])
worst_case_error = 1.0 / (BIN_COUNTS - 1)
empirical_rms = []
for bins in BIN_COUNTS:
index = np.clip(np.round((actions + 1.0) / 2.0 * (bins - 1)), 0, bins - 1)
decoded = index / (bins - 1) * 2.0 - 1.0
empirical_rms.append(float(np.sqrt(((decoded - actions) ** 2).mean())))
empirical_rms = np.array(empirical_rms)
Two bin counts show the scale of the tradeoff. At the RMS error is around 0.0171 normalized action units, and at it is around 0.0021, roughly eight times smaller. These are action-space errors, not position errors; translating them into a millimeter tolerance requires the action semantics and the closed-loop controller.
Discretization also changes what the policy can represent when demonstrations contain two equally valid action strategies.
bimodal_rng = np.random.default_rng(3)
mode_upper = np.clip(bimodal_rng.normal(0.55, 0.08, 700), -1.0, 1.0)
mode_lower = np.clip(bimodal_rng.normal(-0.55, 0.08, 400), -1.0, 1.0)
bimodal_samples = np.concatenate([mode_upper, mode_lower])
bimodal_mean = float(bimodal_samples.mean())
bimodal_median = float(np.median(bimodal_samples))
bimodal_edges = np.linspace(-1.0, 1.0, 25)
bimodal_counts, _ = np.histogram(bimodal_samples, bins=bimodal_edges)
bimodal_centers = 0.5 * (bimodal_edges[:-1] + bimodal_edges[1:])
The mean sits in the empty valley between the two strategies, an action nobody ever demonstrated. A categorical head over bins, by construction, can place probability mass on both peaks, and a diffusion or flow head can sample from either. That is the strongest argument for tokenized or generative action heads, and it has nothing to do with precision.
A miniature world-action model
Now we build the model. A shared state encoder feeds two paths: the policy combines its state features with instruction tokens to choose an action, while the dynamics path combines the same state features with an action to predict the next state. Both losses update the shared encoder, making the predictive objective capable of changing the policy representation.
class WorldActionModel(nn.Module):
"""A language-conditioned policy plus an action-conditioned dynamics model."""
def __init__(self, vocab_size, n_bins=None, d_model=64):
super().__init__()
self.n_bins = n_bins
self.instruction = nn.Embedding(vocab_size, d_model)
self.state_encoder = nn.Sequential(
nn.Linear(4, d_model),
nn.ReLU(),
nn.Linear(d_model, d_model),
nn.ReLU(),
)
self.policy_trunk = nn.Sequential(
nn.Linear(d_model, d_model), nn.ReLU()
)
self.dynamics_trunk = nn.Sequential(
nn.Linear(d_model + 2, d_model),
nn.ReLU(),
nn.Linear(d_model, d_model),
nn.ReLU(),
)
self.action_head = nn.Linear(
d_model, 2 if n_bins is None else 2 * n_bins
)
self.next_state_head = nn.Linear(d_model, 4)
def policy_output(self, states, token_ids):
embedded = self.instruction(token_ids).mean(dim=1)
hidden = self.policy_trunk(self.state_encoder(states)) + embedded
out = self.action_head(hidden)
return out if self.n_bins is None else out.view(-1, 2, self.n_bins)
def next_state(self, states, action):
return self.next_state_head(
self.dynamics_trunk(
torch.cat([self.state_encoder(states), action], dim=-1)
)
)Two design notes. The instruction embedding is added to the policy-path output rather than concatenated. policy_output returns continuous actions or per-dimension logits shaped (batch, dim, n_bins), depending on whether n_bins is set. next_state predicts the state at from the state at and the action at ; it does not receive the instruction. The shared state encoder is the only component updated by both losses, so an auxiliary prediction loss can influence the policy without making the two output heads identical.
The split below holds out 80 episodes, but both sides retain the same four goals, instruction tokens, and nominal start. Its offline metrics test new trajectories within this synthetic task, not transfer to an unseen task family. The perturbed-start rollout later tests a different, local form of generalization.
TRAIN_EPISODES = 320
BATCH = 256
N_BINS = 64
train_mask = episodes < TRAIN_EPISODES
test_mask = ~train_mask
train_states = torch.tensor(states[train_mask], dtype=torch.float32)
train_tokens = torch.tensor(tokens[train_mask], dtype=torch.long)
train_actions = torch.tensor(actions[train_mask], dtype=torch.float32)
train_next_states = torch.tensor(next_states[train_mask], dtype=torch.float32)
test_states = torch.tensor(states[test_mask], dtype=torch.float32)
test_tokens = torch.tensor(tokens[test_mask], dtype=torch.long)
test_actions = torch.tensor(actions[test_mask], dtype=torch.float32)
test_next_states = torch.tensor(next_states[test_mask], dtype=torch.float32)def quantize(action, n_bins):
return torch.clamp(
((action + 1.0) / 2.0 * (n_bins - 1)).round().long(), 0, n_bins - 1
)
def decode_bins(bins, n_bins):
return bins.float() / (n_bins - 1) * 2.0 - 1.0
def train_model(variant, epochs=30, aux_weight=1.0, seed=0):
torch.manual_seed(seed)
model = WorldActionModel(
vocab_size=len(VOCAB),
n_bins=None if variant == "continuous" else N_BINS,
)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
history = []
for _ in range(epochs):
epoch_loss = 0.0
for batch in torch.randperm(train_states.shape[0]).split(BATCH):
s, tok = train_states[batch], train_tokens[batch]
a, ns = train_actions[batch], train_next_states[batch]
out = model.policy_output(s, tok)
if variant == "continuous":
action_loss = ((out - a) ** 2).mean()
else:
target = quantize(a, N_BINS)
action_loss = nn.functional.cross_entropy(
out.reshape(-1, N_BINS), target.reshape(-1)
)
dynamics_loss = ((model.next_state(s, a) - ns) ** 2).mean()
loss = action_loss + aux_weight * dynamics_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
epoch_loss += loss.detach().item() * batch.numel()
history.append(epoch_loss / train_states.shape[0])
return model, history
continuous_model, continuous_history = train_model("continuous")
token_model, token_history = train_model("tokenized")The loss has an action term and a next-state term. aux_weight controls how strongly the next-state error updates the shared state encoder, as well as the separate dynamics path. Each call to train_model initializes and trains a separate model with the same dataset, optimizer settings, and seed, but a different action-output parameterization and loss. An observed performance difference therefore cannot be attributed to action representation alone.
variant action MSE next-state MSE continuous 0.106577 0.00005836 tokenized 0.165972 0.00009249
The two independently trained models both achieve low next-state error, but their dynamics heads are separate and their next-state MSEs differ slightly. The continuous policy also has lower one-step action MSE in this run. Neither metric alone tells us how often either policy will reach its goal when its own actions determine the states it visits. The closed-loop experiment below measures that outcome directly; neither policy consults its learned dynamics head during those rollouts.
We can also run the counterfactual action-sensitivity check described earlier. At one held-out state, reset the toy transition to the same starting point, apply two different actions, and compare the true change in next state with the learned predictor's change. This is a local diagnostic, not evidence that the predictor ranks every useful action correctly.
probe_state = test_states[0].numpy()
probe_actions = np.array([[0.0, 0.0], [0.8, -0.8]], dtype=np.float32)
probe_states = np.repeat(probe_state[None, :], 2, axis=0)
true_next = np.stack([step(probe_state, action) for action in probe_actions])
with torch.no_grad():
predicted_next = continuous_model.next_state(
torch.tensor(probe_states), torch.tensor(probe_actions)
).numpy()
print("true action-effect delta:", np.round(true_next[1] - true_next[0], 4))
print(
"predicted action-effect delta:",
np.round(predicted_next[1] - predicted_next[0], 4),
)true action-effect delta: [ 0.008 -0.008 0.08 -0.08 ] predicted action-effect delta: [ 0.0018 -0.001 0.0632 -0.0792]
The true action-effect delta is ; the fitted model predicts approximately . It captures the velocity-change direction and approximate magnitude, but understates the position change at this state. This one comparison can reveal weak action conditioning at the probe state; it does not measure error across other states or longer counterfactual rollouts.
Closed-loop evaluation in the true environment
Now we roll the policies out in the true environment, including from perturbed initial states so that recovery behavior is tested rather than just nominal tracking. A trial counts as a success if the policy reaches the goal tolerance at least once during its final 20 states; unlike the demonstration-set summary, this is not a final-state settling criterion.
def make_policy_fn(model, token_tensor):
def act(state):
with torch.no_grad():
out = model.policy_output(
torch.tensor(state, dtype=torch.float32).unsqueeze(0),
token_tensor.unsqueeze(0),
)
if model.n_bins is None:
return out.squeeze(0).numpy()
return decode_bins(out.argmax(dim=-1), model.n_bins).squeeze(0).numpy()
return act
def evaluate_policy(
action_fn_factory, trials_per_goal=25, steps=120, start_noise=0.10, seed=11
):
eval_rng = np.random.default_rng(seed)
successes = []
for index, color in enumerate(COLORS):
act = action_fn_factory(index, color)
hits = 0
for _ in range(trials_per_goal):
start = np.concatenate([START, np.zeros(2)]) + eval_rng.normal(
0.0, start_noise, size=4
)
path = rollout(act, steps, start=start)
distance = np.linalg.norm(
path[-20:, :2] - GOALS[index], axis=1
).min()
hits += float(distance < SUCCESS_RADIUS)
successes.append(hits / trials_per_goal)
return float(np.mean(successes))success_rates = {}
success_rates["expert"] = evaluate_policy(
lambda index, color: lambda s: expert_action(s, GOALS[index]), seed=11
)
for name, model, seed in [
("continuous policy", continuous_model, 11),
("tokenized policy", token_model, 11),
]:
success_rates[name] = evaluate_policy(
lambda index, color, model=model: make_policy_fn(
model, torch.tensor(instruction_tokens(color), dtype=torch.long)
),
seed=seed,
)
def wilson_interval(rate, trials, z=1.96):
denom = 1.0 + z**2 / trials
center = (rate + z**2 / (2.0 * trials)) / denom
half_width = (
z
/ denom
* np.sqrt(rate * (1.0 - rate) / trials + z**2 / (4.0 * trials**2))
)
return max(0.0, center - half_width), min(1.0, center + half_width)
success_intervals = {
name: wilson_interval(rate, 25 * len(COLORS))
for name, rate in success_rates.items()
}policy closed-loop success expert 100% continuous policy 100% tokenized policy 84% One-step action MSE (lower is better): continuous < tokenized Closed-loop success (higher is better): continuous > tokenized
Read that output carefully. The one-step MSE ranking and the closed-loop success ranking need not agree in general, because closed-loop success measures whether errors compound or self-correct as a policy visits states induced by its own actions. In this seeded run, the continuous and tokenized policies achieve 100% and 84% success, respectively, across 100 matched perturbed starts each. That is a 16-percentage-point observed gap, not an equivalence result. The separate model fits, one seed, and small synthetic task do not establish that the output parameterization caused the gap. The 100% expert and continuous-policy rates in this finite trial set also do not imply perfect reliability outside it.
Open-loop versus re-anchored model rollouts
The final experiment isolates the compounding problem directly. We take held-out demonstration action sequences and compare two ways of using the learned dynamics model: free-running, where each prediction is fed back as the next input, and re-anchored, where every prediction starts from the true state.
HORIZONS = np.arange(1, 41)
def model_rollout_error(model, episode_list, horizons):
free_running = np.zeros(len(horizons))
reanchored = np.zeros(len(horizons))
count = 0
for episode in episode_list:
block_states = states[
episode * STEPS_PER_EPISODE : (episode + 1) * STEPS_PER_EPISODE
]
block_actions = actions[
episode * STEPS_PER_EPISODE : (episode + 1) * STEPS_PER_EPISODE
]
for t in range(0, STEPS_PER_EPISODE - horizons[-1]):
predicted = torch.tensor(
block_states[t], dtype=torch.float32
).unsqueeze(0)
for step, _ in enumerate(horizons):
action = torch.tensor(
block_actions[t + step], dtype=torch.float32
).unsqueeze(0)
with torch.no_grad():
predicted = model.next_state(predicted, action)
anchored = model.next_state(
torch.tensor(
block_states[t + step], dtype=torch.float32
).unsqueeze(0),
action,
)
target = block_states[t + step + 1]
free_running[step] += float(
np.linalg.norm(predicted.squeeze(0).numpy() - target)
)
reanchored[step] += float(
np.linalg.norm(anchored.squeeze(0).numpy() - target)
)
count += 1
return np.maximum(free_running / count, 1e-12), np.maximum(
reanchored / count, 1e-12
)
held_out_episodes = sorted(set(episodes[test_mask].tolist()))[:20]
free_running_error, reanchored_error = model_rollout_error(
continuous_model, held_out_episodes, HORIZONS
)
The re-anchored curve stays low and declines from about 0.018 to 0.009 across the plotted horizons. Each of its predictions starts from a true state, avoiding recursive prediction feedback; the downward trend also reflects the particular states and actions sampled at later horizons. Free-running error grows from about 0.018 at one step to about 0.654 at forty steps. Its consecutive-step growth ratios trend downward overall, with local increases, rather than remaining constant. This is accumulated drift in this experiment, not a measured constant-factor exponential growth law. The result is specific to this near-linear system and its training coverage; a contact-rich robot may have a different error horizon.
The offline table and closed-loop chart answer different questions. Here the continuous policy has a lower action MSE and a 16-percentage-point higher observed success rate, but this single seeded toy comparison cannot identify a general advantage for either head. Real-robot tasks may change the ranking, especially when actions are multimodal or dynamics include contact discontinuities. The methodological point holds regardless of which side wins: report closed-loop outcomes with uncertainty, and use offline error as a debugging aid.