| |
| """RoboTwin 数据准备脚本：将原始 ZIP 格式转换为 RobotWinDataset 期望的 qpos+videos+metas 格式。 |
| |
| 用法: |
| python scripts/data/prepare_robotwin.py \\ |
| --input /root/autol-tmp/data/robotwin3/dataset \\ |
| --output /root/autol-tmp/data/robotwin_gear |
| |
| RoboTwin 原始结构 (ZIP): |
| <task_dir>/ |
| ├── franka_clean_50.zip |
| │ └── franka_clean_50/ |
| │ ├── scene_info.json |
| │ ├── instructions/episodeN.json # {"seen": [...], "unseen": [...]} |
| │ ├── video/episodeN.mp4 |
| │ └── _traj_data/episodeN.pkl # {arm_key: [seg1, seg2, ...]} |
| ├── aloha-agilex_clean_50.zip |
| └── ... |
| |
| 输出结构: |
| <output_dir>/ |
| ├── <task>_<robot>/ |
| │ ├── qpos/ |
| │ │ ├── episode0.pt # [T_qpos, 14] float32 |
| │ │ └── episode1.pt |
| │ ├── videos/ |
| │ │ ├── episode0.mp4 |
| │ │ └── episode1.mp4 |
| │ └── metas/ |
| │ ├── task_0.txt |
| │ └── task_1.txt |
| └── ... |
| """ |
|
|
| import os, sys, json, argparse, logging, pickle, shutil |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import zipfile |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") |
| logger = logging.getLogger(__name__) |
|
|
| |
| |
| ROBOT_CONFIG = { |
| "franka": {"arm": "right", "dims": 7}, |
| "aloha-agilex": {"arm": "left", "dims": 6}, |
| "arx-x5": {"arm": "left", "dims": 6}, |
| "ur5": {"arm": "left", "dims": 6}, |
| "piper": {"arm": "left", "dims": 6}, |
| } |
|
|
| OUTPUT_STATE_DIM = 14 |
|
|
|
|
| def get_robot_name(zip_path: Path) -> str | None: |
| """从 ZIP 文件名提取机器人名称（去掉 _clean_50.zip 后缀）。""" |
| name = zip_path.stem |
| |
| for suffix in ["_clean_50", "_50"]: |
| if name.endswith(suffix): |
| name = name[:-len(suffix)] |
| break |
| return name |
|
|
|
|
| def get_arm_key(robot: str) -> str: |
| """根据机器人类型返回 pkl 中的 arm key。""" |
| cfg = ROBOT_CONFIG.get(robot) |
| if cfg is None: |
| logger.warning(f"Unknown robot {robot}, trying right_joint_path") |
| return "right_joint_path" |
| return f"{cfg['arm']}_joint_path" |
|
|
|
|
| def merge_trajectory_segments(segments: list) -> np.ndarray: |
| """合并多段轨迹为一个连续数组。 |
| |
| Args: |
| segments: list of dict, each with "position" key [T_i, D] |
| |
| Returns: |
| concatenated position array [sum(T_i), D] |
| """ |
| arrays = [seg["position"] for seg in segments if seg.get("position") is not None] |
| if not arrays: |
| return np.zeros((0, arrays[0].shape[1])) if arrays else np.zeros((0, 1)) |
| return np.concatenate(arrays, axis=0) |
|
|
|
|
| def pad_to_14dim(arr: np.ndarray, joint_dim: int) -> np.ndarray: |
| """Pad joint positions to 14-dim。 |
| |
| 布局：前 7 = right arm, 后 7 = left arm。 |
| - franka (7, right): 放在前 7 维 |
| - aloha (6, left): 放在后 6 维（前补 0） |
| """ |
| if arr.ndim == 1: |
| padded = np.zeros(OUTPUT_STATE_DIM, dtype=np.float32) |
| if joint_dim == 7: |
| padded[:joint_dim] = arr.astype(np.float32) |
| else: |
| padded[OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32) |
| else: |
| padded = np.zeros((arr.shape[0], OUTPUT_STATE_DIM), dtype=np.float32) |
| if joint_dim == 7: |
| padded[:, :joint_dim] = arr.astype(np.float32) |
| else: |
| padded[:, OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32) |
| return padded |
|
|
|
|
| def get_task_description(zf: zipfile.ZipFile, robot_prefix: str, ep_idx: int) -> str: |
| """从 instruction JSON 读取任务描述。 |
| |
| Returns: |
| 第一条 "seen" 指令，或 fallback 文本 |
| """ |
| try: |
| inst_path = f"{robot_prefix}/instructions/episode{ep_idx}.json" |
| with zf.open(inst_path) as f: |
| inst = json.load(f) |
| seen = inst.get("seen", []) |
| if seen: |
| return seen[0] |
| unseen = inst.get("unseen", []) |
| if unseen: |
| return unseen[0] |
| except Exception as e: |
| logger.debug(f" Cannot read instruction: {e}") |
| return "Manipulate the object on the table" |
|
|
|
|
| def process_robot_zip( |
| zip_path: Path, |
| output_dir: Path, |
| delete_after: bool = False, |
| ) -> int: |
| """处理一个 RoboTwin ZIP 文件。 |
| |
| Returns: |
| 成功处理的 episode 数 |
| """ |
| robot = get_robot_name(zip_path) |
| if robot is None: |
| logger.warning(f"Cannot determine robot from {zip_path.name}, skipping") |
| return 0 |
|
|
| |
| task_name = zip_path.parent.name |
| out_subdir = output_dir / f"{task_name}_{robot}" |
| qpos_dir = out_subdir / "qpos" |
| video_dir = out_subdir / "videos" |
| meta_dir = out_subdir / "metas" |
| qpos_dir.mkdir(parents=True, exist_ok=True) |
| video_dir.mkdir(parents=True, exist_ok=True) |
| meta_dir.mkdir(parents=True, exist_ok=True) |
|
|
| cfg = ROBOT_CONFIG.get(robot) |
| if cfg is None: |
| logger.warning(f"Unknown robot {robot}, skipping {zip_path}") |
| return 0 |
|
|
| joint_dim = cfg["dims"] |
| arm_key = get_arm_key(robot) |
| robot_prefix = f"{robot}_clean_50" |
|
|
| try: |
| zf = zipfile.ZipFile(str(zip_path)) |
| except Exception as e: |
| logger.error(f"Cannot open {zip_path}: {e}") |
| return 0 |
|
|
| |
| try: |
| with zf.open(f"{robot_prefix}/scene_info.json") as f: |
| scene_info = json.load(f) |
| except Exception as e: |
| logger.warning(f"No scene_info in {zip_path}: {e}") |
| zf.close() |
| return 0 |
|
|
| |
| episode_keys = sorted([k for k in scene_info if k.startswith("episode_")], |
| key=lambda x: int(x.split("_")[1])) |
| if not episode_keys: |
| logger.warning(f"No episodes in scene_info of {zip_path}") |
| zf.close() |
| return 0 |
|
|
| success_count = 0 |
| for ep_key in episode_keys: |
| ep_idx = int(ep_key.split("_")[1]) |
| ep_name = f"episode{ep_idx}" |
|
|
| |
| traj_path = f"{robot_prefix}/_traj_data/{ep_name}.pkl" |
| video_path_in = f"{robot_prefix}/video/{ep_name}.mp4" |
|
|
| if traj_path not in zf.namelist(): |
| logger.debug(f" No traj data for {ep_name} in {zip_path.name}") |
| continue |
| if video_path_in not in zf.namelist(): |
| logger.debug(f" No video for {ep_name} in {zip_path.name}") |
| continue |
|
|
| try: |
| |
| with zf.open(traj_path) as f: |
| traj_data = pickle.load(f) |
|
|
| segments = traj_data.get(arm_key, []) |
| if not segments: |
| logger.debug(f" No {arm_key} segments for {ep_name}") |
| continue |
|
|
| |
| pos = merge_trajectory_segments(segments) |
| if pos.shape[0] < 24: |
| logger.debug(f" Too few frames ({pos.shape[0]}) for {ep_name}") |
| continue |
|
|
| |
| pos_14 = pad_to_14dim(pos, joint_dim) |
|
|
| |
| qpos_path = qpos_dir / f"{ep_name}.pt" |
| torch.save(torch.from_numpy(pos_14), qpos_path) |
|
|
| |
| video_out_path = video_dir / f"{ep_name}.mp4" |
| with zf.open(video_path_in) as src, open(video_out_path, "wb") as dst: |
| shutil.copyfileobj(src, dst) |
|
|
| |
| task_desc = get_task_description(zf, robot_prefix, ep_idx) |
| meta_path = meta_dir / f"task_{ep_idx}.txt" |
| with open(meta_path, "w") as f: |
| f.write(task_desc) |
|
|
| success_count += 1 |
|
|
| except Exception as e: |
| logger.warning(f" Error processing {ep_name} in {zip_path.name}: {e}") |
| continue |
|
|
| zf.close() |
|
|
| |
| if delete_after and success_count > 0: |
| try: |
| zip_path.unlink() |
| logger.info(f" Deleted {zip_path.name}") |
| except Exception as e: |
| logger.warning(f" Cannot delete {zip_path.name}: {e}") |
|
|
| if success_count > 0: |
| logger.info(f"{zip_path.name}: {success_count}/{len(episode_keys)} episodes extracted -> {out_subdir}") |
|
|
| return success_count |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Prepare RoboTwin data for DreamZero training") |
| parser.add_argument("--input", "-i", required=True, |
| help="RoboTwin 数据目录 (含任务子目录)") |
| parser.add_argument("--output", "-o", required=True, |
| help="输出目录") |
| parser.add_argument("--delete-zip", action="store_true", |
| help="处理完成后删除 ZIP 文件（节省空间）") |
| parser.add_argument("--max-tasks", type=int, default=None, |
| help="最多处理前 N 个任务（用于测试）") |
| parser.add_argument("--num-workers", 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"Found {len(task_dirs)} task directories") |
|
|
| if args.max_tasks: |
| task_dirs = task_dirs[:args.max_tasks] |
| logger.info(f"Limited to {args.max_tasks} tasks") |
|
|
| total_episodes = 0 |
| total_zips = 0 |
|
|
| for task_dir in task_dirs: |
| |
| zip_files = sorted(task_dir.glob("*_clean_50.zip")) |
| if not zip_files: |
| logger.warning(f"No ZIP files in {task_dir}") |
| continue |
|
|
| for zip_path in zip_files: |
| try: |
| ep_count = process_robot_zip( |
| zip_path, output_dir, delete_after=args.delete_zip |
| ) |
| total_episodes += ep_count |
| total_zips += 1 |
| except Exception as e: |
| logger.error(f"Fatal error processing {zip_path}: {e}") |
| continue |
|
|
| logger.info(f"Done! Processed {total_zips} ZIPs, {total_episodes} episodes") |
| logger.info(f"Output: {output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
