added composer flag, and small adjustements to code for code review

Author: vittorione94

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

@@ -47,9 +47,6 @@ def add_markers(model):
 
   Args:
       model: The model to add markers to.
-
-  Returns:
-      None
   """
   bodies = model.find_all("body")
 
@@ -71,7 +68,7 @@ def add_markers(model):
 
           marker_idx += 1
           total_markers += 1
-
+  
   for i in range(total_markers):
     marker_body = model.worldbody.add(
         "body", name="marker_" + str(i), mocap=True)

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

@@ -218,12 +218,7 @@ def add_muscles(
       bones_meshes[bone[4:-4]] = bone_mesh
 
   muscle = model.default.default["muscle"]
-
-  flexors = muscle.add("default", dclass="flexors")
-  flexors.tendon.rgba = [0.5, 0, 0, 1]
-
-  extensors = muscle.add("default", dclass="extensors")
-  extensors.tendon.rgba = [0.5, 0, 0, 1]
+  muscle.tendon.rgba = [0.5, 0, 0, 1]
 
   length_range = model.compiler.lengthrange
   length_range.inttotal = 5
@@ -276,19 +271,9 @@ def add_muscles(
     )
 
     paths = slices2paths(mtu, slices, muscle_length)
-    spatial = model.tendon.add("spatial", name=f"{mtu}_tendon")
+    spatial = model.tendon.add("spatial", name=f"{mtu}_tendon", dclass="muscle")
     spatials.append(spatial)
     used_muscles.append(mtu)
-    if (
-        mtu in muscles_constants.FLEXORS_BACK
-        or mtu in muscles_constants.FLEXORS_FRONT
-    ):
-      spatial.dclass = "flexors"
-    elif (
-        mtu in muscles_constants.EXTENSORS_BACK
-        or mtu in muscles_constants.EXTENSORS_FRONT
-    ):
-      spatial.dclass = "extensors"
 
     counter = 0  # used for site naming
     prev_body = None
@@ -447,15 +432,10 @@ def add_muscles(
           continue
 
       physics.forward()
-
       vector_min = np.minimum(vector_min, physics.data.ten_length[:])
       vector_max = np.maximum(vector_max, physics.data.ten_length[:])
-
       physics.reset()
 
-    # physics.named.model.actuator_lengthrange[:, 0] = vector_min
-    # physics.named.model.actuator_lengthrange[:, 1] = vector_max
-
   for tend_idx, spatial in enumerate(spatials):
     if muscle_dynamics in ["Millard", "Sigmoid"]:
       muscle = model.actuator.add(

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

@@ -20,9 +20,6 @@ def add_wrapping_geoms(model):
 
   Args:
       model: The model to which the wrapping geometries and sites will be added.
-
-  Returns:
-      None
   """
   pelvis = model.find("body", "pelvis")
   pelvis.add(

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

@@ -73,6 +73,12 @@
     1,
     "Number of degrees of freedom vertebra in caudal spine.",
 )
+flags.DEFINE_boolean(
+    "composer",
+    False,
+    "Make slight adjustments to the model to make it compatible for the composer.",
+)
+
 
 FLAGS = flags.FLAGS
 
@@ -174,6 +180,7 @@ def main(argv):
     muscle_dynamics = FLAGS.muscle_dynamics
     add_markers = FLAGS.add_markers
     lengthrange_from_joints = FLAGS.lengthrange_from_joints
+    composer = FLAGS.composer
   else:
     lumbar_dofs_per_vert = FLAGS["lumbar_dofs_per_vertebra"].default
     cervical_dofs_per_vertebra = FLAGS["cervical_dofs_per_vertebra"].default
@@ -185,11 +192,25 @@ def main(argv):
     muscle_dynamics = FLAGS["muscle_dynamics"].default
     add_markers = FLAGS["add_markers"].default
     lengthrange_from_joints = FLAGS["lengthrange_from_joints"].default
+    composer = FLAGS["composer"].default
 
   print("Load base model.")
   with open(os.path.join(CURRENT_DIR, BASE_MODEL), "r") as f:
     model = mjcf.from_file(f)
 
+
+  if not composer:
+    floor = model.worldbody.add(
+          "body", name="floor", pos=(0, 0, 0)
+      )
+    floor.add(
+          "geom",
+          name="floor",
+          type="plane",
+          conaffinity=1,
+          size=[10, 10, 0.1],
+          material="grid"
+      )
   # Helper constants:
   side_sign = {
       "_L": np.array((1.0, -1.0, 1.0)),
@@ -253,6 +274,7 @@ def main(argv):
       lumbar_dofs_per_vert,
       side_sign,
       parent=model.worldbody,
+      composer=composer,
   )
 
   print("Neck, skull, jaw.")
@@ -344,7 +366,7 @@ def main(argv):
 
   if make_skin:
     create_skin.create(
-        model=model, mesh_file=skin_msh, asset_dir=ASSET_DIR, mesh_name="dog_skin"
+        model=model, mesh_file=skin_msh, asset_dir=ASSET_DIR, mesh_name="dog_skin", composer=composer
     )
 
     # Add skin from .skn
@@ -382,6 +404,7 @@ def main(argv):
             mesh_name=filename[:-4],
             tex_coords=False,
             transform=False,
+            composer=composer,
         )
 
         # Add skin from .skn
@@ -397,9 +420,9 @@ def main(argv):
         muscle_msh.remove()
 
   model.option.timestep = 0.005
-  model.option.integrator = "Euler"
+  model.option.integrator = "implicitfast"
   model.option.noslip_iterations = 3
-  model.option.cone = "pyramidal"
+  model.option.cone = "elliptic"
 
   if add_markers:
     print("Add Markers")
@@ -473,6 +496,10 @@ def main(argv):
           )
       )
 
+  if composer:
+    torso = model.find("body", "torso")
+    torso.pos = np.zeros(3)
+    
   print("Finalising and saving model.")
   xml_string = model.to_xml_string("float", precision=4, zero_threshold=1e-7)
   root = etree.XML(xml_string, etree.XMLParser(remove_blank_text=True))
@@ -513,11 +540,12 @@ def main(argv):
   newlines.append(b"")
   xml_string = b"\n".join(newlines)
 
+  name_prefix = "composer_" if composer else ""
   # Save to file.
   if not use_muscles:
-    name = "dog.xml"
+    name = name_prefix + "dog.xml"
   else:
-    name = "dog_muscles_{}_{}.xml".format(
+    name = name_prefix + "dog_muscles_{}_{}.xml".format(
         muscle_strength_scale, muscle_dynamics)
 
   f = open(os.path.join(CURRENT_DIR, name), "wb")

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

@@ -23,7 +23,7 @@
 
 
 def create_torso(
-    model, bones, bone_position, lumbar_dofs_per_vertebra, side_sign, parent
+    model, bones, bone_position, lumbar_dofs_per_vertebra, side_sign, parent, composer
 ):
   """Add torso in the dog model.
 
@@ -36,7 +36,7 @@ def create_torso(
     side_sign: a dictionary with two axis representing the signs of
       translations.
     parent: parent object on which we should start attaching new components.
-
+    composer: boolean to determine if a model used by the composer is being created.
   Returns:
     The tuple `(pelvic_bones, lumbar_joints)`.
   """
@@ -55,7 +55,9 @@ def create_torso(
   sternum = [m for m in bones if "Sternum" in m]
   torso_bones = thoracic_spine + ribs + sternum  # + ['Xiphoid_cartilage']
   torso = parent.add("body", name="torso")
-  torso.add("freejoint", name="root")
+
+  if not composer:
+    torso.add("freejoint", name="root")
 
   torso.add("site", name="root", size=(0.01,), rgba=[0, 1, 0, 1])
   torso.add("light", name="light", mode="trackcom", pos=[0, 0, 3])
@@ -84,7 +86,6 @@ def create_torso(
   torso.pos = torso_pos
   for geom in torso_geoms:
     geom.pos = -torso_pos
-
   # Collision primitive for torso
   torso.add(
       "geom",

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

@@ -24,7 +24,7 @@
 from dm_control.mujoco.wrapper.mjbindings import enums
 
 
-def create(model, mesh_file, mesh_name, asset_dir, tex_coords=True, transform=True):
+def create(model, mesh_file, mesh_name, asset_dir, tex_coords=True, transform=True, composer=False):
   """Create and add skin in the dog model.
 
   Args:
@@ -34,6 +34,7 @@ def create(model, mesh_file, mesh_name, asset_dir, tex_coords=True, transform=Tr
     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.
+    composer: boolean to determine if a model used by the composer is being created.
   """
   # Add skin mesh:
   if transform:
@@ -132,7 +133,8 @@ def box(sx, sy, sz):
   # Find smallest distance between
   # each skin vertex and vertices of all meshes in body i
   distance = np.ones((skin_vertices.shape[0], physics.model.nbody)) * 1e6
-  for i in range(2, physics.model.nbody):
+  start = 1 if composer else 2  # in composer mode we don't have the floor
+  for i in range(start, physics.model.nbody):
     geom_id = np.argwhere(physics.model.geom_bodyid == i).ravel()
     mesh_id = physics.model.geom_dataid[geom_id]
     body_verts = []
@@ -175,7 +177,7 @@ def box(sx, sy, sz):
   weights[weights < threshold] = 0
   weights /= np.atleast_2d(np.sum(weights, axis=1)).T
 
-  for i in range(2, physics.model.nbody):
+  for i in range(start, physics.model.nbody):
     vertweight = weights[weights[:, i - 1] >= threshold, i - 1]
     vertid = np.argwhere(weights[:, i - 1] >= threshold).ravel()
     if vertid.any():