alakazam-robonet/so101-fd-worldmodel

SO-101 Forward-Dynamics World Model (Cosmos 3 Edge fine-tune)

An action-conditioned video world model for the SO-101 arm: give it a camera frame and a sequence of commanded joint angles, and it predicts the video of the arm executing them. Fine-tuned from nvidia/Cosmos3-Edge (4B, causal, streaming) on 254 episodes of a 3D-printed SO-101 doing pick-and-place.

This is the model behind alakazam.gg/so-wam.

What it does

Input: one anchor frame (480 px) plus commanded joint targets, 16 per chunk at 15 fps. Output: the predicted video of that motion. It is a world model, not a policy — it answers "what happens if I command this", not "what should I command".

How well it does it

The test: hand the model an anchor frame and one second of real commanded joint angles, let it imagine the result, and measure where the end effector lands versus the real arm in a held-out episode (25 episodes, never trained on, same task and same robot).

model	end-effector error after 1 s
this fine-tune	5.8 px
a LoRA fine-tune of the same base, same data	14.0 px
stock Cosmos 3 Edge	132.3 px

It wins all 25 held-out episodes. The honest read of that table is the middle row, not the bottom: most of the gap is buying the fine-tune at all, not this particular recipe. The cheaper LoRA sibling gets within 8 px of it. Stock Edge renders clean video the whole way and sends the arm somewhere else entirely.

Error is measured in pixels of the 640×480 frame, against the real arm's end effector in the same episode.

What it is not good at, specifically

These are measured failures on this model, not generic caveats.

It is bound to one camera pose. Every training episode used a single fixed overhead view. Moving the camera degrades it sharply, and the failure is not graceful: a sibling policy in the same project stopped moving entirely when a prop changed colour from red to yellow. Reproduce the training viewpoint and scene colours before expecting behaviour.

A second camera did not help at equal budget. Stacking the wrist view above the overhead view, at matched token budget and matched 7,600 steps, scored 19.15 px against this model's 17.58 px on the same end-effector metric — about 9% worse, with the wrist half rendering as visible smear. Per-unit-of-motion pixel sensitivity to the wrist joints is 5.7× higher in the wrist view (0.810 vs 0.141), so the view carries real signal; the concat layout at a fixed canvas is what failed.

What is in this repo

Two copies of the same trained model in different formats. Picking the wrong one costs an afternoon, so:

path	size	what it is for
export_ft7600/	7.7 GB	Serving. The export format, with its own config.json and the anchor frames in anchors/. This is what the demo and the Reactor pipeline load.
so101_v7_full_posttrain/checkpoints/iter_000007600/	20.2 GB	Training. The raw PyTorch DCP checkpoint, for resuming or further post-training.
known_answer/	small	A fixed input and the output this checkpoint produces for it, for verifying a deployment.
integration/, INTEGRATION.md, Dockerfile	small	The framework files and the environment recipe.

If you only want it running, there is a prebuilt image with the weights already inside and nothing to mount:

docker run --gpus all --network host \
  us-central1-docker.pkg.dev/ewilan-pipeline/so101/so101-fd-reactor:2.2

It serves the model through Reactor Runtime and hands out a browser page on 8899. The image is built from public sources only: the framework at a pinned commit, its environment from that repo's own uv.lock, and these weights. The Dockerfile that builds it is in the repo below if you would rather build it yourself. Source and the longer paths: Alakazam-studios/so101-dream-drive.

Using it

The checkpoint is a PyTorch DCP directory, not a single file. Two things bite immediately:

• The framework infers the experiment name from the checkpoint path. Keep the .../so101_v7_full_posttrain/checkpoints/iter_000007600/ structure; a bare iter_* directory raises Could not infer experiment from checkpoint path.
• The Wan 2.2 VAE is required alongside it, and the framework resolves it by name from a HuggingFace cache at load time. Offline or without that cache, the resolve fails inside a startup hook and reads as a hang rather than an error.

setup.checkpoint_path = ".../checkpoints/iter_000007600"
setup.parallelism_preset = "latency"
setup.guardrails = False           # OpenMDW imposes no guardrail obligation
pipe = setup.build_setup().get_inference_cls().create(setup_args)

Contract: embodiment so101 (domain id 10), 6 joints in order pan, lift, elbow, wrist_flex, wrist_roll, gripper, absolute degrees, 16-action chunks at 15 fps, 480 px. Working directory must be the framework root, since several configs are opened by relative path.

Resolution is not a free parameter. Post-training ran at 480, Edge's native inference resolution, and the sampler's flow shift is keyed off it (shift={"256": 3, "480": 5, "720": 10}), so serving at another size changes the sampler as well and the output degrades.

Measured serving cost on one H100: 12 s model load, ~2 s per rollout warm. On one RTX PRO 6000: 7.6 s load and ~1.2 s per 16-frame chunk, which is roughly real time.

Training

Fine-tuned from nvidia/Cosmos3-Edge, full-rank, 7,600 steps at 480 px on a single GPU. Data: 174 teleoperated SO-101 pick-and-place episodes at 640×480 / 30 fps, plus 80 augmented episodes (environment restyle and distractor insertion) generated by our own pipeline and filtered so no augmentation derives from a held-out source episode. Held-out split is episode_index % 7 == 3.

License

Inherits OpenMDW-1.1 from nvidia/Cosmos3-Edge. Verified against the raw license text: no guardrail obligation and no attribution requirement on outputs. Note this differs from NVIDIA's Open Model License used by Cosmos Transfer/Predict, which carries a self-executing termination clause — do not assume terms transfer between Cosmos SKUs.

Citation

Technical report and evaluation detail: alakazam.gg/so-wam.