#!/usr/bin/env python3 """将 RoboTwin (qpos.pt + video.mp4 + metas/*.txt) 转换为 HuggingFace datasets 格式。 用法: python scripts/data/convert_robotwin_to_hf.py \\ --input /path/to/robotwin_data \\ --output /path/to/hf_robotwin \\ --num-workers 8 RoboTwin 目录结构: / ├── qpos/ │ ├── episode0.pt # [T_state, 14] float32 │ └── episode1.pt ├── videos/ │ ├── episode0.mp4 # [T_video, H, W, C] │ └── episode1.mp4 └── metas/ ├── task_0.txt # 任务描述 └── task_1.txt """ import os, json, argparse, logging from pathlib import Path import numpy as np import pandas as pd import pyarrow as pa import pyarrow.parquet as pq import torch import cv2 from tqdm import tqdm logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") logger = logging.getLogger(__name__) # DreamZero 标准参数 ACTION_HORIZON = 12 MAX_STATE_DIM = 44 MAX_ACTION_DIM = 32 VIDEO_HEIGHT = 160 VIDEO_WIDTH = 320 FPS = 30 def pad_state(state: np.ndarray) -> np.ndarray: 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: 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 def process_episode(args): """处理一个 RoboTwin episode。""" task_dir_str, ep_name, output_dir, ep_idx = args task_dir = Path(task_dir_str) task_name = task_dir.name video_out_dir = output_dir / "videos" / "view_0" video_out_dir.mkdir(parents=True, exist_ok=True) try: # 读取 qpos qpos_path = task_dir / "qpos" / f"{ep_name}.pt" qpos_data = torch.load(str(qpos_path)).numpy() # [T_state, 14] # 读取视频 video_path = task_dir / "videos" / f"{ep_name}.mp4" cap = cv2.VideoCapture(str(video_path)) n_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap.release() # 读取任务描述 meta_dir = task_dir / "metas" task_desc = "Manipulate the object on the table" if meta_dir.exists(): meta_files = sorted(meta_dir.glob("*.txt")) ep_num = int(ep_name.replace("episode", "")) if ep_num < len(meta_files): with open(meta_files[ep_num]) as f: task_desc = f.read().strip() except Exception as e: logger.warning(f" 无法读取 {task_name}/{ep_name}: {e}") return [] if n_video_frames < ACTION_HORIZON + 1 or len(qpos_data) < ACTION_HORIZON + 1: return [] # 视频帧数 vs state 帧数对齐 video_to_qpos_ratio = len(qpos_data) / max(n_video_frames, 1) records = [] # 读取所有视频帧到内存 cap = cv2.VideoCapture(str(video_path)) all_frames = [] while True: ret, frame = cap.read() if not ret: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = cv2.resize(frame, (VIDEO_WIDTH, VIDEO_HEIGHT)) all_frames.append(frame) cap.release() if len(all_frames) < ACTION_HORIZON + 1: return [] # 将整个 episode 编码为一个 mp4 文件 video_filename = f"episode_{task_name}_{ep_idx:06d}.mp4" video_path_out = video_out_dir / video_filename writer = cv2.VideoWriter( str(video_path_out), cv2.VideoWriter_fourcc(*"mp4v"), FPS, (VIDEO_WIDTH, VIDEO_HEIGHT), ) for frame in all_frames: writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) writer.release() # 为每个时间窗口生成样本 for t in range(len(all_frames) - ACTION_HORIZON): # 对齐 qpos 索引 qpos_t = min(int(t * video_to_qpos_ratio), len(qpos_data) - 1) state = pad_state(qpos_data[min(qpos_t, len(qpos_data) - 1)]) # 从 qpos 差值计算 action (velocity) action_chunk = np.zeros((ACTION_HORIZON, 14), dtype=np.float32) for a in range(ACTION_HORIZON): fi = min(qpos_t + a + 1, len(qpos_data) - 1) si = min(qpos_t + a, len(qpos_data) - 1) action_chunk[a] = qpos_data[fi] - qpos_data[si] action_chunk = pad_action(action_chunk) records.append({ "task": task_name, "episode_index": ep_idx, "frame_index": t, "state": state.tolist(), "action": action_chunk.tolist(), "action_mask": [True] * ACTION_HORIZON, "text": task_desc, "video_path": f"videos/view_0/{video_filename}", }) logger.info(f" {task_name}/{ep_name}: {len(records)} samples") return records def main(): parser = argparse.ArgumentParser(description="Convert RoboTwin to HF datasets format") parser.add_argument("--input", "-i", required=True, help="RoboTwin 数据根目录（包含任务子目录）") parser.add_argument("--output", "-o", required=True, help="HF 数据集输出目录") parser.add_argument("--num-workers", "-w", type=int, default=4, help="并行进程数") args = parser.parse_args() input_dir = Path(args.input) output_dir = Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) # 扫描任务目录 task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()]) logger.info(f"找到 {len(task_dirs)} 个任务目录") all_records = [] global_ep_idx = 0 for task_dir in tqdm(task_dirs, desc="Processing tasks"): video_files = sorted((task_dir / "videos").glob("episode*.mp4")) for vf in video_files: ep_name = vf.stem qf = task_dir / "qpos" / f"{ep_name}.pt" if not qf.exists(): continue records = process_episode((str(task_dir), ep_name, output_dir, global_ep_idx)) all_records.extend(records) global_ep_idx += 1 logger.info(f"总共生成 {len(all_records)} 条训练样本") if not all_records: logger.error("未生成任何样本！") return # 写入 parquet df = pd.DataFrame(all_records) data_dir = output_dir / "data" data_dir.mkdir(parents=True, exist_ok=True) table = pa.Table.from_pandas(df) # 分片（每 10 万条一个文件） num_shards = max(1, len(all_records) // 100_000) if num_shards > 1: shard_size = len(all_records) // num_shards for i in range(num_shards): start = i * shard_size end = start + shard_size if i < num_shards - 1 else len(all_records) shard_table = table.slice(start, end - start) pq.write_table(shard_table, data_dir / f"train-{i:05d}.parquet") else: pq.write_table(table, data_dir / "train-00000.parquet") dataset_info = { "description": "RoboTwin benchmark dataset for DreamZero", "features": { "task": {"dtype": "string", "_type": "Value"}, "episode_index": {"dtype": "int64", "_type": "Value"}, "frame_index": {"dtype": "int64", "_type": "Value"}, "state": {"dtype": "float32", "shape": [MAX_STATE_DIM], "_type": "Sequence"}, "action": {"dtype": "float32", "shape": [ACTION_HORIZON, MAX_ACTION_DIM], "_type": "Sequence"}, "action_mask": {"dtype": "bool", "shape": [ACTION_HORIZON], "_type": "Sequence"}, "text": {"dtype": "string", "_type": "Value"}, "video_path": {"dtype": "string", "_type": "Value"}, }, "splits": {"train": {"num_examples": len(all_records)}}, } with open(output_dir / "dataset_info.json", "w") as f: json.dump(dataset_info, f, indent=2) readme = f"""--- license: cc-by-4.0 --- # DreamZero - RoboTwin ## Description RoboTwin benchmark dataset converted for DreamZero training. ## Schema | Column | Type | Description | |--------|------|-------------| | video_path | string | Video file path | | state | float32[{MAX_STATE_DIM}] | Robot state (padded) | | action | float32[{ACTION_HORIZON},{MAX_ACTION_DIM}] | Action chunks (padded, velocity) | | text | string | Task description | ## Statistics - Total samples: {len(all_records)} - Views: 1 """ with open(output_dir / "README.md", "w") as f: f.write(readme) logger.info(f"转换完成！输出: {output_dir}") if __name__ == "__main__": main()