knightnemo/furniturebench-hdf-h100-three-task

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 rollout
• world_model_final.pt: final inference/export checkpoint
• training/checkpoint_latest.pt: full optimizer, EMA, RNG, and sampler state for exact resume
• training/metrics.jsonl: training and validation history
• training/manifest.json: architecture, protocol, metrics, and SHA-256 checksums
• configs/: 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.