diff --git a/alr_envs/__init__.py b/alr_envs/__init__.py index da4204e..3d221b4 100644 --- a/alr_envs/__init__.py +++ b/alr_envs/__init__.py @@ -88,6 +88,16 @@ register( id='ALRBallInACupSimple-v0', entry_point='alr_envs.mujoco:ALRBallInACupEnv', max_episode_steps=4000, + kwargs={ + "simplified": True, + "reward_type": "no_context", + } +) + +register( + id='ALRBallInACupPDSimple-v0', + entry_point='alr_envs.mujoco:ALRBallInACupPDEnv', + max_episode_steps=4000, kwargs={ "simplified": True, "reward_type": "no_context" @@ -368,6 +378,26 @@ register( } ) +register( + id='ALRBallInACupPDSimpleDetPMP-v0', + entry_point='alr_envs.mujoco.ball_in_a_cup.biac_pd:make_detpmp_env', + kwargs={ + "name": "alr_envs:ALRBallInACupPDSimple-v0", + "num_dof": 3, + "num_basis": 5, + "duration": 3.5, + "post_traj_time": 4.5, + "width": 0.0035, + # "off": -0.05, + "policy_type": "motor", + "weights_scale": 0.2, + "zero_start": True, + "zero_goal": True, + "p_gains": np.array([4./3., 2.4, 2.5, 5./3., 2., 2., 1.25]), + "d_gains": np.array([0.0466, 0.12, 0.125, 0.04166, 0.06, 0.06, 0.025]) + } +) + register( id='ALRBallInACupDetPMP-v0', entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env', diff --git a/alr_envs/classic_control/hole_reacher.py b/alr_envs/classic_control/hole_reacher.py index 281a771..d5be111 100644 --- a/alr_envs/classic_control/hole_reacher.py +++ b/alr_envs/classic_control/hole_reacher.py @@ -7,10 +7,10 @@ from gym.utils import seeding from matplotlib import patches from alr_envs.classic_control.utils import check_self_collision -from alr_envs.utils.mps.alr_env import AlrEnv +from mp_env_api.envs.mp_env import MpEnv -class HoleReacherEnv(AlrEnv): +class HoleReacherEnv(MpEnv): def __init__(self, n_links: int, hole_x: Union[None, float] = None, hole_depth: Union[None, float] = None, hole_width: float = 1., random_start: bool = False, allow_self_collision: bool = False, @@ -44,7 +44,7 @@ class HoleReacherEnv(AlrEnv): self._start_pos = np.hstack([[np.pi / 2], np.zeros(self.n_links - 1)]) self._start_vel = np.zeros(self.n_links) - self.dt = 0.01 + self._dt = 0.01 action_bound = np.pi * np.ones((self.n_links,)) state_bound = np.hstack([ @@ -66,6 +66,10 @@ class HoleReacherEnv(AlrEnv): self._steps = 0 self.seed() + @property + def dt(self) -> Union[float, int]: + return self._dt + def step(self, action: np.ndarray): """ A single step with an action in joint velocity space diff --git a/alr_envs/classic_control/simple_reacher.py b/alr_envs/classic_control/simple_reacher.py index 43b0381..f564f89 100644 --- a/alr_envs/classic_control/simple_reacher.py +++ b/alr_envs/classic_control/simple_reacher.py @@ -5,10 +5,10 @@ import numpy as np from gym import spaces from gym.utils import seeding -from alr_envs.utils.mps.alr_env import AlrEnv +from mp_env_api.envs.mp_env import MpEnv -class SimpleReacherEnv(AlrEnv): +class SimpleReacherEnv(MpEnv): """ Simple Reaching Task without any physics simulation. Returns no reward until 150 time steps. This allows the agent to explore the space, but requires precise actions diff --git a/alr_envs/classic_control/viapoint_reacher.py b/alr_envs/classic_control/viapoint_reacher.py index afaa26e..15934a8 100644 --- a/alr_envs/classic_control/viapoint_reacher.py +++ b/alr_envs/classic_control/viapoint_reacher.py @@ -6,10 +6,10 @@ import numpy as np from gym.utils import seeding from alr_envs.classic_control.utils import check_self_collision -from alr_envs.utils.mps.alr_env import AlrEnv +from mp_env_api.envs.mp_env import MpEnv -class ViaPointReacher(AlrEnv): +class ViaPointReacher(MpEnv): def __init__(self, n_links, random_start: bool = True, via_target: Union[None, Iterable] = None, target: Union[None, Iterable] = None, allow_self_collision=False, collision_penalty=1000): diff --git a/alr_envs/mujoco/__init__.py b/alr_envs/mujoco/__init__.py index eb7b41f..82f3784 100644 --- a/alr_envs/mujoco/__init__.py +++ b/alr_envs/mujoco/__init__.py @@ -1,3 +1,4 @@ from alr_envs.mujoco.reacher.alr_reacher import ALRReacherEnv from alr_envs.mujoco.balancing import BalancingEnv -from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv +# from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup import ALRBallInACupEnv +from alr_envs.mujoco.ball_in_a_cup.biac_pd import ALRBallInACupPDEnv diff --git a/alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py b/alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py index daee289..295dcf9 100644 --- a/alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py +++ b/alr_envs/mujoco/ball_in_a_cup/ball_in_a_cup_reward_simple.py @@ -3,9 +3,8 @@ from alr_envs.mujoco import alr_reward_fct class BallInACupReward(alr_reward_fct.AlrReward): - def __init__(self, sim_time): - self.sim_time = sim_time - + def __init__(self, env): + self.env = env self.collision_objects = ["cup_geom1", "cup_geom2", "cup_base_contact_below", "wrist_palm_link_convex_geom", "wrist_pitch_link_convex_decomposition_p1_geom", @@ -32,7 +31,8 @@ class BallInACupReward(alr_reward_fct.AlrReward): self.reset(None) def reset(self, context): - self.ball_traj = np.zeros(shape=(self.sim_time, 3)) + # self.sim_time = self.env.sim.dtsim_time + self.ball_traj = [] # np.zeros(shape=(self.sim_time, 3)) self.dists = [] self.dists_final = [] self.costs = [] @@ -40,23 +40,24 @@ class BallInACupReward(alr_reward_fct.AlrReward): self.angle_costs = [] self.cup_angles = [] - def compute_reward(self, action, env): - self.ball_id = env.sim.model._body_name2id["ball"] - self.ball_collision_id = env.sim.model._geom_name2id["ball_geom"] - self.goal_id = env.sim.model._site_name2id["cup_goal"] - self.goal_final_id = env.sim.model._site_name2id["cup_goal_final"] - self.collision_ids = [env.sim.model._geom_name2id[name] for name in self.collision_objects] + def compute_reward(self, action): + self.ball_id = self.env.sim.model._body_name2id["ball"] + self.ball_collision_id = self.env.sim.model._geom_name2id["ball_geom"] + self.goal_id = self.env.sim.model._site_name2id["cup_goal"] + self.goal_final_id = self.env.sim.model._site_name2id["cup_goal_final"] + self.collision_ids = [self.env.sim.model._geom_name2id[name] for name in self.collision_objects] - ball_in_cup = self.check_ball_in_cup(env.sim, self.ball_collision_id) + ball_in_cup = self.check_ball_in_cup(self.env.sim, self.ball_collision_id) # Compute the current distance from the ball to the inner part of the cup - goal_pos = env.sim.data.site_xpos[self.goal_id] - ball_pos = env.sim.data.body_xpos[self.ball_id] - goal_final_pos = env.sim.data.site_xpos[self.goal_final_id] + goal_pos = self.env.sim.data.site_xpos[self.goal_id] + ball_pos = self.env.sim.data.body_xpos[self.ball_id] + goal_final_pos = self.env.sim.data.site_xpos[self.goal_final_id] self.dists.append(np.linalg.norm(goal_pos - ball_pos)) self.dists_final.append(np.linalg.norm(goal_final_pos - ball_pos)) - self.ball_traj[env._steps, :] = ball_pos - cup_quat = np.copy(env.sim.data.body_xquat[env.sim.model._body_name2id["cup"]]) + # self.ball_traj[self.env._steps, :] = ball_pos + self.ball_traj.append(ball_pos) + cup_quat = np.copy(self.env.sim.data.body_xquat[self.env.sim.model._body_name2id["cup"]]) cup_angle = np.arctan2(2 * (cup_quat[0] * cup_quat[1] + cup_quat[2] * cup_quat[3]), 1 - 2 * (cup_quat[1]**2 + cup_quat[2]**2)) cost_angle = (cup_angle - np.pi / 2) ** 2 @@ -66,9 +67,9 @@ class BallInACupReward(alr_reward_fct.AlrReward): action_cost = np.sum(np.square(action)) self.action_costs.append(action_cost) - self._is_collided = self.check_collision(env.sim) or env.check_traj_in_joint_limits() + self._is_collided = self.check_collision(self.env.sim) or self.env.check_traj_in_joint_limits() - if env._steps == env.sim_steps - 1 or self._is_collided: + if self.env._steps == self.env.ep_length - 1 or self._is_collided: t_min_dist = np.argmin(self.dists) angle_min_dist = self.cup_angles[t_min_dist] # cost_angle = (angle_min_dist - np.pi / 2)**2 diff --git a/alr_envs/mujoco/ball_in_a_cup/biac_pd.py b/alr_envs/mujoco/ball_in_a_cup/biac_pd.py new file mode 100644 index 0000000..1990b33 --- /dev/null +++ b/alr_envs/mujoco/ball_in_a_cup/biac_pd.py @@ -0,0 +1,225 @@ +import gym.envs.mujoco +import mujoco_py.builder +from gym import utils, spaces +import os +import numpy as np +import gym.envs.mujoco as mujoco_env +from mp_env_api.envs.mp_env import MpEnv +from mp_env_api.envs.positional_env import PositionalEnv +from mp_env_api.mp_wrappers.detpmp_wrapper import DetPMPWrapper +from mp_env_api.utils.policies import PDControllerExtend + + +def make_detpmp_env(**kwargs): + name = kwargs.pop("name") + _env = gym.make(name) + policy = PDControllerExtend(_env, p_gains=kwargs.pop('p_gains'), d_gains=kwargs.pop('d_gains')) + kwargs['policy_type'] = policy + return DetPMPWrapper(_env, **kwargs) + + +class ALRBallInACupPDEnv(mujoco_env.MujocoEnv, PositionalEnv, MpEnv, utils.EzPickle): + def __init__(self, frame_skip=4, apply_gravity_comp=True, simplified: bool = False, + reward_type: str = None, context: np.ndarray = None): + self._steps = 0 + + self.xml_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", + "biac_base" + ".xml") + + self.max_ctrl = np.array([150., 125., 40., 60., 5., 5., 2.]) + + self.j_min = np.array([-2.6, -1.985, -2.8, -0.9, -4.55, -1.5707, -2.7]) + self.j_max = np.array([2.6, 1.985, 2.8, 3.14159, 1.25, 1.5707, 2.7]) + + self.context = context + self.apply_gravity_comp = apply_gravity_comp + self.simplified = simplified + + self._start_pos = np.array([0.0, 0.58760536, 0.0, 1.36004913, 0.0, -0.32072943, -1.57]) + self._start_vel = np.zeros(7) + + self.sim_time = 8 # seconds + self.ep_length = 4000 # based on 8 seconds with dt = 0.02 int(self.sim_time / self.dt) + if reward_type == "no_context": + from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward_simple import BallInACupReward + reward_function = BallInACupReward + elif reward_type == "contextual_goal": + from alr_envs.mujoco.ball_in_a_cup.ball_in_a_cup_reward import BallInACupReward + reward_function = BallInACupReward + else: + raise ValueError("Unknown reward type: {}".format(reward_type)) + self.reward_function = reward_function(self) + + mujoco_env.MujocoEnv.__init__(self, self.xml_path, frame_skip) + utils.EzPickle.__init__(self) + + @property + def start_pos(self): + if self.simplified: + return self._start_pos[1::2] + else: + return self._start_pos + + @property + def start_vel(self): + if self.simplified: + return self._start_vel[1::2] + else: + return self._start_vel + + @property + def current_pos(self): + return self.sim.data.qpos[0:7].copy() + + @property + def current_vel(self): + return self.sim.data.qvel[0:7].copy() + + # def _set_action_space(self): + # if self.simplified: + # bounds = self.model.actuator_ctrlrange.copy().astype(np.float32)[1::2] + # else: + # bounds = self.model.actuator_ctrlrange.copy().astype(np.float32) + # low, high = bounds.T + # self.action_space = spaces.Box(low=low, high=high, dtype=np.float32) + # return self.action_space + + def reset(self): + self.reward_function.reset(None) + return super().reset() + + def reset_model(self): + init_pos_all = self.init_qpos.copy() + init_pos_robot = self._start_pos + init_vel = np.zeros_like(init_pos_all) + + self._steps = 0 + self._q_pos = [] + self._q_vel = [] + + start_pos = init_pos_all + start_pos[0:7] = init_pos_robot + + self.set_state(start_pos, init_vel) + + return self._get_obs() + + def step(self, a): + reward_dist = 0.0 + angular_vel = 0.0 + reward_ctrl = - np.square(a).sum() + + # if self.simplified: + # tmp = np.zeros(7) + # tmp[1::2] = a + # a = tmp + + if self.apply_gravity_comp: + a = a + self.sim.data.qfrc_bias[:len(a)].copy() / self.model.actuator_gear[:, 0] + try: + self.do_simulation(a, self.frame_skip) + crash = False + except mujoco_py.builder.MujocoException: + crash = True + # joint_cons_viol = self.check_traj_in_joint_limits() + + ob = self._get_obs() + + if not crash: + reward, success, is_collided = self.reward_function.compute_reward(a) + done = success or is_collided # self._steps == self.sim_steps - 1 + self._steps += 1 + else: + reward = -2000 + success = False + is_collided = False + done = True + + return ob, reward, done, dict(reward_dist=reward_dist, + reward_ctrl=reward_ctrl, + velocity=angular_vel, + # traj=self._q_pos, + action=a, + q_pos=self.sim.data.qpos[0:7].ravel().copy(), + q_vel=self.sim.data.qvel[0:7].ravel().copy(), + is_success=success, + is_collided=is_collided, sim_crash=crash) + + def check_traj_in_joint_limits(self): + return any(self.current_pos > self.j_max) or any(self.current_pos < self.j_min) + + # TODO: extend observation space + def _get_obs(self): + theta = self.sim.data.qpos.flat[:7] + return np.concatenate([ + np.cos(theta), + np.sin(theta), + # self.get_body_com("target"), # only return target to make problem harder + [self._steps], + ]) + + # TODO + @property + def active_obs(self): + return np.hstack([ + [False] * 7, # cos + [False] * 7, # sin + # [True] * 2, # x-y coordinates of target distance + [False] # env steps + ]) + + # These functions are for the task with 3 joint actuations + def extend_des_pos(self, des_pos): + des_pos_full = self._start_pos.copy() + des_pos_full[1] = des_pos[0] + des_pos_full[3] = des_pos[1] + des_pos_full[5] = des_pos[2] + return des_pos_full + + def extend_des_vel(self, des_vel): + des_vel_full = self._start_vel.copy() + des_vel_full[1] = des_vel[0] + des_vel_full[3] = des_vel[1] + des_vel_full[5] = des_vel[2] + return des_vel_full + + def render(self, render_mode, **render_kwargs): + if render_mode == "plot_trajectory": + if self._steps == 1: + import matplotlib.pyplot as plt + # plt.ion() + self.fig, self.axs = plt.subplots(3, 1) + + if self._steps <= 1750: + for ax, cp in zip(self.axs, self.current_pos[1::2]): + ax.scatter(self._steps, cp, s=2, marker=".") + + # self.fig.show() + + else: + super().render(render_mode, **render_kwargs) + + +if __name__ == "__main__": + env = ALRBallInACupPDEnv(reward_type="no_context", simplified=True) + # env = gym.make("alr_envs:ALRBallInACupPDSimpleDetPMP-v0") + # ctxt = np.array([-0.20869846, -0.66376693, 1.18088501]) + + # env.configure(ctxt) + env.reset() + env.render("human") + for i in range(16000): + # test with random actions + ac = 0.02 * env.action_space.sample()[0:7] + # ac = env.start_pos + # ac[0] += np.pi/2 + obs, rew, d, info = env.step(ac) + env.render("human") + + print(rew) + + if d: + break + + env.close() + diff --git a/example.py b/example.py index 20fd165..abf69c9 100644 --- a/example.py +++ b/example.py @@ -110,6 +110,6 @@ if __name__ == '__main__': # example_async("alr_envs:LongSimpleReacherDMP-v0", 4) # example_async_contextual_sampler() # env = gym.make("alr_envs:HoleReacherDetPMP-v1") - env_name = "alr_envs:ALRBallInACupSimpleDetPMP-v0" + env_name = "alr_envs:ALRBallInACupPDSimpleDetPMP-v0" example_async_sampler(env_name) # example_mp(env_name)