add invalid trajectory callback & invalid traj return & register all 3 variantes of table tennis tasks

fancy_gym/black_box/black_box_wrapper.py

@@ -63,14 +63,6 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         self.traj_gen_action_space = self._get_traj_gen_action_space()
         self.action_space = self._get_action_space()
 
-        # no goal learning
-        # tricky_action_upperbound = [np.inf] * (self.traj_gen_action_space.shape[0] - 7)
-        # tricky_action_lowerbound = [-np.inf] * (self.traj_gen_action_space.shape[0] - 7)
-        # self.action_space = spaces.Box(np.array(tricky_action_lowerbound), np.array(tricky_action_upperbound), dtype=np.float32)
-        self.action_space.low[0] = 0.8
-        self.action_space.high[0] = 1.5
-        self.action_space.low[1] = 0.05
-        self.action_space.high[1] = 0.15
         self.observation_space = self._get_observation_space()
 
         # rendering
@@ -93,8 +85,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         return observation.astype(self.observation_space.dtype)
 
     def get_trajectory(self, action: np.ndarray) -> Tuple:
-        # duration = self.duration
-        duration = self.duration - self.current_traj_steps * self.dt
+        duration = self.duration
+        # duration = self.duration - self.current_traj_steps * self.dt
         if self.learn_sub_trajectories:
             duration = None
             # reset  with every new call as we need to set all arguments, such as tau, delay, again.
@@ -157,8 +149,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         # TODO remove this part, right now only needed for beer pong
         # mp_params, env_spec_params, proceed = self.env.episode_callback(action, self.traj_gen)
         position, velocity = self.get_trajectory(action)
-        traj_is_valid = self.env.episode_callback(action, position, velocity)
-
+        traj_is_valid = self.env.preprocessing_and_validity_callback(action, position, velocity)
+        # insert validation here
         trajectory_length = len(position)
         rewards = np.zeros(shape=(trajectory_length,))
         if self.verbose >= 2:
@@ -169,7 +161,11 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         infos = dict()
         done = False
 
-        if traj_is_valid:
+        if not traj_is_valid:
+            obs, trajectory_return, done, infos = self.env.invalid_traj_callback(action, position, velocity,
+                                                                                 self.return_context_observation)
+            return self.observation(obs), trajectory_return, done, infos
+        else:
             self.plan_steps += 1
             for t, (pos, vel) in enumerate(zip(position, velocity)):
                 current_pos = self.current_pos
@@ -215,10 +211,6 @@ class BlackBoxWrapper(gym.ObservationWrapper):
             infos['trajectory_length'] = t + 1
             trajectory_return = self.reward_aggregation(rewards[:t + 1])
             return self.observation(obs), trajectory_return, done, infos
-        else:
-            obs, trajectory_return, done, infos = self.env.invalid_traj_callback(action, position, velocity,
-                                                                                 self.return_context_observation)
-            return self.observation(obs), trajectory_return, done, infos
 
     def render(self, **kwargs):
         """Only set render options here, such that they can be used during the rollout.

fancy_gym/black_box/raw_interface_wrapper.py

@@ -52,6 +52,19 @@ class RawInterfaceWrapper(gym.Wrapper):
         """
         return self.env.dt
 
+    def preprocessing_and_validity_callback(self, action: np.ndarray, pos_traj: np.ndarray, vel_traj: np.ndarray) \
+            -> Tuple[bool, np.ndarray, np.ndarray]:
+        """
+        Used to preprocess the action and check if the desired trajectory is valid.
+        """
+        return True, pos_traj, vel_traj
+
+    def set_episode_arguments(self, action, pos_traj, vel_traj):
+        """
+        Used to set the arguments for env that valid for the whole episode
+        """
+        return pos_traj, vel_traj
+
     def episode_callback(self, action: np.ndarray, pos_traj: np.ndarray, vel_traj: np.array) -> Tuple[bool]:
         """
         Used to extract the parameters for the movement primitive and other parameters from an action array which might
@@ -68,7 +81,6 @@ class RawInterfaceWrapper(gym.Wrapper):
 
     def invalid_traj_callback(self, action: np.ndarray, pos_traj: np.ndarray, vel_traj: np.ndarray) -> Tuple[np.ndarray, float, bool, dict]:
         """
-        Used to return a fake return from the environment if the desired trajectory is invalid.
+        Used to return a artificial return from the env if the desired trajectory is invalid.
         """
-        obs = np.zeros(1)
-        return obs, 0, True, {}
+        return np.zeros(1), 0, True, {}