From b433e07d333ba869ec76467c17558fc80deb9682 Mon Sep 17 00:00:00 2001
From: Erwin Coumans <erwincoumans@google.com>
Date: Mon, 22 Jul 2019 17:18:41 -0700
Subject: [PATCH] pybullet_envs.deep_mimic: use getLinkStates for faster
 'getReward' pybullet_envs.deep_mimic: add profile timers for each thread to
 optimize training

---
 .../deep_mimic/DeepMimic_Optimizer.py              |  9 ++++++++-
 .../deep_mimic/env/humanoid_stable_pd.py           | 14 +++++++++++---
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/examples/pybullet/gym/pybullet_envs/deep_mimic/DeepMimic_Optimizer.py b/examples/pybullet/gym/pybullet_envs/deep_mimic/DeepMimic_Optimizer.py
index b6e9e6607..25427f16d 100644
--- a/examples/pybullet/gym/pybullet_envs/deep_mimic/DeepMimic_Optimizer.py
+++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/DeepMimic_Optimizer.py
@@ -12,6 +12,7 @@ from pybullet_envs.deep_mimic.learning.rl_world import RLWorld
 from pybullet_utils.logger import Logger
 from pybullet_envs.deep_mimic.testrl import update_world, update_timestep, build_world
 import pybullet_utils.mpi_util as MPIUtil
+from mpi4py import MPI
 
 args = []
 world = None
@@ -20,9 +21,15 @@ world = None
 def run():
   global update_timestep
   global world
-
+  rank = MPI.COMM_WORLD.Get_rank()
+  name = "timings"+str(rank)+".json"
+  slot = world.env._pybullet_client.startStateLogging(world.env._pybullet_client.STATE_LOGGING_PROFILE_TIMINGS,name)
+  count=1000
   done = False
   while not (done):
+    count-=1
+    if (count<0):
+      world.env._pybullet_client.stopStateLogging(slot)
     update_world(world, update_timestep)
 
   return
diff --git a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/humanoid_stable_pd.py b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/humanoid_stable_pd.py
index c846ed33b..244112a1e 100644
--- a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/humanoid_stable_pd.py
+++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/humanoid_stable_pd.py
@@ -341,7 +341,7 @@ class HumanoidStablePD(object):
     #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,
@@ -822,6 +822,9 @@ class HumanoidStablePD(object):
       jointIndices = range(num_joints)
       simJointStates = self._pybullet_client.getJointStatesMultiDof(self._sim_model, jointIndices)
       kinJointStates = self._pybullet_client.getJointStatesMultiDof(self._kin_model, jointIndices)
+    if useArray:
+      linkStatesSim = self._pybullet_client.getLinkStates(self._sim_model, jointIndices)
+      linkStatesKin = self._pybullet_client.getLinkStates(self._kin_model, jointIndices)
     for j in range(num_joints):
       curr_pose_err = 0
       curr_vel_err = 0
@@ -859,10 +862,15 @@ class HumanoidStablePD(object):
       vel_err += w * curr_vel_err
 
       is_end_eff = j in self._end_effectors
+      
       if is_end_eff:
 
-        linkStateSim = self._pybullet_client.getLinkState(self._sim_model, j)
-        linkStateKin = self._pybullet_client.getLinkState(self._kin_model, j)
+        if useArray:
+          linkStateSim = linkStatesSim[j]
+          linkStateKin = linkStatesKin[j]
+        else:
+          linkStateSim = self._pybullet_client.getLinkState(self._sim_model, j)
+          linkStateKin = self._pybullet_client.getLinkState(self._kin_model, j)
         linkPosSim = linkStateSim[0]
         linkPosKin = linkStateKin[0]
         linkPosDiff = [