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

Prepare MJX for condim.

Browse Source 
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
This commit is contained in:
Erik Frey
2024-04-18 12:42:10 -07:00
committed by Copybara-Service
parent 8b7f1094f1
commit a4df912018
20 changed files with 868 additions and 914 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+11 -4
View File
@@ -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)
mjx/mujoco/mjx/__init__.py
-3
View File
@@ -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
mjx/mujoco/mjx/_src/collision_base.py
-67
View File
@@ -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
mjx/mujoco/mjx/_src/collision_convex.py
+59 -25
View File
@@ -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
mjx/mujoco/mjx/_src/collision_driver.py
+274 -332
View File
@@ -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 <pair> (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)
mjx/mujoco/mjx/_src/collision_driver_test.py
+21 -15
View File
@@ -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)
mjx/mujoco/mjx/_src/collision_primitive.py
+51 -48
View File
@@ -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)
mjx/mujoco/mjx/_src/collision_sdf.py
+28 -15
View File
@@ -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)
mjx/mujoco/mjx/_src/collision_types.py
+65
View File
@@ -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
mjx/mujoco/mjx/_src/constraint.py
+45 -26
View File
@@ -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."""
mjx/mujoco/mjx/_src/constraint_test.py
+3 -3
View File
@@ -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)
mjx/mujoco/mjx/_src/device.py
+19 -26
View File
@@ -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):
mjx/mujoco/mjx/_src/io.py
+105 -127
View File
@@ -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)
mjx/mujoco/mjx/_src/io_test.py
+1 -6
View File
@@ -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):
mjx/mujoco/mjx/_src/mesh.py
+102 -145
View File
@@ -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)
mjx/mujoco/mjx/_src/mesh_test.py
+10 -17
View File
@@ -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
mjx/mujoco/mjx/_src/solver.py
+8 -13
View File
@@ -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:
mjx/mujoco/mjx/_src/support.py
+5
View File
@@ -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))
mjx/mujoco/mjx/_src/types.py
+54 -41
View File
@@ -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
mjx/mujoco/mjx/integration_test/collision_driver_test.py
+7 -1
View File
@@ -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)