From 541a145e2d0e53fc8cf75b98cad7d492bdd080fd Mon Sep 17 00:00:00 2001
From: Bart Moyaers <bart.moyaers@kuleuven.be>
Date: Thu, 13 Feb 2020 14:39:21 +0100
Subject: [PATCH] enable follow_rot to be set correct grab goal settings enable
 different links for different goal types

---
 .../gym/pybullet_envs/deep_mimic/env/goals.py       | 13 +++++++------
 .../deep_mimic/env/pybullet_deep_mimic_env.py       |  7 +++++--
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py
index 6b4684221..5f7646c28 100644
--- a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py
+++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py
@@ -64,8 +64,9 @@ class Strike(Goal):
                 distance_RB: RandomBounds,
                 height_RB: RandomBounds,
                 rot_RB: RandomBounds,
-                hit_range = 0.2):
-        self.follow_rot = True
+                hit_range = 0.2,
+                follow_rot = True):
+        self.follow_rot = follow_rot
         self.is_hit_prev = False
         self.distance_RB = distance_RB
         self.height_RB = height_RB
@@ -127,12 +128,12 @@ class Kick(Strike):
         super().__init__(distance_RB, height_RB, rot_RB, hit_range=0.2)
 
 class Grab(Strike):
-    def __inite__(self):
+    def __init__(self):
         distance_RB = RandomBounds(0.6, 0.8)
-        height_RB = RandomBounds(0.8, 0.9)
-        rot_RB = RandomBounds(-1, 1)
+        height_RB = RandomBounds(0.8, 1.1)
+        rot_RB = RandomBounds(3.14159 - 0.5, 3.14159 + 0.5)
         hit_range = 0.1
-        super().__init__(distance_RB, height_RB, rot_RB, hit_range=hit_range)
+        super().__init__(distance_RB, height_RB, rot_RB, hit_range=hit_range, follow_rot=False)
 
 class TargetHeadingGoal(Goal):
     def __init__(self):
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 0cba9d279..eb51f20d5 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
@@ -6,7 +6,7 @@ from pybullet_utils import bullet_client
 import time
 from pybullet_envs.deep_mimic.env import motion_capture_data
 from pybullet_envs.deep_mimic.env import humanoid_stable_pd
-from pybullet_envs.deep_mimic.env.goals import GoalType, Goal, createGoal
+from pybullet_envs.deep_mimic.env.goals import GoalType, Goal, createGoal, Kick, Grab
 from pybullet_envs.deep_mimic.env.humanoid_link_ids import HumanoidLinks
 import pybullet_data
 import pybullet as p1
@@ -257,7 +257,10 @@ class PyBulletDeepMimicEnv(Env):
     goal_weight = 0.3
 
     if self.goal.goal_type == GoalType.Strike:
-      linkPos, linkOrient = self._humanoid.getLinkPositionAndOrientation(HumanoidLinks.rightAnkle)
+      if self.goal.__class__ == Kick:
+        linkPos, linkOrient = self._humanoid.getLinkPositionAndOrientation(HumanoidLinks.rightAnkle)
+      else: #elif type(self.goal) == Grab:
+        linkPos, linkOrient = self._humanoid.getLinkPositionAndOrientation(HumanoidLinks.rightWrist)
       reward = mimic_weight * reward + goal_weight * self.calcStrikeGoalReward(linkPos)
 
     return reward