diff --git a/examples/pybullet/gym/pybullet_envs/gym_locomotion_envs.py b/examples/pybullet/gym/pybullet_envs/gym_locomotion_envs.py index 6220b28cf..6e042efdf 100644 --- a/examples/pybullet/gym/pybullet_envs/gym_locomotion_envs.py +++ b/examples/pybullet/gym/pybullet_envs/gym_locomotion_envs.py @@ -14,6 +14,7 @@ class WalkerBaseBulletEnv(MJCFBaseBulletEnv): self.walk_target_x = 1e3 # kilometer away self.walk_target_y = 0 self.stateId=-1 + self.alive = None def create_single_player_scene(self, bullet_client): @@ -39,6 +40,12 @@ class WalkerBaseBulletEnv(MJCFBaseBulletEnv): return r + + def _isDone(self): + if self._alive is not None: + return self._alive < 0 + else: + return False def move_robot(self, init_x, init_y, init_z): "Used by multiplayer stadium to move sideways, to another running lane." @@ -60,8 +67,8 @@ class WalkerBaseBulletEnv(MJCFBaseBulletEnv): state = self.robot.calc_state() # also calculates self.joints_at_limit - alive = float(self.robot.alive_bonus(state[0]+self.robot.initial_z, self.robot.body_rpy[1])) # state[0] is body height above ground, body_rpy[1] is pitch - done = alive < 0 + self.alive = float(self.robot.alive_bonus(state[0]+self.robot.initial_z, self.robot.body_rpy[1])) # state[0] is body height above ground, body_rpy[1] is pitch + done = self._isDone() if not np.isfinite(state).all(): print("~INF~", state) done = True @@ -135,6 +142,9 @@ class HalfCheetahBulletEnv(WalkerBaseBulletEnv): def __init__(self): self.robot = HalfCheetah() WalkerBaseBulletEnv.__init__(self, self.robot) + + def _isDone(self): + return False class AntBulletEnv(WalkerBaseBulletEnv): def __init__(self): @@ -172,4 +182,3 @@ class HumanoidFlagrunHarderBulletEnv(HumanoidBulletEnv): s = HumanoidBulletEnv.create_single_player_scene(self, bullet_client) s.zero_at_running_strip_start_line = False return s -