#!/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()