chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add pyink and isort config. Reformat.

Browse Source 
PiperOrigin-RevId: 704533915
Change-Id: I37e9fd51261bd166b725c7460fc65d02fed2b391
This commit is contained in:
Baruch Tabanpour
2024-12-09 21:10:00 -08:00
committed by Copybara-Service
parent 6f6244b739
commit f3b3024291
41 changed files with 895 additions and 480 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/collision_convex.py
+9 -5
View File
@@ -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
mjx/mujoco/mjx/_src/collision_driver.py
+3 -3
View File
@@ -146,13 +146,13 @@ def geom_pairs(
b_end = b_start + m.body_geomnum
for b1 in range(m.nbody):
if not geom_con[b_start[b1]:b_end[b1]].any():
if not geom_con[b_start[b1] : b_end[b1]].any():
continue
w1 = m.body_weldid[b1]
w1_p = m.body_weldid[m.body_parentid[w1]]
for b2 in range(b1, m.nbody):
if not geom_con[b_start[b2]:b_end[b2]].any():
if not geom_con[b_start[b2] : b_end[b2]].any():
continue
signature = (b1 << 16) + (b2)
if signature in exclude_signature:
@@ -272,7 +272,7 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
jp.clip(m.pair_friction[ip], a_min=eps),
m.pair_solref[ip],
m.pair_solreffriction[ip],
m.pair_solimp[ip]
m.pair_solimp[ip],
))
if geom1.size > 0 and geom2.size > 0:
# other contacts get their params from geom fields
mjx/mujoco/mjx/_src/collision_driver_test.py
+19 -9
View File
@@ -218,7 +218,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5)
dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
)
_ELLIPSOID_ELLIPSOID = """
<mujoco>
@@ -240,7 +241,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5)
dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-5
)
_ELLIPSOID_SPHERE = """
<mujoco>
@@ -263,7 +265,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-3)
dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-3
)
_ELLIPSOID_CAPSULE = """
<mujoco>
@@ -285,7 +288,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-3)
dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-3
)
_ELLIPSOID_CYLINDER = """
<mujoco>
@@ -308,7 +312,8 @@ class EllipsoidCollisionTest(parameterized.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4)
dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
)
class CapsuleCollisionTest(parameterized.TestCase):
@@ -550,7 +555,8 @@ class CylinderTest(absltest.TestCase):
# cylinder is vertical
xml = self._CYLINDER_PLANE.replace(
'<geom fromto="-0.1 0 0 0.1 0 0"', '<geom fromto="0 0 -0.1 0 0 0.1"')
'<geom fromto="-0.1 0 0 0.1 0 0"', '<geom fromto="0 0 -0.1 0 0 0.1"'
)
xml = xml.replace('pos="0 0 0.04"', 'pos="0 0 0.095"')
d, dx = _collide(xml)
@@ -579,7 +585,8 @@ class CylinderTest(absltest.TestCase):
self.assertLess(dx.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4)
dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
)
class ConvexTest(absltest.TestCase):
@@ -760,11 +767,14 @@ class ConvexTest(absltest.TestCase):
np.testing.assert_array_less(0, c.dist[1:])
np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1]))
np.testing.assert_array_almost_equal(
c.pos[0], np.array([0, 2, 1.3155]), decimal=5)
c.pos[0], np.array([0, 2, 1.3155]), decimal=5
)
_, dx = _collide(
self._CONVEX_CONVEX_THIN.replace(
'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'))
'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'
)
)
c = dx.contact
self.assertTrue((c.dist > 0).all())
mjx/mujoco/mjx/_src/collision_primitive.py
+11 -6
View File
@@ -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
mjx/mujoco/mjx/_src/collision_sdf.py
+17 -12
View File
@@ -41,6 +41,7 @@ SDFFn = Callable[[jax.Array], jax.Array]
def collider(ncon: int):
"""Wraps collision functions for use by collision_driver."""
def wrapper(func):
def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
g1, g2 = geom.T
@@ -81,7 +82,7 @@ def _capsule(pos: jax.Array, size: jax.Array):
def _ellipsoid(pos: jax.Array, size: jax.Array) -> jax.Array:
k0 = math.norm(pos / size)
k1 = math.norm(pos / (size*size))
k1 = math.norm(pos / (size * size))
return k0 * (k0 - 1.0) / (k1 + (k1 == 0.0) * 1e-12)
@@ -96,12 +97,12 @@ def _cylinder(pos: jax.Array, size: jax.Array) -> jax.Array:
def _cylinder_grad(x: jax.Array, size: jax.Array) -> jax.Array:
"""Gradient of the cylinder SDF wrt query point and singularities removed."""
c = jp.sqrt(x[0]*x[0]+x[1]*x[1])
c = jp.sqrt(x[0] * x[0] + x[1] * x[1])
e = jp.abs(x[2])
a = jp.array([c - size[0], e - size[1]])
b = jp.array([jp.maximum(a[0], 0), jp.maximum(a[1], 0)])
j = jp.argmax(a)
bnorm = jp.sqrt(b[0]*b[0] + b[1]*b[1])
bnorm = jp.sqrt(b[0] * b[0] + b[1] * b[1])
bnorm += jp.allclose(bnorm, 0) * 1e-12
grada = jp.array([
x[0] / (c + jp.allclose(c, 0) * 1e-12),
@@ -151,7 +152,7 @@ def _gradient_step(objective: SDFFn, state: GradientState) -> GradientState:
"""Performs a step of gradient descent."""
# TODO: find better parameters
amin = 1e-4 # minimum value for line search factor scaling the gradient
amax = 2. # maximum value for line search factor scaling the gradient
amax = 2.0 # maximum value for line search factor scaling the gradient
nlinesearch = 10 # line search points
grad = jax.grad(objective)(state.x)
alpha = jp.geomspace(amin, amax, nlinesearch).reshape(nlinesearch, -1)
@@ -179,7 +180,11 @@ def _gradient_descent(
def _optim(
d1, d2, info1: GeomInfo, info2: GeomInfo, x0: jax.Array,
d1,
d2,
info1: GeomInfo,
info2: GeomInfo,
x0: jax.Array,
) -> Collision:
"""Optimizes the clearance function."""
d1 = functools.partial(d1, size=info1.size)
@@ -198,14 +203,14 @@ def _optim(
@collider(ncon=1)
def sphere_ellipsoid(s: GeomInfo, e: GeomInfo) -> Collision:
""""Calculates contact between a sphere and an ellipsoid."""
"""Calculates contact between a sphere and an ellipsoid."""
x0 = 0.5 * (s.pos + e.pos)
return _optim(_sphere, _ellipsoid, s, e, x0)
@collider(ncon=1)
def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision:
""""Calculates contact between a sphere and a cylinder."""
"""Calculates contact between a sphere and a cylinder."""
# TODO: implement analytical version.
x0 = 0.5 * (s.pos + c.pos)
return _optim(_sphere, _cylinder, s, c, x0)
@@ -213,14 +218,14 @@ def sphere_cylinder(s: GeomInfo, c: GeomInfo) -> Collision:
@collider(ncon=1)
def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
""""Calculates contact between a capsule and an ellipsoid."""
""" "Calculates contact between a capsule and an ellipsoid."""
x0 = 0.5 * (c.pos + e.pos)
return _optim(_capsule, _ellipsoid, c, e, x0)
@collider(ncon=2)
def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision:
""""Calculates contact between a capsule and a cylinder."""
"""Calculates contact between a capsule and a cylinder."""
# TODO: improve robustness
# Near sharp corners, the SDF might give the penetration depth with respect
# to a surface that is not in collision. Possible solutions is to find the
@@ -235,21 +240,21 @@ def capsule_cylinder(ca: GeomInfo, cy: GeomInfo) -> Collision:
@collider(ncon=1)
def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision:
""""Calculates contact between two ellipsoids."""
"""Calculates contact between two ellipsoids."""
x0 = 0.5 * (e1.pos + e2.pos)
return _optim(_ellipsoid, _ellipsoid, e1, e2, x0)
@collider(ncon=1)
def ellipsoid_cylinder(e: GeomInfo, c: GeomInfo) -> Collision:
""""Calculates contact between and ellipsoid and a cylinder."""
"""Calculates contact between and ellipsoid and a cylinder."""
x0 = 0.5 * (e.pos + c.pos)
return _optim(_ellipsoid, _cylinder, e, c, x0)
@collider(ncon=4)
def cylinder_cylinder(c1: GeomInfo, c2: GeomInfo) -> Collision:
""""Calculates contact between a cylinder and a cylinder."""
"""Calculates contact between a cylinder and a cylinder."""
# TODO: improve robustness
# Near sharp corners, the SDF might give the penetration depth with respect
# to a surface that is not in collision. Possible solutions is to find the
mjx/mujoco/mjx/_src/collision_types.py
+1
View File
@@ -73,6 +73,7 @@ class FunctionKey:
resulting constraint jacobian is determined at compile time.
subgrid_size: the size determines the hfield subgrid to collide with
"""
types: Tuple[int, int]
data_ids: Tuple[int, int]
condim: int
mjx/mujoco/mjx/_src/constraint.py
+1
View File
@@ -39,6 +39,7 @@ import numpy as np
class _Efc(PyTreeNode):
"""Support data for creating constraint matrices."""
J: jax.Array
pos_aref: jax.Array
pos_imp: jax.Array
mjx/mujoco/mjx/_src/constraint_test.py
+2 -2
View File
@@ -30,7 +30,7 @@ _TOLERANCE = 5e-5
def _assert_eq(a, b, name):
tol = _TOLERANCE * 10 # avoid test noise
tol = _TOLERANCE * 10 # avoid test noise
err_msg = f'mismatch: {name}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
@@ -75,7 +75,7 @@ class ConstraintTest(parameterized.TestCase):
_assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
_assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
_assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
_assert_eq(dx.efc_pos[order][d.nefc:], 0, 'efc_pos')
_assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos')
_assert_eq(
d.efc_frictionloss,
dx.efc_frictionloss[order][: d.nefc],
mjx/mujoco/mjx/_src/dataclasses.py
+1 -2
View File
@@ -16,7 +16,6 @@
import copy
import dataclasses
import typing
from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
import jax
@@ -57,7 +56,7 @@ def dataclass(clz: _T) -> _T:
meta_fields.append(field)
def replace(self, **updates):
""""Returns a new object replacing the specified fields with new values."""
"""Returns a new object replacing the specified fields with new values."""
return dataclasses.replace(self, **updates)
data_clz.replace = replace
mjx/mujoco/mjx/_src/io.py
+8 -7
View File
@@ -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
mjx/mujoco/mjx/_src/io_test.py
+2 -3
View File
@@ -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
mjx/mujoco/mjx/_src/math.py
+1
View File
@@ -28,6 +28,7 @@ def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array:
Args:
a: left hand of matmul operand
b: right hand of matmul operand
Returns:
the matrix product of the inputs.
"""
mjx/mujoco/mjx/_src/mesh.py
+2 -2
View File
@@ -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 '
mjx/mujoco/mjx/_src/mesh_test.py
+3 -2
View File
@@ -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,
mjx/mujoco/mjx/_src/passive.py
+1
View File
@@ -31,6 +31,7 @@ from mujoco.mjx._src.types import Model
def _spring_damper(m: Model, d: Data) -> jax.Array:
"""Applies joint level spring and damping forces."""
def fn(jnt_typs, stiffness, qpos_spring, qpos):
qpos_i = 0
qfrcs = []
mjx/mujoco/mjx/_src/ray.py
+1 -1
View File
@@ -269,7 +269,7 @@ def ray(
geom_filter_dyn = (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
geom_filter_dyn &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
for geom_type, fn in _RAY_FUNC.items():
id_, = np.nonzero(geom_filter & (m.geom_type == geom_type))
(id_,) = np.nonzero(geom_filter & (m.geom_type == geom_type))
if id_.size == 0:
continue
mjx/mujoco/mjx/_src/scan.py
+11 -15
View File
@@ -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':
mjx/mujoco/mjx/_src/scan_test.py
+3
View File
@@ -90,6 +90,7 @@ class ScanTest(absltest.TestCase):
if tuple(jnt_types) == (JointType.FREE,):
return None
return val + sum(jnt_types)
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
@@ -99,6 +100,7 @@ class ScanTest(absltest.TestCase):
if jnt_types.size == 0:
self.fail('world has no dofs, should not be called')
return val + sum(jnt_types)
v_in = jp.ones((m.nv, 1))
scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in)
@@ -141,6 +143,7 @@ class ScanTest(absltest.TestCase):
return None
carry = jp.zeros_like(val) if carry is None else carry
return carry + val + sum(jnt_types)
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
mjx/mujoco/mjx/_src/sensor_test.py
-2
View File
@@ -17,13 +17,11 @@
from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
from mujoco.mjx._src.types import ConeType
import numpy as np
# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
mjx/mujoco/mjx/_src/smooth.py
+2 -2
View File
@@ -334,7 +334,7 @@ def factor_m(m: Model, d: Data) -> Data:
pivots = []
out = []
for (b, e, madr_d, madr_ij) in updates:
for b, e, madr_d, madr_ij in updates:
width = e - b
rows.append(np.arange(madr_ij, madr_ij + width))
madr_ijs.append(np.full((width,), madr_ij))
@@ -511,7 +511,6 @@ def subtree_vel(m: Model, d: Data) -> Data:
angmom_child, mom_parent_child = carry
return angmom + mom + angmom_child + mom_parent_child, mom_parent
subtree_angmom, _ = scan.body_tree(
m,
_subtree_angmom,
@@ -535,6 +534,7 @@ def subtree_vel(m: Model, d: Data) -> Data:
def rne(m: Model, d: Data) -> Data:
"""Computes inverse dynamics using the recursive Newton-Euler algorithm."""
# forward scan over tree: accumulate link center of mass acceleration
def cacc_fn(cacc, cdof_dot, qvel):
if cacc is None:
mjx/mujoco/mjx/_src/solver.py
+3 -1
View File
@@ -52,6 +52,7 @@ class _Context(PyTreeNode):
u: friction cone (normal and tangents) (num(con.dim > 1), 6)
h: cone hessian (num(con.dim > 1), 6, 6)
"""
qacc: jax.Array
qfrc_constraint: jax.Array
Jaref: jax.Array # pylint: disable=invalid-name
@@ -225,6 +226,7 @@ class _LSContext(PyTreeNode):
def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
"""Scan-based implementation (jit ok, reverse-mode autodiff ok)."""
def _iter(val):
next_val = body_fun(val)
next_cond = cond_fun(next_val)
@@ -382,7 +384,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
# set efc of cone H along diagonal
for i, (condim, addr) in enumerate(zip(dim, efc_address)):
h_cone = ctx.h[i, :condim, :condim]
cm = cm.at[addr:addr+condim, addr:addr+condim].add(h_cone)
cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone)
h = d.efc_J.T @ cm @ d.efc_J
else:
h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J
mjx/mujoco/mjx/_src/test_util.py
+1 -1
View File
@@ -116,7 +116,7 @@ def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray:
if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
num_rows = (dx.contact.dim[i] - 1) * 2
if dx.contact.dist[i] > 0: # move empty contacts to end
order = np.append(order, np.repeat(2 ** 16, num_rows))
order = np.append(order, np.repeat(2**16, num_rows))
continue
contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1)
contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1)
mjx/mujoco/mjx/_src/types.py
+23 -3
View File
@@ -47,6 +47,7 @@ class DisableBit(enum.IntFlag):
REFSAFE: integrator safety: make ref[0]>=2*timestep
SENSOR: sensors
"""
CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY
FRICTIONLOSS = mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS
@@ -73,6 +74,7 @@ class JointType(enum.IntEnum):
SLIDE: sliding distance along body-fixed axis (1,)
HINGE: rotation angle (rad) around body-fixed axis (1,)
"""
FREE = mujoco.mjtJoint.mjJNT_FREE
BALL = mujoco.mjtJoint.mjJNT_BALL
SLIDE = mujoco.mjtJoint.mjJNT_SLIDE
@@ -93,6 +95,7 @@ class IntegratorType(enum.IntEnum):
RK4: 4th-order Runge Kutta
IMPLICITFAST: implicit in velocity, no rne derivative
"""
EULER = mujoco.mjtIntegrator.mjINT_EULER
RK4 = mujoco.mjtIntegrator.mjINT_RK4
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
@@ -113,6 +116,7 @@ class GeomType(enum.IntEnum):
MESH: mesh
SDF: signed distance field
"""
PLANE = mujoco.mjtGeom.mjGEOM_PLANE
HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD
SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE
@@ -134,6 +138,7 @@ class ConvexMesh(PyTreeNode):
edge: edge indexes for all edges in the convex mesh
edge_face_normal: indexes for face normals adjacent to edges in `edge`
"""
vert: jax.Array
face: jax.Array
face_normal: jax.Array
@@ -148,6 +153,7 @@ class ConeType(enum.IntEnum):
PYRAMIDAL: pyramidal
ELLIPTIC: elliptic
"""
PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL
ELLIPTIC = mujoco.mjtCone.mjCONE_ELLIPTIC
@@ -160,6 +166,7 @@ class JacobianType(enum.IntEnum):
SPARSE: sparse
AUTO: sparse if nv>60 and device is TPU, dense otherwise
"""
DENSE = mujoco.mjtJacobian.mjJAC_DENSE
SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
AUTO = mujoco.mjtJacobian.mjJAC_AUTO
@@ -172,6 +179,7 @@ class SolverType(enum.IntEnum):
CG: Conjugate gradient (primal)
NEWTON: Newton (primal)
"""
# unsupported: PGS
CG = mujoco.mjtSolver.mjSOL_CG
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
@@ -186,6 +194,7 @@ class EqType(enum.IntEnum):
JOINT: couple the values of two scalar joints with cubic
TENDON: couple the lengths of two tendons with cubic
"""
CONNECT = mujoco.mjtEq.mjEQ_CONNECT
WELD = mujoco.mjtEq.mjEQ_WELD
JOINT = mujoco.mjtEq.mjEQ_JOINT
@@ -203,6 +212,7 @@ class WrapType(enum.IntEnum):
SPHERE: wrap around sphere
CYLINDER: wrap around (infinite) cylinder
"""
JOINT = mujoco.mjtWrap.mjWRAP_JOINT
PULLEY = mujoco.mjtWrap.mjWRAP_PULLEY
SITE = mujoco.mjtWrap.mjWRAP_SITE
@@ -219,6 +229,7 @@ class TrnType(enum.IntEnum):
TENDON: force on tendon
SITE: force on site
"""
JOINT = mujoco.mjtTrn.mjTRN_JOINT
JOINTINPARENT = mujoco.mjtTrn.mjTRN_JOINTINPARENT
SITE = mujoco.mjtTrn.mjTRN_SITE
@@ -236,6 +247,7 @@ class DynType(enum.IntEnum):
FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration
MUSCLE: piece-wise linear filter with two time constants
"""
NONE = mujoco.mjtDyn.mjDYN_NONE
INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
FILTER = mujoco.mjtDyn.mjDYN_FILTER
@@ -252,6 +264,7 @@ class GainType(enum.IntEnum):
AFFINE: const + kp*length + kv*velocity
MUSCLE: muscle FLV curve computed by muscle_gain
"""
FIXED = mujoco.mjtGain.mjGAIN_FIXED
AFFINE = mujoco.mjtGain.mjGAIN_AFFINE
MUSCLE = mujoco.mjtGain.mjGAIN_MUSCLE
@@ -266,6 +279,7 @@ class BiasType(enum.IntEnum):
AFFINE: const + kp*length + kv*velocity
MUSCLE: muscle passive force computed by muscle_bias
"""
NONE = mujoco.mjtBias.mjBIAS_NONE
AFFINE = mujoco.mjtBias.mjBIAS_AFFINE
MUSCLE = mujoco.mjtBias.mjBIAS_MUSCLE
@@ -282,6 +296,7 @@ class ConstraintType(enum.IntEnum):
CONTACT_FRICTIONLESS: frictionless contact
CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
"""
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
FRICTION_DOF = mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF
FRICTION_TENDON = mujoco.mjtConstraint.mjCNSTR_FRICTION_TENDON
@@ -302,6 +317,7 @@ class CamLightType(enum.IntEnum):
TARGETBODY: pos fixed in body, rot tracks target body
TARGETBODYCOM: pos fixed in body, rot tracks target subtree com
"""
FIXED = mujoco.mjtCamLight.mjCAMLIGHT_FIXED
TRACK = mujoco.mjtCamLight.mjCAMLIGHT_TRACK
TRACKCOM = mujoco.mjtCamLight.mjCAMLIGHT_TRACKCOM
@@ -346,6 +362,7 @@ class SensorType(enum.IntEnum):
FRAMELINACC: 3D linear acceleration
FRAMEANGACC: 3D angular acceleration
"""
MAGNETOMETER = mujoco.mjtSensor.mjSENS_MAGNETOMETER
CAMPROJECTION = mujoco.mjtSensor.mjSENS_CAMPROJECTION
RANGEFINDER = mujoco.mjtSensor.mjSENS_RANGEFINDER
@@ -391,6 +408,7 @@ class ObjType(PyTreeNode):
SITE: site
CAMERA: camera
"""
UNKNOWN = mujoco.mjtObj.mjOBJ_UNKNOWN
BODY = mujoco.mjtObj.mjOBJ_BODY
XBODY = mujoco.mjtObj.mjOBJ_XBODY
@@ -437,7 +455,7 @@ class Option(PyTreeNode):
disableactuator: bit flags for disabling actuators by group id (not used)
sdf_initpoints: number of starting points for gradient descent (not used)
sdf_iterations: max number of iterations for gradient descent (not used)
"""
""" # fmt: skip
timestep: jax.Array
apirate: jax.Array = _restricted_to('mujoco')
impratio: jax.Array
@@ -480,6 +498,7 @@ class Statistic(PyTreeNode):
extent: spatial extent (not used)
center: center of model (not used)
"""
meaninertia: jax.Array
meanmass: jax.Array
meansize: jax.Array
@@ -813,6 +832,7 @@ class Model(PyTreeNode):
name_keyadr: keyframe name pointers (nkey,)
names: names of all objects, 0-terminated (nnames,)
"""
nq: int
nv: int
nu: int
@@ -1161,7 +1181,7 @@ class Contact(PyTreeNode):
geom2: id of geom 2; deprecated, use geom[1]
geom: geom ids (2,)
efc_address: address in efc; -1: not included
"""
""" # fmt: skip
dist: jax.Array
pos: jax.Array
frame: jax.Array
@@ -1306,7 +1326,7 @@ class Data(PyTreeNode):
_qM_sparse: qM in sparse representation (nM,)
_qLD_sparse: qLD in sparse representation (nM,)
_qLDiagInv_sparse: qLDiagInv in sparse representation (nv,)
"""
""" # fmt: skip
# constant sizes:
ne: int
nf: int
mjx/mujoco/mjx/testspeed.py
+3 -1
View File
@@ -22,7 +22,9 @@ from etils import epath
import mujoco
from mujoco import mjx
_MJCF = flags.DEFINE_string('mjcf', None, 'path to model `.xml` or `.mjb`', required=True)
_MJCF = flags.DEFINE_string(
'mjcf', None, 'path to model `.xml` or `.mjb`', required=True
)
_BASE_PATH = flags.DEFINE_string(
'base_path', None, 'base path, defaults to mujoco.mjx resource path'
)
mjx/mujoco/mjx/viewer.py
+3 -2
View File
@@ -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 = {
mjx/pyproject.toml
+21
View File
@@ -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$
)'''