Add mimic and metrics support to sequential composite tasks (#337)

## Summary
Add Mimic datagen and metrics support for sequential composite tasks.

For Mimic:

1. Add feature to automatically combine Mimic subtask configs from
atomic tasks into single config for composite task
2. Add example of a mimic config for a sequential task
(`franka_put_and_close_door_task.py`)
3. Update annotate_demos.py script to call the `success_term` after each
step. This is required to make sure the sequential task's success state
machine is properly updated after each step.

For Metrics:

1. Add support for `get_metrics` in sequential composite tasks.
2. Add sequential task specific success/subtask success metrics.
4. Update existing arena metrics recorder term naming to enable metric
recorder terms to be renamed. This allows sequential composite tasks to
rename metrics with a subtask suffix.

Author: peterd-NV

isaaclab_arena/affordances/turnable.py

@@ -122,7 +122,7 @@ def turn_to_level(
         """
         assert (
             target_level >= -1 and target_level < self.num_levels
-        ), f"target_level must be between -1 and {self.num_levels-1}"
+        ), f"target_level must be between -1 and {self.num_levels - 1}"
 
         if asset_cfg is None:
             asset_cfg = SceneEntityCfg(self.name)

isaaclab_arena/metrics/object_moved.py

@@ -17,10 +17,10 @@
 class ObjectVelocityRecorder(RecorderTerm):
     """Records the linear velocity of an object for each sim step of an episode."""
 
-    name = "object_linear_velocity"
-
     def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv):
         super().__init__(cfg, env)
+        # Get name from config instead of class attribute
+        self.name = cfg.name
         self.object_name = cfg.object_name
 
     def record_post_step(self):
@@ -33,6 +33,7 @@ def record_post_step(self):
 @configclass
 class ObjectVelocityRecorderCfg(RecorderTermCfg):
     class_type: type[RecorderTerm] = ObjectVelocityRecorder
+    name: str = "object_linear_velocity"
     object_name: str = MISSING
 
 
@@ -44,7 +45,7 @@ class ObjectMovedRateMetric(MetricBase):
     """
 
     name = "object_moved_rate"
-    recorder_term_name = ObjectVelocityRecorder.name
+    recorder_term_name = "object_linear_velocity"
 
     def __init__(self, object: Asset, object_velocity_threshold: float = 0.5):
         """Initializes the object-moved rate metric.
@@ -59,7 +60,7 @@ def __init__(self, object: Asset, object_velocity_threshold: float = 0.5):
 
     def get_recorder_term_cfg(self) -> RecorderTermCfg:
         """Return the recorder term configuration for the object-moved rate metric."""
-        return ObjectVelocityRecorderCfg(object_name=self.object.name)
+        return ObjectVelocityRecorderCfg(name=self.recorder_term_name, object_name=self.object.name)
 
     def compute_metric_from_recording(self, recorded_metric_data: list[np.ndarray]) -> float:
         """Computes the object-moved rate from the recorded metric data.

isaaclab_arena/metrics/revolute_joint_moved_rate.py

@@ -18,10 +18,9 @@
 class RevoluteJointStateRecorder(RecorderTerm):
     """Records the openness of an object for each sim step of an episode."""
 
-    name = "revolute_joint_state"
-
     def __init__(self, cfg: RecorderTermCfg, env: ManagerBasedEnv):
         super().__init__(cfg, env)
+        self.name = cfg.name
         self.object = cfg.object
 
     def record_post_step(self):
@@ -32,6 +31,7 @@ def record_post_step(self):
 @configclass
 class JointStateRecorderCfg(RecorderTermCfg):
     class_type: type[RecorderTerm] = RevoluteJointStateRecorder
+    name: str = "revolute_joint_state"
     object: ObjectBase = MISSING
 
 
@@ -43,7 +43,7 @@ class RevoluteJointMovedRateMetric(MetricBase):
     """
 
     name = "revolute_joint_moved_rate"
-    recorder_term_name = RevoluteJointStateRecorder.name
+    recorder_term_name = "revolute_joint_state"
 
     def __init__(self, object: Openable, reset_joint_percentage: float, joint_percentage_delta_threshold: float = 0.05):
         """Initializes the door-moved rate metric.
@@ -62,7 +62,7 @@ def __init__(self, object: Openable, reset_joint_percentage: float, joint_percen
 
     def get_recorder_term_cfg(self) -> RecorderTermCfg:
         """Return the recorder term configuration for the revolute joint moved rate metric."""
-        return JointStateRecorderCfg(object=self.object)
+        return JointStateRecorderCfg(name=self.recorder_term_name, object=self.object)
 
     def compute_metric_from_recording(self, recorded_metric_data: list[np.ndarray]) -> float:
         """Computes the revolute joint moved rate from the recorded metric data.

isaaclab_arena/metrics/success_rate.py

@@ -15,10 +15,10 @@
 class SuccessRecorder(RecorderTerm):
     """Records whether an episode was successful just before the environment is reset."""
 
-    name = "success"
-
     def __init__(self, cfg, env):
         super().__init__(cfg, env)
+        # Get name from config instead of class attribute
+        self.name = cfg.name
         # We track the first reset for each environment
         self.first_reset = True
 
@@ -35,12 +35,14 @@ def record_pre_reset(self, env_ids):
         assert "success" in self._env.termination_manager.active_terms
         success_results = torch.zeros(len(env_ids), dtype=bool, device=self._env.device)
         success_results |= self._env.termination_manager.get_term("success")[env_ids]
+
         return self.name, success_results
 
 
 @configclass
 class SuccessRecorderCfg(RecorderTermCfg):
     class_type: type[RecorderTerm] = SuccessRecorder
+    name: str = "success"
 
 
 class SuccessRateMetric(MetricBase):
@@ -51,11 +53,11 @@ class SuccessRateMetric(MetricBase):
     """
 
     name = "success_rate"
-    recorder_term_name = SuccessRecorder.name
+    recorder_term_name = "success"
 
     def get_recorder_term_cfg(self) -> RecorderTermCfg:
         """Return the recorder term configuration for the success rate metric."""
-        return SuccessRecorderCfg()
+        return SuccessRecorderCfg(name=self.recorder_term_name)
 
     def compute_metric_from_recording(self, recorded_metric_data: list[np.ndarray]) -> float:
         """Gets the average success rate from a list of recorded success flags.

isaaclab_arena/scripts/imitation_learning/annotate_demos.py

@@ -320,6 +320,8 @@ def replay_episode(
                 continue
         action_tensor = torch.Tensor(action).reshape([1, action.shape[0]])
         env.step(torch.Tensor(action_tensor))
+        # Run the success term so that tasks with a success term dependent state machine can update their state.
+        success_term.func(env, **success_term.params)
     if success_term is not None:
         if not bool(success_term.func(env, **success_term.params)[0]):
             return False

isaaclab_arena/tasks/common/open_close_door_mimic.py

@@ -46,7 +46,7 @@ def __post_init__(self):
                 object_ref=self.openable_object_name,
                 # 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",
+                subtask_term_signal="move_to_door",
                 # Specifies time offsets for data generation when splitting a trajectory into
                 # subtask segments. Random offsets are added to the termination boundary.
                 subtask_term_offset_range=(10, 20),

isaaclab_arena/tasks/pick_and_place_task.py

@@ -9,7 +9,7 @@
 import isaaclab.envs.mdp as mdp_isaac_lab
 from isaaclab.envs.common import ViewerCfg
 from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
-from isaaclab.managers import SceneEntityCfg, TerminationTermCfg
+from isaaclab.managers import EventTermCfg, SceneEntityCfg, TerminationTermCfg
 from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg
 from isaaclab.utils import configclass
 
@@ -21,6 +21,7 @@
 from isaaclab_arena.tasks.task_base import TaskBase
 from isaaclab_arena.tasks.terminations import object_on_destination
 from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object
+from isaaclab_arena.utils.pose import PoseRange
 
 
 class PickAndPlaceTask(TaskBase):
@@ -30,19 +31,24 @@ def __init__(
         pick_up_object: Asset,
         destination_location: Asset,
         background_scene: Asset,
+        destination_object: Asset | None = None,
         episode_length_s: float | None = None,
         task_description: str | None = None,
+        reset_pose_range: PoseRange = PoseRange(),
     ):
         super().__init__(episode_length_s=episode_length_s)
         self.pick_up_object = pick_up_object
+        self.destination_object = destination_object
         self.background_scene = background_scene
         self.destination_location = destination_location
         self.scene_config = SceneCfg(
             pick_up_object_contact_sensor=self.pick_up_object.get_contact_sensor_cfg(
                 contact_against_prim_paths=[self.destination_location.get_prim_path()],
             ),
         )
-        self.events_cfg = None
+        self.events_cfg = PickPlaceEventsCfg(
+            pick_up_object=self.pick_up_object, reset_pose_range=reset_pose_range.to_dict()
+        )
         self.termination_cfg = self.make_termination_cfg()
         self.task_description = (
             f"Pick up the {pick_up_object.name}, and place it into the {destination_location.name}"
@@ -85,7 +91,7 @@ def get_mimic_env_cfg(self, arm_mode: ArmMode):
         return PickPlaceMimicEnvCfg(
             arm_mode=arm_mode,
             pick_up_object_name=self.pick_up_object.name,
-            destination_location_name=self.destination_location.name,
+            destination_location_name=self.destination_object.name,
         )
 
     def get_metrics(self) -> list[MetricBase]:
@@ -116,6 +122,28 @@ class TerminationsCfg:
     object_dropped: TerminationTermCfg = MISSING
 
 
+@configclass
+class PickPlaceEventsCfg:
+    """Configuration for Pick Up Object."""
+
+    reset_pick_up_object_pose: EventTermCfg = MISSING
+
+    def __init__(self, pick_up_object: Asset, reset_pose_range: dict[str, tuple[float, float]]):
+        self.reset_pick_up_object_pose = EventTermCfg(
+            func=mdp_isaac_lab.reset_root_state_uniform,
+            mode="reset",
+            params={
+                "pose_range": {
+                    "x": reset_pose_range["x"],
+                    "y": reset_pose_range["y"],
+                    "z": reset_pose_range["z"],
+                },
+                "velocity_range": {},
+                "asset_cfg": SceneEntityCfg(pick_up_object.name),
+            },
+        )
+
+
 @configclass
 class PickPlaceMimicEnvCfg(MimicEnvCfg):
     """

isaaclab_arena/tasks/sequential_composite_tasks/franka_put_and_close_door_task.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2025-2026, The Isaac Lab Arena Project Developers (https://github.com/isaac-sim/IsaacLab-Arena/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+
+from isaaclab.envs.common import ViewerCfg
+from isaaclab.envs.mimic_env_cfg import MimicEnvCfg
+from isaaclab.utils import configclass
+
+from isaaclab_arena.embodiments.common.arm_mode import ArmMode
+from isaaclab_arena.tasks.sequential_task_base import SequentialTaskBase
+from isaaclab_arena.tasks.task_base import TaskBase
+from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object
+
+
+class FrankaPutAndCloseDoorTask(SequentialTaskBase):
+
+    def __init__(
+        self,
+        openable_object,
+        subtasks: list[TaskBase],
+        episode_length_s: float | None = None,
+    ):
+        super().__init__(subtasks=subtasks, episode_length_s=episode_length_s)
+        self.openable_object = openable_object
+
+    def get_viewer_cfg(self) -> ViewerCfg:
+        return get_viewer_cfg_look_at_object(lookat_object=self.openable_object, offset=np.array([-1.3, -1.3, 1.3]))
+
+    def get_prompt(self) -> str:
+        return None
+
+    def get_mimic_env_cfg(self, arm_mode: ArmMode):
+        mimic_env_cfg = FrankaPutAndCloseDoorTaskMimicEnvCfg()
+        mimic_env_cfg.subtask_configs = self.combine_mimic_subtask_configs(arm_mode)
+        return mimic_env_cfg
+
+
+@configclass
+class FrankaPutAndCloseDoorTaskMimicEnvCfg(MimicEnvCfg):
+    """
+    Isaac Lab Mimic environment config class for Franka put and close door task.
+    """
+
+    def __post_init__(self):
+        # post init of parents
+        super().__post_init__()
+
+        # Override the existing values
+        self.datagen_config.name = "franka_put_and_close_door_task_D0"
+        self.datagen_config.generation_guarantee = True
+        self.datagen_config.generation_keep_failed = False
+        self.datagen_config.generation_num_trials = 100
+        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.seed = 1