File size: 5,231 Bytes
ac29381 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 | # 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` |
|