Add pyink and isort config. Reformat.

PiperOrigin-RevId: 704533915
Change-Id: I37e9fd51261bd166b725c7460fc65d02fed2b391

STYLEGUIDE.md

@@ -6,7 +6,7 @@ possible in your code contributions.
 
 ### Scope of this guide
 
-MuJoCo has three main code categories:
+Most of this guide involves C/C++ code. For Python, jump to the section [below](#python-code). For MuJoCo C/C++, code has three main categories:
 
 1. **C code:** MuJoCo's core codebase. It consists of public headers under
 `include/` and C source files and internal headers under `src/`. This style
@@ -158,3 +158,7 @@ example above.
 New code should use the C99 convention. When editing an existing function,
 please move existing variable declarations into local scope. Pull requests
 helping us to complete the migration are very welcome.
+
+### [Python code](#python-code)
+
+For Python code, run `pyink foo.py` to adhere to Google's [Python style guide](https://google.github.io/styleguide/pyguide.html). For sorting and cleaning imports, run `isort foo.py`. Both `pyink` and `isort` can be pip installed via `pip install pyink isort`.

mjx/mujoco/mjx/_src/collision_convex.py

@@ -906,7 +906,8 @@ def _sat_gaussmap(
     return edge_axis * sign, degenerate_edge_axis
 
   edge_axes, degenerate_edge_axes = jax.vmap(get_normals)(
-      edge_a_dir, edge_a_pt, edge_b_dir)
+      edge_a_dir, edge_a_pt, edge_b_dir
+  )
   edge_dist = jax.vmap(jp.dot)(edge_axes, edge_b_pt - edge_a_pt)
   # handle degenerate axis
   edge_dist = jp.where(degenerate_edge_axes, -jp.inf, edge_dist)
@@ -928,11 +929,14 @@ def _sat_gaussmap(
       dist,
   )
   a_closest, b_closest = math.closest_segment_to_segment_points(
-      edge_a_pt[best_edge_idx], edge_a_pt_2[best_edge_idx],
+      edge_a_pt[best_edge_idx],
-      edge_b_pt[best_edge_idx], edge_b_pt_2[best_edge_idx])
+      edge_a_pt_2[best_edge_idx],
+      edge_b_pt[best_edge_idx],
+      edge_b_pt_2[best_edge_idx],
+  )
   pos = jp.where(
-      is_edge_contact,
+      is_edge_contact, jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos
-      jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos)
+  )
 
   return dist, pos, normal
 

mjx/mujoco/mjx/_src/collision_driver.py

@@ -146,13 +146,13 @@ def geom_pairs(
   b_end = b_start + m.body_geomnum
 
   for b1 in range(m.nbody):
-    if not geom_con[b_start[b1]:b_end[b1]].any():
+    if not geom_con[b_start[b1] : b_end[b1]].any():
       continue
     w1 = m.body_weldid[b1]
     w1_p = m.body_weldid[m.body_parentid[w1]]
 
     for b2 in range(b1, m.nbody):
-      if not geom_con[b_start[b2]:b_end[b2]].any():
+      if not geom_con[b_start[b2] : b_end[b2]].any():
         continue
       signature = (b1 << 16) + (b2)
       if signature in exclude_signature:
@@ -272,7 +272,7 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
           jp.clip(m.pair_friction[ip], a_min=eps),
           m.pair_solref[ip],
           m.pair_solreffriction[ip],
-          m.pair_solimp[ip]
+          m.pair_solimp[ip],
       ))
     if geom1.size > 0 and geom2.size > 0:
       # other contacts get their params from geom fields

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -218,7 +218,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5)
+          dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
+      )
 
   _ELLIPSOID_ELLIPSOID = """
     <mujoco>
@@ -240,7 +241,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5)
+          dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5
+      )
 
   _ELLIPSOID_SPHERE = """
     <mujoco>
@@ -263,7 +265,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-3)
+          dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-3
+      )
 
   _ELLIPSOID_CAPSULE = """
     <mujoco>
@@ -285,7 +288,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-3)
+          dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-3
+      )
 
   _ELLIPSOID_CYLINDER = """
     <mujoco>
@@ -308,7 +312,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4)
+          dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
+      )
 
 
 class CapsuleCollisionTest(parameterized.TestCase):
@@ -550,7 +555,8 @@ class CylinderTest(absltest.TestCase):
 
     # cylinder is vertical
     xml = self._CYLINDER_PLANE.replace(
-        '<geom fromto="-0.1 0 0 0.1 0 0"', '<geom fromto="0 0 -0.1 0 0 0.1"')
+        '<geom fromto="-0.1 0 0 0.1 0 0"', '<geom fromto="0 0 -0.1 0 0 0.1"'
+    )
     xml = xml.replace('pos="0 0 0.04"', 'pos="0 0 0.095"')
     d, dx = _collide(xml)
 
@@ -579,7 +585,8 @@ class CylinderTest(absltest.TestCase):
     self.assertLess(dx.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4)
+          dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
+      )
 
 
 class ConvexTest(absltest.TestCase):
@@ -760,11 +767,14 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(0, c.dist[1:])
     np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1]))
     np.testing.assert_array_almost_equal(
-        c.pos[0], np.array([0, 2, 1.3155]), decimal=5)
+        c.pos[0], np.array([0, 2, 1.3155]), decimal=5
+    )
 
     _, dx = _collide(
         self._CONVEX_CONVEX_THIN.replace(
-            'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'))
+            'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'
+        )
+    )
     c = dx.contact
     self.assertTrue((c.dist > 0).all())
 

mjx/mujoco/mjx/_src/collision_primitive.py

@@ -29,6 +29,7 @@ from mujoco.mjx._src.types import Model
 
 def collider(ncon: int):
   """Wraps collision functions for use by collision_driver."""
+
   def wrapper(func):
     def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
       g1, g2 = geom.T
@@ -119,7 +120,7 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision:
       # disk parallel to plane: pick x-axis of cylinder, scale by radius
       cylinder.mat[:, 0] * cylinder.size[0],
       # general configuration: normalize vector, scale by radius
-      vec / len_ * cylinder.size[0]
+      vec / len_ * cylinder.size[0],
   )
 
   # project vector on normal
@@ -138,11 +139,15 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision:
   d1 = dist0 + prjaxis + prjvec
   d2 = dist0 + prjaxis + prjvec1
   dist = jp.array([d1, d2, d2])
-  pos = cylinder.pos + axis + jp.array([
+  pos = (
-      vec  - n * d1 * 0.5,
+      cylinder.pos
-      vec1 + vec * -0.5 - n * d2 * 0.5,
+      + axis
-      -vec1 + vec * -0.5 - n * d2 * 0.5,
+      + jp.array([
-  ])
+          vec - n * d1 * 0.5,
+          vec1 + vec * -0.5 - n * d2 * 0.5,
+          -vec1 + vec * -0.5 - n * d2 * 0.5,
+      ])
+  )
 
   # cylinder parallel to plane
   cond = jp.abs(prjaxis) < 1e-3

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -41,6 +41,7 @@ SDFFn = Callable[[jax.Array], jax.Array]
 
 def collider(ncon: int):
   """Wraps collision functions for use by collision_driver."""
+
   def wrapper(func):
     def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
       g1, g2 = geom.T
@@ -81,7 +82,7 @@ def _capsule(pos: jax.Array, size: jax.Array):
 
 def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array:
   k0 = math.norm(pos / size)
-  k1 = math.norm(pos / (size*size))
+  k1 = math.norm(pos / (size * size))
   return k0 * (k0 - 1.0) / (k1 + (k1 == 0.0) * 1e-12)
 
 
@@ -96,12 +97,12 @@ def _cylinder(pos: jax.Array, size: jax.Array) -> jax.Array:
 
 def _cylinder_grad(x: jax.Array, size: jax.Array) -> jax.Array:
   """Gradient of the cylinder SDF wrt query point and singularities removed."""
-  c = jp.sqrt(x[0]*x[0]+x[1]*x[1])
+  c = jp.sqrt(x[0] * x[0] + x[1] * x[1])
   e = jp.abs(x[2])
   a = jp.array([c - size[0], e - size[1]])
   b = jp.array([jp.maximum(a[0], 0), jp.maximum(a[1], 0)])
   j = jp.argmax(a)
-  bnorm = jp.sqrt(b[0]*b[0] + b[1]*b[1])
+  bnorm = jp.sqrt(b[0] * b[0] + b[1] * b[1])
   bnorm += jp.allclose(bnorm, 0) * 1e-12
   grada = jp.array([
       x[0] / (c + jp.allclose(c, 0) * 1e-12),
@@ -151,7 +152,7 @@ def _gradient_step(objective: SDFFn, state: GradientState) -> GradientState:
   """Performs a step of gradient descent."""
   # TODO: find better parameters
   amin = 1e-4  # minimum value for line search factor scaling the gradient
-  amax = 2.  # maximum value for line search factor scaling the gradient
+  amax = 2.0  # maximum value for line search factor scaling the gradient
   nlinesearch = 10  # line search points
   grad = jax.grad(objective)(state.x)
   alpha = jp.geomspace(amin, amax, nlinesearch).reshape(nlinesearch, -1)
@@ -179,7 +180,11 @@ def _gradient_descent(
 
 
 def _optim(
-    d1, d2, info1: GeomInfo, info2: GeomInfo, x0: jax.Array,
+    d1,
+    d2,
+    info1: GeomInfo,
+    info2: GeomInfo,
+    x0: jax.Array,
 ) -> Collision:
   """Optimizes the clearance function."""
   d1 = functools.partial(d1, size=info1.size)
@@ -198,14 +203,14 @@ def _optim(
 
 @collider(ncon=1)
 def sphere_ellipsoid(s: GeomInfo, e: GeomInfo) -> Collision:
-  """"Calculates contact between a sphere and an ellipsoid."""
+  """Calculates contact between a sphere and an ellipsoid."""
   x0 = 0.5 * (s.pos + e.pos)
   return _optim(_sphere, _ellipsoid, s, e, x0)
 
 
 @collider(ncon=1)
 def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision:
-  """"Calculates contact between a sphere and a cylinder."""
+  """Calculates contact between a sphere and a cylinder."""
   # TODO: implement analytical version.
   x0 = 0.5 * (s.pos + c.pos)
   return _optim(_sphere, _cylinder, s, c, x0)
@@ -213,14 +218,14 @@ def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision:
 
 @collider(ncon=1)
 def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
-  """"Calculates contact between a capsule and an ellipsoid."""
+  """ "Calculates contact between a capsule and an ellipsoid."""
   x0 = 0.5 * (c.pos + e.pos)
   return _optim(_capsule, _ellipsoid, c, e, x0)
 
 
 @collider(ncon=2)
 def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision:
-  """"Calculates contact between a capsule and a cylinder."""
+  """Calculates contact between a capsule and a cylinder."""
   # TODO: improve robustness
   # Near sharp corners, the SDF might give the penetration depth with respect
   # to a surface that is not in collision. Possible solutions is to find the
@@ -235,21 +240,21 @@ def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision:
 
 @collider(ncon=1)
 def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision:
-  """"Calculates contact between two ellipsoids."""
+  """Calculates contact between two ellipsoids."""
   x0 = 0.5 * (e1.pos + e2.pos)
   return _optim(_ellipsoid, _ellipsoid, e1, e2, x0)
 
 
 @collider(ncon=1)
 def ellipsoid_cylinder(e: GeomInfo, c: GeomInfo) -> Collision:
-  """"Calculates contact between and ellipsoid and a cylinder."""
+  """Calculates contact between and ellipsoid and a cylinder."""
   x0 = 0.5 * (e.pos + c.pos)
   return _optim(_ellipsoid, _cylinder, e, c, x0)
 
 
 @collider(ncon=4)
 def cylinder_cylinder(c1: GeomInfo, c2: GeomInfo) -> Collision:
-  """"Calculates contact between a cylinder and a cylinder."""
+  """Calculates contact between a cylinder and a cylinder."""
   # TODO: improve robustness
   # Near sharp corners, the SDF might give the penetration depth with respect
   # to a surface that is not in collision. Possible solutions is to find the

mjx/mujoco/mjx/_src/collision_types.py

@@ -73,6 +73,7 @@ class FunctionKey:
       resulting constraint jacobian is determined at compile time.
     subgrid_size: the size determines the hfield subgrid to collide with
   """
+
   types: Tuple[int, int]
   data_ids: Tuple[int, int]
   condim: int

mjx/mujoco/mjx/_src/constraint.py

@@ -39,6 +39,7 @@ import numpy as np
 
 class _Efc(PyTreeNode):
   """Support data for creating constraint matrices."""
+
   J: jax.Array
   pos_aref: jax.Array
   pos_imp: jax.Array

mjx/mujoco/mjx/_src/constraint_test.py

@@ -30,7 +30,7 @@ _TOLERANCE = 5e-5
 
 
 def _assert_eq(a, b, name):
-  tol = _TOLERANCE * 10   # avoid test noise
+  tol = _TOLERANCE * 10  # avoid test noise
   err_msg = f'mismatch: {name}'
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
@@ -75,7 +75,7 @@ class ConstraintTest(parameterized.TestCase):
       _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
       _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
       _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
-      _assert_eq(dx.efc_pos[order][d.nefc:], 0, 'efc_pos')
+      _assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos')
       _assert_eq(
           d.efc_frictionloss,
           dx.efc_frictionloss[order][: d.nefc],

mjx/mujoco/mjx/_src/dataclasses.py

@@ -16,7 +16,6 @@
 
 import copy
 import dataclasses
-
 import typing
 from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
 import jax
@@ -57,7 +56,7 @@ def dataclass(clz: _T) -> _T:
       meta_fields.append(field)
 
   def replace(self, **updates):
-    """"Returns a new object replacing the specified fields with new values."""
+    """Returns a new object replacing the specified fields with new values."""
     return dataclasses.replace(self, **updates)
 
   data_clz.replace = replace

mjx/mujoco/mjx/_src/io.py

@@ -34,6 +34,7 @@ def _strip_weak_type(tree):
     if isinstance(leaf, jax.Array):
       return leaf.astype(jax.dtypes.canonicalize_dtype(leaf.dtype))
     return leaf
+
   return jax.tree_util.tree_map(f, tree)
 
 
@@ -95,7 +96,7 @@ def put_model(
     m: the model to put onto device
     device: which device to use - if unspecified picks the default device
     _full_compat: put all MjModel fields onto device irrespective of MJX support
-        This is an experimental feature.  Avoid using it for now.
+      This is an experimental feature.  Avoid using it for now.
 
   Returns:
     an mjx.Model placed on device
@@ -215,8 +216,8 @@ def make_data(
     m: the model to use
     device: which device to use - if unspecified picks the default device
     _full_compat: create all MjData fields on device irrespective of MJX support
-        This is an experimental feature.  Avoid using it for now.
+      This is an experimental feature.  Avoid using it for now. If using this
-        If using this flag, also use _full_compat for put_model.
+      flag, also use _full_compat for put_model.
 
   Returns:
     an initialized mjx.Data placed on device
@@ -383,7 +384,7 @@ def make_data(
       contact=contact,
       efc_type=efc_type,
       eq_active=m.eq_active0,
-      **zero_fields
+      **zero_fields,
   )
 
   return d
@@ -556,8 +557,8 @@ def put_data(
     d: the data to put on device
     device: which device to use - if unspecified picks the default device
     _full_compat: put all MjModel fields onto device irrespective of MJX support
-        This is an experimental feature.  Avoid using it for now.
+      This is an experimental feature.  Avoid using it for now. If using this
-        If using this flag, also use _full_compat for put_model.
+      flag, also use _full_compat for put_model.
 
   Returns:
     an mjx.Data placed on device
@@ -646,7 +647,7 @@ def put_data(
       if num_rows > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
         num_rows = (num_rows - 1) * 2
       efc_i, efc_o = d.contact.efc_address[id_from], efc_address[id_to]
-      value[efc_o:efc_o + num_rows] = fields[fname][efc_i:efc_i + num_rows]
+      value[efc_o : efc_o + num_rows] = fields[fname][efc_i : efc_i + num_rows]
 
     fields[fname] = value
 

mjx/mujoco/mjx/_src/io_test.py

@@ -92,12 +92,11 @@ _MULTIPLE_CONSTRAINTS = """
 class ModelIOTest(parameterized.TestCase):
   """IO tests for mjx.Model."""
 
-  @parameterized.parameters(
+  @parameterized.parameters(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS)
-      _MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS
-  )
   def test_put_model(self, xml):
     m = mujoco.MjModel.from_xml_string(xml)
     mx = mjx.put_model(m)
+
     def assert_not_weak_type(x):
       if isinstance(x, jax.Array):
         assert not x.weak_type

mjx/mujoco/mjx/_src/math.py

@@ -28,6 +28,7 @@ def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array:
   Args:
     a: left hand of matmul operand
     b: right hand of matmul operand
+
   Returns:
     the matrix product of the inputs.
   """

mjx/mujoco/mjx/_src/mesh.py

@@ -172,8 +172,8 @@ def _merge_coplanar(
 
     # resize faces that exceed max polygon vertices
     if face.shape[0] > _MAX_HULL_FACE_VERTICES:
-      name = m.names[m.name_meshadr[meshid]:]
+      name = m.names[m.name_meshadr[meshid] :]
-      name = name[:name.find(b'\x00')].decode('utf-8')
+      name = name[: name.find(b'\x00')].decode('utf-8')
       warnings.warn(
           f'Mesh "{name}" has a coplanar face with more than '
           f'{_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '

mjx/mujoco/mjx/_src/mesh_test.py

@@ -53,8 +53,9 @@ class MeshTest(absltest.TestCase):
     map_ = {v: k for k, v in enumerate(vidx)}
     h_face = np.vectorize(map_.get)(convex_face)
     face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()])
-    expected_face_verts = sorted([
+    expected_face_verts = sorted(
-        (0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)])
+        [(0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)]
+    )
     self.assertSequenceEqual(
         face_verts,
         expected_face_verts,

mjx/mujoco/mjx/_src/passive.py

@@ -31,6 +31,7 @@ from mujoco.mjx._src.types import Model
 
 def _spring_damper(m: Model, d: Data) -> jax.Array:
   """Applies joint level spring and damping forces."""
+
   def fn(jnt_typs, stiffness, qpos_spring, qpos):
     qpos_i = 0
     qfrcs = []

mjx/mujoco/mjx/_src/ray.py

@@ -269,7 +269,7 @@ def ray(
   geom_filter_dyn = (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
   geom_filter_dyn &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
   for geom_type, fn in _RAY_FUNC.items():
-    id_, = np.nonzero(geom_filter & (m.geom_type == geom_type))
+    (id_,) = np.nonzero(geom_filter & (m.geom_type == geom_type))
 
     if id_.size == 0:
       continue

mjx/mujoco/mjx/_src/scan.py

@@ -144,13 +144,11 @@ def _check_input(m: Model, args: Any, in_types: str) -> None:
   }
   for idx, (arg, typ) in enumerate(zip(args, in_types)):
     if len(arg) != size[typ]:
-      raise IndexError(
+      raise IndexError((
-          (
+          f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
-              f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
+          f' which does not match the in_types[{idx}] expected length of '
-              f' which does not match the in_types[{idx}] expected length of '
+          f'"{size[typ]}".'
-              f'"{size[typ]}".'
+      ))
-          )
-      )
 
 
 def _check_output(
@@ -158,13 +156,11 @@ def _check_output(
 ) -> None:
   """Checks that scan output has the right shape."""
   if y.shape[0] != take_ids.shape[0]:
-    raise IndexError(
+    raise IndexError((
-        (
+        f'f output "{idx}" with type "{typ}" has shape "{y.shape[0]}" '
-            f'f output "{idx}" with type "{typ}" has shape "{y.shape[0]}" '
+        f'which does not match the out_types[{idx}] expected size of'
-            f'which does not match the out_types[{idx}] expected size of'
+        f' "{take_ids.shape[0]}".'
-            f' "{take_ids.shape[0]}".'
+    ))
-        )
-    )
 
 
 def flat(
@@ -400,7 +396,7 @@ def body_tree(
       if t == 'b':
         continue
       elif t == 'j':
-        key += (tuple(m.jnt_type[np.nonzero(m.jnt_bodyid == id_)[0]]))
+        key += tuple(m.jnt_type[np.nonzero(m.jnt_bodyid == id_)[0]])
       elif t == 'v':
         key += (len(np.nonzero(m.dof_bodyid == id_)[0]),)
       elif t == 'q':

mjx/mujoco/mjx/_src/scan_test.py

@@ -90,6 +90,7 @@ class ScanTest(absltest.TestCase):
       if tuple(jnt_types) == (JointType.FREE,):
         return None
       return val + sum(jnt_types)
+
     b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
     b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in)
     np.testing.assert_equal(np.array(b_out), np.array(b_expect))
@@ -99,6 +100,7 @@ class ScanTest(absltest.TestCase):
       if jnt_types.size == 0:
         self.fail('world has no dofs, should not be called')
       return val + sum(jnt_types)
+
     v_in = jp.ones((m.nv, 1))
     scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in)
 
@@ -141,6 +143,7 @@ class ScanTest(absltest.TestCase):
         return None
       carry = jp.zeros_like(val) if carry is None else carry
       return carry + val + sum(jnt_types)
+
     b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
     b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in)
     np.testing.assert_equal(np.array(b_out), np.array(b_expect))

mjx/mujoco/mjx/_src/sensor_test.py

@@ -17,13 +17,11 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
-
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
 from mujoco.mjx._src.types import ConeType
-
 import numpy as np
 
 # tolerance for difference between MuJoCo and MJX smooth calculations - mostly

mjx/mujoco/mjx/_src/smooth.py

@@ -334,7 +334,7 @@ def factor_m(m: Model, d: Data) -> Data:
     pivots = []
     out = []
 
-    for (b, e, madr_d, madr_ij) in updates:
+    for b, e, madr_d, madr_ij in updates:
       width = e - b
       rows.append(np.arange(madr_ij, madr_ij + width))
       madr_ijs.append(np.full((width,), madr_ij))
@@ -511,7 +511,6 @@ def subtree_vel(m: Model, d: Data) -> Data:
       angmom_child, mom_parent_child = carry
       return angmom + mom + angmom_child + mom_parent_child, mom_parent
 
-
   subtree_angmom, _ = scan.body_tree(
       m,
       _subtree_angmom,
@@ -535,6 +534,7 @@ def subtree_vel(m: Model, d: Data) -> Data:
 
 def rne(m: Model, d: Data) -> Data:
   """Computes inverse dynamics using the recursive Newton-Euler algorithm."""
+
   # forward scan over tree: accumulate link center of mass acceleration
   def cacc_fn(cacc, cdof_dot, qvel):
     if cacc is None:

mjx/mujoco/mjx/_src/solver.py

@@ -52,6 +52,7 @@ class _Context(PyTreeNode):
     u: friction cone (normal and tangents)            (num(con.dim > 1), 6)
     h: cone hessian                                   (num(con.dim > 1), 6, 6)
   """
+
   qacc: jax.Array
   qfrc_constraint: jax.Array
   Jaref: jax.Array  # pylint: disable=invalid-name
@@ -225,6 +226,7 @@ class _LSContext(PyTreeNode):
 
 def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
   """Scan-based implementation (jit ok, reverse-mode autodiff ok)."""
+
   def _iter(val):
     next_val = body_fun(val)
     next_cond = cond_fun(next_val)
@@ -382,7 +384,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
       # set efc of cone H along diagonal
       for i, (condim, addr) in enumerate(zip(dim, efc_address)):
         h_cone = ctx.h[i, :condim, :condim]
-        cm = cm.at[addr:addr+condim, addr:addr+condim].add(h_cone)
+        cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone)
       h = d.efc_J.T @ cm @ d.efc_J
     else:
       h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J

mjx/mujoco/mjx/_src/test_util.py

@@ -116,7 +116,7 @@ def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray:
     if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
       num_rows = (dx.contact.dim[i] - 1) * 2
     if dx.contact.dist[i] > 0:  # move empty contacts to end
-      order = np.append(order, np.repeat(2 ** 16, num_rows))
+      order = np.append(order, np.repeat(2**16, num_rows))
       continue
     contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1)
     contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1)

mjx/mujoco/mjx/_src/types.py

@@ -47,6 +47,7 @@ class DisableBit(enum.IntFlag):
     REFSAFE:      integrator safety: make ref[0]>=2*timestep
     SENSOR:       sensors
   """
+
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
   EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY
   FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS
@@ -73,6 +74,7 @@ class JointType(enum.IntEnum):
     SLIDE: sliding distance along body-fixed axis       (1,)
     HINGE: rotation angle (rad) around body-fixed axis  (1,)
   """
+
   FREE = mujoco.mjtJoint.mjJNT_FREE
   BALL = mujoco.mjtJoint.mjJNT_BALL
   SLIDE = mujoco.mjtJoint.mjJNT_SLIDE
@@ -93,6 +95,7 @@ class IntegratorType(enum.IntEnum):
     RK4: 4th-order Runge Kutta
     IMPLICITFAST: implicit in velocity, no rne derivative
   """
+
   EULER = mujoco.mjtIntegrator.mjINT_EULER
   RK4 = mujoco.mjtIntegrator.mjINT_RK4
   IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
@@ -113,6 +116,7 @@ class GeomType(enum.IntEnum):
     MESH: mesh
     SDF: signed distance field
   """
+
   PLANE = mujoco.mjtGeom.mjGEOM_PLANE
   HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD
   SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE
@@ -134,6 +138,7 @@ class ConvexMesh(PyTreeNode):
     edge: edge indexes for all edges in the convex mesh
     edge_face_normal: indexes for face normals adjacent to edges in `edge`
   """
+
   vert: jax.Array
   face: jax.Array
   face_normal: jax.Array
@@ -148,6 +153,7 @@ class ConeType(enum.IntEnum):
     PYRAMIDAL: pyramidal
     ELLIPTIC: elliptic
   """
+
   PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL
   ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC
 
@@ -160,6 +166,7 @@ class JacobianType(enum.IntEnum):
     SPARSE: sparse
     AUTO: sparse if nv>60 and device is TPU, dense otherwise
   """
+
   DENSE = mujoco.mjtJacobian.mjJAC_DENSE
   SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
   AUTO = mujoco.mjtJacobian.mjJAC_AUTO
@@ -172,6 +179,7 @@ class SolverType(enum.IntEnum):
     CG: Conjugate gradient (primal)
     NEWTON: Newton (primal)
   """
+
   # unsupported: PGS
   CG = mujoco.mjtSolver.mjSOL_CG
   NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
@@ -186,6 +194,7 @@ class EqType(enum.IntEnum):
     JOINT: couple the values of two scalar joints with cubic
     TENDON: couple the lengths of two tendons with cubic
   """
+
   CONNECT = mujoco.mjtEq.mjEQ_CONNECT
   WELD = mujoco.mjtEq.mjEQ_WELD
   JOINT = mujoco.mjtEq.mjEQ_JOINT
@@ -203,6 +212,7 @@ class WrapType(enum.IntEnum):
     SPHERE: wrap around sphere
     CYLINDER: wrap around (infinite) cylinder
   """
+
   JOINT = mujoco.mjtWrap.mjWRAP_JOINT
   PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY
   SITE = mujoco.mjtWrap.mjWRAP_SITE
@@ -219,6 +229,7 @@ class TrnType(enum.IntEnum):
     TENDON: force on tendon
     SITE: force on site
   """
+
   JOINT = mujoco.mjtTrn.mjTRN_JOINT
   JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT
   SITE = mujoco.mjtTrn.mjTRN_SITE
@@ -236,6 +247,7 @@ class DynType(enum.IntEnum):
     FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration
     MUSCLE: piece-wise linear filter with two time constants
   """
+
   NONE = mujoco.mjtDyn.mjDYN_NONE
   INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
   FILTER = mujoco.mjtDyn.mjDYN_FILTER
@@ -252,6 +264,7 @@ class GainType(enum.IntEnum):
     AFFINE: const + kp*length + kv*velocity
     MUSCLE: muscle FLV curve computed by muscle_gain
   """
+
   FIXED = mujoco.mjtGain.mjGAIN_FIXED
   AFFINE = mujoco.mjtGain.mjGAIN_AFFINE
   MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE
@@ -266,6 +279,7 @@ class BiasType(enum.IntEnum):
     AFFINE: const + kp*length + kv*velocity
     MUSCLE: muscle passive force computed by muscle_bias
   """
+
   NONE = mujoco.mjtBias.mjBIAS_NONE
   AFFINE = mujoco.mjtBias.mjBIAS_AFFINE
   MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE
@@ -282,6 +296,7 @@ class ConstraintType(enum.IntEnum):
     CONTACT_FRICTIONLESS: frictionless contact
     CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
   """
+
   EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
   FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF
   FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON
@@ -302,6 +317,7 @@ class CamLightType(enum.IntEnum):
     TARGETBODY: pos fixed in body, rot tracks target body
     TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
   """
+
   FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
   TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
   TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
@@ -346,6 +362,7 @@ class SensorType(enum.IntEnum):
     FRAMELINACC: 3D linear acceleration
     FRAMEANGACC: 3D angular acceleration
   """
+
   MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
   CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION
   RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER
@@ -391,6 +408,7 @@ class ObjType(PyTreeNode):
     SITE: site
     CAMERA: camera
   """
+
   UNKNOWN = mujoco.mjtObj.mjOBJ_UNKNOWN
   BODY = mujoco.mjtObj.mjOBJ_BODY
   XBODY = mujoco.mjtObj.mjOBJ_XBODY
@@ -437,7 +455,7 @@ class Option(PyTreeNode):
     disableactuator:   bit flags for disabling actuators by group id (not used)
     sdf_initpoints:    number of starting points for gradient descent (not used)
     sdf_iterations:    max number of iterations for gradient descent (not used)
-  """
+  """  # fmt: skip
   timestep: jax.Array
   apirate: jax.Array = _restricted_to('mujoco')
   impratio: jax.Array
@@ -480,6 +498,7 @@ class Statistic(PyTreeNode):
     extent: spatial extent (not used)
     center: center of model (not used)
   """
+
   meaninertia: jax.Array
   meanmass: jax.Array
   meansize: jax.Array
@@ -813,6 +832,7 @@ class Model(PyTreeNode):
     name_keyadr: keyframe name pointers                       (nkey,)
     names: names of all objects, 0-terminated                 (nnames,)
   """
+
   nq: int
   nv: int
   nu: int
@@ -1161,7 +1181,7 @@ class Contact(PyTreeNode):
     geom2: id of geom 2; deprecated, use geom[1]
     geom: geom ids                                                    (2,)
     efc_address: address in efc; -1: not included
-  """
+  """  # fmt: skip
   dist: jax.Array
   pos: jax.Array
   frame: jax.Array
@@ -1306,7 +1326,7 @@ class Data(PyTreeNode):
     _qM_sparse: qM in sparse representation                     (nM,)
     _qLD_sparse: qLD in sparse representation                   (nM,)
     _qLDiagInv_sparse: qLDiagInv in sparse representation       (nv,)
-  """
+  """  # fmt: skip
   # constant sizes:
   ne: int
   nf: int

mjx/mujoco/mjx/testspeed.py

@@ -22,7 +22,9 @@ from etils import epath
 import mujoco
 from mujoco import mjx
 
-_MJCF = flags.DEFINE_string('mjcf', None, 'path to model `.xml` or `.mjb`', required=True)
+_MJCF = flags.DEFINE_string(
+    'mjcf', None, 'path to model `.xml` or `.mjb`', required=True
+)
 _BASE_PATH = flags.DEFINE_string(
     'base_path', None, 'base path, defaults to mujoco.mjx resource path'
 )

mjx/mujoco/mjx/viewer.py

@@ -28,8 +28,9 @@ import mujoco.viewer
 
 
 _JIT = flags.DEFINE_bool('jit', True, 'To jit or not to jit.')
-_MODEL_PATH = flags.DEFINE_string('mjcf', None, 'Path to a MuJoCo MJCF file.',
+_MODEL_PATH = flags.DEFINE_string(
-                                  required=True)
+    'mjcf', None, 'Path to a MuJoCo MJCF file.', required=True
+)
 
 
 _VIEWER_GLOBAL_STATE = {

mjx/pyproject.toml

@@ -44,3 +44,24 @@ Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
 Documentation = "https://mujoco.readthedocs.io/en/3.2.7"
 Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
 Changelog = "https://mujoco.readthedocs.io/en/3.2.7/changelog.html"
+
+[tool.isort]
+force_single_line = true
+force_sort_within_sections = true
+lexicographical = true
+single_line_exclusions = ["typing"]
+order_by_type = false
+group_by_package = true
+line_length = 120
+use_parentheses = true
+multi_line_output = 3
+skip_glob = ["**/*.ipynb"]
+
+[tool.pyink]
+line-length = 80
+unstable = true
+pyink-indentation = 2
+pyink-use-majority-quotes = true
+extend-exclude = '''(
+ .ipynb$
+)'''

python/mujoco/bindings_test.py

@@ -120,13 +120,17 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     xml_2 = rb"""<mujoco><geom name="box" type="box" size="1 1 1"/></mujoco>"""
     xml_3 = rb"""<mujoco><geom name="ball" type="sphere" size="1"/></mujoco>"""
     model = mujoco.MjModel.from_xml_string(
-        xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3})
+        xml_1, {'model_.xml': xml_2, 'model__.xml': xml_3}
+    )
     self.assertEqual(
-        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0)
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'plane'), 0
+    )
     self.assertEqual(
-        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1)
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'box'), 1
+    )
     self.assertEqual(
-        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2)
+        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, 'ball'), 2
+    )
 
   def test_load_xml_repeated_asset_name(self):
     # Assets aren't allowed to have the same filename (even if they have
@@ -139,23 +143,25 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_can_read_array(self):
     np.testing.assert_array_equal(
         self.model.body_pos,
-        [[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]])
+        [[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]],
+    )
 
   def test_can_set_array(self):
     self.data.qpos = 0.12345
     np.testing.assert_array_equal(
-        self.data.qpos, [0.12345]*len(self.data.qpos))
+        self.data.qpos, [0.12345] * len(self.data.qpos)
+    )
 
   def test_array_is_a_view(self):
     qpos_ref = self.data.qpos
     self.data.qpos = 0.789
-    np.testing.assert_array_equal(
+    np.testing.assert_array_equal(qpos_ref, [0.789] * len(self.data.qpos))
-        qpos_ref, [0.789]*len(self.data.qpos))
 
   # This test is disabled on PyPy as it uses sys.getrefcount
   # However PyPy is not officially supported by MuJoCo
-  @absltest.skipIf(sys.implementation.name == 'pypy',
+  @absltest.skipIf(
-                   reason='requires sys.getrefcount')
+      sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
+  )
   def test_array_keeps_struct_alive(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
     qpos0 = model.qpos0
@@ -185,11 +191,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
   def test_named_indexing_actuator_ctrl(self):
     actuator_id = mujoco.mj_name2id(
-        self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator')
+        self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator'
-    self.assertIs(self.data.actuator('myactuator'),
+    )
-                  self.data.actuator(actuator_id))
+    self.assertIs(
-    self.assertIs(self.data.actuator('myactuator').ctrl,
+        self.data.actuator('myactuator'), self.data.actuator(actuator_id)
-                  self.data.actuator(actuator_id).ctrl)
+    )
+    self.assertIs(
+        self.data.actuator('myactuator').ctrl,
+        self.data.actuator(actuator_id).ctrl,
+    )
     self.assertEqual(self.data.actuator('myactuator').ctrl.shape, (1,))
 
     # Test that the indexer is returning a view into the underlying struct.
@@ -202,41 +212,49 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_named_indexing_invalid_names_in_model(self):
     with self.assertRaisesRegex(
         KeyError,
-        r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]"):
+        r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]",
+    ):
       self.model.geom('badgeom')
 
   def test_named_indexing_no_name_argument_in_model(self):
     with self.assertRaisesRegex(
         KeyError,
-        r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]"):
+        r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]",
+    ):
       self.model.joint()
 
   def test_named_indexing_invalid_names_in_data(self):
     with self.assertRaisesRegex(
         KeyError,
-        r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]"):
+        r"Invalid name 'badgeom'\. Valid names: \['mybox', 'myplane'\]",
+    ):
       self.data.geom('badgeom')
 
   def test_named_indexing_no_name_argument_in_data(self):
     with self.assertRaisesRegex(
         KeyError,
-        r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]"):
+        r"Invalid name ''\. Valid names: \['myball', 'myfree', 'myhinge'\]",
+    ):
       self.data.jnt()
 
   def test_named_indexing_invalid_index_in_model(self):
     with self.assertRaisesRegex(
-        IndexError, r'Invalid index 3\. Valid indices from 0 to 2'):
+        IndexError, r'Invalid index 3\. Valid indices from 0 to 2'
+    ):
       self.model.geom(3)
     with self.assertRaisesRegex(
-        IndexError, r'Invalid index -1\. Valid indices from 0 to 2'):
+        IndexError, r'Invalid index -1\. Valid indices from 0 to 2'
+    ):
       self.model.geom(-1)
 
   def test_named_indexing_invalid_index_in_data(self):
     with self.assertRaisesRegex(
-        IndexError, r'Invalid index 3\. Valid indices from 0 to 2'):
+        IndexError, r'Invalid index 3\. Valid indices from 0 to 2'
+    ):
       self.data.geom(3)
     with self.assertRaisesRegex(
-        IndexError, r'Invalid index -1\. Valid indices from 0 to 2'):
+        IndexError, r'Invalid index -1\. Valid indices from 0 to 2'
+    ):
       self.data.geom(-1)
 
   def test_named_indexing_geom_size(self):
@@ -267,45 +285,53 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
   def test_named_indexing_ragged_qpos(self):
     balljoint_id = mujoco.mj_name2id(
-        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball')
+        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myball'
+    )
     self.assertIs(self.data.joint('myball'), self.data.joint(balljoint_id))
-    self.assertIs(self.data.joint('myball').qpos,
+    self.assertIs(
-                  self.data.joint(balljoint_id).qpos)
+        self.data.joint('myball').qpos, self.data.joint(balljoint_id).qpos
+    )
     self.assertEqual(self.data.joint('myball').qpos.shape, (4,))
 
     # Test that the indexer is returning a view into the underlying struct.
     qpos_from_indexer = self.data.joint('myball').qpos
     qpos_idx = self.model.jnt_qposadr[balljoint_id]
-    self.data.qpos[qpos_idx:qpos_idx+4] = [4, 5, 6, 7]
+    self.data.qpos[qpos_idx : qpos_idx + 4] = [4, 5, 6, 7]
     np.testing.assert_array_equal(qpos_from_indexer, [4, 5, 6, 7])
     self.data.joint('myball').qpos = [9, 8, 7, 6]
-    np.testing.assert_array_equal(self.data.qpos[qpos_idx:qpos_idx+4],
+    np.testing.assert_array_equal(
-                                  [9, 8, 7, 6])
+        self.data.qpos[qpos_idx : qpos_idx + 4], [9, 8, 7, 6]
+    )
 
   def test_named_indexing_ragged2d_cdof(self):
     freejoint_id = mujoco.mj_name2id(
-        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree')
+        self.model, mujoco.mjtObj.mjOBJ_JOINT, 'myfree'
+    )
     self.assertIs(self.data.joint('myfree'), self.data.joint(freejoint_id))
-    self.assertIs(self.data.joint('myfree').cdof,
+    self.assertIs(
-                  self.data.joint(freejoint_id).cdof)
+        self.data.joint('myfree').cdof, self.data.joint(freejoint_id).cdof
+    )
     self.assertEqual(self.data.joint('myfree').cdof.shape, (6, 6))
 
     # Test that the indexer is returning a view into the underlying struct.
     cdof_from_indexer = self.data.joint('myfree').cdof
     dof_idx = self.model.jnt_dofadr[freejoint_id]
-    self.data.cdof[dof_idx:dof_idx+6, :] = np.reshape(range(36), (6, 6))
+    self.data.cdof[dof_idx : dof_idx + 6, :] = np.reshape(range(36), (6, 6))
-    np.testing.assert_array_equal(cdof_from_indexer,
+    np.testing.assert_array_equal(
-                                  np.reshape(range(36), (6, 6)))
+        cdof_from_indexer, np.reshape(range(36), (6, 6))
+    )
     self.data.joint('myfree').cdof = 42
-    np.testing.assert_array_equal(self.data.cdof[dof_idx:dof_idx+6], [[42]*6]*6)
+    np.testing.assert_array_equal(
+        self.data.cdof[dof_idx : dof_idx + 6], [[42] * 6] * 6
+    )
 
   def test_named_indexing_repr_in_data(self):
-    expected_repr = '''<_MjDataGeomViews
+    expected_repr = """<_MjDataGeomViews
   id: 1
   name: 'mybox'
   xmat: array([0., 0., 0., 0., 0., 0., 0., 0., 0.])
   xpos: array([0., 0., 0.])
->'''
+>"""
     self.assertEqual(expected_repr, repr(self.data.geom('mybox')))
 
   def test_named_indexing_body_repr_in_data(self):
@@ -328,8 +354,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertGreater(self.data._address, 0)
     self.assertGreater(model2._address, 0)
     self.assertGreater(data2._address, 0)
-    self.assertLen({self.model._address, self.data._address,
+    self.assertLen(
-                    model2._address, data2._address}, 4)
+        {
+            self.model._address,
+            self.data._address,
+            model2._address,
+            data2._address,
+        },
+        4,
+    )
 
   def test_mjmodel_can_read_and_write_opt(self):
     self.assertEqual(self.model.opt.timestep, 0.002)
@@ -361,7 +394,9 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_mjmodel_can_access_names_directly(self):
     # mjModel offers direct access to names array, to allow usecases other than
     # id2name
-    model_name = str(self.model.names[0:self.model.names.find(b'\0')], 'utf-8')
+    model_name = str(
+        self.model.names[0 : self.model.names.find(b'\0')], 'utf-8'
+    )
     self.assertEqual(model_name, 'test')
 
     start_index = self.model.name_geomadr[0]
@@ -402,15 +437,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     model_copy = copy.copy(self.model)
 
     self.assertEqual(
-        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0),
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0), 'myfree'
-        'myfree')
+    )
 
     self.assertEqual(
-        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0),
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0), 'myplane'
-        'myplane')
+    )
     self.assertEqual(
-        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1),
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1), 'mybox'
-        'mybox')
+    )
 
     # Make sure it's a copy.
     self.model.geom_size[1] = 0.5
@@ -420,7 +455,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
   def test_mjdata_can_copy(self):
     self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001,
                       np.cos(np.pi/8), np.sin(np.pi/8), 0, 0, 0,
-                      1, 0, 0, 0]
+                      1, 0, 0, 0]  # fmt: skip
     mujoco.mj_forward(self.model, self.data)
 
     data_copy = copy.copy(self.data)
@@ -455,7 +490,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
       contact_copy.append(copy.copy(self.data.contact[i]))
     # Sort contacts in anticlockwise order
     contact_copy = sorted(
-        contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
+        contact_copy, key=lambda x: np.arctan2(x.pos[1], x.pos[0])
+    )
     np.testing.assert_allclose(contact_copy[0].pos[:2], [-0.1, -0.1])
     np.testing.assert_allclose(contact_copy[1].pos[:2], [0.1, -0.1])
     np.testing.assert_allclose(contact_copy[2].pos[:2], [0.1, 0.1])
@@ -502,7 +538,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     # Sort contacts in anticlockwise order
     sorted_contact = sorted(
-        contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0]))
+        contact, key=lambda x: np.arctan2(x.pos[1], x.pos[0])
+    )
     np.testing.assert_allclose(sorted_contact[0].pos[:2], [-0.1, -0.1])
     np.testing.assert_allclose(sorted_contact[1].pos[:2], [0.1, -0.1])
     np.testing.assert_allclose(sorted_contact[2].pos[:2], [0.1, 0.1])
@@ -589,7 +626,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertEqual(data2.ncon, 4)
     self.assertEqual(data2.contact, self.data.contact)
 
-    self.data.qpos[3:7] = [np.cos(np.pi/8), np.sin(np.pi/8), 0, 0]
+    self.data.qpos[3:7] = [np.cos(np.pi / 8), np.sin(np.pi / 8), 0, 0]
     self.data.qpos[2] *= (np.sqrt(2) - 1) * 0.1 - 1e-6
     mujoco.mj_forward(self.model, self.data)
     self.assertEqual(self.data.ncon, 2)
@@ -674,7 +711,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
   def test_mju_rotVecQuat(self):  # pylint: disable=invalid-name
     vec = [1, 0, 0]
-    quat = [np.cos(np.pi/8), 0, 0, np.sin(np.pi/8)]
+    quat = [np.cos(np.pi / 8), 0, 0, np.sin(np.pi / 8)]
     expected = np.array([1, 1, 0]) / np.sqrt(2)
 
     # Check that the output argument works, and that the binding returns None.
@@ -722,7 +759,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     size = mujoco.mj_stateSize(self.model, spec)
 
     state_bad_size = np.empty(size + 1, np.float64)
-    expected_message = ('state size should equal mj_stateSize(m, spec)')
+    expected_message = 'state size should equal mj_stateSize(m, spec)'
     with self.assertRaisesWithLiteralMatch(TypeError, expected_message):
       mujoco.mj_getState(self.model, self.data, state_bad_size, spec)
 
@@ -781,8 +818,9 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     mat = np.empty((3, 10), np.float64)
     mujoco.mj_angmomMat(self.model, self.data, mat, 0)
-    np.testing.assert_almost_equal(mat @ self.data.qvel,
+    np.testing.assert_almost_equal(
-                                   self.data.subtree_angmom[0, :])
+        mat @ self.data.qvel, self.data.subtree_angmom[0, :]
+    )
 
   def test_mj_jacSite(self):  # pylint: disable=invalid-name
     mujoco.mj_forward(self.model, self.data)
@@ -792,20 +830,22 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     jacp = np.empty((3, 10), np.float64)
     mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id)
 
-    expected_jacp = np.array(
+    expected_jacp = np.array([
-        [[0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, -1, 0, 0, 0],
-         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    ])
     np.testing.assert_array_equal(jacp, expected_jacp)
 
     # Call mj_jacSite with only jacr.
     jacr = np.empty((3, 10), np.float64)
     mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id)
 
-    expected_jacr = np.array(
+    expected_jacr = np.array([
-        [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-         [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
-         [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
+        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    ])
     np.testing.assert_array_equal(jacr, expected_jacr)
 
     # Call mj_jacSite with both jacp and jacr.
@@ -818,12 +858,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     # Check that the jacp argument must have the right size.
     with self.assertRaises(TypeError):
       mujoco.mj_jacSite(
-          self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id)
+          self.model, self.data, np.empty((3, 6), jacp.dtype), None, site_id
+      )
 
     # Check that the jacr argument must have the right size.
     with self.assertRaises(TypeError):
       mujoco.mj_jacSite(
-          self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id)
+          self.model, self.data, None, np.empty((4, 7), jacr.dtype), site_id
+      )
 
     # The following two checks need to be done with fully initialized arrays,
     # since pybind11 prints out the array's contents when generating TypeErrors.
@@ -832,12 +874,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     # Check that the jacp argument must have the right dtype.
     with self.assertRaises(TypeError):
       mujoco.mj_jacSite(
-          self.model, self.data, np.zeros(jacp.shape, int), None, site_id)
+          self.model, self.data, np.zeros(jacp.shape, int), None, site_id
+      )
 
     # Check that the jacr argument must have the right dtype.
     with self.assertRaises(TypeError):
       mujoco.mj_jacSite(
-          self.model, self.data, None, np.zeros(jacr.shape, int), site_id)
+          self.model, self.data, None, np.zeros(jacr.shape, int), site_id
+      )
 
   def test_docstrings(self):  # pylint: disable=invalid-name
     self.assertEqual(
@@ -845,13 +889,15 @@ class MuJoCoBindingsTest(parameterized.TestCase):
         """mj_versionString() -> str
 
 Return the current version of MuJoCo as a null-terminated string.
-""")
+""",
+    )
     self.assertEqual(
         mujoco.mj_Euler.__doc__,
         """mj_Euler(m: mujoco._structs.MjModel, d: mujoco._structs.MjData) -> None
 
 Euler integrator, semi-implicit in velocity.
-""")
+""",
+    )
 
   def test_float_constant(self):
     self.assertEqual(mujoco.mjMAXVAL, 1e10)
@@ -866,17 +912,19 @@ Euler integrator, semi-implicit in velocity.
     self.assertLen(mujoco.mjVISSTRING, mujoco.mjtVisFlag.mjNVISFLAG)
     self.assertLen(mujoco.mjRNDSTRING, mujoco.mjtRndFlag.mjNRNDFLAG)
     self.assertEqual(mujoco.mjDISABLESTRING[11], 'Refsafe')
-    self.assertEqual(mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
+    self.assertEqual(
-                     ('Inertia', '0', 'I'))
+        mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
+        ('Inertia', '0', 'I'),
+    )
 
   def test_enum_values(self):
     self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, 0)
     self.assertEqual(mujoco.mjtJoint.mjJNT_BALL, 1)
     self.assertEqual(mujoco.mjtJoint.mjJNT_SLIDE, 2)
     self.assertEqual(mujoco.mjtJoint.mjJNT_HINGE, 3)
-    self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1<<0)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1 << 0)
-    self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1<<1)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1 << 1)
-    self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1 << 2)
     self.assertEqual(mujoco.mjtEnableBit.mjNENABLE, 7)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_PLANE, 0)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_HFIELD, 1)
@@ -899,8 +947,9 @@ Euler integrator, semi-implicit in velocity.
     x = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k']
     self.assertEqual(x[mujoco.mjtFrame.mjFRAME_WORLD], 'h')
     self.assertEqual(
-        x[mujoco.mjtFrame.mjFRAME_GEOM:mujoco.mjtFrame.mjFRAME_CAMERA],
+        x[mujoco.mjtFrame.mjFRAME_GEOM : mujoco.mjtFrame.mjFRAME_CAMERA],
-        ['c', 'd'])
+        ['c', 'd'],
+    )
 
   def test_enum_ops(self):
     # Note: when modifying this test, make sure the enum value is an odd number
@@ -909,10 +958,12 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(mujoco.mjtFrame.mjFRAME_WORLD, 7.0)
     self.assertEqual(7, mujoco.mjtFrame.mjFRAME_WORLD)
     self.assertEqual(7.0, mujoco.mjtFrame.mjFRAME_WORLD)
-    self.assertEqual(mujoco.mjtFrame.mjFRAME_WORLD,
+    self.assertEqual(
-                     mujoco.mjtFrame.mjFRAME_WORLD)
+        mujoco.mjtFrame.mjFRAME_WORLD, mujoco.mjtFrame.mjFRAME_WORLD
-    self.assertNotEqual(mujoco.mjtFrame.mjFRAME_WORLD,
+    )
-                        mujoco.mjtFrame.mjFRAME_NONE)
+    self.assertNotEqual(
+        mujoco.mjtFrame.mjFRAME_WORLD, mujoco.mjtFrame.mjFRAME_NONE
+    )
 
     self.assertEqual(-mujoco.mjtFrame.mjFRAME_WORLD, -7)
     self.assertIsInstance(-mujoco.mjtFrame.mjFRAME_WORLD, int)
@@ -989,22 +1040,28 @@ Euler integrator, semi-implicit in velocity.
 
     self.assertEqual(
         mujoco.mjtDisableBit.mjDSBL_GRAVITY | mujoco.mjtDisableBit.mjDSBL_LIMIT,
-        72)
+        72,
+    )
     self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE | 33, 33)
     self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE & 33, 32)
     self.assertEqual(mujoco.mjtDisableBit.mjDSBL_PASSIVE ^ 33, 1)
     self.assertEqual(33 | mujoco.mjtDisableBit.mjDSBL_PASSIVE, 33)
     self.assertEqual(33 & mujoco.mjtDisableBit.mjDSBL_PASSIVE, 32)
     self.assertEqual(33 ^ mujoco.mjtDisableBit.mjDSBL_PASSIVE, 1)
-    self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1,
+    self.assertEqual(
-                     mujoco.mjtDisableBit.mjDSBL_WARMSTART)
+        mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1,
-    self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3,
+        mujoco.mjtDisableBit.mjDSBL_WARMSTART,
-                     mujoco.mjtDisableBit.mjDSBL_CONTACT)
+    )
+    self.assertEqual(
+        mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3,
+        mujoco.mjtDisableBit.mjDSBL_CONTACT,
+    )
 
   def test_can_raise_error(self):
     self.data.pstack = self.data.narena
-    with self.assertRaisesRegex(mujoco.FatalError,
+    with self.assertRaisesRegex(
-                                r'\Amj_stackAlloc: insufficient memory:'):
+        mujoco.FatalError, r'\Amj_stackAlloc: insufficient memory:'
+    ):
       mujoco.mj_forward(self.model, self.data)
 
   def test_mjcb_time(self):
@@ -1042,7 +1099,8 @@ Euler integrator, semi-implicit in velocity.
       with self.assertRaises(TestError) as e:
         mujoco.mj_forward(self.model, self.data)
       self.assertEqual(
-          e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2}))
+          e.exception.args, ('string', (1, 2, 3), {'a': 1, 'b': 2})
+      )
 
     # Should not raise now that we've cleared the callback.
     mujoco.mj_forward(self.model, self.data)
@@ -1050,12 +1108,14 @@ Euler integrator, semi-implicit in velocity.
   def test_mjcb_time_wrong_return_type(self):
     with temporary_callback(mujoco.set_mjcb_time, lambda: 'string'):
       with self.assertRaisesWithLiteralMatch(
-          TypeError, 'mjcb_time callback did not return a number'):
+          TypeError, 'mjcb_time callback did not return a number'
+      ):
         mujoco.mj_forward(self.model, self.data)
 
   def test_mjcb_time_not_callable(self):
     with self.assertRaisesWithLiteralMatch(
-        TypeError, 'callback is not an Optional[Callable]'):
+        TypeError, 'callback is not an Optional[Callable]'
+    ):
       mujoco.set_mjcb_time(1)
 
   def test_mjcb_sensor(self):
@@ -1088,8 +1148,9 @@ Euler integrator, semi-implicit in velocity.
 
   # This test is disabled on PyPy as it uses sys.getrefcount
   # However PyPy is not officially supported by MuJoCo
-  @absltest.skipIf(sys.implementation.name == 'pypy',
+  @absltest.skipIf(
-                   reason='requires sys.getrefcount')
+      sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
+  )
   def test_mjcb_control_not_leak_memory(self):
     model_instances = []
     data_instances = []
@@ -1110,8 +1171,9 @@ Euler integrator, semi-implicit in velocity.
 
   # This test is disabled on PyPy as it uses sys.getrefcount
   # However PyPy is not officially supported by MuJoCo
-  @absltest.skipIf(sys.implementation.name == 'pypy',
+  @absltest.skipIf(
-                   reason='requires sys.getrefcount')
+      sys.implementation.name == 'pypy', reason='requires sys.getrefcount'
+  )
   def test_mjdata_holds_ref_to_model(self):
     data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
     model = data.model
@@ -1150,9 +1212,15 @@ Euler integrator, semi-implicit in velocity.
     # When the scene is updated, geoms are added to the scene
     # (ngeom is incremented)
     mujoco.mj_forward(self.model, self.data)
-    mujoco.mjv_updateScene(self.model, self.data, mujoco.MjvOption(),
+    mujoco.mjv_updateScene(
-                           None, mujoco.MjvCamera(),
+        self.model,
-                           mujoco.mjtCatBit.mjCAT_ALL, scene)
+        self.data,
+        mujoco.MjvOption(),
+        None,
+        mujoco.MjvCamera(),
+        mujoco.mjtCatBit.mjCAT_ALL,
+        scene,
+    )
     self.assertGreater(scene.ngeom, 0)
 
   def test_mjv_scene_without_model(self):
@@ -1164,10 +1232,19 @@ Euler integrator, semi-implicit in velocity.
     # mj_ray has tricky argument types
     geomid = np.zeros(1, np.int32)
     mujoco.mj_forward(self.model, self.data)
-    mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0,
+    mujoco.mj_ray(
-                  geomid)
+        self.model, self.data, [0, 0, 0], [0, 0, 1], None, 0, 0, geomid
-    mujoco.mj_ray(self.model, self.data, [0, 0, 0], [0, 0, 1],
+    )
-                  [0, 0, 0, 0, 0, 0], 0, 0, geomid)
+    mujoco.mj_ray(
+        self.model,
+        self.data,
+        [0, 0, 0],
+        [0, 0, 1],
+        [0, 0, 0, 0, 0, 0],
+        0,
+        0,
+        geomid,
+    )
     # Check that named arguments work
     mujoco.mj_ray(
         m=self.model,
@@ -1177,7 +1254,8 @@ Euler integrator, semi-implicit in velocity.
         geomgroup=None,
         flg_static=0,
         bodyexclude=0,
-        geomid=geomid)
+        geomid=geomid,
+    )
 
   def test_mj_multi_ray(self):
     nray = 3
@@ -1201,14 +1279,13 @@ Euler integrator, semi-implicit in velocity.
         geomid=geomid,
         dist=dist,
         nray=nray,
-        cutoff=mujoco.mjMAXVAL)
+        cutoff=mujoco.mjMAXVAL,
+    )
 
     for i in range(0, 3):
       self.assertEqual(
           dist[i],
-          mujoco.mj_ray(
+          mujoco.mj_ray(self.model, self.data, pnt, vec[i], None, 1, -1, geom1),
-              self.model, self.data, pnt, vec[i], None, 1, -1, geom1
-          ),
       )
       self.assertEqual(geomid[i], geom1)
       self.assertEqual(geomid[i], geom_ex[i])
@@ -1217,16 +1294,28 @@ Euler integrator, semi-implicit in velocity.
   def test_inverse_fd_none(self):
     eps = 1e-6
     flg_centered = 0
-    mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered,
+    mujoco.mjd_inverseFD(
-                         None, None, None, None, None, None, None)
+        self.model,
+        self.data,
+        eps,
+        flg_centered,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
 
   def test_geom_distance(self):
     mujoco.mj_forward(self.model, self.data)
     fromto = np.empty(6, np.float64)
     dist = mujoco.mj_geomDistance(self.model, self.data, 0, 2, 200, fromto)
     self.assertEqual(dist, 41.9)
-    np.testing.assert_array_equal(fromto,
+    np.testing.assert_array_equal(
-                                  np.array((42., 0., 0., 42., 0., 41.9)))
+        fromto, np.array((42.0, 0.0, 0.0, 42.0, 0.0, 41.9))
+    )
 
   def test_inverse_fd(self):
     eps = 1e-6
@@ -1238,8 +1327,19 @@ Euler integrator, semi-implicit in velocity.
     ds_dv = np.zeros((self.model.nv, self.model.nsensordata))
     ds_da = np.zeros((self.model.nv, self.model.nsensordata))
     dm_dq = np.zeros((self.model.nv, self.model.nM))
-    mujoco.mjd_inverseFD(self.model, self.data, eps, flg_centered,
+    mujoco.mjd_inverseFD(
-                         df_dq, df_dv, df_da, ds_dq, ds_dv, ds_da, dm_dq)
+        self.model,
+        self.data,
+        eps,
+        flg_centered,
+        df_dq,
+        df_dv,
+        df_da,
+        ds_dq,
+        ds_dv,
+        ds_da,
+        dm_dq,
+    )
     self.assertGreater(np.linalg.norm(df_dq), eps)
     self.assertGreater(np.linalg.norm(df_dv), eps)
     self.assertGreater(np.linalg.norm(df_da), eps)
@@ -1272,15 +1372,17 @@ Euler integrator, semi-implicit in velocity.
     n_total = 4
     n_band = 1
     n_dense = 1
-    dense = np.array([[1.0, 0, 0, 0.1],
+    dense = np.array([
-                      [0, 2.0, 0, 0.2],
+        [1.0, 0, 0, 0.1],
-                      [0, 0, 3.0, 0.3],
+        [0, 2.0, 0, 0.2],
-                      [0.1, 0.2, 0.3, 4.0]])
+        [0, 0, 3.0, 0.3],
-    band = np.zeros(n_band*(n_total-n_dense) + n_dense*n_total)
+        [0.1, 0.2, 0.3, 4.0],
+    ])
+    band = np.zeros(n_band * (n_total - n_dense) + n_dense * n_total)
     mujoco.mju_dense2Band(band, dense, n_total, n_band, n_dense)
     for i in range(4):
       index = mujoco.mju_bandDiag(i, n_total, n_band, n_dense)
-      self.assertEqual(band[index], i+1)
+      self.assertEqual(band[index], i + 1)
     dense2 = np.zeros((n_total, n_total))
     flg_sym = 1
     mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym)
@@ -1288,20 +1390,22 @@ Euler integrator, semi-implicit in velocity.
     vec = np.array([[2.0], [2.0], [3.0], [4.0]])
     res = np.zeros_like(vec)
     n_vec = 1
-    mujoco.mju_bandMulMatVec(res, band, vec,
+    mujoco.mju_bandMulMatVec(
-                             n_total, n_band, n_dense, n_vec, flg_sym)
+        res, band, vec, n_total, n_band, n_dense, n_vec, flg_sym
+    )
     np.testing.assert_array_equal(res, dense @ vec)
     diag_add = 0
     diag_mul = 0
-    mujoco.mju_cholFactorBand(band, n_total, n_band, n_dense,
+    mujoco.mju_cholFactorBand(
-                              diag_add, diag_mul)
+        band, n_total, n_band, n_dense, diag_add, diag_mul
+    )
     mujoco.mju_cholSolveBand(res, band, vec, n_total, n_band, n_dense)
     np.testing.assert_almost_equal(res, np.linalg.solve(dense, vec))
 
   def test_mju_box_qp(self):
     n = 5
     res = np.zeros(n)
-    r = np.zeros((n, n+7))
+    r = np.zeros((n, n + 7))
     index = np.zeros(n, np.int32)
     h = np.eye(n)
     g = np.ones((n,))
@@ -1324,7 +1428,7 @@ Euler integrator, semi-implicit in velocity.
     mat = np.linspace(0, 1, 16).reshape(4, 4)
     res = np.empty((4, 4), np.float64)
     mujoco.mju_symmetrize(res, mat)
-    np.testing.assert_array_equal(res, 0.5*(mat + mat.T))
+    np.testing.assert_array_equal(res, 0.5 * (mat + mat.T))
 
   def test_mju_clip(self):
     self.assertEqual(mujoco.mju_clip(1.5, 1.0, 2.0), 1.5)
@@ -1332,14 +1436,14 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(mujoco.mju_clip(1.5, 0.0, 1.0), 1.0)
 
   def test_mju_mul_vec_mat_vec(self):
-    vec1 = np.array([1., 2., 3.])
+    vec1 = np.array([1.0, 2.0, 3.0])
-    vec2 = np.array([3., 2., 1.])
+    vec2 = np.array([3.0, 2.0, 1.0])
-    mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
+    mat = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
-    self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)
+    self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.0)
 
   def test_mju_dense_to_sparse(self):
-    mat = np.array([[0., 1., 0.], [2., 0., 3.]])
+    mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
-    expected_vals = np.array([1., 2., 3.])
+    expected_vals = np.array([1.0, 2.0, 3.0])
     expected_rownnz = np.array([1, 2])
     expected_rowadr = np.array([0, 1])
     expected_colind = np.array([1, 0, 2])
@@ -1355,8 +1459,8 @@ Euler integrator, semi-implicit in velocity.
     np.testing.assert_array_equal(col_ind, expected_colind)
 
   def test_mju_sparse_to_dense(self):
-    expected = np.array([[0., 1., 0.], [2., 0., 3.]])
+    expected = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
-    mat = np.array((1., 2., 3.))
+    mat = np.array((1.0, 2.0, 3.0))
     rownnz = np.array([1, 2])
     rowadr = np.array([0, 1])
     colind = np.array([1, 0, 2])
@@ -1366,10 +1470,10 @@ Euler integrator, semi-implicit in velocity.
 
   def test_mju_euler_to_quat(self):
     quat = np.zeros(4)
-    euler = np.array([0, np.pi/2, 0])
+    euler = np.array([0, np.pi / 2, 0])
     seq = 'xyz'
     mujoco.mju_euler2Quat(quat, euler, seq)
-    expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.])
+    expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.0])
     np.testing.assert_almost_equal(quat, expected_quat)
 
     error = 'mju_euler2Quat: seq must contain exactly 3 characters'
@@ -1377,7 +1481,7 @@ Euler integrator, semi-implicit in velocity.
       mujoco.mju_euler2Quat(quat, euler, 'xy')
     with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error):
       mujoco.mju_euler2Quat(quat, euler, 'xyzy')
-    error = 'mju_euler2Quat: seq[2] is \'p\', should be one of x, y, z, X, Y, Z'
+    error = "mju_euler2Quat: seq[2] is 'p', should be one of x, y, z, X, Y, Z"
     with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error):
       mujoco.mju_euler2Quat(quat, euler, 'xYp')
 
@@ -1396,8 +1500,16 @@ Euler integrator, semi-implicit in velocity.
     mujoco.mj_step(self.model, self.data)
     data2 = pickle.loads(pickle.dumps(self.data))
     attr_to_compare = (
-        'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
+        'time',
-        'warning', 'energy', 'contact', 'efc_J'
+        'qpos',
+        'qvel',
+        'qacc',
+        'xpos',
+        'mocap_pos',
+        'warning',
+        'energy',
+        'contact',
+        'efc_J',
     )
     self._assert_attributes_equal(data2, self.data, attr_to_compare)
     for _ in range(10):
@@ -1410,8 +1522,16 @@ Euler integrator, semi-implicit in velocity.
     mujoco.mj_step(self.model, self.data)
     data2 = pickle.loads(pickle.dumps(self.data))
     attr_to_compare = (
-        'time', 'qpos', 'qvel', 'qacc', 'xpos', 'mocap_pos',
+        'time',
-        'warning', 'energy', 'contact', 'efc_J'
+        'qpos',
+        'qvel',
+        'qacc',
+        'xpos',
+        'mocap_pos',
+        'warning',
+        'energy',
+        'contact',
+        'efc_J',
     )
     self._assert_attributes_equal(data2, self.data, attr_to_compare)
     for _ in range(10):
@@ -1422,7 +1542,10 @@ Euler integrator, semi-implicit in velocity.
   def test_pickle_mjmodel(self):
     model2 = pickle.loads(pickle.dumps(self.model))
     attr_to_compare = (
-        'nq', 'nmat', 'body_pos', 'names',
+        'nq',
+        'nmat',
+        'body_pos',
+        'names',
     )
     self._assert_attributes_equal(model2, self.model, attr_to_compare)
 
@@ -1506,8 +1629,11 @@ Euler integrator, semi-implicit in velocity.
         else:
           self.assertEqual(actual_value, expected_value)
       except AssertionError as e:
-        self.fail("Attribute '{}' differs from expected value: {}".format(
+        self.fail(
-            name, str(e)))
+            "Attribute '{}' differs from expected value: {}".format(
+                name, str(e)
+            )
+        )
 
 
 if __name__ == '__main__':

python/mujoco/memory_leak_test.py

@@ -57,6 +57,7 @@ class MemoryLeakTest(absltest.TestCase):
     soft = -1
     try:
       import resource  # pylint: disable=g-import-not-at-top
+
       soft, hard = resource.getrlimit(resource.RLIMIT_AS)
       resource.setrlimit(resource.RLIMIT_AS, (limit_in_bytes, hard))
     except (ImportError, ValueError):
@@ -65,5 +66,5 @@ class MemoryLeakTest(absltest.TestCase):
     return soft
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
   absltest.main()

python/mujoco/minimize.py

@@ -209,7 +209,7 @@ def least_squares(
 
   # Decrease mu agressively: sequential decreases grow exponentially.
   def decrease_mu(mu, n_reduc):
-    dmu = (1/mu_factor) ** (2**n_reduc)
+    dmu = (1 / mu_factor) ** (2**n_reduc)
     mu = 0.0 if mu * dmu < mu_min else mu * dmu
     n_reduc += 1
     return mu, n_reduc
@@ -427,7 +427,6 @@ def jacobian_fd(
   Returns:
     jac: Jacobian of the residual at x.
     n_res: updated number of residual evaluations (add x.size).
-
   """
   n = x.size
   if bounds is None:
@@ -438,7 +437,7 @@ def jacobian_fd(
   xh = x + np.diag(eps_vec)
   rh = residual(xh)
   jac = (rh - r) / eps_vec
-  return jac, n_res+n
+  return jac, n_res + n
 
 
 def check_jacobian(
@@ -467,14 +466,15 @@ def check_jacobian(
 
   Returns:
     n_res: updated number of residual evaluations.
-
   """
   jac_fd, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds)
   denom = np.abs(jac).sum() + np.abs(jac_fd).sum() + 1e-8
   rel_diff = np.abs(jac - jac_fd) / denom
   if np.any(rel_diff > 1e-5):
-    raise ValueError(f'User-provided {name} does not match finite-differences '
+    raise ValueError(
-                     'to a relative tolerance of 1e-5.')
+        f'User-provided {name} does not match finite-differences '
+        'to a relative tolerance of 1e-5.'
+    )
   print(f'User-provided {name} matches finite-differences.', file=output)
   return n_res
 
@@ -489,8 +489,8 @@ def check_norm(
 
   Args:
     r: residual vector.
-    norm: Norm function returning either the norm scalar or its gradient
+    norm: Norm function returning either the norm scalar or its gradient and
-      and Gauss-Newton Hessian.
+      Gauss-Newton Hessian.
     eps: finite-difference step size.
     output: Optional file or StringIO to which to print messages.
   """
@@ -506,12 +506,16 @@ def check_norm(
   # Check that Hessian is positive-definite.
   if np.any(np.linalg.eigvals(n_h) < 0):
     h_min = np.min(np.linalg.eigvals(n_h))
-    raise ValueError('User-provided norm Hessian is not positive definite. '
+    raise ValueError(
-                     f'Minimum eigenvalue is {h_min:<.4g}')
+        'User-provided norm Hessian is not positive definite. '
+        f'Minimum eigenvalue is {h_min:<.4g}'
+    )
 
   # Local function returning norm values (vectorized).
   def norm_vec(v):
-    norms = [np.atleast_2d(norm.value(v[:, i:i+1])) for i in range(v.shape[1])]
+    norms = [
+        np.atleast_2d(norm.value(v[:, i : i + 1])) for i in range(v.shape[1])
+    ]
     return np.hstack(norms)
 
   # Check the norm gradient.
@@ -519,7 +523,9 @@ def check_norm(
 
   # Local function returning norm gradients (vectorized).
   def grad_vec(v):
-    gradients = [norm.grad_hess(v[:, i:i+1], eye)[0] for i in range(v.shape[1])]
+    gradients = [
+        norm.grad_hess(v[:, i : i + 1], eye)[0] for i in range(v.shape[1])
+    ]
     return np.hstack(gradients)
 
   # Check the norm Hessian.

python/mujoco/minimize_test.py

@@ -56,8 +56,9 @@ class MinimizeTest(absltest.TestCase):
 
     x0 = np.array((0.0, 0.0))
     out = io.StringIO()
-    x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out,
+    x, _ = minimize.least_squares(
-                                  check_derivatives=True)
+        x0, residual, jacobian=jacobian, output=out, check_derivatives=True
+    )
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
     self.assertIn('norm(dx) < tol', out.getvalue())
@@ -67,9 +68,15 @@ class MinimizeTest(absltest.TestCase):
     def bad_jacobian(x, r):
       del r  # Unused.
       return np.array([[-1, 0], [-20 * x[0, 0], 15]])
+
     with self.assertRaisesRegex(ValueError, r'\bJacobian does not match\b'):
-      minimize.least_squares(x0, residual, jacobian=bad_jacobian, output=out,
+      minimize.least_squares(
-                             check_derivatives=True)
+          x0,
+          residual,
+          jacobian=bad_jacobian,
+          output=out,
+          check_derivatives=True,
+      )
 
   def test_max_iter(self) -> None:
     dim = 20  # High-D Rosenbrock
@@ -98,13 +105,16 @@ class MinimizeTest(absltest.TestCase):
     x0 = np.array((0.0, 0.0))
     expected_x = np.array((1.0, 1.0))
 
-    bounds_types = {'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
+    bounds_types = {
-                    'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
+        'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
-                    'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))]}
+        'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
+        'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))],
+    }
 
     # In bounds finds true minimum.
-    x, _ = minimize.least_squares(x0, residual, bounds=bounds_types['inbounds'],
+    x, _ = minimize.least_squares(
-                                  output=out)
+        x0, residual, bounds=bounds_types['inbounds'], output=out
+    )
     np.testing.assert_array_almost_equal(x, expected_x)
     self.assertIn('norm(dx) < tol', out.getvalue())
 
@@ -157,8 +167,9 @@ class MinimizeTest(absltest.TestCase):
       print(f'Hello iteration {len(trace)}!', file=out)
 
     x0 = np.array((0.0, 0.0))
-    x, _ = minimize.least_squares(x0, residual, output=out,
+    x, _ = minimize.least_squares(
-                                  iter_callback=iter_callback)
+        x0, residual, output=out, iter_callback=iter_callback
+    )
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
     self.assertIn('Hello iteration 3!', out.getvalue())
@@ -170,11 +181,12 @@ class MinimizeTest(absltest.TestCase):
     p = 0.01  # Smoothing radius for smooth-L2 norm.
 
     class SmoothL2(minimize.Norm):
+
       def value(self, r):
-        return np.sqrt((r.T @ r).item() + p*p) - p
+        return np.sqrt((r.T @ r).item() + p * p) - p
 
       def grad_hess(self, r, proj):
-        s = np.sqrt((r.T @ r).item() + p*p)
+        s = np.sqrt((r.T @ r).item() + p * p)
         y_r = r / s
         grad = proj.T @ y_r
         y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
@@ -183,8 +195,9 @@ class MinimizeTest(absltest.TestCase):
 
     out = io.StringIO()
     x0 = np.array((0.0, 0.0))
-    x, _ = minimize.least_squares(x0, residual, norm=SmoothL2(), output=out,
+    x, _ = minimize.least_squares(
-                                  check_derivatives=True)
+        x0, residual, norm=SmoothL2(), output=out, check_derivatives=True
+    )
     expected_x = np.array((1.0, 1.0))
     np.testing.assert_array_almost_equal(x, expected_x)
     self.assertIn('norm(dx) < tol', out.getvalue())
@@ -192,11 +205,12 @@ class MinimizeTest(absltest.TestCase):
     self.assertIn('User-provided norm Hessian matches', out.getvalue())
 
     class SmoothL2BadGrad(minimize.Norm):
+
       def value(self, r):
-        return np.sqrt((r.T @ r).item() + p*p) - p
+        return np.sqrt((r.T @ r).item() + p * p) - p
 
       def grad_hess(self, r, proj):
-        s = np.sqrt((r.T @ r).item() + p*p)
+        s = np.sqrt((r.T @ r).item() + p * p)
         y_r = r / s
         grad = proj.T @ (y_r + 0.001)  # 0.001 is erronous.
         y_rr = (np.eye(r.size) - y_r @ y_r.T) / s
@@ -204,15 +218,21 @@ class MinimizeTest(absltest.TestCase):
         return grad, hess
 
     with self.assertRaisesRegex(ValueError, r'\bgradient does not match\b'):
-      minimize.least_squares(x0, residual, norm=SmoothL2BadGrad(), output=out,
+      minimize.least_squares(
-                             check_derivatives=True)
+          x0,
+          residual,
+          norm=SmoothL2BadGrad(),
+          output=out,
+          check_derivatives=True,
+      )
 
     class SmoothL2BadHess(minimize.Norm):
+
       def value(self, r):
-        return np.sqrt((r.T @ r).item() + p*p) - p
+        return np.sqrt((r.T @ r).item() + p * p) - p
 
       def grad_hess(self, r, proj):
-        s = np.sqrt((r.T @ r).item() + p*p)
+        s = np.sqrt((r.T @ r).item() + p * p)
         y_r = r / s
         grad = proj.T @ y_r
         y_rr = (1.001 * np.eye(r.size) - y_r @ y_r.T) / s  # 1.001 is erronous.
@@ -220,15 +240,21 @@ class MinimizeTest(absltest.TestCase):
         return grad, hess
 
     with self.assertRaisesRegex(ValueError, r'\bHessian does not match\b'):
-      minimize.least_squares(x0, residual, norm=SmoothL2BadHess(), output=out,
+      minimize.least_squares(
-                             check_derivatives=True)
+          x0,
+          residual,
+          norm=SmoothL2BadHess(),
+          output=out,
+          check_derivatives=True,
+      )
 
     class SmoothL2AsymHess(minimize.Norm):
+
       def value(self, r):
-        return np.sqrt((r.T @ r).item() + p*p) - p
+        return np.sqrt((r.T @ r).item() + p * p) - p
 
       def grad_hess(self, r, proj):
-        s = np.sqrt((r.T @ r).item() + p*p)
+        s = np.sqrt((r.T @ r).item() + p * p)
         y_r = r / s
         grad = proj.T @ y_r
         y_rr = (np.eye(r.size) - (y_r + 0.0001) @ y_r.T) / s
@@ -236,15 +262,21 @@ class MinimizeTest(absltest.TestCase):
         return grad, hess
 
     with self.assertRaisesRegex(ValueError, r'\bnot symmetric\b'):
-      minimize.least_squares(x0, residual, norm=SmoothL2AsymHess(), output=out,
+      minimize.least_squares(
-                             check_derivatives=True)
+          x0,
+          residual,
+          norm=SmoothL2AsymHess(),
+          output=out,
+          check_derivatives=True,
+      )
 
     class SmoothL2NegHess(minimize.Norm):
+
       def value(self, r):
-        return np.sqrt((r.T @ r).item() + p*p) - p
+        return np.sqrt((r.T @ r).item() + p * p) - p
 
       def grad_hess(self, r, proj):
-        s = np.sqrt((r.T @ r).item() + p*p)
+        s = np.sqrt((r.T @ r).item() + p * p)
         y_r = r / s
         grad = proj.T @ y_r
         y_rr = -(np.eye(r.size) - y_r @ y_r.T) / s  # Negative-definite.
@@ -252,7 +284,14 @@ class MinimizeTest(absltest.TestCase):
         return grad, hess
 
     with self.assertRaisesRegex(ValueError, r'\bnot positive definite\b'):
-      minimize.least_squares(x0, residual, norm=SmoothL2NegHess(), output=out,
+      minimize.least_squares(
-                             check_derivatives=True)
+          x0,
+          residual,
+          norm=SmoothL2NegHess(),
+          output=out,
+          check_derivatives=True,
+      )
+
+
 if __name__ == '__main__':
   absltest.main()

python/mujoco/msh2obj_test.py

@@ -63,7 +63,8 @@ class MshTest(absltest.TestCase):
     obj = msh2obj.msh_to_obj(msh_path)
 
     obj_model = mujoco.MjModel.from_xml_string(
-        _XML, {"abdomen_1_body.obj": obj.encode()})
+        _XML, {"abdomen_1_body.obj": obj.encode()}
+    )
 
     for field in _MESH_FIELDS:
       np.testing.assert_allclose(
@@ -73,5 +74,6 @@ class MshTest(absltest.TestCase):
           err_msg=f"Field {field} does not match between msh and obj models.",
       )
 
+
 if __name__ == "__main__":
   absltest.main()

python/mujoco/render_test.py

@@ -19,8 +19,9 @@ import mujoco
 import numpy as np
 
 
-@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
+@absltest.skipUnless(
-                     'MuJoCo rendering is disabled')
+    hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
+)
 class MuJoCoRenderTest(absltest.TestCase):
 
   def setUp(self):
@@ -48,8 +49,14 @@ class MuJoCoRenderTest(absltest.TestCase):
 
     scene = mujoco.MjvScene(self.model, maxgeom=0)
     mujoco.mjv_updateScene(
-        self.model, self.data, mujoco.MjvOption(), mujoco.MjvPerturb(),
+        self.model,
-        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
+        self.data,
+        mujoco.MjvOption(),
+        mujoco.MjvPerturb(),
+        mujoco.MjvCamera(),
+        mujoco.mjtCatBit.mjCAT_ALL,
+        scene,
+    )
 
     context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
     mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)
@@ -62,7 +69,7 @@ class MuJoCoRenderTest(absltest.TestCase):
     mujoco.mjr_rectangle(blue_rect, 0, 0, 1, 1)
 
     expected_upside_down_image = np.zeros((480, 640, 3), dtype=np.uint8)
-    expected_upside_down_image[67:67+123, 56:56+234, 2] = 255
+    expected_upside_down_image[67 : 67 + 123, 56 : 56 + 234, 2] = 255
 
     upside_down_image = np.empty((480, 640, 3), dtype=np.uint8)
     mujoco.mjr_readPixels(upside_down_image, None, full_rect, context)
@@ -71,7 +78,8 @@ class MuJoCoRenderTest(absltest.TestCase):
     # Check that mjr_readPixels can accept a flattened array.
     upside_down_image[:] = 0
     mujoco.mjr_readPixels(
-        np.reshape(upside_down_image, -1), None, full_rect, context)
+        np.reshape(upside_down_image, -1), None, full_rect, context
+    )
     np.testing.assert_array_equal(upside_down_image, expected_upside_down_image)
     context.free()
 
@@ -81,8 +89,14 @@ class MuJoCoRenderTest(absltest.TestCase):
 
     scene = mujoco.MjvScene(self.model, maxgeom=0)
     mujoco.mjv_updateScene(
-        self.model, self.data, mujoco.MjvOption(), None,
+        self.model,
-        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
+        self.data,
+        mujoco.MjvOption(),
+        None,
+        mujoco.MjvCamera(),
+        mujoco.mjtCatBit.mjCAT_ALL,
+        scene,
+    )
 
     context = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
     mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, context)

python/mujoco/renderer.py

@@ -32,7 +32,7 @@ class Renderer:
       model: _structs.MjModel,
       height: int = 240,
       width: int = 320,
-      max_geom: int = 10000
+      max_geom: int = 10000,
   ) -> None:
     """Initializes a new `Renderer`.
 
@@ -43,6 +43,7 @@ class Renderer:
       max_geom: Optional integer specifying the maximum number of geoms that can
         be rendered in the same scene. If None this will be chosen automatically
         based on the estimated maximum number of renderable geoms in the model.
+
     Raises:
       ValueError: If `camera_id` is outside the valid range, or if `width` or
         `height` exceed the dimensions of MuJoCo's offscreen framebuffer.
@@ -220,9 +221,7 @@ the clause:
       # Convert 3-channel uint8 to 1-channel uint32.
       image3 = out.astype(np.uint32)
       segimage = (
-          image3[:, :, 0]
+          image3[:, :, 0] + image3[:, :, 1] * (2**8) + image3[:, :, 2] * (2**16)
-          + image3[:, :, 1] * (2**8)
-          + image3[:, :, 2] * (2**16)
       )
       # Remap segid to 2-channel (object ID, object type) pair.
       # Seg ID 0 is background -- will be remapped to (-1, -1).
@@ -251,15 +250,15 @@ the clause:
       self,
       data: _structs.MjData,
       camera: Union[int, str, _structs.MjvCamera] = -1,
-      scene_option: Optional[_structs.MjvOption] = None
+      scene_option: Optional[_structs.MjvOption] = None,
-    ):
+  ):
     """Updates geometry used for rendering.
 
     Args:
       data: An instance of `MjData`.
       camera: An instance of `MjvCamera`, a string or an integer
-      scene_option: A custom `MjvOption` instance to use to render
+      scene_option: A custom `MjvOption` instance to use to render the scene
-        the scene instead of the default.
+        instead of the default.
 
     Raises:
       ValueError: If `camera_id` is outside the valid range, or if camera does
@@ -274,8 +273,10 @@ the clause:
         if camera_id == -1:
           raise ValueError(f'The camera "{camera}" does not exist.')
       if camera_id < -1 or camera_id >= self._model.ncam:
-        raise ValueError(f'The camera id {camera_id} is out of'
+        raise ValueError(
-                         f' range [-1, {self._model.ncam}).')
+            f'The camera id {camera_id} is out of'
+            f' range [-1, {self._model.ncam}).'
+        )
 
       # Render camera.
       camera = _structs.MjvCamera()
@@ -295,7 +296,8 @@ the clause:
         data,
         scene_option,
         None,
-        camera, _enums.mjtCatBit.mjCAT_ALL.value,
+        camera,
+        _enums.mjtCatBit.mjCAT_ALL.value,
         self._scene,
     )
 

python/mujoco/renderer_test.py

@@ -20,9 +20,11 @@ import mujoco
 import numpy as np
 
 
-@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
+@absltest.skipUnless(
-                     'MuJoCo rendering is disabled')
+    hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
+)
 class MuJoCoRendererTest(parameterized.TestCase):
+
   def test_renderer_unknown_camera_name(self):
     xml = """
 <mujoco>

python/mujoco/rollout.py

@@ -23,17 +23,19 @@ import numpy as np
 from numpy import typing as npt
 
 
-def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
+def rollout(
-            data: mujoco.MjData,
+    model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
-            initial_state: npt.ArrayLike,
+    data: mujoco.MjData,
-            control: Optional[npt.ArrayLike] = None,
+    initial_state: npt.ArrayLike,
-            *,  # require subsequent arguments to be named
+    control: Optional[npt.ArrayLike] = None,
-            control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
+    *,  # require subsequent arguments to be named
-            skip_checks: bool = False,
+    control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
-            nstep: Optional[int] = None,
+    skip_checks: bool = False,
-            initial_warmstart: Optional[npt.ArrayLike] = None,
+    nstep: Optional[int] = None,
-            state: Optional[npt.ArrayLike] = None,
+    initial_warmstart: Optional[npt.ArrayLike] = None,
-            sensordata: Optional[npt.ArrayLike] = None):
+    state: Optional[npt.ArrayLike] = None,
+    sensordata: Optional[npt.ArrayLike] = None,
+):
   """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
 
   Python wrapper for rollout.cc, see documentation therein.
@@ -66,15 +68,24 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
 
   Raises:
     ValueError: bad shapes or sizes.
-  """
+  """  # fmt: skip
   # skip_checks shortcut:
   #   don't infer nroll/nstep
   #   don't support singleton expansion
   #   don't allocate output arrays
   #   just call rollout and return
   if skip_checks:
-    _rollout.rollout(model, data, nstep, control_spec, initial_state,
+    _rollout.rollout(
-                     initial_warmstart, control, state, sensordata)
+        model,
+        data,
+        nstep,
+        control_spec,
+        initial_state,
+        initial_warmstart,
+        control,
+        state,
+        sensordata,
+    )
     return state, sensordata
 
   if not isinstance(model, mujoco.MjModel):
@@ -92,17 +103,16 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
       initial_warmstart=initial_warmstart,
       control=control,
       state=state,
-      sensordata=sensordata)
+      sensordata=sensordata,
-
+  )
 
   # check number of dimensions
-  _check_number_of_dimensions(2,
+  _check_number_of_dimensions(
-                              initial_state=initial_state,
+      2, initial_state=initial_state, initial_warmstart=initial_warmstart
-                              initial_warmstart=initial_warmstart)
+  )
-  _check_number_of_dimensions(3,
+  _check_number_of_dimensions(
-                              control=control,
+      3, control=control, state=state, sensordata=sensordata
-                              state=state,
+  )
-                              sensordata=sensordata)
 
   # ensure 2D, make contiguous, row-major (C ordering)
   initial_state = _ensure_2d(initial_state)
@@ -114,38 +124,46 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
   sensordata = _ensure_3d(sensordata)
 
   # infer nroll, check for incompatibilities
-  nroll = _infer_dimension(0, 1,
+  nroll = _infer_dimension(
-                           initial_state=initial_state,
+      0,
-                           initial_warmstart=initial_warmstart,
+      1,
-                           control=control,
+      initial_state=initial_state,
-                           state=state,
+      initial_warmstart=initial_warmstart,
-                           sensordata=sensordata)
+      control=control,
+      state=state,
+      sensordata=sensordata,
+  )
   if isinstance(model, list) and nroll == 1:
     nroll = len(model)
 
   if isinstance(model, list) and len(model) != nroll:
-    raise ValueError(f'nroll inferred as {nroll} '
+    raise ValueError(
-                     f'but model is length {len(model)}')
+        f'nroll inferred as {nroll} but model is length {len(model)}'
+    )
   elif not isinstance(model, list):
-    model = [model] # Use a length 1 list to simplify code below
+    model = [model]  # Use a length 1 list to simplify code below
 
   # infer nstep, check for incompatibilities
-  nstep = _infer_dimension(1, nstep or 1,
+  nstep = _infer_dimension(
-                           control=control,
+      1, nstep or 1, control=control, state=state, sensordata=sensordata
-                           state=state,
+  )
-                           sensordata=sensordata)
 
   # get nstate/ncontrol/nv/nsensordata
   # check that they are equal across models
-  nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
+  nstate = mujoco.mj_stateSize(
+      model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
+  )
   ncontrol = mujoco.mj_stateSize(model[0], control_spec)
   nv = model[0].nv
   nsensordata = model[0].nsensordata
   for m in model[1:]:
-    if (nstate != mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
+    if (
+        nstate
+        != mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
         or ncontrol != mujoco.mj_stateSize(m, control_spec)
         or nv != m.nv
-        or nsensordata != m.nsensordata):
+        or nsensordata != m.nsensordata
+    ):
       raise ValueError('models are not compatible')
 
   # check trailing dimensions
@@ -167,8 +185,17 @@ def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
     sensordata = np.empty((nroll, nstep, nsensordata))
 
   # call rollout
-  _rollout.rollout(model, data, nstep, control_spec, initial_state,
+  _rollout.rollout(
-                   initial_warmstart, control, state, sensordata)
+      model,
+      data,
+      nstep,
+      control_spec,
+      initial_state,
+      initial_warmstart,
+      control,
+      state,
+      sensordata,
+  )
 
   # return outputs
   return state, sensordata
@@ -227,8 +254,8 @@ def _infer_dimension(dim, value, **kwargs):
   Args:
     dim: Dimension to be inferred.
     value: Initial guess of inferred value (1: unknown).
-    **kwargs: List of arrays which should all have the same size (or 1)
+    **kwargs: List of arrays which should all have the same size (or 1) along
-      along dimension dim.
+      dimension dim.
 
   Returns:
     Inferred dimension.

python/mujoco/rollout_test.py

@@ -127,10 +127,12 @@ TEST_XML_DIVERGE = r"""
 </mujoco>
 """
 
-ALL_MODELS = {'TEST_XML': TEST_XML,
+ALL_MODELS = {
-              'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
+    'TEST_XML': TEST_XML,
-              'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
+    'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
-              'TEST_XML_EMPTY': TEST_XML_EMPTY}
+    'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
+    'TEST_XML_EMPTY': TEST_XML_EMPTY,
+}
 
 # ------------------------------ tests -----------------------------------------
 
@@ -242,8 +244,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     initial_state = np.random.randn(nstate)
     control = np.random.randn(nstep, model.nu)
     initial_warmstart = np.tile(data.qacc_warmstart.copy(), (nroll, 1))
-    state, sensordata = rollout.rollout(model, data, initial_state, control,
+    state, sensordata = rollout.rollout(
-                                        initial_warmstart=initial_warmstart)
+        model, data, initial_state, control, initial_warmstart=initial_warmstart
+    )
 
     mujoco.mj_resetData(model, data)
     initial_state = np.tile(initial_state, (nroll, 1))
@@ -264,8 +267,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     initial_state = np.random.randn(nstate)
     control = np.random.randn(nstep, model.nu)
     state = np.empty((nroll, nstep, nstate))
-    state, sensordata = rollout.rollout(model, data, initial_state, control,
+    state, sensordata = rollout.rollout(
-                                        state=state)
+        model, data, initial_state, control, state=state
+    )
 
     mujoco.mj_resetData(model, data)
     initial_state = np.tile(initial_state, (nroll, 1))
@@ -286,8 +290,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     initial_state = np.random.randn(nstate)
     control = np.random.randn(nstep, model.nu)
     sensordata = np.empty((nroll, nstep, model.nsensordata))
-    state, sensordata = rollout.rollout(model, data, initial_state, control,
+    state, sensordata = rollout.rollout(
-                                        sensordata=sensordata)
+        model, data, initial_state, control, sensordata=sensordata
+    )
 
     mujoco.mj_resetData(model, data)
     initial_state = np.tile(initial_state, (nroll, 1))
@@ -309,8 +314,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     control = np.random.randn(model.nu)
     state = np.empty((nroll, nstep, nstate))
     sensordata = np.empty((nroll, nstep, model.nsensordata))
-    rollout.rollout(model, data, initial_state, control,
+    rollout.rollout(
-                    state=state, sensordata=sensordata)
+        model, data, initial_state, control, state=state, sensordata=sensordata
+    )
 
     control = np.tile(control, (nstep, 1))
     py_state, py_sensordata = py_rollout(model, data, initial_state, control)
@@ -374,8 +380,9 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     initial_state = np.random.randn(nroll, nstate)
     control = np.random.randn(nroll, 1, model.nu)
     state = np.empty((nroll, nstep, nstate))
-    state, sensordata = rollout.rollout(model, data, initial_state, control,
+    state, sensordata = rollout.rollout(
-                                        state=state)
+        model, data, initial_state, control, state=state
+    )
 
     control = np.repeat(control, nstep, axis=1)
     py_state, py_sensordata = py_rollout(model, data, initial_state, control)
@@ -393,17 +400,21 @@ class MuJoCoRolloutTest(parameterized.TestCase):
 
     initial_state = np.random.randn(nroll, nstate)
 
-    control_spec = (mujoco.mjtState.mjSTATE_CTRL |
+    control_spec = (
-                    mujoco.mjtState.mjSTATE_QFRC_APPLIED |
+        mujoco.mjtState.mjSTATE_CTRL
-                    mujoco.mjtState.mjSTATE_XFRC_APPLIED)
+        | mujoco.mjtState.mjSTATE_QFRC_APPLIED
+        | mujoco.mjtState.mjSTATE_XFRC_APPLIED
+    )
     ncontrol = mujoco.mj_stateSize(model, control_spec)
     control = np.random.randn(nroll, nstep, ncontrol)
 
-    state, sensordata = rollout.rollout(model, data, initial_state, control,
+    state, sensordata = rollout.rollout(
-                                        control_spec=control_spec)
+        model, data, initial_state, control, control_spec=control_spec
+    )
 
-    py_state, py_sensordata = py_rollout(model, data, initial_state, control,
+    py_state, py_sensordata = py_rollout(
-                                         control_spec=control_spec)
+        model, data, initial_state, control, control_spec=control_spec
+    )
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
@@ -416,15 +427,19 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     initial_state = np.empty((nroll, nstate))
 
     # get diverging (0, 2) and non-diverging (1, 3) states
-    mujoco.mj_getState(model, data, initial_state[0],
+    mujoco.mj_getState(
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+        model, data, initial_state[0], mujoco.mjtState.mjSTATE_FULLPHYSICS
-    mujoco.mj_getState(model, data, initial_state[2],
+    )
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(
+        model, data, initial_state[2], mujoco.mjtState.mjSTATE_FULLPHYSICS
+    )
     mujoco.mj_resetDataKeyframe(model, data, 0)  # keyframe 0 does not diverge
-    mujoco.mj_getState(model, data, initial_state[1],
+    mujoco.mj_getState(
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+        model, data, initial_state[1], mujoco.mjtState.mjSTATE_FULLPHYSICS
-    mujoco.mj_getState(model, data, initial_state[3],
+    )
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(
+        model, data, initial_state[3], mujoco.mjtState.mjSTATE_FULLPHYSICS
+    )
 
     nstep = 10000  # divergence after ~15s, timestep = 2e-3
 
@@ -459,27 +474,40 @@ class MuJoCoRolloutTest(parameterized.TestCase):
       thread_local.data = mujoco.MjData(model)
 
     model_list = [model] * nroll
+
     def call_rollout(initial_state, control, state, sensordata):
-      rollout.rollout(model_list, thread_local.data, initial_state, control,
+      rollout.rollout(
-                      skip_checks=True,
+          model_list,
-                      nstep=nstep, state=state, sensordata=sensordata)
+          thread_local.data,
+          initial_state,
+          control,
+          skip_checks=True,
+          nstep=nstep,
+          state=state,
+          sensordata=sensordata,
+      )
 
     n = nroll // num_workers  # integer division
     chunks = []  # a list of tuples, one per worker
-    for i in range(num_workers-1):
+    for i in range(num_workers - 1):
-      chunks.append((initial_state[i*n:(i+1)*n],
+      chunks.append((
-                     control[i*n:(i+1)*n],
+          initial_state[i * n : (i + 1) * n],
-                     state[i*n:(i+1)*n],
+          control[i * n : (i + 1) * n],
-                     sensordata[i*n:(i+1)*n]))
+          state[i * n : (i + 1) * n],
+          sensordata[i * n : (i + 1) * n],
+      ))
 
     # last chunk, absorbing the remainder:
-    chunks.append((initial_state[(num_workers-1)*n:],
+    chunks.append((
-                   control[(num_workers-1)*n:],
+        initial_state[(num_workers - 1) * n :],
-                   state[(num_workers-1)*n:],
+        control[(num_workers - 1) * n :],
-                   sensordata[(num_workers-1)*n:]))
+        state[(num_workers - 1) * n :],
+        sensordata[(num_workers - 1) * n :],
+    ))
 
     with concurrent.futures.ThreadPoolExecutor(
-        max_workers=num_workers, initializer=thread_initializer) as executor:
+        max_workers=num_workers, initializer=thread_initializer
+    ) as executor:
       futures = []
       for chunk in chunks:
         futures.append(executor.submit(call_rollout, *chunk))
@@ -513,12 +541,14 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     state, _ = rollout.rollout(model, data, state1[0], control)
 
     # assert that stepping without warmstarts is not exact
-    np.testing.assert_raises(AssertionError,
+    np.testing.assert_raises(
-                             np.testing.assert_array_equal, state, state2)
+        AssertionError, np.testing.assert_array_equal, state, state2
+    )
 
     # take step using rollout, take warmstart into account
-    state, _ = rollout.rollout(model, data, state1, control,
+    state, _ = rollout.rollout(
-                               initial_warmstart=initial_warmstart)
+        model, data, state1, control, initial_warmstart=initial_warmstart
+    )
 
     # assert exact equality
     np.testing.assert_array_equal(state, np.expand_dims(state2, axis=0))
@@ -530,19 +560,21 @@ class MuJoCoRolloutTest(parameterized.TestCase):
 
     initial_state = np.zeros(nstate)
 
-    control_spec = (mujoco.mjtState.mjSTATE_MOCAP_POS |
+    control_spec = (
-                    mujoco.mjtState.mjSTATE_MOCAP_QUAT)
+        mujoco.mjtState.mjSTATE_MOCAP_POS | mujoco.mjtState.mjSTATE_MOCAP_QUAT
+    )
 
-    pos1 = np.array((1., 2., 3.))
+    pos1 = np.array((1.0, 2.0, 3.0))
-    quat1 = np.array((1., 2., 3., 4.))
+    quat1 = np.array((1.0, 2.0, 3.0, 4.0))
     quat1 /= np.linalg.norm(quat1)
-    pos2 = np.array((2., 3., 4.))
+    pos2 = np.array((2.0, 3.0, 4.0))
-    quat2 = np.array((2., 3., 4., 5.))
+    quat2 = np.array((2.0, 3.0, 4.0, 5.0))
     quat2 /= np.linalg.norm(quat2)
     control = np.hstack((pos1, pos2, quat1, quat2))
 
-    _, sensordata = rollout.rollout(model, data, initial_state, control,
+    _, sensordata = rollout.rollout(
-                                    control_spec=control_spec)
+        model, data, initial_state, control, control_spec=control_spec
+    )
 
     np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1)
     np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1)
@@ -562,7 +594,8 @@ class MuJoCoRolloutTest(parameterized.TestCase):
 
     model.opt.solver = 10  # invalid solver type
     with self.assertRaisesWithLiteralMatch(
-        mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'):
+        mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'
+    ):
       rollout.rollout(model, data, initial_state, ctrl)
 
   def test_invalid(self):
@@ -576,12 +609,14 @@ class MuJoCoRolloutTest(parameterized.TestCase):
 
     control = 'string'
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'control must be a numpy array or float'):
+        ValueError, 'control must be a numpy array or float'
+    ):
       rollout.rollout(model, data, initial_state, control)
 
     control = np.zeros((2, 3, 4, 5))
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'control can have at most 3 dimensions'):
+        ValueError, 'control can have at most 3 dimensions'
+    ):
       rollout.rollout(model, data, initial_state, control)
 
   def test_bad_sizes(self):
@@ -594,28 +629,33 @@ class MuJoCoRolloutTest(parameterized.TestCase):
 
     initial_state = np.random.randn(nroll, nstate + 1)
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'trailing dimension of initial_state must be 6, got 7'):
+        ValueError, 'trailing dimension of initial_state must be 6, got 7'
+    ):
       rollout.rollout(model, data, initial_state)
 
     initial_state = np.random.randn(nroll, nstate)
     control = np.random.randn(1, nstep, model.nu + 1)
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'trailing dimension of control must be 2, got 3'):
+        ValueError, 'trailing dimension of control must be 2, got 3'
+    ):
       rollout.rollout(model, data, initial_state, control)
 
     control = np.random.randn(nroll, nstep, model.nu)
-    state = np.random.randn(nroll, nstep+1, nstate)  # incompatible nstep
+    state = np.random.randn(nroll, nstep + 1, nstate)  # incompatible nstep
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'dimension 1 inferred as 3 but state has 4'):
+        ValueError, 'dimension 1 inferred as 3 but state has 4'
+    ):
       rollout.rollout(model, data, initial_state, control, state=state)
 
     initial_state = np.random.randn(nroll, nstate)
     control = np.random.randn(nroll, nstep, model.nu)
     bad_spec = mujoco.mjtState.mjSTATE_ACT
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'control_spec can only contain bits in mjSTATE_USER'):
+        ValueError, 'control_spec can only contain bits in mjSTATE_USER'
-      rollout.rollout(model, data, initial_state, control,
+    ):
-                      control_spec=bad_spec)
+      rollout.rollout(
+          model, data, initial_state, control, control_spec=bad_spec
+      )
 
   def test_stateless(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
@@ -655,8 +695,9 @@ def get_state(model, data):
   return state.reshape((1, nstate))
 
 
-def step(model, data, state, control,
+def step(
-         control_spec=mujoco.mjtState.mjSTATE_CTRL):
+    model, data, state, control, control_spec=mujoco.mjtState.mjSTATE_CTRL
+):
   if state is not None:
     mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
   mujoco.mj_setState(model, data, control, control_spec)
@@ -664,8 +705,13 @@ def step(model, data, state, control,
   return (get_state(model, data), data.sensordata)
 
 
-def one_rollout(model, data, initial_state, control,
+def one_rollout(
-                control_spec=mujoco.mjtState.mjSTATE_CTRL):
+    model,
+    data,
+    initial_state,
+    control,
+    control_spec=mujoco.mjtState.mjSTATE_CTRL,
+):
   nstep = control.shape[0]
   nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
   state = np.empty((nstep, nstate))
@@ -673,9 +719,9 @@ def one_rollout(model, data, initial_state, control,
 
   mujoco.mj_resetData(model, data)
   for t in range(nstep):
-    state[t], sensordata[t] = step(model, data,
+    state[t], sensordata[t] = step(
-                                   initial_state if t == 0 else None,
+        model, data, initial_state if t == 0 else None, control[t], control_spec
-                                   control[t], control_spec)
+    )
   return state, sensordata
 
 
@@ -700,15 +746,20 @@ def ensure_3d(arg):
     return np.ascontiguousarray(arg, dtype=np.float64)
 
 
-def py_rollout(model, data, initial_state, control,
+def py_rollout(
-               control_spec=mujoco.mjtState.mjSTATE_CTRL):
+    model,
+    data,
+    initial_state,
+    control,
+    control_spec=mujoco.mjtState.mjSTATE_CTRL,
+):
   initial_state = ensure_2d(initial_state)
   control = ensure_3d(control)
   nroll = initial_state.shape[0]
   nstep = control.shape[1]
 
   if isinstance(model, mujoco.MjModel):
-    model = [model]*nroll
+    model = [model] * nroll
 
   nstate = mujoco.mj_stateSize(model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS)
 

python/mujoco/specs_test.py

@@ -104,7 +104,9 @@ class SpecsTest(absltest.TestCase):
     self.assertEqual(model.nuser_site, 6)
     np.testing.assert_array_equal(model.site_user[0], [1, 2, 3, 4, 5, 6])
 
-    self.assertEqual(spec.to_xml(), textwrap.dedent("""\
+    self.assertEqual(
+        spec.to_xml(),
+        textwrap.dedent("""\
         <mujoco model="MuJoCo Model">
           <compiler angle="radian"/>
 
@@ -116,7 +118,8 @@ class SpecsTest(absltest.TestCase):
             </body>
           </worldbody>
         </mujoco>
-    """),)
+    """),
+    )
 
   def test_kwarg(self):
     # Create a spec.
@@ -467,7 +470,7 @@ class SpecsTest(absltest.TestCase):
     # Try to compile, get error.
     expected_error = (
         'Error: size 0 must be positive in geom\n'
-        + f'Element name \'MyGeom\', id 0, geom added on line {added_on_line}'
+        + f"Element name 'MyGeom', id 0, geom added on line {added_on_line}"
     )
     with self.assertRaisesRegex(ValueError, expected_error):
       spec.compile()
@@ -531,7 +534,9 @@ class SpecsTest(absltest.TestCase):
     spec.worldbody.add_geom(main)
 
     spec.compile()
-    self.assertEqual(spec.to_xml(), textwrap.dedent("""\
+    self.assertEqual(
+        spec.to_xml(),
+        textwrap.dedent("""\
         <mujoco model="test">
           <compiler angle="radian"/>
 
@@ -547,7 +552,8 @@ class SpecsTest(absltest.TestCase):
             <geom/>
           </worldbody>
         </mujoco>
-    """))
+    """),
+    )
     spec = mujoco.MjSpec()
     spec.modelname = 'test'
 
@@ -561,7 +567,9 @@ class SpecsTest(absltest.TestCase):
     spec.worldbody.add_geom(main)
 
     spec.compile()
-    self.assertEqual(spec.to_xml(), textwrap.dedent("""\
+    self.assertEqual(
+        spec.to_xml(),
+        textwrap.dedent("""\
         <mujoco model="test">
           <compiler angle="radian"/>
 
@@ -577,7 +585,8 @@ class SpecsTest(absltest.TestCase):
             <geom/>
           </worldbody>
         </mujoco>
-    """))
+    """),
+    )
 
   def test_element_list(self):
     spec = mujoco.MjSpec()
@@ -718,13 +727,17 @@ class SpecsTest(absltest.TestCase):
         </worldbody>
       </mujoco>
     """
-    spec = mujoco.MjSpec.from_string(textwrap.dedent("""
+    spec = mujoco.MjSpec.from_string(
+        textwrap.dedent("""
       <mujoco model="MuJoCo Model">
         <include file="included.xml"/>
       </mujoco>
-    """), {'included.xml': included_xml.encode('utf-8')})
+    """),
-    self.assertEqual(spec.worldbody.first_body().first_geom().type,
+        {'included.xml': included_xml.encode('utf-8')},
-                     mujoco.mjtGeom.mjGEOM_BOX)
+    )
+    self.assertEqual(
+        spec.worldbody.first_body().first_geom().type, mujoco.mjtGeom.mjGEOM_BOX
+    )
 
   def test_delete(self):
     file_path = epath.resource_path("mujoco") / "testdata" / "model.xml"

python/mujoco/viewer.py

@@ -42,7 +42,7 @@ PERCENT_REALTIME = (
     10, 8, 6.6, 5, 4, 3.3, 2.5, 2, 1.6, 1.3,
     1, 0.8, 0.66, 0.5, 0.4, 0.33, 0.25, 0.2, 0.16, 0.13,
     0.1
-)
+)  # fmt: skip
 
 # Maximum time mis-alignment before re-sync.
 MAX_SYNC_MISALIGN = 0.1
@@ -194,12 +194,13 @@ def _file_loader(path: str) -> _LoaderWithPathType:
 
 
 def _reload(
-    simulate: _Simulate, loader: _InternalLoaderType,
+    simulate: _Simulate,
-    notify_loaded: Optional[Callable[[], None]] = None
+    loader: _InternalLoaderType,
+    notify_loaded: Optional[Callable[[], None]] = None,
 ) -> Optional[Tuple[mujoco.MjModel, mujoco.MjData]]:
   """Internal function for reloading a model in the viewer."""
   try:
-    simulate.load_message('') # path is unknown at this point
+    simulate.load_message('')  # path is unknown at this point
     load_tuple = loader()
   except Exception as e:  # pylint: disable=broad-except
     simulate.load_error = str(e)
@@ -275,14 +276,16 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
           # Inject noise.
           if simulate.ctrl_noise_std != 0.0:
             # Convert rate and scale to discrete time (Ornstein–Uhlenbeck).
-            rate = math.exp(-m.opt.timestep /
+            rate = math.exp(
-                            max(simulate.ctrl_noise_rate, mujoco.mjMINVAL))
+                -m.opt.timestep / max(simulate.ctrl_noise_rate, mujoco.mjMINVAL)
+            )
             scale = simulate.ctrl_noise_std * math.sqrt(1 - rate * rate)
 
             for i in range(m.nu):
               # Update noise.
-              ctrl_noise[i] = (rate * ctrl_noise[i] +
+              ctrl_noise[i] = rate * ctrl_noise[
-                               scale * mujoco.mju_standardNormal(None))
+                  i
+              ] + scale * mujoco.mju_standardNormal(None)
 
               # Apply noise.
               d.ctrl[i] = ctrl_noise[i]
@@ -291,12 +294,18 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
           slowdown = 100 / PERCENT_REALTIME[simulate.real_time_index]
 
           # Misalignment: distance from target sim time > MAX_SYNC_MISALIGN.
-          misaligned = abs(elapsedcpu / slowdown -
+          misaligned = (
-                           elapsedsim) > MAX_SYNC_MISALIGN
+              abs(elapsedcpu / slowdown - elapsedsim) > MAX_SYNC_MISALIGN
+          )
 
           # Out-of-sync (for any reason): reset sync times, step.
-          if (elapsedsim < 0 or elapsedcpu < 0 or synccpu == 0 or misaligned or
+          if (
-              simulate.speed_changed):
+              elapsedsim < 0
+              or elapsedcpu < 0
+              or synccpu == 0
+              or misaligned
+              or simulate.speed_changed
+          ):
             # Re-sync.
             synccpu = startcpu
             syncsim = d.time
@@ -312,9 +321,9 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
             prevsim = d.time
             refreshtime = SIM_REFRESH_FRACTION / simulate.refresh_rate
             # Step while sim lags behind CPU and within refreshtime.
-            while (((d.time - syncsim) * slowdown <
+            while (
-                    (time.time() - synccpu)) and
+                (d.time - syncsim) * slowdown < (time.time() - synccpu)
-                   ((time.time() - startcpu) < refreshtime)):
+            ) and ((time.time() - startcpu) < refreshtime):
               # Measure slowdown before first step.
               if not measured and elapsedsim:
                 simulate.measured_slowdown = elapsedcpu / elapsedsim
@@ -329,7 +338,7 @@ def _physics_loop(simulate: _Simulate, loader: Optional[_InternalLoaderType]):
                 break
 
           # save current state to history buffer
-          if (stepped):
+          if stepped:
             simulate.add_to_history()
 
         else:  # simulate.run is False: GUI is paused.
@@ -355,7 +364,8 @@ def _launch_internal(
     raise ValueError('mjData is specified but mjModel is not')
   elif callable(model) and data is not None:
     raise ValueError(
-        'mjData should not be specified when an mjModel loader is used')
+        'mjData should not be specified when an mjModel loader is used'
+    )
   elif loader is not None and model is not None:
     raise ValueError('model and loader are both specified')
   elif run_physics_thread and handle_return is not None:
@@ -398,14 +408,17 @@ def _launch_internal(
 
   if run_physics_thread:
     side_thread = threading.Thread(
-        target=_physics_loop, args=(simulate, loader))
+        target=_physics_loop, args=(simulate, loader)
+    )
   else:
     side_thread = threading.Thread(
-        target=_reload, args=(simulate, loader, notify_loaded))
+        target=_reload, args=(simulate, loader, notify_loaded)
+    )
 
   def make_exit(simulate):
     def exit_simulate():
       simulate.exit()
+
     return exit_simulate
 
   exit_simulate = make_exit(simulate)
@@ -456,8 +469,7 @@ def launch_passive(
   if not isinstance(data, mujoco.MjData):
     raise ValueError(f'`data` is not a mujoco.MjData: got {data!r}')
   if key_callback is not None and not callable(key_callback):
-    raise ValueError(
+    raise ValueError(f'`key_callback` is not callable: got {key_callback!r}')
-        f'`key_callback` is not callable: got {key_callback!r}')
 
   mujoco.mj_forward(model, data)
   handle_return = queue.Queue(1)
@@ -480,7 +492,8 @@ def launch_passive(
     if not isinstance(_MJPYTHON, _MjPythonBase):
       raise RuntimeError(
           '`launch_passive` requires that the Python script be run under '
-          '`mjpython` on macOS')
+          '`mjpython` on macOS'
+      )
     _MJPYTHON.launch_on_ui_thread(
         model,
         data,

python/pyproject.toml

@@ -65,3 +65,24 @@ usd = [
     "usd-core",
     "pillow"
 ]
+
+[tool.isort]
+force_single_line = true
+force_sort_within_sections = true
+lexicographical = true
+single_line_exclusions = ["typing"]
+order_by_type = false
+group_by_package = true
+line_length = 120
+use_parentheses = true
+multi_line_output = 3
+skip_glob = ["**/*.ipynb"]
+
+[tool.pyink]
+line-length = 80
+unstable = true
+pyink-indentation = 2
+pyink-use-majority-quotes = true
+extend-exclude = '''(
+ .ipynb$
+)'''

python/setup.py

@@ -101,15 +101,15 @@ def tokenize_quoted_substr(input_string, quote_char, placeholders=None):
   placeholders = placeholders if placeholders is not None else dict()
   prev_end = -1
   for start, end in start_and_end(quote_positions):
-    output_string += input_string[prev_end+1:start]
+    output_string += input_string[prev_end + 1 : start]
     while True:
       placeholder = ''.join(random.choices(string.ascii_lowercase, k=5))
       if placeholder not in input_string and placeholder not in output_string:
         break
     output_string += placeholder
-    placeholders[placeholder] = input_string[start+1:end]
+    placeholders[placeholder] = input_string[start + 1 : end]
     prev_end = end
-  output_string += input_string[prev_end+1:]
+  output_string += input_string[prev_end + 1 :]
 
   return output_string, placeholders
 
@@ -145,15 +145,17 @@ class BuildCMakeExtension(build_ext.build_ext):
   """Uses CMake to build extensions."""
 
   def run(self):
-    self._is_apple = (platform.system() == 'Darwin')
+    self._is_apple = platform.system() == 'Darwin'
-    (self._mujoco_library_path,
+    (
-     self._mujoco_include_path,
+        self._mujoco_library_path,
-     self._mujoco_plugins_path,
+        self._mujoco_include_path,
-     self._mujoco_framework_path) = self._find_mujoco()
+        self._mujoco_plugins_path,
+        self._mujoco_framework_path,
+    ) = self._find_mujoco()
     self._configure_cmake()
     for ext in self.extensions:
       assert ext.name.startswith(EXT_PREFIX)
-      assert '.' not in ext.name[len(EXT_PREFIX):]
+      assert '.' not in ext.name[len(EXT_PREFIX) :]
       self.build_extension(ext)
     self._copy_external_libraries()
     self._copy_mujoco_headers()
@@ -163,20 +165,22 @@ class BuildCMakeExtension(build_ext.build_ext):
 
   def _find_mujoco(self):
     if MUJOCO_PATH not in os.environ:
-      raise RuntimeError(
+      raise RuntimeError(f'{MUJOCO_PATH} environment variable is not set')
-          f'{MUJOCO_PATH} environment variable is not set')
     if MUJOCO_PLUGIN_PATH not in os.environ:
       raise RuntimeError(
-          f'{MUJOCO_PLUGIN_PATH} environment variable is not set')
+          f'{MUJOCO_PLUGIN_PATH} environment variable is not set'
+      )
     library_path = None
     include_path = None
     plugin_path = os.environ[MUJOCO_PLUGIN_PATH]
     for directory, subdirs, filenames in os.walk(os.environ[MUJOCO_PATH]):
       if self._is_apple and 'mujoco.framework' in subdirs:
-        return (os.path.join(directory, 'mujoco.framework/Versions/A'),
+        return (
-                os.path.join(directory, 'mujoco.framework/Headers'),
+            os.path.join(directory, 'mujoco.framework/Versions/A'),
-                plugin_path,
+            os.path.join(directory, 'mujoco.framework/Headers'),
-                directory)
+            plugin_path,
+            directory,
+        )
       if fnmatch.filter(filenames, get_mujoco_lib_pattern()):
         library_path = directory
       if os.path.exists(os.path.join(directory, 'mujoco/mujoco.h')):
@@ -190,63 +194,78 @@ class BuildCMakeExtension(build_ext.build_ext):
     for directory, _, filenames in os.walk(os.environ[MUJOCO_PATH]):
       for pattern in get_external_lib_patterns():
         for filename in fnmatch.filter(filenames, pattern):
-          shutil.copyfile(os.path.join(directory, filename),
+          shutil.copyfile(
-                          os.path.join(dst, filename))
+              os.path.join(directory, filename), os.path.join(dst, filename)
+          )
 
   def _copy_plugin_libraries(self):
     dst = os.path.join(
         os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
-        'plugin')
+        'plugin',
+    )
     os.makedirs(dst)
     for directory, _, filenames in os.walk(self._mujoco_plugins_path):
       for pattern in get_plugin_lib_patterns():
         for filename in fnmatch.filter(filenames, pattern):
-          shutil.copyfile(os.path.join(directory, filename),
+          shutil.copyfile(
-                          os.path.join(dst, filename))
+              os.path.join(directory, filename), os.path.join(dst, filename)
+          )
 
   def _copy_mujoco_headers(self):
     dst = os.path.join(
         os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
-        'include/mujoco')
+        'include/mujoco',
+    )
     os.makedirs(dst)
     for directory, _, filenames in os.walk(self._mujoco_include_path):
       for filename in fnmatch.filter(filenames, '*.h'):
-        shutil.copyfile(os.path.join(directory, filename),
+        shutil.copyfile(
-                        os.path.join(dst, filename))
+            os.path.join(directory, filename), os.path.join(dst, filename)
+        )
 
   def _copy_mjpython(self):
     src_dir = os.path.join(os.path.dirname(__file__), 'mujoco/mjpython')
     dst_contents_dir = os.path.join(
         os.path.dirname(self.get_ext_fullpath(self.extensions[0].name)),
-        'MuJoCo_(mjpython).app/Contents')
+        'MuJoCo_(mjpython).app/Contents',
+    )
     os.makedirs(dst_contents_dir)
-    shutil.copyfile(os.path.join(src_dir, 'Info.plist'),
+    shutil.copyfile(
-                    os.path.join(dst_contents_dir, 'Info.plist'))
+        os.path.join(src_dir, 'Info.plist'),
+        os.path.join(dst_contents_dir, 'Info.plist'),
+    )
 
     dst_bin_dir = os.path.join(dst_contents_dir, 'MacOS')
     os.makedirs(dst_bin_dir)
-    shutil.copyfile(os.path.join(self.build_temp, 'mjpython'),
+    shutil.copyfile(
-                    os.path.join(dst_bin_dir, 'mjpython'))
+        os.path.join(self.build_temp, 'mjpython'),
+        os.path.join(dst_bin_dir, 'mjpython'),
+    )
     os.chmod(os.path.join(dst_bin_dir, 'mjpython'), 0o755)
 
     dst_resources_dir = os.path.join(dst_contents_dir, 'Resources')
     os.makedirs(dst_resources_dir)
-    shutil.copyfile(os.path.join(src_dir, 'mjpython.icns'),
+    shutil.copyfile(
-                    os.path.join(dst_resources_dir, 'mjpython.icns'))
+        os.path.join(src_dir, 'mjpython.icns'),
+        os.path.join(dst_resources_dir, 'mjpython.icns'),
+    )
 
   def _configure_cmake(self):
     """Check for CMake."""
     cmake = os.environ.get(MUJOCO_CMAKE, 'cmake')
     build_cfg = 'Debug' if self.debug else 'Release'
     cmake_module_path = os.path.join(
-        os.path.dirname(__file__), 'mujoco', 'cmake')
+        os.path.dirname(__file__), 'mujoco', 'cmake'
+    )
     cmake_args = [
         f'-DPython3_ROOT_DIR:PATH={sys.prefix}',
         f'-DPython3_EXECUTABLE:STRING={sys.executable}',
         f'-DCMAKE_MODULE_PATH:PATH={cmake_module_path}',
         f'-DCMAKE_BUILD_TYPE:STRING={build_cfg}',
         f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY:PATH={self.build_temp}',
-        f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL={"OFF" if self.debug else "ON"}',
+        (
+            f'-DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL={"OFF" if self.debug else "ON"}'
+        ),
         '-DCMAKE_Fortran_COMPILER:STRING=',
         '-DBUILD_TESTING:BOOL=OFF',
     ]
@@ -284,14 +303,17 @@ class BuildCMakeExtension(build_ext.build_ext):
     for arg in cmake_args:
       print(f'    {arg}')
     subprocess.check_call(
-        [cmake] + cmake_args +
+        [cmake]
-        [os.path.join(os.path.dirname(__file__), 'mujoco')],
+        + cmake_args
-        cwd=self.build_temp)
+        + [os.path.join(os.path.dirname(__file__), 'mujoco')],
+        cwd=self.build_temp,
+    )
 
     print('Building all extensions with CMake')
     subprocess.check_call(
         [cmake, '--build', '.', f'-j{os.cpu_count()}', '--config', build_cfg],
-        cwd=self.build_temp)
+        cwd=self.build_temp,
+    )
 
   def build_extension(self, ext):
     dest_path = self.get_ext_fullpath(ext.name)
@@ -331,6 +353,7 @@ class InstallScripts(install_scripts.install_scripts):
       else:
         self.outfiles.append(oldfile)
 
+
 setuptools.setup(
     long_description=get_long_description(),
     long_description_content_type='text/markdown',
@@ -350,7 +373,7 @@ setuptools.setup(
         CMakeExtension('mujoco._specs'),
         CMakeExtension('mujoco._structs'),
     ],
-    scripts=[
+    scripts=['mujoco/mjpython/mjpython.py']
-        'mujoco/mjpython/mjpython.py'
+    if platform.system() == 'Darwin'
-    ] if platform.system() == 'Darwin' else [],
+    else [],
 )