diff --git a/doc/changelog.rst b/doc/changelog.rst index d1c7c1ec..010fa6f3 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -14,12 +14,30 @@ General MJX ^^^ -3. Added cylinder plane collisions. +.. admonition:: Breaking API changes + :class: attention + + 3. Removed deprecated ``mjx.device_get_into`` and ``mjx.device_put`` functions as they lack critical new + functionality. + + **Migration:** Use ``mjx.get_data_into`` instead of ``mjx.device_get_into``, and ``mjx.put_data`` instead of + ``mjx.device_put``. + +4. Added cylinder plane collisions. +5. Added ``efc_type`` to ``mjx.Data`` and ``dim``, ``efc_address`` to ``mjx.Contact``. +6. Added ``geom`` to ``mjx.Contact`` and marked ``geom1``, ``geom2`` deprecated. +7. Added ``ne``, ``nf``, ``nl``, ``nefc``, and ``ncon`` to ``mjx.Data`` to match ``mujoco.MjData``. +8. Given the above added fields, removed ``mjx.get_params``, ``mjx.ncon``, and ``mjx.count_constraints``. +9. Changed the way meshes are organized on device to speed up collision detection when a mesh is replicated for many + geoms. +10. 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. Bug introduced in 3.1.4. +11. Defaults of lights were not being saved, now fixed. +12. Prevent overwriting of frame names by body names when saving an XML. Introduced in 3.1.4. +13. 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/doc/models.rst b/doc/models.rst index 9bc624cc..c6bbc815 100644 --- a/doc/models.rst +++ b/doc/models.rst @@ -129,3 +129,17 @@ Drones - Preview * - `Skydio X2 `_ - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/skydio_x2/x2.png + * - `Bitcraze Crazyflie 2 `_ + - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/bitcraze_crazyflie_2/cf2.png + + +Biomechanical +^^^^^^^^^^^^^ + +.. list-table:: + :header-rows: 1 + + * - Model + - Preview + * - `Fruitfly `_ + - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/flybody/flybody.png diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 978f44ab..120db360 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -16,12 +16,7 @@ # 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 from mujoco.mjx._src.forward import euler from mujoco.mjx._src.forward import forward from mujoco.mjx._src.forward import fwd_acceleration 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 726ecade..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_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) @@ -347,7 +389,7 @@ def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: degenerate_edge_dir, edge_closest_pt, cap_closest_pt, - ) = jax.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res) + ) = jax.tree_util.tree_map(lambda x, i=e_idx: jp.take(x, i, axis=0), res) edge_face_normals = edge_face_normal[e_idx] edge_voronoi_front = ((edge_face_normals @ edge_axis) < 0).all() @@ -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 7fd905b4..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_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_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_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_map(jp.concatenate, res) - - # repeat params by the number of contacts per geom pair - geom1, geom2, params = jax.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) - contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts) + # collapse contacts together, ensuring they are grouped by condim + condim_groups = {} + for key, contact in groups.items(): + condim_groups.setdefault(key.condim, []).append(contact) - 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_map(lambda x, idx=idx: jp.take(x, idx, axis=0), 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) 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 d501be5d..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_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_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_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_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) @@ -531,7 +535,8 @@ class ConvexTest(absltest.TestCase): np.testing.assert_array_less(dx.contact.dist[:2], 0) np.testing.assert_array_less(-dx.contact.dist[2:], 0) # extract the contact points with penetration - c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) + 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_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 7407919a..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_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_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_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_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_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_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 999b3e8a..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_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_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 50aee9b3..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 @@ -111,7 +115,7 @@ def _instantiate_equality_connect(m: Model, d: Data) -> Optional[_Efc]: return j, cpos, jp.repeat(math.norm(cpos), 3) # concatenate to drop connect grouping dimension - j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2)) + j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2)) invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0] invweight = jp.repeat(invweight, 3) solref = jp.tile(m.eq_solref[ids], (3, 1)) @@ -164,7 +168,7 @@ def _instantiate_equality_weld(m: Model, d: Data) -> Optional[_Efc]: return j, pos, jp.repeat(math.norm(pos), 6) # concatenate to drop weld grouping dimension - j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2)) + j, pos, pos_norm = jax.tree_util.tree_map(jp.concatenate, fn(data, id1, id2)) invweight = m.body_invweight0[id1] + m.body_invweight0[id2] invweight = jp.repeat(invweight, 3) solref = jp.tile(m.eq_solref[ids], (6, 1)) @@ -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 @@ -308,42 +312,57 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]: res = fn(d.contact) # remove contact grouping dimension: - j, invweight, pos, solref, solimp = jax.tree_map(jp.concatenate, res) + j, invweight, pos, solref, solimp = jax.tree_util.tree_map(jp.concatenate, res) frictionloss = jp.zeros_like(pos) 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.""" @@ -366,7 +385,7 @@ def make_constraint(m: Model, d: Data) -> Data: d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z) return d - efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs) + efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs) @jax.vmap def fn(efc): diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py index a2bd8cc9..99346828 100644 --- a/mjx/mujoco/mjx/_src/constraint_test.py +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -65,14 +65,14 @@ class ConstraintTest(absltest.TestCase): pos = jp.ones(3) m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.REFSAFE - mx = mjx.device_put(m) + mx = mjx.put_model(m) k, *_ = constraint._kbi(mx, solimp, solref, pos) self.assertEqual(k, 1 / (0.99**2 * timeconst**2)) 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 deleted file mode 100644 index 0a8130e2..00000000 --- a/mjx/mujoco/mjx/_src/device.py +++ /dev/null @@ -1,319 +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. -# ============================================================================== -"""Get and put mujoco data on/off device.""" - -import copy -import dataclasses -from typing import Any, Dict, Iterable, List, Union, overload -import warnings - -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 - -_MJ_TYPE_ATTR = { - mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,), - mujoco.mjtDyn: (mujoco.MjModel.actuator_dyntype,), - mujoco.mjtEq: (mujoco.MjModel.eq_type,), - mujoco.mjtGain: (mujoco.MjModel.actuator_gaintype,), - mujoco.mjtTrn: (mujoco.MjModel.actuator_trntype,), - mujoco.mjtCone: ( - mujoco.MjModel.opt, - mujoco.MjOption.cone, - ), - mujoco.mjtIntegrator: ( - mujoco.MjModel.opt, - mujoco.MjOption.integrator, - ), - mujoco.mjtSolver: ( - mujoco.MjModel.opt, - mujoco.MjOption.solver, - ), -} - -_TYPE_MAP = { - mujoco._structs._MjContactList: types.Contact, # pylint: disable=protected-access - mujoco.MjData: types.Data, - mujoco.MjModel: types.Model, - mujoco.MjOption: types.Option, - mujoco.MjStatistic: types.Statistic, - mujoco.mjtBias: types.BiasType, - mujoco.mjtCone: types.ConeType, - mujoco.mjtDisableBit: types.DisableBit, - mujoco.mjtDyn: types.DynType, - mujoco.mjtEq: types.EqType, - mujoco.mjtGain: types.GainType, - mujoco.mjtIntegrator: types.IntegratorType, - mujoco.mjtSolver: types.SolverType, - mujoco.mjtTrn: types.TrnType, -} - -_TRANSFORMS = { - (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), - (types.Contact, 'frame'): ( - lambda x: x.reshape(x.shape[:-1] + (3, 3)) # pylint: disable=g-long-lambda - if x is not None and x.shape[0] - else jp.zeros((0, 3, 3)) - ), -} - -_INVERSE_TRANSFORMS = { - (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Data, 'cam_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Model, 'cam_mat0'): lambda x: x.reshape(x.shape[:-2] + (9,)), - (types.Contact, 'frame'): ( - lambda x: x.reshape(x.shape[:-2] + (9,)) # pylint: disable=g-long-lambda - if x is not None and x.shape[0] - else jp.zeros((0, 9)) - ), -} - -_DERIVED = mesh.DERIVED.union( - # 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()} - - -def _data_derived(value: mujoco.MjData) -> Dict[str, Any]: - return {'efc_J': jax.device_put(value.efc_J)} - - -def _option_derived(value: types.Option) -> Dict[str, Any]: - has_fluid = ( - value.density > 0 or value.viscosity > 0 or (value.wind != 0.0).any() - ) - return {'has_fluid_params': has_fluid} - - -def _validate(m: mujoco.MjModel): - """Validates that an mjModel is compatible with MJX.""" - - # check enum types - for mj_type, attrs in _MJ_TYPE_ATTR.items(): - val = m - for attr in attrs: - val = attr.fget(val) # pytype: disable=attribute-error - - typs = set(val) if isinstance(val, Iterable) else {val} - unsupported_typs = typs - set(_TYPE_MAP[mj_type]) - unsupported = [mj_type(t) for t in unsupported_typs] # pylint: disable=too-many-function-args - if unsupported: - raise NotImplementedError(f'{unsupported} not implemented.') - - if m.ntendon: - raise NotImplementedError('Tendons are not supported.') - - # check condim - if (m.geom_condim != 3).any() or (m.pair_dim != 3).any(): - raise NotImplementedError('Only condim=3 is supported.') - - 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})' - ) - - # TODO(erikfrey): warn for high solver iterations, nefc, etc. - - # mjNDISABLE is not a DisableBit flag, so must be explicitly ignored - disablebit_members = set(mujoco.mjtDisableBit.__members__.values()) - { - mujoco.mjtDisableBit.mjNDISABLE} - unsupported_disable = disablebit_members - { - mujoco.mjtDisableBit(t.value) for t in types.DisableBit - } - for f in unsupported_disable: - if f & m.opt.disableflags: - warnings.warn(f'Ignoring disable flag {f.name}.') - - # mjNENABLE is not an EnableBit flag, so must be explicitly ignored - unsupported_enable = set(mujoco.mjtEnableBit.__members__.values()) - { - mujoco.mjtEnableBit.mjNENABLE - } - for f in unsupported_enable: - if f & m.opt.enableflags: - warnings.warn(f'Ignoring enable flag {f.name}.') - - if not np.allclose(m.dof_frictionloss, 0): - raise NotImplementedError('dof_frictionloss is not implemented.') - - -@overload -def device_put(value: mujoco.MjData) -> types.Data: - ... - - -@overload -def device_put(value: mujoco.MjModel) -> types.Model: - ... - - -def device_put(value): - """Places mujoco data onto a device. - - Args: - value: a mujoco struct to transfer - - Returns: - on-device MJX struct reflecting the input value - """ - warnings.warn( - 'device_put is deprecated, use put_model and put_data instead', - category=DeprecationWarning, - ) - - clz = _TYPE_MAP.get(type(value)) - if clz is None: - raise NotImplementedError(f'{type(value)} is not supported for device_put.') - - if isinstance(value, mujoco.MjModel): - _validate(value) # type: ignore - - init_kwargs = {} - for f in dataclasses.fields(clz): # type: ignore - if (clz, f.name) in _DERIVED: - continue - - field_value = getattr(value, f.name) - if (clz, f.name) in _TRANSFORMS: - field_value = _TRANSFORMS[(clz, f.name)](field_value) - - if f.type is jax.Array: - field_value = jax.device_put(field_value) - elif type(field_value) in _TYPE_MAP.keys(): - field_value = device_put(field_value) - - init_kwargs[f.name] = copy.copy(field_value) - - derived_kwargs = {} - if isinstance(value, mujoco.MjModel): - derived_kwargs = _model_derived(value) - elif isinstance(value, mujoco.MjData): - derived_kwargs = _data_derived(value) - elif isinstance(value, mujoco.MjOption): - derived_kwargs = _option_derived(value) - - return clz(**init_kwargs, **derived_kwargs) # type: ignore - - -@overload -def device_get_into( - result: Union[mujoco.MjData, List[mujoco.MjData]], value: types.Data -): - ... - - -def device_get_into(result, value): - """Transfers data off device into a mujoco MjData. - - Data on device often has a batch dimension which adds (N,) to the beginning - of each array shape where N = batch size. - - If result is a single MjData, arrays are copied over with the batch dimension - intact. If result is a list, the list must be length N and will be populated - with distinct MjData structs where the batch dimension is stripped. - - Args: - result: struct (or list of structs) to transfer into - value: device value to transfer - - Raises: - RuntimeError: if result length doesn't match data batch size - """ - warnings.warn( - 'device_get_into is deprecated, use get_data instead', - category=DeprecationWarning, - ) - - value = jax.device_get(value) - - if isinstance(result, list): - array_shapes = [s.shape for s in jax.tree_util.tree_flatten(value)[0]] - - if any(len(s) < 1 or s[0] != array_shapes[0][0] for s in array_shapes): - raise ValueError('unrecognizable batch dimension in value') - - batch_size = array_shapes[0][0] - - if len(result) != batch_size: - raise ValueError( - f"result length ({len(result)}) doesn't match value batch size" - f' ({batch_size})' - ) - - for i in range(batch_size): - value_i = jax.tree_map(lambda x, i=i: x[i], value) - device_get_into(result[i], value_i) - - else: - if isinstance(result, mujoco.MjData): - ncon = value.contact.dist.shape[0] - nefc = value.efc_J.shape[0] - mujoco._functions._realloc_con_efc( # pylint: disable=protected-access - result, ncon=ncon, nefc=nefc - ) - result.ncon = ncon - result.nefc = nefc - efc_start = nefc - ncon * 4 - result.contact.efc_address[:] = np.arange(efc_start, nefc, 4) - result.contact.dim[:] = 3 - - for f in dataclasses.fields(value): # type: ignore - if (type(value), f.name) in _DERIVED: - continue - - field_value = getattr(value, f.name) - - if (type(value), f.name) in _INVERSE_TRANSFORMS: - field_value = _INVERSE_TRANSFORMS[(type(value), f.name)](field_value) - - if type(field_value) in _TYPE_MAP.values(): - device_get_into(getattr(result, f.name), field_value) - continue - - try: - setattr(result, f.name, field_value) - except AttributeError: - getattr(result, f.name)[:] = field_value diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py deleted file mode 100644 index 6126f707..00000000 --- a/mjx/mujoco/mjx/_src/device_test.py +++ /dev/null @@ -1,203 +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. -# ============================================================================== -"""Tests for moving mujoco structs on and off device.""" - -import dataclasses - -from absl.testing import absltest -from absl.testing import parameterized -import jax -from jax import numpy as jp -import mujoco -from mujoco import mjx -from mujoco.mjx._src import device -from mujoco.mjx._src import test_util -from mujoco.mjx._src import types -# pylint: disable=g-importing-member -from mujoco.mjx._src.dataclasses import PyTreeNode -# pylint: enable=g-importing-member -import numpy as np - - -def _assert_eq(testcase, a, b, attr=None, name=None): - if (type(a), attr) in device._DERIVED: - return - - if attr: - a, b = getattr(a, attr), getattr(b, attr) - - if isinstance(a, PyTreeNode): - for field in dataclasses.fields(a): - _assert_eq(testcase, a, b, field.name, type(a).__name__) - return - - typ = {'Model': types.Model, 'Data': types.Data, - 'Contact': types.Contact}.get(name) - if (typ, attr) in device._TRANSFORMS: - b = device._TRANSFORMS[(typ, attr)](b) - - err_msg = f'mismatch: {attr} in {name}' - if not hasattr(b, 'shape') or not b.shape: - testcase.assertEqual(a, b, err_msg) - return - - a, b = np.array(a), np.array(b) - np.testing.assert_allclose(a, b, err_msg=err_msg, atol=1e-8) - - -class DeviceTest(parameterized.TestCase): - - @parameterized.parameters('constraints.xml', 'pendula.xml') - def testdevice_put(self, fname): - """Test putting MjData and MjModel on device.""" - m = test_util.load_test_file(fname) - # advance state to ensure non-zero fields - d = mujoco.MjData(m) - for _ in range(10): - mujoco.mj_step(m, d) - - _assert_eq(self, mjx.device_put(d), d) - _assert_eq(self, mjx.device_put(m), m) - - @parameterized.parameters('constraints.xml', 'pendula.xml') - def testdevice_get(self, fname): - """Test getting MjData from a device.""" - m = test_util.load_test_file(fname) - m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing - mx = device.device_put(m) - dx = mjx.make_data(mx) - d = mujoco.MjData(m) - device.device_get_into(d, dx) - _assert_eq(self, dx, d) - - @parameterized.parameters('constraints.xml', 'pendula.xml') - def testdevice_get_batched(self, fname): - """Test getting MjData from a device.""" - m = test_util.load_test_file(fname) - m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE # force sparse for testing - mx = device.device_put(m) - batch_size = 32 - - # create mjx_data and batch it - dx = mjx.make_data(mx) - dx = jax.tree_map( - lambda x: jp.repeat(x, batch_size).reshape((batch_size,) + x.shape), - dx, - ) - ds = [mujoco.MjData(m) for _ in range(batch_size - 1)] - - with self.assertRaises(ValueError): - device.device_get_into(ds, dx) - - ds = [mujoco.MjData(m) for _ in range(batch_size)] - device.device_get_into(ds, dx) - dx = jax.device_get(dx) # faster indexing for testing - for i in range(batch_size): - _assert_eq(self, jax.tree_map(lambda x, i=i: x[i], dx), ds[i]) - - -class ValidateInputTest(absltest.TestCase): - - def test_solver(self): - m = mujoco.MjModel.from_xml_string( - '' - ) - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_integrator(self): - m = mujoco.MjModel.from_xml_string( - '' - ) - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_cone(self): - m = mujoco.MjModel.from_xml_string( - '' - ) - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_dyn(self): - m = test_util.load_test_file('pendula.xml') - m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_gain(self): - m = test_util.load_test_file('pendula.xml') - m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_bias(self): - m = test_util.load_test_file('pendula.xml') - m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_condim(self): - m = test_util.load_test_file('constraints.xml') - for i in [1, 4, 6]: - m.geom_condim[0] = i - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_geoms(self): - m = mujoco.MjModel.from_xml_string(""" - - - - - - - - - - - - - """) - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - def test_tendon(self): - m = mujoco.MjModel.from_xml_string(""" - - - - - - - - - - - - - - - - - - - """) - with self.assertRaises(NotImplementedError): - mjx.device_put(m) - - -if __name__ == '__main__': - absltest.main() diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index f04eafe3..905e463a 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -310,7 +310,7 @@ def rungekutta4(m: Model, d: Data) -> Data: kqvel = d.qvel # intermediate RK solution # RK solutions sum - qvel, qacc, act_dot = jax.tree_map( + qvel, qacc, act_dot = jax.tree_util.tree_map( lambda k: B[0] * k, (kqvel, d.qacc, d.act_dot) ) integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep) @@ -318,7 +318,7 @@ def rungekutta4(m: Model, d: Data) -> Data: def f(carry, x): qvel, qacc, act_dot, kqvel, d = carry a, b, t = x # tableau numbers - dqvel, dqacc, dact_dot = jax.tree_map( + dqvel, dqacc, dact_dot = jax.tree_util.tree_map( lambda k: a * k, (kqvel, d.qacc, d.act_dot) ) # get intermediate RK solutions diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 3e90b26e..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): @@ -295,12 +270,11 @@ def get_data_into( j = m.dof_parentid[j] for i in range(batch_size): - d_i = jax.tree_map(lambda x, i=i: x[i], d) if batched else d + d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d 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_map(pad_fn, fields) + 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 36806c2e..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): @@ -435,7 +430,7 @@ class DataIOTest(parameterized.TestCase): mujoco.mj_step(m, d, 2) dx = mjx.put_data(m, d) # second data in batch has contact dist > 0, disables contact - dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) + dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx) ds = mjx.get_data(m, dx_b) self.assertLen(ds, 2) np.testing.assert_allclose(ds[0].qpos, d.qpos) 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/scan.py b/mjx/mujoco/mjx/_src/scan.py index 168637dc..a0cf3037 100644 --- a/mjx/mujoco/mjx/_src/scan.py +++ b/mjx/mujoco/mjx/_src/scan.py @@ -62,7 +62,7 @@ def _take(obj: Y, idx: np.ndarray) -> Y: x = x.take(jp.array(idx), axis=0, mode='wrap') return x - return jax.tree_map(take, obj) + return jax.tree_util.tree_map(take, obj) def _q_bodyid(m: Model) -> np.ndarray: @@ -120,7 +120,7 @@ def _nvmap(f: Callable[..., Y], *args) -> Y: args = [a if n is None else None for n, a in zip(np_args, args)] # remove empty args that we should not vmap over - args = jax.tree_map(lambda a: a if a.shape[0] else None, args) + args = jax.tree_util.tree_map(lambda a: a if a.shape[0] else None, args) in_axes = [None if a is None else 0 for a in args] def outer_f(*args, np_args=np_args): @@ -322,7 +322,7 @@ def flat( [v if typ in flat_ else jp.concatenate(v) for v, typ in zip(y, out_types)] for y in ys ] - ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys) + ys = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *ys) # put concatenated results back in order reordered_ys = [] @@ -465,15 +465,15 @@ def body_tree( def index_sum(x, i=id_map, s=body_ids.size): return jax.ops.segment_sum(x, i, s) - y = jax.tree_map(index_sum, y) - carry = y if carry is None else jax.tree_map(jp.add, carry, y) + y = jax.tree_util.tree_map(index_sum, y) + carry = y if carry is None else jax.tree_util.tree_map(jp.add, carry, y) elif key in key_parents: ys = [key_y[p] for p in key_parents[key]] - y = jax.tree_map(lambda *x: jp.concatenate(x), *ys) + y = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *ys) body_ids = np.concatenate([key_body_ids[p] for p in key_parents[key]]) parent_ids = m.body_parentid[key_body_ids[key]] take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i) - carry = jax.tree_map(take_fn, y) + carry = jax.tree_util.tree_map(take_fn, y) f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])] key_y[key] = _nvmap(f, carry, *f_args) @@ -488,8 +488,8 @@ def body_tree( if len(out_types) > 1: y_typ = [y_[i] for y_ in y_typ] if typ != 'b': - y_typ = jax.tree_map(jp.concatenate, y_typ) - y_typ = jax.tree_map(lambda *x: jp.concatenate(x), *y_typ) + y_typ = jax.tree_util.tree_map(jp.concatenate, y_typ) + y_typ = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *y_typ) y_take = np.argsort(np.concatenate([key_y_take[key][i] for key in keys])) _check_output(y_typ, y_take, typ, i) y.append(_take(y_typ, y_take)) diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py index 456067d6..edb57339 100644 --- a/mjx/mujoco/mjx/_src/scan_test.py +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -55,7 +55,7 @@ class ScanTest(absltest.TestCase): """) - m = mjx.device_put(m) + m = mjx.put_model(m) def fn(body_id): return body_id + 1 @@ -69,7 +69,7 @@ class ScanTest(absltest.TestCase): def test_flat_joints(self): """Tests scanning over bodies with joints of different types.""" m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) - m = mjx.device_put(m) + m = mjx.put_model(m) # we will test two functions: # 1) j_fn receives jnt_types as a jp array @@ -105,7 +105,7 @@ class ScanTest(absltest.TestCase): def test_body_tree(self): """Tests tree scanning over bodies with different joint counts.""" m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) - m = mjx.device_put(m) + m = mjx.put_model(m) # we will test two functions: # 1) j_fn receives jnt_pos which is a jp array @@ -196,7 +196,7 @@ class ScanTest(absltest.TestCase): def test_scan_actuators(self): """Tests scanning over actuators.""" m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML) - m = mjx.device_put(m) + m = mjx.put_model(m) fn = lambda *args: args args = ( diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index dae5a82b..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 @@ -283,14 +278,14 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: # 1) they are not correctly at a bracket boundary (e.g. lo.deriv_0 > 0), OR # 2) if moving to next or mid narrows the bracket swap_lo_next = (lo.deriv_0 > 0) | (lo.deriv_0 < lo_next.deriv_0) - lo = jax.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next) + lo = jax.tree_util.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next) swap_lo_mid = (mid.deriv_0 < 0) & (lo.deriv_0 < mid.deriv_0) - lo = jax.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid) + lo = jax.tree_util.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid) swap_hi_next = (hi.deriv_0 < 0) | (hi.deriv_0 > hi_next.deriv_0) - hi = jax.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next) + hi = jax.tree_util.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next) swap_hi_mid = (mid.deriv_0 > 0) & (hi.deriv_0 > mid.deriv_0) - hi = jax.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid) + hi = jax.tree_util.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid) swap = swap_lo_next | swap_lo_mid | swap_hi_next | swap_hi_mid @@ -302,8 +297,8 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context: p0 = point_fn(jp.array(0.0)) lo = point_fn(p0.alpha - p0.deriv_0 / p0.deriv_1) lesser_fn = lambda x, y: jp.where(lo.deriv_0 < p0.deriv_0, x, y) - hi = jax.tree_map(lesser_fn, p0, lo) - lo = jax.tree_map(lesser_fn, lo, p0) + hi = jax.tree_util.tree_map(lesser_fn, p0, lo) + lo = jax.tree_util.tree_map(lesser_fn, lo, p0) ls_ctx = _LSContext(lo=lo, hi=hi, swap=jp.array(True), ls_iter=0) ls_ctx = _while_loop_scan(cond, body, ls_ctx, 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 a9891656..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) @@ -79,9 +82,12 @@ class CollisionDriverIntegrationTest(parameterized.TestCase): self.assertSequenceEqual(set(idx_mjx), set(idx_mj)) idx = sorted(range(len(idx_mj)), key=lambda x: idx_mj.index(idx_mjx[x])) - mjx_contact = jax.tree_map( + 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) diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb index f57aee4d..a839c308 100644 --- a/mjx/tutorial.ipynb +++ b/mjx/tutorial.ipynb @@ -829,7 +829,7 @@ "\n", " friction, gain, bias = rand(rng)\n", "\n", - " in_axes = jax.tree_map(lambda x: None, sys)\n", + " in_axes = jax.tree_util.tree_map(lambda x: None, sys)\n", " in_axes = in_axes.tree_replace({\n", " 'geom_friction': 0,\n", " 'actuator_gainprm': 0,\n", diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc index 0282f8ca..8cb6b65c 100644 --- a/python/mujoco/functions.cc +++ b/python/mujoco/functions.cc @@ -106,11 +106,11 @@ PYBIND11_MODULE(_functions, pymodule) { Def(pymodule); Def(pymodule); // Skipped: mj_copyModel (have MjModel.__copy__, memory managed by MjModel) - DEF_WITH_OMITTED_PY_ARGS(traits::mj_saveModel, "buffer_sz")( - pymodule, - [](const raw::MjModel* m, const std::optional& filename, + pymodule.def( + "mj_saveModel", + [](const MjModelWrapper& m, const std::optional& filename = std::nullopt, std::optional< - Eigen::Ref>> buffer) { + Eigen::Ref>> buffer = std::nullopt) { void* buffer_ptr = nullptr; int buffer_sz = 0; if (buffer.has_value()) { @@ -118,9 +118,13 @@ PYBIND11_MODULE(_functions, pymodule) { buffer_sz = buffer->size(); } return InterceptMjErrors(::mj_saveModel)( - m, filename.has_value() ? filename->c_str() : nullptr, + m.get(), filename.has_value() ? filename->c_str() : nullptr, buffer_ptr, buffer_sz); - }); + }, + py::arg("m"), py::arg_v("filename", std::nullopt), + py::arg_v("buffer", std::nullopt), + py::doc(traits::mj_saveModel::doc), + py::call_guard()); // Skipped: mj_loadModel (have MjModel.from_binary_path) // Skipped: mj_deleteModel (have MjModel.__del__) Def(pymodule); diff --git a/python/mujoco/renderer.py b/python/mujoco/renderer.py index 973e8076..afe5fc77 100644 --- a/python/mujoco/renderer.py +++ b/python/mujoco/renderer.py @@ -78,8 +78,11 @@ the clause: # Create render contexts. # TODO(nimrod): Figure out why pytype doesn't like gl_context.GLContext - self._gl_context = gl_context.GLContext(width, height) # type: ignore - self._gl_context.make_current() + self._gl_context = None # type: ignore + if gl_context.GLContext is not None: + self._gl_context = gl_context.GLContext(width, height) + if self._gl_context: + self._gl_context.make_current() self._mjr_context = _render.MjrContext( model, _enums.mjtFontScale.mjFONTSCALE_150.value ) @@ -148,9 +151,11 @@ the clause: self._scene.flags[_enums.mjtRndFlag.mjRND_SEGMENT] = True self._scene.flags[_enums.mjtRndFlag.mjRND_IDCOLOR] = True - if self._gl_context is None: + if self._mjr_context is None: raise RuntimeError('render cannot be called after close.') - self._gl_context.make_current() + + if self._gl_context: + self._gl_context.make_current() if self._depth_rendering: out_shape = (self._height, self._width) diff --git a/src/user/user_api.cc b/src/user/user_api.cc index 6220d98a..f06d3de2 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -91,6 +91,17 @@ int mjs_attachBody(mjsFrame* parent, const mjsBody* child, +// attach frame to a parent body +int mjs_attachFrame(mjsBody* parent, const mjsFrame* child, + const char* prefix, const char* suffix) { + mjCBody* body_parent = static_cast(parent->element); + mjCFrame* child_frame = static_cast(child->element); + *body_parent += std::string(prefix) + *child_frame + std::string(suffix); + return 0; +} + + + // get error message from model const char* mjs_getError(mjSpec* s) { mjCModel* modelC = static_cast(s->element); diff --git a/src/user/user_api.h b/src/user/user_api.h index dd03c959..0682191a 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -40,39 +40,39 @@ typedef struct _mjDoubleVec* mjDoubleVec; //---------------------------------- enum types (mjt) ---------------------------------------------- -typedef enum _mjtGeomInertia { // type of inertia inference - mjINERTIA_VOLUME, // mass distributed in the volume - mjINERTIA_SHELL, // mass distributed on the surface +typedef enum _mjtGeomInertia { // type of inertia inference + mjINERTIA_VOLUME, // mass distributed in the volume + mjINERTIA_SHELL, // mass distributed on the surface } mjtGeomInertia; -typedef enum _mjtBuiltin { // type of built-in procedural texture - mjBUILTIN_NONE = 0, // no built-in texture - mjBUILTIN_GRADIENT, // gradient: rgb1->rgb2 - mjBUILTIN_CHECKER, // checker pattern: rgb1, rgb2 - mjBUILTIN_FLAT // 2d: rgb1; cube: rgb1-up, rgb2-side, rgb3-down +typedef enum _mjtBuiltin { // type of built-in procedural texture + mjBUILTIN_NONE = 0, // no built-in texture + mjBUILTIN_GRADIENT, // gradient: rgb1->rgb2 + mjBUILTIN_CHECKER, // checker pattern: rgb1, rgb2 + mjBUILTIN_FLAT // 2d: rgb1; cube: rgb1-up, rgb2-side, rgb3-down } mjtBuiltin; -typedef enum _mjtMark { // mark type for procedural textures - mjMARK_NONE = 0, // no mark - mjMARK_EDGE, // edges - mjMARK_CROSS, // cross - mjMARK_RANDOM // random dots +typedef enum _mjtMark { // mark type for procedural textures + mjMARK_NONE = 0, // no mark + mjMARK_EDGE, // edges + mjMARK_CROSS, // cross + mjMARK_RANDOM // random dots } mjtMark; -typedef enum _mjtLimited { // type of limit specification - mjLIMITED_FALSE = 0, // not limited - mjLIMITED_TRUE, // limited - mjLIMITED_AUTO, // limited inferred from presence of range +typedef enum _mjtLimited { // type of limit specification + mjLIMITED_FALSE = 0, // not limited + mjLIMITED_TRUE, // limited + mjLIMITED_AUTO, // limited inferred from presence of range } mjtLimited; typedef enum _mjtInertiaFromGeom { - mjINERTIAFROMGEOM_FALSE = 0, // do not use; inertial element required - mjINERTIAFROMGEOM_TRUE, // always use; overwrite inertial element - mjINERTIAFROMGEOM_AUTO // use only if inertial element is missing + mjINERTIAFROMGEOM_FALSE = 0, // do not use; inertial element required + mjINERTIAFROMGEOM_TRUE, // always use; overwrite inertial element + mjINERTIAFROMGEOM_AUTO // use only if inertial element is missing } mjtInertiaFromGeom; @@ -84,8 +84,8 @@ typedef struct _mjElement { // element type, do not modify typedef struct _mjSpec { // model specification - mjElement* element; // object type - mjStatistic stat; // statistics override (if defined) + mjElement* element; // element type + mjString modelname; // model name // compiler settings mjtByte autolimits; // infer "limited" attribute based on range @@ -109,33 +109,33 @@ typedef struct _mjSpec { // model specification mjLROpt LRopt; // options for lengthrange computation // engine data - mjString modelname; // model name - mjOption option; // options - mjVisual visual; // visual options - size_t memory; // size of arena+stack memory in bytes - int nemax; // max number of equality constraints - int njmax; // max number of constraints (Jacobian rows) - int nconmax; // max number of detected contacts (mjContact array size) - size_t nstack; // (deprecated) number of fields in mjData stack - int nuserdata; // number extra fields in mjData - int nuser_body; // number of mjtNums in body_user - int nuser_jnt; // number of mjtNums in jnt_user - int nuser_geom; // number of mjtNums in geom_user - int nuser_site; // number of mjtNums in site_user - int nuser_cam; // number of mjtNums in cam_user - int nuser_tendon; // number of mjtNums in tendon_user - int nuser_actuator; // number of mjtNums in actuator_user - int nuser_sensor; // number of mjtNums in sensor_user + mjOption option; // physics options + mjVisual visual; // visual options + mjStatistic stat; // statistics override (if defined) // sizes - int nkey; // number of keyframes + size_t memory; // number of bytes in arena+stack memory + int nemax; // max number of equality constraints + int nuserdata; // number of mjtNums in userdata + int nuser_body; // number of mjtNums in body_user + int nuser_jnt; // number of mjtNums in jnt_user + int nuser_geom; // number of mjtNums in geom_user + int nuser_site; // number of mjtNums in site_user + int nuser_cam; // number of mjtNums in cam_user + int nuser_tendon; // number of mjtNums in tendon_user + int nuser_actuator; // number of mjtNums in actuator_user + int nuser_sensor; // number of mjtNums in sensor_user + int nkey; // number of keyframes + int njmax; // (deprecated) max number of constraints + int nconmax; // (deprecated) max number of detected contacts + size_t nstack; // (deprecated) number of mjtNums in mjData stack // global data - mjString comment; // comment at top of XML - mjString modelfiledir; // path to model file + mjString comment; // comment at top of XML + mjString modelfiledir; // path to model file // other - bool hasImplicitPluginElem; // already encountered an implicit plugin sensor/actuator + bool hasImplicitPluginElem; // already encountered an implicit plugin sensor/actuator } mjSpec; @@ -148,7 +148,7 @@ typedef struct _mjsOrientation { // alternative orientation specifiers typedef struct _mjsPlugin { // plugin specification - mjElement* instance; // object type + mjElement* instance; // element type mjString name; // name mjString instance_name; // instance name int plugin_slot; // global registered slot number of the plugin @@ -158,7 +158,7 @@ typedef struct _mjsPlugin { // plugin specification typedef struct _mjsBody { // body specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString childclass; // childclass name @@ -186,7 +186,7 @@ typedef struct _mjsBody { // body specification typedef struct _mjsFrame { // frame specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString childclass; // childclass name double pos[3]; // position @@ -197,7 +197,7 @@ typedef struct _mjsFrame { // frame specification typedef struct _mjsJoint { // joint specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjtJoint type; // joint type @@ -238,7 +238,7 @@ typedef struct _mjsJoint { // joint specification typedef struct _mjsGeom { // geom specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // classname mjtGeom type; // geom type @@ -287,7 +287,7 @@ typedef struct _mjsGeom { // geom specification typedef struct _mjsSite { // site specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -311,7 +311,7 @@ typedef struct _mjsSite { // site specification typedef struct _mjsCamera { // camera specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -340,7 +340,7 @@ typedef struct _mjsCamera { // camera specification typedef struct _mjsLight { // light specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -368,7 +368,7 @@ typedef struct _mjsLight { // light specification typedef struct _mjsFlex { - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -409,7 +409,7 @@ typedef struct _mjsFlex { typedef struct _mjsMesh { // mesh specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjString content_type; // content type of file @@ -430,7 +430,7 @@ typedef struct _mjsMesh { // mesh specification typedef struct _mjsHField { // height field specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString content_type; // content type of file mjString file; // file: (nrow, ncol, [elevation data]) @@ -444,7 +444,7 @@ typedef struct _mjsHField { // height field specification typedef struct _mjsSkin { // skin specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjString file; // skin file @@ -471,7 +471,7 @@ typedef struct _mjsSkin { // skin specification typedef struct _mjsTexture { // texture specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjtTexture type; // texture type @@ -505,7 +505,7 @@ typedef struct _mjsTexture { // texture specification typedef struct _mjsMaterial { // material specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjString texture; // name of texture (empty: none) @@ -521,7 +521,7 @@ typedef struct _mjsMaterial { // material specification typedef struct _mjsPair { - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjString geomname1; // name of geom 1 @@ -540,7 +540,7 @@ typedef struct _mjsPair { typedef struct _mjsExclude { - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString bodyname1; // name of geom 1 mjString bodyname2; // name of geom 2 @@ -549,7 +549,7 @@ typedef struct _mjsExclude { typedef struct _mjsEquality { // equality specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name mjtEq type; // constraint type @@ -564,7 +564,7 @@ typedef struct _mjsEquality { // equality specification typedef struct _mjsTendon { // tendon specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -596,13 +596,13 @@ typedef struct _mjsTendon { // tendon specification typedef struct _mjsWrap { // wrapping object specification - mjElement* element; // object type + mjElement* element; // element type mjString info; // message appended to errors } mjsWrap; typedef struct _mjsActuator { // actuator specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -646,7 +646,7 @@ typedef struct _mjsActuator { // actuator specification typedef struct _mjsSensor { // sensor specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString classname; // class name @@ -674,7 +674,7 @@ typedef struct _mjsSensor { // sensor specification typedef struct _mjsNumeric { // custom numeric field specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjDoubleVec data; // initialization data int size; // array size, can be bigger than data size @@ -683,7 +683,7 @@ typedef struct _mjsNumeric { // custom numeric field specification typedef struct _mjsText { // custom text specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjString data; // text string mjString info; // message appended to compiler errors @@ -691,7 +691,7 @@ typedef struct _mjsText { // custom text specification typedef struct _mjsTuple { // tuple specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name mjIntVec objtype; // object types mjStringVec objname; // object names @@ -701,7 +701,7 @@ typedef struct _mjsTuple { // tuple specification typedef struct _mjsKey { // keyframe specification - mjElement* element; // object type + mjElement* element; // element type mjString name; // name double time; // time mjDoubleVec qpos; // qpos @@ -716,7 +716,7 @@ typedef struct _mjsKey { // keyframe specification typedef struct _mjsDefault { // default specification mjString name; // name - mjElement* element; // object type + mjElement* element; // element type mjsJoint* joint; // joint defaults mjsGeom* geom; // geom defaults mjsSite* site; // site defaults @@ -732,34 +732,44 @@ typedef struct _mjsDefault { // default specification } mjsDefault; -//---------------------------------- API functions ------------------------------------------------- +//---------------------------------- Top-level spec manipulation ----------------------------------- -// Create model. +// Create spec. MJAPI mjSpec* mjs_createSpec(); -// Copy model. -MJAPI mjSpec* mjs_copySpec(const mjSpec* s); - -// Copy back model. -MJAPI void mjs_copyBack(mjSpec* s, const mjModel* m); - -// Compile model. +// Compile spec to model. MJAPI mjModel* mjs_compile(mjSpec* s, const mjVFS* vfs); -// Attach child body to a frame of the parent, return 0 if success +// Copy spec. +MJAPI mjSpec* mjs_copySpec(const mjSpec* s); + +// Get compiler error message from spec. +MJAPI const char* mjs_getError(mjSpec* s); + +// Return 1 if compiler error is a warning. +MJAPI int mjs_isWarning(mjSpec* s); + +// Copy model fields back into spec. +MJAPI void mjs_copyBack(mjSpec* s, const mjModel* m); + +// Delete spec. +MJAPI void mjs_deleteSpec(mjSpec* s); + + +//---------------------------------- Attachment ---------------------------------------------------- + +// Attach child body to a parent frame, return 0 on success. MJAPI int mjs_attachBody(mjsFrame* parent, const mjsBody* child, const char* prefix, const char* suffix); -// Get error message from model. -MJAPI const char* mjs_getError(mjSpec* s); +// Attach child frame to a parent body, return 0 if success. +MJAPI int mjs_attachFrame(mjsBody* parent, const mjsFrame* child, + const char* prefix, const char* suffix); -// Return 1 if model has warnings. -MJAPI int mjs_isWarning(mjSpec* s); -// Delete model. -MJAPI void mjs_deleteSpec(mjSpec* s); +//---------------------------------- Add tree elements --------------------------------------------- -// Add child body to body, return child spec. +// Add child body to body, return child. MJAPI mjsBody* mjs_addBody(mjsBody* body, mjsDefault* def); // Add site to body, return site spec. @@ -783,34 +793,28 @@ MJAPI mjsLight* mjs_addLight(mjsBody* body, mjsDefault* def); // Add frame to body. MJAPI mjsFrame* mjs_addFrame(mjsBody* body, mjsFrame* parentframe); -// Add flex to model. + +//---------------------------------- Add non-tree elements ----------------------------------------- + +// Add actuator. +MJAPI mjsActuator* mjs_addActuator(mjSpec* s, mjsDefault* def); + +// Add sensor. +MJAPI mjsSensor* mjs_addSensor(mjSpec* s); + +// Add flex. MJAPI mjsFlex* mjs_addFlex(mjSpec* s); -// Add mesh to model. -MJAPI mjsMesh* mjs_addMesh(mjSpec* s, mjsDefault* def); - -// Add height field to model. -MJAPI mjsHField* mjs_addHField(mjSpec* s); - -// Add skin to model. -MJAPI mjsSkin* mjs_addSkin(mjSpec* s); - -// Add texture to model. -MJAPI mjsTexture* mjs_addTexture(mjSpec* s); - -// Add material to model. -MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def); - -// Add pair to model. +// Add contact pair. MJAPI mjsPair* mjs_addPair(mjSpec* s, mjsDefault* def); -// Add excluded body pair to model. +// Add excluded body pair. MJAPI mjsExclude* mjs_addExclude(mjSpec* s); -// Add equality to model. +// Add equality. MJAPI mjsEquality* mjs_addEquality(mjSpec* s, mjsDefault* def); -// Add tendon to model. +// Add tendon. MJAPI mjsTendon* mjs_addTendon(mjSpec* s, mjsDefault* def); // Wrap site using tendon. @@ -825,42 +829,48 @@ MJAPI mjsWrap* mjs_wrapJoint(mjsTendon* tendon, const char* name, double coef); // Wrap pulley using tendon. MJAPI mjsWrap* mjs_wrapPulley(mjsTendon* tendon, double divisor); -// Add actuator to model. -MJAPI mjsActuator* mjs_addActuator(mjSpec* s, mjsDefault* def); - -// Add sensor to model. -MJAPI mjsSensor* mjs_addSensor(mjSpec* s); - -// Add numeric to model. +// Add numeric. MJAPI mjsNumeric* mjs_addNumeric(mjSpec* s); -// Add text to model. +// Add text. MJAPI mjsText* mjs_addText(mjSpec* s); -// Add tuple to model. +// Add tuple. MJAPI mjsTuple* mjs_addTuple(mjSpec* s); -// Add keyframe to model. +// Add keyframe. MJAPI mjsKey* mjs_addKey(mjSpec* s); -// Add plugin to model. +// Add plugin. MJAPI mjsPlugin* mjs_addPlugin(mjSpec* s); -// Add default to model. +// Add default. MJAPI mjsDefault* mjs_addDefault(mjSpec* s, const char* classname, int parentid, int* id); -// Get model spec from body. + +//---------------------------------- Add assets ---------------------------------------------------- + +// Add mesh. +MJAPI mjsMesh* mjs_addMesh(mjSpec* s, mjsDefault* def); + +// Add height field. +MJAPI mjsHField* mjs_addHField(mjSpec* s); + +// Add skin. +MJAPI mjsSkin* mjs_addSkin(mjSpec* s); + +// Add texture. +MJAPI mjsTexture* mjs_addTexture(mjSpec* s); + +// Add material. +MJAPI mjsMaterial* mjs_addMaterial(mjSpec* s, mjsDefault* def); + + +//---------------------------------- Find/get utilities -------------------------------------------- + +// Get spec from body. MJAPI mjSpec* mjs_getSpec(mjsBody* body); -// Get default corresponding to an mjElement. -MJAPI mjsDefault* mjs_getDefault(mjElement* element); - -// Find default in model by class name. -MJAPI mjsDefault* mjs_findDefault(mjSpec* s, const char* classname); - -// Get global default from model. -MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s); - // Find body in model by name. MJAPI mjsBody* mjs_findBody(mjSpec* s, const char* name); @@ -873,9 +883,21 @@ MJAPI mjsMesh* mjs_findMesh(mjSpec* s, const char* name); // Find frame by name. MJAPI mjsFrame* mjs_findFrame(mjSpec* s, const char* name); +// Get default corresponding to an element. +MJAPI mjsDefault* mjs_getDefault(mjElement* element); + +// Find default in model by class name. +MJAPI mjsDefault* mjs_findDefault(mjSpec* s, const char* classname); + +// Get global default from model. +MJAPI mjsDefault* mjs_getSpecDefault(mjSpec* s); + // Get element id. MJAPI int mjs_getId(mjElement* element); + +//---------------------------------- Attribute setters --------------------------------------------- + // Copy text to string. MJAPI void mjs_setString(mjString dest, const char* text); @@ -903,33 +925,39 @@ MJAPI void mjs_appendFloatVec(mjFloatVecVec dest, const float* array, int size); // Copy double array to vector. MJAPI void mjs_setDouble(mjDoubleVec dest, const double* array, int size); +// Set plugin attributes. +MJAPI void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes); + + +//---------------------------------- Attribute getters --------------------------------------------- + // Get string contents. MJAPI const char* mjs_getString(mjString source); // Get double array contents and optionally its size. MJAPI const double* mjs_getDouble(mjDoubleVec source, int* size); -// Set plugin attributes. -MJAPI void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes); + +//---------------------------------- Other utilities ----------------------------------------------- // Set active plugins. MJAPI void mjs_setActivePlugins(mjSpec* s, void* activeplugins); -// Set default. +// Set element's default. MJAPI void mjs_setDefault(mjElement* element, mjsDefault* def); -// Set frame. +// Set element's enlcosing frame. MJAPI void mjs_setFrame(mjElement* dest, mjsFrame* frame); // Resolve alternative orientations to quat. MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence, const mjsOrientation* orientation); -// Compute quat and inertia from body->fullinertia. +// Compute quat and inertia from fullinertia, return error if any. MJAPI const char* mjs_setFullInertia(mjsBody* body, double quat[4], double inertia[3]); -//---------------------------------- Initialization functions -------------------------------------- +//---------------------------------- Initialization ----------------------------------------------- // Default model attributes. MJAPI void mjs_defaultSpec(mjSpec& model); @@ -1006,7 +1034,8 @@ MJAPI void mjs_defaultKey(mjsKey& key); // Default plugin attributes. MJAPI void mjs_defaultPlugin(mjsPlugin& plugin); -//------------------------- Cache functions ------------------------------------ + +//---------------------------------- Compiler cache ------------------------------------------------ typedef struct _mjCache* mjCache; diff --git a/src/user/user_model.cc b/src/user/user_model.cc index 0a539376..20f81f32 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -3118,7 +3118,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // compile objects in kinematic tree for (int i=0; iCompile(); // also compiles joints, geoms, sites, cameras, lights + bodies[i]->Compile(); // also compiles joints, geoms, sites, cameras, lights, frames } // compile all other objects except for keyframes diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 17421d9a..eb24e2b1 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -628,7 +628,6 @@ void mjCBase::SetFrame(mjCFrame* _frame) { return; } frame = _frame; - frame->Compile(); } @@ -729,13 +728,67 @@ mjCBody& mjCBody::operator+=(const mjCBody& other) { +// attach frame to body +mjCBody& mjCBody::operator+=(const mjCFrame& other) { + mjCBody* subtree = other.body; + other.model->prefix = other.prefix; + other.model->suffix = other.suffix; + + // copy input frame + frames.push_back(new mjCFrame(other)); + frames.back()->body = this; + frames.back()->model = model; + frames.back()->frame = other.frame; + int i = frames.size(); + + // map input frames to index in this->frames + std::map fmap; + for (auto frame : subtree->frames) { + if (frame == static_cast(&other)) { + fmap[frame] = frames.size() - 1; + } else if (other.IsAncestor(frame)) { + fmap[frame] = i++; + } + } + + // copy children that are inside the input frame + CopyList(frames, subtree->frames, fmap, &other); // needs to be done first + CopyList(geoms, subtree->geoms, fmap, &other); + CopyList(joints, subtree->joints, fmap, &other); + CopyList(sites, subtree->sites, fmap, &other); + CopyList(cameras, subtree->cameras, fmap, &other); + CopyList(lights, subtree->lights, fmap, &other); + + for (int i=0; ibodies.size(); i++) { + if (!other.IsAncestor(subtree->bodies[i]->frame)) { + continue; + } + bodies.push_back(new mjCBody(*subtree->bodies[i], model)); // triggers recursive call + bodies.back()->frame = + subtree->bodies[i]->frame ? frames[fmap[subtree->bodies[i]->frame]] : nullptr; + } + + // name space + this->NameSpace(other.model); + + // attach referencing elements + *model += *other.model; + + // clear namespace and return body + other.model->prefix.clear(); + other.model->suffix.clear(); + return *this; +} + + + // copy src list of elements into dst; set body, model and frame template void mjCBody::CopyList(std::vector& dst, const std::vector& src, std::map& fmap, const mjCFrame* pframe) { int nsrc = (int)src.size(); for (int i=0; iframe != pframe) { + if (pframe && !pframe->IsAncestor(src[i]->frame)) { continue; // skip if the element is not inside pframe } dst.push_back(new T(*src[i])); @@ -1169,6 +1222,11 @@ void mjCBody::MakeInertialExplicit() { void mjCBody::Compile(void) { CopyFromSpec(); + // compile all frames + for (int i=0; iCompile(); + } + // resize userdata if (userdata_.size() > model->nuser_body) { throw mjCError(this, "user has more values than nuser_body in body '%s' (id = %d)", @@ -1417,6 +1475,21 @@ mjCFrame& mjCFrame::operator+=(const mjCBody& other) { +// return true if child is descendent of this frame +bool mjCFrame::IsAncestor(const mjCFrame* child) const { + if (!child) { + return false; + } + + if (child == this) { + return true; + } + + return IsAncestor(child->frame); +} + + + void mjCFrame::SetParent(mjCBody* _body) { body = _body; } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index db4ea58f..f6e91a81 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -275,6 +275,7 @@ class mjCBody : public mjCBody_, private mjsBody { // API for adding existing objects to body mjCBody& operator+=(const mjCBody& other); + mjCBody& operator+=(const mjCFrame& other); // API for accessing objects int NumObjects(mjtObj type); @@ -360,7 +361,8 @@ class mjCFrame : public mjCFrame_, private mjsFrame { void SetParent(mjCBody* _body); mjCFrame& operator+=(const mjCBody& other); - mjCFrame& operator+=(const mjCFrame& other); + + bool IsAncestor(const mjCFrame* child) const; // true if child is contained in this frame private: mjCFrame(mjCModel* = 0, mjCFrame* = 0); // constructor diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 4b01e1a8..3e712310 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -196,20 +196,20 @@ TEST_F(PluginTest, RecompileCompareCache) { } // -------------------------------- test attach ------------------------------- -TEST_F(MujocoTest, AttachSame) { - std::array er; - mjtNum tol = 0; - std::string field = ""; - - static constexpr char xml[] = R"( +static constexpr char xml_child[] = R"( - - - - - - + + + + + + + + + + + @@ -227,26 +227,33 @@ TEST_F(MujocoTest, AttachSame) { - + )"; +TEST_F(MujocoTest, AttachSame) { + std::array er; + mjtNum tol = 0; + std::string field = ""; + static constexpr char xml_result[] = R"( - - + + + - - + + + @@ -270,13 +277,13 @@ TEST_F(MujocoTest, AttachSame) { - - + + )"; // create parent - mjSpec* parent = ParseSpecFromString(xml, er.data(), er.size()); + mjSpec* parent = ParseSpecFromString(xml_child, er.data(), er.size()); EXPECT_THAT(parent, NotNull()) << er.data(); // get frame @@ -296,10 +303,10 @@ TEST_F(MujocoTest, AttachSame) { EXPECT_THAT(m_attached, NotNull()); // check full name stored in mjModel - EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 4), "attached-body-1"); + EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 5), "attached-body-1"); // check body 3 is attached to the world - EXPECT_THAT(m_attached->body_parentid[3], 0); + EXPECT_THAT(m_attached->body_parentid[4], 0); // compare with expected XML mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size()); @@ -330,32 +337,6 @@ TEST_F(MujocoTest, AttachDifferent) { )"; - static constexpr char xml_child[] = R"( - - - - - - - - - - - - - - - - - - - - - - - - )"; - static constexpr char xml_result[] = R"( @@ -366,14 +347,15 @@ TEST_F(MujocoTest, AttachDifferent) { - + + - + @@ -381,9 +363,12 @@ TEST_F(MujocoTest, AttachDifferent) { - - + + + + + )"; // model with one free sphere and a frame @@ -403,7 +388,7 @@ TEST_F(MujocoTest, AttachDifferent) { EXPECT_THAT(body, NotNull()); // attach child to parent frame - EXPECT_THAT( + EXPECT_EQ( mjs_attachBody(frame, body, /*prefix=*/"attached-", /*suffix=*/"-1"), 0); // compile new model @@ -430,5 +415,102 @@ TEST_F(MujocoTest, AttachDifferent) { mj_deleteModel(m_expected); } +TEST_F(MujocoTest, AttachFrame) { + std::array er; + mjtNum tol = 0; + std::string field = ""; + + static constexpr char xml_parent[] = R"( + + + + + + + + + )"; + + static constexpr char xml_result[] = R"( + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + )"; + + // model with one free sphere and a frame + mjSpec* parent = ParseSpecFromString(xml_parent, er.data(), er.size()); + EXPECT_THAT(parent, NotNull()) << er.data(); + + // get frame + mjsBody* body = mjs_findBody(parent, "sphere"); + EXPECT_THAT(body, NotNull()); + + // model with one cylinder and a hinge + mjSpec* child = ParseSpecFromString(xml_child, er.data(), er.size()); + EXPECT_THAT(child, NotNull()) << er.data(); + + // get subtree + mjsFrame* frame = mjs_findFrame(child, "pframe"); + EXPECT_THAT(frame, NotNull()); + + // attach child to parent frame + EXPECT_THAT( + mjs_attachFrame(body, frame, /*prefix=*/"attached-", /*suffix=*/"-1"), 0); + + // compile new model + mjModel* m_attached = mjs_compile(parent, 0); + EXPECT_THAT(m_attached, NotNull()); + + // check full name stored in mjModel + EXPECT_STREQ(mj_id2name(m_attached, mjOBJ_BODY, 2), "attached-body-1"); + + // check body 2 is attached to body 1 + EXPECT_THAT(m_attached->body_parentid[2], 1); + + // compare with expected XML + mjModel* m_expected = LoadModelFromString(xml_result, er.data(), er.size()); + EXPECT_THAT(m_expected, NotNull()) << er.data(); + EXPECT_LE(CompareModel(m_attached, m_expected, field), tol) + << "Expected and attached models are different!\n" + << "Different field: " << field << '\n';; + + // destroy everything + mjs_deleteSpec(parent); + mjs_deleteSpec(child); + mj_deleteModel(m_attached); + mj_deleteModel(m_expected); +} + } // namespace } // namespace mujoco diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc index 32805357..e86e3254 100644 --- a/test/xml/xml_native_writer_test.cc +++ b/test/xml/xml_native_writer_test.cc @@ -1280,8 +1280,39 @@ TEST_F(XMLWriterTest, WriteReadCompare) { EXPECT_EQ(d->pstack, 0) << "mjData stack memory leak detected in " << p.path().string() << '\n'; - // delete original structures + // delete data mj_deleteData(d); + + // allocate buffer, save m into it + size_t sz = mj_sizeModel(m); + void* buffer = mju_malloc(sz); + mj_saveModel(m, nullptr, buffer, sz); + + // make new VFS add buffer to it + mjVFS* vfs = (mjVFS*)mju_malloc(sizeof(mjVFS)); + mj_defaultVFS(vfs); + int failed = mj_addBufferVFS(vfs, "model.mjb", buffer, sz); + EXPECT_EQ(failed, 0) << "Failed to add buffer to VFS"; + + // load model from VFS + mtemp = mj_loadModel("model.mjb", vfs); + ASSERT_THAT(mtemp, NotNull()); + + // compare with 0 tolerance + std::string field = ""; + mjtNum result = CompareModel(m, mtemp, field); + EXPECT_EQ(result, 0) + << "Loaded and saved binary models are different!\n" + << "Affected file " << p.path().string() << '\n' + << "Different field: " << field << '\n'; + + // clean up + mj_deleteModel(mtemp); + mj_deleteVFS(vfs); + mju_free(vfs); + mju_free(buffer); + + // delete model mj_deleteModel(m); } }