Add pyink and isort config. Reformat.

PiperOrigin-RevId: 704533915 Change-Id: I37e9fd51261bd166b725c7460fc65d02fed2b391
2024-12-09 21:10:00 -08:00
parent 6f6244b739
commit f3b3024291
41 changed files with 895 additions and 480 deletions
@@ -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`.
@@ -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_b_pt[best_edge_idx], edge_b_pt_2[best_edge_idx])
+      edge_a_pt[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,
-      jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos)
+      is_edge_contact, jp.tile(0.5 * (a_closest + b_closest), (4, 1)), pos
+  )

  return dist, pos, normal

@@ -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
@@ -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())

@@ -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([
-      vec  - n * d1 * 0.5,
-      vec1 + vec * -0.5 - n * d2 * 0.5,
-      -vec1 + vec * -0.5 - n * d2 * 0.5,
-  ])
+  pos = (
+      cylinder.pos
+      + axis
+      + 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
@@ -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
@@ -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
@@ -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
@@ -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],
@@ -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
@@ -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.
-        If using this flag, also use _full_compat for put_model.
+      This is an experimental feature.  Avoid using it for now. If using this
+      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.
-        If using this flag, also use _full_compat for put_model.
+      This is an experimental feature.  Avoid using it for now. If using this
+      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

@@ -92,12 +92,11 @@ _MULTIPLE_CONSTRAINTS = """
 class ModelIOTest(parameterized.TestCase):
  """IO tests for mjx.Model."""

-  @parameterized.parameters(
-      _MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS
-  )
+  @parameterized.parameters(_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
@@ -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.
  """
@@ -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 = name[:name.find(b'\x00')].decode('utf-8')
+      name = m.names[m.name_meshadr[meshid] :]
+      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 '
@@ -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([
-        (0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)])
+    expected_face_verts = sorted(
+        [(0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)]
+    )
    self.assertSequenceEqual(
        face_verts,
        expected_face_verts,
@@ -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 = []
@@ -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
@@ -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(
-          (
-              f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
-              f' which does not match the in_types[{idx}] expected length of '
-              f'"{size[typ]}".'
-          )
-      )
+      raise IndexError((
+          f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
+          f' which does not match the in_types[{idx}] expected length of '
+          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(
-        (
-            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' "{take_ids.shape[0]}".'
-        )
-    )
+    raise IndexError((
+        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' "{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':
@@ -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))
@@ -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
@@ -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:
@@ -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
@@ -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)
@@ -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
@@ -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'
 )
@@ -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.',
-                                  required=True)
+_MODEL_PATH = flags.DEFINE_string(
+    'mjcf', None, 'Path to a MuJoCo MJCF file.', required=True
+)


 _VIEWER_GLOBAL_STATE = {
@@ -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$
+)'''
@@ -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(
-        qpos_ref, [0.789]*len(self.data.qpos))
+    np.testing.assert_array_equal(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',
-                   reason='requires sys.getrefcount')
+  @absltest.skipIf(
+      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.assertIs(self.data.actuator('myactuator'),
-                  self.data.actuator(actuator_id))
-    self.assertIs(self.data.actuator('myactuator').ctrl,
-                  self.data.actuator(actuator_id).ctrl)
+        self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, 'myactuator'
+    )
+    self.assertIs(
+        self.data.actuator('myactuator'), self.data.actuator(actuator_id)
+    )
+    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.data.joint(balljoint_id).qpos)
+    self.assertIs(
+        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],
-                                  [9, 8, 7, 6])
+    np.testing.assert_array_equal(
+        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.data.joint(freejoint_id).cdof)
+    self.assertIs(
+        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))
-    np.testing.assert_array_equal(cdof_from_indexer,
-                                  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.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,
-                    model2._address, data2._address}, 4)
+    self.assertLen(
+        {
+            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),
-        'myfree')
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_JOINT, 0), 'myfree'
+    )

    self.assertEqual(
-        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0),
-        'myplane')
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 0), 'myplane'
+    )
    self.assertEqual(
-        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1),
-        'mybox')
+        mujoco.mj_id2name(model_copy, mujoco.mjtObj.mjOBJ_GEOM, 1), '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,
-                                   self.data.subtree_angmom[0, :])
+    np.testing.assert_almost_equal(
+        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(
-        [[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]])
+    expected_jacp = np.array([
+        [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(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(
-        [[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, 0, 0, 0, 0]])
+    expected_jacr = np.array([
+        [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, 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],
-                     ('Inertia', '0', 'I'))
+    self.assertEqual(
+        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_ENERGY, 1<<1)
-    self.assertEqual(mujoco.mjtEnableBit.mjENBL_FWDINV, 1<<2)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_OVERRIDE, 1 << 0)
+    self.assertEqual(mujoco.mjtEnableBit.mjENBL_ENERGY, 1 << 1)
+    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],
-        ['c', 'd'])
+        x[mujoco.mjtFrame.mjFRAME_GEOM : mujoco.mjtFrame.mjFRAME_CAMERA],
+        ['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,
-                     mujoco.mjtFrame.mjFRAME_WORLD)
-    self.assertNotEqual(mujoco.mjtFrame.mjFRAME_WORLD,
-                        mujoco.mjtFrame.mjFRAME_NONE)
+    self.assertEqual(
+        mujoco.mjtFrame.mjFRAME_WORLD, mujoco.mjtFrame.mjFRAME_WORLD
+    )
+    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,
-                     mujoco.mjtDisableBit.mjDSBL_WARMSTART)
-    self.assertEqual(mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL >> 3,
-                     mujoco.mjtDisableBit.mjDSBL_CONTACT)
+    self.assertEqual(
+        mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL << 1,
+        mujoco.mjtDisableBit.mjDSBL_WARMSTART,
+    )
+    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,
-                                r'\Amj_stackAlloc: insufficient memory:'):
+    with self.assertRaisesRegex(
+        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',
-                   reason='requires sys.getrefcount')
+  @absltest.skipIf(
+      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',
-                   reason='requires sys.getrefcount')
+  @absltest.skipIf(
+      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(),
-                           None, mujoco.MjvCamera(),
-                           mujoco.mjtCatBit.mjCAT_ALL, scene)
+    mujoco.mjv_updateScene(
+        self.model,
+        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,
-                  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], 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,
+    )
    # 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(
-              self.model, self.data, pnt, vec[i], None, 1, -1, geom1
-          ),
+          mujoco.mj_ray(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,
-                         None, None, None, None, None, None, None)
+    mujoco.mjd_inverseFD(
+        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.array((42., 0., 0., 42., 0., 41.9)))
+    np.testing.assert_array_equal(
+        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,
-                         df_dq, df_dv, df_da, ds_dq, ds_dv, ds_da, dm_dq)
+    mujoco.mjd_inverseFD(
+        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],
-                      [0, 2.0, 0, 0.2],
-                      [0, 0, 3.0, 0.3],
-                      [0.1, 0.2, 0.3, 4.0]])
-    band = np.zeros(n_band*(n_total-n_dense) + n_dense*n_total)
+    dense = np.array([
+        [1.0, 0, 0, 0.1],
+        [0, 2.0, 0, 0.2],
+        [0, 0, 3.0, 0.3],
+        [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,
-                             n_total, n_band, n_dense, n_vec, flg_sym)
+    mujoco.mju_bandMulMatVec(
+        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,
-                              diag_add, diag_mul)
+    mujoco.mju_cholFactorBand(
+        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.])
-    vec2 = np.array([3., 2., 1.])
-    mat = np.array([[1., 2., 3.], [4., 5., 6.], [7., 8., 9.]])
-    self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.)
+    vec1 = np.array([1.0, 2.0, 3.0])
+    vec2 = np.array([3.0, 2.0, 1.0])
+    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.0)

  def test_mju_dense_to_sparse(self):
-    mat = np.array([[0., 1., 0.], [2., 0., 3.]])
-    expected_vals = np.array([1., 2., 3.])
+    mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
+    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.]])
-    mat = np.array((1., 2., 3.))
+    expected = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
+    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',
-        'warning', 'energy', 'contact', 'efc_J'
+        'time',
+        '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',
-        'warning', 'energy', 'contact', 'efc_J'
+        'time',
+        '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(
-            name, str(e)))
+        self.fail(
+            "Attribute '{}' differs from expected value: {}".format(
+                name, str(e)
+            )
+        )


 if __name__ == '__main__':
@@ -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()
@@ -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 '
-                     'to a relative tolerance of 1e-5.')
+    raise ValueError(
+        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
-      and Gauss-Newton Hessian.
+    norm: Norm function returning either the norm scalar or its gradient and
+      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. '
-                     f'Minimum eigenvalue is {h_min:<.4g}')
+    raise ValueError(
+        '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.
@@ -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,
-                                  check_derivatives=True)
+    x, _ = minimize.least_squares(
+        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,
-                             check_derivatives=True)
+      minimize.least_squares(
+          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))],
-                    '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))]}
+    bounds_types = {
+        'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 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'],
-                                  output=out)
+    x, _ = minimize.least_squares(
+        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,
-                                  iter_callback=iter_callback)
+    x, _ = minimize.least_squares(
+        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,
-                                  check_derivatives=True)
+    x, _ = minimize.least_squares(
+        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,
-                             check_derivatives=True)
+      minimize.least_squares(
+          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,
-                             check_derivatives=True)
+      minimize.least_squares(
+          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,
-                             check_derivatives=True)
+      minimize.least_squares(
+          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,
-                             check_derivatives=True)
+      minimize.least_squares(
+          x0,
+          residual,
+          norm=SmoothL2NegHess(),
+          output=out,
+          check_derivatives=True,
+      )
+
+
 if __name__ == '__main__':
  absltest.main()
@@ -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()
@@ -19,8 +19,9 @@ import mujoco
 import numpy as np


-@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
-                     'MuJoCo rendering is disabled')
+@absltest.skipUnless(
+    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(),
-        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
+        self.model,
+        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,
-        mujoco.MjvCamera(), mujoco.mjtCatBit.mjCAT_ALL, scene)
+        self.model,
+        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)
@@ -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[:, :, 1] * (2**8)
-          + image3[:, :, 2] * (2**16)
+          image3[:, :, 0] + 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
-        the scene instead of the default.
+      scene_option: A custom `MjvOption` instance to use to render the scene
+        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'
-                         f' range [-1, {self._model.ncam}).')
+        raise ValueError(
+            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,
    )

@@ -20,9 +20,11 @@ import mujoco
 import numpy as np


-@absltest.skipUnless(hasattr(mujoco, 'GLContext'),
-                     'MuJoCo rendering is disabled')
+@absltest.skipUnless(
+    hasattr(mujoco, 'GLContext'), 'MuJoCo rendering is disabled'
+)
 class MuJoCoRendererTest(parameterized.TestCase):
+
  def test_renderer_unknown_camera_name(self):
    xml = """
 <mujoco>
@@ -23,17 +23,19 @@ import numpy as np
 from numpy import typing as npt


-def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
-            data: mujoco.MjData,
-            initial_state: npt.ArrayLike,
-            control: Optional[npt.ArrayLike] = None,
-            *,  # require subsequent arguments to be named
-            control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
-            skip_checks: bool = False,
-            nstep: Optional[int] = None,
-            initial_warmstart: Optional[npt.ArrayLike] = None,
-            state: Optional[npt.ArrayLike] = None,
-            sensordata: Optional[npt.ArrayLike] = None):
+def rollout(
+    model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
+    data: mujoco.MjData,
+    initial_state: npt.ArrayLike,
+    control: Optional[npt.ArrayLike] = None,
+    *,  # require subsequent arguments to be named
+    control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
+    skip_checks: bool = False,
+    nstep: Optional[int] = None,
+    initial_warmstart: 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,
-                     initial_warmstart, control, state, sensordata)
+    _rollout.rollout(
+        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,
-                              initial_state=initial_state,
-                              initial_warmstart=initial_warmstart)
-  _check_number_of_dimensions(3,
-                              control=control,
-                              state=state,
-                              sensordata=sensordata)
+  _check_number_of_dimensions(
+      2, initial_state=initial_state, initial_warmstart=initial_warmstart
+  )
+  _check_number_of_dimensions(
+      3, control=control, 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,
-                           initial_state=initial_state,
-                           initial_warmstart=initial_warmstart,
-                           control=control,
-                           state=state,
-                           sensordata=sensordata)
+  nroll = _infer_dimension(
+      0,
+      1,
+      initial_state=initial_state,
+      initial_warmstart=initial_warmstart,
+      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} '
-                     f'but model is length {len(model)}')
+    raise ValueError(
+        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,
-                           control=control,
-                           state=state,
-                           sensordata=sensordata)
+  nstep = _infer_dimension(
+      1, nstep or 1, control=control, 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,
-                   initial_warmstart, control, state, sensordata)
+  _rollout.rollout(
+      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)
-      along dimension dim.
+    **kwargs: List of arrays which should all have the same size (or 1) along
+      dimension dim.

  Returns:
    Inferred dimension.
@@ -127,10 +127,12 @@ TEST_XML_DIVERGE = r"""
 </mujoco>
 """

-ALL_MODELS = {'TEST_XML': TEST_XML,
-              'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
-              'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
-              'TEST_XML_EMPTY': TEST_XML_EMPTY}
+ALL_MODELS = {
+    'TEST_XML': TEST_XML,
+    'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
+    '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,
-                                        initial_warmstart=initial_warmstart)
+    state, sensordata = rollout.rollout(
+        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=state)
+    state, sensordata = rollout.rollout(
+        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,
-                                        sensordata=sensordata)
+    state, sensordata = rollout.rollout(
+        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,
-                    state=state, sensordata=sensordata)
+    rollout.rollout(
+        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=state)
+    state, sensordata = rollout.rollout(
+        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 |
-                    mujoco.mjtState.mjSTATE_QFRC_APPLIED |
-                    mujoco.mjtState.mjSTATE_XFRC_APPLIED)
+    control_spec = (
+        mujoco.mjtState.mjSTATE_CTRL
+        | 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,
-                                        control_spec=control_spec)
+    state, sensordata = rollout.rollout(
+        model, data, initial_state, control, control_spec=control_spec
+    )

-    py_state, py_sensordata = py_rollout(model, data, initial_state, control,
-                                         control_spec=control_spec)
+    py_state, py_sensordata = py_rollout(
+        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.mjtState.mjSTATE_FULLPHYSICS)
-    mujoco.mj_getState(model, data, initial_state[2],
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(
+        model, data, initial_state[0], 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.mjtState.mjSTATE_FULLPHYSICS)
-    mujoco.mj_getState(model, data, initial_state[3],
-                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(
+        model, data, initial_state[1], 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,
-                      skip_checks=True,
-                      nstep=nstep, state=state, sensordata=sensordata)
+      rollout.rollout(
+          model_list,
+          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):
-      chunks.append((initial_state[i*n:(i+1)*n],
-                     control[i*n:(i+1)*n],
-                     state[i*n:(i+1)*n],
-                     sensordata[i*n:(i+1)*n]))
+    for i in range(num_workers - 1):
+      chunks.append((
+          initial_state[i * n : (i + 1) * n],
+          control[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:],
-                   control[(num_workers-1)*n:],
-                   state[(num_workers-1)*n:],
-                   sensordata[(num_workers-1)*n:]))
+    chunks.append((
+        initial_state[(num_workers - 1) * n :],
+        control[(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_array_equal, state, state2)
+    np.testing.assert_raises(
+        AssertionError, np.testing.assert_array_equal, state, state2
+    )

    # take step using rollout, take warmstart into account
-    state, _ = rollout.rollout(model, data, state1, control,
-                               initial_warmstart=initial_warmstart)
+    state, _ = rollout.rollout(
+        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 |
-                    mujoco.mjtState.mjSTATE_MOCAP_QUAT)
+    control_spec = (
+        mujoco.mjtState.mjSTATE_MOCAP_POS | mujoco.mjtState.mjSTATE_MOCAP_QUAT
+    )

-    pos1 = np.array((1., 2., 3.))
-    quat1 = np.array((1., 2., 3., 4.))
+    pos1 = np.array((1.0, 2.0, 3.0))
+    quat1 = np.array((1.0, 2.0, 3.0, 4.0))
    quat1 /= np.linalg.norm(quat1)
-    pos2 = np.array((2., 3., 4.))
-    quat2 = np.array((2., 3., 4., 5.))
+    pos2 = np.array((2.0, 3.0, 4.0))
+    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,
-                                    control_spec=control_spec)
+    _, sensordata = rollout.rollout(
+        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'):
-      rollout.rollout(model, data, initial_state, control,
-                      control_spec=bad_spec)
+        ValueError, 'control_spec can only contain bits in mjSTATE_USER'
+    ):
+      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,
-         control_spec=mujoco.mjtState.mjSTATE_CTRL):
+def step(
+    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,
-                control_spec=mujoco.mjtState.mjSTATE_CTRL):
+def one_rollout(
+    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,
-                                   initial_state if t == 0 else None,
-                                   control[t], control_spec)
+    state[t], sensordata[t] = step(
+        model, data, initial_state if t == 0 else None, 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,
-               control_spec=mujoco.mjtState.mjSTATE_CTRL):
+def py_rollout(
+    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)

@@ -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,
-                     mujoco.mjtGeom.mjGEOM_BOX)
+    """),
+        {'included.xml': included_xml.encode('utf-8')},
+    )
+    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"
@@ -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,
-    notify_loaded: Optional[Callable[[], None]] = None
+    simulate: _Simulate,
+    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 /
-                            max(simulate.ctrl_noise_rate, mujoco.mjMINVAL))
+            rate = math.exp(
+                -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] +
-                               scale * mujoco.mju_standardNormal(None))
+              ctrl_noise[i] = rate * ctrl_noise[
+                  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 -
-                           elapsedsim) > MAX_SYNC_MISALIGN
+          misaligned = (
+              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
-              simulate.speed_changed):
+          if (
+              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 <
-                    (time.time() - synccpu)) and
-                   ((time.time() - startcpu) < refreshtime)):
+            while (
+                (d.time - syncsim) * slowdown < (time.time() - synccpu)
+            ) 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(
-        f'`key_callback` is not callable: got {key_callback!r}')
+    raise ValueError(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,
@@ -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$
+)'''
@@ -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._mujoco_library_path,
-     self._mujoco_include_path,
-     self._mujoco_plugins_path,
-     self._mujoco_framework_path) = self._find_mujoco()
+    self._is_apple = platform.system() == 'Darwin'
+    (
+        self._mujoco_library_path,
+        self._mujoco_include_path,
+        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(
-          f'{MUJOCO_PATH} environment variable is not set')
+      raise RuntimeError(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'),
-                os.path.join(directory, 'mujoco.framework/Headers'),
-                plugin_path,
-                directory)
+        return (
+            os.path.join(directory, 'mujoco.framework/Versions/A'),
+            os.path.join(directory, 'mujoco.framework/Headers'),
+            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),
-                          os.path.join(dst, filename))
+          shutil.copyfile(
+              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),
-                          os.path.join(dst, filename))
+          shutil.copyfile(
+              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),
-                        os.path.join(dst, filename))
+        shutil.copyfile(
+            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'),
-                    os.path.join(dst_contents_dir, 'Info.plist'))
+    shutil.copyfile(
+        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'),
-                    os.path.join(dst_bin_dir, 'mjpython'))
+    shutil.copyfile(
+        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'),
-                    os.path.join(dst_resources_dir, 'mjpython.icns'))
+    shutil.copyfile(
+        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 +
-        [os.path.join(os.path.dirname(__file__), 'mujoco')],
-        cwd=self.build_temp)
+        [cmake]
+        + cmake_args
+        + [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=[
-        'mujoco/mjpython/mjpython.py'
-    ] if platform.system() == 'Darwin' else [],
+    scripts=['mujoco/mjpython/mjpython.py']
+    if platform.system() == 'Darwin'
+    else [],
 )