nyu-wm-v3

NYU DS-GA-3001 — World Models (v3, velocity actions)

Trained world-model checkpoints for the lquan9/collect1 drone trajectory dataset. Eight configurations across two encoders (CNN / DINOv2-ViT-B/14) × two dynamics (MLP / RSSM) × two training splits (filtered _c1 / unfiltered _c1u).

All checkpoints use velocity actions (6-dim: vx, vy, vz, wx, wy, wz) and odom with per-episode position deltas from the start frame.

Trajectory Evaluation (ADE / FDE, meters)

ADE = Average Displacement Error over the rollout horizon. FDE = Final Displacement Error at the last timestep. Both are reported in meters on the held-out validation split.

Files per experiment

Loading

import torch
from huggingface_hub import hf_hub_download
from world_model.architecture import WorldModel
from world_model.data.normalization import NormStats

exp = "dino_rssm_c1"
ckpt_path = hf_hub_download(repo_id="dmody1/nyu-wm-v3", filename=f"{exp}/latest.pt")
ns_path   = hf_hub_download(repo_id="dmody1/nyu-wm-v3", filename=f"{exp}/norm_stats.pt")

ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
ns   = NormStats.load(ns_path)
mcfg, targs = ckpt["model_config"], ckpt["args"]

model = WorldModel(
    odom_dim=mcfg["odom_dim"], imu_dim=mcfg["imu_dim"], act_dim=mcfg["act_dim"],
    state_dim=mcfg["state_dim"],
    backbone=targs.get("backbone", "dinov2"),
    dynamics_type=targs.get("dynamics_type", "rssm"),
)
model.load_state_dict(ckpt["model"])
model.eval()

See cell 35 of docs/notebooks/world_model_pipeline.ipynb for MPC inference (CEM planning with Δ-state cost).

Dataset

Training data: lquan9/collect1. Train/val split: episodes with eid % 5 == 0 are validation, the rest are training.