infra: add PrimitiveFrameCfg; isaaclab: support it

metasim/cfg/objects.py

@@ -131,6 +131,22 @@ def density(self) -> float:
         return self.mass / (4 / 3 * math.pi * self.radius**3)
 
 
+@configclass
+class PrimitiveFrameCfg(RigidObjCfg):
+    """Primitive coordinate frame cfg."""
+
+    # TODO: This is object shouldn't inherit from RigidObjCfg?
+    name: str = MISSING
+    scale: float = 1.0
+    """Scale of the frame"""
+    base_link: str | tuple[str, str] | None = None
+    """Base link to attach the frame.
+        If ``None``, the frame will be attached to the world origin.
+        If a ``str``, the frame will be attached to the root link of the object specified by the name.
+        If a ``tuple[str, str]``, the frame will be attached to the object specified by the first str and the body link specified by the second str.
+    """
+
+
 @configclass
 class PrimitiveCylinderCfg(RigidObjCfg):
     """Primitive cylinder object cfg."""

metasim/sim/isaaclab/isaaclab.py

@@ -7,7 +7,7 @@
 import torch
 from loguru import logger as log
 
-from metasim.cfg.objects import ArticulationObjCfg, BaseObjCfg, RigidObjCfg
+from metasim.cfg.objects import ArticulationObjCfg, BaseObjCfg, PrimitiveFrameCfg, RigidObjCfg
 from metasim.cfg.scenario import ScenarioCfg
 from metasim.sim import BaseSimHandler, EnvWrapper, IdentityEnvWrapper
 from metasim.types import Action, EnvState, Extra, Obs, Reward, Success, TimeOut
@@ -128,6 +128,23 @@ def step(self, action: list[Action]) -> tuple[Obs, Reward, Success, TimeOut, Ext
         time_out = time_out.cpu()
         success = self.checker.check(self)
         states = self.get_states()
+
+        ## TODO: organize this
+        for obj in self.objects:
+            if isinstance(obj, PrimitiveFrameCfg):
+                if obj.base_link is None:
+                    pos = torch.zeros((self.num_envs, 3), device=self.device)
+                    rot = torch.zeros((self.num_envs, 4), device=self.device)
+                    rot[:, 0] = 1.0
+                elif isinstance(obj.base_link, str):
+                    pos, rot = (states.objects | states.robots)[obj.base_link].root_state[:, :7].split([3, 4], dim=-1)
+                else:
+                    base_obj_name = obj.base_link[0]
+                    base_body_name = obj.base_link[1]
+                    merged_states = states.objects | states.robots
+                    body_idx = merged_states[base_obj_name].body_names.index(base_body_name)
+                    pos, rot = merged_states[base_obj_name].body_state[:, body_idx, :7].split([3, 4], dim=-1)
+                self._set_object_pose(obj, pos, rot)
         return states, None, success, time_out, extras
 
     def reset(self, env_ids: list[int] | None = None) -> tuple[list[EnvState], Extra]:

metasim/sim/isaaclab/isaaclab_helper.py

@@ -9,6 +9,7 @@
     BaseObjCfg,
     PrimitiveCubeCfg,
     PrimitiveCylinderCfg,
+    PrimitiveFrameCfg,
     PrimitiveSphereCfg,
     RigidObjCfg,
 )
@@ -91,6 +92,21 @@ def add_object(env: "EmptyEnv", obj: BaseObjCfg) -> None:
                 )
             )
             return
+        if isinstance(obj, PrimitiveFrameCfg):
+            env.scene.rigid_objects[obj.name] = RigidObject(
+                RigidObjectCfg(
+                    prim_path=prim_path,
+                    spawn=sim_utils.UsdFileCfg(
+                        usd_path="metasim/data/quick_start/assets/COMMON/frame/usd/frame.usd",
+                        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+                            disable_gravity=True, kinematic_enabled=True
+                        ),  # fixed
+                        collision_props=None,  # no collision
+                        scale=obj.scale,
+                    ),
+                )
+            )
+            return
 
         ## File-based object
         usd_file_cfg = sim_utils.UsdFileCfg(