Merge branch 'bruce_port_envs' into future

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/box_pushing/box_pushing_env.py

@@ -97,11 +97,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
         }

fancy_gym/envs/mujoco/table_tennis/table_tennis_env.py

@@ -41,6 +41,7 @@ class TableTennisEnv(MujocoEnv, utils.EzPickle):
                  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
@@ -178,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)
@@ -490,8 +502,9 @@ class TableTennisGoalSwitching(TableTennisEnv):
 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):
+                 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)
+                         random_vel_scale=random_vel_scale,
-
+                         **kwargs)