Geometry-grounded video world models for robot manipulation

GEM-4D: Geometry-Enhanced Video World Models
for Robot Manipulation

One mid-layer feature. One frozen geometry teacher. Zero extra cost at inference.

A Geometry DiT that sees the scene only through the video DiT's features turns 4D geometry
into a training signal for the video backbone, then is thrown away.

Kaichen Zhou^1,2,*   Yuzhen Chen^1,*   Fangneng Zhan²   Hang Hua⁴   Grace Chen¹   Xinhai Chang²
Ao Qu²   Yilun Du¹   Zhuang Liu³   Paul Pu Liang^2,†   Mengyu Wang^1,†

¹Harvard University    ²MIT    ³Princeton University    ⁴MIT-IBM Watson AI Lab

^*Equal contribution    ^†Joint supervision

   

Abstract · Method · Results · Paper ↔ code · Quick start · Citation

Abstract

Video world models can generate realistic futures from a single instruction, but they often fail to track the same physical points consistently across time. The generated videos look plausible, yet lack the physical grounding needed for reliable action execution such as robot manipulation.

GEM-4D injects dense 4D correspondence supervision, distilled from a pretrained geometry foundation model, into the video generative backbone during training. The world model learns appearance and geometric structure jointly while keeping a single-stream architecture with no additional inference cost. GEM-4D achieves state-of-the-art video prediction and geometric consistency in simulated and real-world scenarios, and improves real-world manipulation success from 61% to 81%.

Keywords: video world models, 4D generation, geometry foundation models, robot manipulation.

Method

One shared feature, two flows.

A video DiT predicts the velocity of the noised video latent. The output m_t of one of its middle blocks is handed to a small Geometry DiT, which predicts the velocity of a masked, noised geometry latent built from a frozen geometry foundation model (Depth Anything 3). The Geometry DiT has no access to pixels, depth or cameras: everything it knows about the scene arrives through m_t.

first frame + instruction ──► video DiT ──────────────────────────► video velocity ──► L_vid
                                │ m_t  (block `hook_layer`, not detached)
                                ▼
 masked, noised geometry latent ──► Geometry DiT ──► geometry velocity ──► L_geo
 g0 = projector(DA3 features)   (same noise as the video latent)

                        L = L_vid + α · L_geo

Training: distil geometry into the video backbone

L_geo can only be minimised if m_t encodes depth, camera motion and scene flow, the factors that determine inter-frame correspondence. Its gradient α · ∂L_geo/∂m_t · ∂m_t/∂θ flows back into the video backbone and pushes that structure into its representation. The geometry teacher stays frozen.

Inference: drop the branch

The Geometry DiT only reads from m_t and never writes back, so it is removed after training. The deployed model is a plain image-to-video generator, the same size and speed as its backbone.

Results

4D scene generation (paper Tab. 1). R = real (DROID), S = simulated (RLBench).

D	Method	FVD ↓	SSIM ↑	PSNR ↑	AbsRel ↓	δ1 ↑	Chamfer ↓	Tracking ↑
R	CogVideoX	35.56	75.91	20.18	22.33	68.32	0.2670	66.22
R	Wan 2.2-14B	33.43	76.24	20.70	21.39	71.18	0.2349	68.18
R	TesserAct	33.28	75.66	20.08	22.07	66.80	0.2630	67.14
R	GEM-4D	31.82	82.05	21.11	20.13	78.19	0.2001	71.23
S	CogVideoX	40.21	75.51	20.03	15.41	70.99	0.2913	58.32
S	Wan 2.2-14B	49.20	73.01	19.87	17.81	67.07	0.1762	61.99
S	TesserAct	41.97	76.72	19.71	16.02	69.26	0.1813	61.15
S	GEM-4D	27.94	80.27	23.36	14.11	74.13	0.0702	68.18

More videos and real-robot rollouts: gem-4d.github.io.

Paper ↔ code

Paper (Sec. 3.2)	Code
Video DiT E_θ, U_θ; mid-level feature m_t	model/gem4d/models/{cogvideox,wan}.py; forward hook on block geometry.hook_layer (21/42 CogVideoX, 15/30 Wan), not detached (models/gem4d.py)
Frozen geometry model G, representation g_0	geometry/teacher.py: DA3 ViT-L layer 12 → frozen projector to the latent grid
Mask mechanism M on the noised geometry	heads/geo_dit.py: geometry.mask_ratio of tokens replaced by a learned mask token
v_ψ^geo(M(g_t), t, m_t), m_t the only scene input	GeoDiT: cross-attention to m_t only (geometry.use_eps_hat=false)
L = L_vid + α L_geo	models/gem4d.py + training/trainer.py (geometry.alpha, linear ramp geometry.warmup)
Geometry branch dropped at inference	inference/infer.py loads only the exported transformer

Not in this repository: the adaptive inverse-dynamics module (video → robot actions) and the VGGT / depth-target ablations. In the paper Wan 2.2-14B is a fine-tuned baseline; GEM-4D on Wan2.2-TI2V-5B is provided here as an additional backbone without released weights.

Quick start

The same code trains GEM-4D on two backbones:

Backbone	Config	Latent	Weights
CogVideoX-5b-I2V (42 blocks)	training/config/cogvideox_*.yaml	16 × F/4 × H/8 × W/8	paper model, released checkpoints
Wan2.2-TI2V-5B (30 blocks)	training/config/wan22_ti2v_*.yaml	48 × F/4 × H/16 × W/16	train your own

Code layout

model/gem4d/
  models/backbone.py        backbone interface (encode, noise, loss, forward)
  models/cogvideox.py       CogVideoX-5b-I2V (DDPM v-prediction)
  models/wan.py             Wan2.2-TI2V-5B   (rectified flow, first-frame conditioning)
  models/gem4d.py           GEM-4D module: video DiT + Geometry DiT + frozen projector
  heads/geo_dit.py          Geometry DiT (masked geometry flow)
  geometry/teacher.py       frozen DepthAnything-3 teacher + geometry projector
  utils/video_io.py
model/dependency/depth_anything_3/   vendored DepthAnything-3
training/
  config/                   default.yaml + one file per backbone x {gem4d, rgb}; accelerate_zero2.yaml
  data/manifest.py          manifest format (shared by training, validation, inference)
  data/video_dataset.py     generic dataset — the place to customise data loading
  train_utils/              config loader, checkpointing
  trainer.py                single training loop for every backbone
  launch.py                 entry point
training_bash/train.sh      multi-GPU launcher (accelerate + DeepSpeed ZeRO-2)
training_bash/infer.sh      inference launcher
inference/infer.py          checkpoint -> videos
tools/download_weights.py
examples/                   example manifests
tests/test_smoke.py         CPU tests with tiny random models (`PYTHONPATH=model:. python tests/test_smoke.py`)

Environment

conda create -n gem4d python=3.10 -y && conda activate gem4d
pip install -r requirements.txt          # diffusers>=0.35 (Wan2.2 TI2V support)

Weights

python tools/download_weights.py cogvideox da3            # CogVideoX GEM-4D training from the base model
python tools/download_weights.py cogvideox checkpoint-2   # CogVideoX inference
python tools/download_weights.py wan2.2 da3               # Wan2.2-TI2V-5B GEM-4D training

Name	Source	Use as
cogvideox/	zhouk777/gem4d-weight (= CogVideoX-5b-I2V)	backbone.base_model
da3/	same repo (= DA3-LARGE-1.1)	geometry.teacher
checkpoint-1/	same repo — GEM-4D CogVideoX, step 21000	backbone.init_weights (fine-tuning)
checkpoint-2/	same repo — GEM-4D CogVideoX, step 42000	inference --weights
Wan2.2-TI2V-5B-Diffusers/	Wan-AI/Wan2.2-TI2V-5B-Diffusers	backbone.base_model

Data

There are no predefined datasets. List your clips in a manifest (JSON or JSONL):

[
  {"video": "rlbench/close_drawer/ep_0001/front.mp4", "prompt": "close the middle drawer", "robot": "Franka robot arm"},
  {"video": "my_lab/frames_0042", "prompt": "push the red block to the left"},
  ["pick orange from white bowl", "rt1/record_00162/49/rgb.mp4", "google robot"]
]

Field	Meaning
video	.mp4 / .webm, or a folder of frames 0.png, 1.png, … (an image is also accepted for validation / inference). Relative paths resolve against data.data_root, else the manifest's folder.
prompt	the instruction
robot	optional embodiment text appended to the prompt; keep it identical at train and test time

[prompt, video] and [prompt, video, robot] lists are accepted too; a directory of manifests is concatenated. Each clip is center-cropped and resized to data.height × data.width (480 × 640) and brought to data.num_frames (49) frames: long clips are subsampled uniformly, short clips are padded with their last frame (prefer ≥ 49-frame clips; data.min_frames drops short ones). Frame 0 is the condition. No depth files are needed: geometry comes from DA3 on the fly.

To customise loading, subclass training.data.video_dataset.ManifestVideoDataset and override read_frames(entry), resize(frames, entry) or get_prompt(entry).

Training

# GEM-4D on CogVideoX, from the base model, 8 GPUs
bash training_bash/train.sh training/config/cogvideox_gem4d.yaml \
    data.train_manifest=/data/train.json data.val_manifest=/data/val.json

# GEM-4D on Wan2.2-TI2V-5B
bash training_bash/train.sh training/config/wan22_ti2v_gem4d.yaml \
    data.train_manifest=/data/train.json data.val_manifest=/data/val.json

# continue from a saved checkpoint (released or your own export)
bash training_bash/train.sh training/config/cogvideox_gem4d.yaml \
    data.train_manifest=/data/train.json backbone.init_weights=weights/checkpoint-1

# RGB-only fine-tuning baseline (no geometry branch)
bash training_bash/train.sh training/config/cogvideox_rgb.yaml data.train_manifest=/data/train.json

Any key in training/config/default.yaml can be overridden as key.sub=value. Main settings:

Key	Default	Meaning (paper)
geometry.alpha	0.03	weight α of L_geo
geometry.mask_ratio	0.5	fraction of geometry tokens masked (mask mechanism M)
geometry.hook_layer	21 (CogVideoX) / 15 (Wan)	block whose output is m_t
geometry.feat_layer	12	DA3 layer used as the geometry representation g0
geometry.depth	4	Geometry DiT blocks
geometry.warmup	100	linear ramp of α
optim.lr	1e-5	constant with 100 warm-up steps, AdamW (0.9, 0.95), wd 1e-3
train.max_steps	10000	batch 1 per GPU, 49 × 480 × 640, bf16, gradient checkpointing

Starting point. backbone.init_weights accepts a released checkpoint or any exported checkpoints/step_N (its transformer/ is used). The Geometry DiT starts fresh unless geometry.init_geometry=<step_N>/geometry.safetensors is given; a released CogVideoX checkpoint's geometry projector is reused automatically so the geometry space stays the same.

Outputs (experiment.output_dir): checkpoints/step_N/transformer/ (diffusers format, for inference / fine-tuning), checkpoints/step_N/geometry.safetensors (Geometry DiT + projector), checkpoints/state/ (full resumable state of the newest save), vis/step_N/*.mp4 (validation, left GT / right prediction), logs/ (TensorBoard: video_loss, geo_loss, val/*), config.yaml. Exports keep the newest checkpoint.keep_last plus every multiple of checkpoint.keep_every.

Resume. Re-run the same command with the same experiment.output_dir.

Inference

bash training_bash/infer.sh --backbone cogvideox --base_model weights/cogvideox \
    --weights weights/checkpoint-2 --manifest examples/val_manifest.json \
    --output_dir outputs/cogvideox --side_by_side

bash training_bash/infer.sh --backbone wan2.2-ti2v-5b --base_model weights/Wan2.2-TI2V-5B-Diffusers \
    --weights runs/wan22_ti2v_gem4d/checkpoints/step_010000 --manifest examples/val_manifest.json \
    --output_dir outputs/wan

The first frame of each manifest video (or the image) conditions the rollout; the prompt is prompt + robot. Output: <idx>_<clip>.mp4 (49 frames, 480 × 640) and with --side_by_side <idx>_<clip>_gt_pred.mp4. Options: --steps 50, --guidance_scale (6.0 CogVideoX / 5.0 Wan), --num_samples N, --cpu_offload. Exported transformers are plain diffusers models, so they also work directly with CogVideoXImageToVideoPipeline / WanImageToVideoPipeline(expand_timesteps=True).

Implementation notes

• Video objective. CogVideoX keeps its DDPM v-prediction loss (SNR-weighted x0 error); Wan uses rectified flow with shifted timesteps (shift 5) and TI2V first-frame conditioning (clean image latent at frame 0 with timestep 0, excluded from the loss).
• Geometry latent. DA3 ViT-L layer-12 features [F, 24, 24, 1024] → frozen 1×1 conv to 256 → frozen temporal average (stride 4) → resize to the latent grid → frozen 1×1 conv to the latent channels (16 / 48). Geometry and video share one noise sample; the geometry target is the backbone's own velocity parametrisation.
• Reproducing the released CogVideoX checkpoints. They were trained with an extra Geometry DiT input, a stop-grad RGB noise estimate (geometry.use_eps_hat=true). It is off by default, as in the paper; inference is unaffected.

Citation

@inproceedings{zhou2026gem,
  title={Gem-4d: Geometry-enhanced video world models for robot manipulation},
  author={Zhou, Kaichen and Chen, Yuzhen and Zhan, Fangneng and Hua, Hang and Chen, Grace and Chang, Xinhai and Qu, Ao and Du, Yilun and Liu, Zhuang and Liang, Paul Pu and others},
  booktitle={European Conference on Computer Vision},
  pages={130--147},
  year={2026},
  organization={Springer}
}

Acknowledgements

Built on diffusers, CogVideoX, Wan2.2, Depth Anything 3 and TesserAct.