fixed code review issues

Author: vittorione94

dm_control/locomotion/walkers/assets/dog_v2/add_markers.py

@@ -15,23 +15,23 @@
 
 """Add markers to the dog model, used for motion capture tracking."""
 
-markers_per_body = {
-  'torso': [[-0.25, 0, 0.13], [-0.1, 0, 0.13], [0.05, 0, 0.15],
-            [-0.25, -0.06, -0.01], [-0.1, -0.08, -0.02], [0.05, -0.09, -0.02],
-            [-0.25, 0.06, -0.01], [-0.1, 0.08, -0.02], [0.05, 0.09, -0.02]],
-  'C_6': [[0.05, -0.05, 0], [0.05, 0.05, 0]],
-  'upper_arm': [[0.04, 0, 0]],
-  'lower_arm': [[0.025, 0, 0], [0.025, 0, -0.1]],
-  'hand': [[0, -0.02, 0], [0, 0.02, 0], [0.02, -0.02, -0.08],
-           [0.02, 0.02, -0.08]],
-  'finger': [[0.065, 0, 0]],
-  'pelvis': [[0.04, -0.07, 0.02], [0.04, 0.07, 0.02]],
-  'upper_leg': [[-0.03, 0.04, 0]],
-  'lower_leg': [[0.04, 0, 0], [-0.04, 0.035, 0]],
-  'foot': [[0, -0.02, 0], [0, 0.02, 0], [0, -0.02, -0.1], [0, 0.03, -0.1]],
-  'toe': [[0.055, 0, 0]],
-  'Ca_10': [[0, 0, 0.01]],
-  'skull': [[0.01, 0, 0.04], [0, -0.05, 0], [0, 0.05, 0], [-0.1, 0, 0.02]]}
+MARKERS_PER_BODY = {
+  'torso': ((-0.25, 0, 0.13), (-0.1, 0, 0.13), (0.05, 0, 0.15),
+            (-0.25, -0.06, -0.01), (-0.1, -0.08, -0.02), (0.05, -0.09, -0.02),
+            (-0.25, 0.06, -0.01), (-0.1, 0.08, -0.02), (0.05, 0.09, -0.02)),
+  'C_6': ((0.05, -0.05, 0), (0.05, 0.05, 0)),
+  'upper_arm': ((0.04, 0, 0)),
+  'lower_arm': ((0.025, 0, 0), (0.025, 0, -0.1)),
+  'hand': ((0, -0.02, 0), (0, 0.02, 0), (0.02, -0.02, -0.08),
+           (0.02, 0.02, -0.08)),
+  'finger': ((0.065, 0, 0)),
+  'pelvis': ((0.04, -0.07, 0.02), (0.04, 0.07, 0.02)),
+  'upper_leg': ((-0.03, 0.04, 0)),
+  'lower_leg': ((0.04, 0, 0), (-0.04, 0.035, 0)),
+  'foot': ((0, -0.02, 0), (0, 0.02, 0), (0, -0.02, -0.1), (0, 0.03, -0.1)),
+  'toe': ((0.055, 0, 0)),
+  'Ca_10': ((0, 0, 0.01)),
+  'skull': ((0.01, 0, 0.04), (0, -0.05, 0), (0, 0.05, 0), (-0.1, 0, 0.02))}
 
 def add_markers(model):
   bodies = model.find_all('body')
@@ -53,5 +53,5 @@ def add_markers(model):
   for i in range(50):
     marker_body = model.worldbody.add('body', name="marker_" + str(i), mocap=True)
     marker_body.add('site', name="marker_" + str(i), dclass="mocap_marker")
-
-  print("TOT_MARKERS:", TOTAL_MARKERS)
+  
+  print("TOTAL MARKERS ADDED:", TOTAL_MARKERS)

dm_control/locomotion/walkers/assets/dog_v2/add_muscles.py

@@ -19,28 +19,28 @@
 from os.path import isfile, join
 
 import numpy as np
-from scipy.spatial.transform import Rotation as Rot
+from dm_control.utils.transformations import *
 from pykdtree.kdtree import KDTree
 import trimesh
 import pyvista as pv
 from dm_control import mjcf as mjcf_module
 from dm_control.mujoco.wrapper import mjbindings
 
-from muscles import extensors_back, extensors_front, flexors_back, \
-  flexors_front, lateral, neck, tail, torso
+from muscles import EXTENSORS_BACK, EXTENSORS_FRONT, FLEXORS_BACK, \
+  FLEXORS_FRONT, LATERAL, NECK, TAIL, TORSO
 
 from utils import array_to_string, slices2paths
 
-muscle_legs = extensors_back + extensors_front + \
-              flexors_back + flexors_front
+MUSCLE_LEGS = EXTENSORS_BACK + EXTENSORS_FRONT + \
+              FLEXORS_BACK + FLEXORS_FRONT
 
-wrap_geoms_legs = ['shoulder_L_wrapping', 'elbow_L_wrapping',
+WRAP_GEOMS_LEGS = ('shoulder_L_wrapping', 'elbow_L_wrapping',
                    'wrist_L_wrapping', 'finger_L_wrapping', 'hip_L_wrapping',
                    'knee_L_wrapping', 'toe_L_wrapping',
 
                    'shoulder_R_wrapping', 'elbow_R_wrapping',
                    'wrist_R_wrapping', 'finger_R_wrapping', 'hip_R_wrapping',
-                   'knee_R_wrapping', 'toe_R_wrapping']
+                   'knee_R_wrapping', 'toe_R_wrapping')
 
 
 def getClosestGeom(kd_tree, point, mjcf, mtu):
@@ -57,14 +57,14 @@ def getClosestGeom(kd_tree, point, mjcf, mtu):
         closest_geom_body: The closest geometric body (as a Mujoco body object)
         to the target point.
   """
-  dist = 100000000000
+  dist = np.inf
   closest_geom_body = 0
   for bone, tree in kd_tree.items():
     dist_new, i = tree.query(np.array([point]), k=1)
     geom = mjcf.find("geom", bone)
     body = geom.parent
 
-    if mtu in muscle_legs:
+    if mtu in MUSCLE_LEGS:
       if body.name not in ["upper_leg_L",
                            "upper_leg_R", "lower_leg_L", "lower_leg_R",
                            "foot_L", "foot_R", "toe_L", "toe_R",
@@ -116,9 +116,9 @@ def calculate_transformation(element):
       pos = np.zeros(3)
 
     if element.quat is not None:
-      rot = Rot.from_quat(element.quat).as_matrix()
+      rot = quat_to_mat(element.get('quat'))
     else:
-      rot = Rot.identity().as_matrix()
+      rot = np.eye(3)
 
     all_transformations.append(
       np.vstack(
@@ -137,11 +137,35 @@ def calculate_transformation(element):
 
 def add_muscles(model, scale_multiplier, muscle_dynamics,
                 asset_dir, lengthrange_from_joints):
+  """
+    Add muscles to a Mujoco (MJCF) model.
+
+    This function adds muscles to a Mujoco model, 
+    creating MTUs (muscle-tendon units) and defining
+    their properties.
+
+    Args:
+        model (MJCFModel): The input Mujoco model to which muscles will be added.
+        scale_multiplier (float): A scaling factor for muscle forces (0 to disable anatomical scaling).
+        muscle_dynamics (str): Muscle dynamics model, either 'Millard', 'Sigmoid', or 'General'.
+        asset_dir (str): The directory path containing muscle and bone assets.
+        lengthrange_from_joints (bool): If True, compute length ranges from joint limits.
+
+    Returns:
+        None
+
+    Note:
+    - The function modifies the input `model` in place by adding muscles and related properties.
+    - The specific muscles added depend on the provided `muscle_dynamics` and other parameters.
+
+    Raises:
+        NameError: If an unsupported `muscle_dynamics` value is provided.
+    """
   physics = mjcf_module.Physics.from_mjcf_model(model)
 
   mjcf = model
 
-  muscle_meshes_path = asset_dir + '/muscles/'
+  muscle_meshes_path = join(asset_dir, 'muscles/')
   bones_path = asset_dir
 
   bones = [f for f in listdir(bones_path)
@@ -171,8 +195,8 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
 
   mjcf.option.timestep = 0.005
 
-  muscles = extensors_front + flexors_front + \
-            extensors_back + flexors_back + torso + neck + tail
+  muscles = EXTENSORS_FRONT + FLEXORS_FRONT + \
+            EXTENSORS_BACK + FLEXORS_BACK + TORSO + NECK + TAIL
 
   used_muscles = []
   volumes = []
@@ -184,7 +208,7 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
     pv_mesh = pv.read(muscle_meshes_path + mtu + '.stl')
     volumes.append(pv_mesh.volume)
     # check along which axis to slice
-    if mtu not in lateral:
+    if mtu not in LATERAL:
       plane_normal = [0, 0, -1]
       start = np.argmax(m.vertices[:, 2])
       end = np.argmin(m.vertices[:, 2])
@@ -213,9 +237,9 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
     used_muscles.append(mtu)
 
     # print("adding spatial", spatial.name)
-    if mtu in flexors_back or mtu in flexors_front:
+    if mtu in FLEXORS_BACK or mtu in FLEXORS_FRONT:
       spatial.dclass = "flexors"
-    elif mtu in extensors_back or mtu in extensors_front:
+    elif mtu in EXTENSORS_BACK or mtu in EXTENSORS_FRONT:
       spatial.dclass = "extensors"
 
     counter = 0  # used for site naming
@@ -229,9 +253,9 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
       # Add site to tendon
       min_dist = 0.05
       max_dist = 0.11
-      if mtu in tail:
+      if mtu in TAIL:
         max_dist = 0.05
-      elif mtu in neck:
+      elif mtu in NECK:
         min_dist = 0.11
         max_dist = 0.15
 
@@ -266,28 +290,28 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
         # check for wrapping geoms
         if idx != len(paths) - 1:
           geoms = physics.named.data.geom_xpos.axes.row.names
-          dist = 100000000000
+          dist = np.inf
           closest_g = None
           closest_geom_pos = None
           for g_name in geoms:
             if 'wrapping' in g_name:
               pos = physics.named.data.geom_xpos[g_name]
               new_dist = np.linalg.norm(pos - point)
-              if (mtu in lateral or mtu in neck) and \
-                  g_name not in wrap_geoms_legs and \
+              if (mtu in LATERAL or mtu in NECK) and \
+                  g_name not in WRAP_GEOMS_LEGS and \
                   pos[2] < point[2] and \
                   new_dist < dist:
                 closest_g = g_name
                 dist = new_dist
                 closest_geom_pos = pos
-              elif (mtu in flexors_back or
-                    mtu in flexors_front or
-                    mtu in extensors_front or
-                    mtu in extensors_back) and \
+              elif (mtu in FLEXORS_BACK or
+                    mtu in FLEXORS_FRONT or
+                    mtu in EXTENSORS_FRONT or
+                    mtu in EXTENSORS_BACK) and \
                   pos[2] < point[2] and \
                   dist > 0.01 and \
                   new_dist < dist and \
-                  g_name in wrap_geoms_legs:
+                  g_name in WRAP_GEOMS_LEGS:
                 # in the legs we are interested only in wrapping
                 # geometries that are lower than the site we
                 # are adding
@@ -296,15 +320,15 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
                 closest_geom_pos = pos
 
           if dist < 0.1:
-            if (mtu in flexors_back or mtu in flexors_front) \
+            if (mtu in FLEXORS_BACK or mtu in FLEXORS_FRONT) \
                   and closest_geom_pos[2] < point[2]:
               spatial.add('geom', geom=closest_g,
                           sidesite=closest_g + '_backward')
-            elif (mtu in extensors_back or mtu in extensors_front) \
+            elif (mtu in EXTENSORS_BACK or mtu in EXTENSORS_FRONT) \
                     and closest_geom_pos[2] < point[2]:
               spatial.add('geom', geom=closest_g,
                           sidesite=closest_g + '_forward')
-            elif mtu in neck:
+            elif mtu in NECK:
               pos1 = physics.named.data.site_xpos[closest_g + '_forward']
               d1 = np.linalg.norm(pos1 - point)
               pos2 = physics.named.data.site_xpos[closest_g + '_backward']
@@ -315,7 +339,7 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
               else:
                 spatial.add('geom', geom=closest_g,
                             sidesite=closest_g + '_backward')
-            elif mtu in torso and dist < 0.05:
+            elif mtu in TORSO and dist < 0.05:
               spatial.add('geom', geom=closest_g,
                           sidesite=closest_g + '_forward')
 
@@ -325,7 +349,7 @@ def add_muscles(model, scale_multiplier, muscle_dynamics,
                                  gaintype="muscle", biastype="muscle")
 
       prms = [0] * 10
-      prms[0] = 0.04
+      prms[0] = 0.01
       prms[1] = 0.04
       if muscle_dynamics == 'Millard':
         muscle.dynprm = prms

dm_control/locomotion/walkers/assets/dog_v2/build_dog.py

@@ -341,7 +341,7 @@ def main(argv):
 
   if make_skin:
     create_skin.create(model=model, mesh_file=skin_msh, asset_dir=ASSET_DIR,
-                       name_mesh='dog_skin')
+                       mesh_name='dog_skin')
 
     # Add skin from .skn
     print('Adding Skin.')
@@ -364,7 +364,7 @@ def main(argv):
         muscle_msh = model.asset.add('mesh', name=filename[:-4],
                                      file=ASSET_DIR + '/muscles/' + filename)
         create_skin.create(model=model, asset_dir=ASSET_DIR,
-                           mesh_file=muscle_msh, name_mesh=filename[:-4],
+                           mesh_file=muscle_msh, mesh_name=filename[:-4],
                            tex_coords=False, transform=False)
 
         # Add skin from .skn
@@ -392,7 +392,8 @@ def main(argv):
     add_wrapping_geoms.add_wrapping_geoms(model)
     # Add muscles
     add_muscles.add_muscles(model, muscle_strength_scale,
-                            muscle_dynamics, ASSET_DIR, lengthrange_from_joints)
+                            muscle_dynamics, ASSET_DIR, 
+                            lengthrange_from_joints)
   else:
     actuated_joints = add_torque_actuators.add_motors(
       physics, model, lumbar_joints, cervical_joints, caudal_joints

dm_control/locomotion/walkers/assets/dog_v2/create_skin.py

@@ -23,13 +23,13 @@
 from scipy import spatial
 
 
-def create(model, mesh_file, name_mesh, asset_dir, tex_coords=True, transform=True):
+def create(model, mesh_file, mesh_name, asset_dir, tex_coords=True, transform=True):
   """Create and add skin in the dog model.
 
   Args:
     model: model in which we want to add the skin.
     mesh_file: a binary mesh format of the skin.
-    name_mesh: the name of the skin to use in the xml
+    mesh_name: the name of the skin to use in the xml
     asset_dir: asset directory to load the .skn file
     tex_coords: boolean to indicate if the mesh has texture coordinates.
     transform: a boolean to rotate mesh orientation by 90 degrees along the z-axis.
@@ -39,8 +39,8 @@ def create(model, mesh_file, name_mesh, asset_dir, tex_coords=True, transform=Tr
   if transform:
     skinmesh = model.worldbody.add(
       'geom',
-      name=name_mesh,
-      mesh=name_mesh,
+      name=mesh_name,
+      mesh=mesh_name,
       type='mesh',
       contype=0,
       conaffinity=0,
@@ -50,28 +50,28 @@ def create(model, mesh_file, name_mesh, asset_dir, tex_coords=True, transform=Tr
   else:
     skinmesh = model.worldbody.add(
       'geom',
-      name=name_mesh,
-      mesh=name_mesh,
+      name=mesh_name,
+      mesh=mesh_name,
       type='mesh',
       contype=0,
       conaffinity=0)
   physics = mjcf.Physics.from_mjcf_model(model)
 
   # Get skinmesh vertices in global coordinates
-  vertadr = physics.named.model.mesh_vertadr[name_mesh]
-  vertnum = physics.named.model.mesh_vertnum[name_mesh]
+  vertadr = physics.named.model.mesh_vertadr[mesh_name]
+  vertnum = physics.named.model.mesh_vertnum[mesh_name]
   skin_vertices = physics.model.mesh_vert[vertadr:vertadr + vertnum, :]
   skin_vertices = skin_vertices.dot(
-    physics.named.data.geom_xmat[name_mesh].reshape(3, 3).T)
-  skin_vertices += physics.named.data.geom_xpos[name_mesh]
+    physics.named.data.geom_xmat[mesh_name].reshape(3, 3).T)
+  skin_vertices += physics.named.data.geom_xpos[mesh_name]
   skin_normals = physics.model.mesh_normal[vertadr:vertadr + vertnum, :]
   skin_normals = skin_normals.dot(
-    physics.named.data.geom_xmat[name_mesh].reshape(3, 3).T)
-  skin_normals += physics.named.data.geom_xpos[name_mesh]
+    physics.named.data.geom_xmat[mesh_name].reshape(3, 3).T)
+  skin_normals += physics.named.data.geom_xpos[mesh_name]
 
   # Get skinmesh faces
-  faceadr = physics.named.model.mesh_faceadr[name_mesh]
-  facenum = physics.named.model.mesh_facenum[name_mesh]
+  faceadr = physics.named.model.mesh_faceadr[mesh_name]
+  facenum = physics.named.model.mesh_facenum[mesh_name]
   skin_faces = physics.model.mesh_face[faceadr:faceadr + facenum, :]
 
   # Make skin
@@ -191,7 +191,7 @@ def box(sx, sy, sz):
 
   # Convert skin into *.skn file according to
   # https://mujoco.readthedocs.io/en/latest/XMLreference.html#asset-skin
-  f = open(asset_dir + '/skins/' + name_mesh + '.skn', 'w+b')
+  f = open(asset_dir + '/skins/' + mesh_name + '.skn', 'w+b')
 
   nvert = skin.vertex.size // 3
 

dm_control/locomotion/walkers/assets/dog_v2/dog_base.xml

@@ -50,14 +50,14 @@
       <joint group="3" pos="0.012 0 0.01" armature="0.01" damping="3.0"/>
       <general gainprm="40"/>
       <default class="lumbar_extend">
-        <joint stiffness="30.0" range="-24.5 24.5"/>
+        <joint stiffness="30.0" range="-13 13"/>
         <general gainprm="60"/>
       </default>
       <default class="lumbar_bend">
-        <joint stiffness="45.0" range="-21 21"/>
+        <joint stiffness="45.0" range="-15 15"/>
       </default>
       <default class="lumbar_twist">
-        <joint stiffness="45.0" range="-16.333 16.333"/>
+        <joint stiffness="45.0" range="-10 10"/>
       </default>
     </default>
     <default class="cervical">