OGBench cube MLP world models (diffusion-free control)
Deterministic MLP dynamics models trained on OGBench cube-single and
cube-double. They exist as a control: a diffusion-free baseline that
measures what the denoising process in a diffusion world model actually buys.
They are not proposed as the better model — see the results below.
Each model regresses one step of dynamics,
(state_t, action_t) -> state_{t+1}
as a [256, 256, 256] SiLU MLP over the shared normalized state space. It
predicts the normalized state delta from a zero-initialized output head,
so an untrained model sits at the identity transition. Sampling is a single
forward pass per step, against 20 denoising passes per step for the diffusion
model it is compared with.
Headline result
One-step accuracy does not predict rollout quality. The MLP is better than
the diffusion model at one-step prediction on both tasks and worse over a
100-step rollout on both tasks.
Validation normalized MSE after four exhaustive epochs. The diffusion control
is the heterogeneous Diffusion Forcing transformer trained on identical data
under an identical schedule.
task
model
one-step NMSE
100-step rollout NMSE
cube-single
MLP (this repo)
0.000135
0.00735
cube-single
Diffusion Forcing
0.000553
0.00163
cube-double
MLP (this repo)
0.0000487
0.01628 *
cube-double
Diffusion Forcing
0.000657
0.00435
* best checkpoint, epoch 2. The final epoch reaches 2.82 — see below.
The MLP is 4.1x (cube-single) and 13.5x (cube-double) better at one step, and
4.5x and 3.7x worse over 100 steps.
The failure mode
cube-double shows it sharply. Across four epochs its one-step error falls
monotonically while its rollout error blows up in the final epoch:
epoch
one-step NMSE
rollout NMSE
1
0.0000898
0.02135
2
0.0000647
0.01628
3
0.0000525
0.02174
4
0.0000487
2.8221
A one-step regressor with a small consistent bias compounds that bias over 100
autoregressive steps. The denoising step of a diffusion model projects each
predicted state back toward the data manifold; a deterministic MLP has nothing
that does this, so it drifts.
Two practical consequences:
Do not select or early-stop these models on one-step error. It improves
while the thing you care about degrades.
Use world_model_best.pt, not world_model_final.pt, for
cube-double. Best is selected on rollout NMSE. For cube-single the two
are the same checkpoint (epoch 4).
Files
cube-single/world_model_best.pt # epoch 4, rollout NMSE 0.00735
cube-single/world_model_final.pt # identical checkpoint
cube-single/metrics.jsonl # full per-step training + per-epoch validation
cube-double/world_model_best.pt # epoch 2, rollout NMSE 0.01628 <- use this
cube-double/world_model_final.pt # epoch 4, rollout NMSE 2.8221 <- diverged
cube-double/metrics.jsonl
Each .pt is a plain torch.save payload:
key
contents
checkpoint_format
ogbench.deterministic_mlp.export.v1
model_config
the DeterministicMlpDynamicsConfig fields
model
EMA weights (147,550 parameters for cube-single, 152,185 for cube-double)
The format string is deliberate: diffusion exports use
ogbench.heterogeneous_df.export.v2, and consumers dispatch on it, so a model
of a different family must not borrow that label.
Inference — read this before evaluating
The network outputs a state delta that is already in normalized space.
Add it to the normalized current state, then unnormalize. Getting this wrong
does not fail loudly; it silently produces errors 1,500x too large.
python
1import torch
23p = torch.load("cube-single/world_model_best.pt", map_location="cpu")4assert p["checkpoint_format"]=="ogbench.deterministic_mlp.export.v1"5w = p["model"]6s_off, s_scale = w["state_normalizer.offset"], w["state_normalizer.scale"]7a_off, a_scale = w["action_normalizer.offset"], w["action_normalizer.scale"]89net = torch.nn.Sequential(# layers 0/2/4/6 are Linear10 torch.nn.Linear(33,256), torch.nn.SiLU(),11 torch.nn.Linear(256,256), torch.nn.SiLU(),12 torch.nn.Linear(256,256), torch.nn.SiLU(),13 torch.nn.Linear(256,28),14)15net.load_state_dict({k[len("network."):]: v
16for k, v in w.items()if k.startswith("network.")})17net.eval()1819defstep(state, action):# both raw, batch-first20 x =(state - s_off)/ s_scale # normalize state21 u =(action - a_off)/ a_scale # normalize action22 delta = net(torch.cat([x, u], dim=-1))# ALREADY normalized-space23return(x + delta)* s_scale + s_off # add, then unnormalize
Three ways to get this wrong, with the error each produces on
cube-single-play-v0-val (all 100,000 in-episode transitions):
interpretation
normalized one-step MSE
x + delta, then unnormalize (correct)
0.000171
treat output as the absolute normalized next state
0.256
normalize the output again before adding
983
add output to the raw state (normalize(s + out))
0.0814
The second row is the common one: because the output head is zero-initialized
and deltas are small (mean absolute value 0.0272, matching the true mean
absolute normalized step of 0.0272), reading the output as an absolute state
predicts approximately the midpoint of the normalization range, giving an
error of about mean(x^2) = 0.256.
Equivalently, load the weights into DeterministicMlpDynamicsModel and call
rollout, which returns raw (unnormalized) future states and ignores the
diffusion noise arguments it accepts for interface parity.
Sanity check
Against the trivial "predict no change" baseline on the same 100,000
transitions, under the metric's own normalization:
normalized one-step MSE
predict no change
0.00545
this model
0.000171
The model is 31.9x better than predicting no change. If you measure
otherwise, the inference convention above is the first thing to check.
Note the 0.000171 here is over every validation transition, while the
0.000135 in the table above is what training logged from its sampled
validation batches (4 batches of 8 windows). The all-transitions figure is
the more reliable estimate.
where min/max are per-dimension min and max over the training split.
It is a half-range, not a standard deviation and not a delta scale. On
cube-single the mean per-dimension scale is 1.53 against a mean per-dimension
std of 0.595, so (scale/std)^2 averages about 52 — an NMSE quoted against
std is roughly 52x larger than one quoted against this normalizer.
Training
data
OGBench cube-single-play-v0.npz / cube-double-play-v0.npz and their -val.npz splits, read directly
observation construction
none — the state is the raw observations row, unmodified and unreordered (28 dims cube-single, 37 cube-double). No goal conditioning, no reordering, no derived features. actions and terminals are the only other fields read.
schedule
4 exhaustive epochs, every window seen exactly once per epoch
updates
56,316 (identical to the diffusion control)
global batch
64, one rank
window
context 1, prediction horizon 100, stride 1
optimizer
AdamW, lr 1e-4, weight decay 1e-5, grad clip 10.0, EMA 0.999
precision
bfloat16 autocast
hardware
one NVIDIA H200
state / action dims
cube-single 28/5, cube-double 37/5
Every element of the protocol other than the model is shared with the
diffusion control, and both see the same windows for the same number of
gradient updates, so the comparison isolates the model.
Training cost is the other half of the comparison: this model sustains
roughly 300 updates/second against the diffusion model's 0.33 seconds per
update — about two orders of magnitude cheaper per update, minutes rather
than hours for the full schedule.
Caveats
Single seed per task. The cube-double divergence is one run; treat its
magnitude as illustrative rather than precisely estimated.
Validation metrics are dataset-based rollout error, not task success. No
downstream policy-improvement results are included here.
