vjepa2-hwm

vjepa2-hwm — Hierarchical high-level world model for V-JEPA 2-AC

An independent reimplementation of the high-level world model from Hierarchical Planning with Latent World Models (arXiv:2604.03208), built on top of the official V-JEPA 2-AC stack and trained on the public DROID corpus. It provides the two modules the paper's method adds to a frozen latent world model:

• p2.* — the high-level predictor P²(z_{t+h} | z_t, l_t): predicts the latent of a waypoint 0.25–4 s ahead, conditioned on a learned macro-action. Architecture is the exact vit_ac_predictor from the official facebookresearch/vjepa2 repository (24 layers, 1024 dim, ~300 M params) — no architectural changes.
• aenc.* — the macro-action encoder A_ψ (~1.6 M params): compresses a variable-length primitive-action segment into a 4-dim latent macro-action l ∈ [-1, 1]⁴ (transformer encoder, CLS readout — per the paper §3.1).

Together with the frozen V-JEPA 2 ViT-g encoder and the official AC predictor as the low-level model, this enables the paper's two-level planning: CEM over macro-actions in the high level proposes subgoal latents; the unchanged low-level planner pursues them.

Evaluation (held-out DROID episodes, never trained on)

• Prediction: high-level single-step beats the low-level model rolled out autoregressively at all 16 gap lengths (0.25–4 s); at 3.5 s: L1 0.382 vs 0.498, cosine +0.769 vs +0.669 — the paper's error-accumulation argument reproduced. Chaining two macro-steps costs only ~0.003 cosine.
• Conditioning: real-macro-action vs permuted-macro-action margin +0.025 cosine; a linear probe decodes segment xyz displacement from l at R² ≈ 0.999.
• Planning (offline, first-action vs logged action at t=0, 3 s goal): hierarchical +0.116 direction-cosine vs flat +0.079 — the hierarchy overtakes flat planning in the long-horizon regime where flat approaches chance. (Short-horizon planning remains the flat planner's domain.)
• Closed-loop (honest limitation): on a robosuite Lift reach task the hierarchy did not improve task success over flat planning (0/20 for both), despite consistently steering the arm closer to the goal (mean closest approach 9.9 cm from 12.9 cm starts). An oracle test — substituting ground-truth encoder latents as subgoals — performs identically to the model's subgoals (+0.124 vs +0.116 offline), so the subgoal content is at ceiling; the residual bottleneck is the low-level planner's flat energy landscape at episode start (its CEM outputs near-zero-magnitude actions there). If your low-level planner handles start states well, these subgoals should transfer their full offline advantage.

Differences from the paper (read before comparing numbers)

1. P² is warm-started from the official V-JEPA 2-AC predictor weights (the same checkpoint the low-level planner uses). The paper appears to train its high-level model from scratch ("a scaled-up variant of the low-level world model"). Our from-scratch control arm converges to similar conditioning margins but several× slower.
2. The macro-action rides through the predictor's existing 7-dim action input, zero-padded (dims 0–3 = l, 4–6 = 0), and the states input receives the start pose broadcast to every slot (future waypoint poses don't exist at plan time). The paper does not specify its conditioning pathway in this detail.
3. l is tanh-bounded to [-1, 1] and the planning-time CEM samples with init std 0.25 (matched to the trained l std ≈ 0.19). The paper does not document bounds or sampling ranges. Two anti-saturation details matter: LayerNorm before the tanh head and a 0.1× head-weight init — without them the head rails at ±1 and the macro-action collapses irrecoverably.
4. An auxiliary loss (linear decode of segment xyz displacement from l, weight 1.0) was active during this run. The paper uses none. Our ablation indicates it is droppable once the anti-saturation fixes are in place — the training-only aux head is stripped from these weights.
5. Training is teacher-forcing-only (rollout-loss weight 0), matching the paper's γ_roll = 0.
6. Dataset: 21,247 DROID success episodes (public GCS release, left exterior camera only, transcoded to 4 fps / 256 px), waypoints sampled with i.i.d. uniform gaps of 1–16 frames (0.25–4 s), K=3 segments per sample. The first 240 episodes of the deterministic listing order are held out entirely for evaluation. The paper's Franka experiments use their own robot data; ours are DROID-only.
7. Optimization: batch 256, 6,000 iters, AdamW, peak LR 1.5e-4 (warm-start finetune schedule; 5% warmup, 25% cosine anneal), wd 0.04, bf16, A_ψ at 3× the schedule LR. Checkpoint selection: prefer the final checkpoint over best-val-cosine — validation cosine peaks early while the conditioning margins (which drive planning quality) keep improving; the early "best-val" checkpoint plans markedly worse.

Usage

Weights are a flat safetensors file with aenc.* and p2.* key prefixes. P² loads directly into the official repo's predictor:

import torch
from safetensors.torch import load_file

# pip install git+https://github.com/facebookresearch/vjepa2
from src.models.ac_predictor import vit_ac_predictor

sd = load_file("hwm_b200_6k.safetensors")

p2 = vit_ac_predictor(img_size=(256, 256), patch_size=16, num_frames=64,
                      tubelet_size=2, embed_dim=1408)
p2.load_state_dict({k[3:]: v for k, v in sd.items() if k.startswith("p2.")})

A_ψ is small enough to define inline (this is the exact architecture the aenc.* keys expect):

import torch.nn as nn

MACRO_DIM, GAP_MAX, W = 4, 16, 256

class ActionEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(7, W)
        self.pos = nn.Parameter(torch.zeros(1, GAP_MAX + 1, W))
        self.cls = nn.Parameter(torch.zeros(1, 1, W))
        layer = nn.TransformerEncoderLayer(W, nhead=4, dim_feedforward=4 * W,
                                           dropout=0.0, batch_first=True,
                                           norm_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.norm = nn.LayerNorm(W)
        self.head = nn.Linear(W, MACRO_DIM)

    def forward(self, actions, lengths):
        # actions: [B, GAP_MAX, 7] zero-padded 7-DoF deltas at 4 Hz
        #          (xyz deltas scaled by 1/0.075 to O(1)); lengths: [B]
        B, G, _ = actions.shape
        actions = torch.cat([actions[..., :3] / 0.075, actions[..., 3:]], -1)
        x = torch.cat([self.cls.expand(B, -1, -1), self.proj(actions)], 1)
        x = x + self.pos[:, : G + 1]
        pad = torch.arange(G, device=actions.device)[None] >= lengths[:, None]
        pad = torch.cat([pad.new_zeros(B, 1), pad], 1)
        h = self.encoder(x, src_key_padding_mask=pad)
        return torch.tanh(self.head(self.norm(h[:, 0])))

aenc = ActionEncoder()
aenc.load_state_dict({k[5:]: v for k, v in sd.items() if k.startswith("aenc.")})

Conditioning contract for P² (one high-level step): with z_t the layer-normed frozen-encoder latent [B, 256, 1408] of the current frame, l a macro-action [B, MACRO_DIM], and pose the current 7-DoF end-effector state [B, 1, 7]:

import torch.nn.functional as F

cond = F.pad(l, (0, 7 - MACRO_DIM))          # l rides the 7-dim action input
z_next = p2(z_t, cond[:, None], pose)[:, -256:]
z_next = F.layer_norm(z_next, (1408,))       # match the encoder's convention

For multi-step rollout, append each prediction to the latent context and repeat with the pose input broadcast from the start pose (the model was trained with start-pose-only conditioning). For planning, run any sampling optimizer (e.g. CEM, ~400 samples, 6 iterations, Gaussian init σ≈0.25 clipped to [-1,1]) over macro-action sequences against ‖z_goal − ẑ‖₁, and hand the best rollout's intermediate latents to your low-level planner as subgoals.

The frozen encoder and the low-level AC predictor are the official release: load them per the facebookresearch/vjepa2 instructions (vjepa2_ac_vit_giant).

Related

• Paper: Hierarchical Planning with Latent World Models
• Base stack: facebookresearch/vjepa2
• Sibling release: dobri420/vjepa2-cac — a from-scratch reproduction of the low-level AC predictor
• Data: DROID