addressed code review changes

Author: vittorione94

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

@@ -0,0 +1,21 @@
+import build_torso
+import build_neck
+import add_markers
+import build_tail
+import create_skin
+import build_front_legs
+import build_back_legs
+import add_muscles
+import add_wrapping_geoms
+import add_torque_actuators
+
+build_torso = build_torso.create_torso
+build_neck = build_neck.create_neck
+add_markers = add_markers.add_markers
+build_tail = build_tail.create_tail
+create_skin = create_skin.create
+build_back_legs = build_back_legs.create_back_legs
+build_front_legs = build_front_legs.create_front_legs
+add_muscles = add_muscles.add_muscles
+add_wrapping_geoms = add_wrapping_geoms.add_wrapping_geoms
+add_torque_actuators = add_torque_actuators.add_motors

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

@@ -271,7 +271,8 @@ def add_muscles(
     )
 
     paths = slices2paths(mtu, slices, muscle_length)
-    spatial = model.tendon.add("spatial", name=f"{mtu}_tendon", dclass="muscle")
+    spatial = model.tendon.add(
+        "spatial", name=f"{mtu}_tendon", dclass="muscle")
     spatials.append(spatial)
     used_muscles.append(mtu)
 
@@ -453,13 +454,13 @@ def add_muscles(
       if muscle_dynamics == "Millard":
         muscle.dynprm = prms
       else:
-        prms[2] = 0.001
+        prms[2] = 2.0
         muscle.dynprm = prms
 
       if spatial.name[: -len("_tendon")] in muscles_constants.NECK:
         scale = 1500
       else:
-        scale = 3000
+        scale = 8000
 
       # range(2), force, scale, lmin, lmax, vmax, fpmax, fvmax
       gainprm = [0.75, 1.05, -1, scale, 0.5, 1.6, 1.5, 1.3, 1.2, 0]

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

@@ -15,8 +15,7 @@
 
 
 def add_wrapping_geoms(model):
-  """
-  Adds wrapping geometries and sites to the specified model.
+  """ Adds wrapping geometries and sites to the specified model.
 
   Args:
       model: The model to which the wrapping geometries and sites will be added.

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

@@ -22,18 +22,7 @@
 from lxml import etree
 
 from dm_control import mjcf
-from dm_control.locomotion.walkers.assets.dog_v2 import add_markers as add_markers_func
-from dm_control.locomotion.walkers.assets.dog_v2 import (
-    add_muscles,
-    add_torque_actuators,
-    add_wrapping_geoms,
-    build_back_legs,
-    build_front_legs,
-    build_neck,
-    build_tail,
-    build_torso,
-    create_skin,
-)
+from dm_control.locomotion.walkers.assets import dog_v2
 from dm_control.utils import io as resources
 
 flags.DEFINE_boolean("make_skin", True, "Whether to make a new dog_skin.skn")
@@ -78,14 +67,18 @@
     False,
     "Make slight adjustments to the model to make it compatible for the composer.",
 )
-
+flags.DEFINE_string(
+    "output_dir",
+    os.path.dirname(__file__),
+    "Output directory for the created dog model"
+)
 
 FLAGS = flags.FLAGS
 
 BASE_MODEL = "dog_base.xml"
 ASSET_RELPATH = "../../../../suite/dog_assets"
 CURRENT_DIR = os.path.dirname(__file__)
-ASSET_DIR = CURRENT_DIR + "/" + ASSET_RELPATH
+ASSET_DIR = os.path.join(CURRENT_DIR, ASSET_RELPATH)
 
 
 def exclude_contacts(model):
@@ -181,6 +174,7 @@ def main(argv):
     add_markers = FLAGS.add_markers
     lengthrange_from_joints = FLAGS.lengthrange_from_joints
     composer = FLAGS.composer
+    output_dir = FLAGS.output_dir
   else:
     lumbar_dofs_per_vert = FLAGS["lumbar_dofs_per_vertebra"].default
     cervical_dofs_per_vertebra = FLAGS["cervical_dofs_per_vertebra"].default
@@ -193,24 +187,24 @@ def main(argv):
     add_markers = FLAGS["add_markers"].default
     lengthrange_from_joints = FLAGS["lengthrange_from_joints"].default
     composer = FLAGS["composer"].default
+    output_dir = FLAGS["output_dir"].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)
-      )
+        "body", name="floor", pos=(0, 0, 0)
+    )
     floor.add(
-          "geom",
-          name="floor",
-          type="plane",
-          conaffinity=1,
-          size=[10, 10, 0.1],
-          material="grid"
-      )
+        "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)),
@@ -267,7 +261,7 @@ def main(argv):
 
   # Torso
   print("Torso, lumbar spine, pelvis.")
-  pelvic_bones, lumbar_joints = build_torso.create_torso(
+  pelvic_bones, lumbar_joints = dog_v2.build_torso(
       model,
       bones,
       bone_position,
@@ -279,7 +273,7 @@ def main(argv):
 
   print("Neck, skull, jaw.")
   # Cervical spine (neck) bodies:
-  cervical_joints = build_neck.create_neck(
+  cervical_joints = dog_v2.build_neck(
       model,
       bone_position,
       cervical_dofs_per_vertebra,
@@ -290,7 +284,7 @@ def main(argv):
   )
 
   print("Back legs.")
-  nails, sole_sites = build_back_legs.create_back_legs(
+  nails, sole_sites = dog_v2.build_back_legs(
       model,
       primary_axis,
       bone_position,
@@ -302,7 +296,7 @@ def main(argv):
   )
 
   print("Shoulders, front legs.")
-  palm_sites = build_front_legs.create_front_legs(
+  palm_sites = dog_v2.build_front_legs(
       nails,
       model,
       primary_axis,
@@ -312,7 +306,7 @@ def main(argv):
   )
 
   print("Tail.")
-  caudal_joints = build_tail.create_tail(
+  caudal_joints = dog_v2.build_tail(
       caudal_dofs_per_vertebra,
       bone_size,
       model,
@@ -365,9 +359,8 @@ def main(argv):
   exclude_contacts(model)
 
   if make_skin:
-    create_skin.create(
-        model=model, mesh_file=skin_msh, asset_dir=ASSET_DIR, mesh_name="dog_skin", composer=composer
-    )
+    dog_v2.create_skin(model=model, mesh_file=skin_msh,
+                       asset_dir=ASSET_DIR, mesh_name="dog_skin", composer=composer)
 
     # Add skin from .skn
     print("Adding Skin.")
@@ -397,7 +390,7 @@ def main(argv):
         muscle_msh = model.asset.add(
             "mesh", name=filename[:-4], file=ASSET_DIR + "/muscles/" + filename
         )
-        create_skin.create(
+        dog_v2.create_skin(
             model=model,
             asset_dir=ASSET_DIR,
             mesh_file=muscle_msh,
@@ -426,22 +419,22 @@ def main(argv):
 
   if add_markers:
     print("Add Markers")
-    add_markers_func.add_markers(model)
+    dog_v2.add_markers(model)
 
   print("Add Actuators")
   if use_muscles:
     # Add wrapping geometries
-    add_wrapping_geoms.add_wrapping_geoms(model)
+    dog_v2.add_wrapping_geoms(model)
     # Add muscles
-    add_muscles.add_muscles(
+    dog_v2.add_muscles(
         model,
         muscle_strength_scale,
         muscle_dynamics,
         ASSET_DIR,
         lengthrange_from_joints,
     )
   else:
-    actuated_joints = add_torque_actuators.add_motors(
+    actuated_joints = dog_v2.add_torque_actuators(
         physics, model, lumbar_joints, cervical_joints, caudal_joints
     )
 
@@ -499,7 +492,7 @@ 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))
@@ -548,7 +541,7 @@ def main(argv):
     name = name_prefix + "dog_muscles_{}_{}.xml".format(
         muscle_strength_scale, muscle_dynamics)
 
-  f = open(os.path.join(CURRENT_DIR, name), "wb")
+  f = open(os.path.join(output_dir, name), "wb")
   f.write(xml_string)
   f.close()