Merge pull request #4 from ooctipus/flat_multi_task_2

Clean up mdp

Author: rebeccazhang0707

source/isaaclab/isaaclab/envs/manager_based_rl_env.py

@@ -339,7 +339,7 @@ def _configure_gym_env_spaces(self):
                     term_dict[term_name] = gym.spaces.Box(low=low, high=high, shape=term_dim)
                 self.single_observation_space[group_name] = gym.spaces.Dict(term_dict)
         # action space (unbounded since we don't impose any limits)
-        action_dim = self.action_manager.total_action_dim
+        action_dim = sum(self.action_manager.action_term_dim)
         self.single_action_space = gym.spaces.Box(low=-np.inf, high=np.inf, shape=(action_dim,))
 
         # batch the spaces for vectorized environments

source/isaaclab/isaaclab/managers/scene_entity_cfg.py

@@ -7,9 +7,12 @@
 
 from __future__ import annotations
 
+import re
 from dataclasses import MISSING
 from typing import TYPE_CHECKING
 
+import torch
+
 from isaaclab.utils import configclass
 
 if TYPE_CHECKING:
@@ -115,6 +118,29 @@ class for more details.
 
     """
 
+    groups: str | list[str] | None = None
+    """Regex pattern(s) matching task-group names from the :class:`EnvLayout`.
+
+    When set, :meth:`resolve` populates :attr:`env_ids` and :attr:`group_ids`
+    with indices covering only the matched groups' environments.
+    When ``None`` (default), both default to ``slice(None)`` (all envs).
+
+    Patterns are matched against :attr:`EnvLayout.group_names` using
+    :func:`re.fullmatch`.
+    """
+
+    env_ids: slice | torch.Tensor = slice(None)
+    """Global env indices for indexing into ``(num_envs, ...)`` output buffers.
+
+    Populated by :meth:`resolve`. Defaults to ``slice(None)`` (all envs).
+    """
+
+    group_ids: slice | torch.Tensor = slice(None)
+    """Indices into the asset's data buffer (the view).
+
+    Populated by :meth:`resolve`. Defaults to ``slice(None)`` (all envs).
+    """
+
     _resolved: bool = False
     """Internal flag to prevent double resolution."""
 
@@ -157,6 +183,9 @@ def resolve(self, scene: InteractiveScene):
         # convert object collection names to indices based on regex
         self._resolve_object_collection_names(scene)
 
+        # resolve group patterns into env_ids and group_ids
+        self._resolve_groups(scene)
+
     def _resolve_joint_names(self, scene: InteractiveScene):
         # convert joint names to indices based on regex
         if self.joint_names is not None or self.joint_ids != slice(None):
@@ -301,3 +330,28 @@ def _resolve_object_collection_names(self, scene: InteractiveScene):
                 if isinstance(self.object_collection_ids, int):
                     self.object_collection_ids = [self.object_collection_ids]
                 self.object_collection_names = [entity.object_names[i] for i in self.object_collection_ids]
+
+    def _resolve_groups(self, scene: InteractiveScene):
+        """Resolve group patterns into env_ids and group_ids via the scene layout.
+
+        Args:
+            scene: The interactive scene instance.
+
+        Raises:
+            ValueError: If no groups match the specified patterns.
+        """
+        layout = scene.layout
+        if self.groups is None:
+            return
+        if isinstance(self.groups, str):
+            self.groups = [self.groups]
+        matched = []
+        for pattern in self.groups:
+            for name in layout.group_names:
+                if re.fullmatch(pattern, name) and name not in matched:
+                    matched.append(name)
+        if not matched:
+            raise ValueError(f"No groups matched patterns {self.groups}. Available: {list(layout.group_names)}")
+        group_view = layout.get(matched, asset=self.name)
+        self.env_ids = group_view.env_ids
+        self.group_ids = group_view.group_ids

source/isaaclab/isaaclab/scene/__init__.pyi

@@ -7,30 +7,16 @@ __all__ = [
     "CloneCfg",
     "CloneGroup",
     "EnvLayout",
-    "ExclusionSet",
     "GroupView",
     "InclusionSet",
-    "IntersectionGroup",
     "InteractiveScene",
     "InteractiveSceneCfg",
-    "PatternGroup",
-    "PredicateGroup",
-    "PrefixGroup",
-    "SuffixGroup",
-    "UnionGroup",
 ]
 
 from .clone_cfg import (
     CloneCfg,
     CloneGroup,
-    ExclusionSet,
     InclusionSet,
-    IntersectionGroup,
-    PatternGroup,
-    PredicateGroup,
-    PrefixGroup,
-    SuffixGroup,
-    UnionGroup,
 )
 from .env_layout import EnvLayout, GroupView
 from .interactive_scene import InteractiveScene

source/isaaclab/isaaclab/scene/clone_cfg.py

@@ -3,29 +3,10 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Configuration classes for environment cloning partitioning.
-
-Built-in :class:`CloneGroup` descriptors
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-==========================  =============================================
-Descriptor                  Selection logic
-==========================  =============================================
-:class:`InclusionSet`       Explicit list of asset names
-:class:`ExclusionSet`       Everything *except* listed asset names
-:class:`PrefixGroup`        Assets whose name starts with a prefix
-:class:`SuffixGroup`        Assets whose name ends with a suffix
-:class:`PatternGroup`       Regex full-match on asset names
-:class:`PredicateGroup`     Arbitrary ``Callable[[str], bool]``
-:class:`UnionGroup`         Logical OR of child descriptors
-:class:`IntersectionGroup`  Logical AND of child descriptors
-==========================  =============================================
-"""
+"""Configuration classes for environment cloning partitioning."""
 
 from __future__ import annotations
 
-import re
-from collections.abc import Callable
 from dataclasses import MISSING
 
 from isaaclab.cloner.cloner_strategies import random as random_strategy
@@ -96,164 +77,6 @@ def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
         return [a for a in self.assets if a in known]
 
 
-@configclass
-class ExclusionSet(CloneGroup):
-    """Clone group that includes everything *except* the listed assets.
-
-    Useful when a group should contain most scene assets and only a
-    few should be excluded.
-
-    Example::
-
-        ExclusionSet(exclude=["ground_plane", "light"], weight=1)
-    """
-
-    exclude: list[str] = MISSING
-    """Asset names to *exclude* from this group."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        excluded = set(self.exclude)
-        return [a for a in all_asset_names if a not in excluded]
-
-
-@configclass
-class PrefixGroup(CloneGroup):
-    """Clone group that selects assets whose name starts with a prefix.
-
-    Example::
-
-        PrefixGroup(prefix="lift_", weight=1)
-        # matches "lift_table", "lift_object", ...
-    """
-
-    prefix: str = MISSING
-    """Assets whose name starts with this string are included."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        return [a for a in all_asset_names if a.startswith(self.prefix)]
-
-
-@configclass
-class SuffixGroup(CloneGroup):
-    """Clone group that selects assets whose name ends with a suffix.
-
-    Example::
-
-        SuffixGroup(suffix="_frame", weight=1)
-        # matches "ee_frame", "cabinet_frame", ...
-    """
-
-    suffix: str = MISSING
-    """Assets whose name ends with this string are included."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        return [a for a in all_asset_names if a.endswith(self.suffix)]
-
-
-@configclass
-class PatternGroup(CloneGroup):
-    """Clone group that selects assets matching any of the given regex patterns.
-
-    Each pattern is tested as a **full match** against the asset name
-    (equivalent to ``re.fullmatch``).  Standard ``re`` syntax is supported.
-
-    Example::
-
-        PatternGroup(patterns=["lift_.*", "cabinet"], weight=1)
-        # matches "lift_table", "lift_object", "cabinet"
-    """
-
-    patterns: list[str] = MISSING
-    """Regex patterns (full-match) to test against asset names."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        compiled = [re.compile(p) for p in self.patterns]
-        return [a for a in all_asset_names if any(r.fullmatch(a) for r in compiled)]
-
-
-@configclass
-class PredicateGroup(CloneGroup):
-    """Clone group defined by an arbitrary callable predicate.
-
-    The :attr:`predicate` receives each asset name and returns ``True``
-    to include it.  This is the most flexible built-in descriptor.
-
-    Example::
-
-        PredicateGroup(
-            predicate=lambda name: "sensor" not in name,
-            weight=2,
-        )
-    """
-
-    predicate: Callable[[str], bool] = MISSING
-    """Callable that returns ``True`` for asset names to include."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        return [a for a in all_asset_names if self.predicate(a)]
-
-
-@configclass
-class UnionGroup(CloneGroup):
-    """Clone group that takes the union of multiple child descriptors.
-
-    An asset is included if **any** child descriptor claims it.
-
-    Example::
-
-        UnionGroup(
-            groups=[
-                PrefixGroup(prefix="lift_"),
-                InclusionSet(assets=["shared_sensor"]),
-            ],
-            weight=1,
-        )
-    """
-
-    groups: list[CloneGroup] = MISSING
-    """Child descriptors whose results are merged (union)."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        seen: set[str] = set()
-        result: list[str] = []
-        for g in self.groups:
-            for a in g.resolve_assets(all_asset_names):
-                if a not in seen:
-                    seen.add(a)
-                    result.append(a)
-        return result
-
-
-@configclass
-class IntersectionGroup(CloneGroup):
-    """Clone group that takes the intersection of multiple child descriptors.
-
-    An asset is included only if **all** child descriptors claim it.
-
-    Example::
-
-        IntersectionGroup(
-            groups=[
-                PrefixGroup(prefix="lift_"),
-                ExclusionSet(exclude=["lift_debug_viz"]),
-            ],
-            weight=1,
-        )
-    """
-
-    groups: list[CloneGroup] = MISSING
-    """Child descriptors whose results are intersected."""
-
-    def resolve_assets(self, all_asset_names: list[str]) -> list[str]:
-        if not self.groups:
-            return []
-        sets = [set(g.resolve_assets(all_asset_names)) for g in self.groups]
-        common = sets[0]
-        for s in sets[1:]:
-            common &= s
-        return [a for a in all_asset_names if a in common]
-
-
 # ── top-level config ──────────────────────────────────────────────────────────