Merge branch 'master' of https://github.com/ARISE-Initiative/robosuite

Author: Abhiram824

robosuite/controllers/composite/composite_controller.py

@@ -122,6 +122,23 @@ def get_control_dim(self, part_name):
             return self.part_controllers[part_name].control_dim
 
     def get_controller_base_pose(self, controller_name):
+        """
+        Get the base position and orientation of a specified controller's part. Note: this pose may likely differ from
+        the robot base's pose.
+
+        Args:
+            controller_name (str): The name of the controller, used to look up part-specific information.
+
+        Returns:
+            tuple[np.ndarray, np.ndarray]: A tuple containing:
+                - base_pos (np.ndarray): The 3D position of the part's center in world coordinates (shape: (3,)).
+                - base_ori (np.ndarray): The 3x3 rotation matrix representing the part's orientation in world coordinates.
+
+        Details:
+            - Uses the controller's `naming_prefix` and `part_name` to construct the corresponding site name.
+            - Queries the simulation (`self.sim`) for the site's position (`site_xpos`) and orientation (`site_xmat`).
+            - The site orientation matrix is reshaped from a flat array of size 9 to a 3x3 rotation matrix.
+        """
         naming_prefix = self.part_controllers[controller_name].naming_prefix
         part_name = self.part_controllers[controller_name].part_name
         base_pos = np.array(self.sim.data.site_xpos[self.sim.model.site_name2id(f"{naming_prefix}{part_name}_center")])

robosuite/controllers/parts/arm/osc.py

@@ -242,18 +242,19 @@ def set_goal(self, action):
 
         # Parse action based on the impedance mode, and update kp / kd as necessary
         if self.impedance_mode == "variable":
-            damping_ratio, kp, delta = action[:6], action[6:12], action[12:]
+            damping_ratio, kp, goal_update = action[:6], action[6:12], action[12:]
             self.kp = np.clip(kp, self.kp_min, self.kp_max)
             self.kd = 2 * np.sqrt(self.kp) * np.clip(damping_ratio, self.damping_ratio_min, self.damping_ratio_max)
         elif self.impedance_mode == "variable_kp":
-            kp, delta = action[:6], action[6:]
+            kp, goal_update = action[:6], action[6:]
             self.kp = np.clip(kp, self.kp_min, self.kp_max)
             self.kd = 2 * np.sqrt(self.kp)  # critically damped
         else:  # This is case "fixed"
-            delta = action
+            goal_update = action
 
         # If we're using deltas, interpret actions as such
         if self.input_type == "delta":
+            delta = goal_update
             scaled_delta = self.scale_action(delta)
             self.goal_pos = self.compute_goal_pos(scaled_delta[0:3])
             if self.use_ori is True:
@@ -262,9 +263,10 @@ def set_goal(self, action):
                 self.goal_ori = self.compute_goal_ori(np.zeros(3))
         # Else, interpret actions as absolute values
         elif self.input_type == "absolute":
-            self.goal_pos = action[0:3]
+            abs_action = goal_update
+            self.goal_pos = abs_action[0:3]
             if self.use_ori is True:
-                self.goal_ori = Rotation.from_rotvec(action[3:6]).as_matrix()
+                self.goal_ori = Rotation.from_rotvec(abs_action[3:6]).as_matrix()
             else:
                 self.goal_ori = self.compute_goal_ori(np.zeros(3))
         else:

robosuite/renderers/context/egl_context.py

@@ -26,7 +26,7 @@
     raise ImportError(
         "Cannot use EGL rendering platform. "
         "The PYOPENGL_PLATFORM environment variable is set to {!r} "
-        "(should be either unset or 'egl')."
+        "(should be either unset or 'egl').".format(PYOPENGL_PLATFORM)
     )
 
 from mujoco.egl import egl_ext as EGL

robosuite/robots/robot.py

@@ -356,10 +356,14 @@ def joint_pos_sin(obs_cache):
         def joint_vel(obs_cache):
             return np.array([self.sim.data.qvel[x] for x in self._ref_joint_vel_indexes])
 
-        sensors = [joint_pos, joint_pos_cos, joint_pos_sin, joint_vel]
-        names = ["joint_pos", "joint_pos_cos", "joint_pos_sin", "joint_vel"]
+        @sensor(modality=modality)
+        def joint_acc(obs_cache):
+            return np.array([self.sim.data.qacc[x] for x in self._ref_joint_vel_indexes])
+
+        sensors = [joint_pos, joint_pos_cos, joint_pos_sin, joint_vel, joint_acc]
+        names = ["joint_pos", "joint_pos_cos", "joint_pos_sin", "joint_vel", "joint_acc"]
         # We don't want to include the direct joint pos sensor outputs
-        actives = [True, True, True, True]
+        actives = [True, True, True, True, True]
 
         for arm in self.arms:
             arm_sensors, arm_sensor_names = self._create_arm_sensors(arm, modality=modality)

robosuite/wrappers/data_collection_wrapper.py

@@ -77,6 +77,7 @@ def _start_new_episode(self):
 
         # trick for ensuring that we can play MuJoCo demonstrations back
         # deterministically by using the recorded actions open loop
+        self.env.set_ep_meta(self.env.get_ep_meta())
         self.env.reset_from_xml_string(self._current_task_instance_xml)
         self.env.sim.reset()
         self.env.sim.set_state_from_flattened(self._current_task_instance_state)
@@ -144,6 +145,7 @@ def reset(self):
         Returns:
             OrderedDict: Environment observation space after reset occurs
         """
+        self.env.unset_ep_meta()  # unset any episode meta data that was previously set
         ret = super().reset()
         self._start_new_episode()
         return ret