Vizuara/dreamzero-so101-lora

DreamZero-SO101 (LoRA, 70K steps)

A World Action Model (WAM) for the SO-101 robot arm, fine-tuned from DreamZero (Wan2.1-I2V-14B + joint action heads). Given a single camera observation and a natural-language task, it jointly predicts:

• —24 future 6-DOF joint actions (shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper)
• —33 future video frames showing the predicted task execution

Both modalities are denoised in a single forward pass using flow matching, so the model is internally consistent — the predicted actions and the predicted video describe the same imagined rollout.

[image]

2× H100 80GB · 72K steps · ~127 hours · rank-4 LoRA · joint flow matching Final action loss: 0.0015 (166× drop) · Final dynamics loss: 0.0298 (6× drop)

Why a "World Action Model"?

Most robot policies output actions and treat the world as a black box. A World Action Model also predicts what the world will look like during the rollout. This has three useful properties:

1. 1.Self-consistency — actions and predicted video share the same denoising trajectory, so the model has to imagine a coherent future
2. 2.Interpretability — you can literally watch what the policy "thinks" will happen before sending actions to the robot
3. 3.Sim-free evaluation — the predicted video gives you a free imagined rollout you can score offline

Quick Start

pip install huggingface_hub safetensors torch

# Download base + LoRA
huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P
huggingface-cli download Vizuara/dreamzero-so101-lora --local-dir ./checkpoints/dreamzero-so101-lora

# Clone DreamZero codebase + apply SO-101 patch
git clone https://github.com/dreamzero0/dreamzero.git
cd dreamzero && pip install -e .
git clone https://github.com/vizuara/dreamzero-so101.git
cd dreamzero && git apply ../dreamzero-so101/patches/so101_embodiment.patch

# Run inference
python ../dreamzero-so101/scripts/infer_demo.py \
  --model-path ./checkpoints/dreamzero-so101-lora \
  --base-model-path ./checkpoints/Wan2.1-I2V-14B-480P \
  --image ./sample_obs.jpg \
  --prompt "Pick up the red cube and place it in the bowl"

Model Details

LoRA Configuration

Training Recipe

Results

Final Loss

The dynamics loss converged around step 30K and remained flat. The action loss continued to slowly improve through 70K. Training was halted at step ~72K due to a pod migration; the loss curve indicates the model is well-converged and additional steps would have offered marginal returns.

Files

Intended Use

• —✅ Research & education — studying world models and joint video/action prediction for robotics
• —✅ SO-101 manipulation policy bootstrap — fine-tune further on your own data
• —✅ Offline rollout visualization — predict-before-execute to debug task setups
• —⚠️ Real-robot deployment — possible but requires safety wrappers and additional fine-tuning on your specific embodiment / camera setup
• —❌ Other arm types — trained only on SO-101; do not expect zero-shot transfer

Limitations

• —Trained only on whosricky/so101-megamix-v1 (400 episodes, 8 tasks). Out-of-distribution objects/scenes will degrade quality.
• —Joint trajectories are predicted, not torques — your low-level controller must accept position targets.
• —Video predictions are 33-frame snippets (~1 sec at 30 FPS). Longer horizons require chunked rollout.
• —LoRA only — for best quality, do a full fine-tune (see `dreamzero-so101` repo).
• —Trained at 320×176 — higher resolutions need re-training.

Citation

@misc{dreamzero-so101-2026,
  title  = {DreamZero-SO101: A World Action Model for the SO-101 Robot Arm},
  author = {Vizuara AI Labs},
  year   = {2026},
  howpublished = {\url{https://huggingface.co/Vizuara/dreamzero-so101-lora}},
  note   = {LoRA fine-tune of DreamZero on aggregated SO-101 LeRobot datasets}
}

Please also cite the underlying work:

• —DreamZero — Liu et al., "Learning World Action Models from Video", GEAR Lab, 2025
• —Wan2.1 — Wan-AI, "Wan2.1-I2V-14B", 2025
• —SO-101 — TheRobotStudio, "SO-100/SO-101 Open-Source Robot Arm"
• —LeRobot — HuggingFace, 2024

Acknowledgments

• —DreamZero by GEAR Lab — Apache 2.0 codebase
• —Wan2.1 — video generation backbone
• —LeRobot — dataset format and community
• —SO-101 dataset contributors on HuggingFace Hub

License

Apache 2.0 (same as DreamZero)