diff --git a/metasim/cfg/objects.py b/metasim/cfg/objects.py index a65c773e2..aa6cd5305 100644 --- a/metasim/cfg/objects.py +++ b/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.""" diff --git a/metasim/data/quick_start/assets/COMMON/frame/usd/frame.usd b/metasim/data/quick_start/assets/COMMON/frame/usd/frame.usd new file mode 100644 index 000000000..1bb428956 Binary files /dev/null and b/metasim/data/quick_start/assets/COMMON/frame/usd/frame.usd differ diff --git a/metasim/sim/isaaclab/isaaclab.py b/metasim/sim/isaaclab/isaaclab.py index 59a471a4f..942063cc3 100644 --- a/metasim/sim/isaaclab/isaaclab.py +++ b/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]: diff --git a/metasim/sim/isaaclab/isaaclab_helper.py b/metasim/sim/isaaclab/isaaclab_helper.py index 5fca92e76..12b55d5ac 100644 --- a/metasim/sim/isaaclab/isaaclab_helper.py +++ b/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(