diff --git a/docs/docs/docs/features/lerobot_recorder.md b/docs/docs/docs/features/lerobot_recorder.md new file mode 100644 index 00000000..2854e84c --- /dev/null +++ b/docs/docs/docs/features/lerobot_recorder.md @@ -0,0 +1,49 @@ +# LeRobot Recorder +> LeRobot Dataset Recorder: Record LeRobot-format datasets directly during teleoperation + +LeRobot Recorder is an enhanced feature provided by LeIsaac that allows users to record data directly in LeRobot format during teleoperation. This feature seamlessly integrates into the teleoperation workflow, generating LeRobot-standard datasets without requiring additional data conversion steps. + +## How It Works + +LeRobot Recorder replaces the default recorder manager with `LeRobotRecorderManager`. After each environment step: + +1. **Data Collection**: Collects observations (including joint positions, camera images, etc.) and actions from the current step +2. **Format Conversion**: Converts data to LeRobot frame format via the `build_lerobot_frame` method +3. **Buffer Management**: Adds frame data to the LeRobot Dataset buffer +4. **Episode Processing**: When an episode ends, decides whether to save based on task success status: + - **Success**: Calls `flush()` to save the entire episode to the dataset + - **Failure**: Calls `clear()` to clear the buffer without saving data + +LeIsaac automatically skips the first 5 frames of each episode to avoid instability from initial states affecting data quality. + +## Usage + +LeRobot Recorder requires the lerobot dependency. For installation instructions, refer to this [section](../getting_started/installation#optional-install-lerobot). + +To record data in LeRobot format during teleoperation, use the following command: + +```shell +python scripts/environments/teleoperation/teleop_se3_agent.py \ + --task=LeIsaac-SO101-PickOrange-v0 \ + --teleop_device=so101leader \ + --port=/dev/ttyACM0 \ + --num_envs=1 \ + --device=cuda \ + --enable_cameras \ + --record \ + --use_lerobot_recorder \ + --lerobot_dataset_repo_id=EverNorif/test_lerobot_recorder \ + --lerobot_dataset_fps=30 +``` + +Simply enable `--use_lerobot_recorder` and specify the `repo_id` and `dataset_fps` to record data directly in LeRobot Dataset format during teleoperation. + +## Parameters + +- `--use_lerobot_recorder`: Enables the LeRobot format recorder +- `--lerobot_dataset_repo_id`: HuggingFace dataset repository ID (format: `username/repository_name`) +- `--lerobot_dataset_fps`: Dataset frame rate, typically set to 30 FPS + +::::tip +Compared to recording as hdf5, LeRobot Recorder integration may cause slight delays in teleoperation. If you encounter this issue, consider not using `--use_lerobot_recorder`. +:::: diff --git a/docs/docs/docs/getting_started/installation.md b/docs/docs/docs/getting_started/installation.md index efc8adc5..afe91123 100644 --- a/docs/docs/docs/getting_started/installation.md +++ b/docs/docs/docs/getting_started/installation.md @@ -75,6 +75,18 @@ If you are using a 50-series GPU, we recommend using IsaacSim 5.0+ and IsaacLab | PyTorch | 2.5.1 | 2.7.0 | 2.7.0 | :::: +### [Optional] Install Lerobot + +We also provide integration with LeRobot. In certain cases, you may need the lerobot dependency, such as for data conversion, lerobot dataset recorder, lerobot model inference, and envhub support. This is optional; you can install lerobot alongside leisaac when you need these features. + +```bash +# Install with lerobot +pip install -e "source/leisaac[lerobot]" + +# Fix numpy version +pip install numpy==1.26.0 +``` + ## 2. Asset Preparation We provide an example USD asset—a kitchen scene. Please download related scene [here](https://github.com/LightwheelAI/leisaac/releases/tag/v0.1.0) and extract it into the `assets` directory. The directory structure should look like this: diff --git a/docs/docs/docs/getting_started/teleoperation.md b/docs/docs/docs/getting_started/teleoperation.md index d1edfc26..24da1696 100644 --- a/docs/docs/docs/getting_started/teleoperation.md +++ b/docs/docs/docs/getting_started/teleoperation.md @@ -47,6 +47,12 @@ python scripts/environments/teleoperation/teleop_se3_agent.py \ - `--quality`: Whether to enable quality render mode. +- `--use_lerobot_recorder`: Whether to use lerobot recorder. + +- `--lerobot_dataset_repo_id`: LeRobot dataset repository ID. + +- `--lerobot_dataset_fps`: LeRobot dtaset frames per second. + ::::tip diff --git a/docs/docs/docs/tutorials/custom_task.md b/docs/docs/docs/tutorials/custom_task.md index bcdf0a88..04654bc3 100644 --- a/docs/docs/docs/tutorials/custom_task.md +++ b/docs/docs/docs/tutorials/custom_task.md @@ -119,6 +119,8 @@ class CustomTaskEnvCfg(SingleArmTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "pick up the red cube and place it into the box." + def __post_init__(self) -> None: super().__post_init__() diff --git a/docs/sidebars.js b/docs/sidebars.js index 6bd6fc4e..2f753838 100644 --- a/docs/sidebars.js +++ b/docs/sidebars.js @@ -105,6 +105,12 @@ const sidebars = { link: { type: 'doc', id: 'docs/features/envhub_support' }, items: [], }, + { + type: 'category', + label: 'LeRobot Recorder', + link: { type: 'doc', id: 'docs/features/lerobot_recorder' }, + items: [], + }, ], }, 'docs/trouble_shooting', diff --git a/scripts/environments/teleoperation/teleop_se3_agent.py b/scripts/environments/teleoperation/teleop_se3_agent.py index 38e8e093..de518f3d 100644 --- a/scripts/environments/teleoperation/teleop_se3_agent.py +++ b/scripts/environments/teleoperation/teleop_se3_agent.py @@ -11,6 +11,7 @@ if multiprocessing.get_start_method() != "spawn": multiprocessing.set_start_method("spawn", force=True) import argparse +import signal from isaaclab.app import AppLauncher @@ -64,6 +65,9 @@ parser.add_argument("--recalibrate", action="store_true", help="recalibrate SO101-Leader or Bi-SO101Leader") parser.add_argument("--quality", action="store_true", help="whether to enable quality render mode.") +parser.add_argument("--use_lerobot_recorder", action="store_true", help="whether to use lerobot recorder.") +parser.add_argument("--lerobot_dataset_repo_id", type=str, default=None, help="Lerobot Dataset repository ID.") +parser.add_argument("--lerobot_dataset_fps", type=int, default=30, help="Lerobot Dataset frames per second.") # append AppLauncher cli args AppLauncher.add_app_launcher_args(parser) @@ -180,16 +184,22 @@ def main(): # noqa: C901 env_cfg.terminations.success = None # recorder preprocess & manual success terminate preprocess if args_cli.record: - if args_cli.resume: - env_cfg.recorders.dataset_export_mode = EnhanceDatasetExportMode.EXPORT_ALL_RESUME - assert os.path.exists( - args_cli.dataset_file - ), "the dataset file does not exist, please don't use '--resume' if you want to record a new dataset" + if args_cli.use_lerobot_recorder: + if args_cli.resume: + env_cfg.recorders.dataset_export_mode = EnhanceDatasetExportMode.EXPORT_SUCCEEDED_ONLY_RESUME + else: + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY else: - env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL - assert not os.path.exists( - args_cli.dataset_file - ), "the dataset file already exists, please use '--resume' to resume recording" + if args_cli.resume: + env_cfg.recorders.dataset_export_mode = EnhanceDatasetExportMode.EXPORT_ALL_RESUME + assert os.path.exists( + args_cli.dataset_file + ), "the dataset file does not exist, please don't use '--resume' if you want to record a new dataset" + else: + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_ALL + assert not os.path.exists( + args_cli.dataset_file + ), "the dataset file already exists, please use '--resume' to resume recording" env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name if is_direct_env: @@ -205,12 +215,26 @@ def main(): # noqa: C901 # create environment env: ManagerBasedRLEnv | DirectRLEnv = gym.make(task_name, cfg=env_cfg).unwrapped - # replace the original recorder manager with the streaming recorder manager + # replace the original recorder manager with the streaming recorder manager or lerobot recorder manager if args_cli.record: del env.recorder_manager - env.recorder_manager = StreamingRecorderManager(env_cfg.recorders, env) - env.recorder_manager.flush_steps = 100 - env.recorder_manager.compression = "lzf" + if args_cli.use_lerobot_recorder: + from leisaac.enhance.datasets.lerobot_dataset_handler import ( + LeRobotDatasetCfg, + ) + from leisaac.enhance.managers.lerobot_recorder_manager import ( + LeRobotRecorderManager, + ) + + dataset_cfg = LeRobotDatasetCfg( + repo_id=args_cli.lerobot_dataset_repo_id, + fps=args_cli.lerobot_dataset_fps, + ) + env.recorder_manager = LeRobotRecorderManager(env_cfg.recorders, dataset_cfg, env) + else: + env.recorder_manager = StreamingRecorderManager(env_cfg.recorders, env) + env.recorder_manager.flush_steps = 100 + env.recorder_manager.compression = "lzf" # create controller if args_cli.teleop_device == "keyboard": @@ -284,61 +308,84 @@ def reset_task_success(): start_record_state = False - # simulate environment - while simulation_app.is_running(): - # run everything in inference mode - with torch.inference_mode(): - if env.cfg.dynamic_reset_gripper_effort_limit: - dynamic_reset_gripper_effort_limit_sim(env, args_cli.teleop_device) - actions = teleop_interface.advance() - if should_reset_task_success: - print("Task Success!!!") - should_reset_task_success = False - if args_cli.record: - manual_terminate(env, True) - if should_reset_recording_instance: - env.reset() - should_reset_recording_instance = False - if start_record_state: + interrupted = False + + def signal_handler(signum, frame): + """Handle SIGINT (Ctrl+C) signal.""" + nonlocal interrupted + interrupted = True + print("\n[INFO] KeyboardInterrupt (Ctrl+C) detected. Cleaning up resources...") + + original_sigint_handler = signal.signal(signal.SIGINT, signal_handler) + + try: + while simulation_app.is_running() and not interrupted: + # run everything in inference mode + with torch.inference_mode(): + if env.cfg.dynamic_reset_gripper_effort_limit: + dynamic_reset_gripper_effort_limit_sim(env, args_cli.teleop_device) + actions = teleop_interface.advance() + if should_reset_task_success: + print("Task Success!!!") + should_reset_task_success = False if args_cli.record: - print("Stop Recording!!!") - start_record_state = False - if args_cli.record: - manual_terminate(env, False) - # print out the current demo count if it has changed - if ( - args_cli.record - and env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count - > current_recorded_demo_count - ): - current_recorded_demo_count = ( - env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count - ) - print(f"Recorded {current_recorded_demo_count} successful demonstrations.") - if ( - args_cli.record - and args_cli.num_demos > 0 - and env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count - >= args_cli.num_demos - ): - print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") - break - - elif actions is None: - env.render() - # apply actions - else: - if not start_record_state: + manual_terminate(env, True) + if should_reset_recording_instance: + env.reset() + should_reset_recording_instance = False + if start_record_state: + if args_cli.record: + print("Stop Recording!!!") + start_record_state = False if args_cli.record: - print("Start Recording!!!") - start_record_state = True - env.step(actions) - if rate_limiter: - rate_limiter.sleep(env) - - # close the simulator - env.close() - simulation_app.close() + manual_terminate(env, False) + # print out the current demo count if it has changed + if ( + args_cli.record + and env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count + > current_recorded_demo_count + ): + current_recorded_demo_count = ( + env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count + ) + print(f"Recorded {current_recorded_demo_count} successful demonstrations.") + if ( + args_cli.record + and args_cli.num_demos > 0 + and env.recorder_manager.exported_successful_episode_count + resume_recorded_demo_count + >= args_cli.num_demos + ): + print(f"All {args_cli.num_demos} demonstrations recorded. Exiting the app.") + break + + elif actions is None: + env.render() + # apply actions + else: + if not start_record_state: + if args_cli.record: + print("Start Recording!!!") + start_record_state = True + env.step(actions) + if rate_limiter: + rate_limiter.sleep(env) + if interrupted: + break + except Exception as e: + import traceback + + print(f"\n[ERROR] An error occurred: {e}\n") + traceback.print_exc() + print("[INFO] Cleaning up resources...") + finally: + # Restore original signal handler + signal.signal(signal.SIGINT, original_sigint_handler) + # finalize the recorder manager + if args_cli.record and hasattr(env.recorder_manager, "finalize"): + env.recorder_manager.finalize() + # close the simulator + env.close() + simulation_app.close() if __name__ == "__main__": diff --git a/source/leisaac/leisaac/devices/lerobot/bi_so101_leader.py b/source/leisaac/leisaac/devices/lerobot/bi_so101_leader.py index c26702e2..4b49d961 100644 --- a/source/leisaac/leisaac/devices/lerobot/bi_so101_leader.py +++ b/source/leisaac/leisaac/devices/lerobot/bi_so101_leader.py @@ -10,9 +10,9 @@ def __init__( # use left so101 leader as the main device to store state print("Connecting to left_so101_leader...") - self.left_so101_leader = SO101Leader(env, left_port, recalibrate, "left_so101_leader.json", verbose=False) + self.left_so101_leader = SO101Leader(env, left_port, recalibrate, "left_so101_leader.json") print("Connecting to right_so101_leader...") - self.right_so101_leader = SO101Leader(env, right_port, recalibrate, "right_so101_leader.json", verbose=False) + self.right_so101_leader = SO101Leader(env, right_port, recalibrate, "right_so101_leader.json") self.left_so101_leader._stop_keyboard_listener() self.right_so101_leader._stop_keyboard_listener() diff --git a/source/leisaac/leisaac/enhance/datasets/lerobot_dataset_handler.py b/source/leisaac/leisaac/enhance/datasets/lerobot_dataset_handler.py new file mode 100644 index 00000000..82d75c90 --- /dev/null +++ b/source/leisaac/leisaac/enhance/datasets/lerobot_dataset_handler.py @@ -0,0 +1,77 @@ +import copy + +from isaaclab.utils import configclass +from isaaclab.utils.datasets.dataset_file_handler_base import DatasetFileHandlerBase +from isaaclab.utils.datasets.episode_data import EpisodeData +from lerobot.datasets.lerobot_dataset import LeRobotDataset + + +@configclass +class LeRobotDatasetCfg: + """Configuration for the LeRobotDataset.""" + + repo_id: str = None + """Lerobot Dataset repository ID.""" + fps: int = 30 + """Lerobot Dataset frames per second.""" + robot_type: str = "so101_follower" + """Robot type: so101_follower or bi_so101_follower, etc.""" + features: dict = None + """Features for the LeRobotDataset.""" + + +class LeRobotDatasetHandler(DatasetFileHandlerBase): + def __init__(self, cfg: LeRobotDatasetCfg): + self._cfg = copy.deepcopy(cfg) + self._lerobot_dataset = None + self._demo_count = 0 + self._env_args = {} + + def create(self, file_path: str, env_name: str = None, resume: bool = False): + if resume: + self._lerobot_dataset = LeRobotDataset( + repo_id=self._cfg.repo_id, + ) + else: + self._lerobot_dataset = LeRobotDataset.create( + repo_id=self._cfg.repo_id, + fps=self._cfg.fps, + robot_type=self._cfg.robot_type, + features=self._cfg.features, + ) + self._env_args["env_name"] = env_name + + def open(self, file_path: str, mode: str = "r"): + self._lerobot_dataset = LeRobotDataset( + repo_id=self._cfg.repo_id, + ) + + def get_env_name(self) -> str | None: + return self._env_args["env_name"] + + def add_frame(self, frame: dict): + self._lerobot_dataset.add_frame(frame=frame) + + def flush(self): + self._lerobot_dataset.save_episode(parallel_encoding=False) + + def clear(self): + self._lerobot_dataset.clear_episode_buffer() + + def finalize(self): + self._lerobot_dataset.finalize() + + def close(self): + if self._lerobot_dataset is not None: + self.finalize() + self._lerobot_dataset = None + + # not used for now + def write_episode(self, episode: EpisodeData): + raise NotImplementedError("write_episode is not supported for LeRobotDatasetHandler") + + def load_episode(self, episode_name: str) -> EpisodeData | None: + raise NotImplementedError("load_episode is not supported for LeRobotDatasetHandler") + + def get_num_episodes(self) -> int: + raise NotImplementedError("get_num_episodes is not supported for LeRobotDatasetHandler") diff --git a/source/leisaac/leisaac/enhance/managers/lerobot_recorder_manager.py b/source/leisaac/leisaac/enhance/managers/lerobot_recorder_manager.py new file mode 100644 index 00000000..60c24724 --- /dev/null +++ b/source/leisaac/leisaac/enhance/managers/lerobot_recorder_manager.py @@ -0,0 +1,124 @@ +from collections.abc import Sequence + +import torch +from isaaclab.envs import ManagerBasedEnv +from isaaclab.managers import DatasetExportMode, RecorderManager +from isaaclab.utils.datasets.episode_data import EpisodeData +from leisaac.utils.robot_utils import build_feature_from_env + +from ..datasets.lerobot_dataset_handler import LeRobotDatasetCfg, LeRobotDatasetHandler +from .recorder_manager import EnhanceDatasetExportMode + + +class LeRobotRecorderManager(RecorderManager): + def __init__(self, cfg: object, dataset_cfg: LeRobotDatasetCfg, env: ManagerBasedEnv) -> None: + real_export_mode = cfg.dataset_export_mode + cfg.dataset_export_mode = DatasetExportMode.EXPORT_NONE + super().__init__(cfg, env) + cfg.dataset_export_mode = real_export_mode + + assert cfg.dataset_export_mode in [ + DatasetExportMode.EXPORT_SUCCEEDED_ONLY, + EnhanceDatasetExportMode.EXPORT_SUCCEEDED_ONLY_RESUME, + ], "only support EXPORT_SUCCEEDED_ONLY|EXPORT_SUCCEEDED_ONLY_RESUME" + + cfg.dataset_file_handler_class_type = LeRobotDatasetHandler + + resume = cfg.dataset_export_mode == EnhanceDatasetExportMode.EXPORT_SUCCEEDED_ONLY_RESUME + + dataset_cfg.robot_type = env.cfg.robot_name + dataset_cfg.features = build_feature_from_env(env, dataset_cfg) + + self._dataset_cfg = dataset_cfg + + self._dataset_file_handler = cfg.dataset_file_handler_class_type(dataset_cfg) + self._dataset_file_handler.create(None, resume=resume) + + self._skip_frames = 5 + self._env_steps_record = torch.zeros(self._env.num_envs) + + def __str__(self) -> str: + msg = "[Enhanced] LeRobotRecorderManager. \n" + msg += super().__str__() + return msg + + def finalize(self): + self._dataset_file_handler.finalize() + + def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + self._env_steps_record[env_ids] = 0 + return super().reset(env_ids) + + def record_post_step(self) -> None: + """add frame to lerobot dataset after record_post_step""" + super().record_post_step() + + env_idx = 0 + self._env_steps_record[env_idx] += 1 + if self._env_steps_record[env_idx] <= self._skip_frames: + return + frame = self._env.cfg.build_lerobot_frame(self._episodes[env_idx], self._dataset_cfg.features) + self._dataset_file_handler.add_frame(frame) + self._episodes[env_idx]._data.clear() + + def export_episodes(self, env_ids: Sequence[int] | None = None) -> None: + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + # Export episode data through dataset exporter + for env_id in env_ids: + if env_id in self._episodes: + episode_succeeded = self._episodes[env_id].success + target_dataset_file_handler = self._dataset_file_handler + if episode_succeeded: + target_dataset_file_handler.flush() + self._exported_successful_episode_count[env_id] = ( + self._exported_successful_episode_count.get(env_id, 0) + 1 + ) + else: + target_dataset_file_handler.clear() + self._exported_failed_episode_count[env_id] = ( + self._exported_failed_episode_count.get(env_id, -1) + 1 + ) # default to -1 to handle the first reset + # Reset the episode buffer for the given environment after export + self._episodes[env_id] = EpisodeData() + + def record_pre_reset(self, env_ids: Sequence[int] | None, force_export_or_skip=None) -> None: + """ + Modified from super().record_pre_reset() with additional logic to retrieve success status values from _get_dones() + to adapt to RecorderEnhanceDirectRLEnv. + """ + # Do nothing if no active recorder terms are provided + if len(self.active_terms) == 0: + return + + if env_ids is None: + env_ids = list(range(self._env.num_envs)) + if isinstance(env_ids, torch.Tensor): + env_ids = env_ids.tolist() + + for term in self._terms.values(): + key, value = term.record_pre_reset(env_ids) + self.add_to_episodes(key, value, env_ids) + + # Set task success values for the relevant episodes + success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device) + # Check success indicator from termination terms + if hasattr(self._env, "termination_manager"): # for ManagerBasedEnv + if "success" in self._env.termination_manager.active_terms: + success_results |= self._env.termination_manager.get_term("success")[env_ids] + elif hasattr(self._env, "_get_dones"): # for DriectEnv + done, _ = self._env._get_dones() + success_results |= done[env_ids] + self.set_success_to_episodes(env_ids, success_results) + + if force_export_or_skip or (force_export_or_skip is None and self.cfg.export_in_record_pre_reset): + self.export_episodes(env_ids) diff --git a/source/leisaac/leisaac/enhance/managers/recorder_manager.py b/source/leisaac/leisaac/enhance/managers/recorder_manager.py index c12b21ac..5a5fa340 100644 --- a/source/leisaac/leisaac/enhance/managers/recorder_manager.py +++ b/source/leisaac/leisaac/enhance/managers/recorder_manager.py @@ -28,6 +28,8 @@ class EnhanceDatasetExportMode(enum.IntEnum): EXPORT_ALL_RESUME = 0 # Export all episodes to a single dataset file and resume recording # NOTE: the exact value is same as DatasetExportMode.EXPORT_NONE, we don't support DatasetExportMode.EXPORT_NONE when recording + EXPORT_SUCCEEDED_ONLY_RESUME = 4 # Export only succeeded episodes to a single dataset file and resume recording + # NOTE: the exact value is different from DatasetExportMode, which contains 0-3 class StreamingRecorderManager(RecorderManager): @@ -102,7 +104,6 @@ def export_episodes(self, env_ids: Sequence[int] | None = None, from_step: bool and (self._env_steps_record[env_id] >= self._flush_steps or not from_step) ): # NOTE: pre_export() is only available in IsaacLab 2.3.0+ - # TODO: remove this after we use IsaacLab 2.3.0+ by default if _AFTER_ISAACLAB_2_3_0: self._episodes[env_id].pre_export() if self._env.cfg.seed is not None: diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py index a894800f..fb31263f 100644 --- a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_bi_arm_env_cfg.py @@ -31,6 +31,8 @@ class AssembleHamburgerBiArmEnvCfg(BiArmTaskEnvCfg): terminations: BiArmTerminationsCfg = BiArmTerminationsCfg() + task_description: str = "Assemble the hamburger in order." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py index 7c4fc70a..98ab50a4 100644 --- a/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py +++ b/source/leisaac/leisaac/tasks/assemble_hamburger/assemble_hamhurger_env_cfg.py @@ -31,6 +31,8 @@ class AssembleHamburgerEnvCfg(SingleArmTaskEnvCfg): terminations: SingleArmTerminationsCfg = SingleArmTerminationsCfg() + task_description: str = "Assemble the hamburger in order." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_bi_arm_env_cfg.py index 4a0b064f..a79375fc 100644 --- a/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_bi_arm_env_cfg.py +++ b/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_bi_arm_env_cfg.py @@ -46,6 +46,8 @@ class CleanToyTableBiArmEnvCfg(BiArmTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "Pick two letter e objects into the box, and reset the arm to rest state." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_env_cfg.py b/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_env_cfg.py index fb00e817..639b6593 100644 --- a/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_env_cfg.py +++ b/source/leisaac/leisaac/tasks/clean_toy_table/clean_toy_table_env_cfg.py @@ -46,6 +46,8 @@ class CleanToyTableEnvCfg(SingleArmTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "Pick two letter e objects into the box, and reset the arm to rest state." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/cleanup_trash/cleanup_trash_env_cfg.py b/source/leisaac/leisaac/tasks/cleanup_trash/cleanup_trash_env_cfg.py index 86bed89a..95c84fdd 100644 --- a/source/leisaac/leisaac/tasks/cleanup_trash/cleanup_trash_env_cfg.py +++ b/source/leisaac/leisaac/tasks/cleanup_trash/cleanup_trash_env_cfg.py @@ -50,6 +50,8 @@ class CleanupTrashEnvCfg(LeKiwiTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "Pick up tissue trash from the floor and throw it into the trash bin." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/fold_cloth/fold_cloth_bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/fold_cloth/fold_cloth_bi_arm_env_cfg.py index bbc8c3f9..b5963c95 100644 --- a/source/leisaac/leisaac/tasks/fold_cloth/fold_cloth_bi_arm_env_cfg.py +++ b/source/leisaac/leisaac/tasks/fold_cloth/fold_cloth_bi_arm_env_cfg.py @@ -38,6 +38,8 @@ class FoldClothBiArmEnvCfg(BiArmTaskEnvCfg): terminations: BiArmTerminationsCfg = BiArmTerminationsCfg() + task_description: str = "Fold the cloth, and reset the arm to rest state." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/lift_cube/lift_cube_env_cfg.py b/source/leisaac/leisaac/tasks/lift_cube/lift_cube_env_cfg.py index e661f62a..05af275e 100644 --- a/source/leisaac/leisaac/tasks/lift_cube/lift_cube_env_cfg.py +++ b/source/leisaac/leisaac/tasks/lift_cube/lift_cube_env_cfg.py @@ -114,6 +114,8 @@ class LiftCubeEnvCfg(SingleArmTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "Lift the red cube up." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/pick_orange/pick_orange_env_cfg.py b/source/leisaac/leisaac/tasks/pick_orange/pick_orange_env_cfg.py index 69adb977..99647b7d 100644 --- a/source/leisaac/leisaac/tasks/pick_orange/pick_orange_env_cfg.py +++ b/source/leisaac/leisaac/tasks/pick_orange/pick_orange_env_cfg.py @@ -84,6 +84,8 @@ class PickOrangeEnvCfg(SingleArmTaskEnvCfg): terminations: TerminationsCfg = TerminationsCfg() + task_description: str = "Pick three oranges and put them into the plate, then reset the arm to rest state." + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/template/bi_arm_env_cfg.py b/source/leisaac/leisaac/tasks/template/bi_arm_env_cfg.py index e4dc7b89..19bcc770 100644 --- a/source/leisaac/leisaac/tasks/template/bi_arm_env_cfg.py +++ b/source/leisaac/leisaac/tasks/template/bi_arm_env_cfg.py @@ -16,8 +16,10 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sensors import TiledCameraCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets.episode_data import EpisodeData from leisaac.assets.robots.lerobot import SO101_FOLLOWER_CFG from leisaac.devices.action_process import init_action_cfg, preprocess_device_action +from leisaac.utils.constant import BI_ARM_JOINT_NAMES from . import mdp @@ -179,6 +181,13 @@ class BiArmTaskEnvCfg(ManagerBasedRLEnvCfg): dynamic_reset_gripper_effort_limit: bool = True """Whether to dynamically reset the gripper effort limit.""" + robot_name: str = "bi_so101_follower" + """Robot name for lerobot dataset export.""" + default_feature_joint_names: list[str] = MISSING + """Default feature joint names for lerobot dataset export.""" + task_description: str = MISSING + """Task description for lerobot dataset export.""" + def __post_init__(self) -> None: super().__post_init__() @@ -194,9 +203,28 @@ def __post_init__(self) -> None: self.scene.left_arm.init_state.pos = (3.4, -0.65, 0.89) self.scene.right_arm.init_state.pos = (3.8, -0.65, 0.89) + self.default_feature_joint_names = [f"{joint_name}.pos" for joint_name in BI_ARM_JOINT_NAMES] + def use_teleop_device(self, teleop_device) -> None: self.task_type = teleop_device self.actions = init_action_cfg(self.actions, device=teleop_device) def preprocess_device_action(self, action: dict[str, Any], teleop_device) -> torch.Tensor: return preprocess_device_action(action, teleop_device) + + def build_lerobot_frame(self, episode_data: EpisodeData, features: dict) -> dict: + obs_data = episode_data._data["obs"] + frame = { + "action": obs_data["actions"][-1].cpu().numpy(), + "observation.state": ( + torch.cat([obs_data["left_joint_pos"][-1], obs_data["right_joint_pos"][-1]]).cpu().numpy() + ), + "task": self.task_description, + } + for frame_key in features.keys(): + if not frame_key.startswith("observation.images"): + continue + camera_key = frame_key.split(".")[-1] + frame[frame_key] = obs_data[camera_key][-1].cpu().numpy() + + return frame diff --git a/source/leisaac/leisaac/tasks/template/direct/bi_arm_env.py b/source/leisaac/leisaac/tasks/template/direct/bi_arm_env.py index f35f8669..adb3794a 100644 --- a/source/leisaac/leisaac/tasks/template/direct/bi_arm_env.py +++ b/source/leisaac/leisaac/tasks/template/direct/bi_arm_env.py @@ -53,6 +53,8 @@ class BiArmTaskDirectEnvCfg(DirectRLEnvCfg): dynamic_reset_gripper_effort_limit: bool = True """Whether to dynamically reset the gripper effort limit.""" + # TODO: add lerobot dataset export support + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/template/direct/single_arm_env.py b/source/leisaac/leisaac/tasks/template/direct/single_arm_env.py index 662f3791..0d303553 100644 --- a/source/leisaac/leisaac/tasks/template/direct/single_arm_env.py +++ b/source/leisaac/leisaac/tasks/template/direct/single_arm_env.py @@ -47,6 +47,8 @@ class SingleArmTaskDirectEnvCfg(DirectRLEnvCfg): dynamic_reset_gripper_effort_limit: bool = True """Whether to dynamically reset the gripper effort limit.""" + # TODO: add lerobot dataset export support + def __post_init__(self) -> None: super().__post_init__() diff --git a/source/leisaac/leisaac/tasks/template/lekiwi_env_cfg.py b/source/leisaac/leisaac/tasks/template/lekiwi_env_cfg.py index 56884961..61b90d62 100644 --- a/source/leisaac/leisaac/tasks/template/lekiwi_env_cfg.py +++ b/source/leisaac/leisaac/tasks/template/lekiwi_env_cfg.py @@ -16,8 +16,10 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sensors import FrameTransformerCfg, OffsetCfg, TiledCameraCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets.episode_data import EpisodeData from leisaac.assets.robots.lerobot import LEKIWI_CFG from leisaac.devices.action_process import init_action_cfg, preprocess_device_action +from leisaac.utils.constant import LEKIWI_JOINT_NAMES from . import mdp @@ -178,6 +180,13 @@ class LeKiwiTaskEnvCfg(ManagerBasedRLEnvCfg): dynamic_reset_gripper_effort_limit: bool = True """Whether to dynamically reset the gripper effort limit.""" + robot_name: str = "lekiwi" + """Robot name for lerobot dataset export.""" + default_feature_joint_names: list[str] = MISSING + """Default feature joint names for lerobot dataset export.""" + task_description: str = MISSING + """Task description for lerobot dataset export.""" + def __post_init__(self) -> None: super().__post_init__() @@ -192,9 +201,30 @@ def __post_init__(self) -> None: self.scene.ee_frame.visualizer_cfg.markers["frame"].scale = (0.05, 0.05, 0.05) + self.default_feature_joint_names = [f"{joint_name}.pos" for joint_name in LEKIWI_JOINT_NAMES[:-3]] + [ + f"{joint_name}.vel" for joint_name in LEKIWI_JOINT_NAMES[-3:] + ] + def use_teleop_device(self, teleop_device) -> None: self.task_type = teleop_device self.actions = init_action_cfg(self.actions, device=teleop_device) def preprocess_device_action(self, action: dict[str, Any], teleop_device) -> torch.Tensor: return preprocess_device_action(action, teleop_device) + + def build_lerobot_frame(self, episode_data: EpisodeData, features: dict) -> dict: + obs_data = episode_data._data["obs"] + frame = { + "action": obs_data["actions"][-1].cpu().numpy(), + "observation.state": ( + torch.cat([obs_data["joint_pos"][-1][:-3], obs_data["user_vel_cmd"][-1]]).cpu().numpy() + ), + "task": self.task_description, + } + for frame_key in features.keys(): + if not frame_key.startswith("observation.images"): + continue + camera_key = frame_key.split(".")[-1] + frame[frame_key] = obs_data[camera_key][-1].cpu().numpy() + + return frame diff --git a/source/leisaac/leisaac/tasks/template/single_arm_env_cfg.py b/source/leisaac/leisaac/tasks/template/single_arm_env_cfg.py index e1fac90e..e2dfb95c 100644 --- a/source/leisaac/leisaac/tasks/template/single_arm_env_cfg.py +++ b/source/leisaac/leisaac/tasks/template/single_arm_env_cfg.py @@ -16,8 +16,10 @@ from isaaclab.scene import InteractiveSceneCfg from isaaclab.sensors import FrameTransformerCfg, OffsetCfg, TiledCameraCfg from isaaclab.utils import configclass +from isaaclab.utils.datasets.episode_data import EpisodeData from leisaac.assets.robots.lerobot import SO101_FOLLOWER_CFG from leisaac.devices.action_process import init_action_cfg, preprocess_device_action +from leisaac.utils.constant import SINGLE_ARM_JOINT_NAMES from . import mdp @@ -164,6 +166,13 @@ class SingleArmTaskEnvCfg(ManagerBasedRLEnvCfg): dynamic_reset_gripper_effort_limit: bool = True """Whether to dynamically reset the gripper effort limit.""" + robot_name: str = "so101_follower" + """Robot name for lerobot dataset export.""" + default_feature_joint_names: list[str] = MISSING + """Default feature joint names for lerobot dataset export.""" + task_description: str = MISSING + """Task description for lerobot dataset export.""" + def __post_init__(self) -> None: super().__post_init__() @@ -178,6 +187,8 @@ def __post_init__(self) -> None: self.scene.ee_frame.visualizer_cfg.markers["frame"].scale = (0.05, 0.05, 0.05) + self.default_feature_joint_names = [f"{joint_name}.pos" for joint_name in SINGLE_ARM_JOINT_NAMES] + def use_teleop_device(self, teleop_device) -> None: self.task_type = teleop_device self.actions = init_action_cfg(self.actions, device=teleop_device) @@ -186,3 +197,18 @@ def use_teleop_device(self, teleop_device) -> None: def preprocess_device_action(self, action: dict[str, Any], teleop_device) -> torch.Tensor: return preprocess_device_action(action, teleop_device) + + def build_lerobot_frame(self, episode_data: EpisodeData, features: dict) -> dict: + obs_data = episode_data._data["obs"] + frame = { + "action": obs_data["actions"][-1].cpu().numpy(), + "observation.state": obs_data["joint_pos"][-1].cpu().numpy(), + "task": self.task_description, + } + for frame_key in features.keys(): + if not frame_key.startswith("observation.images"): + continue + camera_key = frame_key.split(".")[-1] + frame[frame_key] = obs_data[camera_key][-1].cpu().numpy() + + return frame diff --git a/source/leisaac/leisaac/utils/constant.py b/source/leisaac/leisaac/utils/constant.py index dde62446..2390354d 100644 --- a/source/leisaac/leisaac/utils/constant.py +++ b/source/leisaac/leisaac/utils/constant.py @@ -39,3 +39,14 @@ def _resolve_assets_root() -> str: "right_wrist_roll", "right_gripper", ] +LEKIWI_JOINT_NAMES = [ + "shoulder_pan", + "shoulder_lift", + "elbow_flex", + "wrist_flex", + "wrist_roll", + "gripper", + "x", + "y", + "theta", +] diff --git a/source/leisaac/leisaac/utils/robot_utils.py b/source/leisaac/leisaac/utils/robot_utils.py index ddab74d2..2420f16d 100644 --- a/source/leisaac/leisaac/utils/robot_utils.py +++ b/source/leisaac/leisaac/utils/robot_utils.py @@ -1,10 +1,75 @@ +from dataclasses import asdict, dataclass, field + import numpy as np import torch +from isaaclab.envs import ManagerBasedEnv +from isaaclab.sensors import Camera from leisaac.assets.robots.lerobot import ( SO101_FOLLOWER_MOTOR_LIMITS, SO101_FOLLOWER_REST_POSE_RANGE, SO101_FOLLOWER_USD_JOINT_LIMLITS, ) +from leisaac.enhance.datasets.lerobot_dataset_handler import LeRobotDatasetCfg + + +@dataclass +class StateFeatureItem: + dtype: str = "float32" + shape: tuple = (6,) + names: list[str] = field( + default_factory=lambda: ["joint1.pos", "joint2.pos", "joint3.pos", "joint4.pos", "joint5.pos", "joint6.pos"] + ) + + +@dataclass +class VideoFeatureItem: + dtype: str = "video" + shape: list = field(default_factory=lambda: [480, 640, 3]) # [h, w, c] + names: list[str] = field(default_factory=lambda: ["height", "width", "channels"]) + video_info: dict = field( + default_factory=lambda: { + "video.height": 480, + "video.width": 640, + "video.codec": "av1", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "video.fps": 30.0, + "video.channels": 3, + "has_audio": False, + } + ) + + +def build_feature_from_env(env: ManagerBasedEnv, dataset_cfg: LeRobotDatasetCfg) -> dict: + """ + Build the feature from the environment. + """ + features = {} + + default_feature_joint_names = env.cfg.default_feature_joint_names + action_dim = env.action_manager.total_action_dim + if action_dim != len(default_feature_joint_names): + # [A bit tricky, currently works because the action dimension matches the joints only when we use leader control] + action_joint_names = [f"dim_{index}" for index in range(action_dim)] + else: + action_joint_names = default_feature_joint_names + features["action"] = asdict(StateFeatureItem(dtype="float32", shape=(action_dim,), names=action_joint_names)) + features["observation.state"] = asdict( + StateFeatureItem(dtype="float32", shape=(len(default_feature_joint_names),), names=default_feature_joint_names) + ) + + for camera_key, camera_sensor in env.scene.sensors.items(): + if isinstance(camera_sensor, Camera): + height, width = camera_sensor.image_shape + video_feature_item = VideoFeatureItem( + dtype="video", shape=[height, width, 3], names=["height", "width", "channels"] + ) + video_feature_item.video_info["video.height"] = height + video_feature_item.video_info["video.width"] = width + video_feature_item.video_info["video.fps"] = dataset_cfg.fps + features[f"observation.images.{camera_key}"] = asdict(video_feature_item) + + return features def is_so101_at_rest_pose(joint_pos: torch.Tensor, joint_names: list[str]) -> torch.Tensor: diff --git a/source/leisaac/pyproject.toml b/source/leisaac/pyproject.toml index 39444757..a7566649 100644 --- a/source/leisaac/pyproject.toml +++ b/source/leisaac/pyproject.toml @@ -42,4 +42,5 @@ zip-safe = false isaaclab = ["isaaclab[isaacsim,all]==2.3.0"] gr00t = ["pyzmq>=27.0.0", "pydantic==2.10.6", "msgpack>=1.0.5"] lerobot-async = ["grpcio==1.74.0", "protobuf==6.32.0"] # original isaaclab install protobuf version is 3.20.3 +lerobot = ["lerobot==0.4.2"] openpi = ["dm-tree>=0.1.8", "msgpack>=1.0.5", "numpy>=1.22.4,<2.0.0", "pillow>=9.0.0", "tree>=0.2.4", "websockets>=11.0"]