FurnitureBench H=100 Three-Task World Model
Public research artifact for the heterogeneous diffusion-forcing world model trained jointly on FurnitureBench one_leg, lamp, and round_table state trajectories.
Model
- 9,624,757 parameters
- Transformer width 208, 6 layers, 4 attention heads, FFN width 832
- Canonical state dimension 100; action dimension 10
- Context length 1 and prediction horizon 100 (101 state frames / 100 actions)
- 1,000 diffusion training steps and 20 inference steps
- Semantic-union state schema with task-specific object slots; absent slots are masked
- Rotation contract: proprio uses a 6D end-effector rotation, while furniture parts use quaternions
- Mask-aware, train-split-only affine normalization is stored in the exported model state
Training
The model completed four exhaustive passes over all valid H=100 training windows using four H200 GPUs, BF16, per-rank batch size 64 (global batch size 256), and 1,020 optimizer updates.
Source code: 9fd7d081
Training metrics: Weights & Biases
Final validation
| Metric | Mean |
|---|---|
| Objective | 0.125541 |
| One-step normalized MSE | 0.114518 |
| H=100 rollout normalized MSE | 0.200781 |
| One-step raw MSE | 0.041257 |
| H=100 rollout raw MSE | 0.078556 |
Files
world_model_best.pt: inference/export checkpoint selected by validation rolloutworld_model_final.pt: final inference/export checkpointtraining/checkpoint_latest.pt: full optimizer, EMA, RNG, and sampler state for exact resumetraining/metrics.jsonl: training and validation historytraining/manifest.json: architecture, protocol, metrics, and SHA-256 checksumsconfigs/: exact source configuration files
The export checkpoints are PyTorch dictionaries with model, model_config, schema, trajectory_metadata, environment_names, metrics, and step. Instantiate HeterogeneousDiffusionForcing from the source commit with model_config, then load model as its state dict. Use schema for canonical/native state conversion and slot masking.
import torch
from huggingface_hub import hf_hub_download
path = hf_hub_download(
repo_id="knightnemo/furniturebench-hdf-h100-three-task",
filename="world_model_best.pt",
)
payload = torch.load(path, map_location="cpu", weights_only=False)
print(payload["model_config"])
print(payload["schema"])
Dataset files are not included. This checkpoint is intended for research use with the matching FurnitureBench state/action and normalization contracts.