# 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 到固定维度：

```python
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 类

```python
# 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 配置

```yaml
# 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 训练

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

---

## 4. Embodiment Tag

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

```python
# 内置 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. 数据验证

训练前验证数据格式：

```python
# 快速验证
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` |