From a4df912018d32ea56eeaf6baa946785d59b859bf Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Thu, 18 Apr 2024 12:42:10 -0700 Subject: [PATCH] Prepare MJX for condim. This is a refactor of collision_driver and some of its surrounding code in order to prepare for condim in MJX. In this change we reify types that are needed for condim: dim, efc_address, efc_type. We make explicit the way contacts are organized and grouped to guarantee that dim and efc_type are statically defined. This change simplifies the way meshes are organized on device and slightly speeds up mesh collisions for cases where a single mesh is instanced across many geoms. PiperOrigin-RevId: 626119500 Change-Id: Ic0c8599bcda2326f2e19cd3246a673e56097886b --- doc/changelog.rst | 15 +- mjx/mujoco/mjx/__init__.py | 3 - mjx/mujoco/mjx/_src/collision_base.py | 67 -- mjx/mujoco/mjx/_src/collision_convex.py | 84 ++- mjx/mujoco/mjx/_src/collision_driver.py | 606 ++++++++---------- mjx/mujoco/mjx/_src/collision_driver_test.py | 36 +- mjx/mujoco/mjx/_src/collision_primitive.py | 99 +-- mjx/mujoco/mjx/_src/collision_sdf.py | 43 +- mjx/mujoco/mjx/_src/collision_types.py | 65 ++ mjx/mujoco/mjx/_src/constraint.py | 71 +- mjx/mujoco/mjx/_src/constraint_test.py | 6 +- mjx/mujoco/mjx/_src/device.py | 45 +- mjx/mujoco/mjx/_src/io.py | 232 +++---- mjx/mujoco/mjx/_src/io_test.py | 7 +- mjx/mujoco/mjx/_src/mesh.py | 247 +++---- mjx/mujoco/mjx/_src/mesh_test.py | 27 +- mjx/mujoco/mjx/_src/solver.py | 21 +- mjx/mujoco/mjx/_src/support.py | 5 + mjx/mujoco/mjx/_src/types.py | 95 +-- .../integration_test/collision_driver_test.py | 8 +- 20 files changed, 868 insertions(+), 914 deletions(-) delete mode 100644 mjx/mujoco/mjx/_src/collision_base.py create mode 100644 mjx/mujoco/mjx/_src/collision_types.py diff --git a/doc/changelog.rst b/doc/changelog.rst index fb59d33f..69851a63 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -15,13 +15,20 @@ General MJX ^^^ 3. Added cylinder plane collisions. +4. Added ``efc_type`` to ``mjx.Data`` and ``dim``, ``efc_address`` to ``mjx.Contact``. +5. Added ``geom`` to ``mjx.Contact`` and marked ``geom1``, ``geom2`` deprecated. +6. Added ``ne``, ``nf``, ``nl``, ``nefc``, and ``ncon`` to ``mjx.Data`` to match ``mujoco.MjData``. +7. Given the above added fields, removed ``mjx.get_params``, ``mjx.ncon``, and ``mjx.count_constraints``. +8. Changed the way meshes are organized on device to speed up collision detection when a mesh is replicated for many + geoms. +9. Fixed a bug where capsules might be ignored in broadphase colliision checking. Bug fixes ^^^^^^^^^ -4. Defaults of lights were not being saved, now fixed. -5. Prevent overwriting of frame names by body names when saving an XML. Introduced in 3.1.4. -6. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually - not optional. +10. Defaults of lights were not being saved, now fixed. +11. Prevent overwriting of frame names by body names when saving an XML. Introduced in 3.1.4. +12. Fixed bug in Python binding of :ref:`mj_saveModel`: ``buffer`` argument was documented as optional but was actually + not optional. Version 3.1.4 (April 10th, 2024) diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 978f44ab..01939ab7 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -16,9 +16,6 @@ # pylint:disable=g-importing-member from mujoco.mjx._src.collision_driver import collision -from mujoco.mjx._src.collision_driver import get_params -from mujoco.mjx._src.collision_driver import ncon -from mujoco.mjx._src.constraint import count_constraints from mujoco.mjx._src.constraint import make_constraint from mujoco.mjx._src.device import device_get_into from mujoco.mjx._src.device import device_put diff --git a/mjx/mujoco/mjx/_src/collision_base.py b/mjx/mujoco/mjx/_src/collision_base.py deleted file mode 100644 index 423200ac..00000000 --- a/mjx/mujoco/mjx/_src/collision_base.py +++ /dev/null @@ -1,67 +0,0 @@ -# Copyright 2023 DeepMind Technologies Limited -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================== -"""Collision base.""" - -import dataclasses -from typing import Dict, List, Optional, Tuple - -import jax -# pylint: disable=g-importing-member -from mujoco.mjx._src.dataclasses import PyTreeNode -from mujoco.mjx._src.types import GeomType -# pylint: enable=g-importing-member - -Contact = Tuple[jax.Array, jax.Array, jax.Array] - - -@dataclasses.dataclass -class Candidate: - geom1: int - geom2: int - ipair: int - geomp: int # priority geom - dim: int - - -CandidateSet = Dict[ - Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]], - List[Candidate], -] - - -class GeomInfo(PyTreeNode): - """Collision info for a geom.""" - - geom_id: jax.Array - pos: jax.Array - mat: jax.Array - size: jax.Array - face: Optional[jax.Array] = None - vert: Optional[jax.Array] = None - edge_dir: Optional[jax.Array] = None - facenorm: Optional[jax.Array] = None - edge: Optional[jax.Array] = None - edge_face_normal: Optional[jax.Array] = None - - -class SolverParams(PyTreeNode): - """Contact solver params.""" - - friction: jax.Array - solref: jax.Array - solreffriction: jax.Array - solimp: jax.Array - margin: jax.Array - gap: jax.Array diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 193d999f..3c3efd5c 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -20,12 +20,53 @@ from typing import Tuple import jax from jax import numpy as jp from mujoco.mjx._src import math +from mujoco.mjx._src import mesh # pylint: disable=g-importing-member -from mujoco.mjx._src.collision_base import Contact -from mujoco.mjx._src.collision_base import GeomInfo +from mujoco.mjx._src.collision_types import Collision +from mujoco.mjx._src.collision_types import ConvexInfo +from mujoco.mjx._src.collision_types import FunctionKey +from mujoco.mjx._src.collision_types import GeomInfo +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member +def collider(ncon: int): + """Wraps collision functions for use by collision_driver.""" + + def wrapper(func): + def collide( + m: Model, d: Data, key: FunctionKey, geom: jax.Array + ) -> Collision: + g1, g2 = geom.T + infos = [ + GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]), + GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]), + ] + in_axes = [0, 0] + for i in [0, 1]: + if key.types[i] == GeomType.BOX: + infos[i] = mesh.box(infos[i]) + in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( + pos=0, mat=0, face=0, vert=0 + ) + elif key.types[i] == GeomType.MESH: + infos[i] = mesh.convex(m, key.data_ids[i], infos[i]) + in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( + pos=0, mat=0 + ) + dist, pos, frame = jax.vmap(func, in_axes=in_axes)(*infos) + if ncon > 1: + return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) + return dist, pos, frame + + collide.ncon = ncon + return collide + + return wrapper + + def _closest_segment_point_plane( a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array ) -> jax.Array: @@ -178,7 +219,8 @@ def _manifold_points( return jp.array([a_idx, b_idx, c_idx, d_idx]) -def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact: +@collider(ncon=4) +def plane_convex(plane: GeomInfo, convex: ConvexInfo) -> Collision: """Calculates contacts between a plane and a convex object.""" vert = convex.vert @@ -200,10 +242,11 @@ def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact: return dist, pos, frame -def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact: +@collider(ncon=1) +def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: """Calculates contact between a sphere and a convex object.""" faces = convex.face - normals = convex.facenorm + normals = convex.face_normal # Put sphere in convex frame. sphere_pos = convex.mat.T @ (sphere.pos - convex.pos) @@ -262,16 +305,15 @@ def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact: n = convex.mat @ n pos = convex.mat @ pos + convex.pos - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: +@collider(ncon=2) +def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision: """Calculates contacts between a capsule and a convex object.""" # Get convex transformed normals, faces, and vertices. faces = convex.face - normals = convex.facenorm + normals = convex.face_normal # Put capsule in convex frame. cap_pos = convex.mat.T @ (cap.pos - convex.pos) @@ -869,10 +911,9 @@ def _sat_gaussmap( return dist, pos, normal -def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact: +@collider(ncon=4) +def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: """Calculates contacts between two convex objects.""" - if c1.face is None or c2.face is None or c1.vert is None or c2.vert is None: - raise AssertionError('Mesh info missing.') # pad face vertices so that we can broadcast between geom1 and geom2 # face has shape (n_face, n_vert, 3) nvert1, nvert2 = c1.face.shape[1], c2.face.shape[1] @@ -895,14 +936,14 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact: to_local_mat = c2.mat.T @ c1.mat faces1 = to_local_pos + faces1 @ to_local_mat.T - normals1 = c1.facenorm @ to_local_mat.T - normals2 = c2.facenorm + normals1 = c1.face_normal @ to_local_mat.T + normals2 = c2.face_normal vertices1 = to_local_pos + c1.vert @ to_local_mat.T vertices2 = c2.vert - unique_edges1 = jp.take(vertices1, c1.edge, axis=0) - unique_edges2 = jp.take(vertices2, c2.edge, axis=0) + unique_edges1 = jp.take(vertices1, c1.edge_dir, axis=0) + unique_edges2 = jp.take(vertices2, c2.edge_dir, axis=0) edges1 = jp.take(vertices1, c1.edge, axis=0) edges2 = jp.take(vertices2, c2.edge, axis=0) @@ -944,13 +985,6 @@ def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact: pos = c2.pos + pos @ c2.mat.T normal = normal @ c2.mat.T normal = -normal if swapped else normal - frame = jax.vmap(math.make_frame)(normal) + return dist, pos, frame - - -# store ncon as function attributes -plane_convex.ncon = 4 -sphere_convex.ncon = 1 -capsule_convex.ncon = 2 -convex_convex.ncon = 4 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index 1fabdf2f..c68e0edd 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -12,21 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Collide geometries.""" +"""Runs collision checking for all geoms in a Model. -from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union +To do this, collision_driver builds a collision function table, and then runs +the collision functions serially on the parameters in the table. + +For example, if a Model has three geoms: + +geom | type +--------------- +1 | sphere +2 | capsule +3 | sphere + +collision_driver organizes it into these functions and runs them: + +function | geom pair +-------------------------- +sphere_sphere | (1, 3) +sphere_capsule | (1, 2), (2, 3) + + +Besides collision function, function tables are keyed on mesh id and condim, +in order to guarantee static shapes for contacts and jacobians. +""" + +import itertools +from typing import Dict, Iterator, List, Tuple, Union import jax from jax import numpy as jp import mujoco -from mujoco.mjx._src import collision_base -from mujoco.mjx._src import mesh from mujoco.mjx._src import support # pylint: disable=g-importing-member -from mujoco.mjx._src.collision_base import Candidate -from mujoco.mjx._src.collision_base import CandidateSet -from mujoco.mjx._src.collision_base import GeomInfo -from mujoco.mjx._src.collision_base import SolverParams from mujoco.mjx._src.collision_convex import capsule_convex from mujoco.mjx._src.collision_convex import convex_convex from mujoco.mjx._src.collision_convex import plane_convex @@ -40,12 +58,14 @@ from mujoco.mjx._src.collision_primitive import sphere_capsule from mujoco.mjx._src.collision_primitive import sphere_sphere from mujoco.mjx._src.collision_sdf import capsule_ellipsoid from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid +from mujoco.mjx._src.collision_types import FunctionKey from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import DisableBit from mujoco.mjx._src.types import GeomType from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member +import numpy as np # pair-wise collision functions _COLLISION_FUNC = { @@ -70,371 +90,293 @@ _COLLISION_FUNC = { } -def get_collision_fn( - key: Tuple[Union[GeomType, mujoco.mjtGeom], Union[GeomType, mujoco.mjtGeom]] -) -> Optional[Callable[[GeomInfo, GeomInfo], collision_base.Contact]]: - """Returns a collision function given a pair of geom types.""" - return _COLLISION_FUNC.get(key, None) +# geoms for which we ignore broadphase +_GEOM_NO_BROADPHASE = {GeomType.HFIELD, GeomType.PLANE} -def _add_candidate( - result: CandidateSet, +def has_collision_fn(t1: GeomType, t2: GeomType) -> bool: + """Returns True if a collision function exists for a pair of geom types.""" + return (t1, t2) in _COLLISION_FUNC + + +def geom_pairs( m: Union[Model, mujoco.MjModel], - g1: int, - g2: int, - ipair: int = -1, -): - """Adds a candidate to test for collision.""" - t1, t2 = m.geom_type[g1], m.geom_type[g2] - if t1 > t2: - t1, t2, g1, g2 = t2, t1, g2, g1 +) -> Iterator[Tuple[int, int, int]]: + """Returns geom pairs to check for collisions. - # MuJoCo does not collide planes with other planes or hfields - if t1 == GeomType.PLANE and t2 == GeomType.PLANE: - return - if t1 == GeomType.PLANE and t2 == GeomType.HFIELD: - return + Args: + m: a MuJoCo or MJX model - def mesh_key(i): - convex_data = [[None] * m.ngeom] * 3 - if isinstance(m, Model): - convex_data = [ - m.geom_convex_face, - m.geom_convex_vert, - m.geom_convex_edge_dir, - ] - elif isinstance(m, mujoco.MjModel): - kwargs = mesh.get(m) - convex_data = [ - kwargs['geom_convex_face'], - kwargs['geom_convex_vert'], - kwargs['geom_convex_edge_dir'], - ] - key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data) - return key + Yields: + geom1, geom2, and pair index if defined in (else -1) + """ + pairs = set() - k1, k2 = mesh_key(g1), mesh_key(g2) + for i in range(m.npair): + g1, g2 = m.pair_geom1[i], m.pair_geom2[i] + # order pairs by geom_type for correct function mapping + if m.geom_type[g1] > m.geom_type[g2]: + g1, g2 = g2, g1 + pairs.add((g1, g2)) + yield g1, g2, i - candidates = {(c.geom1, c.geom2) for c in result.get((t1, t2, k1, k2), [])} - if (g1, g2) in candidates: - return - - if ipair > -1: - candidate = Candidate(g1, g2, ipair, -1, m.pair_dim[ipair]) - elif m.geom_priority[g1] != m.geom_priority[g2]: - gp = g1 if m.geom_priority[g1] > m.geom_priority[g2] else g2 - candidate = Candidate(g1, g2, -1, gp, m.geom_condim[gp]) - else: - dim = max(m.geom_condim[g1], m.geom_condim[g2]) - candidate = Candidate(g1, g2, -1, -1, dim) - - result.setdefault((t1, t2, k1, k2), []).append(candidate) - - -def _pair_params( - m: Model, - candidates: Sequence[Candidate], -) -> SolverParams: - """Gets solver params for pair geoms.""" - ipair = jp.array([c.ipair for c in candidates]) - friction = jp.clip(m.pair_friction[ipair], a_min=mujoco.mjMINMU) - solref = m.pair_solref[ipair] - solreffriction = m.pair_solreffriction[ipair] - solimp = m.pair_solimp[ipair] - margin = m.pair_margin[ipair] - gap = m.pair_gap[ipair] - - return SolverParams(friction, solref, solreffriction, solimp, margin, gap) - - -def _priority_params( - m: Model, - candidates: Sequence[Candidate], -) -> SolverParams: - """Gets solver params from priority geoms.""" - geomp = jp.array([c.geomp for c in candidates]) - friction = m.geom_friction[geomp][:, jp.array([0, 0, 1, 2, 2])] - solref = m.geom_solref[geomp] - solreffriction = jp.zeros(geomp.shape + (mujoco.mjNREF,)) - solimp = m.geom_solimp[geomp] - g = jp.array([(c.geom1, c.geom2) for c in candidates]) - margin = jp.amax(m.geom_margin[g.T], axis=0) - gap = jp.amax(m.geom_gap[g.T], axis=0) - - return SolverParams(friction, solref, solreffriction, solimp, margin, gap) - - -def _dynamic_params( - m: Model, - candidates: Sequence[Candidate], -) -> SolverParams: - """Gets solver params for dynamic geoms.""" - g1 = jp.array([c.geom1 for c in candidates]) - g2 = jp.array([c.geom2 for c in candidates]) - - friction = jp.maximum(m.geom_friction[g1], m.geom_friction[g2]) - # copy friction terms for the full geom pair - friction = friction[:, jp.array([0, 0, 1, 2, 2])] - - minval = jp.array(mujoco.mjMINVAL) - solmix1, solmix2 = m.geom_solmix[g1], m.geom_solmix[g2] - mix = solmix1 / (solmix1 + solmix2) - mix = jp.where((solmix1 < minval) & (solmix2 < minval), 0.5, mix) - mix = jp.where((solmix1 < minval) & (solmix2 >= minval), 0.0, mix) - mix_fn = jax.vmap(lambda a, b, m: m * a + (1 - m) * b) - - solref1, solref2 = m.geom_solref[g1], m.geom_solref[g2] - solref = jp.minimum(solref1, solref2) - s_mix = mix_fn(solref1, solref2, mix) - solref = jp.where((solref1[0] > 0) & (solref2[0] > 0), s_mix, solref) - solreffriction = jp.zeros(g1.shape + (mujoco.mjNREF,)) - solimp = mix_fn(m.geom_solimp[g1], m.geom_solimp[g2], mix) - margin = jp.maximum(m.geom_margin[g1], m.geom_margin[g2]) - gap = jp.maximum(m.geom_gap[g1], m.geom_gap[g2]) - - return SolverParams(friction, solref, solreffriction, solimp, margin, gap) - - -def get_params( - m: Union[Model, mujoco.MjModel], candidates: Sequence[Candidate] -) -> Tuple[List[int], List[int], SolverParams]: - """Gets solver params for a list of collision candidates.""" - # group sol params by different candidate types - typ_cands = {} - for c in candidates: - typ = (c.ipair > -1, c.geomp > -1) - typ_cands.setdefault(typ, []).append(c) - - geom1, geom2, params = [], [], [] - for (pair, priority), candidates in typ_cands.items(): - geom1.extend([c.geom1 for c in candidates]) - geom2.extend([c.geom2 for c in candidates]) - if pair: - params.append(_pair_params(m, candidates)) - elif priority: - params.append(_priority_params(m, candidates)) - else: - params.append(_dynamic_params(m, candidates)) - - params = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *params) - return geom1, geom2, params - - -def _pair_info( - m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int] -) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]: - """Returns geom pair info for calculating collision.""" - def mesh_info(geom): - g = jp.array(geom) - info = GeomInfo( - g, - d.geom_xpos[g], - d.geom_xmat[g], - m.geom_size[g], - ) - in_axes = jax.tree_util.tree_map(lambda x: 0, info) - is_mesh = m.geom_convex_face[geom[0]] is not None - if is_mesh: - info = info.replace( - face=jp.stack([m.geom_convex_face[i] for i in geom]), - vert=jp.stack([m.geom_convex_vert[i] for i in geom]), - edge_dir=jp.stack([m.geom_convex_edge_dir[i] for i in geom]), - facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom]), - edge=jp.stack([m.geom_convex_edge[i] for i in geom]), - edge_face_normal=jp.stack( - [m.geom_convex_edge_face_normal[i] for i in geom] - ), - ) - in_axes = in_axes.replace( - face=0, - vert=0, - edge_dir=0, - facenorm=0, - edge=0, - edge_face_normal=0, - ) - return info, in_axes - - info1, in_axes1 = mesh_info(geom1) - info2, in_axes2 = mesh_info(geom2) - return info1, info2, [in_axes1, in_axes2] - - -def _body_pair_filter( - m: Union[Model, mujoco.MjModel], b1: int, b2: int -) -> bool: - """Filters body pairs for collision.""" - dsbl_filterparent = m.opt.disableflags & DisableBit.FILTERPARENT - weld1 = m.body_weldid[b1] - weld2 = m.body_weldid[b2] - parent_weld1 = m.body_weldid[m.body_parentid[weld1]] - parent_weld2 = m.body_weldid[m.body_parentid[weld2]] - - if weld1 == weld2: - # filter out self-collisions - return True - - if ( - not dsbl_filterparent - and weld1 != 0 - and weld2 != 0 - and (weld1 == parent_weld2 or weld2 == parent_weld1) - ): - # filter out parent-child collisions - return True - - return False - - -def _broadphase_enabled( - geom_types: Tuple[GeomType, GeomType], - n_pairs: int, - max_pairs: int, -) -> bool: - return ( - GeomType.PLANE not in geom_types - and max_pairs > -1 - and n_pairs > max_pairs - ) - - -def _collide_geoms( - m: Model, - d: Data, - geom_types: Tuple[GeomType, GeomType], - candidates: Sequence[Candidate], -) -> Contact: - """Collides a geom pair.""" - fn = get_collision_fn(geom_types) - if not fn: - return Contact.zero() - - geom1, geom2, params = get_params(m, candidates) - g1, g2, in_axes = _pair_info(m, d, geom1, geom2) - - # Run a crude version of broadphase. - max_pairs = int(support.get_custom_numeric(m, 'max_geom_pairs')) - run_broadphase = _broadphase_enabled(geom_types, len(geom1), max_pairs) - n_pairs = max_pairs if run_broadphase else len(geom1) - if run_broadphase: - # broadphase over geom pairs, using bounding spheres - size1 = jp.max(m.geom_size[g1.geom_id], axis=-1) - size2 = jp.max(m.geom_size[g2.geom_id], axis=-1) - dists = jax.vmap(jp.linalg.norm)(g2.pos - g1.pos) - (size1 + size2) - _, idx = jax.lax.top_k(-dists, k=n_pairs) - g1, g2, params = jax.tree_util.tree_map( - lambda x, idx=idx: x[idx, ...], (g1, g2, params) - ) - - # call contact function - res = jax.vmap(fn, in_axes=in_axes)(g1, g2) - dist, pos, frame = jax.tree_util.tree_map(jp.concatenate, res) - - # repeat params by the number of contacts per geom pair - geom1, geom2, params = jax.tree_util.tree_map( - lambda x: jp.repeat(x, fn.ncon, axis=0), # pytype: disable=attribute-error - (g1.geom_id, g2.geom_id, params), - ) - - con = Contact( - dist=dist, - pos=pos, - frame=frame, - includemargin=params.margin - params.gap, - friction=params.friction, - solref=params.solref, - solreffriction=params.solreffriction, - solimp=params.solimp, - geom1=geom1, - geom2=geom2, - ) - return con - - -def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: - """Returns candidates for collision checking.""" - candidate_set = {} - - for ipair in range(m.npair): - g1, g2 = m.pair_geom1[ipair], m.pair_geom2[ipair] - _add_candidate(candidate_set, m, g1, g2, ipair) - - body_pairs = [] exclude_signature = set(m.exclude_signature) geom_con = m.geom_contype | m.geom_conaffinity + filterparent = not (m.opt.disableflags & DisableBit.FILTERPARENT) b_start = m.body_geomadr b_end = b_start + m.body_geomnum for b1 in range(m.nbody): 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(): continue signature = (b1 << 16) + (b2) if signature in exclude_signature: continue - if _body_pair_filter(m, b1, b2): + w2 = m.body_weldid[b2] + # ignore self-collisions + if w1 == w2: continue - body_pairs.append((b1, b2)) + w2_p = m.body_weldid[m.body_parentid[w2]] + # ignore parent-child collisions + if filterparent and w1 != 0 and w2 != 0 and (w1 == w2_p or w2 == w1_p): + continue + g1_range = [g for g in range(b_start[b1], b_end[b1]) if geom_con[g]] + g2_range = [g for g in range(b_start[b2], b_end[b2]) if geom_con[g]] - for b1, b2 in body_pairs: - for g1 in range(b_start[b1], b_end[b1]): - if not geom_con[g1]: - continue - for g2 in range(b_start[b2], b_end[b2]): - if not geom_con[g2]: + for g1, g2 in itertools.product(g1_range, g2_range): + t1, t2 = m.geom_type[g1], m.geom_type[g2] + # order pairs by geom_type for correct function mapping + if t1 > t2: + g1, g2, t1, t2 = g2, g1, t2, t1 + # ignore plane<>plane and plane<>hfield + if (t1, t2) == (GeomType.PLANE, GeomType.PLANE): continue + if (t1, t2) == (GeomType.PLANE, GeomType.HFIELD): + continue + # geoms must match contype and conaffinity on some bit mask = m.geom_contype[g1] & m.geom_conaffinity[g2] mask |= m.geom_contype[g2] & m.geom_conaffinity[g1] - if mask != 0: - _add_candidate(candidate_set, m, g1, g2) + if not mask: + continue - return candidate_set + if (g1, g2) not in pairs: + pairs.add((g1, g2)) + yield g1, g2, -1 -def ncon(m: Union[Model, mujoco.MjModel]) -> int: - """Returns the number of contacts computed in MJX given a model.""" +def _geom_groups( + m: Union[Model, mujoco.MjModel], +) -> Dict[FunctionKey, List[Tuple[int, int, int]]]: + """Returns geom pairs to check for collision grouped by collision function. + + The grouping consists of: + - The collision function to run, which is determined by geom types + - For mesh geoms, convex functions are run for each distinct mesh in the + model, because the convex functions expect static mesh size. If a sphere + collides with a cube and a tetrahedron, sphere_convex is called twice. + - The condim of the collision. This ensures that the size of the resulting + constraint jacobian is determined at compile time. + + Args: + m: a MuJoCo or MJX model + + Returns: + a dict with grouping key and values geom1, geom2, pair index + """ + groups = {} + + for g1, g2, ip in geom_pairs(m): + types = m.geom_type[g1], m.geom_type[g2] + data_ids = m.geom_dataid[g1], m.geom_dataid[g2] + if ip > -1: + condim = m.pair_dim[ip] + elif m.geom_priority[g1] > m.geom_priority[g2]: + condim = m.geom_condim[g1] + elif m.geom_priority[g1] < m.geom_priority[g2]: + condim = m.geom_condim[g2] + else: + condim = max(m.geom_condim[g1], m.geom_condim[g2]) + + key = FunctionKey(types, data_ids, condim) + groups.setdefault(key, []).append((g1, g2, ip)) + + return groups + + +def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]: + """Returns contact groups to check for collisions. + + Contacts are grouped the same way as _geom_groups. Only one contact is + emitted per geom pair, even if the collision function emits multiple contacts. + + Args: + m: MJX model + d: MJX data + + Returns: + a dict where the key is the grouping and value is a Contact + """ + groups = {} + eps = mujoco.mjMINVAL + + for key, geom_ids in _geom_groups(m).items(): + geom = np.array(geom_ids) + geom1, geom2, ip = geom.T + geom1, geom2, ip = geom1[ip == -1], geom2[ip == -1], ip[ip != -1] + params = [] + + if ip.size > 0: + # pair contacts get their params from m.pair_* fields + params.append(( + m.pair_margin[ip] - m.pair_gap[ip], + jp.clip(m.pair_friction[ip], a_min=eps), + m.pair_solref[ip], + m.pair_solreffriction[ip], + m.pair_solimp[ip] + )) + if geom1.size > 0 and geom2.size > 0: + # other contacts get their params from geom fields + margin = jp.maximum(m.geom_margin[geom1], m.geom_margin[geom2]) + gap = jp.maximum(m.geom_gap[geom1], m.geom_gap[geom2]) + solmix1, solmix2 = m.geom_solmix[geom1], m.geom_solmix[geom2] + mix = solmix1 / (solmix1 + solmix2) + mix = jp.where((solmix1 < eps) & (solmix2 < eps), 0.5, mix) + mix = jp.where((solmix1 < eps) & (solmix2 >= eps), 0.0, mix) + mix = jp.where((solmix1 >= eps) & (solmix2 < eps), 1.0, mix) + mix = mix[:, None] # for correct broadcasting + # friction: max + friction = jp.maximum(m.geom_friction[geom1], m.geom_friction[geom2]) + solref1, solref2 = m.geom_solref[geom1], m.geom_solref[geom2] + # reference standard: mix + solref_standard = mix * solref1 + (1 - mix) * solref2 + # reference direct: min + solref_direct = jp.minimum(solref1, solref2) + is_standard = (solref1[:, [0, 0]] > 0) & (solref2[:, [0, 0]] > 0) + solref = jp.where(is_standard, solref_standard, solref_direct) + solreffriction = jp.zeros(geom1.shape + (mujoco.mjNREF,)) + # impedance: mix + solimp = mix * m.geom_solimp[geom1] + (1 - mix) * m.geom_solimp[geom2] + + pri = m.geom_priority[geom1] != m.geom_priority[geom2] + if pri.any(): + # use priority geom when specified instead of mixing + gp1, gp2 = m.geom_priority[geom1], m.geom_priority[geom2] + gp = np.where(gp1 > gp2, geom1, geom2)[pri] + friction = friction.at[pri].set(m.geom_friction[gp]) + solref = solref.at[pri].set(m.geom_solref[gp]) + solimp = solimp.at[pri].set(m.geom_solimp[gp]) + + # unpack 5d friction: + friction = friction[:, [0, 0, 1, 2, 2]] + params.append((margin - gap, friction, solref, solreffriction, solimp)) + + params = map(jp.concatenate, zip(*params)) + includemargin, friction, solref, solreffriction, solimp = params + + groups[key] = Contact( + # dist, pos, frame get filled in by collision functions: + dist=None, + pos=None, + frame=None, + includemargin=includemargin, + friction=friction, + solref=solref, + solreffriction=solreffriction, + solimp=solimp, + dim=d.contact.dim, + geom1=jp.array(geom[:, 0]), + geom2=jp.array(geom[:, 1]), + geom=jp.array(geom[:, :2]), + efc_address=d.contact.efc_address, + ) + + return groups + + +def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray: + """Returns the dims of the contacts for a Model.""" if m.opt.disableflags & DisableBit.CONTACT: - return 0 + return np.empty(0, dtype=int) - candidates = collision_candidates(m) - max_count = int(support.get_custom_numeric(m, 'max_contact_points')) - max_pairs = int(support.get_custom_numeric(m, 'max_geom_pairs')) + group_counts = {k: len(v) for k, v in _geom_groups(m).items()} - count = 0 - for k, v in candidates.items(): - fn = get_collision_fn(k[0:2]) - if fn is None: - continue - run_broadphase = _broadphase_enabled((k[0], k[1]), len(v), max_pairs) - n_pair = max_pairs if run_broadphase else len(v) - count += n_pair * fn.ncon # pytype: disable=attribute-error + # max_geom_pairs limits the number of pairs we process in a collision function + # by first running a primitive broad phase culling on the pairs + max_geom_pairs = support.get_custom_int(m, 'max_geom_pairs') + if max_geom_pairs > -1: + for k in group_counts: + if set(k.types) & _GEOM_NO_BROADPHASE: + continue + group_counts[k] = min(group_counts[k], max_geom_pairs) - return min(max_count, count) if max_count > -1 else count + # max_contact_points limits the number of contacts emitted by selecting the + # contacts with the most penetration after calling collision functions + max_contact_points = support.get_custom_int(m, 'max_contact_points') + + condim_counts = {} + for k, v in group_counts.items(): + func = _COLLISION_FUNC[k.types] + num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v # pytype: disable=attribute-error + if max_contact_points > -1: + num_contacts = min(max_contact_points, num_contacts) + condim_counts[k.condim] = num_contacts + + dims = sum(([c] * condim_counts[c] for c in sorted(condim_counts)), []) + + return np.array(dims) def collision(m: Model, d: Data) -> Data: """Collides geometries.""" - if ncon(m) == 0: - return d.replace(contact=Contact.zero()) + if d.ncon == 0: + return d - candidate_set = collision_candidates(m) + groups = _contact_groups(m, d) + max_geom_pairs = support.get_custom_int(m, 'max_geom_pairs') + max_contact_points = support.get_custom_int(m, 'max_contact_points') - contacts = [] - for key, candidates in candidate_set.items(): - geom_types = key[0:2] - contacts.append(_collide_geoms(m, d, geom_types, candidates)) + # run collision functions on groups + for key, contact in groups.items(): + # determine which contacts we'll use for collision testing by running a + # broad phase cull if requested + if max_geom_pairs > -1 and contact.geom.shape[0] > max_geom_pairs: + pos1, pos2 = d.geom_xpos[contact.geom.T] + size1, size2 = m.geom_rbound[contact.geom.T] + dist = jax.vmap(jp.linalg.norm)(pos2 - pos1) - (size1 + size2) + _, idx = jax.lax.top_k(-dist, k=max_geom_pairs) + contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact) - if not contacts: - raise RuntimeError('No contacts found.') + # run the collision function specified by the grouping key + func = _COLLISION_FUNC[key.types] + dist, pos, frame = func(m, d, key, contact.geom) + ncon = func.ncon # pytype: disable=attribute-error + if ncon > 1: + # repeat contacts to match the number of collisions returned + repeat_fn = lambda x, r=ncon: jp.repeat(x, r, axis=0) + contact = jax.tree_util.tree_map(repeat_fn, contact) + groups[key] = contact.replace(dist=dist, pos=pos, frame=frame) + # collapse contacts together, ensuring they are grouped by condim + condim_groups = {} + for key, contact in groups.items(): + condim_groups.setdefault(key.condim, []).append(contact) + + # limit the number of contacts per condim group if requested + if max_contact_points > -1: + for key, contacts in condim_groups.items(): + contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts) + if contact.geom.shape[0] > max_contact_points: + _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) + contact = jax.tree_util.tree_map(lambda x, idx=idx: x[idx], contact) + condim_groups[key] = [contact] + + contacts = sum([condim_groups[k] for k in sorted(condim_groups)], []) contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts) - max_contact_points = int(support.get_custom_numeric(m, 'max_contact_points')) - if max_contact_points > -1 and contact.dist.shape[0] > max_contact_points: - # get top-k contacts - _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) - contact = jax.tree_util.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact) - return d.replace(contact=contact) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 0f475aba..c1d77e99 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -402,7 +402,8 @@ class CapsuleCollisionTest(parameterized.TestCase): self.assertEqual(c.pos.shape[0], 2) self.assertGreater(c.dist[1], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) @@ -437,7 +438,8 @@ class CapsuleCollisionTest(parameterized.TestCase): self.assertEqual(c.pos.shape[0], 2) self.assertGreater(c.dist[1], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4) np.testing.assert_array_almost_equal( @@ -457,7 +459,8 @@ class CapsuleCollisionTest(parameterized.TestCase): self.assertEqual(c.pos.shape[0], 2) self.assertGreater(c.dist[1], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4) np.testing.assert_array_almost_equal( @@ -494,7 +497,8 @@ class CylinderTest(absltest.TestCase): d.contact.pos[:] = d.contact.pos[idx] # extract the contact points with penetration - c = jax.tree_util.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:2], dx.contact) + c = c.replace(dim=c.dim[:2], efc_address=c.efc_address[:2]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'cylinder_plane', 1e-5) @@ -532,6 +536,7 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(-dx.contact.dist[2:], 0) # extract the contact points with penetration c = jax.tree_util.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) + c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-5) @@ -615,7 +620,8 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(dx.contact.dist[:1], 0) np.testing.assert_array_less(-dx.contact.dist[1:], 0) # extract the contact point with penetration - c = jax.tree_util.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact) + c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1]) for field in dataclasses.fields(Contact): _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) @@ -763,28 +769,28 @@ class BodyPairFilterTest(absltest.TestCase): self.assertEqual(dx.contact.pos.shape[0], 1) -class NconTest(parameterized.TestCase): - """Tests ncon.""" +class DimTest(parameterized.TestCase): + """Tests contact dim.""" def test_ncon(self): m = test_util.load_test_file('constraints.xml') - ncon = collision_driver.ncon(m) - self.assertEqual(ncon, 16) + dim = collision_driver.make_condim(m) + np.testing.assert_array_equal(dim, np.array([3] * 16)) def test_disable_contact(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags |= DisableBit.CONTACT - ncon = collision_driver.ncon(m) - self.assertEqual(ncon, 0) + dim = collision_driver.make_condim(m) + self.assertEqual(dim.size, 0) def test_ncon_meshes(self): m = test_util.load_test_file('shadow_hand/scene_right.xml') - ncon = collision_driver.ncon(m) + ncon = collision_driver.make_condim(m).size self.assertEqual(ncon, 15) mx = mjx.put_model(m) - ncon = collision_driver.ncon(mx) + ncon = collision_driver.make_condim(mx).size self.assertEqual(ncon, 15) # get rid of max_contact_points, test only max_geom_pairs @@ -795,11 +801,11 @@ class NconTest(parameterized.TestCase): if name_ == 'max_contact_points': m.numeric_data[m.numeric_adr[i]] = -1 - ncon = collision_driver.ncon(m) + ncon = collision_driver.make_condim(m).size self.assertEqual(ncon, 98) mx = mjx.put_model(m) - ncon = collision_driver.ncon(mx) + ncon = collision_driver.make_condim(mx).size self.assertEqual(ncon, 98) diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py index 9040b22e..f5757853 100644 --- a/mjx/mujoco/mjx/_src/collision_primitive.py +++ b/mjx/mujoco/mjx/_src/collision_primitive.py @@ -20,34 +20,53 @@ import jax from jax import numpy as jp from mujoco.mjx._src import math # pylint: disable=g-importing-member -from mujoco.mjx._src.collision_base import Contact -from mujoco.mjx._src.collision_base import GeomInfo +from mujoco.mjx._src.collision_types import Collision +from mujoco.mjx._src.collision_types import GeomInfo +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member +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 + info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]) + info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]) + dist, pos, frame = jax.vmap(func)(info1, info2) + if ncon > 1: + return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) + return dist, pos, frame + + collide.ncon = ncon + return collide + + return wrapper + + def _plane_sphere( plane_normal: jax.Array, plane_pos: jax.Array, sphere_pos: jax.Array, - radius: jax.Array, + sphere_radius: jax.Array, ) -> Tuple[jax.Array, jax.Array]: - """Returns the penetration and contact point between a plane and sphere.""" - cdist = jp.dot(sphere_pos - plane_pos, plane_normal) - dist = cdist - radius - pos = sphere_pos - plane_normal * (radius + 0.5 * dist) + """Returns the distance and contact point between a plane and sphere.""" + dist = jp.dot(sphere_pos - plane_pos, plane_normal) - sphere_radius + pos = sphere_pos - plane_normal * (sphere_radius + 0.5 * dist) return dist, pos -def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact: +@collider(ncon=1) +def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Collision: """Calculates contact between a plane and a sphere.""" n = plane.mat[:, 2] dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0]) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact: +@collider(ncon=2) +def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Collision: """Calculates two contacts between a capsule and a plane.""" n, axis = plane.mat[:, 2], cap.mat[:, 2] # align contact frames with capsule axis @@ -56,16 +75,17 @@ def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact: b = jp.where(b_norm < 0.5, jp.where((-0.5 < n[1]) & (n[1] < 0.5), y, z), b) frame = jp.array([[n, b, jp.cross(n, b)]]) segment = axis * cap.size[1] - contacts = [] + collisions = [] for offset in [segment, -segment]: dist, pos = _plane_sphere(n, plane.pos, cap.pos + offset, cap.size[0]) dist = jp.expand_dims(dist, axis=0) pos = jp.expand_dims(pos, axis=0) - contacts.append((dist, pos, frame)) - return jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts) + collisions.append((dist, pos, frame)) + return jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *collisions) -def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Contact: +@collider(ncon=1) +def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Collision: """Calculates one contact between an ellipsoid and a plane.""" n = plane.mat[:, 2] size = ellipsoid.size @@ -73,12 +93,11 @@ def plane_ellipsoid(plane: GeomInfo, ellipsoid: GeomInfo) -> Contact: pos = ellipsoid.pos + ellipsoid.mat @ (sphere_support * size) dist = jp.dot(n, pos - plane.pos) pos = pos - n * dist * 0.5 - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Contact: +@collider(ncon=3) +def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Collision: """Calculates one contact between an cylinder and a plane.""" n = plane.mat[:, 2] axis = cylinder.mat[:, 2] @@ -139,7 +158,7 @@ def plane_cylinder(plane: GeomInfo, cylinder: GeomInfo) -> Contact: def _sphere_sphere( pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array -) -> Contact: +) -> Tuple[jax.Array, jax.Array, jax.Array]: """Returns the penetration, contact point, and normal between two spheres.""" n, dist = math.normalize_with_norm(pos2 - pos1) n = jp.where(dist == 0.0, jp.array([1.0, 0.0, 0.0]), n) @@ -148,15 +167,15 @@ def _sphere_sphere( return dist, pos, n -def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Contact: +@collider(ncon=1) +def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Collision: """Calculates contact between two spheres.""" dist, pos, n = _sphere_sphere(s1.pos, s1.size[0], s2.pos, s2.size[0]) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact: +@collider(ncon=1) +def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Collision: """Calculates one contact between a sphere and a capsule.""" axis, length = cap.mat[:, 2], cap.size[1] segment = axis * length @@ -164,19 +183,14 @@ def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact: cap.pos - segment, cap.pos + segment, sphere.pos ) dist, pos, n = _sphere_sphere(sphere.pos, sphere.size[0], pt, cap.size[0]) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact: +@collider(ncon=1) +def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Collision: """Calculates one contact between two capsules.""" - axis1, length1, axis2, length2 = ( - cap1.mat[:, 2], - cap1.size[1], - cap2.mat[:, 2], - cap2.size[1], - ) + axis1, length1 = cap1.mat[:, 2], cap1.size[1] + axis2, length2 = cap2.mat[:, 2], cap2.size[1] seg1, seg2 = axis1 * length1, axis2 * length2 pt1, pt2 = math.closest_segment_to_segment_points( cap1.pos - seg1, @@ -186,15 +200,4 @@ def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact: ) radius1, radius2 = cap1.size[0], cap2.size[0] dist, pos, n = _sphere_sphere(pt1, radius1, pt2, radius2) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) - -# store ncon as function attributes -plane_sphere.ncon = 1 -plane_capsule.ncon = 2 -plane_ellipsoid.ncon = 1 -plane_cylinder.ncon = 3 -sphere_sphere.ncon = 1 -sphere_capsule.ncon = 1 -capsule_capsule.ncon = 1 + return dist, pos, math.make_frame(n) diff --git a/mjx/mujoco/mjx/_src/collision_sdf.py b/mjx/mujoco/mjx/_src/collision_sdf.py index 71a3dbf6..85481245 100644 --- a/mjx/mujoco/mjx/_src/collision_sdf.py +++ b/mjx/mujoco/mjx/_src/collision_sdf.py @@ -28,15 +28,35 @@ import jax from jax import numpy as jp from mujoco.mjx._src import math # pylint: disable=g-importing-member -from mujoco.mjx._src.collision_base import Contact -from mujoco.mjx._src.collision_base import GeomInfo +from mujoco.mjx._src.collision_types import Collision +from mujoco.mjx._src.collision_types import GeomInfo from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member # the SDF function takes position in, and returns a distance or objective 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 + info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]) + info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]) + dist, pos, frame = jax.vmap(func)(info1, info2) + if ncon > 1: + return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) + return dist, pos, frame + + collide.ncon = ncon + return collide + + return wrapper + + def _plane(pos: jax.Array, size: jax.Array) -> jax.Array: del size return pos[2] @@ -128,22 +148,15 @@ def _optim( return pos, dist, n -def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact: +@collider(ncon=1) +def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision: """"Calculates contact between a capsule and an ellipsoid.""" pos, dist, n = _optim(_capsule, _ellipsoid, c, e) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) + return dist, pos, math.make_frame(n) -def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Contact: +@collider(ncon=1) +def ellipsoid_ellipsoid(e1: GeomInfo, e2: GeomInfo) -> Collision: """"Calculates contact between two ellipsoids.""" pos, dist, n = _optim(_ellipsoid, _ellipsoid, e1, e2) - return jax.tree_util.tree_map( - lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) - ) - - -# store ncon as function attributes -capsule_ellipsoid.ncon = 1 -ellipsoid_ellipsoid.ncon = 1 + return dist, pos, math.make_frame(n) diff --git a/mjx/mujoco/mjx/_src/collision_types.py b/mjx/mujoco/mjx/_src/collision_types.py new file mode 100644 index 00000000..49d77438 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_types.py @@ -0,0 +1,65 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collision base types.""" + +import dataclasses +from typing import Tuple +import jax +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +# pylint: enable=g-importing-member + + +class GeomInfo(PyTreeNode): + """Geom propertes of primitive and SDF shapes.""" + pos: jax.Array + mat: jax.Array + size: jax.Array + + +class ConvexInfo(PyTreeNode): + """Geom propertes of convex meshes.""" + pos: jax.Array + mat: jax.Array + vert: jax.Array + face: jax.Array + face_normal: jax.Array + edge: jax.Array + edge_face_normal: jax.Array + edge_dir: jax.Array + + +# Collision returned by collision functions: +# - distance distance between nearest points; neg: penetration +# - position (3,) position of contact point: midpoint between geoms +# - frame (3, 3) normal is in [0, :], points from geom[0] to geom[1] +Collision = Tuple[jax.Array, jax.Array, jax.Array] + + +@dataclasses.dataclass(frozen=True) +class FunctionKey: + """Specifies how geom pairs group into collision_driver's function table. + + Attributes: + types: geom type pair, which determines the collision function + data_ids: geom data id pair: mesh id for mesh geoms, otherwise -1. + Meshes have distinct face/vertex counts, so must occupy distinct + entries in the collision function table. + condim: grouping by condim of the colliision ensures that the size of the + resulting constraint jacobian is determined at compile time. + """ + 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 f08e1988..e4264aee 100644 --- a/mjx/mujoco/mjx/_src/constraint.py +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -24,6 +24,7 @@ from mujoco.mjx._src import math from mujoco.mjx._src import support # pylint: disable=g-importing-member from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import ConstraintType from mujoco.mjx._src.types import Contact from mujoco.mjx._src.types import Data from mujoco.mjx._src.types import DisableBit @@ -34,6 +35,9 @@ from mujoco.mjx._src.types import Model import numpy as np +_CONDIM_EFC_COUNT = {1: 1, 3: 4, 4: 6, 6: 10} + + class _Efc(PyTreeNode): """Support data for creating constraint matrices.""" J: jax.Array @@ -277,7 +281,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]: def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: """Calculates constraint rows for contacts.""" - if collision_driver.ncon(m) == 0: + if d.ncon == 0: return None @jax.vmap @@ -314,36 +318,51 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss) -def count_constraints( - m: Union[Model, mujoco.MjModel], d: Optional[Data] = None -) -> Tuple[int, int, int, int]: +def counts(efc_type: np.ndarray) -> Tuple[int, int, int, int]: """Returns equality, friction, limit, and contact constraint counts.""" - if m.opt.disableflags & DisableBit.CONSTRAINT: - return 0, 0, 0, 0 - - if m.opt.disableflags & DisableBit.EQUALITY: - ne = 0 - else: - ne_connect = (m.eq_type == EqType.CONNECT).sum() - ne_weld = (m.eq_type == EqType.WELD).sum() - ne_joint = (m.eq_type == EqType.JOINT).sum() - ne = ne_connect * 3 + ne_weld * 6 + ne_joint - - nf = 0 - - if m.opt.disableflags & DisableBit.LIMIT: - nl = 0 - else: - nl = int(m.jnt_limited.sum()) - - if d is None: - nc = collision_driver.ncon(m) * 4 - else: - nc = d.efc_J.shape[-2] - ne - nf - nl + ne = (efc_type == ConstraintType.EQUALITY).sum() + nf = 0 # no support for friction loss yet + nl = (efc_type == ConstraintType.LIMIT_JOINT).sum() + nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum() return ne, nf, nl, nc +def make_efc_type( + m: Union[Model, mujoco.MjModel], dim: Optional[np.ndarray] = None +) -> np.ndarray: + """Returns efc_type that outlines the type of each constraint row.""" + if m.opt.disableflags & DisableBit.CONSTRAINT: + return np.empty(0, dtype=int) + + dim = collision_driver.make_condim(m) if dim is None else dim + efc_types = [] + + if not m.opt.disableflags & DisableBit.EQUALITY: + num_rows = (m.eq_type == EqType.CONNECT).sum() * 3 + num_rows += (m.eq_type == EqType.WELD).sum() * 6 + num_rows += (m.eq_type == EqType.JOINT).sum() + efc_types.extend([ConstraintType.EQUALITY] * num_rows) + + if not m.opt.disableflags & DisableBit.LIMIT: + efc_types.extend([ConstraintType.LIMIT_JOINT] * m.jnt_limited.sum()) + + if not m.opt.disableflags & DisableBit.CONTACT: + num_rows = sum(_CONDIM_EFC_COUNT[d] for d in dim) + efc_types.extend([ConstraintType.CONTACT_PYRAMIDAL] * num_rows) + + return np.array(efc_types) + + +def make_efc_address(efc_type: np.ndarray, dim: np.ndarray) -> np.ndarray: + """Returns efc_address that maps contacts to constraint row address.""" + nc = (efc_type == ConstraintType.CONTACT_PYRAMIDAL).sum() + nc_start = efc_type.size - nc + offsets = np.cumsum([0] + [_CONDIM_EFC_COUNT[d] for d in dim])[:-1] + + return nc_start + offsets + + def make_constraint(m: Model, d: Data) -> Data: """Creates constraint jacobians and other supporting data.""" diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index a2bd8cc9..133640e3 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -72,7 +72,7 @@ class ConstraintTest(absltest.TestCase): def test_disable_constraint(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT - ne, nf, nl, nc = mjx.count_constraints(m) + ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) self.assertEqual(ne, 0) self.assertEqual(nf, 0) self.assertEqual(nl, 0) @@ -83,7 +83,7 @@ class ConstraintTest(absltest.TestCase): def test_disable_equality(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY - ne, nf, nl, nc = mjx.count_constraints(m) + ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) self.assertEqual(ne, 0) self.assertEqual(nf, 0) self.assertEqual(nl, 2) @@ -94,7 +94,7 @@ class ConstraintTest(absltest.TestCase): def test_disable_contact(self): m = test_util.load_test_file('constraints.xml') m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT - ne, nf, nl, nc = mjx.count_constraints(m) + ne, nf, nl, nc = constraint.counts(constraint.make_efc_type(m)) self.assertEqual(ne, 10) self.assertEqual(nf, 0) self.assertEqual(nl, 2) diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index 3ddb7a53..43b3d256 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -23,7 +23,6 @@ import jax from jax import numpy as jp import mujoco from mujoco.mjx._src import collision_driver -from mujoco.mjx._src import mesh from mujoco.mjx._src import types import numpy as np @@ -92,14 +91,10 @@ _INVERSE_TRANSFORMS = { ), } -_DERIVED = mesh.DERIVED.union( +_DERIVED = { # efc_J is dense in MJX, sparse in MJ. ignore for now. - {(types.Data, 'efc_J'), (types.Option, 'has_fluid_params')} -) - - -def _model_derived(value: mujoco.MjModel) -> Dict[str, Any]: - return {k: jax.device_put(v) for k, v in mesh.get(value).items()} + (types.Data, 'efc_J'), (types.Option, 'has_fluid_params') +} def _data_derived(value: mujoco.MjData) -> Dict[str, Any]: @@ -138,23 +133,21 @@ def _validate(m: mujoco.MjModel): if m.body_gravcomp.any(): raise NotImplementedError('gravcomp is not supported') - # check collision geom types - for (g1, g2, *_), c in collision_driver.collision_candidates(m).items(): - g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) - if g1 == mujoco.mjtGeom.mjGEOM_PLANE and g2 in ( - mujoco.mjtGeom.mjGEOM_PLANE, - mujoco.mjtGeom.mjGEOM_HFIELD, - ): - # MuJoCo does not collide planes with other planes or hfields - continue - if collision_driver.get_collision_fn((g1, g2)) is None: - raise NotImplementedError(f'({g1}, {g2}) collisions not implemented.') - *_, params = collision_driver.get_params(m, c) - margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0) - if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap: - raise NotImplementedError( - f'Margin and gap not implemented for ({g1}, {g2})' - ) + for g1, g2, ip in collision_driver.geom_pairs(m): + t1, t2 = m.geom_type[[g1, g2]] + # check collision function exists for type pair + if not collision_driver.has_collision_fn(t1, t2): + t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) + raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') + # margin/gap not supported for geoms + if mujoco.mjtGeom.mjGEOM_MESH in (t1, t2): + if ip != -1: + margin = m.pair_margin[ip] + else: + margin = m.geom_margin[g1] + m.geom_margin[g2] + if margin.any(): + t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) + raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') # TODO(erikfrey): warn for high solver iterations, nefc, etc. @@ -229,7 +222,7 @@ def device_put(value): derived_kwargs = {} if isinstance(value, mujoco.MjModel): - derived_kwargs = _model_derived(value) + derived_kwargs = {} elif isinstance(value, mujoco.MjData): derived_kwargs = _data_derived(value) elif isinstance(value, mujoco.MjOption): diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 9d0036b4..f0872fd4 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -22,15 +22,14 @@ from jax import numpy as jp import mujoco from mujoco.mjx._src import collision_driver from mujoco.mjx._src import constraint -from mujoco.mjx._src import mesh from mujoco.mjx._src import support from mujoco.mjx._src import types import numpy as np import scipy -def _put_option(o: mujoco.MjOption, device=None) -> types.Option: - """Puts mujoco.MjOption onto a device, resulting in mjx.Option.""" +def _make_option(o: mujoco.MjOption) -> types.Option: + """Returns mjx.Option given mujoco.MjOption.""" if o.integrator not in set(types.IntegratorType): raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}') @@ -47,38 +46,20 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option: if o.enableflags & 2**i: raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}') - static_fields = { - f.name: copy.copy(getattr(o, f.name)) - for f in types.Option.fields() - if f.type in (int, bytes, np.ndarray) - } - static_fields['integrator'] = types.IntegratorType(o.integrator) - static_fields['cone'] = types.ConeType(o.cone) - static_fields['jacobian'] = types.JacobianType(o.jacobian) - static_fields['solver'] = types.SolverType(o.solver) - static_fields['disableflags'] = types.DisableBit(o.disableflags) + fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()} + fields['integrator'] = types.IntegratorType(o.integrator) + fields['cone'] = types.ConeType(o.cone) + fields['jacobian'] = types.JacobianType(o.jacobian) + fields['solver'] = types.SolverType(o.solver) + fields['disableflags'] = types.DisableBit(o.disableflags) + fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any() - device_fields = { - f.name: copy.copy(getattr(o, f.name)) - for f in types.Option.fields() - if f.type is jax.Array - } - device_fields = jax.device_put(device_fields, device=device) - - has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any() - - return types.Option( - has_fluid_params=has_fluid_params, - **static_fields, - **device_fields, - ) + return types.Option(**fields) -def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic: +def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic: """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic.""" - return types.Statistic( - meaninertia=jax.device_put(s.meaninertia, device=device) - ) + return types.Statistic(meaninertia=s.meaninertia) def put_model(m: mujoco.MjModel, device=None) -> types.Model: @@ -93,17 +74,21 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: if m.body_gravcomp.any(): raise NotImplementedError('gravcomp is not supported') - # check collision geom types - for (g1, g2, *_), c in collision_driver.collision_candidates(m).items(): - g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) - if collision_driver.get_collision_fn((g1, g2)) is None: - raise NotImplementedError(f'({g1}, {g2}) has no collision function') - *_, params = collision_driver.get_params(m, c) - margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0) - if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap: - raise NotImplementedError( - f'Margin and gap not implemented for ({g1}, {g2})' - ) + for g1, g2, ip in collision_driver.geom_pairs(m): + t1, t2 = m.geom_type[[g1, g2]] + # check collision function exists for type pair + if not collision_driver.has_collision_fn(t1, t2): + t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) + raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.') + # margin/gap not supported for geoms + if mujoco.mjtGeom.mjGEOM_MESH in (t1, t2): + if ip != -1: + margin = m.pair_margin[ip] + else: + margin = m.geom_margin[g1] + m.geom_margin[g2] + if margin.any(): + t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) + raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') for enum_field, enum_type, mj_type in ( (m.actuator_biastype, types.BiasType, mujoco.mjtBias), @@ -121,40 +106,24 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: if not np.allclose(m.dof_frictionloss, 0): raise NotImplementedError('dof_frictionloss is not implemented.') - opt = _put_option(m.opt, device=device) - stat = _put_statistic(m.stat, device=device) + fields = {f.name: getattr(m, f.name) for f in types.Model.fields()} + fields['geom_rgba'] = fields['geom_rgba'].reshape((-1, 4)) + fields['mat_rgba'] = fields['mat_rgba'].reshape((-1, 4)) + fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3)) + fields['opt'] = _make_option(m.opt) + fields['stat'] = _make_statistic(m.stat) + model = types.Model(**{k: copy.copy(v) for k, v in fields.items()}) - static_fields = { - f.name: getattr(m, f.name) - for f in types.Model.fields() - if f.type in (int, bytes, np.ndarray) - } - static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4)) - static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4)) - - device_fields = { - f.name: copy.copy(getattr(m, f.name)) # copy because device_put is async - for f in types.Model.fields() - if f.type is jax.Array - } - device_fields['cam_mat0'] = device_fields['cam_mat0'].reshape((-1, 3, 3)) - device_fields.update(mesh.get(m)) - device_fields = jax.device_put(device_fields, device=device) - - return types.Model( - opt=opt, - stat=stat, - **static_fields, - **device_fields, - ) + return jax.device_put(model, device=device) def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: """Allocate and initialize Data.""" - - ncon = collision_driver.ncon(m) - ne, nf, nl, nc = constraint.count_constraints(m) - nefc = ne + nf + nl + nc + dim = collision_driver.make_condim(m) + efc_type = constraint.make_efc_type(m, dim) + efc_address = constraint.make_efc_address(efc_type, dim) + ne, nf, nl, nc = constraint.counts(efc_type) + ncon, nefc = dim.size, ne + nf + nl + nc zero_0 = jp.zeros(0, dtype=float) zero_nv = jp.zeros(m.nv, dtype=float) @@ -170,8 +139,28 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float) zero_nm = jp.zeros(m.nM, dtype=float) - # create first d to get num contacts and nc + contact = types.Contact( + dist=jp.zeros(ncon), + pos=jp.zeros((ncon, 3)), + frame=jp.zeros((ncon, 3, 3)), + includemargin=jp.zeros(ncon), + friction=jp.zeros((ncon, 5)), + solref=jp.zeros((ncon, mujoco.mjNREF)), + solreffriction=jp.zeros((ncon, mujoco.mjNREF)), + solimp=jp.zeros((ncon, mujoco.mjNIMP)), + dim=dim, + geom1=jp.zeros(ncon, dtype=int) - 1, + geom2=jp.zeros(ncon, dtype=int) - 1, + geom=jp.zeros((ncon, 2), dtype=int) - 1, + efc_address=efc_address, + ) + d = types.Data( + ne=ne, + nf=nf, + nl=nl, + nefc=nefc, + ncon=ncon, solver_niter=jp.array(0, dtype=int), time=jp.array(0.0), qpos=jp.array(m.qpos0), @@ -206,7 +195,8 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: qM=zero_nm if support.is_sparse(m) else zero_nv_nv, qLD=zero_nm if support.is_sparse(m) else zero_nv_nv, qLDiagInv=zero_nv if support.is_sparse(m) else zero_0, - contact=types.Contact.zero(ncon), + contact=contact, + efc_type=efc_type, efc_J=jp.zeros((nefc, m.nv), dtype=float), efc_frictionloss=zero_nefc, efc_D=zero_nefc, @@ -228,11 +218,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data: return d -def _get_contact( - c: mujoco._structs._MjContactList, - cx: types.Contact, - efc_start: int, -): +def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact): """Converts mjx.Contact to mujoco._structs._MjContactList.""" con_id = np.nonzero(cx.dist <= 0)[0] for field in types.Contact.fields(): @@ -241,10 +227,6 @@ def _get_contact( value = value.reshape((-1, 9)) getattr(c, field.name)[:] = value - ncon = cx.dist.shape[0] - c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id] - c.dim[:] = 3 - def get_data( m: mujoco.MjModel, d: types.Data @@ -278,13 +260,6 @@ def get_data_into( d = jax.device_get(d) batch_size = d.qpos.shape[0] if batched else 1 - ne, nf, nl, nc = constraint.count_constraints(m, d) - efc_type = np.array([ - mujoco.mjtConstraint.mjCNSTR_EQUALITY, - mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF, - mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT, - mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL, - ]).repeat([ne, nf, nl, nc]) dof_i, dof_j = [], [] for i in range(m.nv): @@ -299,8 +274,7 @@ def get_data_into( result_i = result[i] if batched else result ncon = (d_i.contact.dist <= 0).sum() efc_active = (d_i.efc_J != 0).any(axis=1) - efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL - nefc, nc = int(efc_active.sum()), int((efc_active & efc_con).sum()) + nefc = int(efc_active.sum()) result_i.nnzJ = nefc * m.nv if ncon != result_i.ncon or nefc != result_i.nefc: mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access @@ -310,61 +284,63 @@ def get_data_into( for field in types.Data.fields(): if field.name == 'contact': - _get_contact(result_i.contact, d_i.contact, nefc - nc) + _get_contact(result_i.contact, d_i.contact) + # efc_address must be updated because rows were deleted above: + efc_map = np.cumsum(efc_active) - 1 + result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address] continue value = getattr(d_i, field.name) - if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): + if field.name in ('nefc', 'ncon'): + value = {'nefc': nefc, 'ncon': ncon}[field.name] + elif field.name.endswith('xmat') or field.name == 'ximat': value = value.reshape((-1, 9)) - - if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'): + elif field.name.startswith('efc_'): value = value[efc_active] - - if field.name == 'efc_J': - value = value[efc_active].reshape(-1) - - if field.name == 'qM' and not support.is_sparse(m): + if field.name == 'efc_J': + value = value.reshape(-1) + elif field.name == 'qM' and not support.is_sparse(m): value = value[dof_i, dof_j] - - if field.name == 'qLD' and not support.is_sparse(m): + elif field.name == 'qLD' and not support.is_sparse(m): value = value[dof_i, dof_j] - - if field.name == 'qLDiagInv' and not support.is_sparse(m): + elif field.name == 'qLDiagInv' and not support.is_sparse(m): value = np.ones(m.nv) - if value.shape: + if isinstance(value, np.ndarray) and value.shape: getattr(result_i, field.name)[:] = value else: setattr(result_i, field.name, value) - result_i.efc_type[:] = efc_type[efc_active] - -def _put_contact( - c: mujoco._structs._MjContactList, ncon: int, device=None +def _make_contact( + c: mujoco._structs._MjContactList, + dim: np.ndarray, + efc_address: np.ndarray, ) -> types.Contact: - """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact.""" - fields = { - f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields() - } + """Converts mujoco.structs._MjContactList into mjx.Contact.""" + fields = {f.name: getattr(c, f.name) for f in types.Contact.fields()} fields['frame'] = fields['frame'].reshape((-1, 3, 3)) - pad_size = ncon - c.dist.shape[0] + pad_size = dim.size - c.dist.shape[0] pad_fn = lambda x: np.concatenate( (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype)) ) fields = jax.tree_util.tree_map(pad_fn, fields) fields['dist'][-pad_size:] = np.inf - fields = jax.device_put(fields, device=device) + # TODO(erikfrey): move contacts to appropriate dim index + fields['dim'] = dim + fields['efc_address'] = efc_address return types.Contact(**fields) def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: """Puts mujoco.MjData onto a device, resulting in mjx.Data.""" - ncon = collision_driver.ncon(m) - ne, nf, nl, nc = constraint.count_constraints(m) - nefc = ne + nf + nl + nc + dim = collision_driver.make_condim(m) + efc_type = constraint.make_efc_type(m, dim) + efc_address = constraint.make_efc_address(efc_type, dim) + ne, nf, nl, nc = constraint.counts(efc_type) + ncon, nefc = dim.size, ne + nf + nl + nc for d_val, val, name in ( (d.ncon, ncon, 'ncon'), @@ -376,12 +352,9 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: if d_val > val: raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}') - fields = { - f.name: copy.copy(getattr(d, f.name)) # copy because device_put is async - for f in types.Data.fields() - if f.type is jax.Array - } + fields = {f.name: getattr(d, f.name) for f in types.Data.fields()} + # MJX prefers square matrices for these fields: for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat', 'cam_xmat'): fields[fname] = fields[fname].reshape((-1, 3, 3)) @@ -409,7 +382,7 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: value_beg = sum([ne, nf, nl][:i]) d_beg = sum([d.ne, d.nf, d.nl][:i]) size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i] - value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size] + value[value_beg : value_beg + size] = fields[fname][d_beg : d_beg + size] fields[fname] = value # convert qM and qLD if jacobian is dense @@ -424,7 +397,12 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data: fields['qLD'] = np.zeros((m.nv, m.nv)) fields['qLDiagInv'] = np.zeros(0) - fields = jax.device_put(fields, device=device) - fields['contact'] = _put_contact(d.contact, ncon, device=device) + fields['contact'] = _make_contact(d.contact, dim, efc_address) + fields.update( + dict(ne=ne, nf=nf, nl=nl, nefc=nefc, ncon=ncon, efc_type=efc_type) + ) - return types.Data(**fields) + # copy because device_put is async: + data = types.Data(**{k: copy.copy(v) for k, v in fields.items()}) + + return jax.device_put(data, device=device) diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index dde9c383..323714eb 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -108,10 +108,6 @@ class ModelIOTest(parameterized.TestCase): np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid) np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref) np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos) - self.assertLen(mx.geom_convex_face, 6) - self.assertLen(mx.geom_convex_vert, 6) - self.assertLen(mx.geom_convex_edge_dir, 6) - self.assertLen(mx.geom_convex_facenormal, 6) np.testing.assert_allclose(mx.jnt_type, m.jnt_type) np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr) @@ -256,6 +252,7 @@ class DataIOTest(parameterized.TestCase): nv = 19 nefc = 185 + self.assertEqual(d.nefc, nefc) self.assertEqual(d.qpos.shape, (nq,)) self.assertEqual(d.qvel.shape, (nv,)) self.assertEqual(d.act.shape, (0,)) @@ -423,8 +420,6 @@ class DataIOTest(parameterized.TestCase): np.testing.assert_allclose(d_2.efc_J, d.efc_J) self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc np.testing.assert_allclose(d_2.efc_aref, d.efc_aref) - - # efc_address is created on demand np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address) def test_get_data_batched(self): diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py index 39151edf..2498aae2 100644 --- a/mjx/mujoco/mjx/_src/mesh.py +++ b/mjx/mujoco/mjx/_src/mesh.py @@ -15,14 +15,15 @@ """Mesh processing.""" import collections -import dataclasses import itertools -from typing import Dict, List, Optional, Sequence, Tuple +from typing import Tuple import warnings -import mujoco +import jax +from jax import numpy as jp # pylint: disable=g-importing-member -from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.collision_types import ConvexInfo +from mujoco.mjx._src.collision_types import GeomInfo from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member import numpy as np @@ -30,37 +31,7 @@ from scipy import spatial import trimesh -_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1))) -# pyformat: disable -# Rectangular box faces using a counter-clockwise winding order convention. -_BOX_FACES = [ - 0, 4, 5, 1, # left - 0, 2, 6, 4, # bottom - 6, 7, 5, 4, # front - 2, 3, 7, 6, # right - 1, 5, 7, 3, # top - 0, 1, 3, 2, # back -] -# pyformat: enable _MAX_HULL_FACE_VERTICES = 20 -_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {} -_DERIVED_ARGS = [ - 'geom_convex_face', - 'geom_convex_vert', - 'geom_convex_edge_dir', - 'geom_convex_facenormal', - 'geom_convex_edge', - 'geom_convex_edge_face_normal', -] -DERIVED = {(Model, d) for d in _DERIVED_ARGS} - - -def _box(size: np.ndarray): - """Creates a mesh for a box with rectangular faces.""" - box_corners = np.array(_BOX_CORNERS) - vert = box_corners * size.reshape(-1, 3) - face = np.array([_BOX_FACES]).reshape(-1, 4) - return vert, face def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray: @@ -170,16 +141,7 @@ def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray: return hull_point_idx -@dataclasses.dataclass -class MeshInfo: - name: str - vert: np.ndarray - face: np.ndarray - convex_vert: Optional[np.ndarray] - convex_face: Optional[np.ndarray] - - -def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray: +def _merge_coplanar(m: Model, tm: trimesh.Trimesh, meshid: int) -> np.ndarray: """Merges coplanar facets.""" if not tm.facets: return tm.faces.copy() # no facets @@ -204,9 +166,11 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray: # 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') warnings.warn( - f'Mesh "{mesh_info.name}" has a coplanar face with more than' - f' {_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance ' + f'Mesh "{name}" has a coplanar face with more than ' + f'{_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance ' 'issues and inaccuracies in collision detection. Consider ' 'decimating the mesh.' ) @@ -231,106 +195,99 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray: return np.concatenate([faces, facets]) -def _mesh_info( - m: mujoco.MjModel, -) -> List[MeshInfo]: - """Extracts mesh info from MjModel.""" - mesh_infos = [] - for i in range(m.nmesh): - name = mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_MESH.value, i) - - last = (i + 1) >= m.nmesh - face_start = m.mesh_faceadr[i] - face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0] - face = m.mesh_face[face_start:face_end] - - vert_start = m.mesh_vertadr[i] - vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0] - vert = m.mesh_vert[vert_start:vert_end] - - graphadr = m.mesh_graphadr[i] - if graphadr < 0: - mesh_infos.append(MeshInfo(name, vert, face, None, None)) - continue - - graph = m.mesh_graph[graphadr:] - numvert, numface = graph[0], graph[1] - - # unused vert_edgeadr - # vert_edgeadr = graph[2 : numvert + 2] - last_idx = numvert + 2 - - vert_globalid = graph[last_idx : last_idx + numvert] - last_idx += numvert - - # unused edge_localid - # edge_localid = graph[last_idx : last_idx + numvert + 3 * numface] - last_idx += numvert + 3 * numface - - face_globalid = graph[last_idx : last_idx + 3 * numface] - face_globalid = face_globalid.reshape((numface, 3)) - - convex_vert = vert[vert_globalid] - vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0]))) - convex_face = np.vectorize(vertex_map.get)(face_globalid) - mesh_infos.append(MeshInfo(name, vert, face, convex_vert, convex_face)) - - return mesh_infos - - -def _geom_mesh_kwargs( - mesh_info: MeshInfo, -) -> Dict[str, np.ndarray]: - """Generates convex mesh attributes for mjx.Model.""" - tm_convex = trimesh.Trimesh( - vertices=mesh_info.convex_vert, faces=mesh_info.convex_face +def box(info: GeomInfo) -> ConvexInfo: + """Creates a box with rectangular faces.""" + vert = np.array( + list(itertools.product((-1, 1), (-1, 1), (-1, 1))), dtype=float ) + # pyformat: disable + # rectangular box faces using a counter-clockwise winding order convention: + face = np.array( + [ + 0, 4, 5, 1, # left + 0, 2, 6, 4, # bottom + 6, 7, 5, 4, # front + 2, 3, 7, 6, # right + 1, 5, 7, 3, # top + 0, 1, 3, 2, # back + ] + ).reshape((-1, 4)) + # pyformat: enable + face_normal = _get_face_norm(vert, face) + edge, edge_face_normal = _get_edge_normals(face, face_normal) + edge_dir = _get_unique_edge_dir(vert, face) + face = vert[face] # materialize full nface x nvert matrix + + c = ConvexInfo( + info.pos, + info.mat, + vert, + face, + face_normal, + edge, + edge_face_normal, + edge_dir, + ) + c = jax.tree_util.tree_map(jp.array, c) + vert = jax.vmap(jp.multiply, in_axes=(None, 0))(c.vert, info.size) + face = jax.vmap(jp.multiply, in_axes=(None, 0))(c.face, info.size) + c = c.replace(vert=vert, face=face) + + return c + + +def convex(m: Model, mesh_id: int, info: GeomInfo) -> ConvexInfo: + """Processes a mesh for use in convex collision algorithms. + + Args: + m: an MJX model + mesh_id: the mesh id to process + info: pos, mat, size of this geom + + Returns: + a convex mesh info + """ + vert_beg = m.mesh_vertadr[mesh_id] + vert_end = m.mesh_vertadr[mesh_id + 1] if mesh_id < m.nmesh - 1 else None + vert = m.mesh_vert[vert_beg:vert_end] + + graphadr = m.mesh_graphadr[mesh_id] + graph = m.mesh_graph[graphadr:] + graph_idx = 0 + + numvert, numface = graph[0], graph[1] + graph_idx += 2 + + # skip vert_edgeadr (numvert,) + graph_idx += numvert + vert_globalid = graph[graph_idx : graph_idx + numvert] + graph_idx += numvert + + # skip edge_localid (numvert, 3) + graph_idx += numvert + 3 * numface + face_globalid = graph[graph_idx : graph_idx + 3 * numface].reshape((-1, 3)) + + vert = vert[vert_globalid] + vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0]))) + face = np.vectorize(vertex_map.get)(face_globalid) + + tm_convex = trimesh.Trimesh(vertices=vert, faces=face) vert = np.array(tm_convex.vertices) - face = _merge_coplanar(tm_convex, mesh_info) - facenormal = _get_face_norm(vert, face) - edge, edge_face_normal = _get_edge_normals(face, facenormal) - return { - 'geom_convex_face': vert[face], - 'geom_convex_face_vert_idx': face, - 'geom_convex_vert': vert, - 'geom_convex_edge_dir': _get_unique_edge_dir(vert, face), - 'geom_convex_facenormal': facenormal, - 'geom_convex_edge': edge, - 'geom_convex_edge_face_normal': edge_face_normal, - } + face = _merge_coplanar(m, tm_convex, mesh_id) + face_normal = _get_face_norm(vert, face) + edge, edge_face_normal = _get_edge_normals(face, face_normal) + edge_dir = _get_unique_edge_dir(vert, face) + face = vert[face] # materialize full nface x nvert matrix + c = ConvexInfo( + info.pos, + info.mat, + vert, + face, + face_normal, + edge, + edge_face_normal, + edge_dir, + ) -def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]: - """Derives geom mesh attributes for mjx.Model from MjModel.""" - kwargs = {k: [] for k in _DERIVED_ARGS} - mesh_infos = _mesh_info(m) - geom_con = m.geom_conaffinity | m.geom_contype - for geomid in range(m.ngeom): - mesh_info = None - dataid = m.geom_dataid[geomid] - if not geom_con[geomid]: - # ignore visual-only meshes - kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS} - continue - elif m.geom_type[geomid] == GeomType.BOX: - vert, face = _box(m.geom_size[geomid]) - mesh_info = MeshInfo( - name='box', - vert=vert, - face=face, - convex_vert=vert, - convex_face=face, - ) - elif dataid < 0: - kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS} - continue - - mesh_info = mesh_info or mesh_infos[dataid] - vert, face = mesh_info.vert, mesh_info.face - key = (hash(vert.data.tobytes()), hash(face.data.tobytes())) - if key not in _CONVEX_CACHE: - _CONVEX_CACHE[key] = _geom_mesh_kwargs(mesh_info) - - kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS} - - return kwargs + return jax.tree_util.tree_map(jp.array, c) diff --git a/mjx/mujoco/mjx/_src/mesh_test.py b/mjx/mujoco/mjx/_src/mesh_test.py index 266a81f0..235348cd 100644 --- a/mjx/mujoco/mjx/_src/mesh_test.py +++ b/mjx/mujoco/mjx/_src/mesh_test.py @@ -20,7 +20,7 @@ import numpy as np import trimesh -class GeomMeshKwargsTest(absltest.TestCase): +class MeshTest(absltest.TestCase): def test_pyramid(self): """Tests that a triangulated pyramid converts to merged coplanar faces.""" @@ -37,29 +37,21 @@ class GeomMeshKwargsTest(absltest.TestCase): tm = trimesh.Trimesh(vertices=vert, faces=face) tm_convex = trimesh.convex.convex_hull(tm) convex_vert = np.array(tm_convex.vertices) - convex_face = np.array(tm_convex.faces) - mesh_info = mesh.MeshInfo( - name='test', - vert=vert, - face=face, - convex_vert=convex_vert, - convex_face=convex_face, - ) - h = mesh._geom_mesh_kwargs(mesh_info) + convex_face = mesh._merge_coplanar(None, tm_convex, 0) # get index of vertices in h['geom_convex_vert'] for vertices in vert dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile( - h['geom_convex_vert'], (vert.shape[0], 1) + convex_vert, (vert.shape[0], 1) ) dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1)) vidx = np.argmin(dist, axis=0) # check verts - np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx]) + np.testing.assert_array_equal(convex_vert, vert[vidx]) # check face vertices map_ = {v: k for k, v in enumerate(vidx)} - h_face = np.vectorize(map_.get)(h['geom_convex_face_vert_idx']) + 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)]) @@ -69,7 +61,8 @@ class GeomMeshKwargsTest(absltest.TestCase): ) # check edges - unique_edge = np.vectorize(map_.get)(h['geom_convex_edge_dir']) + edge_dir = mesh._get_unique_edge_dir(convex_vert, convex_face) + unique_edge = np.vectorize(map_.get)(edge_dir) unique_edge = np.array(sorted(unique_edge.tolist())) np.testing.assert_array_equal( unique_edge, @@ -77,10 +70,11 @@ class GeomMeshKwargsTest(absltest.TestCase): ) # face normals - self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3)) + face_normal = mesh._get_face_norm(convex_vert, convex_face) + self.assertEqual(face_normal.shape, (5, 3)) # face edges - edges = h['geom_convex_edge'] + edges, edge_normal = mesh._get_edge_normals(convex_face, face_normal) edges = np.vectorize(map_.get)(edges) mask = edges[:, 0] != edges[:, 1] edges = edges[mask] @@ -103,7 +97,6 @@ class GeomMeshKwargsTest(absltest.TestCase): ) # face edge normals - edge_normal = h['geom_convex_edge_face_normal'] edge_normal = edge_normal[mask] edge_normal = np.take_along_axis( edge_normal, sort_col_idx[..., None], axis=1 diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index 540cd094..195da43a 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/mjx/mujoco/mjx/_src/solver.py @@ -17,7 +17,6 @@ import jax from jax import numpy as jp import mujoco -from mujoco.mjx._src import constraint from mujoco.mjx._src import math from mujoco.mjx._src import smooth from mujoco.mjx._src import support @@ -81,7 +80,7 @@ class _Context(PyTreeNode): prev_cost=0.0, solver_niter=0, ) - ctx = _update_constraint(m, d, ctx) + ctx = _update_constraint(d, ctx) if grad: ctx = _update_gradient(m, d, ctx) ctx = ctx.replace(search=-ctx.Mgrad) # start with preconditioned gradient @@ -107,7 +106,7 @@ class _LSPoint(PyTreeNode): @classmethod def create( cls, - m: Model, + d: Data, ctx: _Context, alpha: jax.Array, jv: jax.Array, @@ -118,8 +117,7 @@ class _LSPoint(PyTreeNode): # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c # TODO(robotics-team): change this to support friction constraints - ne, nf, *_ = constraint.count_constraints(m) - active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True) + active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True) quad = jax.vmap(jp.multiply)(quad, active) # only active quad_total = quad_gauss + jp.sum(quad, axis=0) @@ -161,13 +159,12 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter): return jax.lax.scan(_fun, init, None, length=max_iter)[0][0] -def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: +def _update_constraint(d: Data, ctx: _Context) -> _Context: """Updates constraint force and resulting cost given latst solver iteration. Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c Args: - m: model defining constraints d: data which contains latest qacc and smooth terms ctx: current solver context @@ -177,8 +174,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context: # TODO(robotics-team): add friction constraints # only count active constraints - ne, nf, *_ = constraint.count_constraints(m) - active = (ctx.Jaref < 0).at[:ne + nf].set(True) + active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True) efc_force = d.efc_D * -ctx.Jaref * active qfrc_constraint = d.efc_J.T @ efc_force @@ -217,8 +213,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context: if m.opt.solver == SolverType.CG: mgrad = smooth.solve_m(m, d, grad) elif m.opt.solver == SolverType.NEWTON: - ne, nf, *_ = constraint.count_constraints(m) - active = (ctx.Jaref < 0).at[: ne + nf].set(True) + active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True) h = (d.efc_J.T * d.efc_D * active) @ d.efc_J h = support.full_m(m, d) + h h_ = jax.scipy.linalg.cho_factor(h) @@ -262,7 +257,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv)) quad = (quad * d.efc_D).T - point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss) + point_fn = lambda a: _LSPoint.create(d, ctx, a, jv, quad, quad_gauss) def cond(ctx: _LSContext) -> jax.Array: done = ctx.ls_iter >= m.opt.ls_iterations @@ -336,7 +331,7 @@ def solve(m: Model, d: Data) -> Data: def body(ctx: _Context) -> _Context: ctx = _linesearch(m, d, ctx) prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad # pylint: disable=invalid-name - ctx = _update_constraint(m, d, ctx) + ctx = _update_constraint(d, ctx) ctx = _update_gradient(m, d, ctx) # polak-ribiere: diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 28314def..7e1bfb3d 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -208,3 +208,8 @@ def get_custom_numeric(m: Union[Model, mujoco.MjModel], name: str) -> float: return m.numeric_data[m.numeric_adr[i]] return -1 + + +def get_custom_int(m: Union[Model, mujoco.MjModel], name: str) -> int: + """Returns a custom integer given an MjModel or mjx.Model.""" + return int(get_custom_numeric(m, name)) diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 38d2c4c3..1df81837 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -15,10 +15,8 @@ """Base types used in MJX.""" import enum -from typing import List, Optional import jax -import jax.numpy as jp import mujoco from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member import numpy as np @@ -213,6 +211,22 @@ class BiasType(enum.IntEnum): # unsupported: MUSCLE, USER +class ConstraintType(enum.IntEnum): + """Type of constraint. + + Attributes: + EQUALITY: equality constraint + LIMIT_JOINT: joint limit + CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone + """ + EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY + # unsupported: FRICTION_DOF, FRICTION_TENDON + LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT + # unsupported: LIMIT_TENDON, CONTACT_FRICTIONLESS + CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL + # unsupported: CONTACT_ELLIPTIC + + class CamLightType(enum.IntEnum): """Type of camera light. @@ -371,6 +385,7 @@ class Model(PyTreeNode): geom_solref: constraint solver reference: contact (ngeom, mjNREF) geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) geom_size: geom-specific size parameters (ngeom, 3) + geom_rbound: radius of bounding sphere (ngeom,) geom_pos: local position offset rel. to body (ngeom, 3) geom_quat: local orientation offset rel. to body (ngeom, 4) geom_friction: friction for (slide, spin, roll) (ngeom, 3) @@ -388,17 +403,13 @@ class Model(PyTreeNode): cam_poscom0: global position rel. to sub-com in qpos0 (ncam, 3) cam_pos0: global position rel. to body in qpos0 (ncam, 3) cam_mat0: global orientation in qpos0 (ncam, 9) - mat_rgba: rgba (nmat, 4) - mesh_vertadr: first vertex address (nmesh x 1) - mesh_faceadr: first face address (nmesh x 1) + mesh_vertadr: first vertex address (nmesh,) + mesh_faceadr: first face address (nmesh,) + mesh_graphadr: graph data address; -1: no graph (nmesh,) mesh_vert: vertex positions for all meshes (nmeshvert, 3) mesh_face: vertex face data (nmeshface, 3) - geom_convex_face: vertex face data, MJX only (ngeom,) - geom_convex_vert: vertex data, MJX only (ngeom,) - geom_convex_edge_dir: unique edge direction, MJX only (ngeom,) - geom_convex_facenormal: normal face data, MJX only (ngeom,) - geom_convex_face_edge: edges for each face (ngeom,) - geom_convex_face_edge_normal: edge normals for each face (ngeom,) + mesh_graph: convex graph data (nmeshgraph,) + mat_rgba: rgba (nmat, 4) pair_dim: contact dimensionality (npair,) pair_geom1: id of geom1 (npair,) pair_geom2: id of geom2 (npair,) @@ -435,6 +446,7 @@ class Model(PyTreeNode): actuator_gear: scale length and transmitted force (nu, 6) numeric_adr: address of field in numeric_data (nnumeric,) numeric_data: array of all numeric fields (nnumericdata,) + name_meshadr: mesh name pointers (nmesh,) name_numericadr: numeric name pointers (nnumeric,) names: names of all objects, 0-terminated (nnames,) """ @@ -516,6 +528,7 @@ class Model(PyTreeNode): geom_solref: jax.Array geom_solimp: jax.Array geom_size: jax.Array + geom_rbound: jax.Array geom_pos: jax.Array geom_quat: jax.Array geom_friction: jax.Array @@ -535,18 +548,14 @@ class Model(PyTreeNode): cam_mat0: jax.Array mesh_vertadr: np.ndarray mesh_faceadr: np.ndarray + mesh_graphadr: np.ndarray mesh_vert: np.ndarray mesh_face: np.ndarray + mesh_graph: np.ndarray mat_rgba: np.ndarray pair_dim: np.ndarray pair_geom1: np.ndarray pair_geom2: np.ndarray - geom_convex_face: List[Optional[jax.Array]] - geom_convex_vert: List[Optional[jax.Array]] - geom_convex_edge_dir: List[Optional[jax.Array]] - geom_convex_facenormal: List[Optional[jax.Array]] - geom_convex_edge: List[Optional[jax.Array]] - geom_convex_edge_face_normal: List[Optional[jax.Array]] pair_solref: jax.Array pair_solreffriction: jax.Array pair_solimp: jax.Array @@ -580,6 +589,7 @@ class Model(PyTreeNode): actuator_gear: jax.Array numeric_adr: np.ndarray numeric_data: np.ndarray + name_meshadr: np.ndarray name_numericadr: np.ndarray names: bytes @@ -596,8 +606,11 @@ class Contact(PyTreeNode): solref: constraint solver reference, normal direction (mjNREF,) solreffriction: constraint solver reference, friction directions (mjNREF,) solimp: constraint solver impedance (mjNIMP,) - geom1: id of geom 1 - geom2: id of geom 2 + dim: contact space dimensionality: 1, 3, 4, or 6 + geom1: id of geom 1; deprecated, use geom[0] + geom2: id of geom 2; deprecated, use geom[1] + geom: geom ids (2,) + efc_address: address in efc; -1: not included """ dist: jax.Array pos: jax.Array @@ -607,33 +620,25 @@ class Contact(PyTreeNode): solref: jax.Array solreffriction: jax.Array solimp: jax.Array - # unsupported: mu, H, dim + # unsupported: mu, H + dim: np.ndarray geom1: jax.Array geom2: jax.Array - # unsupported: efc_address, exclude - - @classmethod - def zero(cls, ncon: int = 0) -> 'Contact': - """Returns a contact filled with zeros.""" - return Contact( - dist=jp.zeros(ncon), - pos=jp.zeros((ncon, 3,)), - frame=jp.zeros((ncon, 3, 3)), - includemargin=jp.zeros(ncon), - friction=jp.zeros((ncon, 5)), - solref=jp.zeros((ncon, mujoco.mjNREF)), - solreffriction=jp.zeros((ncon, mujoco.mjNREF)), - solimp=jp.zeros((ncon, mujoco.mjNIMP,)), - geom1=jp.zeros(ncon, dtype=int), - geom2=jp.zeros(ncon, dtype=int), - ) + geom: jax.Array + # unsupported: flex, elem, vert, exclude + efc_address: np.ndarray class Data(PyTreeNode): - r"""Dynamic state that updates each step.\ + r"""Dynamic state that updates each step. Attributes: - solver_niter: number of solver iterations, per island (mjNISLAND,) + ne: number of equality constraints + nf: number of friction constraints + nl: number of limit constraints + nefc: number of constraints + ncon: number of contacts + solver_niter: number of solver iterations time: simulation time qpos: position (nq,) qvel: velocity (nv,) @@ -663,14 +668,15 @@ class Data(PyTreeNode): cinert: com-based body inertia and mass (nbody, 10) actuator_length: actuator lengths (nu,) actuator_moment: actuator moments (nu, nv) - crb: com-based composite inertia and mass (nbody, 10) + crb: com-based composite inertia and mass (nbody, 10) \ qM: total inertia if sparse: (nM,) if dense: (nv, nv) qLD: L'*D*L (or Cholesky) factorization of M. if sparse: (nM,) if dense: (nv, nv) qLDiagInv: 1/diag(D) if sparse: (nv,) if dense: (0,) - contact: list of all detected contacts (ncon,) + contact: all detected contacts (ncon,) + efc_type: constraint type (nefc,) efc_J: constraint Jacobian (nefc, nv) efc_frictionloss: frictionloss (friction) (nefc,) efc_D: constraint mass (nefc,) @@ -689,6 +695,12 @@ class Data(PyTreeNode): efc_force: constraint force in constraint space (nefc,) userdata: user data, not touched by engine (nuserdata,) """ + # constant sizes: + ne: int + nf: int + nl: int + nefc: int + ncon: int # solver statistics: solver_niter: jax.Array # global properties: @@ -732,6 +744,7 @@ class Data(PyTreeNode): qLD: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name contact: Contact + efc_type: np.ndarray efc_J: jax.Array # pylint:disable=invalid-name efc_frictionloss: jax.Array efc_D: jax.Array # pylint:disable=invalid-name diff --git a/mjx/mujoco/mjx/integration_test/collision_driver_test.py b/mjx/mujoco/mjx/integration_test/collision_driver_test.py index 06116e67..3951b80c 100644 --- a/mjx/mujoco/mjx/integration_test/collision_driver_test.py +++ b/mjx/mujoco/mjx/integration_test/collision_driver_test.py @@ -30,7 +30,10 @@ import numpy as np def _assert_attr_eq(mjx_d, mj_d, attr, name, atol): if attr == 'efc_address': - # we do not test efc_address since it gets set in constraint logic + # contact order not guaranteed to match + np.testing.assert_array_equal( + np.sort(mjx_d.efc_address), np.sort(mj_d.efc_address) + ) return err_msg = f'mismatch: {attr} in run: {name}' mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr) @@ -82,6 +85,9 @@ class CollisionDriverIntegrationTest(parameterized.TestCase): mjx_contact = jax.tree_util.tree_map( lambda x: x.take(np.array(idx), axis=0), dx.contact ) + mjx_contact = mjx_contact.replace( + dim=mjx_contact.dim[idx], efc_address=mjx_contact.efc_address[idx] + ) for field in dataclasses.fields(Contact): _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)