MyoHub · Vittorio-Caggiano · Sep 16, 2024 · Aug 26, 2024 · Aug 27, 2024 · Aug 27, 2024
diff --git a/myosuite/agents/config/hydra_myochal_sb3_ppo_config.yaml b/myosuite/agents/config/hydra_myochal_sb3_ppo_config.yaml
@@ -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

diff --git a/myosuite/agents/train_myosuite.sh b/myosuite/agents/train_myosuite.sh
@@ -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

diff --git a/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml b/myosuite/envs/myo/assets/arm/myoarm_bionic_bimanual.xml
@@ -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 ======= -->

diff --git a/myosuite/envs/myo/myochallenge/__init__.py b/myosuite/envs/myo/myochallenge/__init__.py
@@ -33,11 +33,14 @@ def register_env_with_variants(id, entry_point, max_episode_steps, kwargs):
 
 register_env_with_variants(id='myoChallengeBimanual-v0',
         entry_point='myosuite.envs.myo.myochallenge.bimanual_v0:BimanualEnvV1',
-        max_episode_steps=250,
+        max_episode_steps=300,
         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
         }
     )
 

diff --git 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,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.
@@ -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):
@@ -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 = {}
 
@@ -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]
@@ -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"]
 
@@ -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()
@@ -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
 
@@ -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