YunzeLiu/robodojo_dreamzero

RoboDojo DreamZero

Status: Stage 2 training is still in progress. The repository always retains only the newest durable Stage 2 checkpoint. Resolve latest_checkpoint.json instead of hard-coding a checkpoint number.

This model starts from the official GEAR-Dreams/DreamZero-AgiBot checkpoint. Stage 1 trained only the action encoder/decoder MLPs for one complete 5 FPS RoboDojo epoch on 32 GPUs. Stage 2 reloads the official base, overlays those ten MLP tensors, and fine-tunes the full DreamZero training scope for five complete epochs on 64 GPUs.

The Stage 1 checkpoint is not published separately: it is an intermediate adapter already incorporated into Stage 2 and is not needed for inference.

Two-stage loss trend

This durable interim snapshot uses the complete Stage 1 run and Stage 2 through checkpoint-7388. The first-to-last 50-point mean decreases from 0.2166 to 0.1105 in Stage 1 and from 0.0851 to 0.0341 in Stage 2. Stage 2 is still training; the same image path will be replaced by the complete five-epoch curve in the final release.

Install

conda activate robodojo
cd /path/to/your/dreamzero-integration
pip install -e policy/DreamZero/dreamzero
pip install -e .

hf download Wan-AI/Wan2.1-I2V-14B-480P \
  --local-dir /models/Wan2.1-I2V-14B-480P
hf download google/umt5-xxl --local-dir /models/umt5-xxl

The real-sample demo also requires a local RoboDojo LeRobot v3 dataset. The dataset and the frozen Wan/UMT5 assets are not included in this repository.

Download the latest checkpoint

The following code downloads exactly the checkpoint selected by the latest pointer and writes its local path to latest_model_path.txt:

import json
from pathlib import Path

from huggingface_hub import hf_hub_download, snapshot_download

repo_id = "YunzeLiu/robodojo_dreamzero"
local_dir = Path("/models/robodojo_dreamzero")

pointer_file = hf_hub_download(
    repo_id=repo_id,
    filename="latest_checkpoint.json",
    local_dir=local_dir,
)
pointer = json.loads(Path(pointer_file).read_text())
checkpoint = pointer["checkpoint"]

snapshot_download(
    repo_id=repo_id,
    local_dir=local_dir,
    allow_patterns=[
        "latest_checkpoint.json",
        "stage2_full/latest_checkpoint.json",
        checkpoint + "/**",
    ],
)

model_path = local_dir / checkpoint
(local_dir / "latest_model_path.txt").write_text(
    str(model_path.resolve()) + "\n"
)
print(model_path)

Run hf auth login first if the repository requires authentication. Hugging Face stores portable weights and metadata only; local DeepSpeed optimizer state is not required for inference.

Quick artifact test (no GPU)

MODEL_PATH=$(cat /models/robodojo_dreamzero/latest_model_path.txt)
python - "$MODEL_PATH" <<'PY'
import json
from pathlib import Path
import sys

path = Path(sys.argv[1])
marker = json.loads((path / "checkpoint_complete.json").read_text())

assert marker["training_stage"] == "stage2_full"
assert marker["global_step"] > 0
assert (path / "config.json").is_file()
assert (path / "experiment_cfg/conf.yaml").is_file()
assert (path / "model.safetensors.index.json").is_file()
assert len(list(path.glob("model-*.safetensors"))) == 10

print("checkpoint contract OK:", marker["global_step"], "10 model shards")
PY

This checks the portable artifact layout; it does not execute the policy.

Prepare a real 5 FPS sample without loading the model

python policy/DreamZero/robodojo_inference_demo.py \
  --prepare-only \
  --dataset-path /data/RoboDojo_lerobot_v30_video \
  --episode-index 0 \
  --frame-index 15

The output should report frame indices [0, 5, 10, 15], source FPS 25, policy-input FPS 5, raw state dimension 14, and action horizon 24.

Run one-GPU inference

MODEL_PATH=$(cat /models/robodojo_dreamzero/latest_model_path.txt)

python policy/DreamZero/robodojo_inference_demo.py \
  --model-path "$MODEL_PATH" \
  --wan-checkpoint-path /models/Wan2.1-I2V-14B-480P \
  --tokenizer-path /models/umt5-xxl \
  --dataset-path /data/RoboDojo_lerobot_v30_video \
  --episode-index 0 \
  --frame-index 15 \
  --inference-method lazy_joint_forward_causal \
  --output-dir ./robodojo_inference_result

The demo writes:

• inference_result.json
• predicted_actions.npy
• predicted_action_horizon.png

Validate the inference output

python - <<'PY'
import json
import numpy as np

result = json.load(open("robodojo_inference_result/inference_result.json"))
actions = np.load("robodojo_inference_result/predicted_actions.npy")

assert result["policy_input_fps"] == 5
assert result["frame_sampling_stride"] == 5
assert result["native_action_fps"] == 25
assert result["predicted_action_shape"] == [24, 14]
assert actions.shape == (24, 14)
assert np.isfinite(actions).all()

print("real-checkpoint inference contract OK")
PY

Important runtime contracts

• Policy observations are 5 FPS. Use stride 5 for a 25 Hz source and stride 6 for a 30 Hz source.
• Four observation frames span 0.6 seconds: [t-0.6, t-0.4, t-0.2, t].
• The output is 24 consecutive native 25 Hz actions. Do not downsample the action chunk to 5 Hz.
• Raw action order is [left arm 6, left gripper 1, right arm 6, right gripper 1].
• The model inserts locked AgiBot J3 channels at model indices 2 and 9 and removes exactly those channels on output.
• Call Model.reset() at every episode boundary in a persistent process.

The artifact test and inference smoke test are not task-success metrics. Check camera calibration, joint order, limits, control frequency, emergency stops, and closed-loop behavior before commanding real hardware.