dm_control: Import of refs/pull/366/head

PiperOrigin-RevId: 507444578
Change-Id: I6b553ec182e1cade0dca89b1f0d322164eb047c7

Author: copybara-github

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

@@ -27,6 +27,9 @@ def add_motors(physics, model, lumbar_joints, cervical_joints, caudal_joints):
     lumbar_joints: a list of joints objects.
     cervical_joints: a list of joints objects.
     caudal_joints: a list of joints objects.
+
+  Returns:
+    A list of actuated joints.
   """
   # Fixed Tendons:
   spinal_joints = collections.OrderedDict()
@@ -37,31 +40,32 @@ def add_motors(physics, model, lumbar_joints, cervical_joints, caudal_joints):
   for region in spinal_joints.keys():
     for direction in ['extend', 'bend', 'twist']:
       joints = [
-        joint for joint in spinal_joints[region] if direction in joint.name
+          joint for joint in spinal_joints[region] if direction in joint.name
       ]
       if joints:
         tendon = model.tendon.add(
-          'fixed', name=region + direction, dclass=joints[0].dclass)
+            'fixed', name=region + direction, dclass=joints[0].dclass
+        )
         tendons.append(tendon)
         joint_inertia = physics.bind(joints).M0
-        coefs = joint_inertia ** .25
+        coefs = joint_inertia**0.25
         coefs /= coefs.sum()
         coefs *= len(joints)
         for i, joint in enumerate(joints):
           tendon.add('joint', joint=joint, coef=coefs[i])
 
   # Actuators:
-  all_spinal_joints = [
-    joint for region in spinal_joints.values() for joint in region
-    # pylint: disable=g-complex-comprehension
-  ]
+  all_spinal_joints = []
+  for region in spinal_joints.values():
+    all_spinal_joints.extend(region)
   root_joint = model.find('joint', 'root')
   actuated_joints = [
-    joint for joint in model.find_all('joint')
-    if joint not in all_spinal_joints and joint is not root_joint
+      joint
+      for joint in model.find_all('joint')
+      if joint not in all_spinal_joints and joint is not root_joint
   ]
   for tendon in tendons:
-    gain = 0.
+    gain = 0.0
     for joint in tendon.joint:
       # joint.joint.user = physics.bind(joint.joint).damping
       def_joint = model.default.find('default', joint.joint.dclass)
@@ -70,10 +74,12 @@ def add_motors(physics, model, lumbar_joints, cervical_joints, caudal_joints):
     gain /= len(tendon.joint)
 
     model.actuator.add(
-      'general', tendon=tendon, name=tendon.name, dclass=tendon.dclass)
+        'general', tendon=tendon, name=tendon.name, dclass=tendon.dclass
+    )
 
   for joint in actuated_joints:
     model.actuator.add(
-      'general', joint=joint, name=joint.name, dclass=joint.dclass)
+        'general', joint=joint, name=joint.name, dclass=joint.dclass
+    )
 
   return actuated_joints