dinowm-pusht

A DINO-WM: frozen DINOv2 ViT-S/14 patch tokens (256 per 224 px frame, loaded from torch hub and not stored here) with 10-dim proprioception and action embeddings on every patch, and a 6-layer frame-causal ViT predictor. Trained on the 18,685 episodes of le-wm's Push-T expert dataset (stable-worldmodel's pusht_expert_train.h5), replayed in torch_world's MuJoCo PushT from their first states with their recorded actions, at 224 px, with windows of 4 frames 5 steps apart.

• Model: torch_world.models.DINOWM of torch-world: DINOWM(PatchTokens.from_hub('dinov2_vits14', size=None), num_patches=256, embed_dim=384, action_dim=10, proprio_dim=4, history=3, depth=6, num_heads=16, head_dim=64, mlp_dim=2048).
• Inputs: frames in [0, 1] (the encoder applies the ImageNet normalization); proprioception = pusher position and velocity in the reference's pixel units, z-scored.
• Actions: blocks of 5 Push-T actions (10 values), z-scored with normalization.action of config.json, in the reference's pixel axes (y down); torch_world's PushT takes y up, so flip y.
• Recipe: AdamW (lr 5e-4, no weight decay), constant learning rate, batch 32, bf16, gradient clipping at 1, 150,000 steps (2.4 epochs).
• Paper: Zhou et al., DINO-WM, ICML 2025 (https://arxiv.org/abs/2411.04983).
• Intermediate checkpoints: history/ (step_015000, step_030000, step_045000, step_060000, step_075000, step_090000, step_105000, step_120000, step_135000).

Usage

from torch_world import hub

model = hub.from_pretrained("latenthorizon/dinowm-pusht")  # pip install "torch-world[hub]"

config.json holds the model's spec (what from_pretrained builds) and the normalization of actions and proprioception.

Results

Push-T goal reaching under stable-worldmodel's protocol: goal 25 expert steps ahead, CEM (300 samples, 30 elites, 30 iterations) over 5 blocks of 5 actions, whole plan executed before replanning, 50 steps, success when the goal is reached at any step.

evaluation	env	episodes	success
torch_world CEM (planned per task, as the reference)	torch_world PushT (MuJoCo)	50	94.0%
same, history/step_075000	torch_world PushT (MuJoCo)	50	100.0%