From c9b898c735f9431fc5ae6a0e5387116ebb7e9287 Mon Sep 17 00:00:00 2001
From: Bart Moyaers <bart.moyaers@kuleuven.be>
Date: Thu, 20 Jun 2019 10:12:49 +0200
Subject: [PATCH] throw objects deepmimic env

---
 .../deep_mimic/env/humanoid_stable_pd.py      |  8 ++++
 .../deep_mimic/env/pybullet_deep_mimic_env.py | 44 +++++++++++++++++++
 .../gym/pybullet_envs/deep_mimic/testrl.py    |  3 ++
 3 files changed, 55 insertions(+)

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 cbdf9d191..06175224b 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
@@ -149,6 +149,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)
@@ -705,6 +706,9 @@ 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):
@@ -929,3 +933,7 @@ class HumanoidStablePD(object):
     #print("com_reward=",com_reward)
 
     return reward
+
+  def getSimModelBasePosition(self):
+    return  self._pybullet_client\
+                .getBasePositionAndOrientation(self._sim_model)
\ No newline at end of file
diff --git a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py
index 309de425c..b0be7b03b 100644
--- a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py
+++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/pybullet_deep_mimic_env.py
@@ -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()
@@ -325,3 +328,44 @@ 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):
+    for objectId in self._humanoid.thrown_body_ids:
+      self._pybullet_client.removeBody(objectId)
+
+    self._humanoid.thrown_body_ids = []
\ No newline at end of file
diff --git a/examples/pybullet/gym/pybullet_envs/deep_mimic/testrl.py b/examples/pybullet/gym/pybullet_envs/deep_mimic/testrl.py
index 9b9549359..b37ae1bbd 100644
--- a/examples/pybullet/gym/pybullet_envs/deep_mimic/testrl.py
+++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/testrl.py
@@ -92,6 +92,9 @@ if __name__ == '__main__':
       animating = not animating
     if world.env.isKeyTriggered(keys, 'i'):
       step = True
+    if world.env.isKeyTriggered(keys, 'x'):
+      # print("Throwing object.")
+      world.env.hitWithObject()
     if (animating or step):
       update_world(world, timeStep)
       step = False