start refactor and biac dev merge

alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup.py

@@ -41,7 +41,7 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
             from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward
             reward_function = BallInACupReward
         else:
-            raise ValueError("Unknown reward type")
+            raise ValueError("Unknown reward type: {}".format(reward_type))
         self.reward_function = reward_function(self.sim_steps)
 
     @property
@@ -66,6 +66,10 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
     def current_vel(self):
         return self.sim.data.qvel[0:7].copy()
 
+    def reset(self):
+        self.reward_function.reset(None)
+        return super().reset()
+
     def reset_model(self):
         init_pos_all = self.init_qpos.copy()
         init_pos_robot = self._start_pos
@@ -100,14 +104,18 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
             done = success or self._steps == self.sim_steps - 1 or is_collided
             self._steps += 1
         else:
-            reward = -2
+            reward = -2000
             success = False
             is_collided = False
             done = True
         return ob, reward, done, dict(reward_dist=reward_dist,
                                       reward_ctrl=reward_ctrl,
                                       velocity=angular_vel,
-                                      traj=self._q_pos, is_success=success,
+                                      # traj=self._q_pos,
+                                      action=a,
+                                      q_pos=self.sim.data.qpos[0:7].ravel().copy(),
+                                      q_vel=self.sim.data.qvel[0:7].ravel().copy(),
+                                      is_success=success,
                                       is_collided=is_collided, sim_crash=crash)
 
     def check_traj_in_joint_limits(self):
@@ -148,6 +156,22 @@ class ALRBallInACupEnv(alr_mujoco_env.AlrMujocoEnv, utils.EzPickle):
         des_vel_full[5] = des_vel[2]
         return des_vel_full
 
+    def render(self, render_mode, **render_kwargs):
+        if render_mode == "plot_trajectory":
+            if self._steps == 1:
+                import matplotlib.pyplot as plt
+                # plt.ion()
+                self.fig, self.axs = plt.subplots(3, 1)
+
+            if self._steps <= 1750:
+                for ax, cp in zip(self.axs, self.current_pos[1::2]):
+                    ax.scatter(self._steps, cp, s=2, marker=".")
+
+            # self.fig.show()
+
+        else:
+            super().render(render_mode, **render_kwargs)
+
 
 if __name__ == "__main__":
     env = ALRBallInACupEnv()

alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py

@@ -22,7 +22,7 @@ class BallInACupReward(alr_reward_fct.AlrReward):
         self.goal_final_id = None
         self.collision_ids = None
         self._is_collided = False
-        self.collision_penalty = 1
+        self.collision_penalty = 1000
 
         self.ball_traj = None
         self.dists = None
@@ -37,6 +37,7 @@ class BallInACupReward(alr_reward_fct.AlrReward):
         self.dists_final = []
         self.costs = []
         self.action_costs = []
+        self.angle_costs = []
         self.cup_angles = []
 
     def compute_reward(self, action, env):
@@ -56,8 +57,11 @@ class BallInACupReward(alr_reward_fct.AlrReward):
         self.dists_final.append(np.linalg.norm(goal_final_pos - ball_pos))
         self.ball_traj[env._steps, :] = ball_pos
         cup_quat = np.copy(env.sim.data.body_xquat[env.sim.model._body_name2id["cup"]])
-        self.cup_angles.append(np.arctan2(2 * (cup_quat[0] * cup_quat[1] + cup_quat[2] * cup_quat[3]),
-                                          1 - 2 * (cup_quat[1]**2 + cup_quat[2]**2)))
+        cup_angle = np.arctan2(2 * (cup_quat[0] * cup_quat[1] + cup_quat[2] * cup_quat[3]),
+                                          1 - 2 * (cup_quat[1]**2 + cup_quat[2]**2))
+        cost_angle = (cup_angle - np.pi / 2) ** 2
+        self.angle_costs.append(cost_angle)
+        self.cup_angles.append(cup_angle)
 
         action_cost = np.sum(np.square(action))
         self.action_costs.append(action_cost)
@@ -67,20 +71,21 @@ class BallInACupReward(alr_reward_fct.AlrReward):
         if env._steps == env.sim_steps - 1 or self._is_collided:
             t_min_dist = np.argmin(self.dists)
             angle_min_dist = self.cup_angles[t_min_dist]
-            cost_angle = (angle_min_dist - np.pi / 2)**2
+            # cost_angle = (angle_min_dist - np.pi / 2)**2
 
-            min_dist = self.dists[t_min_dist]
+
+            # min_dist = self.dists[t_min_dist]
             dist_final = self.dists_final[-1]
             min_dist_final = np.min(self.dists_final)
 
-            cost = 0.5 * dist_final + 0.05 * cost_angle  # TODO: Increase cost_angle weight  # 0.5 * min_dist +
+            # cost = 0.5 * dist_final + 0.05 * cost_angle  # TODO: Increase cost_angle weight  # 0.5 * min_dist +
             # reward = np.exp(-2 * cost) - 1e-2 * action_cost - self.collision_penalty * int(self._is_collided)
             # reward = - dist_final**2 - 1e-4 * cost_angle - 1e-5 * action_cost - self.collision_penalty * int(self._is_collided)
-            reward = - dist_final**2 - min_dist_final**2 - 1e-4 * cost_angle - 1e-5 * action_cost - self.collision_penalty * int(self._is_collided)
+            reward = - dist_final**2 - min_dist_final**2 - 1e-4 * cost_angle - 1e-3 * action_cost - self.collision_penalty * int(self._is_collided)
             success = dist_final < 0.05 and ball_in_cup and not self._is_collided
             crash = self._is_collided
         else:
-            reward = - 1e-5 * action_cost  # TODO: increase action_cost weight
+            reward = - 1e-3 * action_cost - 1e-4 * cost_angle  # TODO: increase action_cost weight
             success = False
             crash = False