Adds APIs to Isaac Lab Mimic for supporting loco-manipulation datagen (#3992)

peterd-NV · kellyguo11 · web-flow · commit 681368278744 · 2025-12-01T14:05:22.000-08:00
# Description  This PR adds new APIs to Isaac Lab Mimic to expand support for loco-manipulation data generation. It adds: - Processing for body end effector to treat a robot's base as an eef during Mimic data generation. This enables the use of the same Mimic annotation and subtask interface to enable lower body movement. - An optional way to cleanly add custom recorders for Mimic data generation (useful for when users want to record beyond the action/state data provided by Isaac Lab's default recorder). - Interface for enabling a navigation p-controller during Mimic data generation for robot's with mobile bases. ## Type of change  - New feature (non-breaking change which adds functionality) ## Checklist - [x] I have read and understood the [contribution guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html) - [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with `./isaaclab.sh --format` - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [x] I have updated the changelog and the corresponding version in the extension's `config/extension.toml` file - [x] I have added my name to the `CONTRIBUTORS.md` or my name already exists there  --------- Signed-off-by: peterd-NV <peterd@nvidia.com> Co-authored-by: Kelly Guo <kellyg@nvidia.com>
diff --git a/source/isaaclab/config/extension.toml b/source/isaaclab/config/extension.toml
@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.48.8"
+version = "0.48.9"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"
diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst
@@ -1,6 +1,17 @@
 Changelog
 ---------
 
+
+0.48.9 (2025-11-21)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Add navigation state API to IsaacLabManagerBasedRLMimicEnv
+* Add optional custom recorder config to MimicEnvCfg
+
+
 0.48.8 (2025-10-15)
 ~~~~~~~~~~~~~~~~~~~
 
@@ -18,6 +29,7 @@ Changed
 
 * Changed import from ``isaaclab.sim.utils`` to ``isaaclab.sim.utils.stage`` to properly propagate the Isaac Sim stage context.
 
+
 0.48.6 (2025-11-18)
 ~~~~~~~~~~~~~~~~~~~
 
diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_mimic_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_mimic_env.py
@@ -11,6 +11,12 @@
 from isaaclab.envs import ManagerBasedRLEnv
 
 
+def optional_method(func):
+    """Decorator to mark a method as optional."""
+    func.__is_optional__ = True
+    return func
+
+
 class ManagerBasedRLMimicEnv(ManagerBasedRLEnv):
     """The superclass for the Isaac Lab Mimic environments.
 
@@ -156,3 +162,21 @@ def serialize(self):
         and used in utils/env_utils.py.
         """
         return dict(env_name=self.spec.id, type=2, env_kwargs=dict())
+
+    @optional_method
+    def get_navigation_state(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor]:
+        """
+        Optional method. Only required when using navigation controller locomanipulation data generation.
+
+        Gets the navigation state of the robot. Required when use of the navigation controller is
+        enabled. The navigation state includes a boolean flag "is_navigating" to indicate when the
+        robot is under control by the navigation controller, and a boolean flag "navigation_goal_reached"
+        to indicate when the navigation goal has been reached.
+
+        Args:
+            env_ids: The environment index to get the navigation state for. If None, all envs are considered.
+
+        Returns:
+            A dictionary that of navigation state flags (False or True).
+        """
+        raise NotImplementedError
diff --git a/source/isaaclab/isaaclab/envs/mimic_env_cfg.py b/source/isaaclab/isaaclab/envs/mimic_env_cfg.py
@@ -11,6 +11,7 @@
 """
 import enum
 
+from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg
 from isaaclab.utils import configclass
 
 
@@ -76,6 +77,9 @@ class DataGenConfig:
     use_skillgen: bool = False
     """Whether to use skillgen to generate motion trajectories."""
 
+    use_navigation_controller: bool = False
+    """Whether to use a navigation controller to generate loco-manipulation trajectories."""
+
 
 @configclass
 class SubTaskConfig:
@@ -308,3 +312,6 @@ class MimicEnvCfg:
 
     # List of configurations for subtask constraints
     task_constraint_configs: list[SubTaskConstraintConfig] = []
+
+    # Optional recorder configuration
+    mimic_recorder_config: RecorderManagerBaseCfg | None = None
diff --git a/source/isaaclab_mimic/config/extension.toml b/source/isaaclab_mimic/config/extension.toml
@@ -1,7 +1,7 @@
 [package]
 
 # Semantic Versioning is used: https://semver.org/
-version = "1.0.15"
+version = "1.0.16"
 
 # Description
 category = "isaaclab"
diff --git a/source/isaaclab_mimic/docs/CHANGELOG.rst b/source/isaaclab_mimic/docs/CHANGELOG.rst
@@ -1,6 +1,16 @@
 Changelog
 ---------
 
+
+1.0.16 (2025-11-10)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Add body end effector to Mimic data generation to enable loco-manipulation data generation when a navigation p-controller is provided.
+
+
 1.0.15 (2025-09-25)
 
 Fixed
diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/data_generator.py
@@ -7,11 +7,16 @@
 Base class for data generator.
 """
 import asyncio
+import copy
+import logging
 import numpy as np
 import torch
 from typing import Any
 
 import isaaclab.utils.math as PoseUtils
+
+logger = logging.getLogger(__name__)
+
 from isaaclab.envs import (
     ManagerBasedRLMimicEnv,
     MimicEnvCfg,
@@ -688,6 +693,10 @@ async def generate(  # noqa: C901
             eef_subtasks_done[eef_name] = False
 
         prev_src_demo_datagen_info_pool_size = 0
+
+        if self.env_cfg.datagen_config.use_navigation_controller:
+            was_navigating = False
+
         # While loop that runs per time step
         while True:
             async with self.src_demo_datagen_info_pool.asyncio_lock:
@@ -880,8 +889,54 @@ async def generate(  # noqa: C901
                 generated_actions.extend(exec_results["actions"])
                 generated_success = generated_success or exec_results["success"]
 
+            # Get the navigation state
+            if self.env_cfg.datagen_config.use_navigation_controller:
+                processed_nav_subtask = False
+                navigation_state = self.env.get_navigation_state(env_id)
+                assert navigation_state is not None, "Navigation state cannot be None when using navigation controller"
+                is_navigating = navigation_state["is_navigating"]
+                navigation_goal_reached = navigation_state["navigation_goal_reached"]
+
             for eef_name in self.env_cfg.subtask_configs.keys():
                 current_eef_subtask_step_indices[eef_name] += 1
+
+                # Execute locomanip navigation controller if it is enabled via the use_navigation_controller flag
+                if self.env_cfg.datagen_config.use_navigation_controller:
+                    if "body" not in self.env_cfg.subtask_configs.keys():
+                        error_msg = (
+                            'End effector with name "body" not found in subtask configs. "body" must be a valid end'
+                            " effector to use the navigation controller.\n"
+                        )
+                        logger.error(error_msg)
+                        raise RuntimeError(error_msg)
+
+                    # Repeat the last nav subtask action if the robot is navigating and hasn't reached the waypoint goal
+                    if (
+                        current_eef_subtask_step_indices["body"] == len(current_eef_subtask_trajectories["body"]) - 1
+                        and not processed_nav_subtask
+                    ):
+                        if is_navigating and not navigation_goal_reached:
+                            for name in self.env_cfg.subtask_configs.keys():
+                                current_eef_subtask_step_indices[name] -= 1
+                            processed_nav_subtask = True
+
+                    # Else skip to the end of the nav subtask if the robot has reached the waypoint goal before the end
+                    # of the human recorded trajectory
+                    elif was_navigating and not is_navigating and not processed_nav_subtask:
+                        number_of_steps_to_skip = len(current_eef_subtask_trajectories["body"]) - (
+                            current_eef_subtask_step_indices["body"] + 1
+                        )
+                        for name in self.env_cfg.subtask_configs.keys():
+                            if current_eef_subtask_step_indices[name] + number_of_steps_to_skip < len(
+                                current_eef_subtask_trajectories[name]
+                            ):
+                                current_eef_subtask_step_indices[name] = (
+                                    current_eef_subtask_step_indices[name] + number_of_steps_to_skip
+                                )
+                            else:
+                                current_eef_subtask_step_indices[name] = len(current_eef_subtask_trajectories[name]) - 1
+                        processed_nav_subtask = True
+
                 subtask_ind = current_eef_subtask_indices[eef_name]
                 if current_eef_subtask_step_indices[eef_name] == len(
                     current_eef_subtask_trajectories[eef_name]
@@ -923,6 +978,10 @@ async def generate(  # noqa: C901
                     else:
                         current_eef_subtask_step_indices[eef_name] = None
                         current_eef_subtask_indices[eef_name] += 1
+
+            if self.env_cfg.datagen_config.use_navigation_controller:
+                was_navigating = copy.deepcopy(is_navigating)
+
             # Check if all eef_subtasks_done values are True
             if all(eef_subtasks_done.values()):
                 break
diff --git a/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py b/source/isaaclab_mimic/isaaclab_mimic/datagen/generation.py
@@ -5,12 +5,15 @@
 
 import asyncio
 import contextlib
+import sys
 import torch
+import traceback
 from typing import Any
 
 from isaaclab.envs import ManagerBasedRLMimicEnv
 from isaaclab.envs.mdp.recorders.recorders_cfg import ActionStateRecorderManagerCfg
 from isaaclab.managers import DatasetExportMode, TerminationTermCfg
+from isaaclab.managers.recorder_manager import RecorderManagerBaseCfg
 
 from isaaclab_mimic.datagen.data_generator import DataGenerator
 from isaaclab_mimic.datagen.datagen_info_pool import DataGenInfoPool
@@ -47,14 +50,20 @@ async def run_data_generator(
     """
     global num_success, num_failures, num_attempts
     while True:
-        results = await data_generator.generate(
-            env_id=env_id,
-            success_term=success_term,
-            env_reset_queue=env_reset_queue,
-            env_action_queue=env_action_queue,
-            pause_subtask=pause_subtask,
-            motion_planner=motion_planner,
-        )
+        try:
+            results = await data_generator.generate(
+                env_id=env_id,
+                success_term=success_term,
+                env_reset_queue=env_reset_queue,
+                env_action_queue=env_action_queue,
+                pause_subtask=pause_subtask,
+                motion_planner=motion_planner,
+            )
+        except Exception as e:
+            sys.stderr.write(traceback.format_exc())
+            sys.stderr.flush()
+            raise e
+
         if bool(results["success"]):
             num_success += 1
         else:
@@ -141,6 +150,7 @@ def setup_env_config(
     num_envs: int,
     device: str,
     generation_num_trials: int | None = None,
+    recorder_cfg: RecorderManagerBaseCfg | None = None,
 ) -> tuple[Any, Any]:
     """Configure the environment for data generation.
 
@@ -180,7 +190,10 @@ def setup_env_config(
     env_cfg.observations.policy.concatenate_terms = False
 
     # Setup recorders
-    env_cfg.recorders = ActionStateRecorderManagerCfg()
+    if recorder_cfg is None:
+        env_cfg.recorders = ActionStateRecorderManagerCfg()
+    else:
+        env_cfg.recorders = recorder_cfg
     env_cfg.recorders.dataset_export_dir_path = output_dir
     env_cfg.recorders.dataset_filename = output_file_name
 
