DVG-WM: Disentangled Video Generation Enables Efficient Embodied World Model for Robotic Manipulation

_{Vector version: assets/teaser.pdf}

Official implementation of DVG-WM, an efficient video-based embodied world model that disentangles dynamics learning from visual synthesis:

1. Preview stage — a LoRA-adapted CogVideoX-5B-I2V generates a low-resolution (256×384) dynamics preview from an initial observation and a language instruction (50 NFEs at low cost).
2. Refinement stage — a CogVideoX-2B finetuned with flow matching maps the upsampled preview latents directly to high-resolution (480×720) video latents in only 4 Euler steps. A latent degradation mechanism during training makes the refiner regenerate contact-rich details instead of merely upscaling.

This yields high-quality 49-frame video predictions with up to 3.97× faster inference than monolithic video world models — fast enough for iterative planning in robotic manipulation.

initial frame + instruction ──► [Stage 1: 5B-I2V + LoRA, 256×384, 50 steps] ──► LR dynamics z_lr
z_lr ──► upsample + latent degradation ──► [Stage 2: 2B flow matching, 480×720, 4 steps] ──► HR video

Installation

git clone https://github.com/zyshan0929/DVG-WM.git
cd DVG-WM

conda create -n dvg-wm python=3.10 -y
conda activate dvg-wm

# install torch matching your CUDA driver, e.g.
pip install torch==2.6.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt

The base models (zai-org/CogVideoX-5b-I2V, zai-org/CogVideoX-2b) are downloaded automatically from Hugging Face on first use.

Data Preparation (LIBERO)

Download the LIBERO demonstration datasets, then convert them into 49-frame training clips (no-op steps removed, tasks split 8:2 into train/test):

python tools/prepare_libero.py \
    --libero_dir /path/to/LIBERO/libero/datasets \
    --suites libero_spatial libero_object libero_goal libero_10 \
    --output_dir data/libero

This produces the standard CogVideoX finetune layout:

data/libero/{train,test}/
├── prompt.txt   # one language instruction per line
├── videos.txt   # one clip path per line
└── videos/      # 49-frame 480×720 mp4 clips

Any dataset in this format works — to use your own robot data, just export clips + prompts the same way.

Training

Stage 1: preview (dynamics) — 1 GPU

LoRA (rank 128 on attention/FFN/AdaLN) finetuning of CogVideoX-5B-I2V at 256×384. The learned positional embedding is discarded; the model runs on flexible 3D RoPE.

bash scripts/train_stage1_preview.sh
# ≈ paper setting: batch 4, 10k iterations on a single 80G GPU

Stage 2: refinement (visual synthesis) — 8 GPUs

Flow-matching finetuning of CogVideoX-2B at 480×720 with latent degradation (z̃ = α_s·E(Deg_pix(x)) + β_s·ε, degradation index sampled in [650, 750]):

bash scripts/train_stage2_refine.sh
# ≈ paper setting: batch 1 × 8 GPUs, 10 epochs

Useful overrides (both scripts forward extra args): --max_train_steps, --batch_size, --training_type lora (cheap stage-2 variant), NUM_GPUS=4 bash scripts/train_stage2_refine.sh.

Inference

Generate a high-resolution prediction from one image + instruction:

python tools/predict.py \
    --preview_lora checkpoints/stage1_preview/checkpoint-10000 \
    --refine_ckpt checkpoints/stage2_refine/checkpoint-XXXX/transformer \
    --image assets/example_frame.png \
    --prompt "pick up the black bowl and place it on the plate" \
    --output_dir outputs/demo --save_preview

Or run the whole test split (frame 0 of each ground-truth clip is the initial observation):

bash scripts/inference.sh          # writes outputs/pred/*.mp4

Key knobs: --preview_steps 50 --refine_steps 4 --deg_index 675 --refine_shift 2.5.

Evaluation

PSNR / SSIM / LPIPS / FVD against the ground-truth test clips:

bash scripts/evaluate.sh           # or: python tools/evaluate.py --pred_dir outputs/pred --gt_root data/libero/test

Results on LIBERO (from the paper)

Repository Structure

DVG-WM/
├── dvg_wm/
│   ├── data/video_dataset.py      # video-text dataset (CogVideoX finetune format)
│   ├── models/preview.py          # RoPE-only patch + LoRA targets for stage 1
│   ├── pipeline_dvg_wm.py         # end-to-end preview → refine pipeline
│   └── utils.py                   # latent degradation, upsampling, video I/O
├── train/
│   ├── train_stage1_preview.py    # stage-1 LoRA trainer
│   └── train_stage2_refine.py     # stage-2 flow-matching trainer
├── tools/
│   ├── prepare_libero.py          # LIBERO → training clips
│   ├── predict.py                 # inference driver
│   └── evaluate.py                # PSNR / SSIM / LPIPS / FVD
└── scripts/                       # one-line launchers

Acknowledgements

This codebase builds on CogVideo / CogVideoX and follows the two-stage cascade design of FlashVideo. Simulation data comes from LIBERO. We thank the authors for their open-source contributions.

Citation

@article{shan2026dvgwm,
  title   = {DVG-WM: Disentangled Video Generation Enables Efficient Embodied World Model for Robotic Manipulation},
  author  = {Shan, Ziyu and Wu, Zhenyu and Wang, Xiaofeng and Zhu, Zheng and Wang, Ziwei},
  journal = {arXiv preprint arXiv:2606.32028},
  year    = {2026}
}

License

This repository is released under the Apache 2.0 license. The CogVideoX base models are subject to their own model license.