@@ -22,6 +22,7 @@ def __init__(
2222 simuStep ,
2323 rmodelComplete ,
2424 robotPose = [0.0 , 0.0 , 1.01927 ],
25+ robotOrientation = [0 , 0 , 0 ],
2526 inertiaOffset = True ,
2627 ):
2728 p .connect (p .GUI ) # Start the client for PyBullet
@@ -30,7 +31,7 @@ def __init__(
3031
3132 # place CoM of root link
3233 robotStartPosition = robotPose
33- robotStartOrientation = p .getQuaternionFromEuler ([ 0 , 0 , 0 ] )
34+ robotStartOrientation = p .getQuaternionFromEuler (robotOrientation )
3435 p .setAdditionalSearchPath (modelPath )
3536
3637 self .robotId = p .loadURDF (
@@ -45,11 +46,16 @@ def __init__(
4546 p .loadURDF ("plane.urdf" )
4647
4748 self .localInertiaPos = np .array ([0 , 0 , 0 ])
49+ self .localInertiaOrient = np .array ([0 , 0 , 0 , 0 ])
4850 if inertiaOffset :
4951 self .localInertiaPos = p .getDynamicsInfo (self .robotId , - 1 )[
5052 3
5153 ] # of the base link
5254
55+ self .localInertiaOrient = p .getDynamicsInfo (self .robotId , - 1 )[
56+ 4
57+ ]
58+
5359 # leg_left (45-50), leg_right (52-57), torso (0-1), arm_left (11-17),
5460 # gripper_left (21), arm_right (28-34), gripper_right (38), head (3,4).
5561 self .bulletJointNames = [
@@ -61,6 +67,9 @@ def __init__(
6167 for i in range (2 , rmodelComplete .njoints )
6268 ]
6369
70+ # Contact frame for Go2
71+ self .id_contact_bullet = [10 , 3 , 24 , 17 ]
72+
6473 # Joints controlled with crocoddyl
6574 self .bulletControlledJoints = [
6675 i
@@ -90,12 +99,15 @@ def __init__(
9099
91100 def initializeJoints (self , q0CompleteStart ):
92101 # Initialize position in pyBullet
102+ rotation = R .from_quat (q0CompleteStart [3 :7 ])
103+
104+ offset = rotation .as_matrix () @ self .localInertiaPos
93105 p .resetBasePositionAndOrientation (
94106 self .robotId ,
95107 posObj = [
96- q0CompleteStart [0 ] + self . localInertiaPos [0 ],
97- q0CompleteStart [1 ] + self . localInertiaPos [1 ],
98- q0CompleteStart [2 ] + self . localInertiaPos [2 ],
108+ q0CompleteStart [0 ] + offset [0 ],
109+ q0CompleteStart [1 ] + offset [1 ],
110+ q0CompleteStart [2 ] + offset [2 ],
99111 ],
100112 ornObj = q0CompleteStart [3 :7 ],
101113 )
@@ -116,12 +128,15 @@ def setFrictionCoefficients(self, link_id, lateral_friction, spin_friction):
116128
117129 def resetState (self , q0Start ):
118130 # Initialize position in pyBullet
131+ rotation = R .from_quat (q0Start [3 :7 ])
132+
133+ offset = rotation .as_matrix () @ self .localInertiaPos
119134 p .resetBasePositionAndOrientation (
120135 self .robotId ,
121136 posObj = [
122- q0Start [0 ] + self . localInertiaPos [0 ],
123- q0Start [1 ] + self . localInertiaPos [1 ],
124- q0Start [2 ] + self . localInertiaPos [2 ],
137+ q0Start [0 ] + offset [0 ],
138+ q0Start [1 ] + offset [1 ],
139+ q0Start [2 ] + offset [2 ],
125140 ],
126141 ornObj = q0Start [3 :7 ],
127142 )
@@ -169,6 +184,9 @@ def changeCamera(self, cameraDistance, cameraYaw, cameraPitch, cameraTargetPos):
169184 cameraDistance , cameraYaw , cameraPitch , cameraTargetPos
170185 )
171186
187+ def contactPoints (self ):
188+ return p .getContactPoints ()
189+
172190 def measureState (self ):
173191 jointStates = p .getJointStates (
174192 self .robotId , self .JointIndicesComplete
@@ -191,8 +209,15 @@ def measureState(self):
191209 [jointStates [i_joint ][1 ] for i_joint in range (len (jointStates ))],
192210 ]
193211 )
194- rotation = R .from_quat (q [3 :7 ])
195- q [:3 ] -= rotation .as_matrix () @ self .localInertiaPos
212+
213+ # Correct offset due to frame center position
214+ baseRotation = R .from_quat (q [3 :7 ])
215+ q [:3 ] -= baseRotation .as_matrix () @ self .localInertiaPos
216+
217+ # Cast velocity measure in base frame
218+ v [:3 ] = baseRotation .as_matrix ().T @ v [:3 ]
219+ v [3 :6 ] = baseRotation .as_matrix ().T @ v [3 :6 ]
220+
196221 return q , v
197222
198223 def addTable (self , path , position ):
0 commit comments