From 96fae9cffea2152dd06c26ce775d088bb57cb2cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A1lint=20Hodossy?= Date: Thu, 29 Aug 2024 14:45:28 +0100 Subject: [PATCH] Add environmental variations to bimanual task --- .../myo/assets/arm/myoarm_bionic_bimanual.xml | 4 +- myosuite/envs/myo/myochallenge/__init__.py | 13 ++ myosuite/envs/myo/myochallenge/bimanual_v0.py | 160 ++++++++++++------ 3 files changed, 121 insertions(+), 56 deletions(-) diff --git a/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml b/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml index 028167ca..b4f2bf3c 100644 --- a/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml +++ b/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml @@ -53,9 +53,9 @@ - + - + diff --git a/myosuite/envs/myo/myochallenge/__init__.py b/myosuite/envs/myo/myochallenge/__init__.py index 1990fdef..45d75eab 100644 --- a/myosuite/envs/myo/myochallenge/__init__.py +++ b/myosuite/envs/myo/myochallenge/__init__.py @@ -41,6 +41,19 @@ def register_env_with_variants(id, entry_point, max_episode_steps, kwargs): } ) +register_env_with_variants(id='myoChallengeBimanual-v1', + entry_point='myosuite.envs.myo.myochallenge.bimanual_v0:BimanualEnvV1', + max_episode_steps=250, + kwargs={ + 'model_path': curr_dir + '/../assets/arm/myoarm_bionic_bimanual.xml', + 'normalize_act': True, + 'frame_skip': 5, + 'obj_scale_change': [0.1, 0.05, 0.1], # 10%, 5%, 10% scale variations in respective geom directions + 'obj_mass_change': (-0.050, 0.050), # +-50gms + 'obj_friction_change': (0.1, 0.001, 0.00002) # nominal: 1.0, 0.005, 0.0001 + } + ) + # MyoChallenge 2024 envs ============================================== register_env_with_variants(id='myoChallengeRunTrackP1-v0', entry_point='myosuite.envs.myo.myochallenge.run_track_v0:RunTrack', diff --git a/myosuite/envs/myo/myochallenge/bimanual_v0.py b/myosuite/envs/myo/myochallenge/bimanual_v0.py index 17430124..b3e2388e 100644 --- a/myosuite/envs/myo/myochallenge/bimanual_v0.py +++ b/myosuite/envs/myo/myochallenge/bimanual_v0.py @@ -8,6 +8,7 @@ import enum import os, time +from scipy.spatial.transform import Rotation as R import mujoco from myosuite.utils import gym import numpy as np @@ -24,7 +25,7 @@ class BimanualEnvV1(BaseV0): "object_qvel", "touching_body"] DEFAULT_RWD_KEYS_AND_WEIGHTS = { - "reach_dist": -.1, + "reach_dist": -.1, "act": 0, "fin_dis": -0.5, # "fin_open": -1, @@ -55,9 +56,9 @@ def _setup(self, # shift factor for start/goal random generation with z-axis fixed goal_shifts=np.array([0.098, 0.098, 0]), - obj_size_range=None, # Object size range. Nominal 0.022 - obj_mass_range=None, # {'high': 0.50, 'low': 0.050}, # Object weight range. Nominal 43 gms - obj_friction_range=None, + obj_scale_change=None, # object size change (relative to initial size) + obj_mass_change=None, # object size change (relative to initial size) + obj_friction_change=None, # object friction change (relative to initial size) # {'high': [1.2, 0.006, 0.00012], 'low': [0.8, 0.004, 0.00008]}, # friction change task_choice='fixed', # fixed/ random obs_keys: list = DEFAULT_OBS_KEYS, @@ -71,11 +72,6 @@ def _setup(self, self.task_choice = task_choice self.proximity_th = proximity_th - # setup for task randomization - self.obj_mass_range = self.obj_mass_range = obj_mass_range - self.obj_size_range = {'low': obj_size_range[0], 'high': obj_size_range[1]} if obj_size_range else None - self.obj_friction_range = obj_friction_range - # start position centers (before changes) self.start_center = start_center self.goal_center = goal_center @@ -83,12 +79,18 @@ def _setup(self, self.start_shifts = start_shifts self.goal_shifts = goal_shifts - self.start_bid = self.sim.model.body_name2id('start') - self.goal_bid = self.sim.model.body_name2id('goal') + self.id_info = IdInfo(self.sim.model) + + self.start_bid = self.id_info.start_id + self.goal_bid = self.id_info.goal_id - self.object_bid = self.sim.model.body_name2id('manip_object') - self.object_sid = self.sim.model.site_name2id('touch_site') - self.init_obj_z = self.sim.data.site_xpos[self.object_sid][-1] + self.obj_bid = self.id_info.manip_body_id + self.obj_sid = self.sim.model.site_name2id('touch_site') + self.obj_gid = self.sim.model.body(self.obj_bid).geomadr + 1 + self.obj_mid = next(i + for i in range(self.sim.model.nmesh) + if "box" in self.sim.model.mesh(i).name) + self.init_obj_z = self.sim.data.site_xpos[self.obj_sid][-1] self.target_z = 0.2 # define the palm and tip site id. @@ -113,8 +115,22 @@ def _setup(self, self.over_max = False self.max_force = max_force + self.touch_history = [] + # setup for task randomization + self.obj_mass_range = ({'low': self.sim.model.body_mass[self.obj_bid]+obj_mass_change[0], + 'high': self.sim.model.body_mass[self.obj_bid]+obj_mass_change[1]} + if obj_mass_change else None) + self.obj_scale_range = ({'low': -np.array(obj_scale_change), 'high': obj_scale_change} + if obj_scale_change else None) + self.obj_friction_range = ({'low': self.sim.model.geom_friction[self.obj_gid] - obj_friction_change, + 'high': self.sim.model.geom_friction[self.obj_gid] + obj_friction_change} + if obj_friction_change else None) + # We'll center the mesh on the box to have an easier time scaling it: + if obj_scale_change: + self.__center_box_mesh() + super()._setup(obs_keys=obs_keys, weighted_reward_keys=weighted_reward_keys, frame_skip=frame_skip, @@ -122,8 +138,8 @@ def _setup(self, ) self.init_qpos[:] = self.sim.model.key_qpos[2].copy() # adding random disturbance to start and goal positions, coefficients might need to be adaptable - self.initialized_pos = False + self.sim.renderer.setup_renderer(self.sim.model.ptr, 500, 500) def _obj_label_to_obs(self, touching_body): # Function to convert touching body set to a binary observation vector @@ -143,6 +159,26 @@ def _obj_label_to_obs(self, touching_body): return obs_vec + def __center_box_mesh(self): + """ + Adjusts the mesh geom's transform and vertices so scaling is straightforward afterwards. Only makes sense + to call this method within setup after relevant ids have been identified. + """ + self.obj_size0 = self.sim.model.geom_size[self.obj_gid].copy() + self.obj_vert_addr = np.arange(self.sim.model.mesh(self.obj_mid).vertadr, + self.sim.model.mesh(self.obj_mid).vertadr + self.sim.model.mesh(0).vertnum) + + r = R.from_quat(self.sim.model.geom(self.obj_gid - 1).quat, scalar_first=True) + self.sim.model.mesh_vert[self.obj_vert_addr] = r.apply(self.sim.model.mesh_vert[self.obj_vert_addr]) + self.sim.model.mesh_normal[self.obj_vert_addr] = r.apply(self.sim.model.mesh_normal[self.obj_vert_addr]) + self.sim.model.geom(self.obj_gid - 1).quat = [1, 0, 0, 0] + self.sim.model.mesh_vert[self.obj_vert_addr] += (self.sim.model.geom(self.obj_gid - 1).pos + - self.sim.model.geom(self.obj_gid).pos)[None, :] + + self.sim.model.geom(self.obj_gid - 1).pos = self.sim.model.geom(self.obj_gid).pos + self.mesh_vert0 = self.sim.model.mesh_vert[self.obj_vert_addr].copy() + self.ignore_first_scale = True + def get_obs_dict(self, sim): obs_dict = {} @@ -151,28 +187,26 @@ def get_obs_dict(self, sim): obs_dict["qv"] = sim.data.qvel.copy() # MyoHand data - obs_dict["myohand_qpos"] = sim.data.qpos[:38].copy() - obs_dict["myohand_qvel"] = sim.data.qvel[:38].copy() + obs_dict["myohand_qpos"] = sim.data.qpos[self.id_info.myo_joint_range].copy() + obs_dict["myohand_qvel"] = sim.data.qvel[self.id_info.myo_dof_range].copy() # Prosthetic hand data and velocity - obs_dict["pros_hand_qpos"] = sim.data.qpos[38:-6].copy() - obs_dict["pros_hand_qvel"] = sim.data.qvel[38:-6].copy() + obs_dict["pros_hand_qpos"] = sim.data.qpos[self.id_info.prosth_joint_range].copy() + obs_dict["pros_hand_qvel"] = sim.data.qvel[self.id_info.prosth_dof_range].copy() # One more joint for qpos due to - obs_dict["object_qpos"] = sim.data.qpos[-7:].copy() - obs_dict["object_qvel"] = sim.data.qvel[-6:].copy() + obs_dict["object_qpos"] = sim.data.qpos[self.id_info.manip_joint_range].copy() + obs_dict["object_qvel"] = sim.data.qvel[self.id_info.manip_dof_range].copy() obs_dict["start_pos"] = self.start_pos[:2].copy() obs_dict["goal_pos"] = self.goal_pos[:2].copy() obs_dict["elbow_fle"] = self.sim.data.joint('elbow_flexion').qpos.copy() - # TODO : work on the self.sim_model that changes when env resets this_model = sim.model - id_info = BodyIdInfo(this_model) this_data = sim.data # Get touching object in terms of binary encoding - touching_objects = set(get_touching_objects(this_model, this_data, id_info)) + touching_objects = set(get_touching_objects(this_model, this_data, self.id_info)) self.touch_history.append(touching_objects) current_force = sim.data.sensordata[0] @@ -193,7 +227,7 @@ def get_obs_dict(self, sim): obs_dict['MPL_ori'] = mat2euler(np.reshape(self.sim.data.site_xmat[self.Rpalm1_sid], (3, 3))) obs_dict['MPL_ori_err'] = obs_dict['MPL_ori'] - np.array([np.pi, 0, np.pi]) - obs_dict["obj_pos"] = sim.data.site_xpos[self.object_sid] + obs_dict["obj_pos"] = sim.data.site_xpos[self.obj_sid] obs_dict["reach_err"] = obs_dict["palm_pos"] - obs_dict["obj_pos"] obs_dict["pass_err"] = obs_dict["Rpalm_pos"] - obs_dict["obj_pos"] @@ -301,26 +335,25 @@ def reset(self, **kwargs): self.touch_history = [] self.over_max = False - # Following method from baoding challenging + # box mass changes if self.obj_mass_range: - self.sim.model.body_mass[self.object_bid] = self.np_random.uniform(**self.obj_mass_range) - - for gid in range(self.sim.model.body_geomnum[self.object_bid]): - # Calculate the global geometry ID - global_gid = gid + self.sim.model.body_geomadr[self.object_bid] - - # Randomly assign friction if a range is provided - if self.obj_friction_range is not None: - self.sim.model.geom_friction[global_gid] = self.np_random.uniform(**self.obj_friction_range) - - # Randomly determine the rescaling factor - rescale_factor = self.np_random.uniform(0.8, 1.2) - - # Rescale geometry size - self.sim.model.geom_size[global_gid] *= rescale_factor - self.sim.model.geom_aabb[global_gid][3:] = 1.2 # bounding box, (center, size) - self.sim.model.geom_rbound[global_gid] = 2.0 * 1.2 # radius of bounding sphere - + self.sim.model.body_mass[self.obj_bid] = self.np_random.uniform( + **self.obj_mass_range) # call to mj_setConst(m,d) is being ignored. Derive quantities wont be updated. Die is simple shape. So this is reasonable approximation. + + # box friction changes + if self.obj_friction_range: + self.sim.model.geom_friction[self.obj_gid] = self.np_random.uniform(**self.obj_friction_range) + + # box size changes + if self.obj_scale_range and not self.ignore_first_scale: + obj_scales = self.np_random.uniform(**self.obj_scale_range) + 1 + self.sim.model.geom(self.obj_gid).size = self.obj_size0 * obj_scales + + if self.sim.renderer._window: + self.sim.model.mesh_vert[self.obj_vert_addr] = obj_scales[None, :] * self.mesh_vert0 + self.sim.renderer._window.update_mesh(self.obj_mid) + else: + self.ignore_first_scale = False self.sim.forward() self.init_qpos[:] = self.sim.model.key_qpos[2].copy() @@ -329,9 +362,9 @@ def reset(self, **kwargs): obs = super().reset( reset_qpos=self.init_qpos, reset_qvel=self.init_qvel, **kwargs ) - object_qpos_adr = self.sim.model.body(self.object_bid).jntadr[0] + object_qpos_adr = self.sim.model.body(self.obj_bid).jntadr[0] self.sim.data.qpos[object_qpos_adr:object_qpos_adr + 3] = self.start_pos + np.array([0, 0, 0.1]) - self.init_obj_z = self.sim.data.site_xpos[self.object_sid][-1] + self.init_obj_z = self.sim.data.site_xpos[self.obj_sid][-1] self.init_palm_z = self.sim.data.site_xpos[self.palm_sid][-1] return obs @@ -352,24 +385,43 @@ class ContactTrajIssue(enum.Enum): ENV_CONTACT = 3 -# Adding some default values, however this should be updated -class BodyIdInfo: +class IdInfo: def __init__(self, model: mujoco.MjModel): self.manip_body_id = model.body("manip_object").id myo_bodies = [model.body(i).id for i in range(model.nbody) if not model.body(i).name.startswith("prosthesis") and not model.body(i).name in ["start", "goal", "manip_object"]] - self.myo_range = (min(myo_bodies), max(myo_bodies)) + self.myo_body_range = (min(myo_bodies), max(myo_bodies)) prosth_bodies = [model.body(i).id for i in range(model.nbody) if model.body(i).name.startswith("prosthesis/")] - self.prosth_range = (min(prosth_bodies), max(prosth_bodies)) + self.prosth_body_range = (min(prosth_bodies), max(prosth_bodies)) + + self.myo_joint_range = np.concatenate([model.joint(i).qposadr for i in range(model.njnt) + if not model.joint(i).name.startswith("prosthesis") + and not model.joint(i).name == "manip_object/freejoint"]) + + self.myo_dof_range = np.concatenate([model.joint(i).dofadr for i in range(model.njnt) + if not model.joint(i).name.startswith("prosthesis") + and not model.joint(i).name == "manip_object/freejoint"]) + + self.prosth_joint_range = np.concatenate([model.joint(i).qposadr for i in range(model.njnt) + if model.joint(i).name.startswith("prosthesis")]) + + self.prosth_dof_range = np.concatenate([model.joint(i).dofadr for i in range(model.njnt) + if model.joint(i).name.startswith("prosthesis")]) + + self.manip_joint_range = np.arange(model.joint("manip_object/freejoint").qposadr, + model.joint("manip_object/freejoint").qposadr + 7) + + self.manip_dof_range = np.arange(model.joint("manip_object/freejoint").dofadr, + model.joint("manip_object/freejoint").dofadr + 6) self.start_id = model.body("start").id self.goal_id = model.body("goal").id -def get_touching_objects(model: mujoco.MjModel, data: mujoco.MjData, id_info: BodyIdInfo): +def get_touching_objects(model: mujoco.MjModel, data: mujoco.MjData, id_info: IdInfo): for con in data.contact: if model.geom(con.geom1).bodyid == id_info.manip_body_id: yield body_id_to_label(model.geom(con.geom2).bodyid, id_info) @@ -377,10 +429,10 @@ def get_touching_objects(model: mujoco.MjModel, data: mujoco.MjData, id_info: Bo yield body_id_to_label(model.geom(con.geom1).bodyid, id_info) -def body_id_to_label(body_id, id_info: BodyIdInfo): - if id_info.myo_range[0] < body_id < id_info.myo_range[1]: +def body_id_to_label(body_id, id_info: IdInfo): + if id_info.myo_body_range[0] < body_id < id_info.myo_body_range[1]: return ObjLabels.MYO - elif id_info.prosth_range[0] < body_id < id_info.prosth_range[1]: + elif id_info.prosth_body_range[0] < body_id < id_info.prosth_body_range[1]: return ObjLabels.PROSTH elif body_id == id_info.start_id: return ObjLabels.START