faster thrown object removal

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

@@ -19,11 +19,9 @@ jointFrictionForce = 0
 
 class HumanoidStablePD(object):
 
-  def __init__( self, pybullet_client, mocap_data, timeStep, 
+  def __init__(self, pybullet_client, mocap_data, timeStep, useFixedBase=True):
-                useFixedBase=True, arg_parser=None):
     self._pybullet_client = pybullet_client
     self._mocap_data = mocap_data
-    self._arg_parser = arg_parser
     print("LOADING humanoid!")
 
     self._sim_model = self._pybullet_client.loadURDF(
@@ -137,10 +135,7 @@ class HumanoidStablePD(object):
     self._jointDofCounts = [4, 4, 4, 1, 4, 4, 1, 4, 1, 4, 4, 1]
 
     #only those body parts/links are allowed to touch the ground, otherwise the episode terminates
-    fall_contact_bodies = []
+    self._allowed_body_parts = [5, 11]
-    if self._arg_parser is not None:
-      fall_contact_bodies = self._arg_parser.parse_ints("fall_contact_bodies")
-    self._fall_contact_body_parts = fall_contact_bodies
 
     #[x,y,z] base position and [x,y,z,w] base orientation!
     self._totalDofs = 7
@@ -149,6 +144,7 @@ class HumanoidStablePD(object):
     self.setSimTime(0)
 
     self.resetPose()
+    self.thrown_body_ids = []
 
   def resetPose(self):
     #print("resetPose with self._frame=", self._frame, " and self._frameFraction=",self._frameFraction)
@@ -595,11 +591,14 @@ class HumanoidStablePD(object):
       #ignore self-collision
       if (p[1] == p[2]):
         continue
+      # ignore collisions with thrown objects
+      if p[1] in self.thrown_body_ids or p[2] in self.thrown_body_ids:
+        continue
       if (p[1] == self._sim_model):
         part = p[3]
       if (p[2] == self._sim_model):
         part = p[4]
-      if (part >= 0 and part in self._fall_contact_body_parts):
+      if (part >= 0 and part not in self._allowed_body_parts):
         #print("terminating part:", part)
         terminates = True
 
@@ -800,3 +799,7 @@ class HumanoidStablePD(object):
     #print("com_reward=",com_reward)
 
     return reward
+
+  def getSimModelBasePosition(self):
+    return  self._pybullet_client\
+                .getBasePositionAndOrientation(self._sim_model)

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

@@ -55,7 +55,7 @@ class PyBulletDeepMimicEnv(Env):
       timeStep = 1. / 600.
       useFixedBase = False
       self._humanoid = humanoid_stable_pd.HumanoidStablePD(self._pybullet_client, self._mocapData,
-                                                           timeStep, useFixedBase, self._arg_parser)
+                                                           timeStep, useFixedBase)
       self._isInitialized = True
 
       self._pybullet_client.setTimeStep(timeStep)
@@ -82,6 +82,9 @@ class PyBulletDeepMimicEnv(Env):
     startTime = float(rn) / rnrange * self._humanoid.getCycleTime()
     self.t = startTime
 
+    # Remove all the thrown objects
+    self.removeThrownObjects()
+
     self._humanoid.setSimTime(startTime)
 
     self._humanoid.resetPose()
@@ -319,3 +322,49 @@ class PyBulletDeepMimicEnv(Env):
     if o in keys:
       return keys[ord(key)] & self._pybullet_client.KEY_WAS_TRIGGERED
     return False
+
+  def hitWithObject(self):
+    # Spawn an object with velocity hitting the model
+    r = random.random()
+    distance = 3.0
+    rand_angle = r * 2 * math.pi - math.pi
+
+    position, orientation = \
+      self._humanoid.getSimModelBasePosition()
+
+    # Remember, in bullet: Y direction is "up". X and Z direction are
+    # in the horizontal plane.
+    ball_position = [ distance * math.cos(rand_angle) + position[0],
+                      position[1],
+                      distance * math.sin(rand_angle)+ position[2]]
+
+    visualShapeId = self._pybullet_client.createVisualShape(
+      shapeType=self._pybullet_client.GEOM_SPHERE,
+      radius=0.1)
+    
+    collisionShapeId = self._pybullet_client.createCollisionShape(
+      shapeType=self._pybullet_client.GEOM_SPHERE,
+      radius=0.1)
+
+    body = self._pybullet_client.createMultiBody(
+      baseMass=1,
+      baseCollisionShapeIndex=collisionShapeId,
+      baseVisualShapeIndex=visualShapeId,
+      basePosition=ball_position)
+
+    distance_scale = 10
+    ball_velocity = [distance_scale * (position[i] - ball_position[i]) for i in range(3)]
+    self._pybullet_client.resetBaseVelocity(body, linearVelocity=ball_velocity)
+
+    self._humanoid.thrown_body_ids.append(body)
+
+  def removeThrownObjects(self):
+    # Disable gui so that object removal does not cause stutter
+    self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_RENDERING,0)
+
+    for objectId in self._humanoid.thrown_body_ids:
+      self._pybullet_client.removeBody(objectId)
+
+    self._humanoid.thrown_body_ids = []
+
+    self._pybullet_client.configureDebugVisualizer(self._pybullet_client.COV_ENABLE_RENDERING,1)

examples/pybullet/gym/pybullet_envs/deep_mimic/testrl.py

@@ -87,6 +87,9 @@ if __name__ == '__main__':
 
     if world.env.isKeyTriggered(keys, ' '):
       animating = not animating
+    if world.env.isKeyTriggered(keys, 'x'):
+      # print("Throwing object.")
+      world.env.hitWithObject()
     if (animating):
       update_world(world, timeStep)
       #animating=False