1from pybullet_utils import bullet_client
2import math
3
4
5class HumanoidPoseInterpolator(object):
6
7  def __init__(self):
8    pass
9
10  def Reset(self,
11            basePos=[0, 0, 0],
12            baseOrn=[0, 0, 0, 1],
13            chestRot=[0, 0, 0, 1],
14            neckRot=[0, 0, 0, 1],
15            rightHipRot=[0, 0, 0, 1],
16            rightKneeRot=[0],
17            rightAnkleRot=[0, 0, 0, 1],
18            rightShoulderRot=[0, 0, 0, 1],
19            rightElbowRot=[0],
20            leftHipRot=[0, 0, 0, 1],
21            leftKneeRot=[0],
22            leftAnkleRot=[0, 0, 0, 1],
23            leftShoulderRot=[0, 0, 0, 1],
24            leftElbowRot=[0],
25            baseLinVel=[0, 0, 0],
26            baseAngVel=[0, 0, 0],
27            chestVel=[0, 0, 0],
28            neckVel=[0, 0, 0],
29            rightHipVel=[0, 0, 0],
30            rightKneeVel=[0],
31            rightAnkleVel=[0, 0, 0],
32            rightShoulderVel=[0, 0, 0],
33            rightElbowVel=[0],
34            leftHipVel=[0, 0, 0],
35            leftKneeVel=[0],
36            leftAnkleVel=[0, 0, 0],
37            leftShoulderVel=[0, 0, 0],
38            leftElbowVel=[0]):
39
40    self._basePos = basePos
41    self._baseLinVel = baseLinVel
42    #print("HumanoidPoseInterpolator.Reset: baseLinVel = ", baseLinVel)
43    self._baseOrn = baseOrn
44    self._baseAngVel = baseAngVel
45
46    self._chestRot = chestRot
47    self._chestVel = chestVel
48    self._neckRot = neckRot
49    self._neckVel = neckVel
50
51    self._rightHipRot = rightHipRot
52    self._rightHipVel = rightHipVel
53    self._rightKneeRot = rightKneeRot
54    self._rightKneeVel = rightKneeVel
55    self._rightAnkleRot = rightAnkleRot
56    self._rightAnkleVel = rightAnkleVel
57
58    self._rightShoulderRot = rightShoulderRot
59    self._rightShoulderVel = rightShoulderVel
60    self._rightElbowRot = rightElbowRot
61    self._rightElbowVel = rightElbowVel
62
63    self._leftHipRot = leftHipRot
64    self._leftHipVel = leftHipVel
65    self._leftKneeRot = leftKneeRot
66    self._leftKneeVel = leftKneeVel
67    self._leftAnkleRot = leftAnkleRot
68    self._leftAnkleVel = leftAnkleVel
69
70    self._leftShoulderRot = leftShoulderRot
71    self._leftShoulderVel = leftShoulderVel
72    self._leftElbowRot = leftElbowRot
73    self._leftElbowVel = leftElbowVel
74
75  def ComputeLinVel(self, posStart, posEnd, deltaTime):
76    vel = [(posEnd[0] - posStart[0]) / deltaTime, (posEnd[1] - posStart[1]) / deltaTime,
77           (posEnd[2] - posStart[2]) / deltaTime]
78    return vel
79
80  def ComputeAngVel(self, ornStart, ornEnd, deltaTime, bullet_client):
81    dorn = bullet_client.getDifferenceQuaternion(ornStart, ornEnd)
82    axis, angle = bullet_client.getAxisAngleFromQuaternion(dorn)
83    angVel = [(axis[0] * angle) / deltaTime, (axis[1] * angle) / deltaTime,
84              (axis[2] * angle) / deltaTime]
85    return angVel
86
87  def ComputeAngVelRel(self, ornStart, ornEnd, deltaTime, bullet_client):
88    ornStartConjugate = [-ornStart[0], -ornStart[1], -ornStart[2], ornStart[3]]
89    pos_diff, q_diff = bullet_client.multiplyTransforms([0, 0, 0], ornStartConjugate, [0, 0, 0],
90                                                        ornEnd)
91    axis, angle = bullet_client.getAxisAngleFromQuaternion(q_diff)
92    angVel = [(axis[0] * angle) / deltaTime, (axis[1] * angle) / deltaTime,
93              (axis[2] * angle) / deltaTime]
94    return angVel
95
96  def NormalizeVector(self, vec):
97    length2 = orn[0] * orn[0] + orn[1] * orn[1] + orn[2] * orn[2]
98    if (length2 > 0):
99      length = math.sqrt(length2)
100
101  def NormalizeQuaternion(self, orn):
102    length2 = orn[0] * orn[0] + orn[1] * orn[1] + orn[2] * orn[2] + orn[3] * orn[3]
103    if (length2 > 0):
104      length = math.sqrt(length2)
105      orn[0] /= length
106      orn[1] /= length
107      orn[2] /= length
108      orn[3] /= length
109      return orn
110
111    #print("Normalize? length=",length)
112
113  def PostProcessMotionData(self, frameData):
114    baseOrn1Start = [frameData[5], frameData[6], frameData[7], frameData[4]]
115
116    chestRotStart = [frameData[9], frameData[10], frameData[11], frameData[8]]
117
118    neckRotStart = [frameData[13], frameData[14], frameData[15], frameData[12]]
119    rightHipRotStart = [frameData[17], frameData[18], frameData[19], frameData[16]]
120    rightAnkleRotStart = [frameData[22], frameData[23], frameData[24], frameData[21]]
121    rightShoulderRotStart = [frameData[26], frameData[27], frameData[28], frameData[25]]
122    leftHipRotStart = [frameData[31], frameData[32], frameData[33], frameData[30]]
123    leftAnkleRotStart = [frameData[36], frameData[37], frameData[38], frameData[35]]
124    leftShoulderRotStart = [frameData[40], frameData[41], frameData[42], frameData[39]]
125
126  def GetPose(self):
127    pose = [
128        self._basePos[0], self._basePos[1], self._basePos[2], self._baseOrn[0], self._baseOrn[1],
129        self._baseOrn[2], self._baseOrn[3], self._chestRot[0], self._chestRot[1],
130        self._chestRot[2], self._chestRot[3], self._neckRot[0], self._neckRot[1], self._neckRot[2],
131        self._neckRot[3], self._rightHipRot[0], self._rightHipRot[1], self._rightHipRot[2],
132        self._rightHipRot[3], self._rightKneeRot[0], self._rightAnkleRot[0],
133        self._rightAnkleRot[1], self._rightAnkleRot[2], self._rightAnkleRot[3],
134        self._rightShoulderRot[0], self._rightShoulderRot[1], self._rightShoulderRot[2],
135        self._rightShoulderRot[3], self._rightElbowRot[0], self._leftHipRot[0],
136        self._leftHipRot[1], self._leftHipRot[2], self._leftHipRot[3], self._leftKneeRot[0],
137        self._leftAnkleRot[0], self._leftAnkleRot[1], self._leftAnkleRot[2], self._leftAnkleRot[3],
138        self._leftShoulderRot[0], self._leftShoulderRot[1], self._leftShoulderRot[2],
139        self._leftShoulderRot[3], self._leftElbowRot[0]
140    ]
141    return pose
142
143  def Slerp(self, frameFraction, frameData, frameDataNext, bullet_client):
144    keyFrameDuration = frameData[0]
145    basePos1Start = [frameData[1], frameData[2], frameData[3]]
146    basePos1End = [frameDataNext[1], frameDataNext[2], frameDataNext[3]]
147    self._basePos = [
148        basePos1Start[0] + frameFraction * (basePos1End[0] - basePos1Start[0]),
149        basePos1Start[1] + frameFraction * (basePos1End[1] - basePos1Start[1]),
150        basePos1Start[2] + frameFraction * (basePos1End[2] - basePos1Start[2])
151    ]
152    self._baseLinVel = self.ComputeLinVel(basePos1Start, basePos1End, keyFrameDuration)
153    baseOrn1Start = [frameData[5], frameData[6], frameData[7], frameData[4]]
154    baseOrn1Next = [frameDataNext[5], frameDataNext[6], frameDataNext[7], frameDataNext[4]]
155    self._baseOrn = bullet_client.getQuaternionSlerp(baseOrn1Start, baseOrn1Next, frameFraction)
156    self._baseAngVel = self.ComputeAngVel(baseOrn1Start, baseOrn1Next, keyFrameDuration,
157                                          bullet_client)
158
159    ##pre-rotate to make z-up
160    #y2zPos=[0,0,0.0]
161    #y2zOrn = p.getQuaternionFromEuler([1.57,0,0])
162    #basePos,baseOrn = p.multiplyTransforms(y2zPos, y2zOrn,basePos1,baseOrn1)
163
164    chestRotStart = [frameData[9], frameData[10], frameData[11], frameData[8]]
165    chestRotEnd = [frameDataNext[9], frameDataNext[10], frameDataNext[11], frameDataNext[8]]
166    self._chestRot = bullet_client.getQuaternionSlerp(chestRotStart, chestRotEnd, frameFraction)
167    self._chestVel = self.ComputeAngVelRel(chestRotStart, chestRotEnd, keyFrameDuration,
168                                           bullet_client)
169
170    neckRotStart = [frameData[13], frameData[14], frameData[15], frameData[12]]
171    neckRotEnd = [frameDataNext[13], frameDataNext[14], frameDataNext[15], frameDataNext[12]]
172    self._neckRot = bullet_client.getQuaternionSlerp(neckRotStart, neckRotEnd, frameFraction)
173    self._neckVel = self.ComputeAngVelRel(neckRotStart, neckRotEnd, keyFrameDuration,
174                                          bullet_client)
175
176    rightHipRotStart = [frameData[17], frameData[18], frameData[19], frameData[16]]
177    rightHipRotEnd = [frameDataNext[17], frameDataNext[18], frameDataNext[19], frameDataNext[16]]
178    self._rightHipRot = bullet_client.getQuaternionSlerp(rightHipRotStart, rightHipRotEnd,
179                                                         frameFraction)
180    self._rightHipVel = self.ComputeAngVelRel(rightHipRotStart, rightHipRotEnd, keyFrameDuration,
181                                              bullet_client)
182
183    rightKneeRotStart = [frameData[20]]
184    rightKneeRotEnd = [frameDataNext[20]]
185    self._rightKneeRot = [
186        rightKneeRotStart[0] + frameFraction * (rightKneeRotEnd[0] - rightKneeRotStart[0])
187    ]
188    self._rightKneeVel = [(rightKneeRotEnd[0] - rightKneeRotStart[0]) / keyFrameDuration]
189
190    rightAnkleRotStart = [frameData[22], frameData[23], frameData[24], frameData[21]]
191    rightAnkleRotEnd = [frameDataNext[22], frameDataNext[23], frameDataNext[24], frameDataNext[21]]
192    self._rightAnkleRot = bullet_client.getQuaternionSlerp(rightAnkleRotStart, rightAnkleRotEnd,
193                                                           frameFraction)
194    self._rightAnkleVel = self.ComputeAngVelRel(rightAnkleRotStart, rightAnkleRotEnd,
195                                                keyFrameDuration, bullet_client)
196
197    rightShoulderRotStart = [frameData[26], frameData[27], frameData[28], frameData[25]]
198    rightShoulderRotEnd = [
199        frameDataNext[26], frameDataNext[27], frameDataNext[28], frameDataNext[25]
200    ]
201    self._rightShoulderRot = bullet_client.getQuaternionSlerp(rightShoulderRotStart,
202                                                              rightShoulderRotEnd, frameFraction)
203    self._rightShoulderVel = self.ComputeAngVelRel(rightShoulderRotStart, rightShoulderRotEnd,
204                                                   keyFrameDuration, bullet_client)
205
206    rightElbowRotStart = [frameData[29]]
207    rightElbowRotEnd = [frameDataNext[29]]
208    self._rightElbowRot = [
209        rightElbowRotStart[0] + frameFraction * (rightElbowRotEnd[0] - rightElbowRotStart[0])
210    ]
211    self._rightElbowVel = [(rightElbowRotEnd[0] - rightElbowRotStart[0]) / keyFrameDuration]
212
213    leftHipRotStart = [frameData[31], frameData[32], frameData[33], frameData[30]]
214    leftHipRotEnd = [frameDataNext[31], frameDataNext[32], frameDataNext[33], frameDataNext[30]]
215    self._leftHipRot = bullet_client.getQuaternionSlerp(leftHipRotStart, leftHipRotEnd,
216                                                        frameFraction)
217    self._leftHipVel = self.ComputeAngVelRel(leftHipRotStart, leftHipRotEnd, keyFrameDuration,
218                                             bullet_client)
219
220    leftKneeRotStart = [frameData[34]]
221    leftKneeRotEnd = [frameDataNext[34]]
222    self._leftKneeRot = [
223        leftKneeRotStart[0] + frameFraction * (leftKneeRotEnd[0] - leftKneeRotStart[0])
224    ]
225    self._leftKneeVel = [(leftKneeRotEnd[0] - leftKneeRotStart[0]) / keyFrameDuration]
226
227    leftAnkleRotStart = [frameData[36], frameData[37], frameData[38], frameData[35]]
228    leftAnkleRotEnd = [frameDataNext[36], frameDataNext[37], frameDataNext[38], frameDataNext[35]]
229    self._leftAnkleRot = bullet_client.getQuaternionSlerp(leftAnkleRotStart, leftAnkleRotEnd,
230                                                          frameFraction)
231    self._leftAnkleVel = self.ComputeAngVelRel(leftAnkleRotStart, leftAnkleRotEnd,
232                                               keyFrameDuration, bullet_client)
233
234    leftShoulderRotStart = [frameData[40], frameData[41], frameData[42], frameData[39]]
235    leftShoulderRotEnd = [
236        frameDataNext[40], frameDataNext[41], frameDataNext[42], frameDataNext[39]
237    ]
238    self._leftShoulderRot = bullet_client.getQuaternionSlerp(leftShoulderRotStart,
239                                                             leftShoulderRotEnd, frameFraction)
240    self._leftShoulderVel = self.ComputeAngVelRel(leftShoulderRotStart, leftShoulderRotEnd,
241                                                  keyFrameDuration, bullet_client)
242
243    leftElbowRotStart = [frameData[43]]
244    leftElbowRotEnd = [frameDataNext[43]]
245    self._leftElbowRot = [
246        leftElbowRotStart[0] + frameFraction * (leftElbowRotEnd[0] - leftElbowRotStart[0])
247    ]
248    self._leftElbowVel = [(leftElbowRotEnd[0] - leftElbowRotStart[0]) / keyFrameDuration]
249
250    pose = self.GetPose()
251    return pose
252
253  def ConvertFromAction(self, pybullet_client, action):
254    #turn action into pose
255
256    self.Reset()  #?? needed?
257    index = 0
258    angle = action[index]
259    axis = [action[index + 1], action[index + 2], action[index + 3]]
260    index += 4
261    self._chestRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
262    #print("pose._chestRot=",pose._chestRot)
263
264    angle = action[index]
265    axis = [action[index + 1], action[index + 2], action[index + 3]]
266    index += 4
267    self._neckRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
268
269    angle = action[index]
270    axis = [action[index + 1], action[index + 2], action[index + 3]]
271    index += 4
272    self._rightHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
273
274    angle = action[index]
275    index += 1
276    self._rightKneeRot = [angle]
277
278    angle = action[index]
279    axis = [action[index + 1], action[index + 2], action[index + 3]]
280    index += 4
281    self._rightAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
282
283    angle = action[index]
284    axis = [action[index + 1], action[index + 2], action[index + 3]]
285    index += 4
286    self._rightShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
287
288    angle = action[index]
289    index += 1
290    self._rightElbowRot = [angle]
291
292    angle = action[index]
293    axis = [action[index + 1], action[index + 2], action[index + 3]]
294    index += 4
295    self._leftHipRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
296
297    angle = action[index]
298    index += 1
299    self._leftKneeRot = [angle]
300
301    angle = action[index]
302    axis = [action[index + 1], action[index + 2], action[index + 3]]
303    index += 4
304    self._leftAnkleRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
305
306    angle = action[index]
307    axis = [action[index + 1], action[index + 2], action[index + 3]]
308    index += 4
309    self._leftShoulderRot = pybullet_client.getQuaternionFromAxisAngle(axis, angle)
310
311    angle = action[index]
312    index += 1
313    self._leftElbowRot = [angle]
314
315    pose = self.GetPose()
316    return pose
317