Merge branch 'master' into future

.github/workflows/host_docs_static.yml

@@ -1,40 +0,0 @@
-name: Deploy static docs to Pages
-
-on:
-  push:
-    branches: ["release"]
-
-  # Allows you to run this workflow manually from the Actions tab
-  workflow_dispatch:
-
-# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages
-permissions:
-  contents: read
-  pages: write
-  id-token: write
-
-# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued.
-# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete.
-concurrency:
-  group: "pages"
-  cancel-in-progress: false
-
-jobs:
-  # Single deploy job since we're just deploying
-  deploy:
-    environment:
-      name: github-pages
-      url: ${{ steps.deployment.outputs.page_url }}
-    runs-on: ubuntu-latest
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v4
-      - name: Setup Pages
-        uses: actions/configure-pages@v4
-      - name: Upload artifact
-        uses: actions/upload-pages-artifact@v3
-        with:
-          path: 'docs/build/html'
-      - name: Deploy to GitHub Pages
-        id: deployment
-        uses: actions/deploy-pages@v4

fancy_gym/envs/__init__.py

@@ -26,7 +26,7 @@ from .mujoco.walker_2d_jump.walker_2d_jump import MAX_EPISODE_STEPS_WALKERJUMP
 from .mujoco.box_pushing.box_pushing_env import BoxPushingDense, BoxPushingTemporalSparse, \
     BoxPushingTemporalSpatialSparse, MAX_EPISODE_STEPS_BOX_PUSHING
 from .mujoco.table_tennis.table_tennis_env import TableTennisEnv, TableTennisWind, TableTennisGoalSwitching, \
-    MAX_EPISODE_STEPS_TABLE_TENNIS
+    MAX_EPISODE_STEPS_TABLE_TENNIS, MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER
 from .mujoco.table_tennis.mp_wrapper import TT_MPWrapper as MPWrapper_TableTennis
 from .mujoco.table_tennis.mp_wrapper import TT_MPWrapper_Replan as MPWrapper_TableTennis_Replan
 from .mujoco.table_tennis.mp_wrapper import TTVelObs_MPWrapper as MPWrapper_TableTennis_VelObs
@@ -135,6 +135,19 @@ register(
     }
 )
 
+register(
+    id='fancy/HopperJumpMarkov-v0',
+    entry_point='fancy_gym.envs.mujoco:HopperJumpMarkovRew',
+    mp_wrapper=mujoco.hopper_jump.MPWrapper,
+    max_episode_steps=MAX_EPISODE_STEPS_HOPPERJUMP,
+    kwargs={
+        "sparse": False,
+        "healthy_reward": 1.0,
+        "contact_weight": 0.0,
+        "height_weight": 3.0,
+    }
+)
+
 # TODO: Add [MPs] later when finished (old TODO I moved here during refactor)
 register(
     id='fancy/AntJump-v0',
@@ -290,6 +303,34 @@ register(
     }
 )
 
+register(
+    id='fancy/TableTennisRndRobot-v0',
+    entry_point='fancy_gym.envs.mujoco:TableTennisRandomInit',
+    max_episode_steps=MAX_EPISODE_STEPS_TABLE_TENNIS,
+    kwargs={
+        'random_pos_scale': 0.1,
+        'random_vel_scale': 0.0,
+    }
+)
+
+register(
+    id='fancy/TableTennisMarkovian-v0',
+    entry_point='fancy_gym.envs.mujoco:TableTennisMarkovian',
+    max_episode_steps=MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER,
+    kwargs={
+    }
+)
+
+register(
+    id='fancy/TableTennisRndRobotMarkovian-v0',
+    entry_point='fancy_gym.envs.mujoco:TableTennisMarkovian',
+    max_episode_steps=MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER,
+    kwargs={
+        'random_pos_scale': 0.1,
+        'random_vel_scale': 0.0,
+    }
+)
+
 # Air Hockey environments
 for env_mode in ["7dof-hit", "7dof-defend", "3dof-hit", "3dof-defend", "7dof-hit-airhockit2023", "7dof-defend-airhockit2023"]:
     register(

fancy_gym/envs/mujoco/__init__.py

@@ -1,14 +1,14 @@
 from .ant_jump.ant_jump import AntJumpEnv
 from .beerpong.beerpong import BeerPongEnv, BeerPongEnvStepBasedEpisodicReward
 from .half_cheetah_jump.half_cheetah_jump import HalfCheetahJumpEnv
-from .hopper_jump.hopper_jump import HopperJumpEnv
+from .hopper_jump.hopper_jump import HopperJumpEnv, HopperJumpMarkovRew
 from .hopper_jump.hopper_jump_on_box import HopperJumpOnBoxEnv
 from .hopper_throw.hopper_throw import HopperThrowEnv
 from .hopper_throw.hopper_throw_in_basket import HopperThrowInBasketEnv
 from .reacher.reacher import ReacherEnv
 from .walker_2d_jump.walker_2d_jump import Walker2dJumpEnv
 from .box_pushing.box_pushing_env import BoxPushingDense, BoxPushingTemporalSparse, BoxPushingTemporalSpatialSparse
-from .table_tennis.table_tennis_env import TableTennisEnv, TableTennisWind, TableTennisGoalSwitching
+from .table_tennis.table_tennis_env import TableTennisEnv, TableTennisWind, TableTennisGoalSwitching, TableTennisMarkovian, TableTennisRandomInit
 
 try:
     from .air_hockey.air_hockey_env_wrapper import AirHockeyEnv

fancy_gym/envs/mujoco/ant_jump/ant_jump.py

@@ -115,7 +115,6 @@ class AntJumpEnv(AntEnvCustomXML):
                          contact_force_range=contact_force_range,
                          reset_noise_scale=reset_noise_scale,
                          exclude_current_positions_from_observation=exclude_current_positions_from_observation, **kwargs)
-        self.render_active = False
 
     def step(self, action):
         self.current_step += 1
@@ -154,15 +153,8 @@ class AntJumpEnv(AntEnvCustomXML):
         }
         truncated = False
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return obs, reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _get_obs(self):
         return np.append(super()._get_obs(), self.goal)
 

fancy_gym/envs/mujoco/beerpong/beerpong.py

@@ -44,7 +44,6 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
     }
 
     def __init__(self, **kwargs):
-        utils.EzPickle.__init__(self)
         self._steps = 0
         # Small Context -> Easier. Todo: Should we do different versions?
         # self.xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "beerpong_wo_cup.xml")
@@ -90,7 +89,7 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
             observation_space=self.observation_space,
             **kwargs
         )
-        self.render_active = False
+        utils.EzPickle.__init__(self)
 
     @property
     def start_pos(self):
@@ -170,15 +169,8 @@ class BeerPongEnv(MujocoEnv, utils.EzPickle):
 
         truncated = False
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return ob, reward, terminated, truncated, infos
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _get_obs(self):
         theta = self.data.qpos.flat[:7].copy()
         theta_dot = self.data.qvel.flat[:7].copy()

fancy_gym/envs/mujoco/box_pushing/box_pushing_env.py

@@ -4,9 +4,9 @@ import numpy as np
 from gymnasium import utils, spaces
 from gymnasium.envs.mujoco import MujocoEnv
 from fancy_gym.envs.mujoco.box_pushing.box_pushing_utils import rot_to_quat, get_quaternion_error, rotation_distance
-from fancy_gym.envs.mujoco.box_pushing.box_pushing_utils import rot_to_quat, get_quaternion_error, rotation_distance
 from fancy_gym.envs.mujoco.box_pushing.box_pushing_utils import q_max, q_min, q_dot_max, q_torque_max
 from fancy_gym.envs.mujoco.box_pushing.box_pushing_utils import desired_rod_quat
+from fancy_gym.envs.mujoco.box_pushing.box_pushing_utils import calculate_jerk_profile, calculate_mean_squared_jerk, calculate_dimensionless_jerk, calculate_maximum_jerk
 
 import mujoco
 
@@ -50,6 +50,7 @@ class BoxPushingEnvBase(MujocoEnv, utils.EzPickle):
         self._desired_rod_quat = desired_rod_quat
 
         self._episode_energy = 0.
+        self.velocity_profile = []
 
         self.observation_space = spaces.Box(
             low=-np.inf, high=np.inf, shape=(28,), dtype=np.float64
@@ -61,7 +62,6 @@ class BoxPushingEnvBase(MujocoEnv, utils.EzPickle):
                            frame_skip=self.frame_skip,
                            observation_space=self.observation_space, **kwargs)
         self.action_space = spaces.Box(low=-1, high=1, shape=(7,))
-        self.render_active = False
 
     def step(self, action):
         action = 10 * np.clip(action, self.action_space.low, self.action_space.high)
@@ -69,6 +69,8 @@ class BoxPushingEnvBase(MujocoEnv, utils.EzPickle):
 
         unstable_simulation = False
 
+        self.velocity_profile.append(self.data.qvel[:7].copy())
+
         try:
             self.do_simulation(resultant_action, self.frame_skip)
         except Exception as e:
@@ -98,11 +100,15 @@ class BoxPushingEnvBase(MujocoEnv, utils.EzPickle):
         obs = self._get_obs()
         box_goal_pos_dist = 0. if not episode_end else np.linalg.norm(box_pos - target_pos)
         box_goal_quat_dist = 0. if not episode_end else rotation_distance(box_quat, target_quat)
+        mean_squared_jerk, maximum_jerk, dimensionless_jerk = (0.0,0.0,0.0) if not episode_end else self.calculate_smoothness_metrics(np.array(self.velocity_profile), self.dt)
         infos = {
             'episode_end': episode_end,
             'box_goal_pos_dist': box_goal_pos_dist,
             'box_goal_rot_dist': box_goal_quat_dist,
             'episode_energy': 0. if not episode_end else self._episode_energy,
+            'mean_squared_jerk': mean_squared_jerk,
+            'maximum_jerk': maximum_jerk,
+            'dimensionless_jerk': dimensionless_jerk,
             'is_success': True if episode_end and box_goal_pos_dist < 0.05 and box_goal_quat_dist < 0.5 else False,
             'num_steps': self._steps
         }
@@ -110,14 +116,27 @@ class BoxPushingEnvBase(MujocoEnv, utils.EzPickle):
         terminated = episode_end and infos['is_success']
         truncated = episode_end and not infos['is_success']
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return obs, reward, terminated, truncated, infos
 
-    def render(self):
+    def calculate_smoothness_metrics(self, velocity_profile, dt):
-        self.render_active = True
+        """
-        return super().render()
+        Calculates the smoothness metrics for the given velocity profile.
+        param velocity_profile: np.array
+                                The array containing the movement velocity profile.
+        param dt: float
+                  The sampling time interval of the data.
+        return mean_squared_jerk: float
+                                  The mean squared jerk estimate of the given movement's smoothness.
+        return maximum_jerk: float
+                             The maximum jerk estimate of the given movement's smoothness.
+        return dimensionless_jerk: float
+                                   The dimensionless jerk estimate of the given movement's smoothness.
+        """
+        jerk_profile = calculate_jerk_profile(velocity_profile, dt)
+        mean_squared_jerk = calculate_mean_squared_jerk(jerk_profile)
+        maximum_jerk = calculate_maximum_jerk(jerk_profile)
+        dimensionless_jerk = calculate_dimensionless_jerk(jerk_profile, velocity_profile, dt)
+        return mean_squared_jerk, maximum_jerk, dimensionless_jerk
 
     def reset_model(self):
         # rest box to initial position

fancy_gym/envs/mujoco/box_pushing/box_pushing_utils.py

@@ -51,3 +51,19 @@ def rot_to_quat(theta, axis):
     quant[0] = np.sin(theta / 2.)
     quant[1:] = np.cos(theta / 2.) * axis
     return quant
+
+def calculate_jerk_profile(velocity_profile, dt):
+    jerk = np.diff(velocity_profile, 2, 0) / pow(dt, 2)
+    return jerk
+
+def calculate_mean_squared_jerk(jerk_profile):
+    return np.mean(pow(jerk_profile, 2))
+
+def calculate_maximum_jerk(jerk_profile):
+    return np.max(abs(jerk_profile))
+
+def calculate_dimensionless_jerk(jerk_profile, velocity_profile, dt):
+    sum_squared_jerk = np.sum(pow(jerk_profile, 2), 0)
+    duration = len(velocity_profile) * dt
+    peak_velocity = np.max(abs(velocity_profile), 0)
+    return np.mean(sum_squared_jerk * pow(duration, 3) / pow(peak_velocity, 2))

fancy_gym/envs/mujoco/half_cheetah_jump/half_cheetah_jump.py

@@ -60,11 +60,7 @@ class HalfCheetahEnvCustomXML(HalfCheetahEnv):
             default_camera_config=DEFAULT_CAMERA_CONFIG,
             **kwargs,
         )
-        self.render_active = False
 
-    def render(self):
-        self.render_active = True
-        return super().render()
 
 class HalfCheetahJumpEnv(HalfCheetahEnvCustomXML):
     """
@@ -124,9 +120,6 @@ class HalfCheetahJumpEnv(HalfCheetahEnvCustomXML):
             'max_height': self.max_height
         }
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
     def _get_obs(self):

fancy_gym/envs/mujoco/hopper_jump/hopper_jump.py

@@ -88,12 +88,6 @@ class HopperEnvCustomXML(HopperEnv):
             **kwargs,
         )
 
-        self.render_active = False
-
-    def render(self):
-        self.render_active = True
-        return super().render()
-
 
 class HopperJumpEnv(HopperEnvCustomXML):
     """
@@ -207,10 +201,6 @@ class HopperJumpEnv(HopperEnvCustomXML):
             healthy=self.is_healthy,
             contact_dist=self.contact_dist or 0
         )
-
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
     def _get_obs(self):
@@ -272,76 +262,100 @@ class HopperJumpEnv(HopperEnvCustomXML):
                 return True
         return False
 
-# # TODO is that needed? if so test it
+class HopperJumpMarkovRew(HopperJumpEnv):
-# class HopperJumpStepEnv(HopperJumpEnv):
+    def step(self, action):
-#
+        self._steps += 1
-#     def __init__(self,
+
-#                  xml_file='hopper_jump.xml',
+        self.do_simulation(action, self.frame_skip)
-#                  forward_reward_weight=1.0,
+
-#                  ctrl_cost_weight=1e-3,
+        height_after = self.get_body_com("torso")[2]
-#                  healthy_reward=1.0,
+        # site_pos_after = self.data.get_site_xpos('foot_site')
-#                  height_weight=3,
+        site_pos_after = self.data.site('foot_site').xpos
-#                  dist_weight=3,
+        self.max_height = max(height_after, self.max_height)
-#                  terminate_when_unhealthy=False,
+
-#                  healthy_state_range=(-100.0, 100.0),
+        has_floor_contact = self._is_floor_foot_contact() if not self.contact_with_floor else False
-#                  healthy_z_range=(0.5, float('inf')),
+
-#                  healthy_angle_range=(-float('inf'), float('inf')),
+        if not self.init_floor_contact:
-#                  reset_noise_scale=5e-3,
+            self.init_floor_contact = has_floor_contact
-#                  exclude_current_positions_from_observation=False
+        if self.init_floor_contact and not self.has_left_floor:
-#                  ):
+            self.has_left_floor = not has_floor_contact
-#
+        if not self.contact_with_floor and self.has_left_floor:
-#         self._height_weight = height_weight
+            self.contact_with_floor = has_floor_contact
-#         self._dist_weight = dist_weight
+
-#         super().__init__(xml_file, forward_reward_weight, ctrl_cost_weight, healthy_reward, terminate_when_unhealthy,
+        ctrl_cost = self.control_cost(action)
-#                          healthy_state_range, healthy_z_range, healthy_angle_range, reset_noise_scale,
+        costs = ctrl_cost
-#                          exclude_current_positions_from_observation)
+        terminated = False
-#
+        truncated = False
-#     def step(self, action):
+
-#         self._steps += 1
+        goal_dist = np.linalg.norm(site_pos_after - self.goal)
-#
+        if self.contact_dist is None and self.contact_with_floor:
-#         self.do_simulation(action, self.frame_skip)
+            self.contact_dist = goal_dist
-#
+
-#         height_after = self.get_body_com("torso")[2]
+        rewards = 0
-#         site_pos_after = self.data.site('foot_site').xpos.copy()
+        if not self.sparse or (self.sparse and self._steps >= MAX_EPISODE_STEPS_HOPPERJUMP):
-#         self.max_height = max(height_after, self.max_height)
+            healthy_reward = self.healthy_reward
-#
+            distance_reward = -goal_dist * self._dist_weight
-#         ctrl_cost = self.control_cost(action)
+            height_reward = (self.max_height if self.sparse else height_after) * self._height_weight
-#         healthy_reward = self.healthy_reward
+            contact_reward = -(self.contact_dist or 5) * self._contact_weight
-#         height_reward = self._height_weight * height_after
+            rewards = self._forward_reward_weight * (distance_reward + height_reward + contact_reward + healthy_reward)
-#         goal_dist = np.linalg.norm(site_pos_after - np.array([self.goal, 0, 0]))
+
-#         goal_dist_reward = -self._dist_weight * goal_dist
+        observation = self._get_obs()
-#         dist_reward = self._forward_reward_weight * (goal_dist_reward + height_reward)
+
-#
+        # While loop to simulate the process after jump to make the task Markovian
-#         rewards = dist_reward + healthy_reward
+        if self.sparse and self.has_left_floor:
-#         costs = ctrl_cost
+            while self._steps < MAX_EPISODE_STEPS_HOPPERJUMP:
-#         done = False
+                # Simulate to the end of the episode
-#
+                self._steps += 1
-#         # This is only for logging the distance to goal when first having the contact
+
-#         has_floor_contact = self._is_floor_foot_contact() if not self.contact_with_floor else False
+                try:
-#
+                    self.do_simulation(np.zeros_like(action), self.frame_skip)
-#         if not self.init_floor_contact:
+                except Exception as e:
-#             self.init_floor_contact = has_floor_contact
+                    print(e)
-#         if self.init_floor_contact and not self.has_left_floor:
+
-#             self.has_left_floor = not has_floor_contact
+                height_after = self.get_body_com("torso")[2]
-#         if not self.contact_with_floor and self.has_left_floor:
+                #site_pos_after = self.data.get_site_xpos('foot_site')
-#             self.contact_with_floor = has_floor_contact
+                site_pos_after = self.data.site('foot_site').xpos
-#
+                self.max_height = max(height_after, self.max_height)
-#         if self.contact_dist is None and self.contact_with_floor:
+
-#             self.contact_dist = goal_dist
+                has_floor_contact = self._is_floor_foot_contact() if not self.contact_with_floor else False
-#
+
-#         ##############################################################
+                if not self.init_floor_contact:
-#
+                    self.init_floor_contact = has_floor_contact
-#         observation = self._get_obs()
+                if self.init_floor_contact and not self.has_left_floor:
-#         reward = rewards - costs
+                    self.has_left_floor = not has_floor_contact
-#         info = {
+                if not self.contact_with_floor and self.has_left_floor:
-#             'height': height_after,
+                    self.contact_with_floor = has_floor_contact
-#             'x_pos': site_pos_after,
+
-#             'max_height': copy.copy(self.max_height),
+                ctrl_cost = self.control_cost(action)
-#             'goal': copy.copy(self.goal),
+                costs = ctrl_cost
-#             'goal_dist': goal_dist,
+                done = False
-#             'height_rew': height_reward,
+
-#             'healthy_reward': healthy_reward,
+                goal_dist = np.linalg.norm(site_pos_after - self.goal)
-#             'healthy': copy.copy(self.is_healthy),
+                if self.contact_dist is None and self.contact_with_floor:
-#             'contact_dist': copy.copy(self.contact_dist) or 0
+                    self.contact_dist = goal_dist
-#         }
+
-#         return observation, reward, done, info
+                rewards = 0
+
+            # Task has reached the end, compute the sparse reward
+            done = True
+            healthy_reward = self.healthy_reward
+            distance_reward = -goal_dist * self._dist_weight
+            height_reward = (self.max_height if self.sparse else height_after) * self._height_weight
+            contact_reward = -(self.contact_dist or 5) * self._contact_weight
+            rewards = self._forward_reward_weight * (distance_reward + height_reward + contact_reward + healthy_reward)
+
+        reward = rewards - costs
+        info = dict(
+            height=height_after,
+            x_pos=site_pos_after,
+            max_height=self.max_height,
+            goal=self.goal[:1],
+            goal_dist=goal_dist,
+            height_rew=self.max_height,
+            healthy_reward=self.healthy_reward,
+            healthy=self.is_healthy,
+            contact_dist=self.contact_dist or 0,
+            num_steps=self._steps,
+            has_left_floor=self.has_left_floor
+       )
+        return observation, reward, terminated, truncated, info

fancy_gym/envs/mujoco/hopper_jump/hopper_jump_on_box.py

@@ -140,9 +140,6 @@ class HopperJumpOnBoxEnv(HopperEnvCustomXML):
 
         truncated = self.current_step >= self.max_episode_steps and not terminated
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
     def _get_obs(self):

fancy_gym/envs/mujoco/hopper_throw/hopper_throw.py

@@ -61,8 +61,6 @@ class HopperThrowEnv(HopperEnvCustomXML):
                          exclude_current_positions_from_observation=exclude_current_positions_from_observation,
                          **kwargs)
 
-        self.render_active = False
-
     def step(self, action):
         self.current_step += 1
         self.do_simulation(action, self.frame_skip)
@@ -96,15 +94,8 @@ class HopperThrowEnv(HopperEnvCustomXML):
         }
         truncated = False
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _get_obs(self):
         return np.append(super()._get_obs(), self.goal)
 

fancy_gym/envs/mujoco/hopper_throw/hopper_throw_in_basket.py

@@ -68,7 +68,6 @@ class HopperThrowInBasketEnv(HopperEnvCustomXML):
                          reset_noise_scale=reset_noise_scale,
                          exclude_current_positions_from_observation=exclude_current_positions_from_observation,
                          **kwargs)
-        self.render_active = False
 
     def step(self, action):
 
@@ -119,15 +118,8 @@ class HopperThrowInBasketEnv(HopperEnvCustomXML):
         }
         truncated = False
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _get_obs(self):
         return np.append(super()._get_obs(), self.basket_x)
 

fancy_gym/envs/mujoco/reacher/reacher.py

@@ -47,8 +47,6 @@ class ReacherEnv(MujocoEnv, utils.EzPickle):
                            **kwargs
                            )
 
-        self.render_active = False
-
     def step(self, action):
         self._steps += 1
 
@@ -79,15 +77,8 @@ class ReacherEnv(MujocoEnv, utils.EzPickle):
             goal=self.goal if hasattr(self, "goal") else None
         )
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return ob, reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def distance_reward(self):
         vec = self.get_body_com("fingertip") - self.get_body_com("target")
         return -self._reward_weight * np.linalg.norm(vec)

fancy_gym/envs/mujoco/table_tennis/table_tennis_env.py

@@ -5,11 +5,12 @@ from gymnasium import utils, spaces
 from gymnasium.envs.mujoco import MujocoEnv
 
 from fancy_gym.envs.mujoco.table_tennis.table_tennis_utils import is_init_state_valid, magnus_force
-from fancy_gym.envs.mujoco.table_tennis.table_tennis_utils import jnt_pos_low, jnt_pos_high
+from fancy_gym.envs.mujoco.table_tennis.table_tennis_utils import jnt_pos_low, jnt_pos_high, jnt_vel_low, jnt_vel_high
 
 import mujoco
 
 MAX_EPISODE_STEPS_TABLE_TENNIS = 350
+MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER = 300
 
 CONTEXT_BOUNDS_2DIMS = np.array([[-1.0, -0.65], [-0.2, 0.65]])
 CONTEXT_BOUNDS_4DIMS = np.array([[-1.0, -0.65, -1.0, -0.65],
@@ -18,6 +19,9 @@ CONTEXT_BOUNDS_SWICHING = np.array([[-1.0, -0.65, -1.0, 0.],
                                     [-0.2, 0.65, -0.2, 0.65]])
 
 
+DEFAULT_ROBOT_INIT_POS = np.array([0.0, 0.0, 0.0, 1.5, 0.0, 0.0, 1.5])
+DEFAULT_ROBOT_INIT_VEL = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0])
+
 class TableTennisEnv(MujocoEnv, utils.EzPickle):
     """
     7 DoF table tennis environment
@@ -34,7 +38,11 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
 
     def __init__(self, ctxt_dim: int = 4, frame_skip: int = 4,
                  goal_switching_step: int = None,
-                 enable_artificial_wind: bool = False, **kwargs):
+                 enable_artificial_wind: bool = False,
+                 random_pos_scale: float = 0.0,
+                 random_vel_scale: float = 0.0,
+                 **kwargs,
+                ):
         utils.EzPickle.__init__(**locals())
         self._steps = 0
 
@@ -48,6 +56,10 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
 
         self._id_set = False
 
+        # initial robot state
+        self._random_pos_scale = random_pos_scale
+        self._random_vel_scale = random_vel_scale
+
         # reward calculation
         self.ball_landing_pos = None
         self._goal_pos = np.zeros(2)
@@ -71,8 +83,6 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
                            observation_space=self.observation_space,
                            **kwargs)
 
-        self.render_active = False
-
         if ctxt_dim == 2:
             self.context_bounds = CONTEXT_BOUNDS_2DIMS
         elif ctxt_dim == 4:
@@ -158,17 +168,10 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
             "num_steps": self._steps,
         }
 
-        terminated, truncated = self._terminated, False
+        terminated, truncated = self._terminated, self._steps == MAX_EPISODE_STEPS_TABLE_TENNIS
-
-        if self.render_active and self.render_mode=='human':
-            self.render()
 
         return self._get_obs(), reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _contact_checker(self, id_1, id_2):
         for coni in range(0, self.data.ncon):
             con = self.data.contact[coni]
@@ -176,6 +179,17 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
                 return True
         return False
 
+    def get_initial_robot_state(self):
+
+        robot_init_pos = DEFAULT_ROBOT_INIT_POS + \
+                         self.np_random.uniform(-1.0, 1.0, size=7) *\
+                         np.array([5.2, 4.0, 5.6, 4.0, 6.1, 3.2, 4.4]) *\
+                         self._random_pos_scale
+
+        robot_init_vel = DEFAULT_ROBOT_INIT_VEL + self.np_random.uniform(-1.0, 1.0, size=7) * self._random_vel_scale
+
+        return np.clip(robot_init_pos, jnt_pos_low, jnt_pos_high), np.clip(robot_init_vel, jnt_vel_low, jnt_vel_high)
+
     def reset_model(self):
         self._steps = 0
         self._init_ball_state = self._generate_valid_init_ball(random_pos=True, random_vel=False)
@@ -192,8 +206,10 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
 
         self.model.body_pos[5] = np.concatenate([self._goal_pos, [0.77]])
 
-        self.data.qpos[:7] = np.array([0., 0., 0., 1.5, 0., 0., 1.5])
+        robot_init_pos, robot_init_vel = self.get_initial_robot_state()
-        self.data.qvel[:7] = np.zeros(7)
+
+        self.data.qpos[:7] = robot_init_pos
+        self.data.qvel[:7] = robot_init_vel
 
         mujoco.mj_forward(self.model, self.data)
 
@@ -266,7 +282,7 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
     def get_invalid_traj_step_return(self, action, pos_traj, contextual_obs, tau_bound, delay_bound):
         obs = self._get_obs() if contextual_obs else np.concatenate([self._get_obs(), np.array([0])])  # 0 for invalid traj
         penalty = self._get_traj_invalid_penalty(action, pos_traj, tau_bound, delay_bound)
-        return obs, penalty, True, False, {
+        return obs, penalty, False, True, {
             "hit_ball": [False],
             "ball_returned_success": [False],
             "land_dist_error": [10.],
@@ -283,6 +299,179 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
             return False, pos_traj, vel_traj
         return True, pos_traj, vel_traj
 
+class TableTennisMarkovian(TableTennisEnv):
+    def _get_reward2(self, hit_now, land_now):
+
+        # Phase 1 not hit ball
+        if not self._hit_ball:
+            # Not hit ball
+            min_r_b_dist = np.min(np.linalg.norm(np.array(self._ball_traj) - np.array(self._racket_traj), axis=1))
+            return 0.005 * (1 - np.tanh(min_r_b_dist**2))
+
+        # Phase 2 hit ball now
+        elif self._hit_ball and hit_now:
+            return 2
+
+        # Phase 3 hit ball already and not land yet
+        elif self._hit_ball and self._ball_landing_pos is None:
+            min_b_des_b_dist = np.min(np.linalg.norm(np.array(self._ball_traj)[:,:2] - self._goal_pos[:2], axis=1))
+            return 0.02 * (1 - np.tanh(min_b_des_b_dist**2))
+
+        # Phase 4 hit ball already and land now
+        elif self._hit_ball and land_now:
+            over_net_bonus = int(self._ball_landing_pos[0] < 0)
+            min_b_des_b_land_dist = np.linalg.norm(self._goal_pos[:2] - self._ball_landing_pos[:2])
+            return 4 * (1 - np.tanh(min_b_des_b_land_dist ** 2)) + over_net_bonus
+
+        # Phase 5 hit ball already and land already
+        elif self._hit_ball and not land_now and self._ball_landing_pos is not None:
+            return 0
+
+        else:
+            raise NotImplementedError
+
+    def _get_reward(self, terminated):
+        # if not terminated:
+        #     return 0
+
+        min_r_b_dist = np.min(np.linalg.norm(np.array(self._ball_traj) - np.array(self._racket_traj), axis=1))
+        if not self._hit_ball:
+            # Not hit ball
+            return 0.2 * (1 - np.tanh(min_r_b_dist**2))
+        elif self._ball_landing_pos is None:
+            # Hit ball but not landing pos
+            min_b_des_b_dist = np.min(np.linalg.norm(np.array(self._ball_traj)[:,:2] - self._goal_pos[:2], axis=1))
+            return 2 + (1 - np.tanh(min_b_des_b_dist**2))
+        else:
+            # Hit ball and land
+            min_b_des_b_land_dist = np.linalg.norm(self._goal_pos[:2] - self._ball_landing_pos[:2])
+            over_net_bonus = int(self._ball_landing_pos[0] < 0)
+            return 2 + 4 * (1 - np.tanh(min_b_des_b_land_dist ** 2)) + over_net_bonus
+
+
+    def _get_traj_invalid_penalty(self, action, pos_traj, tau_bound, delay_bound):
+        tau_invalid_penalty = 3 * (np.max([0, action[0] - tau_bound[1]]) + np.max([0, tau_bound[0] - action[0]]))
+        delay_invalid_penalty = 3 * (np.max([0, action[1] - delay_bound[1]]) + np.max([0, delay_bound[0] - action[1]]))
+        violate_high_bound_error = np.mean(np.maximum(pos_traj - jnt_pos_high, 0))
+        violate_low_bound_error = np.mean(np.maximum(jnt_pos_low - pos_traj, 0))
+        invalid_penalty = tau_invalid_penalty + delay_invalid_penalty + \
+                          violate_high_bound_error + violate_low_bound_error
+        return -invalid_penalty
+
+    def get_invalid_traj_step_penalty(self, pos_traj):
+        violate_high_bound_error = (
+            np.maximum(pos_traj - jnt_pos_high, 0).mean())
+        violate_low_bound_error = (
+            np.maximum(jnt_pos_low - pos_traj, 0).mean())
+        invalid_penalty = violate_high_bound_error + violate_low_bound_error
+
+
+    def _update_game_state(self, action):
+        for _ in range(self.frame_skip):
+            if self._enable_artificial_wind:
+                self.data.qfrc_applied[-2] = self._artificial_force
+            try:
+                self.do_simulation(action, 1)
+            except Exception as e:
+                print("Simulation get unstable return with MujocoException: ", e)
+                unstable_simulation = True
+                self._terminated = True
+                break
+
+            # Update game state
+            if not self._terminated:
+                if not self._hit_ball:
+                    self._hit_ball = self._contact_checker(self._ball_contact_id, self._bat_front_id) or \
+                                    self._contact_checker(self._ball_contact_id, self._bat_back_id)
+                    if not self._hit_ball:
+                        ball_land_on_floor_no_hit = self._contact_checker(self._ball_contact_id, self._floor_contact_id)
+                        if ball_land_on_floor_no_hit:
+                            self._ball_landing_pos = self.data.body("target_ball").xpos.copy()
+                            self._terminated = True
+                if self._hit_ball and not self._ball_contact_after_hit:
+                    if self._contact_checker(self._ball_contact_id, self._floor_contact_id):  # first check contact with floor
+                        self._ball_contact_after_hit = True
+                        self._ball_landing_pos = self.data.geom("target_ball_contact").xpos.copy()
+                        self._terminated = True
+                    elif self._contact_checker(self._ball_contact_id, self._table_contact_id):  # second check contact with table
+                        self._ball_contact_after_hit = True
+                        self._ball_landing_pos = self.data.geom("target_ball_contact").xpos.copy()
+                        if self._ball_landing_pos[0] < 0.:  # ball lands on the opponent side
+                            self._ball_return_success = True
+                        self._terminated = True
+
+            # update ball trajectory & racket trajectory
+            self._ball_traj.append(self.data.body("target_ball").xpos.copy())
+            self._racket_traj.append(self.data.geom("bat").xpos.copy())
+
+    def ball_racket_contact(self):
+        return self._contact_checker(self._ball_contact_id, self._bat_front_id) or \
+               self._contact_checker(self._ball_contact_id, self._bat_back_id)
+
+    def step(self, action):
+        if not self._id_set:
+            self._set_ids()
+
+        unstable_simulation = False
+        hit_already = self._hit_ball
+        if self._steps == self._goal_switching_step and self.np_random.uniform() < 0.5:
+                new_goal_pos = self._generate_goal_pos(random=True)
+                new_goal_pos[1] = -new_goal_pos[1]
+                self._goal_pos = new_goal_pos
+                self.model.body_pos[5] = np.concatenate([self._goal_pos, [0.77]])
+                mujoco.mj_forward(self.model, self.data)
+
+        self._update_game_state(action)
+        self._steps += 1
+
+        obs = self._get_obs()
+
+        # Compute reward
+        if unstable_simulation:
+            reward = -25
+        else:
+            # reward = self._get_reward(self._terminated)
+            # hit_now = not hit_already and self._hit_ball
+            hit_finish = self._hit_ball and not self.ball_racket_contact()
+
+            if hit_finish:
+                # Clean the ball and racket traj before hit
+                self._ball_traj = []
+                self._racket_traj = []
+
+                # Simulate the rest of the traj
+                reward = self._get_reward2(True, False)
+                while self._steps < MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER:
+                    land_already = self._ball_landing_pos is not None
+                    self._update_game_state(np.zeros_like(action))
+                    self._steps += 1
+
+                    land_now = (not land_already
+                                and self._ball_landing_pos is not None)
+                    temp_reward = self._get_reward2(False, land_now)
+                    # print(temp_reward)
+                    reward += temp_reward
+
+                    # Uncomment the line below to visualize the sim after hit
+                    # self.render(mode="human")
+            else:
+                reward = self._get_reward2(False, False)
+
+        # Update ball landing error
+        land_dist_err = np.linalg.norm(self._ball_landing_pos[:-1] - self._goal_pos) \
+                            if self._ball_landing_pos is not None else 10.
+
+        info = {
+            "hit_ball": self._hit_ball,
+            "ball_returned_success": self._ball_return_success,
+            "land_dist_error": land_dist_err,
+            "is_success": self._ball_return_success and land_dist_err < 0.2,
+            "num_steps": self._steps,
+        }
+
+        terminated, truncated = self._terminated, self._steps == MAX_EPISODE_STEPS_TABLE_TENNIS_MARKOV_VER
+
+        return obs, reward, terminated, truncated, info
 
 class TableTennisWind(TableTennisEnv):
     def __init__(self, ctxt_dim: int = 4, frame_skip: int = 4, **kwargs):
@@ -305,7 +494,17 @@ class TableTennisWind(TableTennisEnv):
         ])
         return obs
 
-
 class TableTennisGoalSwitching(TableTennisEnv):
     def __init__(self, frame_skip: int = 4, goal_switching_step: int = 99, **kwargs):
         super().__init__(frame_skip=frame_skip, goal_switching_step=goal_switching_step, **kwargs)
+
+
+class TableTennisRandomInit(TableTennisEnv):
+    def __init__(self, ctxt_dim: int = 4, frame_skip: int = 4,
+                 random_pos_scale: float = 1.0,
+                 random_vel_scale: float = 0.0,
+                 **kwargs):
+        super().__init__(ctxt_dim=ctxt_dim, frame_skip=frame_skip,
+                         random_pos_scale=random_pos_scale,
+                         random_vel_scale=random_vel_scale,
+                         **kwargs)

fancy_gym/envs/mujoco/table_tennis/table_tennis_utils.py

@@ -2,6 +2,10 @@ import numpy as np
 
 jnt_pos_low = np.array([-2.6, -2.0, -2.8, -0.9, -4.8, -1.6, -2.2])
 jnt_pos_high = np.array([2.6, 2.0, 2.8, 3.1, 1.3, 1.6, 2.2])
+
+jnt_vel_low = np.ones(7) * -7
+jnt_vel_high = np.ones(7) * 7
+
 delay_bound = [0.05, 0.15]
 tau_bound = [0.5, 1.5]
 
@@ -48,4 +52,4 @@ def magnus_force(top_spin=0.0, side_spin=0.0, v_ball=np.zeros(3), v_wind=np.zero
     C_l = 4.68 * 10e-4 - 2.0984 * 10e-5 * (np.linalg.norm(v_ball) - 50)  # Lift force coeffient or simply 1.23
     w = np.array([0.0, top_spin, side_spin]) # Angular velocity of ball
     f_m = 0.5 * rho * A * C_l * np.linalg.norm(v_ball-v_wind) * np.cross(w, v_ball-v_wind)
-    return f_m
+    return f_m

fancy_gym/envs/mujoco/walker_2d_jump/walker_2d_jump.py

@@ -79,8 +79,6 @@ class Walker2dEnvCustomXML(Walker2dEnv):
             **kwargs,
         )
 
-        self.render_active = False
-
 
 class Walker2dJumpEnv(Walker2dEnvCustomXML):
     """
@@ -147,15 +145,8 @@ class Walker2dJumpEnv(Walker2dEnvCustomXML):
         }
         truncated = False
 
-        if self.render_active and self.render_mode=='human':
-            self.render()
-
         return observation, reward, terminated, truncated, info
 
-    def render(self):
-        self.render_active = True
-        return super().render()
-
     def _get_obs(self):
         return np.append(super()._get_obs(), self.goal)
 

fancy_gym/examples/example_replanning_envs.py

@@ -3,14 +3,14 @@ import fancy_gym
 
 
 def example_run_replanning_env(env_name="fancy_ProDMP/BoxPushingDenseReplan-v0", seed=1, iterations=1, render=False):
-    env = gym.make(env_name, render_mode='human' if render else None)
+    env = gym.make(env_name)
     env.reset(seed=seed)
     for i in range(iterations):
         while True:
             ac = env.action_space.sample()
             obs, reward, terminated, truncated, info = env.step(ac)
             if render:
-                env.render()
+                env.render(mode="human")
             if terminated or truncated:
                 env.reset()
                 break
@@ -38,13 +38,13 @@ def example_custom_replanning_envs(seed=0, iteration=100, render=True):
                         'replanning_schedule': lambda pos, vel, obs, action, t: t % 25 == 0,
                         'condition_on_desired': True}
 
-    base_env = gym.make(base_env_id, render_mode='human' if render else None)
+    base_env = gym.make(base_env_id)
     env = fancy_gym.make_bb(env=base_env, wrappers=wrappers, black_box_kwargs=black_box_kwargs,
                             traj_gen_kwargs=trajectory_generator_kwargs, controller_kwargs=controller_kwargs,
                             phase_kwargs=phase_generator_kwargs, basis_kwargs=basis_generator_kwargs,
                             seed=seed)
     if render:
-        env.render()
+        env.render(mode="human")
 
     obs = env.reset()
 

fancy_gym/examples/examples_dmc.py

@@ -17,7 +17,7 @@ def example_dmc(env_id="dm_control/fish-swim", seed=1, iterations=1000, render=T
     Returns:
 
     """
-    env = gym.make(env_id, render_mode='human' if render else None)
+    env = gym.make(env_id)
     rewards = 0
     obs = env.reset(seed=seed)
     print("observation shape:", env.observation_space.shape)
@@ -26,7 +26,7 @@ def example_dmc(env_id="dm_control/fish-swim", seed=1, iterations=1000, render=T
     for i in range(iterations):
         ac = env.action_space.sample()
         if render:
-            env.render()
+            env.render(mode="human")
         obs, reward, terminated, truncated, info = env.step(ac)
         rewards += reward
 
@@ -84,7 +84,7 @@ def example_custom_dmc_and_mp(seed=1, iterations=1, render=True):
     # basis_generator_kwargs = {'basis_generator_type': 'rbf',
     #                           'num_basis': 5
     #                           }
-    base_env = gym.make(base_env_id, render_mode='human' if render else None)
+    base_env = gym.make(base_env_id)
     env = fancy_gym.make_bb(env=base_env, wrappers=wrappers, black_box_kwargs={},
                             traj_gen_kwargs=trajectory_generator_kwargs, controller_kwargs=controller_kwargs,
                             phase_kwargs=phase_generator_kwargs, basis_kwargs=basis_generator_kwargs,
@@ -96,7 +96,7 @@ def example_custom_dmc_and_mp(seed=1, iterations=1, render=True):
     # It is also possible to change them mode multiple times when
     # e.g. only every nth trajectory should be displayed.
     if render:
-        env.render()
+        env.render(mode="human")
 
     rewards = 0
     obs = env.reset()
@@ -115,7 +115,7 @@ def example_custom_dmc_and_mp(seed=1, iterations=1, render=True):
     env.close()
     del env
 
-def main(render = False):
+def main(render = True):
     # # Standard DMC Suite tasks
     example_dmc("dm_control/fish-swim", seed=10, iterations=1000, render=render)
     #

fancy_gym/examples/examples_general.py

@@ -21,7 +21,7 @@ def example_general(env_id="Pendulum-v1", seed=1, iterations=1000, render=True):
 
     """
 
-    env = gym.make(env_id, render_mode='human' if render else None)
+    env = gym.make(env_id)
     rewards = 0
     obs = env.reset(seed=seed)
     print("Observation shape: ", env.observation_space.shape)
@@ -85,7 +85,7 @@ def example_async(env_id="fancy/HoleReacher-v0", n_cpu=4, seed=int('533D', 16),
     # do not return values above threshold
     return *map(lambda v: np.stack(v)[:n_samples], buffer.values()),
 
-def main(render = False):
+def main(render = True):
     # Basic gym task
     example_general("Pendulum-v1", seed=10, iterations=200, render=render)
 

fancy_gym/examples/examples_metaworld.py

@@ -2,7 +2,7 @@ import gymnasium as gym
 import fancy_gym
 
 
-def example_meta(env_id="metaworld/button-press-v2", seed=1, iterations=1000, render=True):
+def example_meta(env_id="fish-swim", seed=1, iterations=1000, render=True):
     """
     Example for running a MetaWorld based env in the step based setting.
     The env_id has to be specified as `task_name-v2`. V1 versions are not supported and we always
@@ -18,7 +18,7 @@ def example_meta(env_id="metaworld/button-press-v2", seed=1, iterations=1000, re
     Returns:
 
     """
-    env = gym.make(env_id, render_mode='human' if render else None)
+    env = gym.make(env_id)
     rewards = 0
     obs = env.reset(seed=seed)
     print("observation shape:", env.observation_space.shape)
@@ -27,7 +27,9 @@ def example_meta(env_id="metaworld/button-press-v2", seed=1, iterations=1000, re
     for i in range(iterations):
         ac = env.action_space.sample()
         if render:
-            env.render()
+            # THIS NEEDS TO BE SET TO FALSE FOR NOW, BECAUSE THE INTERFACE FOR RENDERING IS DIFFERENT TO BASIC GYM
+            # TODO: Remove this, when Metaworld fixes its interface.
+            env.render(False)
         obs, reward, terminated, truncated, info = env.step(ac)
         rewards += reward
         if terminated or truncated:
@@ -79,7 +81,7 @@ def example_custom_meta_and_mp(seed=1, iterations=1, render=True):
     basis_generator_kwargs = {'basis_generator_type': 'rbf',
                               'num_basis': 5
                               }
-    base_env = gym.make(base_env_id, render_mode='human' if render else None)
+    base_env = gym.make(base_env_id)
     env = fancy_gym.make_bb(env=base_env, wrappers=wrappers, black_box_kwargs={},
                             traj_gen_kwargs=trajectory_generator_kwargs, controller_kwargs=controller_kwargs,
                             phase_kwargs=phase_generator_kwargs, basis_kwargs=basis_generator_kwargs,
@@ -91,7 +93,7 @@ def example_custom_meta_and_mp(seed=1, iterations=1, render=True):
     # It is also possible to change them mode multiple times when
     # e.g. only every nth trajectory should be displayed.
     if render:
-        env.render()
+        env.render(mode="human")
 
     rewards = 0
     obs = env.reset(seed=seed)

fancy_gym/examples/examples_open_ai.py

@@ -13,13 +13,15 @@ def example_mp(env_name, seed=1, render=True):
     Returns:
 
     """
-    env = gym.make(env_name, render_mode='human' if render else None)
+    env = gym.make(env_name)
 
     returns = 0
     obs = env.reset(seed=seed)
     # number of samples/full trajectories (multiple environment steps)
     for i in range(10):
         if render and i % 2 == 0:
+            env.render(mode="human")
+        else:
             env.render()
         ac = env.action_space.sample()
         obs, reward, terminated, truncated, info = env.step(ac)

test/test_gym_builtin_envs.py

@@ -20,7 +20,6 @@ GYM_IDS = [spec.id for spec in gym.envs.registry.values() if
 GYM_MP_IDS = fancy_gym.ALL_DMC_MOVEMENT_PRIMITIVE_ENVIRONMENTS['all']
 SEED = 1
 
-
 @pytest.mark.parametrize('env_id', GYM_IDS)
 def test_step_gym_functionality(env_id: str):
     """Tests that step environments run without errors using random actions."""