work

mani_skill2/agents/__init__.py

@@ -0,0 +1 @@
+from .registration import REGISTERED_AGENTS

mani_skill2/agents/registration.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import Dict
+
+from mani_skill2 import logger
+from mani_skill2.agents.base_agent import BaseAgent
+
+
+@dataclass
+class AgentSpec:
+    """Agent specifications. At the moment it is a simple wrapper around the agent_cls but the dataclass is used in case we may need additional metadata"""
+
+    agent_cls: type[BaseAgent]
+
+
+REGISTERED_AGENTS: Dict[str, AgentSpec] = {}
+
+
+def register_agent(override=False):
+    """A decorator to register agents into ManiSkill so they can be used easily by string uid.
+
+    Args:
+        uid (str): unique id of the agent.
+        override (bool): whether to override the agent if it is already registered.
+    """
+
+    def _register_agent(agent_cls: type[BaseAgent]):
+        if agent_cls.uid in REGISTERED_AGENTS:
+            if override:
+                logger.warn(f"Overriding registered agent {agent_cls.uid}")
+                REGISTERED_AGENTS.pop(agent_cls.uid)
+            else:
+                logger.warn(
+                    f"Agent {agent_cls.uid} is already registered. Skip registration."
+                )
+            return agent_cls
+
+        REGISTERED_AGENTS[agent_cls.uid] = AgentSpec(agent_cls=agent_cls)
+        return agent_cls
+
+    return _register_agent

mani_skill2/agents/robots/__init__.py

@@ -1,24 +1,24 @@
 from mani_skill2.agents.robots.anymal.anymal_c import ANYmalC
+
+from .allegro_hand import AllegroHandLeft, AllegroHandRight, AllegroHandRightTouch
 from .dclaw import DClaw
 from .fetch import Fetch
-from .mobile_panda import MobilePandaDualArm, MobilePandaSingleArm
 from .panda import Panda
 from .xarm import XArm7Ability
 from .xmate3 import Xmate3Robotiq
-from .allegro_hand import AllegroHandRightTouch, AllegroHandRight, AllegroHandLeft
 
-ROBOTS = {
-    "panda": Panda,
-    "mobile_panda_dual_arm": MobilePandaDualArm,
-    "mobile_panda_single_arm": MobilePandaSingleArm,
-    "xmate3_robotiq": Xmate3Robotiq,
-    "fetch": Fetch,
-    # Dexterous Hand
-    "dclaw": DClaw,
-    "xarm7_ability": XArm7Ability,
-    "allegro_hand_right": AllegroHandRight,
-    "allegro_hand_left": AllegroHandLeft,
-    "allegro_hand_right_touch": AllegroHandRightTouch,
-    # Locomotion
-    "anymal-c": ANYmalC,
-}
+# = {
+#     "panda": Panda,
+#     "mobile_panda_dual_arm": MobilePandaDualArm,
+#     "mobile_panda_single_arm": MobilePandaSingleArm,
+#     "xmate3_robotiq": Xmate3Robotiq,
+#     "fetch": Fetch,
+#     # Dexterous Hand
+#     "dclaw": DClaw,
+#     "xarm7_ability": XArm7Ability,
+#     "allegro_hand_right": AllegroHandRight,
+#     "allegro_hand_left": AllegroHandLeft,
+#     "allegro_hand_right_touch": AllegroHandRightTouch,
+#     # Locomotion
+#     "anymal-c": ANYmalC,
+# }

mani_skill2/agents/robots/allegro_hand/allegro.py

@@ -7,13 +7,12 @@
 from mani_skill2 import PACKAGE_ASSET_DIR
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
-from mani_skill2.utils.sapien_utils import (
-    get_obj_by_name,
-)
-from mani_skill2.utils.sapien_utils import get_objs_by_names
+from mani_skill2.agents.registration import register_agent
+from mani_skill2.utils.sapien_utils import get_obj_by_name, get_objs_by_names
 from mani_skill2.utils.structs.pose import vectorize_pose
 
 
+@register_agent()
 class AllegroHandRight(BaseAgent):
     uid = "allegro_hand_right"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_right_glb.urdf"
@@ -146,6 +145,7 @@ def palm_pose(self):
         return vectorize_pose(self.palm_link.pose)
 
 
+@register_agent()
 class AllegroHandLeft(AllegroHandRight):
     uid = "allegro_hand_left"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/allegro_hand_left.urdf"

mani_skill2/agents/robots/allegro_hand/allegro_touch.py

@@ -1,22 +1,24 @@
 import itertools
-from typing import List, Dict, Tuple, Optional
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import sapien
 import torch
 from sapien import physx
 
 from mani_skill2 import PACKAGE_ASSET_DIR
+from mani_skill2.agents.registration import register_agent
 from mani_skill2.agents.robots.allegro_hand.allegro import AllegroHandRight
 from mani_skill2.utils.sapien_utils import (
     compute_total_impulse,
-    get_multiple_pairwise_contacts,
     get_actors_contacts,
+    get_multiple_pairwise_contacts,
+    get_objs_by_names,
 )
-from mani_skill2.utils.sapien_utils import get_objs_by_names
 from mani_skill2.utils.structs.actor import Actor
 
 
+@register_agent()
 class AllegroHandRightTouch(AllegroHandRight):
     uid = "allegro_hand_right_touch"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/allegro/variation/allegro_hand_right_fsr_simple.urdf"

mani_skill2/agents/robots/anymal/anymal_c.py

@@ -3,6 +3,7 @@
 from mani_skill2 import PACKAGE_ASSET_DIR, format_path
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
+from mani_skill2.agents.registration import register_agent
 from mani_skill2.utils.sapien_utils import (
     apply_urdf_config,
     check_urdf_config,
@@ -11,6 +12,7 @@
 from mani_skill2.utils.structs.articulation import Articulation
 
 
+@register_agent()
 class ANYmalC(BaseAgent):
     uid = "anymal-c"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/anymal-c/urdf/anymal.urdf"

mani_skill2/agents/robots/dclaw/dclaw.py

@@ -6,17 +6,15 @@
 from mani_skill2 import PACKAGE_ASSET_DIR
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
-from mani_skill2.agents.utils import (
-    get_active_joint_indices,
-)
-from mani_skill2.utils.sapien_utils import (
-    get_objs_by_names,
-)
+from mani_skill2.agents.registration import register_agent
+from mani_skill2.agents.utils import get_active_joint_indices
+from mani_skill2.utils.sapien_utils import get_objs_by_names
 from mani_skill2.utils.structs.joint import Joint
 from mani_skill2.utils.structs.link import Link
 from mani_skill2.utils.structs.pose import vectorize_pose
 
 
+@register_agent()
 class DClaw(BaseAgent):
     uid = "dclaw"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/dclaw/dclaw_gripper_glb.urdf"

mani_skill2/agents/robots/fetch/fetch.py

@@ -9,6 +9,7 @@
 from mani_skill2 import PACKAGE_ASSET_DIR
 from mani_skill2.agents.base_agent import BaseAgent
 from mani_skill2.agents.controllers import *
+from mani_skill2.agents.registration import register_agent
 from mani_skill2.sensors.camera import CameraConfig
 from mani_skill2.utils.common import compute_angle_between, np_compute_angle_between
 from mani_skill2.utils.sapien_utils import (
@@ -28,6 +29,7 @@
 FETCH_UNIQUE_COLLISION_BIT = 1 << 30
 
 
+@register_agent()
 class Fetch(BaseAgent):
     uid = "fetch"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/fetch/fetch.urdf"
@@ -66,10 +68,28 @@ class Fetch(BaseAgent):
             near=0.01,
             far=10,
             entity_uid="gripper_link",
-        )
+        ),
     ]
     REACHABLE_DIST = 1.5
-    RESTING_QPOS = np.array([0, 0, 0, 0.386, 0, -0.370, 0.562, -1.032, 0.695, 0.955, -0.1, 2.077, 0, 0.015, 0.015])
+    RESTING_QPOS = np.array(
+        [
+            0,
+            0,
+            0,
+            0.386,
+            0,
+            -0.370,
+            0.562,
+            -1.032,
+            0.695,
+            0.955,
+            -0.1,
+            2.077,
+            0,
+            0.015,
+            0.015,
+        ]
+    )
 
     def __init__(self, *args, **kwargs):
         self.arm_joint_names = [