Robotic Manipulation › A Minimal Manipulation World Model in Code
World Models HandbookAll chapters Robotic Manipulation › A Minimal Manipulation World Model in CodeOne 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 will build a miniature 2D pushing predictor. The state contains puck and pusher positions , and the action is a commanded pusher displacement . A post-move distance test gates an ad hoc puck coupling gain. This is a deliberately synthetic rule, not a physical friction model or a reproduction of a planar-pushing benchmark. It lets us inspect feature choice, prediction error, open-loop rollout, and a one-step effect map; it does not demonstrate grasping, dexterity, or visual foresight.
We will do six things:
- Generate random-play transition data from the contact-switch simulator and inspect how rarely the pre-move proximity feature is appreciably nonzero.
- Fit a naive linear world model that ignores contact.
- Fit a contact-aware world model that has access to a contact feature.
- Compare one-step and multi-step rollout error.
- Plan a push to a target using random shooting against the learned model.
- Query the model for a one-step effect field, an affordance proxy.
The core comparison gives one regression additional contact-proximity features. We report both one-step and recursive rollout error; neither alone measures task success.
Before building the models, we fix a random seed so every experiment is reproducible, and we keep each model's input features explicit so the comparison is fair.
Everything in this chapter runs in a single notebook, and the notebook follows a strict computation and rendering boundary: data generation, model fitting, rollout evaluation, planning, and affordance-map computation all happen in ordinary (non-theme-rerun) cells that run once. The figures are produced by separate presentation-only cells, each of which declares #| theme-rerun: true, imports the shared plotting style locally, calls use_book_style(), and renders only values that were prepared in earlier cells. No cell that draws a figure trains a model, samples from a simulator, or mutates shared data.
Setup
This cell seeds the examples and locates the repository root by walking upward from the working directory. Launch it from the repository or a descendant directory; it will not find the style module from an unrelated working directory.
Simulating contact-rich dynamics
The dynamics fit on one screen: the pusher moves by the commanded action subject to workspace clipping; the puck moves by a fixed fraction of the command when the pusher's post-move position passes a distance test. The gain is not a friction coefficient.
The simulator and feature function are separate. The latter measures pre-move proximity and therefore does not encode the simulator's post-move contact decision exactly.
import matplotlib.pyplot as plt
import numpy as np
from book_plot_style import PALETTE, polish_axes, use_book_style
R_CONTACT = 0.35 # pusher-puck center-distance switch threshold
ALPHA = 0.5 # puck follows pusher at half the commanded displacement
BOUND = 2.0 # workspace half-width
def step(state, action):
px, py, qx, qy = state
ux, uy = action
# Pusher moves by the command, clipped to the toy workspace
qx_new = np.clip(qx + ux, -BOUND, BOUND)
qy_new = np.clip(qy + uy, -BOUND, BOUND)
# Post-move pusher position gates puck motion. The full command is used
# even if the pusher was clipped: a nonphysical toy artifact.
dist = np.hypot(qx_new - px, qy_new - py)
if dist < R_CONTACT:
px_new = px + ALPHA * ux
py_new = py + ALPHA * uy
else:
px_new, py_new = px, py
return np.array([px_new, py_new, qx_new, qy_new])
def contact_feature(state):
px, py, qx, qy = state
d = np.hypot(qx - px, qy - py)
return float(np.clip(1.0 - d / R_CONTACT, 0.0, 1.0))The step function advances the pusher and then checks whether its new position is within R_CONTACT of the puck; only then does the puck inherit ALPHA times the full commanded action. If clipping shortened the pusher's move, the puck update is still based on the command—a nonphysical artifact of this toy. The feature decreases linearly from one at zero pre-move separation to zero at the distance threshold. It is an input to a fitted regression, not the simulator's contact rule. A deliberately simplified soft-contact illustration would have displacement
where:
- : the puck's displacement during step , computed as the change in puck position
- : the ad hoc coupling gain introduced above
- : the contact feature, equal to 1 at zero pusher-puck separation and decreasing linearly to 0 at the distance threshold
- : the pusher's commanded displacement
This soft formula is not the fitted model's exact equation: Ridge learns coefficients for state, action, , and , including an intercept. As a function of two-dimensional relative position, the supplied proximity feature is continuous but non-differentiable at coincident positions and at the distance-threshold boundary. As a function of scalar separation , its interior kink is at . It may help a regression represent contact dependence, but timing mismatch can limit that help.
The next plot compares two hypothetical weights at pre-move separation. It illustrates the feature's shape; the simulator itself decides contact from post-move separation.
The sweep that feeds that plot is prepared here, in a computation-only cell, so the figure cell below only renders the precomputed arrays.
## Sweep pre-move separation from zero to beyond the distance threshold.
## The hard curve is a pre-move illustration, not the simulator's post-move rule.
## The ramp is continuous but has a kink at the threshold.
separation = np.linspace(0.0, 0.6, 240)
hard_switch = (separation < R_CONTACT).astype(float)
smooth_feature = np.clip(1.0 - separation / R_CONTACT, 0.0, 1.0)
We will give this feature to one of the two world models and withhold it from the other, so we can see the difference. The contact-aware model uses it as an extra input feature, while the naive model sees only the raw state and action.
Collecting a training set
We generate synthetic random-action episodes from the toy simulator. This is not a faithful model of robot data collection; in this sampling scheme, pre-move proximity is rare, which matters for an unweighted regression loss.
Each transition is stored as an independent (state, action, next_state) tuple, which keeps the regression dataset simple: random play is the data source, not a trained controller.
rng = np.random.default_rng(0)
N_EPISODES = 400
T_EPISODE = 30
states = []
actions = []
next_states = []
for _ in range(N_EPISODES):
s = np.array(
[
rng.uniform(-0.5, 0.5),
rng.uniform(-0.5, 0.5), # puck
rng.uniform(-1.5, -1.0)
if rng.random() < 0.5
else rng.uniform(1.0, 1.5),
rng.uniform(-1.5, 1.5), # pusher
]
)
for _ in range(T_EPISODE):
a = rng.normal(scale=0.15, size=2)
s_next = step(s, a)
states.append(s)
actions.append(a)
next_states.append(s_next)
s = s_next
states = np.array(states)
actions = np.array(actions)
next_states = np.array(next_states)Transitions: 12000 Transitions with contact feature > 0.1: 64 (0.5%)
The printed summary counts states whose pre-move contact feature exceeds 0.1. It is a proximity statistic, not the fraction of transitions in which the simulator's post-move contact test succeeds. Most sampled states have a feature near zero. An unweighted loss over this dataset can therefore be dominated by common non-proximal states; whether those are also the easiest transitions depends on the sampled actions and outcomes. Targeted contact evaluation or deliberate contact-rich sampling would test the rare regime more directly.
That concentration of the data is easy to see directly. Plotting the contact feature over every transition shows how few samples carry contact information.
## Contact feature value for every transition in the random-play dataset. This
## reuses the states already generated and performs no new sampling.
dataset_contact = np.array([contact_feature(s) for s in states])
Two world models
We compare a naive linear model that regresses the next state from the current state and action using only a fixed affine map, with no contact term of any kind:
where:
- : the state at time , the concatenation of the puck's position and the pusher's position
- : the pusher's commanded displacement during step
- : the learned state-transition matrix, capturing how the current state propagates in the absence of action
- : the learned control matrix, mapping the commanded displacement into the next state
- : the learned intercept vector (Ridge fits an intercept by default)
- : the model's predicted next state, whose error is measured against the simulator's true
- The affine prediction approximates the true next state; its errors are measured below.
against a contact-aware model that augments the input with the contact feature and its product with the action. The fitted coefficients remain fixed. The extra nonlinear features let this linear regression represent an action-dependent contact response without an explicit latent contact mode:
from sklearn.linear_model import Ridge
from sklearn.metrics import mean_squared_error
def augmentation_naive(state, action):
return np.concatenate([state, action])
def augmentation_contact(state, action):
c = contact_feature(state)
return np.concatenate([state, action, [c], c * action])
## The two feature builders share the current state and action. The contact-aware
## version adds the scalar contact feature and its product with the action, so the
## regression can represent a contact-sensitive interaction without a hidden mode variable.
X_naive = np.array([augmentation_naive(s, a) for s, a in zip(states, actions)])
X_contact = np.array(
[augmentation_contact(s, a) for s, a in zip(states, actions)]
)
Y = next_states
split = int(0.8 * len(X_naive))
train = slice(0, split)
test = slice(split, len(X_naive))
model_naive = Ridge(alpha=1e-4).fit(X_naive[train], Y[train])
model_contact = Ridge(alpha=1e-4).fit(X_contact[train], Y[train])
## Both models use the same ridge penalty so the comparison isolates the effect of
## the contact feature rather than differences in regularization strength.Naive one-step RMSE: 0.0186 Contact one-step RMSE: 0.0183
Both models are small regressions on hand-crafted features; only one receives the smoothed contact feature and its action products. In this seeded run, the contact-aware model's one-step RMSE is only marginally lower. Contact is rare in the sampled random-play data, and the feature uses pre-move separation while the simulator checks post-move separation. The added feature therefore does not isolate the simulator's hard switch. This comparison tests what the feature contributes under this data distribution; it does not establish that contact-aware modeling is generally better or worse.
Multi-step rollout
One-step and recursive errors answer different questions. A planner may feed predictions back into the model; recursive rollout exposes the resulting distribution shift. The following 15-step test does not measure closed-loop control or establish an asymptotic growth law.
The first 9,600 transitions come from 320 complete episodes and the remaining 2,400 from 80 complete episodes, so the one-step train/test slices do not split episodes. The recursive evaluation below draws fresh start states and actions (50 trials), rather than reusing those test episodes. Both models see identical actions within each trial.
def rollout(model, aug_fn, s0, actions):
preds = [s0]
s = s0
for a in actions:
s = model.predict(aug_fn(s, a).reshape(1, -1))[0]
preds.append(s)
return np.array(preds)
## The earlier contiguous split aligns with 30-step episode boundaries.
## The following recursive trials are freshly sampled.
HORIZON = 15
n_eval = 50
rollout_errors = {"naive": [], "contact": []}
for _ in range(n_eval):
s0 = np.array(
[
rng.uniform(-0.5, 0.5),
rng.uniform(-0.5, 0.5),
rng.uniform(-1.2, 1.2),
rng.uniform(-1.2, 1.2),
]
)
acts = rng.normal(scale=0.15, size=(HORIZON, 2))
truth = [s0]
s = s0
for a in acts:
s = step(s, a)
truth.append(s)
truth = np.array(truth)
for name, mdl, aug in [
("naive", model_naive, augmentation_naive),
("contact", model_contact, augmentation_contact),
]:
preds = rollout(mdl, aug, s0, acts)
err = np.sqrt(mean_squared_error(truth, preds))
rollout_errors[name].append(err)
rollout_errors = {k: float(np.mean(v)) for k, v in rollout_errors.items()}naive mean 15-step RMSE: 0.0452 contact mean 15-step RMSE: 0.0480
The printed whole-trajectory RMSE averages over sixteen states: the exact initial state and fifteen recursively predicted next states. The initial zero-error row is therefore included. The separate per-step curves below show how prediction error changes across the fifteen-step horizon. In this seeded run, the contact-aware model has slightly higher mean rollout RMSE than the naive baseline despite its slightly lower one-step error. Reusing predicted states changes the input distribution at every step; an extra feature does not guarantee a better rollout when contact examples are sparse and its timing differs from the simulator's contact test. For a planner, this mismatch matters because candidate actions are scored on predicted consequences. Neither rollout RMSE nor one-step RMSE, however, substitutes for closed-loop task success.
The per-step plot makes the horizon dependence visible. Both error curves rise; their relative ordering varies across steps rather than separating consistently in favor of the contact-aware model.
## Fresh recursive rollouts to measure error across this 15-step horizon.
## Evaluation design: 30 start states and action sequences drawn from a dedicated
## seeded generator (seed 7), independent of the training and planning samples.
## Both models see identical start states and actions. Per-step error is the
## root-mean-square over the four state dimensions averaged across the 30
## rollouts. No tie-breaking is needed because each step yields one scalar error.
rng_roll = np.random.default_rng(7)
N_ROLL = 30
per_step_naive = np.zeros(HORIZON)
per_step_contact = np.zeros(HORIZON)
for _ in range(N_ROLL):
s0 = np.array(
[
rng_roll.uniform(-0.5, 0.5),
rng_roll.uniform(-0.5, 0.5),
rng_roll.uniform(-1.2, 1.2),
rng_roll.uniform(-1.2, 1.2),
]
)
acts = rng_roll.normal(scale=0.15, size=(HORIZON, 2))
truth = [s0]
s = s0
for a in acts:
s = step(s, a)
truth.append(s)
truth = np.array(truth)
preds_naive = rollout(model_naive, augmentation_naive, s0, acts)
preds_contact = rollout(model_contact, augmentation_contact, s0, acts)
for h in range(1, HORIZON + 1):
per_step_naive[h - 1] += np.sqrt(
np.mean((truth[h] - preds_naive[h]) ** 2)
)
per_step_contact[h - 1] += np.sqrt(
np.mean((truth[h] - preds_contact[h]) ** 2)
)
per_step_naive /= N_ROLL
per_step_contact /= N_ROLL
rollout_horizon = np.arange(1, HORIZON + 1)
Planning a push with random shooting
We use the contact-feature model for a small random-shooting demonstration (Part VII). It samples 400 candidate 12-step action sequences, scores each by the model-predicted final puck distance to a target, and selects one. The chosen sequence is then replayed open-loop in the simulator. This is not receding-horizon control, and the score does not establish task success.
The planner searches only over actions; the contact-aware model supplies the prediction that scores each candidate sequence, which keeps the planning loop separate from the learned dynamics.
def plan_push(
model,
aug_fn,
s0,
target,
n_samples=400,
horizon=12,
action_scale=0.15,
plan_rng=None,
):
if plan_rng is None:
plan_rng = rng
best_score = np.inf
best_actions = None
best_preds = None
for _ in range(n_samples):
acts = plan_rng.normal(scale=action_scale, size=(horizon, 2))
preds = rollout(model, aug_fn, s0, acts)
final_puck = preds[-1, :2]
score = np.linalg.norm(final_puck - target)
if score < best_score:
best_score = score
best_actions = acts
best_preds = preds
return best_actions, best_preds, best_score
start_state = np.array([0.0, 0.0, -1.0, 0.0])
target = np.array([0.6, 0.4])
best_actions, best_preds, best_score = plan_push(
model_contact, augmentation_contact, start_state, target
)Final puck distance to target: 0.4156 Puck start: [0. 0.] Puck end: [0.265 0.154] Target: [0.6 0.4]
The printed score is the model's predicted final distance to the target, not the achieved task error. In this seeded run it remains much larger than the action sampling scale, so the selected sequence is only a coarse approach, not a successful placement. The gap can reflect candidate coverage, the learned transition, and their interaction; the run does not isolate one cause. Replaying the sequence in the simulator below tests whether the predicted trajectory survives execution. A practical controller would execute a short prefix, observe the result, and replan. That closed-loop outcome, rather than the planning score alone, would determine task success.
Replaying the chosen actions through the simulator lets us compare predicted and true puck paths. The pusher path drawn below is also predicted by the learned model; the simulator's pusher path is available separately in true_preds.
## Roll the chosen action sequence through the true simulator so the learned
## prediction can be compared against ground truth. No sampling happens here; the
## actions are exactly the ones returned by the planner above.
true_preds = [start_state]
s = start_state
for a in best_actions:
s = step(s, a)
true_preds.append(s)
true_preds = np.array(true_preds)
pred_puck = best_preds[:, :2]
true_puck = true_preds[:, :2]
pusher_path = best_preds[:, 2:]
Affordance map from the same data
The fitted model can be queried for a one-step effect field: from each selected pusher start, apply the same rightward command and plot the predicted puck displacement. This is one deterministic model query per cell, not an average, probability, confidence estimate, or goal-conditioned success score. Arrow magnitudes may vary for reasons that include sparse contact data and model error; their strongest values need not occur to the left of the puck.
The map displays only start cells whose pre-move contact feature exceeds 0.05. This is a visualization filter, not a guarantee that the simulator's post-move contact test would accept or reject the same cell.
## Grid of pusher start positions around a puck at the origin
puck = np.array([0.0, 0.0])
grid_vals = np.linspace(-0.6, 0.6, 15)
push_action = np.array([0.15, 0.0]) # always push right
aff_x, aff_y, aff_dx, aff_dy = [], [], [], []
for gx in grid_vals:
for gy in grid_vals:
s0 = np.array([puck[0], puck[1], gx, gy])
# Filter by pre-move proximity; the simulator tests post-move contact.
if contact_feature(s0) <= 0.05:
continue
s1 = model_contact.predict(
augmentation_contact(s0, push_action).reshape(1, -1)
)[0]
aff_x.append(gx)
aff_y.append(gy)
aff_dx.append(s1[0] - puck[0])
aff_dy.append(s1[1] - puck[1])
aff_x = np.array(aff_x)
aff_y = np.array(aff_y)
aff_dx = np.array(aff_dx)
aff_dy = np.array(aff_dy)
max_mag = max(np.hypot(aff_dx, aff_dy).max(), 1e-6)
The arrow field queries the learned model at selected start positions without running the simulator again. Its pre-move filter keeps the display near the puck but can omit starts that would make contact after the commanded move. Each arrow encodes the model's predicted puck displacement for the same rightward action. A task planner could rank candidate starts with these predictions, then validate selected actions in closed loop. This map alone does not establish that the strongest-looking arrow corresponds to a successful real push.
The arrow magnitude is the model's predicted effect, not confidence. A large vector can be confidently right, confidently wrong, or uncertain; this plot cannot distinguish those cases. Comparing queries with held-out simulator transitions or real actions would test whether these effects are useful for choosing a push.
Key Parameters
The key parameters used in the miniature pushing experiment are:
- R_CONTACT: Pusher-puck center-distance threshold for the toy simulator's contact switch.
- N_EPISODES: Number of random-play episodes collected for training and evaluation.
- ALPHA: Fraction of the commanded displacement inherited by the puck when the simulator's contact test succeeds.
- BOUND: Workspace half-width used to clip pusher positions.
- grid_vals: Sampled pusher start positions for the illustrative affordance field. Increase this to visualize the affordance field at finer spatial resolution.
- push_action: Fixed displacement used to query that field.
- max_mag: Maximum predicted displacement magnitude, used to normalize displayed arrow lengths.
- T_EPISODE: 30 transitions per episode; with 400 episodes, the dataset contains 12,000 transitions.
- Train/test split and Ridge alpha: An episode-aligned 80/20 split (320/80 episodes); each linear model uses
Ridge(alpha=1e-4). - HORIZON and evaluation counts: Recursive rollouts use 15 actions, with 50 fresh trials for the aggregate RMSE and 30 fresh trials for the per-step plot.
- Planner search settings: Random shooting samples 400 action sequences of length 12 with Gaussian action scale 0.15.