From 451c11ecb39500633ee5dc86300a6fba983e2de8 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Tue, 24 Jun 2025 08:51:07 -0700 Subject: [PATCH 01/14] so-100 prototype --- docs/source/how-to/record_lerobot_datasets.md | 230 ++++++++++++++++++ .../isaaclab_mimic/annotate_demos.py | 1 + .../isaaclab_mimic/generate_dataset.py | 37 ++- scripts/tools/record_demos.py | 92 ++++++- .../isaaclab/isaaclab/controllers/pink_ik.py | 3 + .../isaaclab/envs/manager_based_env.py | 7 +- .../mdp/actions/pink_task_space_actions.py | 5 +- .../isaaclab/managers/recorder_manager.py | 4 +- .../isaaclab/utils/datasets/__init__.py | 6 + .../datasets/dataset_file_handler_base.py | 2 +- .../datasets/hdf5_dataset_file_handler.py | 7 +- .../test_lerobot_dataset_file_handler.py | 121 +++++++++ .../isaaclab_assets/robots/__init__.py | 1 + .../isaaclab_assets/robots/so_100.py | 81 ++++++ .../envs/pinocchio_envs/__init__.py | 11 + .../pinocchio_envs/stack_so100_mimic_env.py | 142 +++++++++++ .../stack_so100_mimic_env_cfg.py | 96 ++++++++ .../stack/config/so_100/__init__.py | 54 ++++ .../stack/config/so_100/modality.json | 48 ++++ .../config/so_100/stack_ik_rel_env_cfg.py | 50 ++++ .../config/so_100/stack_joint_pos_env_cfg.py | 159 ++++++++++++ .../so_100/stack_pink_ik_abs_env_cfg.py | 105 ++++++++ .../stack_pink_ik_abs_visuomotor_env_cfg.py | 194 +++++++++++++++ .../manipulation/stack/mdp/observations.py | 13 +- .../stack/mdp/so100_stack_events.py | 97 ++++++++ .../manipulation/stack/mdp/terminations.py | 29 ++- 26 files changed, 1554 insertions(+), 41 deletions(-) create mode 100644 docs/source/how-to/record_lerobot_datasets.md create mode 100644 source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py create mode 100644 source/isaaclab_assets/isaaclab_assets/robots/so_100.py create mode 100644 source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env.py create mode 100644 source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/modality.json create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py create mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/so100_stack_events.py diff --git a/docs/source/how-to/record_lerobot_datasets.md b/docs/source/how-to/record_lerobot_datasets.md new file mode 100644 index 00000000000..1f79b33f056 --- /dev/null +++ b/docs/source/how-to/record_lerobot_datasets.md @@ -0,0 +1,230 @@ +# Recording Datasets in LeRobot Format + +This guide explains how to record demonstration datasets in LeRobot format instead of the default HDF5 format. LeRobot format is designed for training Vision-Language-Action (VLA) models and is compatible with the [🤗 LeRobot](https://github.com/huggingface/lerobot) ecosystem. + +## Why Use LeRobot Format? + +The LeRobot format offers several advantages over traditional HDF5 storage: + +- **Standardized Format**: Compatible with the growing LeRobot ecosystem +- **Efficient Storage**: Uses MP4 videos for camera observations and Parquet files for tabular data +- **Rich Metadata**: Includes task descriptions, episode information, and standardized feature naming +- **Easy Sharing**: Can be uploaded to HuggingFace Hub for community sharing +- **Training Ready**: Direct compatibility with LeRobot training pipelines + +## Installation + +First, install the required dependencies for LeRobot format: + +```bash +pip install datasets opencv-python imageio[ffmpeg] +``` + +## Basic Usage + +### Using the LeRobot Dataset Handler + +```python +from isaaclab.utils.datasets import LeRobotDatasetFileHandler, EpisodeData + +# Create a new LeRobot dataset +handler = LeRobotDatasetFileHandler() +handler.create("./datasets/my_dataset.lerobot", env_name="Isaac-Stack-Cube-Franka-IK-Rel-v0") + +# Record episode data (example) +episode = EpisodeData() +# ... populate episode with observations and actions ... + +# Write episode to dataset +handler.write_episode(episode) +handler.flush() +handler.close() +``` + +### Modifying Recording Scripts + +To use LeRobot format instead of HDF5 in your recording scripts, modify the recorder manager configuration: + +```python +from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg, DatasetExportMode +from isaaclab.utils.datasets import LeRobotDatasetFileHandler + +# Configure recorder to use LeRobot format +recorder_cfg = RecorderManagerBaseCfg() +recorder_cfg.dataset_file_handler_class_type = LeRobotDatasetFileHandler +recorder_cfg.dataset_export_dir_path = "./datasets" +recorder_cfg.dataset_filename = "my_lerobot_dataset" +recorder_cfg.dataset_export_mode = DatasetExportMode.EXPORT_ALL +``` + +## Dataset Structure + +LeRobot datasets follow a standardized structure: + +``` +dataset.lerobot/ +├── dataset_info.json # HuggingFace dataset metadata +├── state.json # Dataset state information +├── data/ # Parquet files with episode data +│ ├── train-00000-of-00001.parquet +│ └── ... +├── videos/ # Video files for camera observations +│ ├── episode_000000/ +│ │ ├── front.mp4 +│ │ ├── wrist.mp4 +│ │ └── ... +│ └── episode_000001/ +│ └── ... +└── meta/ # Additional metadata + └── info.json # Isaac Lab specific metadata +``` + +## Feature Naming Conventions + +LeRobot uses standardized naming conventions for observations: + +- **Camera observations**: `observation.images.{camera_position}` + - Examples: `observation.images.front`, `observation.images.wrist`, `observation.images.top` +- **Robot state**: `observation.state` +- **Actions**: `action` +- **Episode metadata**: `episode_index`, `frame_index`, `timestamp`, `task` + +## Using the Recording Script + +The easiest way to record demonstrations in LeRobot format is using the built-in recording script: + +```bash +# Record in LeRobot format with SpaceMouse +./isaaclab.sh -p scripts/tools/record_demos.py \ + --task Isaac-Stack-Cube-Franka-IK-Rel-v0 \ + --teleop_device spacemouse \ + --dataset_file ./datasets/my_dataset.lerobot \ + --lerobot_format \ + --num_demos 10 + +# Record in LeRobot format with hand tracking +./isaaclab.sh -p scripts/tools/record_demos.py \ + --task Isaac-Stack-Cube-Franka-IK-Rel-v0 \ + --teleop_device handtracking \ + --lerobot_format \ + --num_demos 5 + +# Record in LeRobot format with keyboard +./isaaclab.sh -p scripts/tools/record_demos.py \ + --task Isaac-Stack-Cube-Franka-IK-Rel-v0 \ + --teleop_device keyboard \ + --lerobot_format \ + --num_demos 3 +``` + +The script automatically handles: +- Dependency checking for LeRobot format +- File extension conversion (.hdf5 → .lerobot) +- Recorder manager configuration +- All teleoperation devices (keyboard, spacemouse, handtracking) + +## Programmatic Usage + +For custom integration into your own scripts: + +```python +import gymnasium as gym +from isaaclab.utils.datasets import LeRobotDatasetFileHandler +from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg +from isaaclab.managers import DatasetExportMode + +def setup_lerobot_recording(env_cfg): + """Configure environment for LeRobot recording.""" + + # Configure recorder for LeRobot format + env_cfg.recorders = ActionStateRecorderManagerCfg() + env_cfg.recorders.dataset_export_dir_path = "./datasets" + env_cfg.recorders.dataset_filename = "my_lerobot_dataset" + env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY + env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler + + return env_cfg + +# Use in your environment +env_cfg = setup_lerobot_recording(env_cfg) +env = gym.make("Isaac-Stack-Cube-Franka-IK-Rel-v0", cfg=env_cfg) +``` + +## Uploading to HuggingFace Hub + +Once you've recorded your dataset, you can share it with the community: + +```bash +# Install HuggingFace CLI +pip install huggingface_hub + +# Login to HuggingFace +huggingface-cli login + +# Upload dataset +huggingface-cli upload / ./datasets/my_dataset.lerobot +``` + +## Training with LeRobot + +After uploading your dataset, you can train models using the LeRobot pipeline: + +```bash +# Install LeRobot +pip install lerobot + +# Train a model +python lerobot/scripts/train.py \ + --dataset-repo-id / \ + --policy-name diffusion \ + --env-name +``` + +## Comparison with HDF5 Format + +| Feature | HDF5 Format | LeRobot Format | +|---------|-------------|----------------| +| Storage | Single HDF5 file | Directory with Parquet + MP4 | +| Video compression | Raw arrays | Efficient MP4 encoding | +| Metadata | Basic | Rich metadata with task descriptions | +| Sharing | Manual file transfer | HuggingFace Hub integration | +| Training compatibility | Isaac Lab only | LeRobot ecosystem | +| File size | Larger | Smaller (compressed videos) | + +## Best Practices + +1. **Use descriptive task names**: Provide clear task descriptions that will help with training +2. **Consistent camera naming**: Follow LeRobot conventions for camera positions +3. **Quality data**: Ensure clean demonstrations with clear success/failure labels +4. **Metadata**: Include rich metadata about the task, environment, and recording conditions +5. **Version control**: Use semantic versioning for dataset uploads + +## Troubleshooting + +### Common Issues + +1. **Import Error**: Ensure LeRobot dependencies are installed + ```bash + pip install datasets opencv-python imageio[ffmpeg] + ``` + +2. **Video encoding issues**: Make sure FFmpeg is properly installed + ```bash + # On Ubuntu/Debian + sudo apt-get install ffmpeg + + # On macOS + brew install ffmpeg + ``` + +3. **Large file sizes**: Videos are compressed by default, but you can adjust quality settings in the handler + +### Getting Help + +- [LeRobot Documentation](https://huggingface.co/docs/lerobot) +- [Isaac Lab Documentation](https://isaac-sim.github.io/IsaacLab/) +- [HuggingFace Datasets Documentation](https://huggingface.co/docs/datasets/) + +--- + +The LeRobot format enables seamless integration with the broader robotics AI ecosystem while maintaining compatibility with Isaac Lab's recording infrastructure. This makes it easier to share datasets, collaborate with the community, and leverage state-of-the-art VLA models for your robotics applications. \ No newline at end of file diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index b40d898a93f..d047d4f1ea0 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -313,6 +313,7 @@ def replay_episode( continue action_tensor = torch.Tensor(action).reshape([1, action.shape[0]]) env.step(torch.Tensor(action_tensor)) + # input("Press Enter to continue...") if success_term is not None: if not bool(success_term.func(env, **success_term.params)[0]): return False diff --git a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py index 05ef9af49ee..f52d3468660 100644 --- a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py +++ b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py @@ -26,7 +26,7 @@ "--output_file", type=str, default="./datasets/output_dataset.hdf5", - help="File path to export recorded and generated episodes.", + help="File path to export recorded and generated episodes. Format determined by extension: .hdf5 or .lerobot", ) parser.add_argument( "--pause_subtask", @@ -63,6 +63,7 @@ import torch import omni +import omni.log from isaaclab.envs import ManagerBasedRLMimicEnv @@ -75,6 +76,14 @@ import isaaclab_tasks # noqa: F401 +# Import dataset handlers +from isaaclab.utils.datasets import HDF5DatasetFileHandler +try: + from isaaclab.utils.datasets import LeRobotDatasetFileHandler + LEROBOT_AVAILABLE = True +except ImportError: + LEROBOT_AVAILABLE = False + def main(): num_envs = args_cli.num_envs @@ -86,6 +95,24 @@ def main(): task_name = args_cli.task.split(":")[-1] env_name = task_name or get_env_name_from_dataset(args_cli.input_file) + # Configure dataset format based on file extension + use_lerobot_format = args_cli.output_file.endswith('.lerobot') + + if use_lerobot_format: + if not LEROBOT_AVAILABLE: + omni.log.error("LeRobot format requested but dependencies not available.") + omni.log.error("Please install: pip install datasets opencv-python imageio[ffmpeg]") + exit(1) + + omni.log.info(f"Generating dataset in LeRobot format: {args_cli.output_file}") + omni.log.info("LeRobot format benefits:") + omni.log.info(" • Compatible with HuggingFace LeRobot ecosystem") + omni.log.info(" • Efficient video compression with MP4") + omni.log.info(" • Standardized naming conventions") + omni.log.info(" • Easy sharing via HuggingFace Hub") + else: + omni.log.info(f"Generating dataset in HDF5 format: {args_cli.output_file}") + # Configure environment env_cfg, success_term = setup_env_config( env_name=env_name, @@ -96,13 +123,19 @@ def main(): generation_num_trials=args_cli.generation_num_trials, ) + # Set dataset file handler based on format + if use_lerobot_format: + env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler + else: + env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler + # create environment env = gym.make(env_name, cfg=env_cfg).unwrapped if not isinstance(env, ManagerBasedRLMimicEnv): raise ValueError("The environment should be derived from ManagerBasedRLMimicEnv") - # check if the mimic API from this environment contains decprecated signatures + # check if the mimic API from this environment contains deprecated signatures if "action_noise_dict" not in inspect.signature(env.target_eef_pose_to_action).parameters: omni.log.warn( f'The "noise" parameter in the "{env_name}" environment\'s mimic API "target_eef_pose_to_action", ' diff --git a/scripts/tools/record_demos.py b/scripts/tools/record_demos.py index 0729f047614..4c36f94d24a 100644 --- a/scripts/tools/record_demos.py +++ b/scripts/tools/record_demos.py @@ -6,8 +6,10 @@ Script to record demonstrations with Isaac Lab environments using human teleoperation. This script allows users to record demonstrations operated by human teleoperation for a specified task. -The recorded demonstrations are stored as episodes in a hdf5 file. Users can specify the task, teleoperation -device, dataset directory, and environment stepping rate through command-line arguments. +The recorded demonstrations are stored as episodes in either HDF5 format or LeRobot format, determined +by the file extension of the dataset_file argument (.hdf5 for HDF5, .lerobot for LeRobot format). +Users can specify the task, teleoperation device, dataset directory, and environment stepping rate +through command-line arguments. required arguments: --task Name of the task. @@ -15,10 +17,17 @@ optional arguments: -h, --help Show this help message and exit --teleop_device Device for interacting with environment. (default: keyboard) - --dataset_file File path to export recorded demos. (default: "./datasets/dataset.hdf5") + --dataset_file File path to export recorded demos. Format determined by extension: .hdf5 or .lerobot (default: "./datasets/dataset.hdf5") --step_hz Environment stepping rate in Hz. (default: 30) --num_demos Number of demonstrations to record. (default: 0) --num_success_steps Number of continuous steps with task success for concluding a demo as successful. (default: 10) + +Examples: + # Record in HDF5 format (default) + ./isaaclab.sh -p scripts/tools/record_demos.py --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --teleop_device spacemouse + + # Record in LeRobot format for VLA training + ./isaaclab.sh -p scripts/tools/record_demos.py --task Isaac-Stack-Cube-Franka-IK-Rel-v0 --teleop_device spacemouse --dataset_file ./datasets/demo.lerobot """ """Launch Isaac Sim Simulator first.""" @@ -42,8 +51,9 @@ parser.add_argument("--task", type=str, default=None, help="Name of the task.") parser.add_argument("--teleop_device", type=str, default="keyboard", help="Device for interacting with environment.") parser.add_argument( - "--dataset_file", type=str, default="./datasets/dataset.hdf5", help="File path to export recorded demos." + "--dataset_file", type=str, default="./datasets/dataset.hdf5", help="File path to export recorded demos. Format determined by extension: .hdf5 or .lerobot" ) + parser.add_argument("--step_hz", type=int, default=30, help="Environment stepping rate in Hz.") parser.add_argument( "--num_demos", type=int, default=0, help="Number of demonstrations to record. Set to 0 for infinite." @@ -104,6 +114,14 @@ import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils.parse_cfg import parse_env_cfg +# Import dataset handlers +from isaaclab.utils.datasets import HDF5DatasetFileHandler +try: + from isaaclab.utils.datasets import LeRobotDatasetFileHandler + LEROBOT_AVAILABLE = True +except ImportError: + LEROBOT_AVAILABLE = False + class RateLimiter: """Convenience class for enforcing rates in loops.""" @@ -222,10 +240,34 @@ def main(): env_cfg.observations.policy.concatenate_terms = False - env_cfg.recorders: ActionStateRecorderManagerCfg = ActionStateRecorderManagerCfg() + # Configure dataset format based on file extension + use_lerobot_format = args_cli.dataset_file.endswith('.lerobot') + + if use_lerobot_format: + if not LEROBOT_AVAILABLE: + omni.log.error("LeRobot format requested but dependencies not available.") + omni.log.error("Please install: pip install datasets opencv-python imageio[ffmpeg]") + exit(1) + + omni.log.info(f"Recording dataset in LeRobot format: {args_cli.dataset_file}") + omni.log.info("LeRobot format benefits:") + omni.log.info(" • Compatible with HuggingFace LeRobot ecosystem") + omni.log.info(" • Efficient video compression with MP4") + omni.log.info(" • Standardized naming conventions") + omni.log.info(" • Easy sharing via HuggingFace Hub") + else: + omni.log.info(f"Recording dataset in HDF5 format: {args_cli.dataset_file}") + + env_cfg.recorders = ActionStateRecorderManagerCfg() env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY + + # Set dataset file handler based on format + if use_lerobot_format: + env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler + else: + env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler # create environment env = gym.make(args_cli.task, cfg=env_cfg).unwrapped @@ -350,7 +392,15 @@ def create_teleop_device(device_name: str, env: gym.Env): # reset before starting env.sim.reset() - env.reset() + obv, _ = env.reset() + + actions = torch.cat([obv["policy"]["eef_pos"][0], obv["policy"]["eef_quat"][0], torch.tensor([1.0], device=env.device)], dim=0).unsqueeze(0) + print("actions", actions) + # actions tensor([[ 0.2050, -0.0051, 0.2188, 0.0571, -0.0446, 0.7064, 0.7041, 1.0000]], + # device='cuda:0') + + actions[0, 2] += 0.1 + # print("actions", actions) teleop_interface.reset() # simulate environment -- run everything in inference mode @@ -377,13 +427,31 @@ def create_teleop_device(device_name: str, env: gym.Env): # perform action on environment if running_recording_instance: # compute actions based on environment - actions = pre_process_actions(teleop_data, env.num_envs, env.device) + # actions = pre_process_actions(teleop_data, env.num_envs, env.device) + delta_pose, gripper_command = teleop_data + + # print("obv", obv) + # # Convert actions[0, -1] to a 1D tensor before concatenation + # gripper_action = actions[0][-1].unsqueeze(0) + # actions = torch.cat([obv["policy"]["eef_pos"][0], obv["policy"]["eef_quat"][0], gripper_action], dim=0).unsqueeze(0) + + # print("actions", actions) + + # actions[0, 2] += 0.01 + # compute actions + control_gain = 0.1 + actions[0, 0] += delta_pose[0] * control_gain + actions[0, 1] += delta_pose[1] * control_gain + actions[0, 2] += delta_pose[2] * control_gain + actions[0, 7] = -1 if gripper_command else 1 + print("actions", actions) obv = env.step(actions) - if subtasks is not None: - if subtasks == {}: - subtasks = obv[0].get("subtask_terms") - elif subtasks: - show_subtask_instructions(instruction_display, subtasks, obv, env.cfg) + obv = obv[0] + # if subtasks is not None: + # if subtasks == {}: + # subtasks = obv[0].get("subtask_terms") + # elif subtasks: + # show_subtask_instructions(instruction_display, subtasks, obv, env.cfg) else: env.sim.render() diff --git a/source/isaaclab/isaaclab/controllers/pink_ik.py b/source/isaaclab/isaaclab/controllers/pink_ik.py index 3657fa6a0fe..561240482cc 100644 --- a/source/isaaclab/isaaclab/controllers/pink_ik.py +++ b/source/isaaclab/isaaclab/controllers/pink_ik.py @@ -105,6 +105,9 @@ def compute( # Update Pink's robot configuration with the current joint positions self.pink_configuration.update(joint_positions_pink) + # for task in self.cfg.variable_input_tasks: + # print("compute_error", task.compute_error(self.pink_configuration)) + # pink.solve_ik can raise an exception if the solver fails try: velocity = solve_ik( diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 1febf07d70a..b0faf2e76b2 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -233,15 +233,16 @@ def load_managers(self): # prepare the managers # -- event manager (we print it here to make the logging consistent) print("[INFO] Event Manager: ", self.event_manager) - # -- recorder manager - self.recorder_manager = RecorderManager(self.cfg.recorders, self) - print("[INFO] Recorder Manager: ", self.recorder_manager) # -- action manager self.action_manager = ActionManager(self.cfg.actions, self) print("[INFO] Action Manager: ", self.action_manager) # -- observation manager self.observation_manager = ObservationManager(self.cfg.observations, self) print("[INFO] Observation Manager:", self.observation_manager) + # -- recorder manager + # Recorder manager depends on the action and observation managers being created first + self.recorder_manager = RecorderManager(self.cfg.recorders, self) + print("[INFO] Recorder Manager: ", self.recorder_manager) # perform events at the start of the simulation # in-case a child implementation creates other managers, the randomization should happen diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 11c3ff6cedf..3b76f5bcfd3 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -180,6 +180,8 @@ def process_actions(self, actions: torch.Tensor): controlled_frame_in_base_link_frame ) + print("controlled_frame_in_base_link_frame_pos", controlled_frame_in_base_link_frame_pos) + # Loop through each task and set the target for env_index, ik_controller in enumerate(self._ik_controllers): for task_index, task in enumerate(ik_controller.cfg.variable_input_tasks): @@ -198,7 +200,8 @@ def apply_actions(self): all_envs_joint_pos_des.append(joint_pos_des) all_envs_joint_pos_des = torch.stack(all_envs_joint_pos_des) # Combine IK joint positions with hand joint positions - all_envs_joint_pos_des = torch.cat((all_envs_joint_pos_des, self._target_hand_joint_positions), dim=1) + # all_envs_joint_pos_des = torch.cat((all_envs_joint_pos_des, self._target_hand_joint_positions), dim=1) + # print("self._target_hand_joint_positions", self._target_hand_joint_positions) self._asset.set_joint_position_target(all_envs_joint_pos_des, self._joint_ids) diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index 4b6ba98f1e1..ea2c8101204 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -161,14 +161,14 @@ def __init__(self, cfg: object, env: ManagerBasedEnv): if cfg.dataset_export_mode != DatasetExportMode.EXPORT_NONE: self._dataset_file_handler = cfg.dataset_file_handler_class_type() self._dataset_file_handler.create( - os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name + os.path.join(cfg.dataset_export_dir_path, cfg.dataset_filename), env_name=env_name, env=env ) self._failed_episode_dataset_file_handler = None if cfg.dataset_export_mode == DatasetExportMode.EXPORT_SUCCEEDED_FAILED_IN_SEPARATE_FILES: self._failed_episode_dataset_file_handler = cfg.dataset_file_handler_class_type() self._failed_episode_dataset_file_handler.create( - os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name + os.path.join(cfg.dataset_export_dir_path, f"{cfg.dataset_filename}_failed"), env_name=env_name, env=env ) self._exported_successful_episode_count = {} diff --git a/source/isaaclab/isaaclab/utils/datasets/__init__.py b/source/isaaclab/isaaclab/utils/datasets/__init__.py index a410fa0a443..711d8daa108 100644 --- a/source/isaaclab/isaaclab/utils/datasets/__init__.py +++ b/source/isaaclab/isaaclab/utils/datasets/__init__.py @@ -15,3 +15,9 @@ from .dataset_file_handler_base import DatasetFileHandlerBase from .episode_data import EpisodeData from .hdf5_dataset_file_handler import HDF5DatasetFileHandler + +try: + from .lerobot_dataset_file_handler import LeRobotDatasetFileHandler +except ImportError: + # LeRobot dependencies not available + pass diff --git a/source/isaaclab/isaaclab/utils/datasets/dataset_file_handler_base.py b/source/isaaclab/isaaclab/utils/datasets/dataset_file_handler_base.py index dc953c0a3c6..ea626687810 100644 --- a/source/isaaclab/isaaclab/utils/datasets/dataset_file_handler_base.py +++ b/source/isaaclab/isaaclab/utils/datasets/dataset_file_handler_base.py @@ -28,7 +28,7 @@ def open(self, file_path: str, mode: str = "r"): return NotImplementedError @abstractmethod - def create(self, file_path: str, env_name: str = None): + def create(self, file_path: str, env_name: str | None = None, env = None): """Create a new file.""" return NotImplementedError diff --git a/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py index 2fa35ca1533..48cea4d5912 100644 --- a/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/hdf5_dataset_file_handler.py @@ -37,7 +37,7 @@ def open(self, file_path: str, mode: str = "r"): self._hdf5_data_group = self._hdf5_file_stream["data"] self._demo_count = len(self._hdf5_data_group) - def create(self, file_path: str, env_name: str = None): + def create(self, file_path: str, env_name: str | None = None, env = None): """Create a new dataset file.""" if self._hdf5_file_stream is not None: raise RuntimeError("HDF5 dataset file stream is already in use") @@ -107,7 +107,7 @@ def demo_count(self) -> int: def load_episode(self, episode_name: str, device: str) -> EpisodeData | None: """Load episode data from the file.""" self._raise_if_not_initialized() - if episode_name not in self._hdf5_data_group: + if self._hdf5_data_group is None or episode_name not in self._hdf5_data_group: return None episode = EpisodeData() h5_episode_group = self._hdf5_data_group[episode_name] @@ -132,7 +132,8 @@ def load_dataset_helper(group): if "success" in h5_episode_group.attrs: episode.success = h5_episode_group.attrs["success"] - episode.env_id = self.get_env_name() + # Note: env_id expects an integer, but we have environment name as string + # The environment name can be accessed via self.get_env_name() if needed return episode diff --git a/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py new file mode 100644 index 00000000000..074bff97380 --- /dev/null +++ b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py @@ -0,0 +1,121 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Test LeRobot dataset file handler functionality.""" + +import os +import shutil +import tempfile +import torch +import uuid +import pytest +from typing import TYPE_CHECKING + +from isaaclab.utils.datasets import EpisodeData + +if TYPE_CHECKING: + from isaaclab.utils.datasets import LeRobotDatasetFileHandler + +try: + from isaaclab.utils.datasets import LeRobotDatasetFileHandler # type: ignore + LEROBOT_AVAILABLE = True +except ImportError: + LEROBOT_AVAILABLE = False + + +@pytest.mark.skipif(not LEROBOT_AVAILABLE, reason="LeRobot dependencies not available") +class TestLeRobotDatasetFileHandler: + """Test LeRobot dataset file handler.""" + + def create_test_episode(self, device): + """Create a test episode with dummy data.""" + test_episode = EpisodeData() + + test_episode.seed = 42 + test_episode.success = True + + # Add some dummy observations and actions + test_episode.add("obs/joint_pos", torch.tensor([1.0, 2.0, 3.0], device=device)) + test_episode.add("obs/joint_pos", torch.tensor([1.1, 2.1, 3.1], device=device)) + test_episode.add("obs/joint_pos", torch.tensor([1.2, 2.2, 3.2], device=device)) + + test_episode.add("actions", torch.tensor([0.1, 0.2], device=device)) + test_episode.add("actions", torch.tensor([0.3, 0.4], device=device)) + test_episode.add("actions", torch.tensor([0.5, 0.6], device=device)) + + return test_episode + + @pytest.fixture + def temp_dir(self): + """Create a temporary directory for test datasets.""" + temp_dir = tempfile.mkdtemp() + yield temp_dir + # cleanup after tests + shutil.rmtree(temp_dir) + + def test_import_available(self): + """Test that LeRobot handler can be imported.""" + assert LEROBOT_AVAILABLE, "LeRobot dependencies should be available for testing" + handler = LeRobotDatasetFileHandler() + assert handler is not None + + def test_create_dataset_file(self, temp_dir): + """Test creating a new LeRobot dataset file.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + handler = LeRobotDatasetFileHandler() + + # Test creating with .lerobot extension + handler.create(dataset_file_path, "test_env_name") + assert handler.get_env_name() == "test_env_name" + handler.close() + + # Test creating without extension (should add .lerobot) + dataset_file_path_no_ext = os.path.join(temp_dir, f"{uuid.uuid4()}") + handler = LeRobotDatasetFileHandler() + handler.create(dataset_file_path_no_ext, "test_env_name") + assert handler.get_env_name() == "test_env_name" + handler.close() + + @pytest.mark.parametrize("device", ["cpu"]) # Only test CPU for CI compatibility + def test_write_episode(self, temp_dir, device): + """Test writing an episode to the LeRobot dataset.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + handler = LeRobotDatasetFileHandler() + handler.create(dataset_file_path, "test_env_name") + + test_episode = self.create_test_episode(device) + + # Write the episode to the dataset + handler.write_episode(test_episode) + assert handler.get_num_episodes() == 1 + + # Write another episode + handler.write_episode(test_episode) + assert handler.get_num_episodes() == 2 + + handler.flush() + handler.close() + + def test_get_properties(self, temp_dir): + """Test getting dataset properties.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + handler = LeRobotDatasetFileHandler() + handler.create(dataset_file_path, "test_env_name") + + # Test environment name + assert handler.get_env_name() == "test_env_name" + + # Test episode count (should be 0 for new dataset) + assert handler.get_num_episodes() == 0 + + # Test episode names (should be empty for new dataset) + episode_names = list(handler.get_episode_names()) + assert len(episode_names) == 0 + + handler.close() + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/source/isaaclab_assets/isaaclab_assets/robots/__init__.py b/source/isaaclab_assets/isaaclab_assets/robots/__init__.py index a4515156081..8220657c14c 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/__init__.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/__init__.py @@ -21,6 +21,7 @@ from .ridgeback_franka import * from .sawyer import * from .shadow_hand import * +from .so_100 import * from .spot import * from .unitree import * from .universal_robots import * diff --git a/source/isaaclab_assets/isaaclab_assets/robots/so_100.py b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py new file mode 100644 index 00000000000..f126fef00a1 --- /dev/null +++ b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py @@ -0,0 +1,81 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Configuration for the Franka Emika robots. + +The following configurations are available: + +* :obj:`SO_100_CFG`: SO-100 robot + +Reference: https://github.com/TheRobotStudio/SO-ARM100 +""" + +import isaaclab.sim as sim_utils +from isaaclab.actuators import ImplicitActuatorCfg +from isaaclab.assets.articulation import ArticulationCfg + +## +# Configuration +## + +SO_100_CFG = ArticulationCfg( + spawn=sim_utils.UsdFileCfg( + usd_path=f"omniverse://isaac-dev.ov.nvidia.com/Users/ashwinvk@nvidia.com/so-100/so-100_pure_usd.usd", + activate_contact_sensors=False, + rigid_props=sim_utils.RigidBodyPropertiesCfg( + disable_gravity=True, + retain_accelerations=False, + linear_damping=0.0, + angular_damping=0.0, + max_linear_velocity=1000.0, + max_angular_velocity=1000.0, + max_depenetration_velocity=1.0, + ), + articulation_props=sim_utils.ArticulationRootPropertiesCfg( + enabled_self_collisions=True, solver_position_iteration_count=8, solver_velocity_iteration_count=4 + ), + # collision_props=sim_utils.CollisionPropertiesCfg(contact_offset=0.005, rest_offset=0.0), + ), + init_state=ArticulationCfg.InitialStateCfg( + joint_pos={ + "a_1": 0.0, + "a_2": 0.0, + "a_3": 0.0, + "a_4": 0.0, + "a_5": 0.0, + "a_6": 0.0, + }, + ), + actuators={ + "all_joints": ImplicitActuatorCfg( + joint_names_expr=["a_[1-5]"], + effort_limit=None, + velocity_limit=None, + stiffness=None, + damping=None, + armature=0.0, + ), + "gripper": ImplicitActuatorCfg( + joint_names_expr=["a_6"], + effort_limit_sim=0.1, + velocity_limit_sim=2.175, + stiffness=8.0, + damping=0.4, + armature=0.0, + ), + }, + soft_joint_pos_limit_factor=1.0, +) +"""Configuration of Franka Emika Panda robot.""" + + +SO_100_HIGH_PD_CFG = SO_100_CFG.copy() +SO_100_HIGH_PD_CFG.spawn.rigid_props.disable_gravity = True +# SO_100_HIGH_PD_CFG.actuators["all_joints"].stiffness = 400.0 +# SO_100_HIGH_PD_CFG.actuators["all_joints"].damping = 80.0 +"""Configuration of Franka Emika Panda robot with stiffer PD control. + +This configuration is useful for task-space control using differential IK. +""" diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/__init__.py b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/__init__.py index 519f5630dac..b3976be3127 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/__init__.py +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/__init__.py @@ -14,6 +14,8 @@ from .pickplace_gr1t2_mimic_env import PickPlaceGR1T2MimicEnv from .pickplace_gr1t2_mimic_env_cfg import PickPlaceGR1T2MimicEnvCfg +from .stack_so100_mimic_env import StackSO100MimicEnv +from .stack_so100_mimic_env_cfg import StackSO100MimicEnvCfg gym.register( id="Isaac-PickPlace-GR1T2-Abs-Mimic-v0", @@ -23,3 +25,12 @@ }, disable_env_checker=True, ) + +gym.register( + id="Isaac-Stack-SO100-Abs-Mimic-v0", + entry_point="isaaclab_mimic.envs.pinocchio_envs:StackSO100MimicEnv", + kwargs={ + "env_cfg_entry_point": stack_so100_mimic_env_cfg.StackSO100MimicEnvCfg, + }, + disable_env_checker=True, +) diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env.py b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env.py new file mode 100644 index 00000000000..aec122e8152 --- /dev/null +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env.py @@ -0,0 +1,142 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + + +import torch +from collections.abc import Sequence + +import isaaclab.utils.math as PoseUtils +from isaaclab.envs import ManagerBasedRLMimicEnv + + +class StackSO100MimicEnv(ManagerBasedRLMimicEnv): + + def get_robot_eef_pose(self, eef_name: str, env_ids: Sequence[int] | None = None) -> torch.Tensor: + """ + Get current robot end effector pose. Should be the same frame as used by the robot end-effector controller. + + Args: + eef_name: Name of the end effector (for SO100, this is typically "gripper"). + env_ids: Environment indices to get the pose for. If None, all envs are considered. + + Returns: + A torch.Tensor eef pose matrix. Shape is (len(env_ids), 4, 4) + """ + if env_ids is None: + env_ids = slice(None) + + # For SO100, the end effector pose is stored in policy observations + eef_pos_name = "eef_pos" + eef_quat_name = "eef_quat" + + target_eef_position = self.obs_buf["policy"][eef_pos_name][env_ids] + target_rot_mat = PoseUtils.matrix_from_quat(self.obs_buf["policy"][eef_quat_name][env_ids]) + + return PoseUtils.make_pose(target_eef_position, target_rot_mat) + + def target_eef_pose_to_action( + self, + target_eef_pose_dict: dict, + gripper_action_dict: dict, + action_noise_dict: dict | None = None, + env_id: int = 0, + ) -> torch.Tensor: + """ + Takes a target pose and gripper action for the end effector controller and returns an action + (usually a normalized delta pose action) to try and achieve that target pose. + Noise is added to the target pose action if specified. + + Args: + target_eef_pose_dict: Dictionary containing target eef pose. For SO100, expects key "gripper". + gripper_action_dict: Dictionary containing gripper action. For SO100, expects key "gripper". + action_noise_dict: Noise to add to the action. If None, no noise is added. + env_id: Environment index to get the action for. + + Returns: + An action torch.Tensor that's compatible with env.step(). + """ + + # target position and rotation for single arm + target_eef_pos, target_rot = PoseUtils.unmake_pose(target_eef_pose_dict["gripper"]) + target_eef_rot_quat = PoseUtils.quat_from_matrix(target_rot) + + # gripper action + gripper_action = gripper_action_dict["gripper"] + + if action_noise_dict is not None: + pos_noise = action_noise_dict["gripper"] * torch.randn_like(target_eef_pos) + quat_noise = action_noise_dict["gripper"] * torch.randn_like(target_eef_rot_quat) + + target_eef_pos += pos_noise + target_eef_rot_quat += quat_noise + + return torch.cat( + ( + target_eef_pos, + target_eef_rot_quat, + gripper_action, + ), + dim=0, + ) + + def action_to_target_eef_pose(self, action: torch.Tensor) -> dict[str, torch.Tensor]: + """ + Converts action (compatible with env.step) to a target pose for the end effector controller. + Inverse of @target_eef_pose_to_action. Usually used to infer a sequence of target controller poses + from a demonstration trajectory using the recorded actions. + + Args: + action: Environment action. Shape is (num_envs, action_dim). + + Returns: + A dictionary of eef pose torch.Tensor that @action corresponds to. + """ + target_poses = {} + + # For SO100: action dimensions are [pos(3), quat(4), gripper(1)] + target_eef_position = action[:, 0:3] + target_rot_mat = PoseUtils.matrix_from_quat(action[:, 3:7]) + target_pose = PoseUtils.make_pose(target_eef_position, target_rot_mat) + target_poses["gripper"] = target_pose + + return target_poses + + def actions_to_gripper_actions(self, actions: torch.Tensor) -> dict[str, torch.Tensor]: + """ + Extracts the gripper actuation part from a sequence of env actions (compatible with env.step). + + Args: + actions: environment actions. The shape is (num_envs, num steps in a demo, action_dim). + + Returns: + A dictionary of torch.Tensor gripper actions. Key is "gripper". + """ + print("actions", actions) + print("actions", actions.shape) + return {"gripper": actions[:, -1:]} # Gripper is at index 7 (single dimension) + + def get_subtask_term_signals(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]: + """ + Gets a dictionary of termination signal flags for each subtask in a task. The flag is 1 + when the subtask has been completed and 0 otherwise. + + Args: + env_ids: Environment indices to get the termination signals for. If None, all envs are considered. + + Returns: + A dictionary termination signal flags (False or True) for each subtask. + """ + if env_ids is None: + env_ids = slice(None) + + signals = dict() + subtask_terms = self.obs_buf["subtask_terms"] + signals["grasp_1"] = subtask_terms["grasp_1"][env_ids] + return signals \ No newline at end of file diff --git a/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env_cfg.py b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env_cfg.py new file mode 100644 index 00000000000..ce10a97cbbe --- /dev/null +++ b/source/isaaclab_mimic/isaaclab_mimic/envs/pinocchio_envs/stack_so100_mimic_env_cfg.py @@ -0,0 +1,96 @@ +# Copyright (c) 2024-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +# Copyright (c) 2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig +from isaaclab.utils import configclass + +from isaaclab_tasks.manager_based.manipulation.stack.config.so_100.stack_pink_ik_abs_visuomotor_env_cfg import SO100CubeStackPinkIKAbsVisuomotorEnvCfg + + +@configclass +class StackSO100MimicEnvCfg(SO100CubeStackPinkIKAbsVisuomotorEnvCfg, MimicEnvCfg): + + def __post_init__(self): + # Calling post init of parents + super().__post_init__() + + # Override the existing values + self.datagen_config.name = "demo_src_so100_stack_task_D0" + self.datagen_config.generation_guarantee = True + self.datagen_config.generation_keep_failed = False + self.datagen_config.generation_num_trials = 1000 + self.datagen_config.generation_select_src_per_subtask = False + self.datagen_config.generation_select_src_per_arm = False + self.datagen_config.generation_relative = False + self.datagen_config.generation_joint_pos = False + self.datagen_config.generation_transform_first_robot_pose = False + self.datagen_config.generation_interpolate_from_last_target_pose = True + self.datagen_config.max_num_failures = 25 + self.datagen_config.num_demo_to_render = 10 + self.datagen_config.num_fail_demo_to_render = 25 + self.datagen_config.seed = 1 + + # The following are the subtask configurations for the stack task. + # For SO100, we only have one "gripper" arm, so all subtasks are for this single arm + subtask_configs = [] + + # Subtask 1: Pick up cube_2 + subtask_configs.append( + SubTaskConfig( + # Each subtask involves manipulation with respect to a single object frame. + # For picking up cube_2, the object of interest is cube_2 + object_ref="cube_2", + # This key corresponds to the binary indicator in "datagen_info" that signals + # when this subtask is finished (e.g., on a 0 to 1 edge). + subtask_term_signal="grasp_1", + first_subtask_start_offset_range=(0, 0), + # Randomization range for starting index of the first subtask + subtask_term_offset_range=(0, 0), + # Selection strategy for the source subtask segment during data generation + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.0, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=0, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) + + # Subtask 2: Place cube_2 on cube_1 + subtask_configs.append( + SubTaskConfig( + # For placing cube_2 on cube_1, the object of interest is cube_1 (target location) + object_ref="cube_1", + # Corresponding key for the binary indicator in "datagen_info" for completion + subtask_term_signal=None, # This is the final subtask, so no termination signal + # Time offsets for data generation when splitting a trajectory + subtask_term_offset_range=(0, 0), + # Selection strategy for source subtask segment + selection_strategy="nearest_neighbor_object", + # Optional parameters for the selection strategy function + selection_strategy_kwargs={"nn_k": 3}, + # Amount of action noise to apply during this subtask + action_noise=0.0, + # Number of interpolation steps to bridge to this subtask segment + num_interpolation_steps=3, + # Additional fixed steps for the robot to reach the necessary pose + num_fixed_steps=0, + # If True, apply action noise during the interpolation phase and execution + apply_noise_during_interpolation=False, + ) + ) + + # For SO100, we only have one arm "gripper", so we assign all subtasks to it + self.subtask_configs["gripper"] = subtask_configs \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py new file mode 100644 index 00000000000..3342c12ca61 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py @@ -0,0 +1,54 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import gymnasium as gym + +from . import stack_joint_pos_env_cfg +from . import stack_ik_rel_env_cfg +from . import stack_pink_ik_abs_env_cfg +from . import stack_pink_ik_abs_visuomotor_env_cfg +## +# Register Gym environments. +## + +## +# Joint Position Control +## + +gym.register( + id="Isaac-Stack-Cube-SO100-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_joint_pos_env_cfg.SO100CubeStackJointPosEnvCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-SO100-IK-Rel-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_ik_rel_env_cfg.SO100CubeStackIKRelEnvCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-SO100-Pink-IK-Abs-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_pink_ik_abs_env_cfg.SO100CubeStackPinkIKAbsEnvCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Isaac-Stack-Cube-SO100-Pink-IK-Abs-Visuomotor-v0", + entry_point="isaaclab.envs:ManagerBasedRLEnv", + kwargs={ + "env_cfg_entry_point": stack_pink_ik_abs_visuomotor_env_cfg.SO100CubeStackPinkIKAbsVisuomotorEnvCfg, + }, + disable_env_checker=True, +) \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/modality.json b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/modality.json new file mode 100644 index 00000000000..ae07b862b23 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/modality.json @@ -0,0 +1,48 @@ +{ + "state": { + "arm_joint_pos": { + "start": 0, + "end": 5, + "description": "SO100 arm joint positions (a_1 to a_5)" + }, + "gripper_pos": { + "start": 5, + "end": 6, + "description": "Gripper joint position" + } + }, + "action": { + "eef_pose": { + "start": 0, + "end": 7 + }, + "gripper_pos": { + "start": 7, + "end": 8 + } + }, + "video": { + "table_cam": { + "original_key": "observation.table_cam", + "description": "External table camera view of workspace" + } + }, + "annotation": { + "human.action.task_description": { + "description": "Human-provided task description" + } + }, + "metadata": { + "robot_type": "SO100", + "control_mode": "Pink_IK_Absolute", + "task": "cube_stacking", + "cameras": ["table_cam"], + "camera_resolution": [512, 512], + "dof": 6, + "workspace_bounds": { + "x": [-0.5, 0.5], + "y": [-0.5, 0.5], + "z": [0.0, 0.5] + } + } +} \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py new file mode 100644 index 00000000000..cdef31dbd8b --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py @@ -0,0 +1,50 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from isaaclab.assets import ArticulationCfg +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg +from isaaclab.utils import configclass + +from . import stack_joint_pos_env_cfg + +## +# Pre-defined configs +## +from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip + + +@configclass +class SO100CubeStackIKRelEnvCfg(stack_joint_pos_env_cfg.SO100CubeStackJointPosEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = SO_100_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + pos=(0, 0, 0.0), + # rot=(0.7071, 0, 0, 0.7071), + rot=(1.0, 0, 0, 0.0), + joint_pos={ + # right-arm + "a_1": 0.0, + "a_2": 1.5708, + "a_3": -1.5708, + "a_4": 1.2, + "a_5": 0.0, + }, + joint_vel={".*": 0.0}, + )) + + # Set actions for the specific robot type (franka) + self.actions.arm_action = DifferentialInverseKinematicsActionCfg( + asset_name="robot", + joint_names=["a_[1-4]"], + body_name="wrist", + controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), + scale=1.0, + body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py new file mode 100644 index 00000000000..08000055684 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py @@ -0,0 +1,159 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import importlib +from isaaclab.assets import ArticulationCfg, RigidObjectCfg +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import FrameTransformerCfg +from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg +from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg +from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + +from isaaclab_tasks.manager_based.manipulation.stack import mdp +from isaaclab_tasks.manager_based.manipulation.stack.stack_env_cfg import StackEnvCfg +from isaaclab_tasks.manager_based.manipulation.stack.mdp import so100_stack_events + +## +# Pre-defined configs +## +from isaaclab.markers.config import FRAME_MARKER_CFG # isort: skip +from isaaclab_assets.robots.so_100 import SO_100_CFG # isort: skip + +import torch + + +@configclass +class EventCfg: + """Configuration for events.""" + + reset_robot_joints = EventTerm( + func=mdp.reset_joints_by_offset, + mode="reset", + params={ + "position_range": (-0.1, 0.1), + "velocity_range": (0.0, 0.0), + "asset_cfg": SceneEntityCfg("robot"), + }, + ) + + randomize_cubes_in_focus = EventTerm( + func=so100_stack_events.randomize_single_rigid_objects, + mode="reset", + params={ + "asset_cfgs": [SceneEntityCfg("cube_1"), SceneEntityCfg("cube_2"), SceneEntityCfg("cube_3")], + "pose_range": {"x": (0.2, 0.3), "y": (-0.10, 0.10), "z": (0.01, 0.01), "yaw": (0.0, 0.0)}, + "min_separation": 0.1, + }, + ) + + +@configclass +class SO100CubeStackJointPosEnvCfg(StackEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Set events + self.events = EventCfg() + + # Set SO-100 as robot + self.scene.robot = SO_100_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + pos=(0, 0, 0.0), + rot=(0.7071, 0, 0, 0.7071), + joint_pos={ + # right-arm + "a_1": 0.0, + "a_2": 1.5708, + "a_3": -1.5708, + "a_4": 0.0, + "a_5": 0.0, + }, + joint_vel={".*": 0.0}, + )) + self.scene.robot.spawn.semantic_tags = [("class", "robot")] + + # Add semantics to table + self.scene.table.spawn.semantic_tags = [("class", "table")] + + # Add semantics to ground + self.scene.plane.semantic_tags = [("class", "ground")] + + # Set actions for the specific robot type (SO-100) + self.actions.arm_action = mdp.JointPositionActionCfg( + asset_name="robot", joint_names=["a_[1-5]"], scale=0.5, use_default_offset=False + ) + # # SO-100 doesn't have gripper, but we need to provide an action for compatibility + # # Use the last joint as a dummy gripper action + self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( + asset_name="robot", + joint_names=["a_6"], + open_command_expr={"a_6": 0.5}, + close_command_expr={"a_6": -0.1}, + ) + + # Rigid body properties of each cube + cube_properties = RigidBodyPropertiesCfg( + solver_position_iteration_count=16, + solver_velocity_iteration_count=1, + max_angular_velocity=1000.0, + max_linear_velocity=1000.0, + max_depenetration_velocity=5.0, + disable_gravity=False, + ) + + # Set each stacking cube deterministically + self.scene.cube_1 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_1", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.2, 0.0, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/blue_block.usd", + scale=(0.5, 0.5, 0.5), + rigid_props=cube_properties, + semantic_tags=[("class", "cube_1")], + ), + ) + self.scene.cube_2 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_2", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.2, 0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/red_block.usd", + scale=(0.5, 0.5, 0.5), + rigid_props=cube_properties, + semantic_tags=[("class", "cube_2")], + ), + ) + self.scene.cube_3 = RigidObjectCfg( + prim_path="{ENV_REGEX_NS}/Cube_3", + init_state=RigidObjectCfg.InitialStateCfg(pos=[0.3, -0.05, 0.0203], rot=[1, 0, 0, 0]), + spawn=UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/green_block.usd", + scale=(0.5, 0.5, 0.5), + rigid_props=cube_properties, + semantic_tags=[("class", "cube_3")], + ), + ) + + # Listens to the required transforms + marker_cfg = FRAME_MARKER_CFG.copy() + marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + marker_cfg.prim_path = "/Visuals/FrameTransformer" + self.scene.ee_frame = FrameTransformerCfg( + prim_path="{ENV_REGEX_NS}/Robot/base", + debug_vis=True, + visualizer_cfg=marker_cfg, + target_frames=[ + FrameTransformerCfg.FrameCfg( + prim_path="{ENV_REGEX_NS}/Robot/gripper", + name="end_effector", + offset=OffsetCfg( + pos=[0.0, -0.1, 0.0], + ), + ), + ], + ) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py new file mode 100644 index 00000000000..f3bec7a7952 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py @@ -0,0 +1,105 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause +from isaaclab.assets import ArticulationCfg +from isaaclab.controllers.pink_ik_cfg import PinkIKControllerCfg + +import isaaclab.controllers.utils as ControllerUtils +from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg +from isaaclab.envs.mdp.actions.pink_actions_cfg import PinkInverseKinematicsActionCfg +from isaaclab.utils import configclass + +from pink.tasks import FrameTask + + +from . import stack_joint_pos_env_cfg + +import tempfile + +## +# Pre-defined configs +## +from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip + + +@configclass +class SO100CubeStackPinkIKAbsEnvCfg(stack_joint_pos_env_cfg.SO100CubeStackJointPosEnvCfg): + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Temporary directory for URDF files + self.temp_urdf_dir = tempfile.gettempdir() + # Convert USD to URDF and change revolute joints to fixed + temp_urdf_output_path, temp_urdf_meshes_output_path = ControllerUtils.convert_usd_to_urdf( + self.scene.robot.spawn.usd_path, self.temp_urdf_dir, force_conversion=True + ) + + # Set Franka as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = SO_100_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + pos=(0, 0, 0.0), + rot=(0.7071, 0, 0, 0.7071), + # rot=(1.0, 0, 0, 0.0), + joint_pos={ + # right-arm + "a_1": 0.0, + "a_2": 1.5708, + "a_3": -1.5708, + "a_4": 1.2, + "a_5": 0.0, + }, + joint_vel={".*": 0.0}, + )) + + # Set actions for the specific robot type (franka) + self.actions.arm_action = PinkInverseKinematicsActionCfg( + pink_controlled_joint_names=[ + "a_1", + "a_2", + "a_3", + "a_4", + "a_5", + ], + # Joints to be locked in URDF + ik_urdf_fixed_joint_names=["a_6"], + hand_joint_names=[], + # the robot in the sim scene we are controlling + asset_name="robot", + # Configuration for the IK controller + # The frames names are the ones present in the URDF file + # The urdf has to be generated from the USD that is being used in the scene + controller=PinkIKControllerCfg( + articulation_name="robot", + base_link_name="base", + num_hand_joints=0, + show_ik_warnings=True, + variable_input_tasks=[ + FrameTask( + "gripper", + position_cost=1.0, # [cost] / [m] + orientation_cost=1.0, # [cost] / [rad] + lm_damping=10, # dampening for solver for step jumps + gain=0.1, + ) + ], + fixed_input_tasks=[ + # COMMENT OUT IF LOCKING WAIST/HEAD + # FrameTask( + # "GR1T2_fourier_hand_6dof_head_yaw_link", + # position_cost=1.0, # [cost] / [m] + # orientation_cost=0.05, # [cost] / [rad] + # ), + ], + ), + ) + ControllerUtils.change_revolute_to_fixed( + temp_urdf_output_path, self.actions.arm_action.ik_urdf_fixed_joint_names + ) + + # Set the URDF and mesh paths for the IK controller + self.actions.arm_action.controller.urdf_path = temp_urdf_output_path + self.actions.arm_action.controller.mesh_path = temp_urdf_meshes_output_path + diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py new file mode 100644 index 00000000000..110858b6fab --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py @@ -0,0 +1,194 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import tempfile + +from pink.tasks import FrameTask + +import isaaclab.controllers.utils as ControllerUtils +import isaaclab.sim as sim_utils +from isaaclab.assets import ArticulationCfg +from isaaclab.controllers.pink_ik_cfg import PinkIKControllerCfg +from isaaclab.envs.mdp.actions.pink_actions_cfg import PinkInverseKinematicsActionCfg +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.managers import SceneEntityCfg +from isaaclab.sensors import CameraCfg +from isaaclab.utils import configclass + +import isaaclab.envs.mdp as base_mdp +from ... import mdp +from . import stack_joint_pos_env_cfg + +## +# Pre-defined configs +## +from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip + + +@configclass +class ObservationsCfg: + """Observation specifications for the MDP.""" + + @configclass + class PolicyCfg(ObsGroup): + """Observations for policy group with state values.""" + + actions = ObsTerm(func=base_mdp.last_action) + joint_pos = ObsTerm(func=base_mdp.joint_pos_rel) + joint_vel = ObsTerm(func=base_mdp.joint_vel_rel) + object = ObsTerm(func=mdp.object_obs) + cube_positions = ObsTerm(func=mdp.cube_positions_in_world_frame) + cube_orientations = ObsTerm(func=mdp.cube_orientations_in_world_frame) + eef_pos = ObsTerm(func=mdp.ee_frame_pos) + eef_quat = ObsTerm(func=mdp.ee_frame_quat) + gripper_pos = ObsTerm(func=mdp.gripper_pos) + table_cam = ObsTerm( + func=base_mdp.image, params={"sensor_cfg": SceneEntityCfg("table_cam"), "data_type": "rgb", "normalize": False} + ) + # wrist_cam = ObsTerm( + # func=base_mdp.image, params={"sensor_cfg": SceneEntityCfg("wrist_cam"), "data_type": "rgb", "normalize": False} + # ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + + @configclass + class SubtaskCfg(ObsGroup): + """Observations for subtask group.""" + + grasp_1 = ObsTerm( + func=mdp.object_grasped, + params={ + "robot_cfg": SceneEntityCfg("robot"), + "ee_frame_cfg": SceneEntityCfg("ee_frame"), + "object_cfg": SceneEntityCfg("cube_2"), + }, + ) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + # observation groups + policy: PolicyCfg = PolicyCfg() + subtask_terms: SubtaskCfg = SubtaskCfg() + + +@configclass +class SO100CubeStackPinkIKAbsVisuomotorEnvCfg(stack_joint_pos_env_cfg.SO100CubeStackJointPosEnvCfg): + observations: ObservationsCfg = ObservationsCfg() + + def __post_init__(self): + # post init of parent + super().__post_init__() + + # Temporary directory for URDF files + self.temp_urdf_dir = tempfile.gettempdir() + # Convert USD to URDF and change revolute joints to fixed + temp_urdf_output_path, temp_urdf_meshes_output_path = ControllerUtils.convert_usd_to_urdf( + self.scene.robot.spawn.usd_path, self.temp_urdf_dir, force_conversion=True + ) + + # Set SO100 as robot + # We switch here to a stiffer PD controller for IK tracking to be better. + self.scene.robot = SO_100_HIGH_PD_CFG.replace( + prim_path="{ENV_REGEX_NS}/Robot", + init_state=ArticulationCfg.InitialStateCfg( + pos=(0, 0, 0.0), + rot=(0.7071, 0, 0, 0.7071), + joint_pos={ + # SO100 joints + "a_1": 0.0, + "a_2": 1.5708, + "a_3": -1.5708, + "a_4": 1.2, + "a_5": 0.0, + }, + joint_vel={".*": 0.0}, + ), + ) + + # Set actions for the specific robot type (SO100) + self.actions.arm_action = PinkInverseKinematicsActionCfg( + pink_controlled_joint_names=[ + "a_1", + "a_2", + "a_3", + "a_4", + "a_5", + ], + # Joints to be locked in URDF + ik_urdf_fixed_joint_names=["a_6"], + hand_joint_names=[], + # the robot in the sim scene we are controlling + asset_name="robot", + # Configuration for the IK controller + # The frames names are the ones present in the URDF file + # The urdf has to be generated from the USD that is being used in the scene + controller=PinkIKControllerCfg( + articulation_name="robot", + base_link_name="base", + num_hand_joints=0, + show_ik_warnings=True, + variable_input_tasks=[ + FrameTask( + "gripper", + position_cost=1.0, # [cost] / [m] + orientation_cost=1.0, # [cost] / [rad] + lm_damping=10, # dampening for solver for step jumps + gain=0.1, + ) + ], + fixed_input_tasks=[], + ), + ) + ControllerUtils.change_revolute_to_fixed( + temp_urdf_output_path, self.actions.arm_action.ik_urdf_fixed_joint_names + ) + + # Set the URDF and mesh paths for the IK controller + self.actions.arm_action.controller.urdf_path = temp_urdf_output_path + self.actions.arm_action.controller.mesh_path = temp_urdf_meshes_output_path + + # # Set wrist camera + # self.scene.wrist_cam = CameraCfg( + # prim_path="{ENV_REGEX_NS}/Robot/gripper/wrist_cam", + # update_period=0.0, + # height=512, + # width=512, + # data_types=["rgb", "distance_to_image_plane"], + # spawn=sim_utils.PinholeCameraCfg( + # focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 2) + # ), + # offset=CameraCfg.OffsetCfg( + # pos=(0.08, 0.0, -0.1), rot=(-0.70614, 0.03701, 0.03701, -0.70614), convention="ros" + # ), + # ) + + # Set table view camera + self.scene.table_cam = CameraCfg( + prim_path="{ENV_REGEX_NS}/table_cam", + update_period=0.0, + height=512, + width=512, + data_types=["rgb", "distance_to_image_plane"], + spawn=sim_utils.PinholeCameraCfg( + focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 2) + ), + offset=CameraCfg.OffsetCfg( + pos=(1.0, 0.0, 0.4), rot=(0.35355, -0.61237, -0.61237, 0.35355), convention="ros" + ), + ) + + # Set settings for camera rendering + self.rerender_on_reset = True + self.sim.render.antialiasing_mode = "OFF" # disable dlss + + # List of image observations in policy observations + self.image_obs_list = ["table_cam"] + # self.image_obs_list = ["table_cam", "wrist_cam"] + \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py index 0d9d087a9c2..5c1473b3d81 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -270,7 +270,7 @@ def object_grasped( ee_frame_cfg: SceneEntityCfg, object_cfg: SceneEntityCfg, diff_threshold: float = 0.06, - gripper_open_val: torch.tensor = torch.tensor([0.04]), + gripper_open_val: torch.tensor = torch.tensor([0.5]), gripper_threshold: float = 0.005, ) -> torch.Tensor: """Check if an object is grasped by the specified robot.""" @@ -283,14 +283,19 @@ def object_grasped( end_effector_pos = ee_frame.data.target_pos_w[:, 0, :] pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) + print(f"Object position: {object_pos}") + print(f"End effector position: {end_effector_pos}") + # print(f"Pose difference: {pose_diff}") + # print(f"Gripper joint positions: {robot.data.joint_pos[:, -1]}") + # print(f"Gripper open value: {gripper_open_val}") + + # TODO: Move this to an observation that is specififc to so-100 grasped = torch.logical_and( pose_diff < diff_threshold, torch.abs(robot.data.joint_pos[:, -1] - gripper_open_val.to(env.device)) > gripper_threshold, ) - grasped = torch.logical_and( - grasped, torch.abs(robot.data.joint_pos[:, -2] - gripper_open_val.to(env.device)) > gripper_threshold - ) + print(f"Grasped status: {grasped}") return grasped diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/so100_stack_events.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/so100_stack_events.py new file mode 100644 index 00000000000..852705907e5 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/so100_stack_events.py @@ -0,0 +1,97 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import math +import random +import torch +from typing import TYPE_CHECKING + +import isaaclab.utils.math as math_utils +from isaaclab.managers import SceneEntityCfg + +if TYPE_CHECKING: + from isaaclab.envs import ManagerBasedEnv + + +def sample_object_poses( + num_objects: int, + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + """Sample object poses with minimum separation constraint.""" + range_list = [pose_range.get(key, (0.0, 0.0)) for key in ["x", "y", "z", "roll", "pitch", "yaw"]] + pose_list = [] + + for i in range(num_objects): + for j in range(max_sample_tries): + sample = [random.uniform(range[0], range[1]) for range in range_list] + + # Accept pose if it is the first one, or if reached max num tries + if len(pose_list) == 0 or j == max_sample_tries - 1: + pose_list.append(sample) + break + + # Check if pose of object is sufficiently far away from all other objects + separation_check = [math.dist(sample[:3], pose[:3]) > min_separation for pose in pose_list] + if False not in separation_check: + pose_list.append(sample) + break + + return pose_list + + +def randomize_single_rigid_objects( + env: ManagerBasedEnv, + env_ids: torch.Tensor, + asset_cfgs: list[SceneEntityCfg], + min_separation: float = 0.0, + pose_range: dict[str, tuple[float, float]] = {}, + max_sample_tries: int = 5000, +): + """ + Randomize poses of individual rigid objects (not multi-object rigid bodies). + + This is a version of randomize_rigid_objects_in_focus adapted for single RigidObjects. + Unlike the Franka version which handles multi-object rigid bodies with num_objects attribute, + this function works with individual RigidObject instances in the scene. + + Args: + env: The manager-based environment. + env_ids: Environment indices to apply the randomization to. + asset_cfgs: List of scene entity configurations for the rigid objects to randomize. + min_separation: Minimum separation distance between objects. + pose_range: Dictionary specifying the range for each pose component (x, y, z, roll, pitch, yaw). + max_sample_tries: Maximum number of attempts to sample valid poses. + """ + if env_ids is None: + return + + # Randomize poses in each environment independently + for cur_env in env_ids.tolist(): + pose_list = sample_object_poses( + num_objects=len(asset_cfgs), + min_separation=min_separation, + pose_range=pose_range, + max_sample_tries=max_sample_tries, + ) + + # Randomize pose for each object + for i in range(len(asset_cfgs)): + asset_cfg = asset_cfgs[i] + asset = env.scene[asset_cfg.name] + + # Write pose to simulation + pose_tensor = torch.tensor([pose_list[i]], device=env.device) + positions = pose_tensor[:, 0:3] + env.scene.env_origins[cur_env, 0:3] + orientations = math_utils.quat_from_euler_xyz(pose_tensor[:, 3], pose_tensor[:, 4], pose_tensor[:, 5]) + asset.write_root_pose_to_sim( + torch.cat([positions, orientations], dim=-1), env_ids=torch.tensor([cur_env], device=env.device) + ) + asset.write_root_velocity_to_sim( + torch.zeros(1, 6, device=env.device), env_ids=torch.tensor([cur_env], device=env.device) + ) \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py index 91a6237cee7..b183b6da3a9 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -29,8 +29,8 @@ def cubes_stacked( cube_3_cfg: SceneEntityCfg = SceneEntityCfg("cube_3"), xy_threshold: float = 0.05, height_threshold: float = 0.005, - height_diff: float = 0.0468, - gripper_open_val: torch.tensor = torch.tensor([0.04]), + height_diff: float = 0.0233, + gripper_open_val: torch.tensor = torch.tensor([0.5]), atol=0.0001, rtol=0.0001, ): @@ -40,27 +40,30 @@ def cubes_stacked( cube_3: RigidObject = env.scene[cube_3_cfg.name] pos_diff_c12 = cube_1.data.root_pos_w - cube_2.data.root_pos_w - pos_diff_c23 = cube_2.data.root_pos_w - cube_3.data.root_pos_w # Compute cube position difference in x-y plane xy_dist_c12 = torch.norm(pos_diff_c12[:, :2], dim=1) - xy_dist_c23 = torch.norm(pos_diff_c23[:, :2], dim=1) # Compute cube height difference h_dist_c12 = torch.norm(pos_diff_c12[:, 2:], dim=1) - h_dist_c23 = torch.norm(pos_diff_c23[:, 2:], dim=1) # Check cube positions - stacked = torch.logical_and(xy_dist_c12 < xy_threshold, xy_dist_c23 < xy_threshold) - stacked = torch.logical_and(h_dist_c12 - height_diff < height_threshold, stacked) - stacked = torch.logical_and(h_dist_c23 - height_diff < height_threshold, stacked) - + stacked = xy_dist_c12 < xy_threshold + print("stacked xy", stacked) + print("xy_dist_c12", xy_dist_c12) + print("xy_threshold", xy_threshold) + # stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) + stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) + print("stacked h", stacked) + print("h_dist_c12", h_dist_c12) + print("height_diff", height_diff) + print("height_threshold", height_threshold) + # Check gripper positions stacked = torch.logical_and( - torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked - ) - stacked = torch.logical_and( - torch.isclose(robot.data.joint_pos[:, -2], gripper_open_val.to(env.device), atol=atol, rtol=rtol), stacked + torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=0.0001, rtol=0.0001), stacked ) + print("stacked gripper", stacked) + print("stacked", stacked) return stacked From 6336c858abc0c4909ff09f822ed19fe4f9a8d324 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Thu, 26 Jun 2025 14:42:27 -0700 Subject: [PATCH 02/14] update usd path --- source/isaaclab_assets/isaaclab_assets/robots/so_100.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/isaaclab_assets/isaaclab_assets/robots/so_100.py b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py index f126fef00a1..da9ed0d03c4 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/so_100.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py @@ -22,7 +22,7 @@ SO_100_CFG = ArticulationCfg( spawn=sim_utils.UsdFileCfg( - usd_path=f"omniverse://isaac-dev.ov.nvidia.com/Users/ashwinvk@nvidia.com/so-100/so-100_pure_usd.usd", + usd_path=f"omniverse://isaac-dev.ov.nvidia.com/Users/ashwinvk@nvidia.com/so-100/so-100.usd", activate_contact_sensors=False, rigid_props=sim_utils.RigidBodyPropertiesCfg( disable_gravity=True, From ed669e16834e65386a91bfec351858db2b23c278 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Thu, 26 Jun 2025 15:19:34 -0700 Subject: [PATCH 03/14] remove ik relative env --- .../config/so_100/stack_ik_rel_env_cfg.py | 50 ------------------- 1 file changed, 50 deletions(-) delete mode 100644 source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py deleted file mode 100644 index cdef31dbd8b..00000000000 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_ik_rel_env_cfg.py +++ /dev/null @@ -1,50 +0,0 @@ -# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). -# All rights reserved. -# -# SPDX-License-Identifier: BSD-3-Clause -from isaaclab.assets import ArticulationCfg -from isaaclab.controllers.differential_ik_cfg import DifferentialIKControllerCfg -from isaaclab.envs.mdp.actions.actions_cfg import DifferentialInverseKinematicsActionCfg -from isaaclab.utils import configclass - -from . import stack_joint_pos_env_cfg - -## -# Pre-defined configs -## -from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip - - -@configclass -class SO100CubeStackIKRelEnvCfg(stack_joint_pos_env_cfg.SO100CubeStackJointPosEnvCfg): - def __post_init__(self): - # post init of parent - super().__post_init__() - - # Set Franka as robot - # We switch here to a stiffer PD controller for IK tracking to be better. - self.scene.robot = SO_100_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", - init_state=ArticulationCfg.InitialStateCfg( - pos=(0, 0, 0.0), - # rot=(0.7071, 0, 0, 0.7071), - rot=(1.0, 0, 0, 0.0), - joint_pos={ - # right-arm - "a_1": 0.0, - "a_2": 1.5708, - "a_3": -1.5708, - "a_4": 1.2, - "a_5": 0.0, - }, - joint_vel={".*": 0.0}, - )) - - # Set actions for the specific robot type (franka) - self.actions.arm_action = DifferentialInverseKinematicsActionCfg( - asset_name="robot", - joint_names=["a_[1-4]"], - body_name="wrist", - controller=DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls"), - scale=1.0, - body_offset=DifferentialInverseKinematicsActionCfg.OffsetCfg(pos=[0.0, 0.0, 0.0]), - ) From 500c1d97a5b76492ee8eb787fb360e1868c0dc26 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Thu, 26 Jun 2025 15:35:04 -0700 Subject: [PATCH 04/14] add udpates for new usd model --- .../isaaclab_assets/robots/so_100.py | 28 ++++++------------- .../stack/config/so_100/__init__.py | 10 ------- .../config/so_100/stack_joint_pos_env_cfg.py | 20 ++++++------- .../so_100/stack_pink_ik_abs_env_cfg.py | 28 +++++++++---------- .../stack_pink_ik_abs_visuomotor_env_cfg.py | 27 +++++++++--------- 5 files changed, 47 insertions(+), 66 deletions(-) diff --git a/source/isaaclab_assets/isaaclab_assets/robots/so_100.py b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py index da9ed0d03c4..4d4067bfa37 100644 --- a/source/isaaclab_assets/isaaclab_assets/robots/so_100.py +++ b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py @@ -40,17 +40,17 @@ ), init_state=ArticulationCfg.InitialStateCfg( joint_pos={ - "a_1": 0.0, - "a_2": 0.0, - "a_3": 0.0, - "a_4": 0.0, - "a_5": 0.0, - "a_6": 0.0, + "shoulder_pan_joint": 0.0, + "shoulder_lift_joint": 0.0, + "elbow_flex_joint": 0.0, + "wrist_flex_joint": 0.0, + "wrist_roll_joint": 0.0, + "gripper_joint": 0.0, }, ), actuators={ "all_joints": ImplicitActuatorCfg( - joint_names_expr=["a_[1-5]"], + joint_names_expr=["shoulder_pan_joint", "shoulder_lift_joint", "elbow_flex_joint", "wrist_flex_joint", "wrist_roll_joint"], effort_limit=None, velocity_limit=None, stiffness=None, @@ -58,7 +58,7 @@ armature=0.0, ), "gripper": ImplicitActuatorCfg( - joint_names_expr=["a_6"], + joint_names_expr=["gripper_joint"], effort_limit_sim=0.1, velocity_limit_sim=2.175, stiffness=8.0, @@ -68,14 +68,4 @@ }, soft_joint_pos_limit_factor=1.0, ) -"""Configuration of Franka Emika Panda robot.""" - - -SO_100_HIGH_PD_CFG = SO_100_CFG.copy() -SO_100_HIGH_PD_CFG.spawn.rigid_props.disable_gravity = True -# SO_100_HIGH_PD_CFG.actuators["all_joints"].stiffness = 400.0 -# SO_100_HIGH_PD_CFG.actuators["all_joints"].damping = 80.0 -"""Configuration of Franka Emika Panda robot with stiffer PD control. - -This configuration is useful for task-space control using differential IK. -""" +"""Configuration of SO‑100 robot arm.""" diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py index 3342c12ca61..b413fcc1f49 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py @@ -6,7 +6,6 @@ import gymnasium as gym from . import stack_joint_pos_env_cfg -from . import stack_ik_rel_env_cfg from . import stack_pink_ik_abs_env_cfg from . import stack_pink_ik_abs_visuomotor_env_cfg ## @@ -26,15 +25,6 @@ disable_env_checker=True, ) -gym.register( - id="Isaac-Stack-Cube-SO100-IK-Rel-v0", - entry_point="isaaclab.envs:ManagerBasedRLEnv", - kwargs={ - "env_cfg_entry_point": stack_ik_rel_env_cfg.SO100CubeStackIKRelEnvCfg, - }, - disable_env_checker=True, -) - gym.register( id="Isaac-Stack-Cube-SO100-Pink-IK-Abs-v0", entry_point="isaaclab.envs:ManagerBasedRLEnv", diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py index 08000055684..95100d5f6c4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py @@ -67,12 +67,12 @@ def __post_init__(self): pos=(0, 0, 0.0), rot=(0.7071, 0, 0, 0.7071), joint_pos={ - # right-arm - "a_1": 0.0, - "a_2": 1.5708, - "a_3": -1.5708, - "a_4": 0.0, - "a_5": 0.0, + "shoulder_pan_joint": 0.0, + "shoulder_lift_joint": 1.5708, + "elbow_flex_joint": -1.5708, + "wrist_flex_joint": 0.0, + "wrist_roll_joint": 0.0, + "gripper_joint": 0.0, }, joint_vel={".*": 0.0}, )) @@ -86,15 +86,15 @@ def __post_init__(self): # Set actions for the specific robot type (SO-100) self.actions.arm_action = mdp.JointPositionActionCfg( - asset_name="robot", joint_names=["a_[1-5]"], scale=0.5, use_default_offset=False + asset_name="robot", joint_names=["shoulder_pan_joint", "shoulder_lift_joint", "elbow_flex_joint", "wrist_flex_joint", "wrist_roll_joint"], scale=1.0, use_default_offset=False ) # # SO-100 doesn't have gripper, but we need to provide an action for compatibility # # Use the last joint as a dummy gripper action self.actions.gripper_action = mdp.BinaryJointPositionActionCfg( asset_name="robot", - joint_names=["a_6"], - open_command_expr={"a_6": 0.5}, - close_command_expr={"a_6": -0.1}, + joint_names=["gripper_joint"], + open_command_expr={"gripper_joint": 0.5}, + close_command_expr={"gripper_joint": -0.1}, ) # Rigid body properties of each cube diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py index f3bec7a7952..b76e3d0bcf4 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_env_cfg.py @@ -20,7 +20,7 @@ ## # Pre-defined configs ## -from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip +from isaaclab_assets.robots.so_100 import SO_100_CFG # isort: skip @configclass @@ -38,18 +38,18 @@ def __post_init__(self): # Set Franka as robot # We switch here to a stiffer PD controller for IK tracking to be better. - self.scene.robot = SO_100_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", + self.scene.robot = SO_100_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot", init_state=ArticulationCfg.InitialStateCfg( pos=(0, 0, 0.0), rot=(0.7071, 0, 0, 0.7071), - # rot=(1.0, 0, 0, 0.0), joint_pos={ # right-arm - "a_1": 0.0, - "a_2": 1.5708, - "a_3": -1.5708, - "a_4": 1.2, - "a_5": 0.0, + "shoulder_pan_joint": 0.0, + "shoulder_lift_joint": 1.5708, + "elbow_flex_joint": -1.5708, + "wrist_flex_joint": 1.2, + "wrist_roll_joint": 0.0, + "gripper_joint": 0.0, }, joint_vel={".*": 0.0}, )) @@ -57,14 +57,14 @@ def __post_init__(self): # Set actions for the specific robot type (franka) self.actions.arm_action = PinkInverseKinematicsActionCfg( pink_controlled_joint_names=[ - "a_1", - "a_2", - "a_3", - "a_4", - "a_5", + "shoulder_pan_joint", + "shoulder_lift_joint", + "elbow_flex_joint", + "wrist_flex_joint", + "wrist_roll_joint", ], # Joints to be locked in URDF - ik_urdf_fixed_joint_names=["a_6"], + ik_urdf_fixed_joint_names=["gripper_joint"], hand_joint_names=[], # the robot in the sim scene we are controlling asset_name="robot", diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py index 110858b6fab..b43155a7678 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py @@ -25,7 +25,7 @@ ## # Pre-defined configs ## -from isaaclab_assets.robots.so_100 import SO_100_HIGH_PD_CFG # isort: skip +from isaaclab_assets.robots.so_100 import SO_100_CFG # isort: skip @configclass @@ -95,18 +95,19 @@ def __post_init__(self): # Set SO100 as robot # We switch here to a stiffer PD controller for IK tracking to be better. - self.scene.robot = SO_100_HIGH_PD_CFG.replace( + self.scene.robot = SO_100_CFG.replace( prim_path="{ENV_REGEX_NS}/Robot", init_state=ArticulationCfg.InitialStateCfg( pos=(0, 0, 0.0), rot=(0.7071, 0, 0, 0.7071), joint_pos={ # SO100 joints - "a_1": 0.0, - "a_2": 1.5708, - "a_3": -1.5708, - "a_4": 1.2, - "a_5": 0.0, + "shoulder_pan_joint": 0.0, + "shoulder_lift_joint": 1.5708, + "elbow_flex_joint": -1.5708, + "wrist_flex_joint": 1.2, + "wrist_roll_joint": 0.0, + "gripper_joint": 0.0, }, joint_vel={".*": 0.0}, ), @@ -115,14 +116,14 @@ def __post_init__(self): # Set actions for the specific robot type (SO100) self.actions.arm_action = PinkInverseKinematicsActionCfg( pink_controlled_joint_names=[ - "a_1", - "a_2", - "a_3", - "a_4", - "a_5", + "shoulder_pan_joint", + "shoulder_lift_joint", + "elbow_flex_joint", + "wrist_flex_joint", + "wrist_roll_joint", ], # Joints to be locked in URDF - ik_urdf_fixed_joint_names=["a_6"], + ik_urdf_fixed_joint_names=["gripper_joint"], hand_joint_names=[], # the robot in the sim scene we are controlling asset_name="robot", From 1368e66cca2f223e44adde15ac1a4f51fcc18ed0 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 09:34:38 -0700 Subject: [PATCH 05/14] remove ignoring all datasets folders and add lerobot file handler --- .gitignore | 2 +- .../datasets/lerobot_dataset_file_handler.py | 254 ++++++++++++++++++ 2 files changed, 255 insertions(+), 1 deletion(-) create mode 100644 source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py diff --git a/.gitignore b/.gitignore index 1472f06b87f..48706d07dbe 100644 --- a/.gitignore +++ b/.gitignore @@ -62,7 +62,7 @@ _build /.pretrained_checkpoints/ # Teleop Recorded Dataset -datasets +/datasets # Tests tests/ diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py new file mode 100644 index 00000000000..c114c25c06f --- /dev/null +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -0,0 +1,254 @@ +# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +# Copyright (c) 2024-2025, The Isaac Lab Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +import json +import os +import shutil +import tempfile +import torch +import numpy as np +from pathlib import Path +from typing import Any, Dict, List, Optional +from collections.abc import Iterable + +try: + from lerobot.common.datasets.lerobot_dataset import LeRobotDataset + from lerobot.common.datasets.utils import ( + create_empty_dataset_info, + get_hf_features_from_features, + DEFAULT_FEATURES + ) + LEROBOT_AVAILABLE = True +except ImportError: + LEROBOT_AVAILABLE = False + +from .dataset_file_handler_base import DatasetFileHandlerBase +from .episode_data import EpisodeData + + +class LeRobotDatasetFileHandler(DatasetFileHandlerBase): + """LeRobot dataset file handler for storing and loading episode data.""" + + def __init__(self): + """Initializes the LeRobot dataset file handler.""" + if not LEROBOT_AVAILABLE: + raise ImportError( + "LeRobot dependencies not available. Please install following the documentation here: " + "https://github.com/huggingface/lerobot" + ) + + self._dataset = None + self._dataset_path = None + self._env_name = None + self._episode_count = 0 + + def create(self, file_path: str, env_name: str | None = None, env = None): + """Create a new dataset file by automatically extracting features from environment.""" + if not file_path.endswith(".lerobot"): + file_path += ".lerobot" + + self._dataset_path = Path(file_path) + + # Delete existing dataset if it exists + if self._dataset_path.exists(): + print(f"Removing existing dataset at {self._dataset_path}") + shutil.rmtree(self._dataset_path) + + # Extract repo_id from file path + repo_id = self._dataset_path.name.replace('.lerobot', '') + + # Initialize environment name + self._env_name = env_name or "isaac_lab_env" + + # Extract features from environment + features = self._extract_features_from_env(env) + + # Calculate FPS from environment timestep + fps = int(1 / env.step_dt) + + # Create LeRobot dataset + try: + self._dataset = LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + features=features, + root=self._dataset_path, + robot_type="isaac_lab_robot", + use_videos=True, + tolerance_s=1e-4 + ) + + # Add environment name to metadata + self._dataset.meta.info["env_name"] = self._env_name + + except Exception as e: + raise RuntimeError(f"Failed to create LeRobot dataset: {e}") + + self._episode_count = 0 + + def _extract_features_from_env(self, env) -> Dict[str, Dict]: + """Extract features schema from environment observations and actions.""" + if env is None: + raise ValueError("Environment must be provided to extract features") + + features = {} + + features["action"] = { + "dtype": "float32", + "shape": (env.action_manager.total_action_dim,), + "names": None + } + + # Add hardcoded features + features["annotation.human.action.task_description"] = { + "dtype": "int64", + "shape": (1,), + "names": None + } + + # Get observation features from observation manager active terms + obs_sample = env.observation_manager.compute() + + print("obs_sample", obs_sample) + # Extract features from nested observation structure + for group_key, group_data in obs_sample.items(): + if isinstance(group_data, dict): + for obs_key, obs_value in group_data.items(): + if isinstance(obs_value, torch.Tensor): + if "cam" in obs_key: + print(f"Processing camera observation: {obs_key}") + print(f"Camera observation shape: {obs_value.shape}") + # For camera observations, remove batch dimension and use [C, H, W] format + # obs_value shape is typically [batch_size, H, W, C] or [H, W, C] + if obs_value.ndim == 4: # [batch_size, H, W, C] + height, width, channels = obs_value.shape[1], obs_value.shape[2], obs_value.shape[3] + print(f"4D camera observation - batch size: {obs_value.shape[0]}") + else: # [H, W, C] + height, width, channels = obs_value.shape[0], obs_value.shape[1], obs_value.shape[2] + print("3D camera observation - no batch dimension") + + print(f"Extracted dimensions - channels: {channels}, height: {height}, width: {width}") + features[f"observation.{obs_key}"] = { + "dtype": "video", + "shape": (channels, height, width), # LeRobot expects [C, H, W] + "names": ["channel", "height", "width"], + "video_info": { + "video.fps": int(1 / env.step_dt) + } + } + print(f"Added video feature for {obs_key} with shape {(channels, height, width)}") + elif "joint_pos" in obs_key: + print(f"Processing joint position observation: {obs_key}") + print(f"Joint position shape: {obs_value.shape}") + # State observation - remove batch dimension + if obs_value.ndim > 1: + shape = obs_value.shape[1:] # Remove batch dimension + print(f"Multi-dimensional joint position - original shape: {obs_value.shape}, reduced shape: {shape}") + else: + shape = obs_value.shape + print(f"Single-dimensional joint position - shape: {shape}") + features["observation.state"] = { + "dtype": "float32", + "shape": shape, + "names": None + } + + return features + + def open(self, file_path: str, mode: str = "r"): + """Open an existing dataset file.""" + raise NotImplementedError("Open not implemented for LeRobot handler") + + def get_env_name(self) -> str | None: + """Get the environment name.""" + return self._env_name + + def get_episode_names(self) -> Iterable[str]: + """Get the names of the episodes in the file.""" + if self._dataset is None: + return [] + return [f"episode_{i:06d}" for i in range(self._episode_count)] + + def get_num_episodes(self) -> int: + """Get number of episodes in the file.""" + return self._episode_count + + def write_episode(self, episode: EpisodeData): + """Add an episode to the dataset.""" + if self._dataset is None or episode.is_empty(): + return + + # Convert Isaac Lab episode data to LeRobot format and save + self._convert_and_save_episode(episode) + self._episode_count += 1 + + def _convert_and_save_episode(self, episode: EpisodeData): + """Convert Isaac Lab episode data to LeRobot format and save it.""" + episode_dict = episode.data + + print("episode_dict", episode_dict) + + # Determine number of frames + num_frames = len(episode_dict["actions"]) + + + # Generate a task description + task = f"Pick the red cube and place it on the blue cube" + + # Add frames one by one to the LeRobot dataset + for frame_idx in range(num_frames): + frame_data = {} + + # Process episode data + for key, value in episode_dict.items(): + if key == "actions": + # Handle actions + frame_data["action"] = value[frame_idx].cpu().numpy() + elif key == "obs": + # Handle observations - value is already the obs dictionary + for obs_key, obs_value in value.items(): + if "cam" in obs_key: + # Convert camera data from [H, W, C] to [C, H, W] format for LeRobot + camera_frame = obs_value[frame_idx].cpu().numpy() + if camera_frame.ndim == 3: # [H, W, C] + camera_frame = camera_frame.transpose(2, 0, 1) # Convert to [C, H, W] + frame_data[f"observation.{obs_key}"] = camera_frame + elif "joint_pos" in obs_key: + frame_data["observation.state"] = obs_value[frame_idx].cpu().numpy() + + # Add hardcoded feature data + frame_data["annotation.human.action.task_description"] = np.array([0], dtype=np.int64) + + # Add frame to the dataset + self._dataset.add_frame(frame_data, task) + + # Save the episode + try: + self._dataset.save_episode() + except Exception as e: + print(f"Warning: Failed to save episode: {e}") + + def load_episode(self, episode_name: str) -> EpisodeData | None: + """Load episode data from the file.""" + raise NotImplementedError("Load episode not implemented for LeRobot handler") + + def flush(self): + """Flush any pending data to disk.""" + # LeRobot dataset handles flushing automatically when save_episode() is called + pass + + def close(self): + """Close the dataset file handler.""" + # Stop any async image writers + if self._dataset and hasattr(self._dataset, 'image_writer') and self._dataset.image_writer: + self._dataset.stop_image_writer() + + # Clear references + self._dataset = None From d2955b646097a10fcfb5e4d300ec83ee0334a234 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 15:29:16 -0700 Subject: [PATCH 06/14] remove print statements and viz debug --- .../stack/config/so_100/stack_joint_pos_env_cfg.py | 2 +- .../manager_based/manipulation/stack/mdp/terminations.py | 9 --------- 2 files changed, 1 insertion(+), 10 deletions(-) diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py index 95100d5f6c4..a9fa60b5cef 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_joint_pos_env_cfg.py @@ -145,7 +145,7 @@ def __post_init__(self): marker_cfg.prim_path = "/Visuals/FrameTransformer" self.scene.ee_frame = FrameTransformerCfg( prim_path="{ENV_REGEX_NS}/Robot/base", - debug_vis=True, + debug_vis=False, visualizer_cfg=marker_cfg, target_frames=[ FrameTransformerCfg.FrameCfg( diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py index b183b6da3a9..14f3a930ba8 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/terminations.py @@ -49,21 +49,12 @@ def cubes_stacked( # Check cube positions stacked = xy_dist_c12 < xy_threshold - print("stacked xy", stacked) - print("xy_dist_c12", xy_dist_c12) - print("xy_threshold", xy_threshold) # stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) - print("stacked h", stacked) - print("h_dist_c12", h_dist_c12) - print("height_diff", height_diff) - print("height_threshold", height_threshold) # Check gripper positions stacked = torch.logical_and( torch.isclose(robot.data.joint_pos[:, -1], gripper_open_val.to(env.device), atol=0.0001, rtol=0.0001), stacked ) - print("stacked gripper", stacked) - print("stacked", stacked) return stacked From 6db76cf83edaf0e6126e559c336b67eca403c671 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 15:36:44 -0700 Subject: [PATCH 07/14] retune the action_gain --- scripts/tools/record_demos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/tools/record_demos.py b/scripts/tools/record_demos.py index 4c36f94d24a..3d4034bba1b 100644 --- a/scripts/tools/record_demos.py +++ b/scripts/tools/record_demos.py @@ -439,7 +439,7 @@ def create_teleop_device(device_name: str, env: gym.Env): # actions[0, 2] += 0.01 # compute actions - control_gain = 0.1 + control_gain = 0.01 actions[0, 0] += delta_pose[0] * control_gain actions[0, 1] += delta_pose[1] * control_gain actions[0, 2] += delta_pose[2] * control_gain From e8ce16d8648f0544d8db168d04b8297a57cdce71 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 15:37:30 -0700 Subject: [PATCH 08/14] restructure lerobot dataset converter --- docs/source/how-to/record_lerobot_datasets.md | 43 + .../isaaclab/envs/manager_based_env_cfg.py | 29 + .../datasets/lerobot_dataset_file_handler.py | 739 +++++++++++++++--- .../test_lerobot_dataset_file_handler.py | 180 ++++- .../stack_pink_ik_abs_visuomotor_env_cfg.py | 15 + 5 files changed, 900 insertions(+), 106 deletions(-) diff --git a/docs/source/how-to/record_lerobot_datasets.md b/docs/source/how-to/record_lerobot_datasets.md index 1f79b33f056..712cc79d620 100644 --- a/docs/source/how-to/record_lerobot_datasets.md +++ b/docs/source/how-to/record_lerobot_datasets.md @@ -86,9 +86,52 @@ LeRobot uses standardized naming conventions for observations: - **Camera observations**: `observation.images.{camera_position}` - Examples: `observation.images.front`, `observation.images.wrist`, `observation.images.top` - **Robot state**: `observation.state` +- **Regular observations**: `observation.{obs_key}` - **Actions**: `action` - **Episode metadata**: `episode_index`, `frame_index`, `timestamp`, `task` +## Configuration Options + +The LeRobot dataset handler supports flexible configuration through the `LeRobotDatasetCfg`: + +```python +from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + +# Configure which observations to record +env_cfg.lerobot_dataset = LeRobotDatasetCfg() + +# Regular observations (saved as "observation.{key}") - REQUIRED +env_cfg.lerobot_dataset.observation_keys_to_record = [ + "policy/camera_rgb", + "policy/end_effector_pos", + "policy/gripper_state" +] + +# State observations (combined into "observation.state") - REQUIRED +env_cfg.lerobot_dataset.state_observation_keys = [ + "policy/joint_pos", + "policy/joint_vel", + "policy/joint_torque" +] + +``` + +**Important**: At least one of `observation_keys_to_record` or `state_observation_keys` must be configured with at least one observation. If both are empty, an error will be raised. If you don't want to record any state observations, use an empty list `[]` for `state_observation_keys` but ensure `observation_keys_to_record` has at least one entry. + +### State Observations + +State observations are special observation keys that are combined into a single `observation.state` feature in the LeRobot format. This is useful for: + +- **Robot state information**: Joint positions, velocities, torques +- **Privileged information**: Ground truth object poses, task-specific state +- **Combined features**: Multiple related state variables that should be treated as a single observation + +When multiple state observations are specified, they are concatenated into a single feature vector. For example, if you have: +- `joint_pos` with 7 dimensions +- `joint_vel` with 7 dimensions + +The resulting `observation.state` will have 14 dimensions (7 + 7). + ## Using the Recording Script The easiest way to record demonstrations in LeRobot format is using the built-in recording script: diff --git a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py index f119b66e487..b0dda651f1a 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py @@ -10,6 +10,7 @@ """ from dataclasses import MISSING +from typing import Dict, List, Optional import isaaclab.envs.mdp as mdp from isaaclab.devices.openxr import XrCfg @@ -34,6 +35,27 @@ class DefaultEventManagerCfg: reset_scene_to_default = EventTerm(func=mdp.reset_scene_to_default, mode="reset") +@configclass +class LeRobotDatasetCfg: + """Configuration for LeRobot dataset recording. + + This configuration specifies which observation keys should be recorded to the LeRobot dataset. + """ + + # List of (group_name, observation_key) tuples to record as regular observations + # e.g., [("policy", "joint_pos"), ("policy", "camera_rgb"), ("critic", "joint_vel")] + # These will be saved as "observation.{obs_key}" in the LeRobot format + observation_keys_to_record: List[tuple[str, str]] = MISSING + + # List of (group_name, observation_key) tuples that should be treated as state observations + # e.g., [("policy", "joint_pos"), ("policy", "joint_vel")] + # These will be saved as "observation.state" in the LeRobot format + state_observation_keys: List[tuple[str, str]] = MISSING + + # Custom task description generator function name + task_description: str = MISSING + + @configclass class ManagerBasedEnvCfg: """Base configuration of the environment.""" @@ -121,3 +143,10 @@ class ManagerBasedEnvCfg: xr: XrCfg | None = None """Configuration for viewing and interacting with the environment through an XR device.""" + + # LeRobot dataset configuration + lerobot_dataset: LeRobotDatasetCfg | None = None + """Configuration for LeRobot dataset recording. Defaults to LeRobotDatasetCfg(). + + This configuration specifies which observation keys should be recorded to the LeRobot dataset. + """ diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py index c114c25c06f..c9fcf1c2ca1 100644 --- a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -8,22 +8,126 @@ # # SPDX-License-Identifier: BSD-3-Clause -import json -import os +""" +LeRobot Dataset File Handler + +This module provides a configuration-driven LeRobot dataset file handler that works with +all manager-based environments in Isaac Lab. + +DESIGN OVERVIEW: +================== + +The LeRobotDatasetFileHandler is designed to automatically extract and record episode data +from Isaac Lab environments to LeRobot dataset format. It uses a configuration-based approach +to determine which observations and actions to record. + +The Lerobot format expects the dataset to contain +- observation.example_1 +- observation.example_2 +- observation.example_3 +- observation.example_4 +- observation.state +- action + +The action is extract using the action manager. + +The observation that correspondes to the state is configured using the state_observation_keys attribute. + +The state is extract using the observation manager. +The state is a concatenation of the state_observation_keys. + +The observation that are not part of the state are configured using the observation_keys_to_record attribute. + +The task description is configured using the task_description attribute. + + +KEY FEATURES: +============ + +1. CONFIGURATION-DRIVEN: + - Uses LeRobotDatasetCfg to specify which observations to record + - Supports both regular observations and state observations + +2. AUTOMATIC FEATURE EXTRACTION: + - Analyzes environment's observation and action managers automatically + - Handles nested observation structures with group-based access + - Automatically detects and processes video/image features + - Supports different action term types + +3. FLEXIBLE OBSERVATION HANDLING: + - Regular observations: saved as "observation.{key}" + - State observations: combined into "observation.state" + - Support for observations from different groups (policy, critic, etc.) + - Automatic tensor shape analysis and feature specification + +4. UNIVERSAL COMPATIBILITY: + - Works with any manager-based environment + - No hardcoded assumptions about observation or action structure + - Adapts to different environment types automatically + +USAGE PATTERNS: +============== + +1. Basic Configuration: + ```python + # Configure the environment first + env.cfg.lerobot_dataset.observation_keys_to_record = [("policy", "joint_pos"), ("policy", "camera_rgb")] + env.cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_vel")] + env.cfg.lerobot_dataset.task_description = "Stack the red cube on top of the blue cube" + + handler = LeRobotDatasetFileHandler() + handler.create("dataset.lerobot", env=env) + ``` + +2. State Observations: + ```python + # Configure state observations (combined into "observation.state") + env.cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos"), ("policy", "joint_vel")] + env.cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb"), ("policy", "end_effector_pos")] + ``` + +3. Multi-Group Observations: + ```python + # Configure observations from different groups + env.cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "joint_pos"), + ("policy", "camera_rgb"), + ("critic", "joint_vel") + ] + ``` + +4. Video/Image Support: + ```python + # Automatically detects and processes video/image data + env.cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb")] + # Handles [B, H, W, C] format and converts to [C, H, W] for LeRobot + ``` + +REQUIRED CONFIGURATION: +===================== + +The environment must have a LeRobotDatasetCfg configuration with at least one observation configured: + +```python +env.cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb"), ("policy", "end_effector_pos")] +env.cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos"), ("policy", "joint_vel")] +env.cfg.lerobot_dataset.task_description = "Custom task description" +``` + +This handler provides a streamlined way to record Isaac Lab environments to LeRobot datasets +with minimal configuration and maximum flexibility. +""" + import shutil -import tempfile import torch import numpy as np from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Callable, Union from collections.abc import Iterable try: from lerobot.common.datasets.lerobot_dataset import LeRobotDataset from lerobot.common.datasets.utils import ( - create_empty_dataset_info, - get_hf_features_from_features, - DEFAULT_FEATURES ) LEROBOT_AVAILABLE = True except ImportError: @@ -34,10 +138,58 @@ class LeRobotDatasetFileHandler(DatasetFileHandlerBase): - """LeRobot dataset file handler for storing and loading episode data.""" + """LeRobot dataset file handler for storing and loading episode data. + + This handler is designed to work with all manager-based environments by automatically + extracting features from the environment's observation and action managers. It provides + flexible configuration options for customizing how observations and actions are mapped + to LeRobot dataset features. + + Key Features: + - Automatic feature extraction from observation and action managers + - Support for nested observation groups and terms + - Flexible video/image detection and processing + - Customizable feature mapping + - Support for different action term types + - Configurable task description generation + - Support for state observations (saved as "observation.state") + - Configurable observation group selection + + Configuration Options: + - observation_keys_to_record: List of (group_name, observation_key) tuples to save as "observation.{key}" + - state_observation_keys: List of (group_name, observation_key) tuples to combine into "observation.state" + - task_description: String to use as the task description for all episodes + + Example Usage: + + ```python + # Basic usage with automatic feature extraction + handler = LeRobotDatasetFileHandler() + handler.create("my_dataset.lerobot", env_name="my_env", env=env) + + # Configure task description + env.cfg.lerobot_dataset.task_description = "Stack the red cube on top of the blue cube" + + # Configure state observations + env.cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos"), ("policy", "joint_vel")] + env.cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb"), ("policy", "end_effector_pos")] + + # Configure observations from different groups + env.cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "joint_pos"), + ("policy", "camera_rgb"), + ("critic", "joint_vel") + ] + ``` + """ - def __init__(self): - """Initializes the LeRobot dataset file handler.""" + def __init__(self, + config: Optional[Any] = None): + """Initialize the LeRobot dataset file handler. + + Args: + config: Optional configuration object from the environment. + """ if not LEROBOT_AVAILABLE: raise ImportError( "LeRobot dependencies not available. Please install following the documentation here: " @@ -48,16 +200,32 @@ def __init__(self): self._dataset_path = None self._env_name = None self._episode_count = 0 + + # Store configuration from environment + self._config = config def create(self, file_path: str, env_name: str | None = None, env = None): - """Create a new dataset file by automatically extracting features from environment.""" + """Create a new dataset file by automatically extracting features from environment. + + This method analyzes the environment's observation and action managers to automatically + create a comprehensive feature schema for the LeRobot dataset. + + Args: + file_path: Path to the dataset file (will be created with .lerobot extension if not present) + env_name: Optional name for the environment (used in metadata) + env: The manager-based environment instance + """ if not file_path.endswith(".lerobot"): - file_path += ".lerobot" + raise ValueError("Dataset file path must end with .lerobot") self._dataset_path = Path(file_path) # Delete existing dataset if it exists if self._dataset_path.exists(): + # get confirmation from user + confirm = input(f"Dataset at {self._dataset_path} already exists. Do you want to remove it? (y/n): ") + if confirm != "y": + raise ValueError("Dataset already exists. Please remove it or use a different file path.") print(f"Removing existing dataset at {self._dataset_path}") shutil.rmtree(self._dataset_path) @@ -67,6 +235,20 @@ def create(self, file_path: str, env_name: str | None = None, env = None): # Initialize environment name self._env_name = env_name or "isaac_lab_env" + # Get configuration from environment if available + if env is not None and hasattr(env, 'cfg') and hasattr(env.cfg, 'lerobot_dataset'): + env_config = env.cfg.lerobot_dataset + + # Store the merged configuration + self._config = env_config + else: + # Error out if configuration does not exist + raise ValueError( + "LeRobot dataset configuration not found. " + "The environment must have a 'lerobot_dataset' configuration in env.cfg. " + "Please ensure the environment is properly configured with LeRobotDatasetCfg." + ) + # Extract features from environment features = self._extract_features_from_env(env) @@ -94,74 +276,237 @@ def create(self, file_path: str, env_name: str | None = None, env = None): self._episode_count = 0 def _extract_features_from_env(self, env) -> Dict[str, Dict]: - """Extract features schema from environment observations and actions.""" + """Extract features schema from environment observations and actions. + + This method automatically analyzes the environment's observation and action managers + to create a comprehensive feature schema for the LeRobot dataset. + + Args: + env: The manager-based environment instance + + Returns: + Dictionary mapping feature names to their specifications + """ if env is None: raise ValueError("Environment must be provided to extract features") features = {} + # Extract action features + features.update(self._extract_action_features(env)) + + # Extract observation features + features.update(self._extract_observation_features(env)) + + # Add annotation features + features.update(self._extract_annotation_features(env)) + + return features + + def _extract_action_features(self, env) -> Dict[str, Dict]: + """Extract action features from the action manager. + + This method handles both the main action tensor and individual action terms. + """ + features = {} + + # Add main action feature features["action"] = { "dtype": "float32", "shape": (env.action_manager.total_action_dim,), "names": None } - # Add hardcoded features - features["annotation.human.action.task_description"] = { - "dtype": "int64", - "shape": (1,), - "names": None - } + return features + + def _extract_observation_features(self, env) -> Dict[str, Dict]: + """Extract observation features from the observation manager. + + This method handles both concatenated observation groups and individual terms + within groups, automatically detecting video/image features. + Uses configuration to determine which observation keys to record. + """ + features = {} - # Get observation features from observation manager active terms + # Get observation sample to analyze structure obs_sample = env.observation_manager.compute() - print("obs_sample", obs_sample) - # Extract features from nested observation structure - for group_key, group_data in obs_sample.items(): - if isinstance(group_data, dict): - for obs_key, obs_value in group_data.items(): - if isinstance(obs_value, torch.Tensor): - if "cam" in obs_key: - print(f"Processing camera observation: {obs_key}") - print(f"Camera observation shape: {obs_value.shape}") - # For camera observations, remove batch dimension and use [C, H, W] format - # obs_value shape is typically [batch_size, H, W, C] or [H, W, C] - if obs_value.ndim == 4: # [batch_size, H, W, C] - height, width, channels = obs_value.shape[1], obs_value.shape[2], obs_value.shape[3] - print(f"4D camera observation - batch size: {obs_value.shape[0]}") - else: # [H, W, C] - height, width, channels = obs_value.shape[0], obs_value.shape[1], obs_value.shape[2] - print("3D camera observation - no batch dimension") - - print(f"Extracted dimensions - channels: {channels}, height: {height}, width: {width}") - features[f"observation.{obs_key}"] = { - "dtype": "video", - "shape": (channels, height, width), # LeRobot expects [C, H, W] - "names": ["channel", "height", "width"], - "video_info": { - "video.fps": int(1 / env.step_dt) - } - } - print(f"Added video feature for {obs_key} with shape {(channels, height, width)}") - elif "joint_pos" in obs_key: - print(f"Processing joint position observation: {obs_key}") - print(f"Joint position shape: {obs_value.shape}") - # State observation - remove batch dimension - if obs_value.ndim > 1: - shape = obs_value.shape[1:] # Remove batch dimension - print(f"Multi-dimensional joint position - original shape: {obs_value.shape}, reduced shape: {shape}") - else: - shape = obs_value.shape - print(f"Single-dimensional joint position - shape: {shape}") - features["observation.state"] = { - "dtype": "float32", - "shape": shape, - "names": None - } - + # Get the lists of observation keys to record from configuration + observation_keys_to_record = self._config.observation_keys_to_record + state_observation_keys = self._config.state_observation_keys + + # Validate configuration - ensure at least some observations are configured + if not observation_keys_to_record or not state_observation_keys: + raise ValueError( + "LeRobotDatasetCfg must have at least one observation configured. " + "Please set either observation_keys_to_record or state_observation_keys (or both). " + "Format: [('group_name', 'observation_key'), ...]" + ) + + # Track state observations to combine them + state_observations = [] + + # Process each (group_name, observation_key) tuple + for group_name, obs_key in observation_keys_to_record: + # Validate that the group exists + if group_name not in obs_sample: + available_groups = list(obs_sample.keys()) + raise ValueError( + f"Observation group '{group_name}' not found. " + f"Available groups: {available_groups}" + ) + + # Validate that the observation key exists in the group + if obs_key not in obs_sample[group_name]: + available_keys = list(obs_sample[group_name].keys()) + raise ValueError( + f"Observation key '{obs_key}' not found in group '{group_name}'. " + f"Available keys: {available_keys}" + ) + + value = obs_sample[group_name][obs_key] + if isinstance(value, torch.Tensor): + print(f"Processing observation: {group_name}.{obs_key}") + feature_name = f"observation.{obs_key}" + features[feature_name] = self._analyze_tensor_feature(value, env) + else: + raise ValueError(f"Observation {group_name}.{obs_key} is not a tensor") + + # Process state observations + for group_name, obs_key in state_observation_keys: + # Validate that the group exists + if group_name not in obs_sample: + available_groups = list(obs_sample.keys()) + raise ValueError( + f"State observation group '{group_name}' not found. " + f"Available groups: {available_groups}" + ) + + # Validate that the observation key exists in the group + if obs_key not in obs_sample[group_name]: + available_keys = list(obs_sample[group_name].keys()) + raise ValueError( + f"State observation key '{obs_key}' not found in group '{group_name}'. " + f"Available keys: {available_keys}" + ) + + value = obs_sample[group_name][obs_key] + if isinstance(value, torch.Tensor): + print(f"Processing state observation: {group_name}.{obs_key}") + state_observations.append((obs_key, value)) + else: + raise ValueError(f"State observation {group_name}.{obs_key} is not a tensor") + + # Create combined state feature if we have state observations + if state_observations: + if len(state_observations) == 1: + # Single state observation + key, value = state_observations[0] + features["observation.state"] = self._analyze_tensor_feature(value, env) + else: + # Multiple state observations - combine their features + total_dim = 0 + for key, value in state_observations: + # Calculate the flattened dimension for this state observation + if value.ndim > 0: + dim = value.shape[1] if value.ndim > 1 else 1 + else: + dim = 1 + total_dim += dim + + # Create combined state feature + features["observation.state"] = { + "dtype": "float32", + "shape": (total_dim,), + "names": None + } + print(f"Combined {len(state_observations)} state observations into single 'observation.state' feature with {total_dim} dimensions") + return features + def _extract_annotation_features(self, env) -> Dict[str, Dict]: + """Extract annotation features.""" + return { + "annotation.human.action.task_description": { + "dtype": "int64", + "shape": (1,), + "names": None + } + } + + def _analyze_tensor_feature(self, tensor: torch.Tensor, env) -> Dict[str, Any]: + """Analyze a tensor to determine its LeRobot feature specification. + + Automatically detects video/image features and handles different tensor shapes. + """ + # Remove batch dimension for feature specification + if tensor.ndim > 0: + shape = tensor.shape[1:] if tensor.ndim > 1 else (1,) + else: + shape = (1,) + + # Determine if this is a video/image feature + if self._is_video_feature(tensor): + return { + "dtype": "video", + "shape": self._get_video_shape(tensor), + "names": ["channel", "height", "width"], + "video_info": { + "video.fps": int(1 / env.step_dt) + } + } + else: + return { + "dtype": "float32", + "shape": shape, + "names": None + } + + def _is_video_feature(self, tensor: torch.Tensor) -> bool: + """Determine if a tensor represents video/image data. + + Image data is expected to be exactly 4 dimensions in [B, H, W, C] format. + Raises an error if tensor is not exactly 4D or not in [B, H, W, C] format. + """ + # Check if tensor has exactly 4 dimensions + if tensor.ndim == 4: + # Check if the last 2-3 dimensions could be spatial (height, width, channels) + spatial_dims = tensor.shape[-2:] + # Assume it's video if spatial dimensions are reasonable for images + is_likely_image = all(dim > 1 and dim < 10000 for dim in spatial_dims) + + if is_likely_image: + # Check if format is [B, H, W, C] + if not (tensor.shape[1] > 1 and tensor.shape[2] > 1 and tensor.shape[3] <= 4): + raise ValueError( + f"Image data must be in [B, H, W, C] format, but got shape {tensor.shape}. " + f"Expected format: batch_size > 0, height > 1, width > 1, channels <= 4" + ) + + return True + return False + + def _get_video_shape(self, tensor: torch.Tensor) -> tuple: + """Get the video shape in [C, H, W] format for LeRobot. + + Expects exactly 4D tensors in [B, H, W, C] format. + """ + if tensor.ndim != 4: + raise ValueError( + f"Video tensor must have exactly 4 dimensions, but got {tensor.ndim} " + f"dimensions with shape {tensor.shape}" + ) + + # Check if format is [B, H, W, C] + if not (tensor.shape[1] > 1 and tensor.shape[2] > 1 and tensor.shape[3] <= 4): + raise ValueError( + f"Video tensor must be in [B, H, W, C] format, but got shape {tensor.shape}. " + f"Expected format: batch_size > 0, height > 1, width > 1, channels <= 4" + ) + + # Convert from [B, H, W, C] to [C, H, W] for LeRobot + return (tensor.shape[3], tensor.shape[1], tensor.shape[2]) + def open(self, file_path: str, mode: str = "r"): """Open an existing dataset file.""" raise NotImplementedError("Open not implemented for LeRobot handler") @@ -181,7 +526,10 @@ def get_num_episodes(self) -> int: return self._episode_count def write_episode(self, episode: EpisodeData): - """Add an episode to the dataset.""" + """Add an episode to the dataset. + + Converts Isaac Lab episode data to LeRobot format and saves it to the dataset. + """ if self._dataset is None or episode.is_empty(): return @@ -190,44 +538,56 @@ def write_episode(self, episode: EpisodeData): self._episode_count += 1 def _convert_and_save_episode(self, episode: EpisodeData): - """Convert Isaac Lab episode data to LeRobot format and save it.""" + """Convert Isaac Lab episode data to LeRobot format and save it. + + This method processes each frame of the episode, converting observations and actions + to the appropriate LeRobot format. + """ episode_dict = episode.data - - print("episode_dict", episode_dict) - # Determine number of frames - num_frames = len(episode_dict["actions"]) - - - # Generate a task description - task = f"Pick the red cube and place it on the blue cube" + # Determine number of frames from actions + if "actions" not in episode_dict: + raise ValueError("No actions found in episode data") + if "obs" not in episode_dict: + raise ValueError("No observations found in episode data") + + # Get the number of frames from the actions tensor + actions_tensor = episode_dict["actions"] + num_frames = actions_tensor.shape[0] + + # Generate task description + task = self._config.task_description # Add frames one by one to the LeRobot dataset for frame_idx in range(num_frames): - frame_data = {} - - # Process episode data - for key, value in episode_dict.items(): - if key == "actions": - # Handle actions - frame_data["action"] = value[frame_idx].cpu().numpy() - elif key == "obs": - # Handle observations - value is already the obs dictionary - for obs_key, obs_value in value.items(): - if "cam" in obs_key: - # Convert camera data from [H, W, C] to [C, H, W] format for LeRobot - camera_frame = obs_value[frame_idx].cpu().numpy() - if camera_frame.ndim == 3: # [H, W, C] - camera_frame = camera_frame.transpose(2, 0, 1) # Convert to [C, H, W] - frame_data[f"observation.{obs_key}"] = camera_frame - elif "joint_pos" in obs_key: - frame_data["observation.state"] = obs_value[frame_idx].cpu().numpy() - - # Add hardcoded feature data - frame_data["annotation.human.action.task_description"] = np.array([0], dtype=np.int64) - - # Add frame to the dataset - self._dataset.add_frame(frame_data, task) + try: + frame_data = {} + + # Process actions if available + if "actions" in episode_dict: + actions_tensor = episode_dict["actions"] + frame_action = actions_tensor[frame_idx] + frame_data.update(self._process_actions(frame_action)) + + # Process observations if available + if "obs" in episode_dict: + obs_dict = episode_dict["obs"] + # Extract frame-specific observations + frame_obs = self._extract_frame_observations(obs_dict, frame_idx) + frame_data.update(self._process_observations(frame_obs)) + + # Add annotation data + frame_data["annotation.human.action.task_description"] = np.array([0], dtype=np.int64) + + print(f"Frame data: {frame_data}") + + # Add frame to the dataset + self._dataset.add_frame(frame_data, task) + + except Exception as e: + print(f"Error processing frame {frame_idx}: {e}") + print(f"Frame data keys: {list(frame_data.keys()) if 'frame_data' in locals() else 'N/A'}") + continue # Save the episode try: @@ -235,6 +595,189 @@ def _convert_and_save_episode(self, episode: EpisodeData): except Exception as e: print(f"Warning: Failed to save episode: {e}") + def _extract_frame_observations(self, obs_dict: Dict[str, Any], frame_idx: int) -> Dict[str, Any]: + """Extract observations for a specific frame from the batch tensor. + + Args: + obs_dict: Dictionary containing observation tensors with batch dimension + frame_idx: Index of the frame to extract + + Returns: + Dictionary containing observations for the specific frame + """ + frame_obs = {} + + # Get the lists of observation keys to record from configuration + observation_keys_to_record = self._config.observation_keys_to_record + state_observation_keys = self._config.state_observation_keys + + # Track state observations to combine them + state_observations = [] + + # Collect all unique observation keys that need to be extracted + all_obs_keys = set() + for group_name, obs_key in observation_keys_to_record: + all_obs_keys.add(obs_key) + for group_name, obs_key in state_observation_keys: + all_obs_keys.add(obs_key) + + # Extract observations from the correct groups + for group_name, obs_key in observation_keys_to_record + state_observation_keys: + if group_name in obs_dict and obs_key in obs_dict[group_name]: + try: + value = obs_dict[group_name][obs_key] + if value.ndim > 0 and frame_idx < value.shape[0]: + # Extract the frame from the batch dimension + frame_obs[obs_key] = value[frame_idx] + else: + # Handle 0D tensors or tensors without batch dimension + frame_obs[obs_key] = value + except Exception as e: + print(f"Error extracting observation for key '{obs_key}' from group '{group_name}' at frame {frame_idx}: {e}") + print(f"Value shape: {value.shape}") + # Skip this observation if there's an error + continue + else: + print(f"Warning: Observation key '{obs_key}' not found in group '{group_name}' in episode data") + + return frame_obs + + def _extract_frame_states(self, states_dict: Dict[str, Any], frame_idx: int) -> Dict[str, Any]: + """Extract states for a specific frame from the batch tensor. + + Args: + states_dict: Dictionary containing state tensors with batch dimension + frame_idx: Index of the frame to extract + + Returns: + Dictionary containing states for the specific frame + """ + frame_states = {} + + for key, value in states_dict.items(): + try: + if value.ndim > 0 and frame_idx < value.shape[0]: + # Extract the frame from the batch dimension + frame_states[key] = value[frame_idx] + else: + # Handle 0D tensors or tensors without batch dimension + frame_states[key] = value + except Exception as e: + print(f"Error extracting state for key '{key}' at frame {frame_idx}: {e}") + print(f"Value shape: {value.shape}") + # Skip this state if there's an error + continue + + return frame_states + + def _process_actions(self, action_tensor: torch.Tensor) -> Dict[str, np.ndarray]: + """Process actions for a single frame. + + Handles both the main action tensor and individual action terms. + """ + frame_data = {} + + # Convert tensor to numpy array + frame_data["action"] = action_tensor.cpu().numpy() + + # Process individual action terms if available + # Note: This would require access to the action manager to split actions + # For now, we'll just add the main action + + return frame_data + + def _process_observations(self, obs_dict: Dict[str, Any]) -> Dict[str, np.ndarray]: + """Process observations for a single frame. + + Uses configuration to determine which observation keys to record. + """ + frame_data = {} + + # Get the lists of observation keys to record from configuration + observation_keys_to_record = self._config.observation_keys_to_record + state_observation_keys = self._config.state_observation_keys + + # Track state observations to combine them + state_observations = [] + + # Process regular observations + for group_name, obs_key in observation_keys_to_record: + if obs_key in obs_dict: + try: + print(f"Processing observation: {group_name}.{obs_key}") + feature_name = f"observation.{obs_key}" + processed_value = self._process_observation_term(obs_key, obs_dict[obs_key]) + frame_data[feature_name] = processed_value + except Exception as e: + print(f"Error processing observation '{group_name}.{obs_key}': {e}") + continue + else: + print(f"Warning: Observation key '{obs_key}' not found in frame data for group '{group_name}'") + + # Process state observations + for group_name, obs_key in state_observation_keys: + if obs_key in obs_dict: + try: + print(f"Processing state observation: {group_name}.{obs_key}") + processed_value = self._process_observation_term(obs_key, obs_dict[obs_key]) + state_observations.append(processed_value) + except Exception as e: + print(f"Error processing state observation '{group_name}.{obs_key}': {e}") + continue + else: + print(f"Warning: State observation key '{obs_key}' not found in frame data for group '{group_name}'") + + # Combine state observations into a single "observation.state" feature + if state_observations: + if len(state_observations) == 1: + # Single state observation + frame_data["observation.state"] = state_observations[0] + else: + # Multiple state observations - concatenate them + # Assuming all state observations are 1D arrays + concatenated_state = np.concatenate(state_observations) + frame_data["observation.state"] = concatenated_state + print(f"Combined {len(state_observations)} state observations into single 'observation.state' feature") + + return frame_data + + def _process_states(self, states_dict: Dict[str, Any]) -> Dict[str, np.ndarray]: + """Process states for a single frame. + + States are typically privileged information that may not be available + in real-world scenarios. + """ + frame_data = {} + + for key, value in states_dict.items(): + try: + feature_name = f"state.{key}" + frame_data[feature_name] = value.cpu().numpy() + except Exception as e: + print(f"Error processing state '{key}': {e}") + continue + + return frame_data + + def _process_observation_term(self, term_name: str, tensor: torch.Tensor) -> np.ndarray: + """Process a single observation term. + + Uses default processing for all observation terms. + """ + # Default processing + numpy_array = tensor.cpu().numpy() + + # Handle video/image data + if self._is_video_feature(tensor): + # _is_video_feature already ensures tensor is 4D in [B, H, W, C] format + if numpy_array.ndim == 4: # [B, H, W, C] + numpy_array = numpy_array.transpose(0, 3, 1, 2) # Convert to [B, C, H, W] + else: + # This should never happen since _is_video_feature ensures 4D + raise ValueError(f"Unexpected video tensor dimensions for {term_name}: {numpy_array.ndim}, expected 4") + + return numpy_array + def load_episode(self, episode_name: str) -> EpisodeData | None: """Load episode data from the file.""" raise NotImplementedError("Load episode not implemented for LeRobot handler") diff --git a/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py index 074bff97380..ca5e2095298 100644 --- a/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py +++ b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py @@ -37,9 +37,9 @@ def create_test_episode(self, device): test_episode.success = True # Add some dummy observations and actions - test_episode.add("obs/joint_pos", torch.tensor([1.0, 2.0, 3.0], device=device)) - test_episode.add("obs/joint_pos", torch.tensor([1.1, 2.1, 3.1], device=device)) - test_episode.add("obs/joint_pos", torch.tensor([1.2, 2.2, 3.2], device=device)) + test_episode.add("obs/policy/joint_pos", torch.tensor([1.0, 2.0, 3.0], device=device)) + test_episode.add("obs/policy/joint_pos", torch.tensor([1.1, 2.1, 3.1], device=device)) + test_episode.add("obs/policy/joint_pos", torch.tensor([1.2, 2.2, 3.2], device=device)) test_episode.add("actions", torch.tensor([0.1, 0.2], device=device)) test_episode.add("actions", torch.tensor([0.3, 0.4], device=device)) @@ -58,13 +58,27 @@ def temp_dir(self): def test_import_available(self): """Test that LeRobot handler can be imported.""" assert LEROBOT_AVAILABLE, "LeRobot dependencies should be available for testing" - handler = LeRobotDatasetFileHandler() + + # Create handler with required configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [] + + handler = LeRobotDatasetFileHandler(config=config) assert handler is not None def test_create_dataset_file(self, temp_dir): """Test creating a new LeRobot dataset file.""" dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") - handler = LeRobotDatasetFileHandler() + + # Create handler with required configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [] + + handler = LeRobotDatasetFileHandler(config=config) # Test creating with .lerobot extension handler.create(dataset_file_path, "test_env_name") @@ -73,7 +87,7 @@ def test_create_dataset_file(self, temp_dir): # Test creating without extension (should add .lerobot) dataset_file_path_no_ext = os.path.join(temp_dir, f"{uuid.uuid4()}") - handler = LeRobotDatasetFileHandler() + handler = LeRobotDatasetFileHandler(config=config) handler.create(dataset_file_path_no_ext, "test_env_name") assert handler.get_env_name() == "test_env_name" handler.close() @@ -82,7 +96,14 @@ def test_create_dataset_file(self, temp_dir): def test_write_episode(self, temp_dir, device): """Test writing an episode to the LeRobot dataset.""" dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") - handler = LeRobotDatasetFileHandler() + + # Create handler with required configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [("policy", "joint_pos"), ("policy", "camera_rgb")] + config.state_observation_keys = [] + + handler = LeRobotDatasetFileHandler(config=config) handler.create(dataset_file_path, "test_env_name") test_episode = self.create_test_episode(device) @@ -98,10 +119,57 @@ def test_write_episode(self, temp_dir, device): handler.flush() handler.close() + @pytest.mark.parametrize("device", ["cpu"]) # Only test CPU for CI compatibility + def test_state_observations(self, temp_dir, device): + """Test that state observations are properly handled.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + + # Create handler with state observation configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.state_observation_keys = [("policy", "joint_pos"), ("policy", "joint_vel")] + config.observation_keys_to_record = [("policy", "camera_rgb")] + + handler = LeRobotDatasetFileHandler(config=config) + handler.create(dataset_file_path, "test_env_name") + + # Create test episode with state observations + test_episode = EpisodeData() + test_episode.seed = 42 + test_episode.success = True + + # Add state observations + test_episode.add("obs/policy/joint_pos", torch.tensor([1.0, 2.0, 3.0], device=device)) + test_episode.add("obs/policy/joint_pos", torch.tensor([1.1, 2.1, 3.1], device=device)) + test_episode.add("obs/policy/joint_vel", torch.tensor([0.1, 0.2, 0.3], device=device)) + test_episode.add("obs/policy/joint_vel", torch.tensor([0.11, 0.21, 0.31], device=device)) + + # Add regular observations + test_episode.add("obs/policy/camera_rgb", torch.tensor([[[[0.1, 0.2, 0.3]]]], device=device)) + test_episode.add("obs/policy/camera_rgb", torch.tensor([[[[0.11, 0.21, 0.31]]]], device=device)) + + # Add actions + test_episode.add("actions", torch.tensor([0.1, 0.2], device=device)) + test_episode.add("actions", torch.tensor([0.3, 0.4], device=device)) + + # Write the episode to the dataset + handler.write_episode(test_episode) + assert handler.get_num_episodes() == 1 + + handler.flush() + handler.close() + def test_get_properties(self, temp_dir): """Test getting dataset properties.""" dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") - handler = LeRobotDatasetFileHandler() + + # Create handler with required configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [] + + handler = LeRobotDatasetFileHandler(config=config) handler.create(dataset_file_path, "test_env_name") # Test environment name @@ -116,6 +184,102 @@ def test_get_properties(self, temp_dir): handler.close() + def test_missing_configuration_error(self, temp_dir): + """Test that appropriate errors are raised when configuration is missing.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + + # Test with both observation_keys_to_record and state_observation_keys empty (should cause an error) + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [] # Empty list + config.state_observation_keys = [] # Empty list + + handler = LeRobotDatasetFileHandler(config=config) + + # Create a mock environment for testing + class MockEnv: + def __init__(self): + self.step_dt = 0.01 + self.action_manager = type('ActionManager', (), { + 'total_action_dim': 7, + '_terms': {} + })() + self.observation_manager = type('ObservationManager', (), { + 'compute': lambda: {'policy': {'joint_pos': torch.tensor([[1.0, 2.0, 3.0]])}} + })() + + mock_env = MockEnv() + + # This should raise an error since both lists are empty + with pytest.raises(ValueError, match="must have at least one observation configured"): + handler.create(dataset_file_path, "test_env_name", env=mock_env) + + # Test with only observation_keys_to_record set (should work) + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [] # Empty list should work if other is set + + handler = LeRobotDatasetFileHandler(config=config) + + # This should work since we have at least one observation configured + handler.create(dataset_file_path, "test_env_name", env=mock_env) + handler.close() + + # Test with only state_observation_keys set (should work) + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [] # Empty list + config.state_observation_keys = [("policy", "joint_pos")] # Should work if other is set + + handler = LeRobotDatasetFileHandler(config=config) + + # This should work since we have at least one observation configured + handler.create(dataset_file_path, "test_env_name", env=mock_env) + handler.close() + + @pytest.mark.parametrize("device", ["cpu"]) # Only test CPU for CI compatibility + def test_multi_group_observations(self, temp_dir, device): + """Test that observations from multiple groups are properly handled.""" + dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") + + # Create handler with multi-group observation configuration + from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg + config = LeRobotDatasetCfg() + config.observation_keys_to_record = [ + ("policy", "joint_pos"), + ("policy", "camera_rgb"), + ("critic", "joint_vel") + ] + config.state_observation_keys = [("policy", "joint_pos"), ("critic", "joint_vel")] + + handler = LeRobotDatasetFileHandler(config=config) + handler.create(dataset_file_path, "test_env_name") + + # Create test episode with observations from multiple groups + test_episode = EpisodeData() + test_episode.seed = 42 + test_episode.success = True + + # Add observations from policy group + test_episode.add("obs/policy/joint_pos", torch.tensor([1.0, 2.0, 3.0], device=device)) + test_episode.add("obs/policy/joint_pos", torch.tensor([1.1, 2.1, 3.1], device=device)) + test_episode.add("obs/policy/camera_rgb", torch.tensor([[[[0.1, 0.2, 0.3]]]], device=device)) + test_episode.add("obs/policy/camera_rgb", torch.tensor([[[[0.11, 0.21, 0.31]]]], device=device)) + + # Add observations from critic group + test_episode.add("obs/critic/joint_vel", torch.tensor([0.1, 0.2, 0.3], device=device)) + test_episode.add("obs/critic/joint_vel", torch.tensor([0.11, 0.21, 0.31], device=device)) + + # Add actions + test_episode.add("actions", torch.tensor([0.1, 0.2], device=device)) + test_episode.add("actions", torch.tensor([0.3, 0.4], device=device)) + + # Write the episode to the dataset + handler.write_episode(test_episode) + assert handler.get_num_episodes() == 1 + + handler.flush() + handler.close() + if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py index b43155a7678..2f96795544c 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py @@ -12,6 +12,7 @@ from isaaclab.assets import ArticulationCfg from isaaclab.controllers.pink_ik_cfg import PinkIKControllerCfg from isaaclab.envs.mdp.actions.pink_actions_cfg import PinkInverseKinematicsActionCfg +from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg from isaaclab.managers import ObservationGroupCfg as ObsGroup from isaaclab.managers import ObservationTermCfg as ObsTerm from isaaclab.managers import SceneEntityCfg @@ -192,4 +193,18 @@ def __post_init__(self): # List of image observations in policy observations self.image_obs_list = ["table_cam"] # self.image_obs_list = ["table_cam", "wrist_cam"] + + # Configure LeRobot dataset recording + self.lerobot_dataset = LeRobotDatasetCfg( + # Record specific observation keys that are useful for training + observation_keys_to_record=[ + ["policy", "table_cam"] + ], + # State observations that should be combined into "observation.state" + state_observation_keys=[ + ["policy", "joint_pos"] + ], + # Task description for the dataset + task_description="Stack the red cube on top of the blue cube", + ) \ No newline at end of file From c8fb0155370382a0c62c4ecccf54b05082a3f6dc Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 16:17:33 -0700 Subject: [PATCH 09/14] add temporary design document --- LEROBOT_DATASET_HANDLER_DESIGN.md | 334 ++++++++++++++++++ .../datasets/lerobot_dataset_file_handler.py | 4 +- 2 files changed, 335 insertions(+), 3 deletions(-) create mode 100644 LEROBOT_DATASET_HANDLER_DESIGN.md diff --git a/LEROBOT_DATASET_HANDLER_DESIGN.md b/LEROBOT_DATASET_HANDLER_DESIGN.md new file mode 100644 index 00000000000..d68aea679f9 --- /dev/null +++ b/LEROBOT_DATASET_HANDLER_DESIGN.md @@ -0,0 +1,334 @@ +# LeRobot Dataset File Handler Design Document + +## Overview + +The LeRobot Dataset File Handler (`LeRobotDatasetFileHandler`) is a configuration-driven system for automatically extracting and recording episode data from Isaac Lab environments to the LeRobot dataset format. It provides a seamless bridge between Isaac Lab's manager-based environments and the HuggingFace LeRobot ecosystem, enabling efficient dataset creation for Vision-Language-Action (VLA) model training. + +## Architecture + +### Core Components + +1. **Configuration System** (`LeRobotDatasetCfg`) + - Defines which observations to record and how to organize them + - Supports both regular observations and state observations + - Configurable task descriptions + +2. **Feature Extraction Engine** + - Automatically analyzes environment observation and action managers + - Handles nested observation structures with group-based access + - Detects and processes video/image features automatically + +3. **Data Conversion Pipeline** + - Converts Isaac Lab episode data to LeRobot format + - Handles tensor shape transformations and data type conversions + - Manages video/image format conversions ([B, H, W, C] → [C, H, W]) + +4. **Dataset Management** + - Creates and manages LeRobot dataset files + - Handles episode writing and metadata management + - Provides efficient storage with MP4 videos and Parquet files + +## Configuration System + +### LeRobotDatasetCfg Structure + +```python +@configclass +class LeRobotDatasetCfg: + # Regular observations saved as "observation.{key}" + observation_keys_to_record: List[tuple[str, str]] = MISSING + + # State observations combined into "observation.state" + state_observation_keys: List[tuple[str, str]] = MISSING + + # Task description for all episodes + task_description: str = MISSING +``` + +### Configuration Patterns + +#### Basic Configuration +```python +env.cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "joint_pos"), + ("policy", "camera_rgb") +] +env.cfg.lerobot_dataset.state_observation_keys = [ + ("policy", "joint_vel") +] +env.cfg.lerobot_dataset.task_description = "Stack the red cube on top of the blue cube" +``` + +#### Multi-Group Observations +```python +env.cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "joint_pos"), # From policy group + ("policy", "camera_rgb"), # From policy group + ("critic", "joint_vel") # From critic group +] +``` + +#### Video/Image Support +```python +env.cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "camera_rgb") # Automatically detected as video +] +# Handles [B, H, W, C] format and converts to [C, H, W] for LeRobot +``` + +## Feature Extraction Process + +### 1. Environment Analysis +The handler automatically analyzes the environment's structure: + +```python +def _extract_features_from_env(self, env) -> Dict[str, Dict]: + features = {} + + # Extract action features + features.update(self._extract_action_features(env)) + + # Extract observation features + features.update(self._extract_observation_features(env)) + + # Add annotation features + features.update(self._extract_annotation_features(env)) + + return features +``` + +### 2. Action Feature Extraction +```python +def _extract_action_features(self, env) -> Dict[str, Dict]: + return { + "action": { + "dtype": "float32", + "shape": (env.action_manager.total_action_dim,), + "names": None + } + } +``` + +### 3. Observation Feature Extraction +The system processes observations based on configuration: + +- **Regular Observations**: Saved as `observation.{key}` +- **State Observations**: Combined into `observation.state` +- **Video Detection**: Automatically detects 4D tensors in [B, H, W, C] format + +### 4. Video/Image Processing +```python +def _is_video_feature(self, tensor: torch.Tensor) -> bool: + # Check if tensor has exactly 4 dimensions + if tensor.ndim == 4: + # Validate [B, H, W, C] format + if (tensor.shape[1] > 1 and tensor.shape[2] > 1 and tensor.shape[3] <= 4): + return True + return False +``` + +## Data Flow + +### Episode Recording Process + +1. **Episode Creation** + ```python + handler = LeRobotDatasetFileHandler() + handler.create("dataset.lerobot", env_name="my_env", env=env) + ``` + +2. **Feature Schema Generation** + - Analyzes environment observation and action managers + - Creates LeRobot-compatible feature specifications + - Validates configuration against available observations + +3. **Episode Writing** + ```python + handler.write_episode(episode_data) + ``` + +4. **Frame-by-Frame Processing** + ```python + def _convert_and_save_episode(self, episode: EpisodeData): + for frame_idx in range(num_frames): + frame_data = {} + + # Process actions + frame_data.update(self._process_actions(frame_action)) + + # Process observations + frame_obs = self._extract_frame_observations(obs_dict, frame_idx) + frame_data.update(self._process_observations(frame_obs)) + + # Add to dataset + self._dataset.add_frame(frame_data, task) + ``` + +## Integration with Isaac Lab + +### Environment Configuration +The handler integrates seamlessly with Isaac Lab's manager-based environments: + +```python +# In environment configuration +env_cfg.lerobot_dataset = LeRobotDatasetCfg() +env_cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb")] +env_cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos")] +env_cfg.lerobot_dataset.task_description = "Custom task" + +# In recorder configuration +env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler +``` + +### Recording Script Integration +The `record_demos.py` script automatically detects LeRobot format: + +```python +# Configure dataset format based on file extension +use_lerobot_format = args_cli.dataset_file.endswith('.lerobot') + +if use_lerobot_format: + env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler +else: + env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler +``` + +## Dataset Structure + +### LeRobot Format Organization +``` +dataset.lerobot/ +├── dataset_info.json # HuggingFace dataset metadata +├── state.json # Dataset state information +├── data/ # Parquet files with episode data +│ ├── train-00000-of-00001.parquet +│ └── ... +├── videos/ # Video files for camera observations +│ ├── episode_000000/ +│ │ ├── front.mp4 +│ │ ├── wrist.mp4 +│ │ └── ... +│ └── episode_000001/ +│ └── ... +└── meta/ # Additional metadata + └── info.json # Isaac Lab specific metadata +``` + +### Feature Naming Conventions +- **Camera observations**: `observation.images.{camera_position}` +- **Robot state**: `observation.state` +- **Regular observations**: `observation.{obs_key}` +- **Actions**: `action` +- **Episode metadata**: `episode_index`, `frame_index`, `timestamp`, `task` + +## Error Handling and Validation + +### Configuration Validation +```python +# Validate that required configuration exists +if not observation_keys_to_record or not state_observation_keys: + raise ValueError( + "LeRobotDatasetCfg must have at least one observation configured. " + "Please set either observation_keys_to_record or state_observation_keys (or both)." + ) + +# Validate observation groups and keys +if group_name not in obs_sample: + available_groups = list(obs_sample.keys()) + raise ValueError(f"Observation group '{group_name}' not found. Available groups: {available_groups}") +``` + +### Video Format Validation +```python +def _is_video_feature(self, tensor: torch.Tensor) -> bool: + if tensor.ndim == 4: + if not (tensor.shape[1] > 1 and tensor.shape[2] > 1 and tensor.shape[3] <= 4): + raise ValueError( + f"Image data must be in [B, H, W, C] format, but got shape {tensor.shape}. " + f"Expected format: batch_size > 0, height > 1, width > 1, channels <= 4" + ) + return True + return False +``` + +## Performance Considerations + +### Memory Management +- Efficient tensor processing with minimal memory overhead +- Automatic cleanup of temporary data structures +- Background video writing for large datasets + +### Storage Optimization +- MP4 compression for video observations +- Parquet format for efficient tabular data storage +- Incremental episode writing to avoid memory accumulation + +## Usage Examples + +### Basic Recording +```python +# Configure environment +env_cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb")] +env_cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos")] +env_cfg.lerobot_dataset.task_description = "Stack cubes" + +# Record dataset +./isaaclab.sh -p scripts/tools/record_demos.py \ + --task Isaac-Stack-Cube-Franka-IK-Rel-v0 \ + --teleop_device spacemouse \ + --dataset_file ./datasets/demo.lerobot +``` + +### Advanced Configuration +```python +# Multi-camera setup +env_cfg.lerobot_dataset.observation_keys_to_record = [ + ("policy", "camera_rgb"), # Main camera + ("policy", "camera_depth"), # Depth camera + ("policy", "wrist_camera"), # Wrist camera + ("policy", "end_effector_pos") # End effector position +] + +# Comprehensive state representation +env_cfg.lerobot_dataset.state_observation_keys = [ + ("policy", "joint_pos"), # Joint positions + ("policy", "joint_vel"), # Joint velocities + ("policy", "gripper_state") # Gripper state +] +``` + +## Dependencies and Installation + +### Required Dependencies +```bash +pip install datasets opencv-python imageio[ffmpeg] +``` + +### Optional Dependencies +- `huggingface_hub` for dataset sharing +- `lerobot` for training pipeline integration + +## Future Enhancements + +### Planned Features +1. **Dynamic Task Descriptions**: Support for episode-specific task descriptions +2. **Multi-Modal Support**: Enhanced support for audio and other sensor modalities +3. **Compression Options**: Configurable video compression settings +4. **Streaming Support**: Real-time dataset writing for long recording sessions +5. **Validation Tools**: Enhanced dataset validation and quality checks + +### Integration Improvements +1. **Direct LeRobot Training**: Seamless integration with LeRobot training pipelines +2. **HuggingFace Hub Integration**: Automated dataset upload and versioning +3. **Dataset Versioning**: Support for dataset versioning and incremental updates + +## Conclusion + +The LeRobot Dataset File Handler provides a robust, configuration-driven solution for creating LeRobot-compatible datasets from Isaac Lab environments. Its automatic feature extraction, flexible configuration system, and seamless integration with the Isaac Lab ecosystem make it an essential tool for VLA model training and dataset sharing within the robotics community. + +The handler's design emphasizes: +- **Ease of Use**: Minimal configuration required for basic usage +- **Flexibility**: Support for complex observation structures and multi-modal data +- **Performance**: Efficient storage and processing of large datasets +- **Compatibility**: Full compatibility with the LeRobot ecosystem +- **Extensibility**: Easy to extend for new data types and use cases \ No newline at end of file diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py index c9fcf1c2ca1..fef3ada8446 100644 --- a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -127,8 +127,6 @@ try: from lerobot.common.datasets.lerobot_dataset import LeRobotDataset - from lerobot.common.datasets.utils import ( - ) LEROBOT_AVAILABLE = True except ImportError: LEROBOT_AVAILABLE = False @@ -554,7 +552,7 @@ def _convert_and_save_episode(self, episode: EpisodeData): # Get the number of frames from the actions tensor actions_tensor = episode_dict["actions"] num_frames = actions_tensor.shape[0] - + # Generate task description task = self._config.task_description From 3c59b3235bbc8c64b51bc578ab2cb37000789692 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 16:19:58 -0700 Subject: [PATCH 10/14] add dependency --- LEROBOT_DATASET_HANDLER_DESIGN.md | 42 +++++-------------------------- 1 file changed, 6 insertions(+), 36 deletions(-) diff --git a/LEROBOT_DATASET_HANDLER_DESIGN.md b/LEROBOT_DATASET_HANDLER_DESIGN.md index d68aea679f9..d56a74e73ab 100644 --- a/LEROBOT_DATASET_HANDLER_DESIGN.md +++ b/LEROBOT_DATASET_HANDLER_DESIGN.md @@ -4,6 +4,12 @@ The LeRobot Dataset File Handler (`LeRobotDatasetFileHandler`) is a configuration-driven system for automatically extracting and recording episode data from Isaac Lab environments to the LeRobot dataset format. It provides a seamless bridge between Isaac Lab's manager-based environments and the HuggingFace LeRobot ecosystem, enabling efficient dataset creation for Vision-Language-Action (VLA) model training. +## Dependencies + +- `lerobot` for dataset interfacing: https://github.com/huggingface/lerobot + +https://github.com/huggingface/lerobot/issues/1398 + ## Architecture ### Core Components @@ -296,39 +302,3 @@ env_cfg.lerobot_dataset.state_observation_keys = [ ("policy", "gripper_state") # Gripper state ] ``` - -## Dependencies and Installation - -### Required Dependencies -```bash -pip install datasets opencv-python imageio[ffmpeg] -``` - -### Optional Dependencies -- `huggingface_hub` for dataset sharing -- `lerobot` for training pipeline integration - -## Future Enhancements - -### Planned Features -1. **Dynamic Task Descriptions**: Support for episode-specific task descriptions -2. **Multi-Modal Support**: Enhanced support for audio and other sensor modalities -3. **Compression Options**: Configurable video compression settings -4. **Streaming Support**: Real-time dataset writing for long recording sessions -5. **Validation Tools**: Enhanced dataset validation and quality checks - -### Integration Improvements -1. **Direct LeRobot Training**: Seamless integration with LeRobot training pipelines -2. **HuggingFace Hub Integration**: Automated dataset upload and versioning -3. **Dataset Versioning**: Support for dataset versioning and incremental updates - -## Conclusion - -The LeRobot Dataset File Handler provides a robust, configuration-driven solution for creating LeRobot-compatible datasets from Isaac Lab environments. Its automatic feature extraction, flexible configuration system, and seamless integration with the Isaac Lab ecosystem make it an essential tool for VLA model training and dataset sharing within the robotics community. - -The handler's design emphasizes: -- **Ease of Use**: Minimal configuration required for basic usage -- **Flexibility**: Support for complex observation structures and multi-modal data -- **Performance**: Efficient storage and processing of large datasets -- **Compatibility**: Full compatibility with the LeRobot ecosystem -- **Extensibility**: Easy to extend for new data types and use cases \ No newline at end of file From 62c495f0255d9810d937c467c1d8210b088d4773 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Fri, 27 Jun 2025 16:36:39 -0700 Subject: [PATCH 11/14] add lerobot extension name --- .../isaaclab/utils/datasets/lerobot_dataset_file_handler.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py index fef3ada8446..c58a6943831 100644 --- a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -214,7 +214,8 @@ def create(self, file_path: str, env_name: str | None = None, env = None): env: The manager-based environment instance """ if not file_path.endswith(".lerobot"): - raise ValueError("Dataset file path must end with .lerobot") + # add .lerobot extension + file_path += ".lerobot" self._dataset_path = Path(file_path) From 3fbc855eac2d4ea25397e2013371c102e69776f7 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Mon, 7 Jul 2025 13:07:35 -0700 Subject: [PATCH 12/14] remove lerbootcfg file --- LEROBOT_DATASET_HANDLER_DESIGN.md | 67 +++++++++++++------ docs/source/how-to/record_lerobot_datasets.md | 16 ++--- .../isaaclab_mimic/annotate_demos.py | 16 ++++- .../isaaclab_mimic/generate_dataset.py | 13 ++++ scripts/tools/record_demos.py | 16 +++++ .../isaaclab/envs/manager_based_env_cfg.py | 28 -------- .../mdp/actions/pink_task_space_actions.py | 2 - .../isaaclab/managers/recorder_manager.py | 27 +++++++- .../datasets/lerobot_dataset_file_handler.py | 47 ++++++++----- .../test_lerobot_dataset_file_handler.py | 36 +++++----- .../stack_pink_ik_abs_visuomotor_env_cfg.py | 24 +++---- 11 files changed, 184 insertions(+), 108 deletions(-) diff --git a/LEROBOT_DATASET_HANDLER_DESIGN.md b/LEROBOT_DATASET_HANDLER_DESIGN.md index d56a74e73ab..5cb91bd34f6 100644 --- a/LEROBOT_DATASET_HANDLER_DESIGN.md +++ b/LEROBOT_DATASET_HANDLER_DESIGN.md @@ -36,38 +36,60 @@ https://github.com/huggingface/lerobot/issues/1398 ## Configuration System -### LeRobotDatasetCfg Structure +### RecorderManagerBaseCfg Structure + +The LeRobot dataset configuration is now integrated into the `RecorderManagerBaseCfg` class: ```python @configclass -class LeRobotDatasetCfg: - # Regular observations saved as "observation.{key}" - observation_keys_to_record: List[tuple[str, str]] = MISSING +class RecorderManagerBaseCfg: + # Standard recorder configuration + dataset_file_handler_class_type: type = HDF5DatasetFileHandler + dataset_export_dir_path: str = "/tmp/isaaclab/logs" + dataset_filename: str = "dataset" + dataset_export_mode: DatasetExportMode = DatasetExportMode.EXPORT_ALL + export_in_record_pre_reset: bool = True + + # LeRobot dataset specific configuration + observation_keys_to_record: Optional[List[tuple[str, str]]] = None + """List of (group_name, observation_key) tuples to record as regular observations. - # State observations combined into "observation.state" - state_observation_keys: List[tuple[str, str]] = MISSING + These will be saved as "observation.{obs_key}" in the LeRobot format. + Example: [("policy", "joint_pos"), ("policy", "camera_rgb"), ("critic", "joint_vel")] + """ + + state_observation_keys: Optional[List[tuple[str, str]]] = None + """List of (group_name, observation_key) tuples that should be treated as state observations. + + These will be combined and saved as "observation.state" in the LeRobot format. + Example: [("policy", "joint_pos"), ("policy", "joint_vel")] + """ + + task_description: Optional[str] = None + """Task description for the LeRobot dataset. - # Task description for all episodes - task_description: str = MISSING + This description will be used for all episodes in the dataset. + Example: "Stack the red cube on top of the blue cube" + """ ``` ### Configuration Patterns #### Basic Configuration ```python -env.cfg.lerobot_dataset.observation_keys_to_record = [ +env.cfg.recorders.observation_keys_to_record = [ ("policy", "joint_pos"), ("policy", "camera_rgb") ] -env.cfg.lerobot_dataset.state_observation_keys = [ +env.cfg.recorders.state_observation_keys = [ ("policy", "joint_vel") ] -env.cfg.lerobot_dataset.task_description = "Stack the red cube on top of the blue cube" +env.cfg.recorders.task_description = "Stack the red cube on top of the blue cube" ``` #### Multi-Group Observations ```python -env.cfg.lerobot_dataset.observation_keys_to_record = [ +env.cfg.recorders.observation_keys_to_record = [ ("policy", "joint_pos"), # From policy group ("policy", "camera_rgb"), # From policy group ("critic", "joint_vel") # From critic group @@ -76,7 +98,7 @@ env.cfg.lerobot_dataset.observation_keys_to_record = [ #### Video/Image Support ```python -env.cfg.lerobot_dataset.observation_keys_to_record = [ +env.cfg.recorders.observation_keys_to_record = [ ("policy", "camera_rgb") # Automatically detected as video ] # Handles [B, H, W, C] format and converts to [C, H, W] for LeRobot @@ -177,10 +199,9 @@ The handler integrates seamlessly with Isaac Lab's manager-based environments: ```python # In environment configuration -env_cfg.lerobot_dataset = LeRobotDatasetCfg() -env_cfg.lerobot_dataset.observation_keys_to_record = [("policy", "camera_rgb")] -env_cfg.lerobot_dataset.state_observation_keys = [("policy", "joint_pos")] -env_cfg.lerobot_dataset.task_description = "Custom task" +env_cfg.recorders.observation_keys_to_record = [("policy", "camera_rgb")] +env_cfg.recorders.state_observation_keys = [("policy", "joint_pos")] +env_cfg.recorders.task_description = "Custom task" # In recorder configuration env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler @@ -232,10 +253,16 @@ dataset.lerobot/ ### Configuration Validation ```python # Validate that required configuration exists -if not observation_keys_to_record or not state_observation_keys: +if not observation_keys_to_record: + raise ValueError( + "RecorderManagerBaseCfg must have observation_keys_to_record configured. " + "Please set observation_keys_to_record with format: [('group_name', 'observation_key'), ...]" + ) + +if not state_observation_keys: raise ValueError( - "LeRobotDatasetCfg must have at least one observation configured. " - "Please set either observation_keys_to_record or state_observation_keys (or both)." + "RecorderManagerBaseCfg must have state_observation_keys configured. " + "Please set state_observation_keys with format: [('group_name', 'observation_key'), ...]" ) # Validate observation groups and keys diff --git a/docs/source/how-to/record_lerobot_datasets.md b/docs/source/how-to/record_lerobot_datasets.md index 712cc79d620..68d348b049c 100644 --- a/docs/source/how-to/record_lerobot_datasets.md +++ b/docs/source/how-to/record_lerobot_datasets.md @@ -92,28 +92,26 @@ LeRobot uses standardized naming conventions for observations: ## Configuration Options -The LeRobot dataset handler supports flexible configuration through the `LeRobotDatasetCfg`: +The LeRobot dataset handler supports flexible configuration through the `RecorderManagerBaseCfg`: ```python -from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg +from isaaclab.managers import RecorderManagerBaseCfg # Configure which observations to record -env_cfg.lerobot_dataset = LeRobotDatasetCfg() - -# Regular observations (saved as "observation.{key}") - REQUIRED -env_cfg.lerobot_dataset.observation_keys_to_record = [ +env_cfg.recorders.observation_keys_to_record = [ "policy/camera_rgb", "policy/end_effector_pos", "policy/gripper_state" ] # State observations (combined into "observation.state") - REQUIRED -env_cfg.lerobot_dataset.state_observation_keys = [ +env_cfg.recorders.state_observation_keys = [ "policy/joint_pos", - "policy/joint_vel", - "policy/joint_torque" + "policy/joint_vel" ] +# Task description for the dataset +env_cfg.recorders.task_description = "Stack the red cube on top of the blue cube" ``` **Important**: At least one of `observation_keys_to_record` or `state_observation_keys` must be configured with at least one observation. If both are empty, an error will be raised. If you don't want to record any state observations, use an empty list `[]` for `state_observation_keys` but ensure `observation_keys_to_record` has at least one entry. diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index d047d4f1ea0..cba4329ceeb 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -188,8 +188,22 @@ def main(): # Disable all termination terms env_cfg.terminations = None + # Store existing LeRobot configuration if present + existing_observation_keys = getattr(env_cfg.recorders, 'observation_keys_to_record', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_state_keys = getattr(env_cfg.recorders, 'state_observation_keys', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_task_description = getattr(env_cfg.recorders, 'task_description', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + # Set up recorder terms for mimic annotations - env_cfg.recorders: MimicRecorderManagerCfg = MimicRecorderManagerCfg() + env_cfg.recorders = MimicRecorderManagerCfg() + + # Restore LeRobot configuration if it existed + if existing_observation_keys is not None: + env_cfg.recorders.observation_keys_to_record = existing_observation_keys + if existing_state_keys is not None: + env_cfg.recorders.state_observation_keys = existing_state_keys + if existing_task_description is not None: + env_cfg.recorders.task_description = existing_task_description + if not args_cli.auto: # disable subtask term signals recorder term if in manual mode env_cfg.recorders.record_pre_step_subtask_term_signals = None diff --git a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py index f52d3468660..d606c646726 100644 --- a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py +++ b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py @@ -123,12 +123,25 @@ def main(): generation_num_trials=args_cli.generation_num_trials, ) + # Store existing LeRobot configuration if present + existing_observation_keys = getattr(env_cfg.recorders, 'observation_keys_to_record', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_state_keys = getattr(env_cfg.recorders, 'state_observation_keys', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_task_description = getattr(env_cfg.recorders, 'task_description', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + # Set dataset file handler based on format if use_lerobot_format: env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler else: env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler + # Restore LeRobot configuration if it existed + if existing_observation_keys is not None: + env_cfg.recorders.observation_keys_to_record = existing_observation_keys + if existing_state_keys is not None: + env_cfg.recorders.state_observation_keys = existing_state_keys + if existing_task_description is not None: + env_cfg.recorders.task_description = existing_task_description + # create environment env = gym.make(env_name, cfg=env_cfg).unwrapped diff --git a/scripts/tools/record_demos.py b/scripts/tools/record_demos.py index 3d4034bba1b..450b50e4e97 100644 --- a/scripts/tools/record_demos.py +++ b/scripts/tools/record_demos.py @@ -258,7 +258,23 @@ def main(): else: omni.log.info(f"Recording dataset in HDF5 format: {args_cli.dataset_file}") + # Store existing LeRobot configuration if present + existing_observation_keys = getattr(env_cfg.recorders, 'observation_keys_to_record', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_state_keys = getattr(env_cfg.recorders, 'state_observation_keys', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_task_description = getattr(env_cfg.recorders, 'task_description', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + + # Create new recorder configuration env_cfg.recorders = ActionStateRecorderManagerCfg() + + # Restore LeRobot configuration if it existed + if existing_observation_keys is not None: + env_cfg.recorders.observation_keys_to_record = existing_observation_keys + if existing_state_keys is not None: + env_cfg.recorders.state_observation_keys = existing_state_keys + if existing_task_description is not None: + env_cfg.recorders.task_description = existing_task_description + + # Set recorder configuration env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name env_cfg.recorders.dataset_export_mode = DatasetExportMode.EXPORT_SUCCEEDED_ONLY diff --git a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py index b0dda651f1a..dba246af090 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py @@ -35,27 +35,6 @@ class DefaultEventManagerCfg: reset_scene_to_default = EventTerm(func=mdp.reset_scene_to_default, mode="reset") -@configclass -class LeRobotDatasetCfg: - """Configuration for LeRobot dataset recording. - - This configuration specifies which observation keys should be recorded to the LeRobot dataset. - """ - - # List of (group_name, observation_key) tuples to record as regular observations - # e.g., [("policy", "joint_pos"), ("policy", "camera_rgb"), ("critic", "joint_vel")] - # These will be saved as "observation.{obs_key}" in the LeRobot format - observation_keys_to_record: List[tuple[str, str]] = MISSING - - # List of (group_name, observation_key) tuples that should be treated as state observations - # e.g., [("policy", "joint_pos"), ("policy", "joint_vel")] - # These will be saved as "observation.state" in the LeRobot format - state_observation_keys: List[tuple[str, str]] = MISSING - - # Custom task description generator function name - task_description: str = MISSING - - @configclass class ManagerBasedEnvCfg: """Base configuration of the environment.""" @@ -143,10 +122,3 @@ class ManagerBasedEnvCfg: xr: XrCfg | None = None """Configuration for viewing and interacting with the environment through an XR device.""" - - # LeRobot dataset configuration - lerobot_dataset: LeRobotDatasetCfg | None = None - """Configuration for LeRobot dataset recording. Defaults to LeRobotDatasetCfg(). - - This configuration specifies which observation keys should be recorded to the LeRobot dataset. - """ diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 3b76f5bcfd3..111395220ed 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -180,8 +180,6 @@ def process_actions(self, actions: torch.Tensor): controlled_frame_in_base_link_frame ) - print("controlled_frame_in_base_link_frame_pos", controlled_frame_in_base_link_frame_pos) - # Loop through each task and set the target for env_index, ik_controller in enumerate(self._ik_controllers): for task_index, task in enumerate(ik_controller.cfg.variable_input_tasks): diff --git a/source/isaaclab/isaaclab/managers/recorder_manager.py b/source/isaaclab/isaaclab/managers/recorder_manager.py index ea2c8101204..45407fb6df1 100644 --- a/source/isaaclab/isaaclab/managers/recorder_manager.py +++ b/source/isaaclab/isaaclab/managers/recorder_manager.py @@ -11,7 +11,7 @@ import torch from collections.abc import Sequence from prettytable import PrettyTable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, List from isaaclab.utils import configclass from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler @@ -50,6 +50,28 @@ class RecorderManagerBaseCfg: export_in_record_pre_reset: bool = True """Whether to export episodes in the record_pre_reset call.""" + # LeRobot dataset specific configuration + observation_keys_to_record: List[tuple[str, str]] = None + """List of (group_name, observation_key) tuples to record as regular observations. + + These will be saved as "observation.{obs_key}" in the LeRobot format. + Example: [("policy", "joint_pos"), ("policy", "camera_rgb"), ("critic", "joint_vel")] + """ + + state_observation_keys: List[tuple[str, str]] = None + """List of (group_name, observation_key) tuples that should be treated as state observations. + + These will be combined and saved as "observation.state" in the LeRobot format. + Example: [("policy", "joint_pos"), ("policy", "joint_vel")] + """ + + task_description: str = None + """Task description for the LeRobot dataset. + + This description will be used for all episodes in the dataset. + Example: "Stack the red cube on top of the blue cube" + """ + class RecorderTerm(ManagerTermBase): """Base class for recorder terms. @@ -475,6 +497,9 @@ def _prepare_terms(self): "dataset_export_dir_path", "dataset_export_mode", "export_in_record_pre_reset", + "observation_keys_to_record", + "state_observation_keys", + "task_description", ]: continue # check if term config is None diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py index c58a6943831..08e27e4cf67 100644 --- a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -234,18 +234,27 @@ def create(self, file_path: str, env_name: str | None = None, env = None): # Initialize environment name self._env_name = env_name or "isaac_lab_env" - # Get configuration from environment if available - if env is not None and hasattr(env, 'cfg') and hasattr(env.cfg, 'lerobot_dataset'): - env_config = env.cfg.lerobot_dataset + # Get configuration from environment's recorder manager if available + if env is not None and hasattr(env, 'cfg') and hasattr(env.cfg, 'recorders'): + recorder_config = env.cfg.recorders - # Store the merged configuration - self._config = env_config + # Check if this is a RecorderManagerBaseCfg with LeRobot configuration + if hasattr(recorder_config, 'observation_keys_to_record') and hasattr(recorder_config, 'state_observation_keys'): + # Store the configuration from recorder manager + self._config = recorder_config + else: + # Error out if configuration does not exist + raise ValueError( + "LeRobot dataset configuration not found in recorder manager. " + "The recorder manager must have 'observation_keys_to_record' and 'state_observation_keys' " + "attributes. Please ensure the recorder manager is properly configured with LeRobot dataset settings." + ) else: - # Error out if configuration does not exist + # Error out if environment or recorder configuration does not exist raise ValueError( - "LeRobot dataset configuration not found. " - "The environment must have a 'lerobot_dataset' configuration in env.cfg. " - "Please ensure the environment is properly configured with LeRobotDatasetCfg." + "Environment or recorder configuration not found. " + "The environment must have a 'recorders' configuration with LeRobot dataset settings. " + "Please ensure the environment is properly configured." ) # Extract features from environment @@ -333,13 +342,21 @@ def _extract_observation_features(self, env) -> Dict[str, Dict]: # Get the lists of observation keys to record from configuration observation_keys_to_record = self._config.observation_keys_to_record state_observation_keys = self._config.state_observation_keys + + print(f"observation_keys_to_record: {observation_keys_to_record}") + print(f"state_observation_keys: {state_observation_keys}") + + # Validate configuration - ensure both observation types are configured + if not observation_keys_to_record: + raise ValueError( + "RecorderManagerBaseCfg must have observation_keys_to_record configured. " + "Please set observation_keys_to_record with format: [('group_name', 'observation_key'), ...]" + ) - # Validate configuration - ensure at least some observations are configured - if not observation_keys_to_record or not state_observation_keys: + if not state_observation_keys: raise ValueError( - "LeRobotDatasetCfg must have at least one observation configured. " - "Please set either observation_keys_to_record or state_observation_keys (or both). " - "Format: [('group_name', 'observation_key'), ...]" + "RecorderManagerBaseCfg must have state_observation_keys configured. " + "Please set state_observation_keys with format: [('group_name', 'observation_key'), ...]" ) # Track state observations to combine them @@ -555,7 +572,7 @@ def _convert_and_save_episode(self, episode: EpisodeData): num_frames = actions_tensor.shape[0] # Generate task description - task = self._config.task_description + task = self._config.task_description or "Isaac Lab task" # Add frames one by one to the LeRobot dataset for frame_idx in range(num_frames): diff --git a/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py index ca5e2095298..fa526136072 100644 --- a/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py +++ b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py @@ -60,10 +60,10 @@ def test_import_available(self): assert LEROBOT_AVAILABLE, "LeRobot dependencies should be available for testing" # Create handler with required configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [("policy", "joint_pos")] - config.state_observation_keys = [] + config.state_observation_keys = [("policy", "joint_vel")] handler = LeRobotDatasetFileHandler(config=config) assert handler is not None @@ -73,10 +73,10 @@ def test_create_dataset_file(self, temp_dir): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Create handler with required configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [("policy", "joint_pos")] - config.state_observation_keys = [] + config.state_observation_keys = [("policy", "joint_vel")] handler = LeRobotDatasetFileHandler(config=config) @@ -98,8 +98,8 @@ def test_write_episode(self, temp_dir, device): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Create handler with required configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [("policy", "joint_pos"), ("policy", "camera_rgb")] config.state_observation_keys = [] @@ -125,8 +125,8 @@ def test_state_observations(self, temp_dir, device): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Create handler with state observation configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.state_observation_keys = [("policy", "joint_pos"), ("policy", "joint_vel")] config.observation_keys_to_record = [("policy", "camera_rgb")] @@ -164,8 +164,8 @@ def test_get_properties(self, temp_dir): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Create handler with required configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [("policy", "joint_pos")] config.state_observation_keys = [] @@ -189,8 +189,8 @@ def test_missing_configuration_error(self, temp_dir): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Test with both observation_keys_to_record and state_observation_keys empty (should cause an error) - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [] # Empty list config.state_observation_keys = [] # Empty list @@ -215,7 +215,7 @@ def __init__(self): handler.create(dataset_file_path, "test_env_name", env=mock_env) # Test with only observation_keys_to_record set (should work) - config = LeRobotDatasetCfg() + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [("policy", "joint_pos")] config.state_observation_keys = [] # Empty list should work if other is set @@ -226,7 +226,7 @@ def __init__(self): handler.close() # Test with only state_observation_keys set (should work) - config = LeRobotDatasetCfg() + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [] # Empty list config.state_observation_keys = [("policy", "joint_pos")] # Should work if other is set @@ -242,8 +242,8 @@ def test_multi_group_observations(self, temp_dir, device): dataset_file_path = os.path.join(temp_dir, f"{uuid.uuid4()}.lerobot") # Create handler with multi-group observation configuration - from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg - config = LeRobotDatasetCfg() + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() config.observation_keys_to_record = [ ("policy", "joint_pos"), ("policy", "camera_rgb"), diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py index 2f96795544c..c01c62c9f80 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/stack_pink_ik_abs_visuomotor_env_cfg.py @@ -12,7 +12,6 @@ from isaaclab.assets import ArticulationCfg from isaaclab.controllers.pink_ik_cfg import PinkIKControllerCfg from isaaclab.envs.mdp.actions.pink_actions_cfg import PinkInverseKinematicsActionCfg -from isaaclab.envs.manager_based_env_cfg import LeRobotDatasetCfg from isaaclab.managers import ObservationGroupCfg as ObsGroup from isaaclab.managers import ObservationTermCfg as ObsTerm from isaaclab.managers import SceneEntityCfg @@ -194,17 +193,14 @@ def __post_init__(self): self.image_obs_list = ["table_cam"] # self.image_obs_list = ["table_cam", "wrist_cam"] - # Configure LeRobot dataset recording - self.lerobot_dataset = LeRobotDatasetCfg( - # Record specific observation keys that are useful for training - observation_keys_to_record=[ - ["policy", "table_cam"] - ], - # State observations that should be combined into "observation.state" - state_observation_keys=[ - ["policy", "joint_pos"] - ], - # Task description for the dataset - task_description="Stack the red cube on top of the blue cube", - ) + # Configure LeRobot dataset recording in recorder manager + self.recorders.observation_keys_to_record = [ + ("policy", "table_cam") + ] + # State observations that should be combined into "observation.state" + self.recorders.state_observation_keys = [ + ("policy", "joint_pos") + ] + # Task description for the dataset + self.recorders.task_description = "Stack the red cube on top of the blue cube" \ No newline at end of file From 8d425cf0e4de79833365c12c1cb4532b89beff3a Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Mon, 7 Jul 2025 14:26:44 -0700 Subject: [PATCH 13/14] move lerobot config to function --- .../isaaclab_mimic/generate_dataset.py | 13 ------------- .../isaaclab_mimic/datagen/generation.py | 6 ++++++ 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py index d606c646726..f52d3468660 100644 --- a/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py +++ b/scripts/imitation_learning/isaaclab_mimic/generate_dataset.py @@ -123,25 +123,12 @@ def main(): generation_num_trials=args_cli.generation_num_trials, ) - # Store existing LeRobot configuration if present - existing_observation_keys = getattr(env_cfg.recorders, 'observation_keys_to_record', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None - existing_state_keys = getattr(env_cfg.recorders, 'state_observation_keys', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None - existing_task_description = getattr(env_cfg.recorders, 'task_description', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None - # Set dataset file handler based on format if use_lerobot_format: env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler else: env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler - # Restore LeRobot configuration if it existed - if existing_observation_keys is not None: - env_cfg.recorders.observation_keys_to_record = existing_observation_keys - if existing_state_keys is not None: - env_cfg.recorders.state_observation_keys = existing_state_keys - if existing_task_description is not None: - env_cfg.recorders.task_description = existing_task_description - # create environment env = gym.make(env_name, cfg=env_cfg).unwrapped diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py index 70f7c7d19ce..be45e5ae5e3 100644 --- a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py +++ b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py @@ -182,7 +182,13 @@ def setup_env_config( env_cfg.observations.policy.concatenate_terms = False # Setup recorders + existing_observation_keys = getattr(env_cfg.recorders, 'observation_keys_to_record', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_state_keys = getattr(env_cfg.recorders, 'state_observation_keys', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None + existing_task_description = getattr(env_cfg.recorders, 'task_description', None) if hasattr(env_cfg, 'recorders') and env_cfg.recorders is not None else None env_cfg.recorders = ActionStateRecorderManagerCfg() + env_cfg.recorders.observation_keys_to_record = existing_observation_keys + env_cfg.recorders.state_observation_keys = existing_state_keys + env_cfg.recorders.task_description = existing_task_description env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name From 27f930c6d02fd35a454b6aec4d53ae6b372fed47 Mon Sep 17 00:00:00 2001 From: Ashwin Varghese Kuruttukulam Date: Mon, 7 Jul 2025 16:07:02 -0700 Subject: [PATCH 14/14] add recorder to annotation and remove prints --- .../isaaclab_mimic/annotate_demos.py | 15 +++- .../datasets/lerobot_dataset_file_handler.py | 88 +++++++------------ .../manipulation/stack/mdp/observations.py | 8 -- 3 files changed, 48 insertions(+), 63 deletions(-) diff --git a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py index cba4329ceeb..a683dba0188 100644 --- a/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py +++ b/scripts/imitation_learning/isaaclab_mimic/annotate_demos.py @@ -68,7 +68,7 @@ from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg from isaaclab.managers import RecorderTerm, RecorderTermCfg, TerminationTermCfg from isaaclab.utils import configclass -from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler +from isaaclab.utils.datasets import EpisodeData, HDF5DatasetFileHandler, LeRobotDatasetFileHandler import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils.parse_cfg import parse_env_cfg @@ -156,6 +156,13 @@ def main(): dataset_file_handler.open(args_cli.input_file) env_name = dataset_file_handler.get_env_name() episode_count = dataset_file_handler.get_num_episodes() + # Configure dataset format based on file extension + use_lerobot_format = args_cli.output_file.endswith('.lerobot') + + if use_lerobot_format: + print(f"Recording dataset in LeRobot format: {args_cli.output_file}") + else: + print(f"Recording dataset in HDF5 format: {args_cli.dataset_file}") if episode_count == 0: print("No episodes found in the dataset.") @@ -207,6 +214,12 @@ def main(): if not args_cli.auto: # disable subtask term signals recorder term if in manual mode env_cfg.recorders.record_pre_step_subtask_term_signals = None + + # Set dataset file handler based on format + if use_lerobot_format: + env_cfg.recorders.dataset_file_handler_class_type = LeRobotDatasetFileHandler + else: + env_cfg.recorders.dataset_file_handler_class_type = HDF5DatasetFileHandler env_cfg.recorders.dataset_export_dir_path = output_dir env_cfg.recorders.dataset_filename = output_file_name diff --git a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py index 08e27e4cf67..7613bfe6427 100644 --- a/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -308,7 +308,7 @@ def _extract_features_from_env(self, env) -> Dict[str, Dict]: # Add annotation features features.update(self._extract_annotation_features(env)) - + return features def _extract_action_features(self, env) -> Dict[str, Dict]: @@ -408,35 +408,33 @@ def _extract_observation_features(self, env) -> Dict[str, Dict]: value = obs_sample[group_name][obs_key] if isinstance(value, torch.Tensor): - print(f"Processing state observation: {group_name}.{obs_key}") state_observations.append((obs_key, value)) else: raise ValueError(f"State observation {group_name}.{obs_key} is not a tensor") # Create combined state feature if we have state observations - if state_observations: - if len(state_observations) == 1: - # Single state observation - key, value = state_observations[0] - features["observation.state"] = self._analyze_tensor_feature(value, env) - else: - # Multiple state observations - combine their features - total_dim = 0 - for key, value in state_observations: - # Calculate the flattened dimension for this state observation - if value.ndim > 0: - dim = value.shape[1] if value.ndim > 1 else 1 - else: - dim = 1 - total_dim += dim - - # Create combined state feature - features["observation.state"] = { - "dtype": "float32", - "shape": (total_dim,), - "names": None - } - print(f"Combined {len(state_observations)} state observations into single 'observation.state' feature with {total_dim} dimensions") + if len(state_observations) == 1: + # Single state observation + key, value = state_observations[0] + features["observation.state"] = self._analyze_tensor_feature(value, env) + else: + # Multiple state observations - combine their features + total_dim = 0 + for key, value in state_observations: + # Calculate the flattened dimension for this state observation + if value.ndim > 0: + dim = value.shape[1] if value.ndim > 1 else 1 + else: + dim = 1 + total_dim += dim + + # Create combined state feature + features["observation.state"] = { + "dtype": "float32", + "shape": (total_dim,), + "names": None + } + print(f"Combined {len(state_observations)} state observations into single 'observation.state' feature with {total_dim} dimensions") return features @@ -595,8 +593,6 @@ def _convert_and_save_episode(self, episode: EpisodeData): # Add annotation data frame_data["annotation.human.action.task_description"] = np.array([0], dtype=np.int64) - print(f"Frame data: {frame_data}") - # Add frame to the dataset self._dataset.add_frame(frame_data, task) @@ -627,34 +623,20 @@ def _extract_frame_observations(self, obs_dict: Dict[str, Any], frame_idx: int) observation_keys_to_record = self._config.observation_keys_to_record state_observation_keys = self._config.state_observation_keys - # Track state observations to combine them - state_observations = [] - - # Collect all unique observation keys that need to be extracted - all_obs_keys = set() - for group_name, obs_key in observation_keys_to_record: - all_obs_keys.add(obs_key) - for group_name, obs_key in state_observation_keys: - all_obs_keys.add(obs_key) - # Extract observations from the correct groups for group_name, obs_key in observation_keys_to_record + state_observation_keys: - if group_name in obs_dict and obs_key in obs_dict[group_name]: + if obs_key in obs_dict: try: - value = obs_dict[group_name][obs_key] - if value.ndim > 0 and frame_idx < value.shape[0]: - # Extract the frame from the batch dimension - frame_obs[obs_key] = value[frame_idx] - else: - # Handle 0D tensors or tensors without batch dimension - frame_obs[obs_key] = value + value = obs_dict[obs_key] + # Extract the frame from the batch dimension + frame_obs[obs_key] = value[frame_idx] + except Exception as e: - print(f"Error extracting observation for key '{obs_key}' from group '{group_name}' at frame {frame_idx}: {e}") + print(f"Error extracting observation for key '{obs_key}' at frame {frame_idx}: {e}") print(f"Value shape: {value.shape}") - # Skip this observation if there's an error - continue + raise Exception(f"Error extracting observation for key '{obs_key}' at frame {frame_idx}: {e}") else: - print(f"Warning: Observation key '{obs_key}' not found in group '{group_name}' in episode data") + print(f"Warning: Observation key '{obs_key}' not found in episode data") return frame_obs @@ -720,28 +702,26 @@ def _process_observations(self, obs_dict: Dict[str, Any]) -> Dict[str, np.ndarra for group_name, obs_key in observation_keys_to_record: if obs_key in obs_dict: try: - print(f"Processing observation: {group_name}.{obs_key}") feature_name = f"observation.{obs_key}" processed_value = self._process_observation_term(obs_key, obs_dict[obs_key]) frame_data[feature_name] = processed_value except Exception as e: - print(f"Error processing observation '{group_name}.{obs_key}': {e}") + print(f"Error processing observation '{obs_key}': {e}") continue else: - print(f"Warning: Observation key '{obs_key}' not found in frame data for group '{group_name}'") + print(f"Warning: Observation key '{obs_key}' not found in frame data'") # Process state observations for group_name, obs_key in state_observation_keys: if obs_key in obs_dict: try: - print(f"Processing state observation: {group_name}.{obs_key}") processed_value = self._process_observation_term(obs_key, obs_dict[obs_key]) state_observations.append(processed_value) except Exception as e: print(f"Error processing state observation '{group_name}.{obs_key}': {e}") continue else: - print(f"Warning: State observation key '{obs_key}' not found in frame data for group '{group_name}'") + print(f"Warning: State observation key '{obs_key}' not found in frame data'") # Combine state observations into a single "observation.state" feature if state_observations: diff --git a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py index 5c1473b3d81..b9da9f84c75 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/mdp/observations.py @@ -282,20 +282,12 @@ def object_grasped( object_pos = object.data.root_pos_w end_effector_pos = ee_frame.data.target_pos_w[:, 0, :] pose_diff = torch.linalg.vector_norm(object_pos - end_effector_pos, dim=1) - - print(f"Object position: {object_pos}") - print(f"End effector position: {end_effector_pos}") - # print(f"Pose difference: {pose_diff}") - # print(f"Gripper joint positions: {robot.data.joint_pos[:, -1]}") - # print(f"Gripper open value: {gripper_open_val}") - # TODO: Move this to an observation that is specififc to so-100 grasped = torch.logical_and( pose_diff < diff_threshold, torch.abs(robot.data.joint_pos[:, -1] - gripper_open_val.to(env.device)) > gripper_threshold, ) - print(f"Grasped status: {grasped}") return grasped