Improved joint definitions. Arm training does not use the lift anymore.

Author: hello-robot-shehab

examples/rlearning/dm_control_ppo.ipynb

@@ -249,6 +249,9 @@
         "\"\"\"\n",
         "This is a training task that tries to push a cube to a target location using all the joints and actuators on the robot.\n",
         "\"\"\"\n",
+        "from functools import cache\n",
+        "\n",
+        "\n",
         "class StretchPushCubeTraining:\n",
         "    def __init__(self, physics: mujoco.Physics, push_cube_by:tuple[float,float,float]):\n",
         "        self.physics = physics\n",
@@ -269,14 +272,26 @@
         "\n",
         "        self.current_distance_to_target = float('inf')\n",
         "\n",
+        "    @cache\n",
         "    def _get_joints(self):\n",
-        "        \"\"\"Gets joints, but removes unnamed ones (which we probably don't care about)\"\"\"\n",
-        "        return [name for name in self.physics.named.model.name_jntadr.axes.row.names if name != \"\"]\n",
+        "        \"\"\"Gets joint names in MJCF\"\"\"\n",
+        "        return [name for j in self._get_actuators() for name in j.get_joint_names_in_mjcf()]\n",
+        "    \n",
+        "    @cache\n",
+        "    def _get_actuator_names(self):\n",
+        "        return [j.name for j in self._get_actuators()]\n",
+        "    \n",
+        "    @cache\n",
+        "    def _get_actuators(self):\n",
+        "        return Actuators.get_actuated_joints()\n",
         "\n",
+        "    @cache\n",
         "    def _get_cube_id(self):\n",
         "        return self.physics.model.name2id(\"object1\", \"body\")\n",
         "    def _get_cube_pos(self):\n",
         "        return self.physics.data.xpos[self._get_cube_id()]\n",
+        "    \n",
+        "    @cache\n",
         "    def _get_cube_original_pos(self):\n",
         "        return self.physics.model.body(\"object1\").pos\n",
         "    \n",
@@ -320,7 +335,7 @@
         "            time.sleep(time_until_next_step)\n",
         "\n",
         "        # Apply the action to the joints\n",
-        "        for index, name in enumerate([j.name for j in Actuators.get_arm_joints()]):\n",
+        "        for index, name in enumerate(self._get_actuator_names()):\n",
         "            self.physics.data.actuator(name).ctrl = action[index]\n",
         "        \n",
         "        # Step the simulation forward\n",
@@ -381,11 +396,11 @@
         "\n",
         "        self.current_distance_to_target = float('inf')\n",
         "\n",
-        "    def _get_joints(self):\n",
-        "        # As defined in stretch.xml MJCF:\n",
-        "        arm_joints = ['joint_arm_l0', 'joint_arm_l1', 'joint_arm_l2', 'joint_arm_l3', 'joint_gripper_slide', 'joint_lift', 'joint_wrist_pitch', 'joint_wrist_roll', 'joint_wrist_yaw']\n",
-        "        finger_joints = ['joint_gripper_finger_left_open', 'joint_gripper_finger_right_open',] #'rubber_left_x', 'rubber_left_y', 'rubber_right_x', 'rubber_right_y']\n",
-        "        return arm_joints"
+        "\n",
+        "    @cache\n",
+        "    def _get_actuators(self):\n",
+        "        return Actuators.get_arm_joints()\n",
+        "    \n"
       ]
     },
     {
@@ -704,7 +719,7 @@
         "spec = mujoco.MjSpec.from_file(xml_path)\n",
         "spec = mujoco.MjSpec.from_file(xml_path)\n",
         "spec.find_body(\"object1\").pos = stretchPushCubeTrainingArmOnly.target_position\n",
-        "spec.meshdir = \"../models/assets/\"\n",
+        "spec.meshdir = \"../../stretch_mujoco/models/assets/\"\n",
         "\"\"\"meshdir is relative here, it should be the same as spec.modelfiledir, but mujoco expects meshdir to be a relative dir? If you get a Not Found error, your relative path may be wrong.\"\"\"\n",
         "spec.texturedir = spec.meshdir\n",
         "spec.compile()\n",

stretch_mujoco/enums/actuators.py

@@ -5,7 +5,7 @@
 
 class Actuators(Enum):
     """
-    An enum for the joints defined in the URDF.
+    An enum for the joints defined in the MJCF (stretch.xml).
     """
 
     arm = 0
@@ -24,12 +24,33 @@ class Actuators(Enum):
     @classmethod
     def get_arm_joints(cls) -> list["Actuators"]:
         return [
-            cls.lift,
             cls.arm,
             cls.wrist_pitch,
             cls.wrist_roll,
             cls.wrist_yaw,
         ]
+    
+    @classmethod
+    def get_actuated_joints(cls) -> list["Actuators"]:
+        return [actuator for actuator in cls if actuator != cls.base_rotate and actuator!= cls.base_translate]
+    
+    def get_joint_names_in_mjcf(self):
+        """
+        An actuator may have multiple joints. Return their names here.
+        """
+        if self == Actuators.left_wheel_vel: return [ "joint_left_wheel"]
+        if self == Actuators.right_wheel_vel: return  ["joint_right_wheel"]
+        if self == Actuators.lift: return  ["joint_lift"]
+        if self == Actuators.arm: return ['joint_arm_l0', 'joint_arm_l1', 'joint_arm_l2', 'joint_arm_l3']
+        if self == Actuators.wrist_yaw: return  ["joint_wrist_yaw"]
+        if self == Actuators.wrist_pitch: return [ "joint_wrist_pitch"]
+        if self == Actuators.wrist_roll: return  ["joint_wrist_roll"]
+        if self == Actuators.gripper: return  ["joint_gripper_slide"]
+        if self == Actuators.head_pan: return [ "joint_head_pan"]
+        if self == Actuators.head_tilt: return  ["joint_head_tilt"]
+
+        raise NotImplementedError(f"Joint names for {self} are not defined.")
+
     def _get_status_attribute(self, is_position: bool, status: StretchStatus) -> float:
         attribute_name = "pos" if is_position else "vel"
         if self == Actuators.arm: