diff --git a/alr_envs/alr/__init__.py b/alr_envs/alr/__init__.py index 67031f1..0181f62 100644 --- a/alr_envs/alr/__init__.py +++ b/alr_envs/alr/__init__.py @@ -375,7 +375,7 @@ for _v in _versions: ALL_ALR_MOTION_PRIMITIVE_ENVIRONMENTS["ProMP"].append(_env_id) ## Beerpong -_versions = ["v0", "v1", "v2", "v3"] +_versions = ["v0", "v1"] for _v in _versions: _env_id = f'BeerpongProMP-{_v}' register( diff --git a/alr_envs/alr/mujoco/beerpong/beerpong.py b/alr_envs/alr/mujoco/beerpong/beerpong.py index 2fb435f..9092ef1 100644 --- a/alr_envs/alr/mujoco/beerpong/beerpong.py +++ b/alr_envs/alr/mujoco/beerpong/beerpong.py @@ -13,9 +13,15 @@ CUP_POS_MAX = np.array([0.32, -1.2]) class ALRBeerBongEnv(MujocoEnv, utils.EzPickle): def __init__(self, frame_skip=1, apply_gravity_comp=True, noisy=False, - rndm_goal=False, cup_goal_pos=[-0.3, -1.2]): - self._steps = 0 + rndm_goal=False, cup_goal_pos=None): + if cup_goal_pos is None: + cup_goal_pos = [-0.3, -1.2, 0.840] + elif len(cup_goal_pos)==2: + cup_goal_pos = np.array(cup_goal_pos) + cup_goal_pos = np.insert(cup_goal_pos, 2, 0.80) + self.cup_goal_pos = np.array(cup_goal_pos) + self._steps = 0 self.xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "beerpong_wo_cup" + ".xml") @@ -43,7 +49,7 @@ class ALRBeerBongEnv(MujocoEnv, utils.EzPickle): else: self.noise_std = 0 - self.cup_goal_pos = np.array(cup_goal_pos.append(0.840)) + reward_function = BeerPongReward self.reward_function = reward_function()