smaller ctxt range hopper jump + step-based rew for bp

alr_envs/alr/mujoco/beerpong/beerpong.py

@@ -187,35 +187,65 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle):
 
 
 class ALRBeerBongEnvStepBased(ALRBeerBongEnv):
+    def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, rndm_goal=False, cup_goal_pos=None):
+        super().__init__(frame_skip, apply_gravity_comp, noisy, rndm_goal, cup_goal_pos)
+        self.release_step = 62  # empirically evaluated for frame_skip=2!
 
-    def _set_action_space(self):
-        bounds = super(ALRBeerBongEnvStepBased, self)._set_action_space()
-        min_bound = np.concatenate(([-1], bounds.low), dtype=bounds.dtype)
-        max_bound = np.concatenate(([1], bounds.high), dtype=bounds.dtype)
-        self.action_space = spaces.Box(low=min_bound, high=max_bound, dtype=bounds.dtype)
-        return self.action_space
+    # def _set_action_space(self):
+    #     bounds = super(ALRBeerBongEnvStepBased, self)._set_action_space()
+    #     min_bound = np.concatenate(([-1], bounds.low), dtype=bounds.dtype)
+    #     max_bound = np.concatenate(([1], bounds.high), dtype=bounds.dtype)
+    #     self.action_space = spaces.Box(low=min_bound, high=max_bound, dtype=bounds.dtype)
+    #     return self.action_space
+
+    # def step(self, a):
+    #     self.release_step = self._steps if a[0]>=0 and self.release_step >= self._steps else self.release_step
+    #     return super(ALRBeerBongEnvStepBased, self).step(a[1:])
+    #
+    # def reset(self):
+    #     ob = super(ALRBeerBongEnvStepBased, self).reset()
+    #     self.release_step = self.ep_length + 1
+    #     return ob
 
     def step(self, a):
-        self.release_step = self._steps if a[0]>=0 and self.release_step >= self._steps else self.release_step
-        return super(ALRBeerBongEnvStepBased, self).step(a[1:])
+        if self._steps < self.release_step:
+            return super(ALRBeerBongEnvStepBased, self).step(a)
+        else:
+            reward = 0
+            done = False
+            while not done:
+                sub_ob, sub_reward, done, sub_infos = super(ALRBeerBongEnvStepBased, self).step(np.zeros(a.shape))
+                if not done or sub_infos['sim_crash']:
+                    reward += sub_reward
+                else:
+                    ball_pos = self.sim.data.body_xpos[self.sim.model._body_name2id["ball"]].copy()
+                    cup_goal_dist_final = np.linalg.norm(ball_pos - self.sim.data.site_xpos[
+                        self.sim.model._site_name2id["cup_goal_final_table"]].copy())
+                    cup_goal_dist_top = np.linalg.norm(ball_pos - self.sim.data.site_xpos[
+                        self.sim.model._site_name2id["cup_goal_table"]].copy())
+                    if sub_infos['success']:
+                        dist_rew = -cup_goal_dist_final**2
+                    else:
+                        dist_rew = -0.5*cup_goal_dist_final**2 - cup_goal_dist_top**2
+                    reward = reward - sub_infos['action_cost'] + dist_rew
+            infos = sub_infos
+            ob = sub_ob
+        return ob, reward, done, infos
+
 
-    def reset(self):
-        ob = super(ALRBeerBongEnvStepBased, self).reset()
-        self.release_step = self.ep_length + 1
-        return ob
 
 if __name__ == "__main__":
     # env = ALRBeerBongEnv(rndm_goal=True)
-    env = ALRBeerBongEnvStepBased(rndm_goal=True)
+    env = ALRBeerBongEnvStepBased(frame_skip=2, rndm_goal=True)
     import time
     env.reset()
     env.render("human")
     for i in range(1500):
-        # ac = 10 * env.action_space.sample()[0:7]
-        ac = np.zeros(8)
-        ac[0] = -1
-        if env._steps > 150:
-            ac[0] = 1
+        ac = 10 * env.action_space.sample()
+        # ac = np.zeros(7)
+        # ac[0] = -1
+        # if env._steps > 150:
+        #     ac[0] = 1
         obs, rew, d, info = env.step(ac)
         env.render("human")
         print(env.dt)

alr_envs/alr/mujoco/beerpong/beerpong_reward_staged.py

@@ -90,7 +90,7 @@ class BeerPongReward:
         self.dist_ground_cup = np.linalg.norm(ball_pos-goal_pos) \
             if self.ball_ground_contact_first and self.dist_ground_cup == -1 else self.dist_ground_cup
         action_cost = np.sum(np.square(action))
-        self.action_costs.append(action_cost)
+        self.action_costs.append(np.copy(action_cost))
         # # ##################### Reward function which does not force to bounce once on the table (quad dist) ############
 
         self._is_collided = self._check_collision_with_itself(env.sim, self.robot_collision_ids)
@@ -110,8 +110,9 @@ class BeerPongReward:
                 else:
                     min_dist_coeff, final_dist_coeff, ground_contact_dist_coeff, rew_offset = 0, 1, 0 ,0
             # dist_ground_cup = 1 * self.dist_ground_cup
+            action_cost = 1e-4 * np.mean(action_cost)
             reward = rew_offset - min_dist_coeff * min_dist ** 2 - final_dist_coeff * final_dist ** 2 - \
-                     1e-4 * np.mean(action_cost) - ground_contact_dist_coeff*self.dist_ground_cup ** 2
+                     action_cost - ground_contact_dist_coeff*self.dist_ground_cup ** 2
             # 1e-7*np.mean(action_cost)
             # release step punishment
             min_time_bound = 0.1
@@ -124,7 +125,8 @@ class BeerPongReward:
             # print('release time :', release_time)
             # print('dist_ground_cup :', dist_ground_cup)
         else:
-            reward = - 1e-2 * action_cost
+            action_cost = 1e-2 * action_cost
+            reward = - action_cost
             # reward = - 1e-4 * action_cost
             # reward = 0
             success = False