DreamZero Documentation

DreamZero executes policies using a repeated "observe $\rightarrow$ predict $\rightarrow$ execute" loop that leverages a KV cache to maintain efficiency and causality.

1. Observe: The robot receives a new observation (image history + state).
2. Predict: The policy (e.g., via lazy_joint_video_action) runs a diffusion loop for the current block. It uses the KV cache containing keys/values from previous blocks so it only processes the new block's tokens and the new action register.
3. Execute: The robot executes the resulting action chunk (e.g., 24 actions) without further model calls.
4. Repeat: Once execution is complete, the current_start_frame is incremented by num_frame_per_block, the KV cache is reused, and the loop starts again with the next block.

If the task or language changes, or if the KV cache is full, the cache and current_start_frame must be reset.

To train a new embodiment using LoRA fine-tuning, create a shell script at scripts/train/<EMBODIMENT>_training.sh. The script uses torchrun to launch groot/vla/experiment/experiment.py with specific Hydra configurations.

Key environment variables required:

• DATA_ROOT: Path to your GEAR-converted dataset.
• OUTPUT_DIR: Where to save checkpoints (defaults to ./checkpoints/dreamzero_<EMBODIMENT>_lora).
• NUM_GPUS: Number of GPUs to use (defaults to count from nvidia-smi).
• WAN_CKPT_DIR: Path to Wan2.1 weights (defaults to ./checkpoints/Wan2.1-I2V-14B-480P).
• TOKENIZER_DIR: Path to umt5-xxl weights (defaults to ./checkpoints/umt5-xxl).

The script automatically attempts to download missing Wan2.1 and umt5-xxl weights via huggingface-cli if the directories are empty.

#!/bin/bash
export HYDRA_FULL_ERROR=1

# ============ CONFIGURATION ==========================================
DATA_ROOT=${DATA_ROOT:?"Set DATA_ROOT to your GEAR-converted dataset"}
OUTPUT_DIR=${OUTPUT_DIR:-"./checkpoints/dreamzero_<EMBODIMENT>_lora"}

if [ -z "${NUM_GPUS:-}" ]; then
  NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
fi
NUM_GPUS=${NUM_GPUS:-8}

WAN_CKPT_DIR=${WAN_CKPT_DIR:-"./checkpoints/Wan2.1-I2V-14B-480P"}
TOKENIZER_DIR=${TOKENIZER_DIR:-"./checkpoints/umt5-xxl"}
# =====================================================================

# Auto-download weights if missing
if [ ! -d "$WAN_CKPT_DIR" ] || [ -z "$(ls -A "$WAN_CKPT_DIR" 2>/dev/null)" ]; then
    huggingface-cli download Wan-AI/Wan2.1-I2V-14B-480P --local-dir "$WAN_CKPT_DIR"
fi
if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
    huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
fi

if [ ! -d "$DATA_ROOT" ]; then
    echo "ERROR: Dataset not found at $DATA_ROOT"
    exit 1
fi
if [ ! -f "$DATA_ROOT/meta/embodiment.json" ]; then
    echo "ERROR: meta/embodiment.json missing — run convert_lerobot_to_gear.py first"
    exit 1
fi

torchrun --nproc_per_node $NUM_GPUS --standalone \
    groot/vla/experiment/experiment.py \
    report_to=wandb \
    data=dreamzero/<EMBODIMENT>_relative \
    wandb_project=dreamzero \
    train_architecture=lora \
    num_frames=33 \
    action_horizon=24 \
    num_views=3 \
    model=dreamzero/vla \
    model/dreamzero/action_head=wan_flow_matching_action_tf \
    model/dreamzero/transform=dreamzero_cotrain \
    num_frame_per_block=2 \
    num_action_per_block=24 \
    num_state_per_block=1 \
    seed=42 \
    training_args.learning_rate=1e-5 \
    training_args.deepspeed="groot/vla/configs/deepspeed/zero2.json" \
    save_steps=10000 \
    training_args.warmup_ratio=0.05 \
    output_dir=$OUTPUT_DIR \
    per_device_train_batch_size=4 \
    max_steps=100000 \
    weight_decay=1e-5 \
    save_total_limit=10 \
    upload_checkpoints=false \
    bf16=true \
    tf32=true \
    eval_bf16=true \
    dataloader_pin_memory=false \
    dataloader_num_workers=1 \
    image_resolution_width=320 \
    image_resolution_height=176 \
    save_lora_only=true \
    max_chunk_size=4 \
    frame_seqlen=880 \
    save_strategy=steps \
    <EMBODIMENT>_data_root=$DATA_ROOT \
    dit_version=$WAN_CKPT_DIR \
    text_encoder_pretrained_path=$WAN_CKPT_DIR/models_t5_umt5-xxl-enc-bf16.pth \
    image_encoder_pretrained_path=$WAN_CKPT_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth \
    vae_pretrained_path=$WAN_CKPT_DIR/Wan2.1_VAE.pth \
    tokenizer_path=$TOKENIZER_DIR \
    pretrained_model_path=./checkpoints/DreamZero-AgiBot \
    ++action_head_cfg.config.skip_component_loading=true \
    ++action_head_cfg.config.defer_lora_injection=true

Modality connects dataset columns to the training pipeline. You must edit groot/vla/configs/data/dreamzero/base_48_wan_fine_aug_relative.yaml to include your embodiment's configuration.

1. Modality Configuration

Define modality_config_<EMBODIMENT> using groot.vla.data.dataset.ModalityConfig. The keys used here must match the names in your generated meta/modality.json with type prefixes:

• State: state.<name>
• Action: action.<name>
• Video: video.<name>
• Language: annotation.<name>

delta_indices usage:

• Video: Frame offsets to sample (e.g., [0, 1, ..., 24]).
• State / Language: [0] for current timestep only.
• Action: Future offsets (e.g., [0, 1, ..., 23]) for action chunking.

2. Transform Configuration

Define transform_<EMBODIMENT> using groot.vla.data.transform.ComposedModalityTransform.

• Normalization: Every state and action key must appear in normalization_modes (typically using q99).
• Concat: Use ConcatTransform to define the order of concatenation for video, state, and action keys.

3. Global Registration

Register your new configs in the base YAML's global maps:

• modality_configs
• transforms
• metadata_versions
• fps

modality_config_<EMBODIMENT>:
  video:
    _target_: groot.vla.data.dataset.ModalityConfig
    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
    eval_delta_indices: [0]
    modality_keys:
      - video.cam0
      - video.cam1
      - video.cam2
  state:
    _target_: groot.vla.data.dataset.ModalityConfig
    delta_indices: [0]
    modality_keys:
      - state.joint_pos
      - state.gripper_pos
  action:
    _target_: groot.vla.data.dataset.ModalityConfig
    delta_indices: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
    modality_keys:
      - action.joint_pos
      - action.gripper_pos
  language:
    _target_: groot.vla.data.dataset.ModalityConfig
    delta_indices: [0]
    modality_keys:
      - annotation.task

transform_<EMBODIMENT>:
  _target_: groot.vla.data.transform.ComposedModalityTransform
  transforms:
    - <<: *totensor_cfg
      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
    - <<: *crop_cfg
      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
    - <<: *resize_cfg
      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
    - <<: *color_jitter_cfg
      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
    - <<: *to_numpy_cfg
      apply_to: ${modality_config_<EMBODIMENT>.video.modality_keys}
    - _target_: groot.vla.data.transform.StateActionToTensor
      apply_to: ${modality_config_<EMBODIMENT>.state.modality_keys}
    - _target_: groot.vla.data.transform.StateActionTransform
      apply_to: ${modality_config_<EMBODIMENT>.state.modality_keys}
      normalization_modes:
        state.joint_pos: q99
        state.gripper_pos: q99
    - _target_: groot.vla.data.transform.StateActionToTensor
      apply_to: ${modality_config_<EMBODIMENT>.action.modality_keys}
    - _target_: groot.vla.data.transform.StateActionTransform
      apply_to: ${modality_config_<EMBODIMENT>.action.modality_keys}
      normalization_modes:
        action.joint_pos: q99
        action.gripper_pos: q99
    - _target_: groot.vla.data.transform.ConcatTransform
      video_concat_order: ${modality_config_<EMBODIMENT>.video.modality_keys}
      state_concat_order: ${modality_config_<EMBODIMENT>.state.modality_keys}
      action_concat_order: ${modality_config_<EMBODIMENT>.action.modality_keys}
    - ${model_specific_transform}

To reproduce the DreamZero DROID dataset conversion or modify the filtering, follow these steps:

1. Install dependencies: Install tensorflow, tensorflow-datasets, polars, and av via pip.
2. Download raw DROID 1.0.1: Use gsutil to download the raw dataset from Google Cloud Storage. Note: You must use version 1.0.1 to ensure you get the complete set of language annotations (~75k episodes).
3. Download idle filter ranges: Download the droid_sample_ranges_v1_0_1.json file from the openpi-assets bucket. This file maps episodes to non-idle frame ranges.
4. Run conversion script: Execute scripts/data/convert_droid.py with the raw data path, output path, and the filter ranges path.

# Step 1: Install dependencies
pip install tensorflow tensorflow-datasets polars av

# Step 2: Download raw DROID 1.0.1 (~1.7TB)
gsutil -m cp -r gs://gresearch/robotics/droid/1.0.1 ./data/droid/1.0.1

# Step 3: Download idle filter ranges
gsutil cp gs://openpi-assets/droid/droid_sample_ranges_v1_0_1.json ./data/keep_ranges.json

# Step 4: Run the conversion
python scripts/data/convert_droid.py \
    ./data/droid/1.0.1 \
    ./data/droid_lerobot \
    --keep-ranges-path ./data/keep_ranges.json \
    --filter-failed \
    -n 16

To add a new robot to DreamZero, you must first convert a LeRobot v2 dataset into the GEAR metadata format. The converter reads the dataset and writes metadata files to a meta/ directory without modifying the original parquet files or videos.

Expected Input Structure

Your dataset must follow this structure:

your_dataset/
├── data/
│   └── chunk-000/
│       ├── episode_000000.parquet
│       └── ...
├── videos/
│   └── chunk-000/
│       ├── observation.images.cam0/
│       │   ├── episode_000000.mp4
│       │   └── ...
│       └── observation.images.cam1/
│           └── ...
└── meta/
    └── info.json          # must contain: features, total_episodes, fps

Conversion Command

Run the convert_lerobot_to_gear.py script. Use --state-keys and --action-keys to map packed vector columns to named sub-keys using a JSON mapping of sub-key name -> [start_index, end_index].

python scripts/data/convert_lerobot_to_gear.py \
    --dataset-path /path/to/your_dataset \
    --embodiment-tag <EMBODIMENT> \
    --state-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \
    --action-keys '{"joint_pos": [0, 6], "gripper_pos": [6, 7]}' \
    --relative-action-keys joint_pos gripper_pos \
    --task-key annotation.task

To train DreamZero, you must first download the base model weights (Wan2.1-I2V-14B-480P and umt5-xxl tokenizer) and the preprocessed DROID dataset.

1. Download Weights and Dataset

pip install "huggingface_hub[cli]"

# Download Wan2.1 model weights (~28GB)
hf download Wan-AI/Wan2.1-I2V-14B-480P --local-dir ./checkpoints/Wan2.1-I2V-14B-480P

# Download umt5-xxl tokenizer
hf download google/umt5-xxl --local-dir ./checkpoints/umt5-xxl

# Download preprocessed DROID dataset (~131GB)
huggingface-cli download GEAR-Dreams/DreamZero-DROID-Data --repo-type dataset --local-dir ./data/droid_lerobot

2. Launch Training

Configure the environment variables and run the training script:

export DROID_DATA_ROOT="./data/droid_lerobot"
export OUTPUT_DIR="./checkpoints/dreamzero_droid"
export NUM_GPUS=4
export WAN_CKPT_DIR="./checkpoints/Wan2.1-I2V-14B-480P"
export TOKENIZER_DIR="./checkpoints/umt5-xxl"

bash scripts/train/droid_training.sh

Training Configuration (Hydra/DeepSpeed)

Prerequisites

• Python: 3.11
• Hardware: Multi-GPU setup (minimum 2 GPUs for distributed inference)
• CUDA: Compatible GPU with CUDA 12.9+

Installation Steps

1. Create and activate a conda environment:

conda create -n dreamzero python=3.11
conda activate dreamzero

2. Install dependencies (PyTorch 2.8+ with CUDA 12.9+):

pip install -e . --extra-index-url https://download.pytorch.org/whl/cu129

3. Install flash attention:

MAX_JOBS=8 pip install --no-build-isolation flash-attn

4. [GB200 ONLY] Install Transformer Engine:

pip install --no-build-isolation transformer_engine[pytorch]

5. [GB200 ONLY FOR TENSORRT] Install Tensorrt:

pip install tensorrt==10.13.2.6 tensorrt_cu13==10.13.2.6 tensorrt_cu13_libs==10.13.2.6 tensorrt_cu13_bindings==10.13.2.6 --no-deps
pip install transformer_engine==2.10.0 transformer_engine_cu12==2.10.0 transformer_engine_torch==2.10.0