Implement faster array versions of PyBullet: getJointStatesMultiDof, getLinkStates, setJointMotorControlMultiDofArray, resetJointStatesMultiDof,

Implement StablePD in C++ through setJointMotorControlMultiDofArray method for pybullet_envs.deep_mimic, see testHumanoid.py and examples/pybullet/examples/humanoidMotionCapture.py
Minor fix in ChromeTraceUtil in case startTime>endTime (why would it happen?)

examples/pybullet/gym/pybullet_envs/deep_mimic/env/humanoid_stable_pd.py

@@ -152,58 +152,84 @@ class HumanoidStablePD(object):
     self.initializePose(self._poseInterpolator, self._kin_model, initBase=False)
 
   def initializePose(self, pose, phys_model, initBase, initializeVelocity=True):
-
+    
+    useArray = True
     if initializeVelocity:
       if initBase:
         self._pybullet_client.resetBasePositionAndOrientation(phys_model, pose._basePos,
                                                               pose._baseOrn)
         self._pybullet_client.resetBaseVelocity(phys_model, pose._baseLinVel, pose._baseAngVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, chest, pose._chestRot,
-                                                    pose._chestVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, neck, pose._neckRot, pose._neckVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightHip, pose._rightHipRot,
-                                                    pose._rightHipVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightKnee, pose._rightKneeRot,
-                                                    pose._rightKneeVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightAnkle, pose._rightAnkleRot,
-                                                    pose._rightAnkleVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightShoulder,
-                                                    pose._rightShoulderRot, pose._rightShoulderVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightElbow, pose._rightElbowRot,
-                                                    pose._rightElbowVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftHip, pose._leftHipRot,
-                                                    pose._leftHipVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftKnee, pose._leftKneeRot,
-                                                    pose._leftKneeVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftAnkle, pose._leftAnkleRot,
-                                                    pose._leftAnkleVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftShoulder,
-                                                    pose._leftShoulderRot, pose._leftShoulderVel)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftElbow, pose._leftElbowRot,
-                                                    pose._leftElbowVel)
+      if useArray:
+        indices = [chest,neck,rightHip,rightKnee,
+                  rightAnkle, rightShoulder, rightElbow,leftHip,
+                  leftKnee, leftAnkle, leftShoulder,leftElbow]
+        jointPositions = [pose._chestRot, pose._neckRot, pose._rightHipRot, pose._rightKneeRot,
+                          pose._rightAnkleRot, pose._rightShoulderRot, pose._rightElbowRot, pose._leftHipRot,
+                          pose._leftKneeRot, pose._leftAnkleRot, pose._leftShoulderRot, pose._leftElbowRot]
+        
+        jointVelocities = [pose._chestVel, pose._neckVel, pose._rightHipVel, pose._rightKneeVel,
+                          pose._rightAnkleVel, pose._rightShoulderVel, pose._rightElbowVel, pose._leftHipVel,
+                          pose._leftKneeVel, pose._leftAnkleVel, pose._leftShoulderVel, pose._leftElbowVel]
+        self._pybullet_client.resetJointStatesMultiDof(phys_model, indices,
+                                                    jointPositions, jointVelocities)
+      else:
+        self._pybullet_client.resetJointStateMultiDof(phys_model, chest, pose._chestRot,
+                                                      pose._chestVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, neck, pose._neckRot, pose._neckVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightHip, pose._rightHipRot,
+                                                      pose._rightHipVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightKnee, pose._rightKneeRot,
+                                                      pose._rightKneeVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightAnkle, pose._rightAnkleRot,
+                                                      pose._rightAnkleVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightShoulder,
+                                                      pose._rightShoulderRot, pose._rightShoulderVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightElbow, pose._rightElbowRot,
+                                                      pose._rightElbowVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftHip, pose._leftHipRot,
+                                                      pose._leftHipVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftKnee, pose._leftKneeRot,
+                                                      pose._leftKneeVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftAnkle, pose._leftAnkleRot,
+                                                      pose._leftAnkleVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftShoulder,
+                                                      pose._leftShoulderRot, pose._leftShoulderVel)
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftElbow, pose._leftElbowRot,
+                                                      pose._leftElbowVel)
     else:
+      
       if initBase:
         self._pybullet_client.resetBasePositionAndOrientation(phys_model, pose._basePos,
                                                               pose._baseOrn)
-      self._pybullet_client.resetJointStateMultiDof(phys_model, chest, pose._chestRot, [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, neck, pose._neckRot, [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightHip, pose._rightHipRot,
-                                                    [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightKnee, pose._rightKneeRot, [0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightAnkle, pose._rightAnkleRot,
-                                                    [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightShoulder,
-                                                    pose._rightShoulderRot, [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, rightElbow, pose._rightElbowRot,
-                                                    [0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftHip, pose._leftHipRot,
-                                                    [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftKnee, pose._leftKneeRot, [0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftAnkle, pose._leftAnkleRot,
-                                                    [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftShoulder,
-                                                    pose._leftShoulderRot, [0, 0, 0])
-      self._pybullet_client.resetJointStateMultiDof(phys_model, leftElbow, pose._leftElbowRot, [0])
+      if useArray:
+        indices = [chest,neck,rightHip,rightKnee,
+                  rightAnkle, rightShoulder, rightElbow,leftHip,
+                  leftKnee, leftAnkle, leftShoulder,leftElbow]
+        jointPositions = [pose._chestRot, pose._neckRot, pose._rightHipRot, pose._rightKneeRot,
+                          pose._rightAnkleRot, pose._rightShoulderRot, pose._rightElbowRot, pose._leftHipRot,
+                          pose._leftKneeRot, pose._leftAnkleRot, pose._leftShoulderRot, pose._leftElbowRot]
+        self._pybullet_client.resetJointStatesMultiDof(phys_model, indices,jointPositions)
+        
+      else:
+        self._pybullet_client.resetJointStateMultiDof(phys_model, chest, pose._chestRot, [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, neck, pose._neckRot, [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightHip, pose._rightHipRot,
+                                                      [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightKnee, pose._rightKneeRot, [0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightAnkle, pose._rightAnkleRot,
+                                                      [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightShoulder,
+                                                      pose._rightShoulderRot, [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, rightElbow, pose._rightElbowRot,
+                                                      [0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftHip, pose._leftHipRot,
+                                                      [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftKnee, pose._leftKneeRot, [0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftAnkle, pose._leftAnkleRot,
+                                                      [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftShoulder,
+                                                      pose._leftShoulderRot, [0, 0, 0])
+        self._pybullet_client.resetJointStateMultiDof(phys_model, leftElbow, pose._leftElbowRot, [0])
 
   def calcCycleCount(self, simTime, cycleTime):
     phases = simTime / cycleTime
@@ -275,6 +301,57 @@ class HumanoidStablePD(object):
     pose = self._poseInterpolator.ConvertFromAction(self._pybullet_client, action)
     return pose
 
+  def computeAndApplyPDForces(self, desiredPositions, maxForces):
+    dofIndex = 7
+    scaling = 1
+    indices = []
+    forces = []
+    targetPositions=[]
+    targetVelocities=[]
+    kps = []
+    kds = []
+    
+    for index in range(len(self._jointIndicesAll)):
+      jointIndex = self._jointIndicesAll[index]
+      indices.append(jointIndex)
+      kps.append(self._kpOrg[dofIndex])
+      kds.append(self._kdOrg[dofIndex])
+      if self._jointDofCounts[index] == 4:
+        force = [
+            scaling * maxForces[dofIndex + 0],
+            scaling * maxForces[dofIndex + 1],
+            scaling * maxForces[dofIndex + 2]
+        ]
+        targetVelocity = [0,0,0]
+        targetPosition = [
+            desiredPositions[dofIndex + 0],
+            desiredPositions[dofIndex + 1],
+            desiredPositions[dofIndex + 2],
+            desiredPositions[dofIndex + 3]
+        ]
+      if self._jointDofCounts[index] == 1:
+        force = [scaling * maxForces[dofIndex]]
+        targetPosition = [desiredPositions[dofIndex+0]]
+        targetVelocity = [0]
+      forces.append(force)
+      targetPositions.append(targetPosition)
+      targetVelocities.append(targetVelocity)
+      dofIndex += self._jointDofCounts[index]
+      
+    #static char* kwlist[] = { "bodyUniqueId", 
+    #"jointIndices", 
+    #"controlMode", "targetPositions", "targetVelocities", "forces", "positionGains", "velocityGains", "maxVelocities", "physicsClientId", NULL };
+	
+    self._pybullet_client.setJointMotorControlMultiDofArray(self._sim_model,
+                                                           indices,
+                                                           self._pybullet_client.STABLE_PD_CONTROL,
+                                                           targetPositions = targetPositions,
+                                                           targetVelocities = targetVelocities,
+                                                           forces=forces,
+                                                           positionGains = kps,
+                                                           velocityGains = kds,
+                                                           )
+
   def computePDForces(self, desiredPositions, desiredVelocities, maxForces):
     if desiredVelocities == None:
       desiredVelocities = [0] * self._totalDofs
@@ -292,27 +369,50 @@ class HumanoidStablePD(object):
   def applyPDForces(self, taus):
     dofIndex = 7
     scaling = 1
-    for index in range(len(self._jointIndicesAll)):
-      jointIndex = self._jointIndicesAll[index]
-      if self._jointDofCounts[index] == 4:
-        force = [
-            scaling * taus[dofIndex + 0], scaling * taus[dofIndex + 1],
-            scaling * taus[dofIndex + 2]
-        ]
-        #print("force[", jointIndex,"]=",force)
-        self._pybullet_client.setJointMotorControlMultiDof(self._sim_model,
-                                                           jointIndex,
-                                                           self._pybullet_client.TORQUE_CONTROL,
-                                                           force=force)
-      if self._jointDofCounts[index] == 1:
-        force = [scaling * taus[dofIndex]]
-        #print("force[", jointIndex,"]=",force)
-        self._pybullet_client.setJointMotorControlMultiDof(
-            self._sim_model,
-            jointIndex,
-            controlMode=self._pybullet_client.TORQUE_CONTROL,
-            force=force)
-      dofIndex += self._jointDofCounts[index]
+    useArray = True
+    indices = []
+    forces = []
+    
+    if (useArray):
+      for index in range(len(self._jointIndicesAll)):
+        jointIndex = self._jointIndicesAll[index]
+        indices.append(jointIndex)
+        if self._jointDofCounts[index] == 4:
+          force = [
+              scaling * taus[dofIndex + 0], scaling * taus[dofIndex + 1],
+              scaling * taus[dofIndex + 2]
+          ]
+        if self._jointDofCounts[index] == 1:
+          force = [scaling * taus[dofIndex]]
+          #print("force[", jointIndex,"]=",force)
+        forces.append(force)
+        dofIndex += self._jointDofCounts[index]
+      self._pybullet_client.setJointMotorControlMultiDofArray(self._sim_model,
+                                                             indices,
+                                                             self._pybullet_client.TORQUE_CONTROL,
+                                                             forces=forces)
+    else:
+      for index in range(len(self._jointIndicesAll)):
+        jointIndex = self._jointIndicesAll[index]
+        if self._jointDofCounts[index] == 4:
+          force = [
+              scaling * taus[dofIndex + 0], scaling * taus[dofIndex + 1],
+              scaling * taus[dofIndex + 2]
+          ]
+          #print("force[", jointIndex,"]=",force)
+          self._pybullet_client.setJointMotorControlMultiDof(self._sim_model,
+                                                             jointIndex,
+                                                             self._pybullet_client.TORQUE_CONTROL,
+                                                             force=force)
+        if self._jointDofCounts[index] == 1:
+          force = [scaling * taus[dofIndex]]
+          #print("force[", jointIndex,"]=",force)
+          self._pybullet_client.setJointMotorControlMultiDof(
+              self._sim_model,
+              jointIndex,
+              controlMode=self._pybullet_client.TORQUE_CONTROL,
+              force=force)
+        dofIndex += self._jointDofCounts[index]
 
   def setJointMotors(self, desiredPositions, maxForces):
     controlMode = self._pybullet_client.POSITION_CONTROL
@@ -531,10 +631,17 @@ class HumanoidStablePD(object):
     #self.pb2dmJoints=[0,1,2,9,10,11,3,4,5,12,13,14,6,7,8]
     self.pb2dmJoints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]
 
+    linkIndicesSim = []
+    for pbJoint in range(self._pybullet_client.getNumJoints(self._sim_model)):
+      linkIndicesSim.append(self.pb2dmJoints[pbJoint])
+      
+    linkStatesSim = self._pybullet_client.getLinkStates(self._sim_model, linkIndicesSim, computeForwardKinematics=True, computeLinkVelocity=True)
+    
     for pbJoint in range(self._pybullet_client.getNumJoints(self._sim_model)):
       j = self.pb2dmJoints[pbJoint]
       #print("joint order:",j)
-      ls = self._pybullet_client.getLinkState(self._sim_model, j, computeForwardKinematics=True)
+      #ls = self._pybullet_client.getLinkState(self._sim_model, j, computeForwardKinematics=True)
+      ls = linkStatesSim[pbJoint]
       linkPos = ls[0]
       linkOrn = ls[1]
       linkPosLocal, linkOrnLocal = self._pybullet_client.multiplyTransforms(
@@ -560,9 +667,12 @@ class HumanoidStablePD(object):
       stateVector.append(linkOrnLocal[1])
       stateVector.append(linkOrnLocal[2])
 
+       
     for pbJoint in range(self._pybullet_client.getNumJoints(self._sim_model)):
       j = self.pb2dmJoints[pbJoint]
-      ls = self._pybullet_client.getLinkState(self._sim_model, j, computeLinkVelocity=True)
+      #ls = self._pybullet_client.getLinkState(self._sim_model, j, computeLinkVelocity=True)
+      ls = linkStatesSim[pbJoint]
+      
       linkLinVel = ls[6]
       linkAngVel = ls[7]
       linkLinVelLocal, unused = self._pybullet_client.multiplyTransforms([0, 0, 0], rootTransOrn,
@@ -706,16 +816,27 @@ class HumanoidStablePD(object):
     root_ang_vel_err = self.calcRootAngVelErr(angVelSim, angVelKin)
     vel_err += root_rot_w * root_ang_vel_err
 
+    useArray = False
+    
+    if useArray:
+      jointIndices = range(num_joints)
+      simJointStates = self._pybullet_client.getJointStatesMultiDof(self._sim_model, jointIndices)
+      kinJointStates = self._pybullet_client.getJointStatesMultiDof(self._kin_model, jointIndices)
     for j in range(num_joints):
       curr_pose_err = 0
       curr_vel_err = 0
       w = mJointWeights[j]
-
-      simJointInfo = self._pybullet_client.getJointStateMultiDof(self._sim_model, j)
+      if useArray:
+        simJointInfo = simJointStates[j]
+      else:
+        simJointInfo = self._pybullet_client.getJointStateMultiDof(self._sim_model, j)
 
       #print("simJointInfo.pos=",simJointInfo[0])
       #print("simJointInfo.vel=",simJointInfo[1])
-      kinJointInfo = self._pybullet_client.getJointStateMultiDof(self._kin_model, j)
+      if useArray:
+        kinJointInfo = kinJointStates[j]
+      else:
+        kinJointInfo = self._pybullet_client.getJointStateMultiDof(self._kin_model, j)
       #print("kinJointInfo.pos=",kinJointInfo[0])
       #print("kinJointInfo.vel=",kinJointInfo[1])
       if (len(simJointInfo[0]) == 1):

examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py

@@ -52,7 +52,7 @@ class PyBulletDeepMimicEnv(Env):
       motionPath = pybullet_data.getDataPath() + "/" + motion_file[0]
       #motionPath = pybullet_data.getDataPath()+"/motions/humanoid3d_backflip.txt"
       self._mocapData.Load(motionPath)
-      timeStep = 1. / 600.
+      timeStep = 1. / 240.
       useFixedBase = False
       self._humanoid = humanoid_stable_pd.HumanoidStablePD(self._pybullet_client, self._mocapData,
                                                            timeStep, useFixedBase)
@@ -283,10 +283,16 @@ class PyBulletDeepMimicEnv(Env):
         ]
 
         if self._useStablePD:
-          taus = self._humanoid.computePDForces(self.desiredPose,
+          usePythonStablePD = False
+          if usePythonStablePD:
+            taus = self._humanoid.computePDForces(self.desiredPose,
                                                 desiredVelocities=None,
                                                 maxForces=maxForces)
-          self._humanoid.applyPDForces(taus)
+            #taus = [0]*43
+            self._humanoid.applyPDForces(taus)
+          else:
+            self._humanoid.computeAndApplyPDForces(self.desiredPose,
+                                                maxForces=maxForces)
         else:
           self._humanoid.setJointMotors(self.desiredPose, maxForces=maxForces)
 

examples/pybullet/gym/pybullet_envs/deep_mimic/env/testHumanoid.py

@@ -47,7 +47,7 @@ def isKeyTriggered(keys, key):
 
 animating = False
 singleStep = False
-
+humanoid.resetPose()
 t = 0
 while (1):
 
@@ -66,7 +66,7 @@ while (1):
     #print("t=",t)
     for i in range(1):
 
-      print("t=", t)
+      #print("t=", t)
       humanoid.setSimTime(t)
 
       humanoid.computePose(humanoid._frameFraction)
@@ -75,9 +75,11 @@ while (1):
       #humanoid.resetPose()
 
       desiredPose = humanoid.computePose(humanoid._frameFraction)
+      
       #desiredPose = desiredPose.GetPose()
       #curPose = HumanoidPoseInterpolator()
       #curPose.reset()
+      
       s = humanoid.getState()
       #np.savetxt("pb_record_state_s.csv", s, delimiter=",")
       maxForces = [
@@ -85,10 +87,14 @@ while (1):
           90, 90, 100, 100, 100, 100, 60, 200, 200, 200, 200, 150, 90, 90, 90, 90, 100, 100, 100,
           100, 60
       ]
-      taus = humanoid.computePDForces(desiredPose, desiredVelocities=None, maxForces=maxForces)
-
-      #print("taus=",taus)
-      humanoid.applyPDForces(taus)
+      
+      usePythonStablePD = False
+      if usePythonStablePD:
+        taus = humanoid.computePDForces(desiredPose, desiredVelocities=None, maxForces=maxForces)
+        #Print("taus=",taus)
+        humanoid.applyPDForces(taus)
+      else:
+        humanoid.computeAndApplyPDForces(desiredPose,maxForces=maxForces)
 
       pybullet_client.stepSimulation()
       t += 1. / 600.

examples/pybullet/gym/pybullet_utils/pd_controller_stable.py

@@ -62,8 +62,15 @@ class PDControllerStableMultiDof(object):
     qIndex = 7
     qdotIndex = 7
     zeroAccelerations = [0, 0, 0, 0, 0, 0, 0]
+    useArray = True
+    if useArray:
+      jointStates = self._pb.getJointStatesMultiDof(bodyUniqueId,jointIndices)
+    
     for i in range(numJoints):
-      js = self._pb.getJointStateMultiDof(bodyUniqueId, jointIndices[i])
+      if useArray:
+        js = jointStates[i]
+      else:
+        js = self._pb.getJointStateMultiDof(bodyUniqueId, jointIndices[i])
 
       jointPos = js[0]
       jointVel = js[1]
@@ -153,8 +160,6 @@ class PDControllerStableMultiDof(object):
     if useNumpySolver:
       qddot = np.linalg.solve(A, b)
     else:
-      dofCount = len(b)
-      print(dofCount)
       qddot = self._pb.ldltSolve(bodyUniqueId, jointPositions=q1, b=b.tolist(), kd=kds, t=timeStep)
 
     tau = p + d - Kd.dot(qddot) * timeStep