scripts/data/convert_libero_to_hf.py

#!/usr/bin/env python3
"""将 LIBERO（LeRobot parquet 格式）转换为 HuggingFace datasets 格式。

用法:
    python scripts/data/convert_libero_to_hf.py \\
        --input /path/to/libero_data \\
        --output /path/to/hf_libero \\
        --num-workers 8

输出结构:
    hf_libero/
    ├── data/
    │   ├── train-00000.parquet   # {episode_idx, frame_idx, state, action, action_mask, text, video_path}
    │   └── ...
    ├── videos/
    │   └── view_0/
    │       ├── episode_000000.mp4
    │       └── ...
    ├── dataset_info.json
    └── README.md
"""

import os, json, argparse, logging, multiprocessing as mp
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import cv2
from tqdm import tqdm

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# DreamZero 标准参数
NUM_FRAMES = 12
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, max_dim: int = MAX_STATE_DIM) -> np.ndarray:
    """Pad state vector to max_dim with zeros."""
    d = state.shape[-1]
    if d >= max_dim:
        return state[..., :max_dim].astype(np.float32)
    padded = np.zeros((max_dim,), dtype=np.float32)
    padded[:d] = state.astype(np.float32)
    return padded


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


def process_episode(args):
    """处理单个 episode: 读取视频帧、state、写入 HF 格式。"""
    ep_idx, parquet_path, video_dir, output_dir = args

    try:
        df = pd.read_parquet(parquet_path)
    except Exception as e:
        logger.warning(f"无法读取 {parquet_path}: {e}")
        return None

    output_video_dir = output_dir / "videos" / "view_0"
    output_video_dir.mkdir(parents=True, exist_ok=True)

    # 检测 episode 边界
    episode_indices = []
    if "episode_index" in df.columns:
        episode_indices = df["episode_index"].unique()
    else:
        episode_indices = [0]

    records = []
    for local_ep_idx, ep_val in enumerate(episode_indices):
        ep_mask = df["episode_index"] == ep_val if "episode_index" in df.columns else slice(None)
        ep_df = df[ep_mask].reset_index(drop=True)
        n_frames = len(ep_df)
        if n_frames < NUM_FRAMES + ACTION_HORIZON:
            continue

        global_ep_idx = ep_idx * 1000 + local_ep_idx

        # 提取视频帧并编码为 mp4
        frames = []
        for i in range(n_frames):
            row = ep_df.iloc[i]
            # LeRobot 格式: 图像存为 {'bytes': b'...'} dict
            raw_img = row.get("image", row.get("observation.image", None))
            if isinstance(raw_img, dict) and "bytes" in raw_img:
                img_bytes = raw_img["bytes"]
            elif isinstance(raw_img, bytes):
                img_bytes = raw_img
            else:
                continue

            img_array = np.frombuffer(img_bytes, dtype=np.uint8)
            img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            if img is None:
                continue
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (VIDEO_WIDTH, VIDEO_HEIGHT))
            frames.append(img)

        if len(frames) < NUM_FRAMES + ACTION_HORIZON:
            continue

        # 写入 mp4
        video_filename = f"episode_{global_ep_idx:06d}.mp4"
        video_path = output_video_dir / video_filename
        out_writer = cv2.VideoWriter(
            str(video_path),
            cv2.VideoWriter_fourcc(*"mp4v"),
            FPS,
            (VIDEO_WIDTH, VIDEO_HEIGHT),
        )
        for f in frames:
            out_writer.write(cv2.cvtColor(f, cv2.COLOR_RGB2BGR))
        out_writer.release()

        # 提取 state/action
        for t in range(n_frames - ACTION_HORIZON):
            state = pad_state(ep_df.iloc[t].get("state", ep_df.iloc[t].get("observation.state", np.zeros(14))).astype(np.float32))
            action_chunk = []
            for a in range(ACTION_HORIZON):
                act = ep_df.iloc[t + a].get("action", ep_df.iloc[t + a].get("action.joint_position", np.zeros(14))).astype(np.float32)
                action_chunk.append(act)
            action_chunk = np.stack(action_chunk)  # [horizon, D]
            action_chunk = pad_action(action_chunk)

            records.append({
                "episode_index": global_ep_idx,
                "frame_index": t,
                "state": state.tolist(),
                "action": action_chunk.tolist(),
                "action_mask": [True] * ACTION_HORIZON,
                "text": ep_df.iloc[t].get(
                    "task_description",
                    ep_df.iloc[t].get("annotation.human.action.task_description",
                                      "Perform the task")),
                "video_path": f"videos/view_0/{video_filename}",
            })

    return records


def main():
    parser = argparse.ArgumentParser(description="Convert LIBERO to HF datasets format")
    parser.add_argument("--input", "-i", required=True, help="LIBERO 数据目录 (parquet)")
    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)

    # 扫描 parquet 文件
    parquet_files = sorted(input_dir.glob("file-*.parquet"))
    if not parquet_files:
        # 尝试其他命名模式
        parquet_files = sorted(input_dir.glob("*.parquet"))
    logger.info(f"找到 {len(parquet_files)} 个 parquet 文件")

    # 并行处理
    video_dir = input_dir / "videos"
    tasks = [(i, str(pf), str(video_dir), output_dir) for i, pf in enumerate(parquet_files)]

    all_records = []
    with mp.Pool(args.num_workers) as pool:
        for result in tqdm(
            pool.imap_unordered(process_episode, tasks),
            total=len(tasks),
            desc="Converting episodes",
        ):
            if result:
                all_records.extend(result)

    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)
    pq.write_table(table, data_dir / "train-00000.parquet")

    # 写入 dataset_info.json
    dataset_info = {
        "description": "LIBERO benchmark dataset for DreamZero",
        "features": {
            "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)}},
        "homepage": "https://huggingface.co/datasets/dreamzero/libero",
    }
    with open(output_dir / "dataset_info.json", "w") as f:
        json.dump(dataset_info, f, indent=2)

    # 写入 README.md (dataset card)
    readme_template = """---
license: cc-by-4.0
task_categories:
- robotics
tags:
- robot-vla
- flow-matching
- libero
---

# DreamZero - LIBERO

## Description
LIBERO benchmark dataset converted to HuggingFace datasets format for DreamZero training.

## Schema
| Column | Type | Shape | Description |
|--------|------|-------|-------------|
| video_path | string | - | Path to video file |
| state | float32 | [{STATE_DIM}] | Robot state (padded) |
| action | float32 | [{ACT_HORIZON}, {ACT_DIM}] | Action chunks (padded) |
| action_mask | bool | [{ACT_HORIZON}] | Valid action mask |
| text | string | - | Task instruction |
| episode_index | int64 | - | Episode ID |

## Statistics
- Total samples: {TOTAL_SAMPLES}
- Views: 1
- Video resolution: {W}x{H}
- Frames per sample: {FRAMES}

## Citation
```
@inproceedings{libero2023,
  title={LIBERO: Benchmarking Knowledge Transfer in Lifelong Robot Learning},
  ...
}
```
"""
    readme = readme_template.format(
        STATE_DIM=MAX_STATE_DIM,
        ACT_HORIZON=ACTION_HORIZON,
        ACT_DIM=MAX_ACTION_DIM,
        TOTAL_SAMPLES=len(all_records),
        W=VIDEO_WIDTH,
        H=VIDEO_HEIGHT,
        FRAMES=NUM_FRAMES,
    )
    with open(output_dir / "README.md", "w") as f:
        f.write(readme)

    logger.info(f"转换完成！输出目录: {output_dir}")
    logger.info(f"  数据: {data_dir}/train-00000.parquet")
    logger.info(f"  视频: {output_dir}/videos/view_0/")


if __name__ == "__main__":
    main()