From f3b3024291cbe7efd91c4871b598a1226114167e Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Mon, 9 Dec 2024 21:10:00 -0800 Subject: [PATCH] Add pyink and isort config. Reformat. PiperOrigin-RevId: 704533915 Change-Id: I37e9fd51261bd166b725c7460fc65d02fed2b391 --- STYLEGUIDE.md | 6 +- mjx/mujoco/mjx/_src/collision_convex.py | 14 +- mjx/mujoco/mjx/_src/collision_driver.py | 6 +- mjx/mujoco/mjx/_src/collision_driver_test.py | 28 +- mjx/mujoco/mjx/_src/collision_primitive.py | 17 +- mjx/mujoco/mjx/_src/collision_sdf.py | 29 +- mjx/mujoco/mjx/_src/collision_types.py | 1 + mjx/mujoco/mjx/_src/constraint.py | 1 + mjx/mujoco/mjx/_src/constraint_test.py | 4 +- mjx/mujoco/mjx/_src/dataclasses.py | 3 +- mjx/mujoco/mjx/_src/io.py | 15 +- mjx/mujoco/mjx/_src/io_test.py | 5 +- mjx/mujoco/mjx/_src/math.py | 1 + mjx/mujoco/mjx/_src/mesh.py | 4 +- mjx/mujoco/mjx/_src/mesh_test.py | 5 +- mjx/mujoco/mjx/_src/passive.py | 1 + mjx/mujoco/mjx/_src/ray.py | 2 +- mjx/mujoco/mjx/_src/scan.py | 26 +- mjx/mujoco/mjx/_src/scan_test.py | 3 + mjx/mujoco/mjx/_src/sensor_test.py | 2 - mjx/mujoco/mjx/_src/smooth.py | 4 +- mjx/mujoco/mjx/_src/solver.py | 4 +- mjx/mujoco/mjx/_src/test_util.py | 2 +- mjx/mujoco/mjx/_src/types.py | 26 +- mjx/mujoco/mjx/testspeed.py | 4 +- mjx/mujoco/mjx/viewer.py | 5 +- mjx/pyproject.toml | 21 + python/mujoco/bindings_test.py | 410 ++++++++++++------- python/mujoco/memory_leak_test.py | 3 +- python/mujoco/minimize.py | 30 +- python/mujoco/minimize_test.py | 101 +++-- python/mujoco/msh2obj_test.py | 4 +- python/mujoco/render_test.py | 30 +- python/mujoco/renderer.py | 24 +- python/mujoco/renderer_test.py | 6 +- python/mujoco/rollout.py | 113 +++-- python/mujoco/rollout_test.py | 199 +++++---- python/mujoco/specs_test.py | 35 +- python/mujoco/viewer.py | 57 ++- python/pyproject.toml | 21 + python/setup.py | 103 +++-- 41 files changed, 895 insertions(+), 480 deletions(-) diff --git a/STYLEGUIDE.md b/STYLEGUIDE.md index c0f12f52..dfecb4ab 100644 --- a/STYLEGUIDE.md +++ b/STYLEGUIDE.md @@ -6,7 +6,7 @@ possible in your code contributions. ### Scope of this guide -MuJoCo has three main code categories: +Most of this guide involves C/C++ code. For Python, jump to the section [below](#python-code). For MuJoCo C/C++, code has three main categories: 1. **C code:** MuJoCo's core codebase. It consists of public headers under `include/` and C source files and internal headers under `src/`. This style @@ -158,3 +158,7 @@ example above. New code should use the C99 convention. When editing an existing function, please move existing variable declarations into local scope. Pull requests helping us to complete the migration are very welcome. + +### [Python code](#python-code) + +For Python code, run `pyink foo.py` to adhere to Google's [Python style guide](https://google.github.io/styleguide/pyguide.html). For sorting and cleaning imports, run `isort foo.py`. Both `pyink` and `isort` can be pip installed via `pip install pyink isort`. diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index b6876c9d..465ec2bb 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -906,7 +906,8 @@ def _sat_gaussmap( return edge_axis * sign, degenerate_edge_axis edge_axes, degenerate_edge_axes = jax.vmap(get_normals)( - edge_a_dir, edge_a_pt, edge_b_dir) + edge_a_dir, edge_a_pt, edge_b_dir + ) edge_dist = jax.vmap(jp.dot)(edge_axes, edge_b_pt - edge_a_pt) # handle degenerate axis edge_dist = jp.where(degenerate_edge_axes, -jp.inf, edge_dist) @@ -928,11 +929,14 @@ def _sat_gaussmap( dist, ) a_closest, b_closest = math.closest_segment_to_segment_points( - edge_a_pt[best_edge_idx], edge_a_pt_2[best_edge_idx], - edge_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 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 129bcf2a..7b947ce1 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -146,13 +146,13 @@ def geom_pairs( b_end = b_start + m.body_geomnum for b1 in range(m.nbody): - if not geom_con[b_start[b1]:b_end[b1]].any(): + if not geom_con[b_start[b1] : b_end[b1]].any(): continue w1 = m.body_weldid[b1] w1_p = m.body_weldid[m.body_parentid[w1]] for b2 in range(b1, m.nbody): - if not geom_con[b_start[b2]:b_end[b2]].any(): + if not geom_con[b_start[b2] : b_end[b2]].any(): continue signature = (b1 << 16) + (b2) if signature in exclude_signature: @@ -272,7 +272,7 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]: jp.clip(m.pair_friction[ip], a_min=eps), m.pair_solref[ip], m.pair_solreffriction[ip], - m.pair_solimp[ip] + m.pair_solimp[ip], )) if geom1.size > 0 and geom2.size > 0: # other contacts get their params from geom fields diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index d5c893d0..068856d7 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -218,7 +218,8 @@ class EllipsoidCollisionTest(parameterized.TestCase): self.assertLess(dx.contact.dist[0], 0) for field in dataclasses.fields(Contact): _assert_attr_eq( - dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5) + dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5 + ) _ELLIPSOID_ELLIPSOID = """ @@ -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 = """ @@ -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 = """ @@ -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 = """ @@ -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( - ' 0).all()) diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py index f5757853..eebd3eb3 100644 --- a/mjx/mujoco/mjx/_src/collision_primitive.py +++ b/mjx/mujoco/mjx/_src/collision_primitive.py @@ -29,6 +29,7 @@ from mujoco.mjx._src.types import Model def collider(ncon: int): """Wraps collision functions for use by collision_driver.""" + def wrapper(func): def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision: g1, g2 = geom.T @@ -119,7 +120,7 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision: # disk parallel to plane: pick x-axis of cylinder, scale by radius cylinder.mat[:, 0] * cylinder.size[0], # general configuration: normalize vector, scale by radius - vec / len_ * cylinder.size[0] + vec / len_ * cylinder.size[0], ) # project vector on normal @@ -138,11 +139,15 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision: d1 = dist0 + prjaxis + prjvec d2 = dist0 + prjaxis + prjvec1 dist = jp.array([d1, d2, d2]) - pos = cylinder.pos + axis + jp.array([ - 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 diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py index 43bcff21..1874211d 100644 --- a/mjx/mujoco/mjx/_src/collision_sdf.py +++ b/mjx/mujoco/mjx/_src/collision_sdf.py @@ -41,6 +41,7 @@ SDFFn = Callable[[jax.Array], jax.Array] def collider(ncon: int): """Wraps collision functions for use by collision_driver.""" + def wrapper(func): def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision: g1, g2 = geom.T @@ -81,7 +82,7 @@ def _capsule(pos: jax.Array, size: jax.Array): def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array: k0 = math.norm(pos / size) - k1 = math.norm(pos / (size*size)) + k1 = math.norm(pos / (size * size)) return k0 * (k0 - 1.0) / (k1 + (k1 == 0.0) * 1e-12) @@ -96,12 +97,12 @@ def _cylinder(pos: jax.Array, size: jax.Array) -> jax.Array: def _cylinder_grad(x: jax.Array, size: jax.Array) -> jax.Array: """Gradient of the cylinder SDF wrt query point and singularities removed.""" - c = jp.sqrt(x[0]*x[0]+x[1]*x[1]) + c = jp.sqrt(x[0] * x[0] + x[1] * x[1]) e = jp.abs(x[2]) a = jp.array([c - size[0], e - size[1]]) b = jp.array([jp.maximum(a[0], 0), jp.maximum(a[1], 0)]) j = jp.argmax(a) - bnorm = jp.sqrt(b[0]*b[0] + b[1]*b[1]) + bnorm = jp.sqrt(b[0] * b[0] + b[1] * b[1]) bnorm += jp.allclose(bnorm, 0) * 1e-12 grada = jp.array([ x[0] / (c + jp.allclose(c, 0) * 1e-12), @@ -151,7 +152,7 @@ def _gradient_step(objective: SDFFn, state: GradientState) -> GradientState: """Performs a step of gradient descent.""" # TODO: find better parameters amin = 1e-4 # minimum value for line search factor scaling the gradient - amax = 2. # maximum value for line search factor scaling the gradient + amax = 2.0 # maximum value for line search factor scaling the gradient nlinesearch = 10 # line search points grad = jax.grad(objective)(state.x) alpha = jp.geomspace(amin, amax, nlinesearch).reshape(nlinesearch, -1) @@ -179,7 +180,11 @@ def _gradient_descent( def _optim( - d1, d2, info1: GeomInfo, info2: GeomInfo, x0: jax.Array, + d1, + d2, + info1: GeomInfo, + info2: GeomInfo, + x0: jax.Array, ) -> Collision: """Optimizes the clearance function.""" d1 = functools.partial(d1, size=info1.size) @@ -198,14 +203,14 @@ def _optim( @collider(ncon=1) def sphere_ellipsoid(s: GeomInfo, e: GeomInfo) -> Collision: - """"Calculates contact between a sphere and an ellipsoid.""" + """Calculates contact between a sphere and an ellipsoid.""" x0 = 0.5 * (s.pos + e.pos) return _optim(_sphere, _ellipsoid, s, e, x0) @collider(ncon=1) def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision: - """"Calculates contact between a sphere and a cylinder.""" + """Calculates contact between a sphere and a cylinder.""" # TODO: implement analytical version. x0 = 0.5 * (s.pos + c.pos) return _optim(_sphere, _cylinder, s, c, x0) @@ -213,14 +218,14 @@ def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision: @collider(ncon=1) def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision: - """"Calculates contact between a capsule and an ellipsoid.""" + """ "Calculates contact between a capsule and an ellipsoid.""" x0 = 0.5 * (c.pos + e.pos) return _optim(_capsule, _ellipsoid, c, e, x0) @collider(ncon=2) def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision: - """"Calculates contact between a capsule and a cylinder.""" + """Calculates contact between a capsule and a cylinder.""" # TODO: improve robustness # Near sharp corners, the SDF might give the penetration depth with respect # to a surface that is not in collision. Possible solutions is to find the @@ -235,21 +240,21 @@ def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision: @collider(ncon=1) def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision: - """"Calculates contact between two ellipsoids.""" + """Calculates contact between two ellipsoids.""" x0 = 0.5 * (e1.pos + e2.pos) return _optim(_ellipsoid, _ellipsoid, e1, e2, x0) @collider(ncon=1) def ellipsoid_cylinder(e: GeomInfo, c: GeomInfo) -> Collision: - """"Calculates contact between and ellipsoid and a cylinder.""" + """Calculates contact between and ellipsoid and a cylinder.""" x0 = 0.5 * (e.pos + c.pos) return _optim(_ellipsoid, _cylinder, e, c, x0) @collider(ncon=4) def cylinder_cylinder(c1: GeomInfo, c2: GeomInfo) -> Collision: - """"Calculates contact between a cylinder and a cylinder.""" + """Calculates contact between a cylinder and a cylinder.""" # TODO: improve robustness # Near sharp corners, the SDF might give the penetration depth with respect # to a surface that is not in collision. Possible solutions is to find the diff --git a/mjx/mujoco/mjx/_src/collision_types.py b/mjx/mujoco/mjx/_src/collision_types.py index e7af8f23..44957711 100644 --- a/mjx/mujoco/mjx/_src/collision_types.py +++ b/mjx/mujoco/mjx/_src/collision_types.py @@ -73,6 +73,7 @@ class FunctionKey: resulting constraint jacobian is determined at compile time. subgrid_size: the size determines the hfield subgrid to collide with """ + types: Tuple[int, int] data_ids: Tuple[int, int] condim: int diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py index 158d9157..e778bcd8 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -39,6 +39,7 @@ import numpy as np class _Efc(PyTreeNode): """Support data for creating constraint matrices.""" + J: jax.Array pos_aref: jax.Array pos_imp: jax.Array diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index 76d9bc43..4e53fb13 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -30,7 +30,7 @@ _TOLERANCE = 5e-5 def _assert_eq(a, b, name): - tol = _TOLERANCE * 10 # avoid test noise + tol = _TOLERANCE * 10 # avoid test noise err_msg = f'mismatch: {name}' np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol) @@ -75,7 +75,7 @@ class ConstraintTest(parameterized.TestCase): _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref') _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D') _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos') - _assert_eq(dx.efc_pos[order][d.nefc:], 0, 'efc_pos') + _assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos') _assert_eq( d.efc_frictionloss, dx.efc_frictionloss[order][: d.nefc], diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index a96bb966..0db4a141 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -16,7 +16,6 @@ import copy import dataclasses - import typing from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union import jax @@ -57,7 +56,7 @@ def dataclass(clz: _T) -> _T: meta_fields.append(field) def replace(self, **updates): - """"Returns a new object replacing the specified fields with new values.""" + """Returns a new object replacing the specified fields with new values.""" return dataclasses.replace(self, **updates) data_clz.replace = replace diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 9e09a3f6..7ace198b 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -34,6 +34,7 @@ def _strip_weak_type(tree): if isinstance(leaf, jax.Array): return leaf.astype(jax.dtypes.canonicalize_dtype(leaf.dtype)) return leaf + return jax.tree_util.tree_map(f, tree) @@ -95,7 +96,7 @@ def put_model( m: the model to put onto device device: which device to use - if unspecified picks the default device _full_compat: put all MjModel fields onto device irrespective of MJX support - This is an experimental feature. Avoid using it for now. + This is an experimental feature. Avoid using it for now. Returns: an mjx.Model placed on device @@ -215,8 +216,8 @@ def make_data( m: the model to use device: which device to use - if unspecified picks the default device _full_compat: create all MjData fields on device irrespective of MJX support - This is an experimental feature. Avoid using it for now. - 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 diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 467b9020..8950c677 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/math.py b/mjx/mujoco/mjx/_src/math.py index 00c9f7fd..c7fffbcf 100644 --- a/mjx/mujoco/mjx/_src/math.py +++ b/mjx/mujoco/mjx/_src/math.py @@ -28,6 +28,7 @@ def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array: Args: a: left hand of matmul operand b: right hand of matmul operand + Returns: the matrix product of the inputs. """ diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py index 31f45571..4c60ee1e 100644 --- a/mjx/mujoco/mjx/_src/mesh.py +++ b/mjx/mujoco/mjx/_src/mesh.py @@ -172,8 +172,8 @@ def _merge_coplanar( # resize faces that exceed max polygon vertices if face.shape[0] > _MAX_HULL_FACE_VERTICES: - name = m.names[m.name_meshadr[meshid]:] - name = 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 ' diff --git a/mjx/mujoco/mjx/_src/mesh_test.py b/mjx/mujoco/mjx/_src/mesh_test.py index 235348cd..f50bffc6 100644 --- a/mjx/mujoco/mjx/_src/mesh_test.py +++ b/mjx/mujoco/mjx/_src/mesh_test.py @@ -53,8 +53,9 @@ class MeshTest(absltest.TestCase): map_ = {v: k for k, v in enumerate(vidx)} h_face = np.vectorize(map_.get)(convex_face) face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()]) - expected_face_verts = sorted([ - (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, diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py index 0975b129..c289690f 100644 --- a/mjx/mujoco/mjx/_src/passive.py +++ b/mjx/mujoco/mjx/_src/passive.py @@ -31,6 +31,7 @@ from mujoco.mjx._src.types import Model def _spring_damper(m: Model, d: Data) -> jax.Array: """Applies joint level spring and damping forces.""" + def fn(jnt_typs, stiffness, qpos_spring, qpos): qpos_i = 0 qfrcs = [] diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py index 6d000373..88db31be 100644 --- a/mjx/mujoco/mjx/_src/ray.py +++ b/mjx/mujoco/mjx/_src/ray.py @@ -269,7 +269,7 @@ def ray( geom_filter_dyn = (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0) geom_filter_dyn &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0) for geom_type, fn in _RAY_FUNC.items(): - id_, = np.nonzero(geom_filter & (m.geom_type == geom_type)) + (id_,) = np.nonzero(geom_filter & (m.geom_type == geom_type)) if id_.size == 0: continue diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py index cbc9bd0e..496ba4cd 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -144,13 +144,11 @@ def _check_input(m: Model, args: Any, in_types: str) -> None: } for idx, (arg, typ) in enumerate(zip(args, in_types)): if len(arg) != size[typ]: - raise IndexError( - ( - 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': diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py index edb57339..74bf3224 100644 --- a/mjx/mujoco/mjx/_src/scan_test.py +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -90,6 +90,7 @@ class ScanTest(absltest.TestCase): if tuple(jnt_types) == (JointType.FREE,): return None return val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in) np.testing.assert_equal(np.array(b_out), np.array(b_expect)) @@ -99,6 +100,7 @@ class ScanTest(absltest.TestCase): if jnt_types.size == 0: self.fail('world has no dofs, should not be called') return val + sum(jnt_types) + v_in = jp.ones((m.nv, 1)) scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in) @@ -141,6 +143,7 @@ class ScanTest(absltest.TestCase): return None carry = jp.zeros_like(val) if carry is None else carry return carry + val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in) np.testing.assert_equal(np.array(b_out), np.array(b_expect)) diff --git a/mjx/mujoco/mjx/_src/sensor_test.py b/mjx/mujoco/mjx/_src/sensor_test.py index 37b1e13e..fce6225c 100644 --- a/mjx/mujoco/mjx/_src/sensor_test.py +++ b/mjx/mujoco/mjx/_src/sensor_test.py @@ -17,13 +17,11 @@ from absl.testing import absltest from absl.testing import parameterized import jax - from jax import numpy as jp import mujoco from mujoco import mjx from mujoco.mjx._src import test_util from mujoco.mjx._src.types import ConeType - import numpy as np # tolerance for difference between MuJoCo and MJX smooth calculations - mostly diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index dd04ee0d..9ad010c1 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -334,7 +334,7 @@ def factor_m(m: Model, d: Data) -> Data: pivots = [] out = [] - for (b, e, madr_d, madr_ij) in updates: + for b, e, madr_d, madr_ij in updates: width = e - b rows.append(np.arange(madr_ij, madr_ij + width)) madr_ijs.append(np.full((width,), madr_ij)) @@ -511,7 +511,6 @@ def subtree_vel(m: Model, d: Data) -> Data: angmom_child, mom_parent_child = carry return angmom + mom + angmom_child + mom_parent_child, mom_parent - subtree_angmom, _ = scan.body_tree( m, _subtree_angmom, @@ -535,6 +534,7 @@ def subtree_vel(m: Model, d: Data) -> Data: def rne(m: Model, d: Data) -> Data: """Computes inverse dynamics using the recursive Newton-Euler algorithm.""" + # forward scan over tree: accumulate link center of mass acceleration def cacc_fn(cacc, cdof_dot, qvel): if cacc is None: diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index 52db56f7..6d3d4f21 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -52,6 +52,7 @@ class _Context(PyTreeNode): u: friction cone (normal and tangents) (num(con.dim > 1), 6) h: cone hessian (num(con.dim > 1), 6, 6) """ + qacc: jax.Array qfrc_constraint: jax.Array Jaref: jax.Array # pylint: disable=invalid-name @@ -225,6 +226,7 @@ class _LSContext(PyTreeNode): def _while_loop_scan(cond_fun, body_fun, init_val, max_iter): """Scan-based implementation (jit ok, reverse-mode autodiff ok).""" + def _iter(val): next_val = body_fun(val) next_cond = cond_fun(next_val) @@ -382,7 +384,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: # set efc of cone H along diagonal for i, (condim, addr) in enumerate(zip(dim, efc_address)): h_cone = ctx.h[i, :condim, :condim] - cm = cm.at[addr:addr+condim, addr:addr+condim].add(h_cone) + cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone) h = d.efc_J.T @ cm @ d.efc_J else: h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py index a5a3af01..f91ee9af 100644 --- a/mjx/mujoco/mjx/_src/test_util.py +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -116,7 +116,7 @@ def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray: if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: num_rows = (dx.contact.dim[i] - 1) * 2 if dx.contact.dist[i] > 0: # move empty contacts to end - order = np.append(order, np.repeat(2 ** 16, num_rows)) + order = np.append(order, np.repeat(2**16, num_rows)) continue contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1) contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 83fe9f4a..4ec13c72 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -47,6 +47,7 @@ class DisableBit(enum.IntFlag): REFSAFE: integrator safety: make ref[0]>=2*timestep SENSOR: sensors """ + CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS @@ -73,6 +74,7 @@ class JointType(enum.IntEnum): SLIDE: sliding distance along body-fixed axis (1,) HINGE: rotation angle (rad) around body-fixed axis (1,) """ + FREE = mujoco.mjtJoint.mjJNT_FREE BALL = mujoco.mjtJoint.mjJNT_BALL SLIDE = mujoco.mjtJoint.mjJNT_SLIDE @@ -93,6 +95,7 @@ class IntegratorType(enum.IntEnum): RK4: 4th-order Runge Kutta IMPLICITFAST: implicit in velocity, no rne derivative """ + EULER = mujoco.mjtIntegrator.mjINT_EULER RK4 = mujoco.mjtIntegrator.mjINT_RK4 IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST @@ -113,6 +116,7 @@ class GeomType(enum.IntEnum): MESH: mesh SDF: signed distance field """ + PLANE = mujoco.mjtGeom.mjGEOM_PLANE HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE @@ -134,6 +138,7 @@ class ConvexMesh(PyTreeNode): edge: edge indexes for all edges in the convex mesh edge_face_normal: indexes for face normals adjacent to edges in `edge` """ + vert: jax.Array face: jax.Array face_normal: jax.Array @@ -148,6 +153,7 @@ class ConeType(enum.IntEnum): PYRAMIDAL: pyramidal ELLIPTIC: elliptic """ + PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC @@ -160,6 +166,7 @@ class JacobianType(enum.IntEnum): SPARSE: sparse AUTO: sparse if nv>60 and device is TPU, dense otherwise """ + DENSE = mujoco.mjtJacobian.mjJAC_DENSE SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE AUTO = mujoco.mjtJacobian.mjJAC_AUTO @@ -172,6 +179,7 @@ class SolverType(enum.IntEnum): CG: Conjugate gradient (primal) NEWTON: Newton (primal) """ + # unsupported: PGS CG = mujoco.mjtSolver.mjSOL_CG NEWTON = mujoco.mjtSolver.mjSOL_NEWTON @@ -186,6 +194,7 @@ class EqType(enum.IntEnum): JOINT: couple the values of two scalar joints with cubic TENDON: couple the lengths of two tendons with cubic """ + CONNECT = mujoco.mjtEq.mjEQ_CONNECT WELD = mujoco.mjtEq.mjEQ_WELD JOINT = mujoco.mjtEq.mjEQ_JOINT @@ -203,6 +212,7 @@ class WrapType(enum.IntEnum): SPHERE: wrap around sphere CYLINDER: wrap around (infinite) cylinder """ + JOINT = mujoco.mjtWrap.mjWRAP_JOINT PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY SITE = mujoco.mjtWrap.mjWRAP_SITE @@ -219,6 +229,7 @@ class TrnType(enum.IntEnum): TENDON: force on tendon SITE: force on site """ + JOINT = mujoco.mjtTrn.mjTRN_JOINT JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT SITE = mujoco.mjtTrn.mjTRN_SITE @@ -236,6 +247,7 @@ class DynType(enum.IntEnum): FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration MUSCLE: piece-wise linear filter with two time constants """ + NONE = mujoco.mjtDyn.mjDYN_NONE INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR FILTER = mujoco.mjtDyn.mjDYN_FILTER @@ -252,6 +264,7 @@ class GainType(enum.IntEnum): AFFINE: const + kp*length + kv*velocity MUSCLE: muscle FLV curve computed by muscle_gain """ + FIXED = mujoco.mjtGain.mjGAIN_FIXED AFFINE = mujoco.mjtGain.mjGAIN_AFFINE MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE @@ -266,6 +279,7 @@ class BiasType(enum.IntEnum): AFFINE: const + kp*length + kv*velocity MUSCLE: muscle passive force computed by muscle_bias """ + NONE = mujoco.mjtBias.mjBIAS_NONE AFFINE = mujoco.mjtBias.mjBIAS_AFFINE MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE @@ -282,6 +296,7 @@ class ConstraintType(enum.IntEnum): CONTACT_FRICTIONLESS: frictionless contact CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone """ + EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON @@ -302,6 +317,7 @@ class CamLightType(enum.IntEnum): TARGETBODY: pos fixed in body, rot tracks target body TARGETBODYCOM: pos fixed in body, rot tracks target subtree com """ + FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM @@ -346,6 +362,7 @@ class SensorType(enum.IntEnum): FRAMELINACC: 3D linear acceleration FRAMEANGACC: 3D angular acceleration """ + MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER @@ -391,6 +408,7 @@ class ObjType(PyTreeNode): SITE: site CAMERA: camera """ + UNKNOWN = mujoco.mjtObj.mjOBJ_UNKNOWN BODY = mujoco.mjtObj.mjOBJ_BODY XBODY = mujoco.mjtObj.mjOBJ_XBODY @@ -437,7 +455,7 @@ class Option(PyTreeNode): disableactuator: bit flags for disabling actuators by group id (not used) sdf_initpoints: number of starting points for gradient descent (not used) sdf_iterations: max number of iterations for gradient descent (not used) - """ + """ # fmt: skip timestep: jax.Array apirate: jax.Array = _restricted_to('mujoco') impratio: jax.Array @@ -480,6 +498,7 @@ class Statistic(PyTreeNode): extent: spatial extent (not used) center: center of model (not used) """ + meaninertia: jax.Array meanmass: jax.Array meansize: jax.Array @@ -813,6 +832,7 @@ class Model(PyTreeNode): name_keyadr: keyframe name pointers (nkey,) names: names of all objects, 0-terminated (nnames,) """ + nq: int nv: int nu: int @@ -1161,7 +1181,7 @@ class Contact(PyTreeNode): geom2: id of geom 2; deprecated, use geom[1] geom: geom ids (2,) efc_address: address in efc; -1: not included - """ + """ # fmt: skip dist: jax.Array pos: jax.Array frame: jax.Array @@ -1306,7 +1326,7 @@ class Data(PyTreeNode): _qM_sparse: qM in sparse representation (nM,) _qLD_sparse: qLD in sparse representation (nM,) _qLDiagInv_sparse: qLDiagInv in sparse representation (nv,) - """ + """ # fmt: skip # constant sizes: ne: int nf: int diff --git a/mjx/mujoco/mjx/testspeed.py b/mjx/mujoco/mjx/testspeed.py index 9ccba845..4d59ad33 100644 --- a/mjx/mujoco/mjx/testspeed.py +++ b/mjx/mujoco/mjx/testspeed.py @@ -22,7 +22,9 @@ from etils import epath import mujoco from mujoco import mjx -_MJCF = flags.DEFINE_string('mjcf', None, 'path to model `.xml` or `.mjb`', required=True) +_MJCF = flags.DEFINE_string( + 'mjcf', None, 'path to model `.xml` or `.mjb`', required=True +) _BASE_PATH = flags.DEFINE_string( 'base_path', None, 'base path, defaults to mujoco.mjx resource path' ) diff --git a/mjx/mujoco/mjx/viewer.py b/mjx/mujoco/mjx/viewer.py index 7ad564c9..5cdce300 100644 --- a/mjx/mujoco/mjx/viewer.py +++ b/mjx/mujoco/mjx/viewer.py @@ -28,8 +28,9 @@ import mujoco.viewer _JIT = flags.DEFINE_bool('jit', True, 'To jit or not to jit.') -_MODEL_PATH = flags.DEFINE_string('mjcf', None, 'Path to a MuJoCo MJCF file.', - required=True) +_MODEL_PATH = flags.DEFINE_string( + 'mjcf', None, 'Path to a MuJoCo MJCF file.', required=True +) _VIEWER_GLOBAL_STATE = { diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml index a5370380..4d276d40 100644 --- a/mjx/pyproject.toml +++ b/mjx/pyproject.toml @@ -44,3 +44,24 @@ Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx" Documentation = "https://mujoco.readthedocs.io/en/3.2.7" Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx" Changelog = "https://mujoco.readthedocs.io/en/3.2.7/changelog.html" + +[tool.isort] +force_single_line = true +force_sort_within_sections = true +lexicographical = true +single_line_exclusions = ["typing"] +order_by_type = false +group_by_package = true +line_length = 120 +use_parentheses = true +multi_line_output = 3 +skip_glob = ["**/*.ipynb"] + +[tool.pyink] +line-length = 80 +unstable = true +pyink-indentation = 2 +pyink-use-majority-quotes = true +extend-exclude = '''( + .ipynb$ +)''' diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py index ffdfb75b..daed5d2c 100644 --- a/python/mujoco/bindings_test.py +++ b/python/mujoco/bindings_test.py @@ -120,13 +120,17 @@ class MuJoCoBindingsTest(parameterized.TestCase): xml_2 = rb"""""" xml_3 = rb"""""" 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('')) 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__': diff --git a/python/mujoco/memory_leak_test.py b/python/mujoco/memory_leak_test.py index 55667a7b..4a713085 100644 --- a/python/mujoco/memory_leak_test.py +++ b/python/mujoco/memory_leak_test.py @@ -57,6 +57,7 @@ class MemoryLeakTest(absltest.TestCase): soft = -1 try: import resource # pylint: disable=g-import-not-at-top + soft, hard = resource.getrlimit(resource.RLIMIT_AS) resource.setrlimit(resource.RLIMIT_AS, (limit_in_bytes, hard)) except (ImportError, ValueError): @@ -65,5 +66,5 @@ class MemoryLeakTest(absltest.TestCase): return soft -if __name__ == '__main__': +if __name__ == "__main__": absltest.main() diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py index 60ffae98..449d4bb8 100644 --- a/python/mujoco/minimize.py +++ b/python/mujoco/minimize.py @@ -209,7 +209,7 @@ def least_squares( # Decrease mu agressively: sequential decreases grow exponentially. def decrease_mu(mu, n_reduc): - dmu = (1/mu_factor) ** (2**n_reduc) + dmu = (1 / mu_factor) ** (2**n_reduc) mu = 0.0 if mu * dmu < mu_min else mu * dmu n_reduc += 1 return mu, n_reduc @@ -427,7 +427,6 @@ def jacobian_fd( Returns: jac: Jacobian of the residual at x. n_res: updated number of residual evaluations (add x.size). - """ n = x.size if bounds is None: @@ -438,7 +437,7 @@ def jacobian_fd( xh = x + np.diag(eps_vec) rh = residual(xh) jac = (rh - r) / eps_vec - return jac, n_res+n + return jac, n_res + n def check_jacobian( @@ -467,14 +466,15 @@ def check_jacobian( Returns: n_res: updated number of residual evaluations. - """ jac_fd, n_res = jacobian_fd(residual, x, r, eps, n_res, bounds) denom = np.abs(jac).sum() + np.abs(jac_fd).sum() + 1e-8 rel_diff = np.abs(jac - jac_fd) / denom if np.any(rel_diff > 1e-5): - raise ValueError(f'User-provided {name} does not match finite-differences ' - '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. diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py index 6f4cd3cc..5cbf1cfd 100644 --- a/python/mujoco/minimize_test.py +++ b/python/mujoco/minimize_test.py @@ -56,8 +56,9 @@ class MinimizeTest(absltest.TestCase): x0 = np.array((0.0, 0.0)) out = io.StringIO() - x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out, - 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() diff --git a/python/mujoco/msh2obj_test.py b/python/mujoco/msh2obj_test.py index cbba8d84..97166db4 100644 --- a/python/mujoco/msh2obj_test.py +++ b/python/mujoco/msh2obj_test.py @@ -63,7 +63,8 @@ class MshTest(absltest.TestCase): obj = msh2obj.msh_to_obj(msh_path) obj_model = mujoco.MjModel.from_xml_string( - _XML, {"abdomen_1_body.obj": obj.encode()}) + _XML, {"abdomen_1_body.obj": obj.encode()} + ) for field in _MESH_FIELDS: np.testing.assert_allclose( @@ -73,5 +74,6 @@ class MshTest(absltest.TestCase): err_msg=f"Field {field} does not match between msh and obj models.", ) + if __name__ == "__main__": absltest.main() diff --git a/python/mujoco/render_test.py b/python/mujoco/render_test.py index af5705ae..2df0a813 100644 --- a/python/mujoco/render_test.py +++ b/python/mujoco/render_test.py @@ -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) diff --git a/python/mujoco/renderer.py b/python/mujoco/renderer.py index afe5fc77..68175755 100644 --- a/python/mujoco/renderer.py +++ b/python/mujoco/renderer.py @@ -32,7 +32,7 @@ class Renderer: model: _structs.MjModel, height: int = 240, width: int = 320, - max_geom: int = 10000 + max_geom: int = 10000, ) -> None: """Initializes a new `Renderer`. @@ -43,6 +43,7 @@ class Renderer: max_geom: Optional integer specifying the maximum number of geoms that can be rendered in the same scene. If None this will be chosen automatically based on the estimated maximum number of renderable geoms in the model. + Raises: ValueError: If `camera_id` is outside the valid range, or if `width` or `height` exceed the dimensions of MuJoCo's offscreen framebuffer. @@ -220,9 +221,7 @@ the clause: # Convert 3-channel uint8 to 1-channel uint32. image3 = out.astype(np.uint32) segimage = ( - image3[:, :, 0] - + image3[:, :, 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, ) diff --git a/python/mujoco/renderer_test.py b/python/mujoco/renderer_test.py index e8022006..bb5f62f1 100644 --- a/python/mujoco/renderer_test.py +++ b/python/mujoco/renderer_test.py @@ -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 = """ diff --git a/python/mujoco/rollout.py b/python/mujoco/rollout.py index f314e795..98eaa3f2 100644 --- a/python/mujoco/rollout.py +++ b/python/mujoco/rollout.py @@ -23,17 +23,19 @@ import numpy as np from numpy import typing as npt -def rollout(model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]], - 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. diff --git a/python/mujoco/rollout_test.py b/python/mujoco/rollout_test.py index 32c670e6..3cc0d062 100644 --- a/python/mujoco/rollout_test.py +++ b/python/mujoco/rollout_test.py @@ -127,10 +127,12 @@ TEST_XML_DIVERGE = r""" """ -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) diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py index e18db9b7..bd9997e1 100644 --- a/python/mujoco/specs_test.py +++ b/python/mujoco/specs_test.py @@ -104,7 +104,9 @@ class SpecsTest(absltest.TestCase): self.assertEqual(model.nuser_site, 6) np.testing.assert_array_equal(model.site_user[0], [1, 2, 3, 4, 5, 6]) - self.assertEqual(spec.to_xml(), textwrap.dedent("""\ + self.assertEqual( + spec.to_xml(), + textwrap.dedent("""\ @@ -116,7 +118,8 @@ class SpecsTest(absltest.TestCase): - """),) + """), + ) 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("""\ @@ -547,7 +552,8 @@ class SpecsTest(absltest.TestCase): - """)) + """), + ) 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("""\ @@ -577,7 +585,8 @@ class SpecsTest(absltest.TestCase): - """)) + """), + ) def test_element_list(self): spec = mujoco.MjSpec() @@ -718,13 +727,17 @@ class SpecsTest(absltest.TestCase): """ - spec = mujoco.MjSpec.from_string(textwrap.dedent(""" + spec = mujoco.MjSpec.from_string( + textwrap.dedent(""" - """), {'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" diff --git a/python/mujoco/viewer.py b/python/mujoco/viewer.py index 28c55e44..31ed22a8 100644 --- a/python/mujoco/viewer.py +++ b/python/mujoco/viewer.py @@ -42,7 +42,7 @@ PERCENT_REALTIME = ( 10, 8, 6.6, 5, 4, 3.3, 2.5, 2, 1.6, 1.3, 1, 0.8, 0.66, 0.5, 0.4, 0.33, 0.25, 0.2, 0.16, 0.13, 0.1 -) +) # fmt: skip # Maximum time mis-alignment before re-sync. MAX_SYNC_MISALIGN = 0.1 @@ -194,12 +194,13 @@ def _file_loader(path: str) -> _LoaderWithPathType: def _reload( - simulate: _Simulate, loader: _InternalLoaderType, - 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, diff --git a/python/pyproject.toml b/python/pyproject.toml index 94067175..600229f4 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -65,3 +65,24 @@ usd = [ "usd-core", "pillow" ] + +[tool.isort] +force_single_line = true +force_sort_within_sections = true +lexicographical = true +single_line_exclusions = ["typing"] +order_by_type = false +group_by_package = true +line_length = 120 +use_parentheses = true +multi_line_output = 3 +skip_glob = ["**/*.ipynb"] + +[tool.pyink] +line-length = 80 +unstable = true +pyink-indentation = 2 +pyink-use-majority-quotes = true +extend-exclude = '''( + .ipynb$ +)''' diff --git a/python/setup.py b/python/setup.py index 8bac24ff..cebdec84 100644 --- a/python/setup.py +++ b/python/setup.py @@ -101,15 +101,15 @@ def tokenize_quoted_substr(input_string, quote_char, placeholders=None): placeholders = placeholders if placeholders is not None else dict() prev_end = -1 for start, end in start_and_end(quote_positions): - output_string += input_string[prev_end+1:start] + output_string += input_string[prev_end + 1 : start] while True: placeholder = ''.join(random.choices(string.ascii_lowercase, k=5)) if placeholder not in input_string and placeholder not in output_string: break output_string += placeholder - placeholders[placeholder] = input_string[start+1:end] + placeholders[placeholder] = input_string[start + 1 : end] prev_end = end - output_string += input_string[prev_end+1:] + output_string += input_string[prev_end + 1 :] return output_string, placeholders @@ -145,15 +145,17 @@ class BuildCMakeExtension(build_ext.build_ext): """Uses CMake to build extensions.""" def run(self): - self._is_apple = (platform.system() == 'Darwin') - (self._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 [], )