Add dataclass and state machine for fine grained subtask tracking

isaaclab_arena/environments/arena_env_builder.py

@@ -125,12 +125,14 @@ def compose_manager_cfg(self) -> IsaacLabArenaManagerBasedRLEnvCfg:
             self.arena_env.scene.get_events_cfg(),
             task.get_events_cfg(),
             placement_event_cfg,
+            task.get_fine_grained_subtask_events_cfg(),
         )
         termination_cfg = combine_configclass_instances(
             "TerminationCfg",
             task.get_termination_cfg(),
             self.arena_env.scene.get_termination_cfg(),
             embodiment.get_termination_cfg(),
+            task.get_fine_grained_subtask_termination_cfg(),
         )
         actions_cfg = embodiment.get_action_cfg()
         xr_cfg = embodiment.get_xr_cfg()

isaaclab_arena/tasks/composite_task_base.py

@@ -4,6 +4,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import copy
+import dataclasses
 import numpy as np
 import torch
 import warnings
@@ -20,6 +21,7 @@
 from isaaclab_arena.metrics.metric_base import MetricBase
 from isaaclab_arena.metrics.metric_term_cfg import MetricTermCfg
 from isaaclab_arena.tasks.common.mimic_default_params import MIMIC_DATAGEN_CONFIG_DEFAULTS
+from isaaclab_arena.tasks.fine_grained_subtask import FineGrainedSubtask
 from isaaclab_arena.tasks.task_base import TaskBase
 from isaaclab_arena.utils.configclass import (
     check_configclass_field_duplicates,
@@ -360,6 +362,31 @@ def get_metrics(self) -> list[MetricBase]:
 
         return subtask_metrics
 
+    def get_own_fine_grained_subtasks(self) -> list[FineGrainedSubtask]:
+        """Composite-level FineGrainedSubtasks.
+
+        These are added on top of whatever FGSs the child subtasks declare and are not gated.
+        """
+        return []
+
+    def get_fine_grained_subtasks(self) -> list[FineGrainedSubtask]:
+        """Concatenate child subtasks's FineGrainedSubtasks with namespace prefixes.
+
+        Each child's FGS gets a new name (subtask_{i}/{original_name}) and a parent_subtask_idx = i tag.
+        """
+        fgs_list: list[FineGrainedSubtask] = []
+        for i, child in enumerate(self.subtasks):
+            for fgs in child.get_fine_grained_subtasks():
+                fgs_list.append(
+                    dataclasses.replace(
+                        fgs,
+                        name=f"subtask_{i}/{fgs.name}",
+                        parent_subtask_idx=i,
+                    )
+                )
+        fgs_list.extend(self.get_own_fine_grained_subtasks())
+        return fgs_list
+
     def _validate_consistent_mimic_eef_names(self, arm_mode: ArmMode) -> set[str]:
         "Check that all subtasks have the same Mimic eef_names."
         mimic_eef_names = set(self.subtasks[0].get_mimic_env_cfg(arm_mode).subtask_configs.keys())