From 297b67cbc22d5ee13d9a552c7d27317aee8852ec Mon Sep 17 00:00:00 2001 From: Qizhi Chen Date: Thu, 12 Jun 2025 20:13:33 +0800 Subject: [PATCH] make script compatible with lerobot (b536f47) (#38) * bump openx2lerobot script * bump agibot2lerobot script * bump robomind2lerobot script --- agibot2lerobot/README.md | 6 ++-- agibot2lerobot/agibot_h5.py | 26 +++++--------- agibot2lerobot/agibot_utils/agibot_utils.py | 3 +- openx2lerobot/README.md | 19 +++------- openx2lerobot/openx_rlds.py | 15 +++----- robomind2lerobot/README.md | 8 ++--- robomind2lerobot/robomind_h5.py | 39 ++++++--------------- 7 files changed, 36 insertions(+), 80 deletions(-) diff --git a/agibot2lerobot/README.md b/agibot2lerobot/README.md index b6fed09..73ce22d 100644 --- a/agibot2lerobot/README.md +++ b/agibot2lerobot/README.md @@ -174,7 +174,7 @@ python convert.py \ #### For single node ```bash -cd agibot2lerobot && bash convert.sh +bash convert.sh ``` #### For multi nodes @@ -197,7 +197,7 @@ On either Node, check the ray cluster status, and start the script ```bash ray status -cd agibot2lerobot && bash convert.sh +bash convert.sh ``` **Slurm-managed System** @@ -257,7 +257,7 @@ done sleep 10 -cd agibot2lerobot && bash convert.sh +bash convert.sh ``` **Other Community Supported Cluster Managers** diff --git a/agibot2lerobot/agibot_h5.py b/agibot2lerobot/agibot_h5.py index cc8fae0..2bbdece 100644 --- a/agibot2lerobot/agibot_h5.py +++ b/agibot2lerobot/agibot_h5.py @@ -25,7 +25,6 @@ from lerobot.common.datasets.utils import ( write_info, ) from lerobot.common.datasets.video_utils import get_safe_default_codec -from lerobot.common.robot_devices.robots.utils import Robot from ray.runtime_env import RuntimeEnv @@ -72,10 +71,9 @@ class AgiBotDataset(LeRobotDataset): cls, repo_id: str, fps: int, + features: dict, root: str | Path | None = None, - robot: Robot | None = None, robot_type: str | None = None, - features: dict | None = None, use_videos: bool = True, tolerance_s: float = 1e-4, image_writer_processes: int = 0, @@ -87,10 +85,9 @@ class AgiBotDataset(LeRobotDataset): obj.meta = AgiBotDatasetMetadata.create( repo_id=repo_id, fps=fps, - root=root, - robot=robot, robot_type=robot_type, features=features, + root=root, use_videos=use_videos, ) obj.repo_id = obj.meta.repo_id @@ -114,7 +111,7 @@ class AgiBotDataset(LeRobotDataset): obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec() return obj - def add_frame(self, frame: dict) -> None: + def add_frame(self, frame: dict, task: str, timestamp: float | None = None) -> None: """ This function only adds the frame to the episode_buffer. Apart from images — which are written in a temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method @@ -133,17 +130,14 @@ class AgiBotDataset(LeRobotDataset): # Automatically add frame_index and timestamp to episode buffer frame_index = self.episode_buffer["size"] - timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps + if timestamp is None: + timestamp = frame_index / self.fps self.episode_buffer["frame_index"].append(frame_index) self.episode_buffer["timestamp"].append(timestamp) + self.episode_buffer["task"].append(task) # Add frame features to episode_buffer for key, value in frame.items(): - if key == "task": - # Note: we associate the task in natural language to its task index during `save_episode` - self.episode_buffer["task"].append(frame["task"]) - continue - if key not in self.features: raise ValueError( f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'." @@ -246,7 +240,7 @@ def save_as_lerobot_dataset(agibot_world_config, task: tuple[Path, Path], num_th if not save_depth: features.pop("observation.images.head_depth") - dataset = AgiBotDataset.create( + dataset: AgiBotDataset = AgiBotDataset.create( repo_id=json_file.stem, root=local_dir, fps=30, @@ -268,7 +262,6 @@ def save_as_lerobot_dataset(agibot_world_config, task: tuple[Path, Path], num_th eid, src_path=src_path, task_id=task_id, - task_instruction=task_instruction, save_depth=save_depth, AgiBotWorld_CONFIG=agibot_world_config, ) @@ -278,7 +271,7 @@ def save_as_lerobot_dataset(agibot_world_config, task: tuple[Path, Path], num_th continue for frame_data in frames: - dataset.add_frame(frame_data) + dataset.add_frame(frame_data, task_instruction) try: dataset.save_episode(videos=videos, action_config=action_config) except Exception as e: @@ -300,7 +293,6 @@ def save_as_lerobot_dataset(agibot_world_config, task: tuple[Path, Path], num_th eid, src_path=src_path, task_id=task_id, - task_instruction=task_instruction, save_depth=save_depth, AgiBotWorld_CONFIG=agibot_world_config, ) @@ -313,7 +305,7 @@ def save_as_lerobot_dataset(agibot_world_config, task: tuple[Path, Path], num_th continue action_config = task_info[eid]["label_info"]["action_config"] for frame_data in frames: - dataset.add_frame(frame_data) + dataset.add_frame(frame_data, task_instruction) try: dataset.save_episode(videos=videos, action_config=action_config) except Exception as e: diff --git a/agibot2lerobot/agibot_utils/agibot_utils.py b/agibot2lerobot/agibot_utils/agibot_utils.py index 15a7351..54fd3c5 100644 --- a/agibot2lerobot/agibot_utils/agibot_utils.py +++ b/agibot2lerobot/agibot_utils/agibot_utils.py @@ -20,7 +20,7 @@ def load_depths(root_dir: str, camera_name: str): def load_local_dataset( - episode_id: int, src_path: str, task_id: int, task_instruction: str, save_depth: bool, AgiBotWorld_CONFIG: dict + episode_id: int, src_path: str, task_id: int, save_depth: bool, AgiBotWorld_CONFIG: dict ) -> tuple[list, dict]: """Load local dataset and return a dict with observations and actions""" ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}" @@ -81,7 +81,6 @@ def load_local_dataset( ) for key, value in action.items() }, - "task": task_instruction, } for i in range(num_frames) ] diff --git a/openx2lerobot/README.md b/openx2lerobot/README.md index 6ee4aa6..65d01f6 100644 --- a/openx2lerobot/README.md +++ b/openx2lerobot/README.md @@ -1,5 +1,7 @@ # OpenX to LeRobot +Open X-Embodiment assembles a dataset from 22 different robots collected through a collaboration between 21 institutions, demonstrating 527 skills (160266 tasks). (Copied from [docs](https://robotics-transformer-x.github.io/)) + ## 🚀 What's New in This Script In this dataset, we have made several key improvements: @@ -69,18 +71,7 @@ Dataset Structure of `meta/info.json`: ## Get started > [!IMPORTANT] -> 1.Before running the following code, modify `save_episode()` function in lerobot. -> ```python -> def save_episode(self, episode_data: dict | None = None, keep_images: bool | None = False) -> None: -> ... -> # delete images -> if not keep_images: -> img_dir = self.root / "images" -> if img_dir.is_dir(): -> shutil.rmtree(self.root / "images") -> ... -> ``` -> 2.for `bc_z` dataset, modify `encode_video_frames()` in `lerobot/common/datasets/video_utils.py`. +> 1.for `bc_z` dataset, modify `encode_video_frames()` in `lerobot/common/datasets/video_utils.py`. > > ```python > # add the following content to line 141: @@ -99,7 +90,7 @@ Dataset Structure of `meta/info.json`: git clone https://github.com/Tavish9/any4lerobot.git ``` -2. Modify path in `openx2lerobot/convert.sh`: +2. Modify path in `convert.sh`: ```bash python openx_rlds.py \ @@ -113,7 +104,7 @@ Dataset Structure of `meta/info.json`: 3. Execute the script: ```bash - cd openx2lerobot && bash convert.sh + bash convert.sh ``` ## Available OpenX_LeRobot Dataset diff --git a/openx2lerobot/openx_rlds.py b/openx2lerobot/openx_rlds.py index b2d6e9f..2f7af34 100644 --- a/openx2lerobot/openx_rlds.py +++ b/openx2lerobot/openx_rlds.py @@ -37,10 +37,8 @@ from pathlib import Path import numpy as np import tensorflow as tf import tensorflow_datasets as tfds -from huggingface_hub import HfApi from lerobot.common.constants import HF_LEROBOT_HOME from lerobot.common.datasets.lerobot_dataset import LeRobotDataset - from oxe_utils.configs import OXE_DATASET_CONFIGS, ActionEncoding, StateEncoding from oxe_utils.transforms import OXE_STANDARDIZATION_TRANSFORMS @@ -149,10 +147,10 @@ def save_as_lerobot_dataset(lerobot_dataset: LeRobotDataset, raw_dataset: tf.dat **image_dict, "observation.state": traj["proprio"][i], "action": traj["action"][i], - "task": traj["task"][0].decode(), - } + }, + task=traj["task"][0].decode(), ) - lerobot_dataset.save_episode(keep_images=kwargs.get("keep_images", False)) + lerobot_dataset.save_episode() def create_lerobot_dataset( @@ -209,7 +207,7 @@ def create_lerobot_dataset( repo_id=repo_id, robot_type=robot_type, root=local_dir, - fps=fps, + fps=int(fps), use_videos=use_videos, features=features, image_writer_threads=image_writer_threads, @@ -287,11 +285,6 @@ def main(): default=10, help="Number of threads per process of image writer for saving images.", ) - parser.add_argument( - "--keep-images", - action="store_true", - help="Whether to keep the cached images.", - ) args = parser.parse_args() create_lerobot_dataset(**vars(args)) diff --git a/robomind2lerobot/README.md b/robomind2lerobot/README.md index 36117aa..74195db 100644 --- a/robomind2lerobot/README.md +++ b/robomind2lerobot/README.md @@ -184,7 +184,7 @@ Dataset Structure of `meta/info.json`: > ``` > [!NOTE] -> The conversion speed of this script is limited by the performance of the physical machine running it, including **CPU cores and memory**. We recommend using **2 CPU cores per task** for optimal performance. However, each task requires approximately 20 GiB of memory. To avoid running out of memory, you may need to increase the number of CPU cores per task depending on your system’s available memory. +> The conversion speed of this script is limited by the performance of the physical machine running it, including **CPU cores and memory**. We recommend using **2 CPU cores per task** for optimal performance. However, each task requires approximately 10 GiB of memory. To avoid running out of memory, you may need to increase the number of CPU cores per task depending on your system’s available memory. ### Download source code: @@ -210,7 +210,7 @@ python robomind_h5.py \ #### For single node ```bash -cd robomind2lerobot && bash convert.sh +bash convert.sh ``` #### For multi nodes @@ -233,7 +233,7 @@ On either Node, check the ray cluster status, and start the script ```bash ray status -cd robomind2lerobot && bash convert.sh +bash convert.sh ``` **Slurm-managed System** @@ -293,7 +293,7 @@ done sleep 10 -cd robomind2lerobot && bash convert.sh +bash convert.sh ``` **Other Community Supported Cluster Managers** diff --git a/robomind2lerobot/robomind_h5.py b/robomind2lerobot/robomind_h5.py index cbf12f0..e04ed9c 100644 --- a/robomind2lerobot/robomind_h5.py +++ b/robomind2lerobot/robomind_h5.py @@ -5,7 +5,6 @@ import logging import shutil from pathlib import Path -import datasets import numpy as np import pandas as pd import ray @@ -15,8 +14,6 @@ from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatas from lerobot.common.datasets.utils import ( check_timestamps_sync, get_episode_data_index, - get_hf_features_from_features, - hf_transform_to_torch, validate_episode_buffer, validate_frame, write_episode, @@ -24,7 +21,6 @@ from lerobot.common.datasets.utils import ( write_info, ) from lerobot.common.datasets.video_utils import get_safe_default_codec -from lerobot.common.robot_devices.robots.utils import Robot from ray.runtime_env import RuntimeEnv from robomind_uitls.configs import ROBOMIND_CONFIG from robomind_uitls.lerobot_uitls import compute_episode_stats, generate_features_from_config @@ -81,10 +77,9 @@ class RoboMINDDataset(LeRobotDataset): cls, repo_id: str, fps: int, + features: dict, root: str | Path | None = None, - robot: Robot | None = None, robot_type: str | None = None, - features: dict | None = None, use_videos: bool = True, tolerance_s: float = 1e-4, image_writer_processes: int = 0, @@ -96,10 +91,9 @@ class RoboMINDDataset(LeRobotDataset): obj.meta = RoboMINDDatasetMetadata.create( repo_id=repo_id, fps=fps, - root=root, - robot=robot, robot_type=robot_type, features=features, + root=root, use_videos=use_videos, ) obj.repo_id = obj.meta.repo_id @@ -123,16 +117,7 @@ class RoboMINDDataset(LeRobotDataset): obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec() return obj - def create_hf_dataset(self) -> datasets.Dataset: - features = get_hf_features_from_features(self.features) - ft_dict = {col: [] for col in features} - hf_dataset = datasets.Dataset.from_dict(ft_dict, features=features, split="train") - - # TODO(aliberts): hf_dataset.set_format("torch") - hf_dataset.set_transform(hf_transform_to_torch) - return hf_dataset - - def add_frame(self, frame: dict) -> None: + def add_frame(self, frame: dict, task: str, timestamp: float | None = None) -> None: """ This function only adds the frame to the episode_buffer. Apart from images — which are written in a temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method @@ -150,17 +135,14 @@ class RoboMINDDataset(LeRobotDataset): # Automatically add frame_index and timestamp to episode buffer frame_index = self.episode_buffer["size"] - timestamp = frame.pop("timestamp") if "timestamp" in frame else frame_index / self.fps + if timestamp is None: + timestamp = frame_index / self.fps self.episode_buffer["frame_index"].append(frame_index) self.episode_buffer["timestamp"].append(timestamp) + self.episode_buffer["task"].append(task) # Add frame features to episode_buffer for key, value in frame.items(): - if key == "task": - # Note: we associate the task in natural language to its task index during `save_episode` - self.episode_buffer["task"].append(frame["task"]) - continue - if key not in self.features: raise ValueError( f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'." @@ -276,11 +258,11 @@ def save_as_lerobot_dataset(task: tuple[dict, Path, str], src_path, benchmark, e # 1. not consistent image shape... # 2. franka and ur image is bgr... bgr2rgb = False + if embodiment in ["franka_1rgb", "franka_3rgb", "franka_fr3_dual", "ur_1rgb"]: + bgr2rgb = True + if "1_0" in benchmark: match embodiment: - case "franka_1rgb" | "franka_3rgb" | "franka_fr3_dual" | "ur_1rgb": - bgr2rgb = True - case "tienkung_gello_1rgb": if task_type in ( "clean_table_2_241211", @@ -331,8 +313,7 @@ def save_as_lerobot_dataset(task: tuple[dict, Path, str], src_path, benchmark, e status, raw_dataset, err = load_local_dataset(episode_path, config, save_depth, bgr2rgb) if status and len(raw_dataset) >= 50: for frame_data in raw_dataset: - frame_data.update({"task": task_instruction}) - dataset.add_frame(frame_data) + dataset.add_frame(frame_data, task_instruction) dataset.save_episode(split, action_config.get(episode_path.parent.parent.name, {})) logging.info(f"process done for {path}, len {len(raw_dataset)}") else: