Add options to get/set all robot joints, not just gripper/arm joints

Author: kevin-thankyou-lin

robosuite/robots/robot.py

@@ -2,7 +2,7 @@
 import json
 import os
 from collections import OrderedDict
-from typing import Optional
+from typing import Optional, Dict
 
 import numpy as np
 
@@ -328,6 +328,10 @@ def setup_references(self):
         self._ref_torso_joint_pos_indexes = [
             self.sim.model.get_joint_qpos_addr(x) for x in self.robot_model._torso_joints
         ]
+        # indices for all joints (includes arms, base, torso, but not gripper joints, as those are handled separately)
+        self._ref_all_joint_pos_indexes = [
+            self.sim.model.get_joint_qpos_addr(x) for x in self.robot_model.all_joints
+        ]
 
     def setup_observables(self):
         """
@@ -638,6 +642,58 @@ def set_gripper_joint_positions(self, jpos: np.ndarray, gripper_arm: Optional[st
         else:
             raise ValueError(f"No gripper found for arm {gripper_arm}")
 
+    def set_all_robot_joint_positions(self, joint_positions: Dict[str, float]):
+        """
+        Helper method to force robot joint positions to the passed values.
+        Assumes valid joint names are passed.
+
+        Args:
+            joint_positions (dict): joint name -> joint position (in angles / radians)
+        """
+        for joint_name, position in joint_positions.items():
+            self.sim.data.qpos[self.sim.model.joint_name2id(joint_name)] = position
+        self.sim.forward()
+
+    def get_robot_joint_positions(self):
+        """
+        Returns:
+            np.array: joint positions (in angles / radians)
+        """
+        return self.sim.data.qpos[self._ref_joint_pos_indexes]
+
+    def get_all_robot_joint_positions(self) -> Dict[str, float]:
+        """
+        Returns all joint positions including robot joints (arms, base, torso) and gripper joints.
+        
+        Returns:
+            dict: joint name -> joint position (in angles / radians)
+        """
+        joint_positions = {}
+        
+        # Get all robot joint positions (includes arms, base, torso, etc.)
+        all_joint_names = self.robot_model.all_joints
+        for i, joint_name in enumerate(all_joint_names):
+            joint_positions[joint_name] = self.sim.data.qpos[self._ref_all_joint_pos_indexes[i]]
+        
+        # Add gripper joint positions for all arms
+        for arm in self.arms:
+            if self.has_gripper[arm]:
+                gripper_joint_names = self.gripper_joints[arm]
+                gripper_positions = self.sim.data.qpos[self._ref_gripper_joint_pos_indexes[arm]]
+                for joint_name, position in zip(gripper_joint_names, gripper_positions):
+                    joint_positions[joint_name] = position
+        
+        return joint_positions
+
+    def get_gripper_joint_positions(self, gripper_arm: Optional[str] = None):
+        """
+        Returns:
+            np.array: gripper joint positions (in angles / radians)
+        """
+        if gripper_arm is None:
+            gripper_arm = self.arms[0]
+        return self.sim.data.qpos[self._ref_gripper_joint_pos_indexes[gripper_arm]]
+
     @property
     def js_energy(self):
         """