knightnemo/robocl-ogbench-mlp-world-model

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)
schema	canonical state schema
environment_names, trajectory_metadata, metrics, step	provenance

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.

import torch

p = torch.load("cube-single/world_model_best.pt", map_location="cpu")
assert p["checkpoint_format"] == "ogbench.deterministic_mlp.export.v1"
w = p["model"]
s_off, s_scale = w["state_normalizer.offset"], w["state_normalizer.scale"]
a_off, a_scale = w["action_normalizer.offset"], w["action_normalizer.scale"]

net = torch.nn.Sequential(                      # layers 0/2/4/6 are Linear
    torch.nn.Linear(33, 256), torch.nn.SiLU(),
    torch.nn.Linear(256, 256), torch.nn.SiLU(),
    torch.nn.Linear(256, 256), torch.nn.SiLU(),
    torch.nn.Linear(256, 28),
)
net.load_state_dict({k[len("network."):]: v
                     for k, v in w.items() if k.startswith("network.")})
net.eval()

def step(state, action):                        # both raw, batch-first
    x = (state - s_off) / s_scale               # normalize state
    u = (action - a_off) / a_scale              # normalize action
    delta = net(torch.cat([x, u], dim=-1))      # ALREADY normalized-space
    return (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.

What the normalization is

normalize(v) = (v - offset) / scale, with

offset = (min + max) / 2
scale  = max(max - min, 1e-2) / 2

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.
• Trained only on cube-single and cube-double.