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/LEROBOT_DATASET_HANDLER_DESIGN.md b/LEROBOT_DATASET_HANDLER_DESIGN.md new file mode 100644 index 00000000000..5cb91bd34f6 --- /dev/null +++ b/LEROBOT_DATASET_HANDLER_DESIGN.md @@ -0,0 +1,331 @@ +# 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. + +## Dependencies + +- `lerobot` for dataset interfacing: https://github.com/huggingface/lerobot + +https://github.com/huggingface/lerobot/issues/1398 + +## 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 + +### RecorderManagerBaseCfg Structure + +The LeRobot dataset configuration is now integrated into the `RecorderManagerBaseCfg` class: + +```python +@configclass +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. + + 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. + + 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.recorders.observation_keys_to_record = [ + ("policy", "joint_pos"), + ("policy", "camera_rgb") +] +env.cfg.recorders.state_observation_keys = [ + ("policy", "joint_vel") +] +env.cfg.recorders.task_description = "Stack the red cube on top of the blue cube" +``` + +#### Multi-Group Observations +```python +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 +] +``` + +#### Video/Image Support +```python +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 +``` + +## 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.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 +``` + +### 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: + 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( + "RecorderManagerBaseCfg must have state_observation_keys configured. " + "Please set state_observation_keys with format: [('group_name', 'observation_key'), ...]" + ) + +# 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 +] +``` 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..68d348b049c --- /dev/null +++ b/docs/source/how-to/record_lerobot_datasets.md @@ -0,0 +1,271 @@ +# 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` +- **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 `RecorderManagerBaseCfg`: + +```python +from isaaclab.managers import RecorderManagerBaseCfg + +# Configure which observations 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.recorders.state_observation_keys = [ + "policy/joint_pos", + "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. + +### 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: + +```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..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.") @@ -188,11 +195,31 @@ 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 + + # 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 @@ -313,6 +340,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..450b50e4e97 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,50 @@ 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}") + + # 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 + + # 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 +408,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 +443,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.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 + 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/manager_based_env_cfg.py b/source/isaaclab/isaaclab/envs/manager_based_env_cfg.py index f119b66e487..dba246af090 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 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..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 @@ -198,7 +198,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..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. @@ -161,14 +183,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 = {} @@ -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/__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/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..7613bfe6427 --- /dev/null +++ b/source/isaaclab/isaaclab/utils/datasets/lerobot_dataset_file_handler.py @@ -0,0 +1,793 @@ +# 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 + +""" +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 torch +import numpy as np +from pathlib import Path +from typing import Any, Dict, List, Optional, Callable, Union +from collections.abc import Iterable + +try: + from lerobot.common.datasets.lerobot_dataset import LeRobotDataset + 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. + + 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, + 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: " + "https://github.com/huggingface/lerobot" + ) + + self._dataset = None + 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. + + 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"): + # add .lerobot extension + file_path += ".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) + + # 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" + + # 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 + + # 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 environment or recorder configuration does not exist + raise ValueError( + "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 + 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. + + 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 + } + + 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 sample to analyze structure + obs_sample = env.observation_manager.compute() + + # 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'), ...]" + ) + + if not state_observation_keys: + raise ValueError( + "RecorderManagerBaseCfg must have state_observation_keys configured. " + "Please set state_observation_keys with 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): + 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 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") + + 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. + + Converts Isaac Lab episode data to LeRobot format and saves it 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. + + This method processes each frame of the episode, converting observations and actions + to the appropriate LeRobot format. + """ + episode_dict = episode.data + + # 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 or "Isaac Lab task" + + # Add frames one by one to the LeRobot dataset + for frame_idx in range(num_frames): + 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) + + # 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: + self._dataset.save_episode() + 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 + + # Extract observations from the correct groups + for group_name, obs_key in observation_keys_to_record + state_observation_keys: + if obs_key in obs_dict: + try: + 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}' at frame {frame_idx}: {e}") + print(f"Value shape: {value.shape}") + 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 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: + 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 '{obs_key}': {e}") + continue + else: + 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: + 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'") + + # 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") + + 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 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..fa526136072 --- /dev/null +++ b/source/isaaclab/test/utils/test_lerobot_dataset_file_handler.py @@ -0,0 +1,285 @@ +# 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/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)) + 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" + + # Create handler with required configuration + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [("policy", "joint_vel")] + + 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") + + # Create handler with required configuration + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + config.observation_keys_to_record = [("policy", "joint_pos")] + config.state_observation_keys = [("policy", "joint_vel")] + + handler = LeRobotDatasetFileHandler(config=config) + + # 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(config=config) + 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") + + # Create handler with required configuration + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + 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) + + # 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() + + @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.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + 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") + + # Create handler with required configuration + from isaaclab.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + 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 + 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() + + 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.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + 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 = RecorderManagerBaseCfg() + 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 = RecorderManagerBaseCfg() + 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.managers import RecorderManagerBaseCfg + config = RecorderManagerBaseCfg() + 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_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..4d4067bfa37 --- /dev/null +++ b/source/isaaclab_assets/isaaclab_assets/robots/so_100.py @@ -0,0 +1,71 @@ +# 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.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={ + "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=["shoulder_pan_joint", "shoulder_lift_joint", "elbow_flex_joint", "wrist_flex_joint", "wrist_roll_joint"], + effort_limit=None, + velocity_limit=None, + stiffness=None, + damping=None, + armature=0.0, + ), + "gripper": ImplicitActuatorCfg( + joint_names_expr=["gripper_joint"], + 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 SO‑100 robot arm.""" 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 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..b413fcc1f49 --- /dev/null +++ b/source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/so_100/__init__.py @@ -0,0 +1,44 @@ +# 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_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-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_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..a9fa60b5cef --- /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={ + "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}, + )) + 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=["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=["gripper_joint"], + open_command_expr={"gripper_joint": 0.5}, + close_command_expr={"gripper_joint": -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=False, + 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..b76e3d0bcf4 --- /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_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_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 + "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}, + )) + + # Set actions for the specific robot type (franka) + self.actions.arm_action = PinkInverseKinematicsActionCfg( + pink_controlled_joint_names=[ + "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=["gripper_joint"], + 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..c01c62c9f80 --- /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,206 @@ +# 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_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_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 + "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}, + ), + ) + + # Set actions for the specific robot type (SO100) + self.actions.arm_action = PinkInverseKinematicsActionCfg( + pink_controlled_joint_names=[ + "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=["gripper_joint"], + 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"] + + # 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 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..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 @@ -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.""" @@ -282,14 +282,11 @@ 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) - + # 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 - ) 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..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 @@ -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,21 @@ 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 + # stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) + stacked = torch.logical_and(height_diff - h_dist_c12 < height_threshold, stacked) + # 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 ) return stacked