Skip to content

Commit

Permalink
Merge pull request #216 from Balint-H/myochallenge/manipulation-varia…
Browse files Browse the repository at this point in the history 
…tion

Add environmental variations to bimanual task
  • Loading branch information
@Vittorio-Caggiano
Vittorio-Caggiano authored Sep 7, 2024
2 parents 76f3925 + 4e4df20 commit 5c8bb0e
Show file tree
Hide file tree
Showing 5 changed files with 122 additions and 58 deletions.
2 changes: 1 addition & 1 deletion myosuite/agents/config/hydra_myochal_sb3_ppo_config.yaml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ batch_size : 64
gamma : 0.99

# PPO.learn function
total_timesteps : 150000
total_timesteps : 1500000
log_interval : 10000

eval_freq : 1000000
Expand Down
2 changes: 1 addition & 1 deletion myosuite/agents/train_myosuite.sh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,7 +55,7 @@ if [ "$4" == 'baseline' ]; then

checkpoint_path="${base_dir}/${filename}"

model_url="https://drive.google.com/uc?export=download&id=1e7bj2Lrk50FU7nSRnIDg2JmYjGLBNChs"
model_url="https://drive.google.com/uc?export=download&id=1P5ip5yjtL4ynbxwDmEOkuF163TG8Vaml"

# Use curl with -L to follow redirects and -o to specify output file
curl -L $model_url -o $checkpoint_path
Expand Down
4 changes: 2 additions & 2 deletions myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@

<!-- ======= Object ======= -->
<body pos="-0.4 -0.2 1.2" euler="1.57 0 3.14" name="manip_object">
<freejoint/>
<freejoint name="manip_object/freejoint"/>
<include file="../../../../simhive/YCB_sim/includes/body_009_gelatin_box.xml"/>
<site name="touch_site" pos = "0 0 0.015" type="sphere" size="0.1" group="5" rgba="0.1 0.1 0.1 0.1"/>
<site name="touch_site" pos = "0 0 0.015" type="sphere" size="0.1" group="5" rgba="0.1 0.1 0.1 0.1"/>
</body>

<!-- ======= Start ======= -->
Expand Down
13 changes: 13 additions & 0 deletions myosuite/envs/myo/myochallenge/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
159 changes: 105 additions & 54 deletions myosuite/envs/myo/myochallenge/bimanual_v0.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -71,24 +72,25 @@ 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

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.
Expand All @@ -113,16 +115,29 @@ def _setup(self,
self.over_max = False
self.max_force = 0


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,
**kwargs,
)
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

def _obj_label_to_obs(self, touching_body):
Expand All @@ -143,6 +158,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)
q = self.sim.model.geom(self.obj_gid - 1).quat
r = R.from_quat([q[1], q[2], q[3], q[0]])
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 = {}

Expand All @@ -151,28 +186,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 </freejoint>
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]
Expand All @@ -194,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"]

Expand Down Expand Up @@ -301,26 +334,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()
Expand All @@ -329,9 +361,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

Expand All @@ -352,35 +384,54 @@ 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)
elif model.geom(con.geom2).bodyid == id_info.manip_body_id:
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
Expand Down

0 comments on commit 5c8bb0e

Please sign in to comment.