Merge branch 'bruce_port_envs' into future

This commit is contained in:
Dominik Moritz Roth 2024-02-02 16:51:50 +01:00
commit e89fe9e07f
3 changed files with 22 additions and 5 deletions

View File

@ -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

View File

@ -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
}

View File

@ -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)