Merge pull request #56 from HongyiZhouCN/55-table-tennis-dev

#55 add table tennis environment to fancy_gym

fancy_gym/black_box/black_box_wrapper.py

@@ -137,9 +137,9 @@ class BlackBoxWrapper(gym.ObservationWrapper):
     def step(self, action: np.ndarray):
         """ This function generates a trajectory based on a MP and then does the usual loop over reset and step"""
 
-        # TODO remove this part, right now only needed for beer pong
-        mp_params, env_spec_params = self.env.episode_callback(action, self.traj_gen)
-        position, velocity = self.get_trajectory(mp_params)
+        position, velocity = self.get_trajectory(action)
+        position, velocity = self.env.set_episode_arguments(action, position, velocity)
+        traj_is_valid, position, velocity = self.env.preprocessing_and_validity_callback(action, position, velocity)
 
         trajectory_length = len(position)
         rewards = np.zeros(shape=(trajectory_length,))
@@ -151,6 +151,11 @@ class BlackBoxWrapper(gym.ObservationWrapper):
         infos = dict()
         done = False
 
+        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
+
         self.plan_steps += 1
         for t, (pos, vel) in enumerate(zip(position, velocity)):
             step_action = self.tracking_controller.get_action(pos, vel, self.current_pos, self.current_vel)
@@ -171,8 +176,8 @@ class BlackBoxWrapper(gym.ObservationWrapper):
                 self.env.render(**self.render_kwargs)
 
             if done or (self.replanning_schedule(self.current_pos, self.current_vel, obs, c_action,
-                                                 t + 1 + self.current_traj_steps)
-                        and self.plan_steps < self.max_planning_times):
+                                                t + 1 + self.current_traj_steps)
+                    and self.plan_steps < self.max_planning_times):
 
                 if self.condition_on_desired:
                     self.condition_pos = pos

fancy_gym/black_box/raw_interface_wrapper.py

@@ -52,8 +52,38 @@ class RawInterfaceWrapper(gym.Wrapper):
         """
         return self.env.dt
 
-    def episode_callback(self, action: np.ndarray, traj_gen: MPInterface) -> Tuple[
-        np.ndarray, Union[np.ndarray, None]]:
+    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.
+        Args:
+            action:  a vector instance of the whole action space, includes traj_gen parameters and additional parameters if
+            specified, else only traj_gen parameters
+            pos_traj: a vector instance of the raw position trajectory
+            vel_traj: a vector instance of the raw velocity trajectory
+        Returns:
+            validity flag: bool, True if the raw trajectory is valid, False if not
+            pos_traj: a vector instance of the preprocessed position trajectory 
+            vel_traj: a vector instance of the preprocessed velocity trajectory
+        """
+        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
+        deprecated, replaced by preprocessing_and_validity_callback
+        Args:
+            action:  a vector instance of the whole action space, includes traj_gen parameters and additional parameters if
+            specified, else only traj_gen parameters
+            pos_traj: a vector instance of the raw position trajectory
+            vel_traj: a vector instance of the raw velocity trajectory
+        Returns:
+            pos_traj: a vector instance of the preprocessed position trajectory
+            vel_traj: a vector instance of the preprocessed velocity trajectory
+        """
+        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
         include other actions like ball releasing time for the beer pong environment.
@@ -65,4 +95,20 @@ class RawInterfaceWrapper(gym.Wrapper):
         Returns:
             Tuple: mp_arguments and other arguments
         """
-        return action, None
+        return True
+
+    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 artificial return from the env if the desired trajectory is invalid.
+        Args:
+            action:  a vector instance of the whole action space, includes traj_gen parameters and additional parameters if
+            specified, else only traj_gen parameters
+            pos_traj: a vector instance of the raw position trajectory
+            vel_traj: a vector instance of the raw velocity trajectory
+        Returns:
+            obs: artificial observation if the trajectory is invalid, by default a zero vector
+            reward: artificial reward if the trajectory is invalid, by default 0
+            done: artificial done if the trajectory is invalid, by default True
+            info: artificial info if the trajectory is invalid, by default empty dict
+        """
+        return np.zeros(1), 0, True, {}