fix site error

misoshiruseijin · misoshiruseijin · commit 9ce616023a3d · 2022-05-11T15:38:03.000-07:00
diff --git a/robosuite/environments/manipulation/single_arm_env.py b/robosuite/environments/manipulation/single_arm_env.py
@@ -52,11 +52,12 @@ def _eef_xmat(self):
         Returns:
             np.array: (3,3) End Effector orientation matrix
         """
-        pf = self.robots[0].robot_model.naming_prefix
+        pf = self.robots[0].gripper.naming_prefix
+        
         if self.env_configuration == "bimanual":
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_ee")]).reshape(3, 3)
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_grip_site")]).reshape(3, 3)
         else:
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3)
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "grip_site")]).reshape(3, 3)
 
     @property
     def _eef_xquat(self):
@@ -68,4 +69,4 @@ def _eef_xquat(self):
         Returns:
             np.array: (x,y,z,w) End Effector quaternion
         """
-        return mat2quat(self._eef_xmat)
+        return mat2quat(self._eef_xmat)
diff --git a/robosuite/environments/manipulation/two_arm_env.py b/robosuite/environments/manipulation/two_arm_env.py
@@ -86,11 +86,13 @@ def _eef0_xmat(self):
         Returns:
             np.array: (3,3) orientation matrix for EEF0
         """
-        pf = self.robots[0].robot_model.naming_prefix
+        pf = self.robots[0].gripper.naming_prefix
+
         if self.env_configuration == "bimanual":
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_ee")]).reshape(3, 3)
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "right_grip_site")]).reshape(3, 3)
+        
         else:
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3)
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "grip_site")]).reshape(3, 3)
 
     @property
     def _eef1_xmat(self):
@@ -103,11 +105,11 @@ def _eef1_xmat(self):
             np.array: (3,3) orientation matrix for EEF1
         """
         if self.env_configuration == "bimanual":
-            pf = self.robots[0].robot_model.naming_prefix
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "left_ee")]).reshape(3, 3)
+            pf = self.robots[0].gripper.naming_prefix
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "left_grip_site")]).reshape(3, 3)
         else:
-            pf = self.robots[1].robot_model.naming_prefix
-            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "ee")]).reshape(3, 3)
+            pf = self.robots[1].gripper.naming_prefix
+            return np.array(self.sim.data.site_xmat[self.sim.model.site_name2id(pf + "grip_site")]).reshape(3, 3)
 
     @property
     def _eef0_xquat(self):
@@ -131,4 +133,4 @@ def _eef1_xquat(self):
         Returns:
             np.array: (x,y,z,w) quaternion for EEF1
         """
-        return mat2quat(self._eef1_xmat)
+        return mat2quat(self._eef1_xmat)
