scripts/data/convert_robotwin_to_hf.py

#!/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 目录结构:
    <task_dir>/
    ├── 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()