docs/data-preparation.md

DreamZero 数据准备指南

---

1. 数据格式要求

DreamZero 期望每个训练样本包含以下字段：

字段	类型	形状	说明
video	uint8 numpy	[T, V, H, W, 3]	T=帧数, V=视角数
state	float32 numpy	[state_horizon, max_state_dim]	机器人状态（需 padding）
action	float32 numpy	[action_horizon, max_action_dim]	动作序列（需 padding）
language	str	-	任务描述

参数约束

action_horizon / (lat_T - 1) = num_action_per_block / num_frame_per_block
(lat_T - 1) / state_horizon = num_frame_per_block / num_state_per_block
lat_T = num_frames // 4           # Wan2.2 VAE 4x 时间下采样

标准参数（已验证）:

• num_frames=12, action_horizon=12, state_horizon=1
• num_frame_per_block=2, num_action_per_block=12, num_state_per_block=1
• max_state_dim=44, max_action_dim=32

---

2. State/Action Padding

State 和 Action 统一 padding 到固定维度：

import numpy as np

MAX_STATE_DIM = 44
MAX_ACTION_DIM = 32

def pad_state(state: np.ndarray) -> np.ndarray:
    """Pad state to MAX_STATE_DIM."""
    d = state.shape[-1]
    padded = np.zeros((MAX_STATE_DIM,), dtype=np.float32)
    padded[:d] = state.astype(np.float32)
    return padded

def pad_action(action_chunk: np.ndarray) -> np.ndarray:
    """Pad action to [horizon, MAX_ACTION_DIM]."""
    d = action_chunk.shape[-1]
    padded = np.zeros((*action_chunk.shape[:-1], MAX_ACTION_DIM), dtype=np.float32)
    padded[..., :d] = action_chunk.astype(np.float32)
    return padded

---

3. 添加新数据集步骤

3.1 创建 Dataset 类

# groot/vla/data/dataset/my_dataset.py
from pathlib import Path
import numpy as np
from groot.vla.model.dreamzero.transform.dreamzero_cotrain import DreamTransform

class MyDataset:
    def __init__(self, dataset_dir, num_frames=12, action_horizon=12,
                 state_horizon=1, num_views=1, ...):
        # 1. 扫描数据文件
        # 2. 构建 episode 列表
        # 3. 创建 DreamTransform 实例
        
        self.transform = DreamTransform(
            default_instruction="Perform the task.",
            max_state_dim=44, max_action_dim=32,
            state_horizon=state_horizon,
            action_horizon=action_horizon,
            num_views=num_views,
            embodiment_tag_mapping={"my_robot": 17},
            tokenizer_path="/path/to/umt5-xxl",
        )
        # 必须设置 metadata
        self.transform.set_metadata(self.merged_metadata["my_robot"])
        self.transform.train()
    
    def __getitem__(self, idx):
        # 返回格式:
        return dict(self.transform({
            "video": video,           # [T, V, H, W, C] uint8
            "state": state,           # [T, D] float32
            "action": action,         # [T, D] float32
            "annotation.human.action.task_description": text,
        }))

3.2 创建 Hydra 配置

# groot/vla/configs/data/dreamzero/my_dataset.yaml
defaults:
  - dreamzero/base_48_wan_fine_aug_relative
  - _self_

my_dataset_dir: /path/to/data

num_frames: 12
action_horizon: 12
state_horizon: 1
num_views: 1
num_frame_per_block: 2
num_action_per_block: 12
num_state_per_block: 1
max_state_dim: 44
max_action_dim: 32
max_chunk_size: 4
image_resolution_height: 160
image_resolution_width: 320
frame_seqlen: 50

train_dataset:
  _target_: groot.vla.data.dataset.my_dataset.MyDataset
  _convert_: object
  dataset_dir: ${my_dataset_dir}
  num_frames: ${num_frames}
  action_horizon: ${action_horizon}
  state_horizon: ${state_horizon}
  num_views: ${num_views}
  video_height: ${image_resolution_height}
  video_width: ${image_resolution_width}

3.3 训练

torchrun ... experiment.py \
  data=dreamzero/my_dataset \
  num_frames=12 action_horizon=12 num_views=1

---

4. Embodiment Tag

每个数据集需要注册一个 embodiment tag（用于识别机器人种类）：

# 内置 tag 映射 (部分)
embodiment_tag_mapping = {
    "oxe_droid": 17,   # DROID 机器人
    "libero": 18,      # LIBERO
    "panda": 19,       # Franka Panda
    "manifeel": 20,    # ManiFeel
    "robotwin": 21,    # RoboTwin
}

如果不确定用哪个 tag，使用 "oxe_droid": 17。

---

5. 数据验证

训练前验证数据格式：

# 快速验证
ds = MyDataset(dataset_dir="/path/to/data", max_episodes=3)
sample = ds[0]
print(sample.keys())
for k, v in sample.items():
    if hasattr(v, 'shape'):
        print(f"  {k}: {v.shape}, {v.dtype}")

# 验证 DreamTransform 输出
transformed = dict(ds.transform({
    "video": sample["video"],
    "state": sample["state"],
    "action": sample["action"],
    "annotation.human.action.task_description": "test",
}))

---

6. 常见数据问题

问题	原因	解决
lat_T = 0	num_frames < 4	设置 num_frames >= 4
reshape error	num_frame_per_block 不对齐	确保 (lat_T-1) % num_frame_per_block == 0
InterpolationKeyError	config 缺少字段	检查 YAML 包含所有 ${...} 引用
VAE dtype error	bf16 vs float32 不匹配	VAE 加载时用 dtype=self.dtype