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 85c49cc61..e7faddebe 100644 --- a/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py +++ b/examples/pybullet/gym/pybullet_envs/deep_mimic/env/goals.py @@ -1,6 +1,8 @@ from enum import Enum import random import math +import pybullet as pb +import numpy def randomVal(lowerBound, upperbound): difference = upperbound - lowerBound @@ -26,10 +28,12 @@ class GoalType(Enum): class Goal: def __init__(self, goal_type: GoalType): self.goal_type = goal_type - self.generateGoalData() + self.follow_rot = True self.is_hit_prev = False - def generateGoalData(self, modelPos=[0,0]): + self.generateGoalData() + + def generateGoalData(self, modelPos=[0,0,0], modelOrient=[0,0,0,1]): if self.goal_type == GoalType.NoGoal: self.goal_data = [] @@ -51,7 +55,24 @@ class Goal: # Y axis up, z axis in different direction self.goal_data = [-x, z, y] - self.world_pos = [-x + modelPos[0], z, y + modelPos[1]] + + if self.follow_rot: + # Take rotation of human model into account + eulerAngles = pb.getEulerFromQuaternion(modelOrient) + # Only Y angle matters + eulerAngles = [0, eulerAngles[1], 0] + yQuat = pb.getQuaternionFromEuler(eulerAngles) + rotMatList = pb.getMatrixFromQuaternion(yQuat) + rotMat = numpy.array([rotMatList[0:3], rotMatList[3:6], rotMatList[6:9]]) + vec = numpy.array(self.goal_data) + rotatedVec = numpy.dot(rotMat, vec) + self.world_pos = rotatedVec.tolist() + + self.world_pos = [ self.world_pos[0] + modelPos[0], + self.world_pos[1], + self.world_pos[2] + modelPos[2]] + else: + self.world_pos = [-x + modelPos[0], z, y + modelPos[2]] elif self.goal_type == GoalType.TargetHeading: # Direction: 2D unit vector 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 e00d67115..4a0ab3a67 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 @@ -809,5 +809,5 @@ class HumanoidStablePD(object): return self._pybullet_client\ .getBasePositionAndOrientation(self._sim_model) - def getLinkPosition(self, link_id): - return self._pybullet_client.getLinkState(self._sim_model, link_id)[0] + def getLinkPositionAndOrientation(self, link_id): + return tuple(self._pybullet_client.getLinkState(self._sim_model, link_id)[0:2]) 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 8a93dd170..cde6a5a76 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 @@ -101,8 +101,8 @@ class PyBulletDeepMimicEnv(Env): #self._pybullet_client.stepSimulation() self._humanoid.resetPose() # generate new goal - humanPos = self._humanoid.getLinkPosition(0) - self.goal.generateGoalData([humanPos[0], humanPos[2]]) + humanPos, humanOrient = self._humanoid.getLinkPositionAndOrientation(0) + self.goal.generateGoalData(humanPos, humanOrient) self.needs_update_time = self.t - 1 #force update def get_num_agents(self): @@ -257,7 +257,7 @@ class PyBulletDeepMimicEnv(Env): goal_weight = 0.3 if self.goal.goal_type == GoalType.Strike: - linkPos = self._humanoid.getLinkPosition(HumanoidLinks.rightAnkle) + linkPos, linkOrient = self._humanoid.getLinkPositionAndOrientation(HumanoidLinks.rightAnkle) reward = mimic_weight * reward + goal_weight * self.calcStrikeGoalReward(linkPos) return reward @@ -289,7 +289,7 @@ class PyBulletDeepMimicEnv(Env): #print("pybullet_deep_mimic_env:update timeStep=",timeStep," t=",self.t) self._pybullet_client.setTimeStep(timeStep) self._humanoid._timeStep = timeStep - self.updateGoal(self._humanoid.getLinkPosition(HumanoidLinks.rightAnkle)) + self.updateGoal(self._humanoid.getLinkPositionAndOrientation(HumanoidLinks.rightAnkle)[0]) if self.target_id is not None: # TODO: check goal type self.updateDrawStrikeGoal() @@ -411,8 +411,8 @@ class PyBulletDeepMimicEnv(Env): if self.goal.goal_type == GoalType.Strike: if self.prevCycleCount != self._humanoid.cycleCount: # generate new goal - humanPos = self._humanoid.getLinkPosition(0) - self.goal.generateGoalData([humanPos[0], humanPos[2]]) + humanPos, humanOrient = self._humanoid.getLinkPositionAndOrientation(0) + self.goal.generateGoalData(humanPos, humanOrient) self.prevCycleCount = self._humanoid.cycleCount goalPos = self.goal.world_pos