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

Merge branch 'google-deepmind:main' into main

Browse Source
This commit is contained in:
Andrew Luo
2024-04-21 17:15:31 +02:00
committed by GitHub
parent ab7d8fe167 34a61ae42f
commit f3a9d7790a
35 changed files with 1346 additions and 1647 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+21 -3
View File
@@ -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)
doc/models.rst
+14
View File
@@ -129,3 +129,17 @@ Drones
- Preview
* - `Skydio X2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/skydio_x2>`_
- .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/skydio_x2/x2.png
* - `Bitcraze Crazyflie 2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/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 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/flybody>`_
- .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/flybody/flybody.png
mjx/mujoco/mjx/__init__.py
-5
View File
@@ -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
mjx/mujoco/mjx/_src/collision_base.py
-67
View File
@@ -1,67 +0,0 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Collision base."""
import dataclasses
from typing import Dict, List, Optional, Tuple
import jax
# pylint: disable=g-importing-member
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import GeomType
# pylint: enable=g-importing-member
Contact = Tuple[jax.Array, jax.Array, jax.Array]
@dataclasses.dataclass
class Candidate:
geom1: int
geom2: int
ipair: int
geomp: int # priority geom
dim: int
CandidateSet = Dict[
Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]],
List[Candidate],
]
class GeomInfo(PyTreeNode):
"""Collision info for a geom."""
geom_id: jax.Array
pos: jax.Array
mat: jax.Array
size: jax.Array
face: Optional[jax.Array] = None
vert: Optional[jax.Array] = None
edge_dir: Optional[jax.Array] = None
facenorm: Optional[jax.Array] = None
edge: Optional[jax.Array] = None
edge_face_normal: Optional[jax.Array] = None
class SolverParams(PyTreeNode):
"""Contact solver params."""
friction: jax.Array
solref: jax.Array
solreffriction: jax.Array
solimp: jax.Array
margin: jax.Array
gap: jax.Array
mjx/mujoco/mjx/_src/collision_convex.py
+60 -26
View File
@@ -20,12 +20,53 @@ from typing import Tuple
import jax
from jax import numpy as jp
from mujoco.mjx._src import math
from mujoco.mjx._src import mesh
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Contact
from mujoco.mjx._src.collision_base import GeomInfo
from mujoco.mjx._src.collision_types import Collision
from mujoco.mjx._src.collision_types import ConvexInfo
from mujoco.mjx._src.collision_types import FunctionKey
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
def collider(ncon: int):
"""Wraps collision functions for use by collision_driver."""
def wrapper(func):
def collide(
m: Model, d: Data, key: FunctionKey, geom: jax.Array
) -> Collision:
g1, g2 = geom.T
infos = [
GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]),
]
in_axes = [0, 0]
for i in [0, 1]:
if key.types[i] == GeomType.BOX:
infos[i] = mesh.box(infos[i])
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, face=0, vert=0
)
elif key.types[i] == GeomType.MESH:
infos[i] = mesh.convex(m, key.data_ids[i], infos[i])
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0
)
dist, pos, frame = jax.vmap(func, in_axes=in_axes)(*infos)
if ncon > 1:
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
return collide
return wrapper
def _closest_segment_point_plane(
a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array
) -> jax.Array:
@@ -178,7 +219,8 @@ def _manifold_points(
return jp.array([a_idx, b_idx, c_idx, d_idx])
def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
@collider(ncon=4)
def plane_convex(plane: GeomInfo, convex: ConvexInfo) -> Collision:
"""Calculates contacts between a plane and a convex object."""
vert = convex.vert
@@ -200,10 +242,11 @@ def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
return dist, pos, frame
def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
@collider(ncon=1)
def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision:
"""Calculates contact between a sphere and a convex object."""
faces = convex.face
normals = convex.facenorm
normals = convex.face_normal
# Put sphere in convex frame.
sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)
@@ -262,16 +305,15 @@ def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
n = convex.mat @ n
pos = convex.mat @ pos + convex.pos
return jax.tree_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
mjx/mujoco/mjx/_src/collision_driver.py
+274 -332
View File
@@ -12,21 +12,39 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Collide geometries."""
"""Runs collision checking for all geoms in a Model.
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
To do this, collision_driver builds a collision function table, and then runs
the collision functions serially on the parameters in the table.
For example, if a Model has three geoms:
geom | type
---------------
1 | sphere
2 | capsule
3 | sphere
collision_driver organizes it into these functions and runs them:
function | geom pair
--------------------------
sphere_sphere | (1, 3)
sphere_capsule | (1, 2), (2, 3)
Besides collision function, function tables are keyed on mesh id and condim,
in order to guarantee static shapes for contacts and jacobians.
"""
import itertools
from typing import Dict, Iterator, List, Tuple, Union
import jax
from jax import numpy as jp
import mujoco
from mujoco.mjx._src import collision_base
from mujoco.mjx._src import mesh
from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Candidate
from mujoco.mjx._src.collision_base import CandidateSet
from mujoco.mjx._src.collision_base import GeomInfo
from mujoco.mjx._src.collision_base import SolverParams
from mujoco.mjx._src.collision_convex import capsule_convex
from mujoco.mjx._src.collision_convex import convex_convex
from mujoco.mjx._src.collision_convex import plane_convex
@@ -40,12 +58,14 @@ from mujoco.mjx._src.collision_primitive import sphere_capsule
from mujoco.mjx._src.collision_primitive import sphere_sphere
from mujoco.mjx._src.collision_sdf import capsule_ellipsoid
from mujoco.mjx._src.collision_sdf import ellipsoid_ellipsoid
from mujoco.mjx._src.collision_types import FunctionKey
from mujoco.mjx._src.types import Contact
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
import numpy as np
# pair-wise collision functions
_COLLISION_FUNC = {
@@ -70,371 +90,293 @@ _COLLISION_FUNC = {
}
def get_collision_fn(
key: Tuple[Union[GeomType, mujoco.mjtGeom], Union[GeomType, mujoco.mjtGeom]]
) -> Optional[Callable[[GeomInfo, GeomInfo], collision_base.Contact]]:
"""Returns a collision function given a pair of geom types."""
return _COLLISION_FUNC.get(key, None)
# geoms for which we ignore broadphase
_GEOM_NO_BROADPHASE = {GeomType.HFIELD, GeomType.PLANE}
def _add_candidate(
result: CandidateSet,
def has_collision_fn(t1: GeomType, t2: GeomType) -> bool:
"""Returns True if a collision function exists for a pair of geom types."""
return (t1, t2) in _COLLISION_FUNC
def geom_pairs(
m: Union[Model, mujoco.MjModel],
g1: int,
g2: int,
ipair: int = -1,
):
"""Adds a candidate to test for collision."""
t1, t2 = m.geom_type[g1], m.geom_type[g2]
if t1 > t2:
t1, t2, g1, g2 = t2, t1, g2, g1
) -> Iterator[Tuple[int, int, int]]:
"""Returns geom pairs to check for collisions.
# MuJoCo does not collide planes with other planes or hfields
if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
return
if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
return
Args:
m: a MuJoCo or MJX model
def mesh_key(i):
convex_data = [[None] * m.ngeom] * 3
if isinstance(m, Model):
convex_data = [
m.geom_convex_face,
m.geom_convex_vert,
m.geom_convex_edge_dir,
]
elif isinstance(m, mujoco.MjModel):
kwargs = mesh.get(m)
convex_data = [
kwargs['geom_convex_face'],
kwargs['geom_convex_vert'],
kwargs['geom_convex_edge_dir'],
]
key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data)
return key
Yields:
geom1, geom2, and pair index if defined in <pair> (else -1)
"""
pairs = set()
k1, k2 = mesh_key(g1), mesh_key(g2)
for i in range(m.npair):
g1, g2 = m.pair_geom1[i], m.pair_geom2[i]
# order pairs by geom_type for correct function mapping
if m.geom_type[g1] > m.geom_type[g2]:
g1, g2 = g2, g1
pairs.add((g1, g2))
yield g1, g2, i
candidates = {(c.geom1, c.geom2) for c in result.get((t1, t2, k1, k2), [])}
if (g1, g2) in candidates:
return
if ipair > -1:
candidate = Candidate(g1, g2, ipair, -1, m.pair_dim[ipair])
elif m.geom_priority[g1] != m.geom_priority[g2]:
gp = g1 if m.geom_priority[g1] > m.geom_priority[g2] else g2
candidate = Candidate(g1, g2, -1, gp, m.geom_condim[gp])
else:
dim = max(m.geom_condim[g1], m.geom_condim[g2])
candidate = Candidate(g1, g2, -1, -1, dim)
result.setdefault((t1, t2, k1, k2), []).append(candidate)
def _pair_params(
m: Model,
candidates: Sequence[Candidate],
) -> SolverParams:
"""Gets solver params for pair geoms."""
ipair = jp.array([c.ipair for c in candidates])
friction = jp.clip(m.pair_friction[ipair], a_min=mujoco.mjMINMU)
solref = m.pair_solref[ipair]
solreffriction = m.pair_solreffriction[ipair]
solimp = m.pair_solimp[ipair]
margin = m.pair_margin[ipair]
gap = m.pair_gap[ipair]
return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
def _priority_params(
m: Model,
candidates: Sequence[Candidate],
) -> SolverParams:
"""Gets solver params from priority geoms."""
geomp = jp.array([c.geomp for c in candidates])
friction = m.geom_friction[geomp][:, jp.array([0, 0, 1, 2, 2])]
solref = m.geom_solref[geomp]
solreffriction = jp.zeros(geomp.shape + (mujoco.mjNREF,))
solimp = m.geom_solimp[geomp]
g = jp.array([(c.geom1, c.geom2) for c in candidates])
margin = jp.amax(m.geom_margin[g.T], axis=0)
gap = jp.amax(m.geom_gap[g.T], axis=0)
return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
def _dynamic_params(
m: Model,
candidates: Sequence[Candidate],
) -> SolverParams:
"""Gets solver params for dynamic geoms."""
g1 = jp.array([c.geom1 for c in candidates])
g2 = jp.array([c.geom2 for c in candidates])
friction = jp.maximum(m.geom_friction[g1], m.geom_friction[g2])
# copy friction terms for the full geom pair
friction = friction[:, jp.array([0, 0, 1, 2, 2])]
minval = jp.array(mujoco.mjMINVAL)
solmix1, solmix2 = m.geom_solmix[g1], m.geom_solmix[g2]
mix = solmix1 / (solmix1 + solmix2)
mix = jp.where((solmix1 < minval) & (solmix2 < minval), 0.5, mix)
mix = jp.where((solmix1 < minval) & (solmix2 >= minval), 0.0, mix)
mix_fn = jax.vmap(lambda a, b, m: m * a + (1 - m) * b)
solref1, solref2 = m.geom_solref[g1], m.geom_solref[g2]
solref = jp.minimum(solref1, solref2)
s_mix = mix_fn(solref1, solref2, mix)
solref = jp.where((solref1[0] > 0) & (solref2[0] > 0), s_mix, solref)
solreffriction = jp.zeros(g1.shape + (mujoco.mjNREF,))
solimp = mix_fn(m.geom_solimp[g1], m.geom_solimp[g2], mix)
margin = jp.maximum(m.geom_margin[g1], m.geom_margin[g2])
gap = jp.maximum(m.geom_gap[g1], m.geom_gap[g2])
return SolverParams(friction, solref, solreffriction, solimp, margin, gap)
def get_params(
m: Union[Model, mujoco.MjModel], candidates: Sequence[Candidate]
) -> Tuple[List[int], List[int], SolverParams]:
"""Gets solver params for a list of collision candidates."""
# group sol params by different candidate types
typ_cands = {}
for c in candidates:
typ = (c.ipair > -1, c.geomp > -1)
typ_cands.setdefault(typ, []).append(c)
geom1, geom2, params = [], [], []
for (pair, priority), candidates in typ_cands.items():
geom1.extend([c.geom1 for c in candidates])
geom2.extend([c.geom2 for c in candidates])
if pair:
params.append(_pair_params(m, candidates))
elif priority:
params.append(_priority_params(m, candidates))
else:
params.append(_dynamic_params(m, candidates))
params = jax.tree_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)
mjx/mujoco/mjx/_src/collision_driver_test.py
+22 -16
View File
@@ -402,7 +402,8 @@ class CapsuleCollisionTest(parameterized.TestCase):
self.assertEqual(c.pos.shape[0], 2)
self.assertGreater(c.dist[1], 0)
# extract the contact point with penetration
c = jax.tree_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)
mjx/mujoco/mjx/_src/collision_primitive.py
+51 -48
View File
@@ -20,34 +20,53 @@ import jax
from jax import numpy as jp
from mujoco.mjx._src import math
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Contact
from mujoco.mjx._src.collision_base import GeomInfo
from mujoco.mjx._src.collision_types import Collision
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
def collider(ncon: int):
"""Wraps collision functions for use by collision_driver."""
def wrapper(func):
def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
g1, g2 = geom.T
info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1])
info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2])
dist, pos, frame = jax.vmap(func)(info1, info2)
if ncon > 1:
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
return collide
return wrapper
def _plane_sphere(
plane_normal: jax.Array,
plane_pos: jax.Array,
sphere_pos: jax.Array,
radius: jax.Array,
sphere_radius: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
"""Returns the penetration and contact point between a plane and sphere."""
cdist = jp.dot(sphere_pos - plane_pos, plane_normal)
dist = cdist - radius
pos = sphere_pos - plane_normal * (radius + 0.5 * dist)
"""Returns the distance and contact point between a plane and sphere."""
dist = jp.dot(sphere_pos - plane_pos, plane_normal) - sphere_radius
pos = sphere_pos - plane_normal * (sphere_radius + 0.5 * dist)
return dist, pos
def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact:
@collider(ncon=1)
def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Collision:
"""Calculates contact between a plane and a sphere."""
n = plane.mat[:, 2]
dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0])
return jax.tree_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)
mjx/mujoco/mjx/_src/collision_sdf.py
+28 -15
View File
@@ -28,15 +28,35 @@ import jax
from jax import numpy as jp
from mujoco.mjx._src import math
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Contact
from mujoco.mjx._src.collision_base import GeomInfo
from mujoco.mjx._src.collision_types import Collision
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
# the SDF function takes position in, and returns a distance or objective
SDFFn = Callable[[jax.Array], jax.Array]
def collider(ncon: int):
"""Wraps collision functions for use by collision_driver."""
def wrapper(func):
def collide(m: Model, d: Data, _, geom: jax.Array) -> Collision:
g1, g2 = geom.T
info1 = GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1])
info2 = GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2])
dist, pos, frame = jax.vmap(func)(info1, info2)
if ncon > 1:
return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame))
return dist, pos, frame
collide.ncon = ncon
return collide
return wrapper
def _plane(pos: jax.Array, size: jax.Array) -> jax.Array:
del size
return pos[2]
@@ -128,22 +148,15 @@ def _optim(
return pos, dist, n
def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Contact:
@collider(ncon=1)
def capsule_ellipsoid(c: GeomInfo, e: GeomInfo) -> Collision:
""""Calculates contact between a capsule and an ellipsoid."""
pos, dist, n = _optim(_capsule, _ellipsoid, c, e)
return jax.tree_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)
mjx/mujoco/mjx/_src/collision_types.py
+65
View File
@@ -0,0 +1,65 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Collision base types."""
import dataclasses
from typing import Tuple
import jax
# pylint: disable=g-importing-member
from mujoco.mjx._src.dataclasses import PyTreeNode
# pylint: enable=g-importing-member
class GeomInfo(PyTreeNode):
"""Geom propertes of primitive and SDF shapes."""
pos: jax.Array
mat: jax.Array
size: jax.Array
class ConvexInfo(PyTreeNode):
"""Geom propertes of convex meshes."""
pos: jax.Array
mat: jax.Array
vert: jax.Array
face: jax.Array
face_normal: jax.Array
edge: jax.Array
edge_face_normal: jax.Array
edge_dir: jax.Array
# Collision returned by collision functions:
# - distance distance between nearest points; neg: penetration
# - position (3,) position of contact point: midpoint between geoms
# - frame (3, 3) normal is in [0, :], points from geom[0] to geom[1]
Collision = Tuple[jax.Array, jax.Array, jax.Array]
@dataclasses.dataclass(frozen=True)
class FunctionKey:
"""Specifies how geom pairs group into collision_driver's function table.
Attributes:
types: geom type pair, which determines the collision function
data_ids: geom data id pair: mesh id for mesh geoms, otherwise -1.
Meshes have distinct face/vertex counts, so must occupy distinct
entries in the collision function table.
condim: grouping by condim of the colliision ensures that the size of the
resulting constraint jacobian is determined at compile time.
"""
types: Tuple[int, int]
data_ids: Tuple[int, int]
condim: int
mjx/mujoco/mjx/_src/constraint.py
+49 -30
View File
@@ -24,6 +24,7 @@ from mujoco.mjx._src import math
from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import ConstraintType
from mujoco.mjx._src.types import Contact
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
@@ -34,6 +35,9 @@ from mujoco.mjx._src.types import Model
import numpy as np
_CONDIM_EFC_COUNT = {1: 1, 3: 4, 4: 6, 6: 10}
class _Efc(PyTreeNode):
"""Support data for creating constraint matrices."""
J: jax.Array
@@ -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):
mjx/mujoco/mjx/_src/constraint_test.py
+4 -4
View File
@@ -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)
mjx/mujoco/mjx/_src/device.py
-319
View File
@@ -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
mjx/mujoco/mjx/_src/device_test.py
-203
View File
@@ -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(
'<mujoco><option solver="PGS"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
def test_integrator(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
def test_cone(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
)
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("""
<mujoco>
<worldbody>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="box"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1" type="cylinder"/>
</body>
</worldbody>
</mujoco>
""")
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
def test_tendon(self):
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<body name="left_thigh" pos="0 0.1 -0.04">
<joint axis="0 1 0" name="left_hip_y" type="hinge"/>
<geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
<body name="left_shin" pos="0 -0.01 -0.403">
<joint axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
<geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
</body>
</body>
</worldbody>
<tendon>
<fixed name="left_hipknee">
<joint coef="-1" joint="left_hip_y"/>
<joint coef="1" joint="left_knee"/>
</fixed>
</tendon>
</mujoco>
""")
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/forward.py
+2 -2
View File
@@ -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
mjx/mujoco/mjx/_src/io.py
+107 -129
View File
@@ -22,15 +22,14 @@ from jax import numpy as jp
import mujoco
from mujoco.mjx._src import collision_driver
from mujoco.mjx._src import constraint
from mujoco.mjx._src import mesh
from mujoco.mjx._src import support
from mujoco.mjx._src import types
import numpy as np
import scipy
def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
"""Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
def _make_option(o: mujoco.MjOption) -> types.Option:
"""Returns mjx.Option given mujoco.MjOption."""
if o.integrator not in set(types.IntegratorType):
raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
@@ -47,38 +46,20 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
if o.enableflags & 2**i:
raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
static_fields = {
f.name: copy.copy(getattr(o, f.name))
for f in types.Option.fields()
if f.type in (int, bytes, np.ndarray)
}
static_fields['integrator'] = types.IntegratorType(o.integrator)
static_fields['cone'] = types.ConeType(o.cone)
static_fields['jacobian'] = types.JacobianType(o.jacobian)
static_fields['solver'] = types.SolverType(o.solver)
static_fields['disableflags'] = types.DisableBit(o.disableflags)
fields = {f.name: getattr(o, f.name, None) for f in types.Option.fields()}
fields['integrator'] = types.IntegratorType(o.integrator)
fields['cone'] = types.ConeType(o.cone)
fields['jacobian'] = types.JacobianType(o.jacobian)
fields['solver'] = types.SolverType(o.solver)
fields['disableflags'] = types.DisableBit(o.disableflags)
fields['has_fluid_params'] = o.density > 0 or o.viscosity > 0 or o.wind.any()
device_fields = {
f.name: copy.copy(getattr(o, f.name))
for f in types.Option.fields()
if f.type is jax.Array
}
device_fields = jax.device_put(device_fields, device=device)
has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
return types.Option(
has_fluid_params=has_fluid_params,
**static_fields,
**device_fields,
)
return types.Option(**fields)
def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
def _make_statistic(s: mujoco.MjStatistic) -> types.Statistic:
"""Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
return types.Statistic(
meaninertia=jax.device_put(s.meaninertia, device=device)
)
return types.Statistic(meaninertia=s.meaninertia)
def put_model(m: mujoco.MjModel, device=None) -> types.Model:
@@ -93,17 +74,21 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
if m.body_gravcomp.any():
raise NotImplementedError('gravcomp is not supported')
# check collision geom types
for (g1, g2, *_), c in collision_driver.collision_candidates(m).items():
g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
if collision_driver.get_collision_fn((g1, g2)) is None:
raise NotImplementedError(f'({g1}, {g2}) has no collision function')
*_, params = collision_driver.get_params(m, c)
margin_gap = not np.allclose(np.concatenate([params.margin, params.gap]), 0)
if mujoco.mjtGeom.mjGEOM_MESH in (g1, g2) and margin_gap:
raise NotImplementedError(
f'Margin and gap not implemented for ({g1}, {g2})'
)
for g1, g2, ip in collision_driver.geom_pairs(m):
t1, t2 = m.geom_type[[g1, g2]]
# check collision function exists for type pair
if not collision_driver.has_collision_fn(t1, t2):
t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) collisions not implemented.')
# margin/gap not supported for geoms
if mujoco.mjtGeom.mjGEOM_MESH in (t1, t2):
if ip != -1:
margin = m.pair_margin[ip]
else:
margin = m.geom_margin[g1] + m.geom_margin[g2]
if margin.any():
t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2)
raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.')
for enum_field, enum_type, mj_type in (
(m.actuator_biastype, types.BiasType, mujoco.mjtBias),
@@ -121,40 +106,24 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
if not np.allclose(m.dof_frictionloss, 0):
raise NotImplementedError('dof_frictionloss is not implemented.')
opt = _put_option(m.opt, device=device)
stat = _put_statistic(m.stat, device=device)
fields = {f.name: getattr(m, f.name) for f in types.Model.fields()}
fields['geom_rgba'] = fields['geom_rgba'].reshape((-1, 4))
fields['mat_rgba'] = fields['mat_rgba'].reshape((-1, 4))
fields['cam_mat0'] = fields['cam_mat0'].reshape((-1, 3, 3))
fields['opt'] = _make_option(m.opt)
fields['stat'] = _make_statistic(m.stat)
model = types.Model(**{k: copy.copy(v) for k, v in fields.items()})
static_fields = {
f.name: getattr(m, f.name)
for f in types.Model.fields()
if f.type in (int, bytes, np.ndarray)
}
static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4))
static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4))
device_fields = {
f.name: copy.copy(getattr(m, f.name)) # copy because device_put is async
for f in types.Model.fields()
if f.type is jax.Array
}
device_fields['cam_mat0'] = device_fields['cam_mat0'].reshape((-1, 3, 3))
device_fields.update(mesh.get(m))
device_fields = jax.device_put(device_fields, device=device)
return types.Model(
opt=opt,
stat=stat,
**static_fields,
**device_fields,
)
return jax.device_put(model, device=device)
def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
"""Allocate and initialize Data."""
ncon = collision_driver.ncon(m)
ne, nf, nl, nc = constraint.count_constraints(m)
nefc = ne + nf + nl + nc
dim = collision_driver.make_condim(m)
efc_type = constraint.make_efc_type(m, dim)
efc_address = constraint.make_efc_address(efc_type, dim)
ne, nf, nl, nc = constraint.counts(efc_type)
ncon, nefc = dim.size, ne + nf + nl + nc
zero_0 = jp.zeros(0, dtype=float)
zero_nv = jp.zeros(m.nv, dtype=float)
@@ -170,8 +139,28 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float)
zero_nm = jp.zeros(m.nM, dtype=float)
# create first d to get num contacts and nc
contact = types.Contact(
dist=jp.zeros(ncon),
pos=jp.zeros((ncon, 3)),
frame=jp.zeros((ncon, 3, 3)),
includemargin=jp.zeros(ncon),
friction=jp.zeros((ncon, 5)),
solref=jp.zeros((ncon, mujoco.mjNREF)),
solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
solimp=jp.zeros((ncon, mujoco.mjNIMP)),
dim=dim,
geom1=jp.zeros(ncon, dtype=int) - 1,
geom2=jp.zeros(ncon, dtype=int) - 1,
geom=jp.zeros((ncon, 2), dtype=int) - 1,
efc_address=efc_address,
)
d = types.Data(
ne=ne,
nf=nf,
nl=nl,
nefc=nefc,
ncon=ncon,
solver_niter=jp.array(0, dtype=int),
time=jp.array(0.0),
qpos=jp.array(m.qpos0),
@@ -206,7 +195,8 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
contact=types.Contact.zero(ncon),
contact=contact,
efc_type=efc_type,
efc_J=jp.zeros((nefc, m.nv), dtype=float),
efc_frictionloss=zero_nefc,
efc_D=zero_nefc,
@@ -228,11 +218,7 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
return d
def _get_contact(
c: mujoco._structs._MjContactList,
cx: types.Contact,
efc_start: int,
):
def _get_contact(c: mujoco._structs._MjContactList, cx: types.Contact):
"""Converts mjx.Contact to mujoco._structs._MjContactList."""
con_id = np.nonzero(cx.dist <= 0)[0]
for field in types.Contact.fields():
@@ -241,10 +227,6 @@ def _get_contact(
value = value.reshape((-1, 9))
getattr(c, field.name)[:] = value
ncon = cx.dist.shape[0]
c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
c.dim[:] = 3
def get_data(
m: mujoco.MjModel, d: types.Data
@@ -278,13 +260,6 @@ def get_data_into(
d = jax.device_get(d)
batch_size = d.qpos.shape[0] if batched else 1
ne, nf, nl, nc = constraint.count_constraints(m, d)
efc_type = np.array([
mujoco.mjtConstraint.mjCNSTR_EQUALITY,
mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT,
mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
]).repeat([ne, nf, nl, nc])
dof_i, dof_j = [], []
for i in range(m.nv):
@@ -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)
mjx/mujoco/mjx/_src/io_test.py
+2 -7
View File
@@ -108,10 +108,6 @@ class ModelIOTest(parameterized.TestCase):
np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref)
np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
self.assertLen(mx.geom_convex_face, 6)
self.assertLen(mx.geom_convex_vert, 6)
self.assertLen(mx.geom_convex_edge_dir, 6)
self.assertLen(mx.geom_convex_facenormal, 6)
np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
@@ -256,6 +252,7 @@ class DataIOTest(parameterized.TestCase):
nv = 19
nefc = 185
self.assertEqual(d.nefc, nefc)
self.assertEqual(d.qpos.shape, (nq,))
self.assertEqual(d.qvel.shape, (nv,))
self.assertEqual(d.act.shape, (0,))
@@ -423,8 +420,6 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(d_2.efc_J, d.efc_J)
self.assertEqual(d_2.efc_aref.shape, (8,)) # nefc
np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
# efc_address is created on demand
np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
def test_get_data_batched(self):
@@ -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)
mjx/mujoco/mjx/_src/mesh.py
+102 -145
View File
@@ -15,14 +15,15 @@
"""Mesh processing."""
import collections
import dataclasses
import itertools
from typing import Dict, List, Optional, Sequence, Tuple
from typing import Tuple
import warnings
import mujoco
import jax
from jax import numpy as jp
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.collision_types import ConvexInfo
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
import numpy as np
@@ -30,37 +31,7 @@ from scipy import spatial
import trimesh
_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1)))
# pyformat: disable
# Rectangular box faces using a counter-clockwise winding order convention.
_BOX_FACES = [
0, 4, 5, 1, # left
0, 2, 6, 4, # bottom
6, 7, 5, 4, # front
2, 3, 7, 6, # right
1, 5, 7, 3, # top
0, 1, 3, 2, # back
]
# pyformat: enable
_MAX_HULL_FACE_VERTICES = 20
_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {}
_DERIVED_ARGS = [
'geom_convex_face',
'geom_convex_vert',
'geom_convex_edge_dir',
'geom_convex_facenormal',
'geom_convex_edge',
'geom_convex_edge_face_normal',
]
DERIVED = {(Model, d) for d in _DERIVED_ARGS}
def _box(size: np.ndarray):
"""Creates a mesh for a box with rectangular faces."""
box_corners = np.array(_BOX_CORNERS)
vert = box_corners * size.reshape(-1, 3)
face = np.array([_BOX_FACES]).reshape(-1, 4)
return vert, face
def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
@@ -170,16 +141,7 @@ def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
return hull_point_idx
@dataclasses.dataclass
class MeshInfo:
name: str
vert: np.ndarray
face: np.ndarray
convex_vert: Optional[np.ndarray]
convex_face: Optional[np.ndarray]
def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
def _merge_coplanar(m: Model, tm: trimesh.Trimesh, meshid: int) -> np.ndarray:
"""Merges coplanar facets."""
if not tm.facets:
return tm.faces.copy() # no facets
@@ -204,9 +166,11 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
# resize faces that exceed max polygon vertices
if face.shape[0] > _MAX_HULL_FACE_VERTICES:
name = m.names[m.name_meshadr[meshid]:]
name = name[:name.find(b'\x00')].decode('utf-8')
warnings.warn(
f'Mesh "{mesh_info.name}" has a coplanar face with more than'
f' {_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '
f'Mesh "{name}" has a coplanar face with more than '
f'{_MAX_HULL_FACE_VERTICES} vertices. This may lead to performance '
'issues and inaccuracies in collision detection. Consider '
'decimating the mesh.'
)
@@ -231,106 +195,99 @@ def _merge_coplanar(tm: trimesh.Trimesh, mesh_info: MeshInfo) -> np.ndarray:
return np.concatenate([faces, facets])
def _mesh_info(
m: mujoco.MjModel,
) -> List[MeshInfo]:
"""Extracts mesh info from MjModel."""
mesh_infos = []
for i in range(m.nmesh):
name = mujoco.mj_id2name(m, mujoco.mjtObj.mjOBJ_MESH.value, i)
last = (i + 1) >= m.nmesh
face_start = m.mesh_faceadr[i]
face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0]
face = m.mesh_face[face_start:face_end]
vert_start = m.mesh_vertadr[i]
vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0]
vert = m.mesh_vert[vert_start:vert_end]
graphadr = m.mesh_graphadr[i]
if graphadr < 0:
mesh_infos.append(MeshInfo(name, vert, face, None, None))
continue
graph = m.mesh_graph[graphadr:]
numvert, numface = graph[0], graph[1]
# unused vert_edgeadr
# vert_edgeadr = graph[2 : numvert + 2]
last_idx = numvert + 2
vert_globalid = graph[last_idx : last_idx + numvert]
last_idx += numvert
# unused edge_localid
# edge_localid = graph[last_idx : last_idx + numvert + 3 * numface]
last_idx += numvert + 3 * numface
face_globalid = graph[last_idx : last_idx + 3 * numface]
face_globalid = face_globalid.reshape((numface, 3))
convex_vert = vert[vert_globalid]
vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0])))
convex_face = np.vectorize(vertex_map.get)(face_globalid)
mesh_infos.append(MeshInfo(name, vert, face, convex_vert, convex_face))
return mesh_infos
def _geom_mesh_kwargs(
mesh_info: MeshInfo,
) -> Dict[str, np.ndarray]:
"""Generates convex mesh attributes for mjx.Model."""
tm_convex = trimesh.Trimesh(
vertices=mesh_info.convex_vert, faces=mesh_info.convex_face
def box(info: GeomInfo) -> ConvexInfo:
"""Creates a box with rectangular faces."""
vert = np.array(
list(itertools.product((-1, 1), (-1, 1), (-1, 1))), dtype=float
)
# pyformat: disable
# rectangular box faces using a counter-clockwise winding order convention:
face = np.array(
[
0, 4, 5, 1, # left
0, 2, 6, 4, # bottom
6, 7, 5, 4, # front
2, 3, 7, 6, # right
1, 5, 7, 3, # top
0, 1, 3, 2, # back
]
).reshape((-1, 4))
# pyformat: enable
face_normal = _get_face_norm(vert, face)
edge, edge_face_normal = _get_edge_normals(face, face_normal)
edge_dir = _get_unique_edge_dir(vert, face)
face = vert[face] # materialize full nface x nvert matrix
c = ConvexInfo(
info.pos,
info.mat,
vert,
face,
face_normal,
edge,
edge_face_normal,
edge_dir,
)
c = jax.tree_util.tree_map(jp.array, c)
vert = jax.vmap(jp.multiply, in_axes=(None, 0))(c.vert, info.size)
face = jax.vmap(jp.multiply, in_axes=(None, 0))(c.face, info.size)
c = c.replace(vert=vert, face=face)
return c
def convex(m: Model, mesh_id: int, info: GeomInfo) -> ConvexInfo:
"""Processes a mesh for use in convex collision algorithms.
Args:
m: an MJX model
mesh_id: the mesh id to process
info: pos, mat, size of this geom
Returns:
a convex mesh info
"""
vert_beg = m.mesh_vertadr[mesh_id]
vert_end = m.mesh_vertadr[mesh_id + 1] if mesh_id < m.nmesh - 1 else None
vert = m.mesh_vert[vert_beg:vert_end]
graphadr = m.mesh_graphadr[mesh_id]
graph = m.mesh_graph[graphadr:]
graph_idx = 0
numvert, numface = graph[0], graph[1]
graph_idx += 2
# skip vert_edgeadr (numvert,)
graph_idx += numvert
vert_globalid = graph[graph_idx : graph_idx + numvert]
graph_idx += numvert
# skip edge_localid (numvert, 3)
graph_idx += numvert + 3 * numface
face_globalid = graph[graph_idx : graph_idx + 3 * numface].reshape((-1, 3))
vert = vert[vert_globalid]
vertex_map = dict(zip(vert_globalid, np.arange(vert_globalid.shape[0])))
face = np.vectorize(vertex_map.get)(face_globalid)
tm_convex = trimesh.Trimesh(vertices=vert, faces=face)
vert = np.array(tm_convex.vertices)
face = _merge_coplanar(tm_convex, mesh_info)
facenormal = _get_face_norm(vert, face)
edge, edge_face_normal = _get_edge_normals(face, facenormal)
return {
'geom_convex_face': vert[face],
'geom_convex_face_vert_idx': face,
'geom_convex_vert': vert,
'geom_convex_edge_dir': _get_unique_edge_dir(vert, face),
'geom_convex_facenormal': facenormal,
'geom_convex_edge': edge,
'geom_convex_edge_face_normal': edge_face_normal,
}
face = _merge_coplanar(m, tm_convex, mesh_id)
face_normal = _get_face_norm(vert, face)
edge, edge_face_normal = _get_edge_normals(face, face_normal)
edge_dir = _get_unique_edge_dir(vert, face)
face = vert[face] # materialize full nface x nvert matrix
c = ConvexInfo(
info.pos,
info.mat,
vert,
face,
face_normal,
edge,
edge_face_normal,
edge_dir,
)
def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
"""Derives geom mesh attributes for mjx.Model from MjModel."""
kwargs = {k: [] for k in _DERIVED_ARGS}
mesh_infos = _mesh_info(m)
geom_con = m.geom_conaffinity | m.geom_contype
for geomid in range(m.ngeom):
mesh_info = None
dataid = m.geom_dataid[geomid]
if not geom_con[geomid]:
# ignore visual-only meshes
kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
continue
elif m.geom_type[geomid] == GeomType.BOX:
vert, face = _box(m.geom_size[geomid])
mesh_info = MeshInfo(
name='box',
vert=vert,
face=face,
convex_vert=vert,
convex_face=face,
)
elif dataid < 0:
kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
continue
mesh_info = mesh_info or mesh_infos[dataid]
vert, face = mesh_info.vert, mesh_info.face
key = (hash(vert.data.tobytes()), hash(face.data.tobytes()))
if key not in _CONVEX_CACHE:
_CONVEX_CACHE[key] = _geom_mesh_kwargs(mesh_info)
kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS}
return kwargs
return jax.tree_util.tree_map(jp.array, c)
mjx/mujoco/mjx/_src/mesh_test.py
+10 -17
View File
@@ -20,7 +20,7 @@ import numpy as np
import trimesh
class GeomMeshKwargsTest(absltest.TestCase):
class MeshTest(absltest.TestCase):
def test_pyramid(self):
"""Tests that a triangulated pyramid converts to merged coplanar faces."""
@@ -37,29 +37,21 @@ class GeomMeshKwargsTest(absltest.TestCase):
tm = trimesh.Trimesh(vertices=vert, faces=face)
tm_convex = trimesh.convex.convex_hull(tm)
convex_vert = np.array(tm_convex.vertices)
convex_face = np.array(tm_convex.faces)
mesh_info = mesh.MeshInfo(
name='test',
vert=vert,
face=face,
convex_vert=convex_vert,
convex_face=convex_face,
)
h = mesh._geom_mesh_kwargs(mesh_info)
convex_face = mesh._merge_coplanar(None, tm_convex, 0)
# get index of vertices in h['geom_convex_vert'] for vertices in vert
dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile(
h['geom_convex_vert'], (vert.shape[0], 1)
convex_vert, (vert.shape[0], 1)
)
dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1))
vidx = np.argmin(dist, axis=0)
# check verts
np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx])
np.testing.assert_array_equal(convex_vert, vert[vidx])
# check face vertices
map_ = {v: k for k, v in enumerate(vidx)}
h_face = np.vectorize(map_.get)(h['geom_convex_face_vert_idx'])
h_face = np.vectorize(map_.get)(convex_face)
face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()])
expected_face_verts = sorted([
(0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)])
@@ -69,7 +61,8 @@ class GeomMeshKwargsTest(absltest.TestCase):
)
# check edges
unique_edge = np.vectorize(map_.get)(h['geom_convex_edge_dir'])
edge_dir = mesh._get_unique_edge_dir(convex_vert, convex_face)
unique_edge = np.vectorize(map_.get)(edge_dir)
unique_edge = np.array(sorted(unique_edge.tolist()))
np.testing.assert_array_equal(
unique_edge,
@@ -77,10 +70,11 @@ class GeomMeshKwargsTest(absltest.TestCase):
)
# face normals
self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3))
face_normal = mesh._get_face_norm(convex_vert, convex_face)
self.assertEqual(face_normal.shape, (5, 3))
# face edges
edges = h['geom_convex_edge']
edges, edge_normal = mesh._get_edge_normals(convex_face, face_normal)
edges = np.vectorize(map_.get)(edges)
mask = edges[:, 0] != edges[:, 1]
edges = edges[mask]
@@ -103,7 +97,6 @@ class GeomMeshKwargsTest(absltest.TestCase):
)
# face edge normals
edge_normal = h['geom_convex_edge_face_normal']
edge_normal = edge_normal[mask]
edge_normal = np.take_along_axis(
edge_normal, sort_col_idx[..., None], axis=1
mjx/mujoco/mjx/_src/scan.py
+9 -9
View File
@@ -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))
mjx/mujoco/mjx/_src/scan_test.py
+4 -4
View File
@@ -55,7 +55,7 @@ class ScanTest(absltest.TestCase):
<worldbody/>
</mujoco>
""")
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 = (
mjx/mujoco/mjx/_src/solver.py
+14 -19
View File
@@ -17,7 +17,6 @@
import jax
from jax import numpy as jp
import mujoco
from mujoco.mjx._src import constraint
from mujoco.mjx._src import math
from mujoco.mjx._src import smooth
from mujoco.mjx._src import support
@@ -81,7 +80,7 @@ class _Context(PyTreeNode):
prev_cost=0.0,
solver_niter=0,
)
ctx = _update_constraint(m, d, ctx)
ctx = _update_constraint(d, ctx)
if grad:
ctx = _update_gradient(m, d, ctx)
ctx = ctx.replace(search=-ctx.Mgrad) # start with preconditioned gradient
@@ -107,7 +106,7 @@ class _LSPoint(PyTreeNode):
@classmethod
def create(
cls,
m: Model,
d: Data,
ctx: _Context,
alpha: jax.Array,
jv: jax.Array,
@@ -118,8 +117,7 @@ class _LSPoint(PyTreeNode):
# roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
# TODO(robotics-team): change this to support friction constraints
ne, nf, *_ = constraint.count_constraints(m)
active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True)
active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
quad = jax.vmap(jp.multiply)(quad, active) # only active
quad_total = quad_gauss + jp.sum(quad, axis=0)
@@ -161,13 +159,12 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]
def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
def _update_constraint(d: Data, ctx: _Context) -> _Context:
"""Updates constraint force and resulting cost given latst solver iteration.
Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
Args:
m: model defining constraints
d: data which contains latest qacc and smooth terms
ctx: current solver context
@@ -177,8 +174,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
# TODO(robotics-team): add friction constraints
# only count active constraints
ne, nf, *_ = constraint.count_constraints(m)
active = (ctx.Jaref < 0).at[:ne + nf].set(True)
active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
efc_force = d.efc_D * -ctx.Jaref * active
qfrc_constraint = d.efc_J.T @ efc_force
@@ -217,8 +213,7 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
if m.opt.solver == SolverType.CG:
mgrad = smooth.solve_m(m, d, grad)
elif m.opt.solver == SolverType.NEWTON:
ne, nf, *_ = constraint.count_constraints(m)
active = (ctx.Jaref < 0).at[: ne + nf].set(True)
active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
h = support.full_m(m, d) + h
h_ = jax.scipy.linalg.cho_factor(h)
@@ -262,7 +257,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
quad = (quad * d.efc_D).T
point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss)
point_fn = lambda a: _LSPoint.create(d, ctx, a, jv, quad, quad_gauss)
def cond(ctx: _LSContext) -> jax.Array:
done = ctx.ls_iter >= m.opt.ls_iterations
@@ -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:
mjx/mujoco/mjx/_src/support.py
+5
View File
@@ -208,3 +208,8 @@ def get_custom_numeric(m: Union[Model, mujoco.MjModel], name: str) -> float:
return m.numeric_data[m.numeric_adr[i]]
return -1
def get_custom_int(m: Union[Model, mujoco.MjModel], name: str) -> int:
"""Returns a custom integer given an MjModel or mjx.Model."""
return int(get_custom_numeric(m, name))
mjx/mujoco/mjx/_src/types.py
+54 -41
View File
@@ -15,10 +15,8 @@
"""Base types used in MJX."""
import enum
from typing import List, Optional
import jax
import jax.numpy as jp
import mujoco
from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member
import numpy as np
@@ -213,6 +211,22 @@ class BiasType(enum.IntEnum):
# unsupported: MUSCLE, USER
class ConstraintType(enum.IntEnum):
"""Type of constraint.
Attributes:
EQUALITY: equality constraint
LIMIT_JOINT: joint limit
CONTACT_PYRAMIDAL: frictional contact, pyramidal friction cone
"""
EQUALITY = mujoco.mjtConstraint.mjCNSTR_EQUALITY
# unsupported: FRICTION_DOF, FRICTION_TENDON
LIMIT_JOINT = mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT
# unsupported: LIMIT_TENDON, CONTACT_FRICTIONLESS
CONTACT_PYRAMIDAL = mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
# unsupported: CONTACT_ELLIPTIC
class CamLightType(enum.IntEnum):
"""Type of camera light.
@@ -371,6 +385,7 @@ class Model(PyTreeNode):
geom_solref: constraint solver reference: contact (ngeom, mjNREF)
geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP)
geom_size: geom-specific size parameters (ngeom, 3)
geom_rbound: radius of bounding sphere (ngeom,)
geom_pos: local position offset rel. to body (ngeom, 3)
geom_quat: local orientation offset rel. to body (ngeom, 4)
geom_friction: friction for (slide, spin, roll) (ngeom, 3)
@@ -388,17 +403,13 @@ class Model(PyTreeNode):
cam_poscom0: global position rel. to sub-com in qpos0 (ncam, 3)
cam_pos0: global position rel. to body in qpos0 (ncam, 3)
cam_mat0: global orientation in qpos0 (ncam, 9)
mat_rgba: rgba (nmat, 4)
mesh_vertadr: first vertex address (nmesh x 1)
mesh_faceadr: first face address (nmesh x 1)
mesh_vertadr: first vertex address (nmesh,)
mesh_faceadr: first face address (nmesh,)
mesh_graphadr: graph data address; -1: no graph (nmesh,)
mesh_vert: vertex positions for all meshes (nmeshvert, 3)
mesh_face: vertex face data (nmeshface, 3)
geom_convex_face: vertex face data, MJX only (ngeom,)
geom_convex_vert: vertex data, MJX only (ngeom,)
geom_convex_edge_dir: unique edge direction, MJX only (ngeom,)
geom_convex_facenormal: normal face data, MJX only (ngeom,)
geom_convex_face_edge: edges for each face (ngeom,)
geom_convex_face_edge_normal: edge normals for each face (ngeom,)
mesh_graph: convex graph data (nmeshgraph,)
mat_rgba: rgba (nmat, 4)
pair_dim: contact dimensionality (npair,)
pair_geom1: id of geom1 (npair,)
pair_geom2: id of geom2 (npair,)
@@ -435,6 +446,7 @@ class Model(PyTreeNode):
actuator_gear: scale length and transmitted force (nu, 6)
numeric_adr: address of field in numeric_data (nnumeric,)
numeric_data: array of all numeric fields (nnumericdata,)
name_meshadr: mesh name pointers (nmesh,)
name_numericadr: numeric name pointers (nnumeric,)
names: names of all objects, 0-terminated (nnames,)
"""
@@ -516,6 +528,7 @@ class Model(PyTreeNode):
geom_solref: jax.Array
geom_solimp: jax.Array
geom_size: jax.Array
geom_rbound: jax.Array
geom_pos: jax.Array
geom_quat: jax.Array
geom_friction: jax.Array
@@ -535,18 +548,14 @@ class Model(PyTreeNode):
cam_mat0: jax.Array
mesh_vertadr: np.ndarray
mesh_faceadr: np.ndarray
mesh_graphadr: np.ndarray
mesh_vert: np.ndarray
mesh_face: np.ndarray
mesh_graph: np.ndarray
mat_rgba: np.ndarray
pair_dim: np.ndarray
pair_geom1: np.ndarray
pair_geom2: np.ndarray
geom_convex_face: List[Optional[jax.Array]]
geom_convex_vert: List[Optional[jax.Array]]
geom_convex_edge_dir: List[Optional[jax.Array]]
geom_convex_facenormal: List[Optional[jax.Array]]
geom_convex_edge: List[Optional[jax.Array]]
geom_convex_edge_face_normal: List[Optional[jax.Array]]
pair_solref: jax.Array
pair_solreffriction: jax.Array
pair_solimp: jax.Array
@@ -580,6 +589,7 @@ class Model(PyTreeNode):
actuator_gear: jax.Array
numeric_adr: np.ndarray
numeric_data: np.ndarray
name_meshadr: np.ndarray
name_numericadr: np.ndarray
names: bytes
@@ -596,8 +606,11 @@ class Contact(PyTreeNode):
solref: constraint solver reference, normal direction (mjNREF,)
solreffriction: constraint solver reference, friction directions (mjNREF,)
solimp: constraint solver impedance (mjNIMP,)
geom1: id of geom 1
geom2: id of geom 2
dim: contact space dimensionality: 1, 3, 4, or 6
geom1: id of geom 1; deprecated, use geom[0]
geom2: id of geom 2; deprecated, use geom[1]
geom: geom ids (2,)
efc_address: address in efc; -1: not included
"""
dist: jax.Array
pos: jax.Array
@@ -607,33 +620,25 @@ class Contact(PyTreeNode):
solref: jax.Array
solreffriction: jax.Array
solimp: jax.Array
# unsupported: mu, H, dim
# unsupported: mu, H
dim: np.ndarray
geom1: jax.Array
geom2: jax.Array
# unsupported: efc_address, exclude
@classmethod
def zero(cls, ncon: int = 0) -> 'Contact':
"""Returns a contact filled with zeros."""
return Contact(
dist=jp.zeros(ncon),
pos=jp.zeros((ncon, 3,)),
frame=jp.zeros((ncon, 3, 3)),
includemargin=jp.zeros(ncon),
friction=jp.zeros((ncon, 5)),
solref=jp.zeros((ncon, mujoco.mjNREF)),
solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
geom1=jp.zeros(ncon, dtype=int),
geom2=jp.zeros(ncon, dtype=int),
)
geom: jax.Array
# unsupported: flex, elem, vert, exclude
efc_address: np.ndarray
class Data(PyTreeNode):
r"""Dynamic state that updates each step.\
r"""Dynamic state that updates each step.
Attributes:
solver_niter: number of solver iterations, per island (mjNISLAND,)
ne: number of equality constraints
nf: number of friction constraints
nl: number of limit constraints
nefc: number of constraints
ncon: number of contacts
solver_niter: number of solver iterations
time: simulation time
qpos: position (nq,)
qvel: velocity (nv,)
@@ -663,14 +668,15 @@ class Data(PyTreeNode):
cinert: com-based body inertia and mass (nbody, 10)
actuator_length: actuator lengths (nu,)
actuator_moment: actuator moments (nu, nv)
crb: com-based composite inertia and mass (nbody, 10)
crb: com-based composite inertia and mass (nbody, 10) \
qM: total inertia if sparse: (nM,)
if dense: (nv, nv)
qLD: L'*D*L (or Cholesky) factorization of M. if sparse: (nM,)
if dense: (nv, nv)
qLDiagInv: 1/diag(D) if sparse: (nv,)
if dense: (0,)
contact: list of all detected contacts (ncon,)
contact: all detected contacts (ncon,)
efc_type: constraint type (nefc,)
efc_J: constraint Jacobian (nefc, nv)
efc_frictionloss: frictionloss (friction) (nefc,)
efc_D: constraint mass (nefc,)
@@ -689,6 +695,12 @@ class Data(PyTreeNode):
efc_force: constraint force in constraint space (nefc,)
userdata: user data, not touched by engine (nuserdata,)
"""
# constant sizes:
ne: int
nf: int
nl: int
nefc: int
ncon: int
# solver statistics:
solver_niter: jax.Array
# global properties:
@@ -732,6 +744,7 @@ class Data(PyTreeNode):
qLD: jax.Array # pylint:disable=invalid-name
qLDiagInv: jax.Array # pylint:disable=invalid-name
contact: Contact
efc_type: np.ndarray
efc_J: jax.Array # pylint:disable=invalid-name
efc_frictionloss: jax.Array
efc_D: jax.Array # pylint:disable=invalid-name
mjx/mujoco/mjx/integration_test/collision_driver_test.py
+8 -2
View File
@@ -30,7 +30,10 @@ import numpy as np
def _assert_attr_eq(mjx_d, mj_d, attr, name, atol):
if attr == 'efc_address':
# we do not test efc_address since it gets set in constraint logic
# contact order not guaranteed to match
np.testing.assert_array_equal(
np.sort(mjx_d.efc_address), np.sort(mj_d.efc_address)
)
return
err_msg = f'mismatch: {attr} in run: {name}'
mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr)
@@ -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)
mjx/tutorial.ipynb
+1 -1
View File
@@ -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",
python/mujoco/functions.cc
+10 -6
View File
@@ -106,11 +106,11 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_defaultOption>(pymodule);
Def<traits::mj_defaultVisual>(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<std::string>& filename,
pymodule.def(
"mj_saveModel",
[](const MjModelWrapper& m, const std::optional<std::string>& filename = std::nullopt,
std::optional<
Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> buffer) {
Eigen::Ref<Eigen::Vector<std::uint8_t, Eigen::Dynamic>>> 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<py::gil_scoped_release>());
// Skipped: mj_loadModel (have MjModel.from_binary_path)
// Skipped: mj_deleteModel (have MjModel.__del__)
Def<traits::mj_sizeModel>(pymodule);
python/mujoco/renderer.py
+9 -4
View File
@@ -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)
src/user/user_api.cc
+11
View File
@@ -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<mjCBody*>(parent->element);
mjCFrame* child_frame = static_cast<mjCFrame*>(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<mjCModel*>(s->element);
src/user/user_api.h
+164 -135
View File
@@ -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;
src/user/user_model.cc
+1 -1
View File
@@ -3118,7 +3118,7 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
// compile objects in kinematic tree
for (int i=0; i<bodies.size(); i++) {
bodies[i]->Compile(); // 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
src/user/user_objects.cc
+75 -2
View File
@@ -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<mjCFrame*, int> fmap;
for (auto frame : subtree->frames) {
if (frame == static_cast<const mjCFrame*>(&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; i<subtree->bodies.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 <typename T>
void mjCBody::CopyList(std::vector<T*>& dst, const std::vector<T*>& src,
std::map<mjCFrame*, int>& fmap, const mjCFrame* pframe) {
int nsrc = (int)src.size();
for (int i=0; i<nsrc; i++) {
if (pframe && src[i]->frame != 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; i<frames.size(); i++) {
frames[i]->Compile();
}
// 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;
}
src/user/user_objects.h
+3 -1
View File
@@ -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
test/user/user_api_test.cc
+135 -53
View File
@@ -196,20 +196,20 @@ TEST_F(PluginTest, RecompileCompareCache) {
}
// -------------------------------- test attach -------------------------------
TEST_F(MujocoTest, AttachSame) {
std::array<char, 1000> er;
mjtNum tol = 0;
std::string field = "";
static constexpr char xml[] = R"(
static constexpr char xml_child[] = R"(
<mujoco>
<worldbody>
<body name="body">
<joint type="hinge" name="hinge"/>
<geom type="cylinder" size=".1 1 0"/>
<light mode="targetbody" target="child"/>
<body name="child"/>
</body>
<frame name="pframe">
<frame name="cframe">
<body name="body">
<joint type="hinge" name="hinge"/>
<geom type="cylinder" size=".1 1 0"/>
<light mode="targetbody" target="targetbody"/>
<body name="targetbody"/>
<body/>
</body>
</frame>
</frame>
<body name="ignore"/>
<frame name="frame" pos=".1 0 0" euler="0 90 0"/>
</worldbody>
@@ -227,26 +227,33 @@ TEST_F(MujocoTest, AttachSame) {
<position name="fixed" tendon="fixed"/>
</actuator>
<contact>
<exclude body1="body" body2="child"/>
<exclude body1="body" body2="targetbody"/>
</contact>
</mujoco>)";
TEST_F(MujocoTest, AttachSame) {
std::array<char, 1000> er;
mjtNum tol = 0;
std::string field = "";
static constexpr char xml_result[] = R"(
<mujoco>
<worldbody>
<body name="body">
<joint type="hinge" name="hinge"/>
<geom type="cylinder" size=".1 1 0"/>
<light mode="targetbody" target="child"/>
<body name="child"/>
<light mode="targetbody" target="targetbody"/>
<body name="targetbody"/>
<body/>
</body>
<body name="ignore"/>
<frame name="frame" pos=".1 0 0" euler="0 90 0">
<body name="attached-body-1">
<joint type="hinge" name="attached-hinge-1"/>
<geom type="cylinder" size=".1 1 0"/>
<light mode="targetbody" target="attached-child-1"/>
<body name="attached-child-1"/>
<light mode="targetbody" target="attached-targetbody-1"/>
<body name="attached-targetbody-1"/>
<body/>
</body>
</frame>
</worldbody>
@@ -270,13 +277,13 @@ TEST_F(MujocoTest, AttachSame) {
<position name="attached-fixed-1" tendon="attached-fixed-1"/>
</actuator>
<contact>
<exclude body1="body" body2="child"/>
<exclude body1="attached-body-1" body2="attached-child-1"/>
<exclude body1="body" body2="targetbody"/>
<exclude body1="attached-body-1" body2="attached-targetbody-1"/>
</contact>
</mujoco>)";
// 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) {
</worldbody>
</mujoco>)";
static constexpr char xml_child[] = R"(
<mujoco>
<worldbody>
<body name="body">
<joint type="hinge" name="hinge"/>
<geom type="cylinder" size=".1 1 0"/>
<body name="named"/>
<body/>
</body>
<body name="discard"/>
</worldbody>
<sensor>
<framepos name="keep" objtype="body" objname="body"/>
<framepos name="discard" objtype="body" objname="discard"/>
</sensor>
<tendon>
<fixed name="fixed">
<joint joint="hinge" coef="2"/>
</fixed>
</tendon>
<actuator>
<position joint="hinge"/>
<position tendon="fixed"/>
</actuator>
</mujoco>)";
static constexpr char xml_result[] = R"(
<mujoco>
<worldbody>
@@ -366,14 +347,15 @@ TEST_F(MujocoTest, AttachDifferent) {
<body name="attached-body-1">
<joint type="hinge" name="attached-hinge-1"/>
<geom type="cylinder" size=".1 1 0"/>
<body name="attached-named-1"/>
<light mode="targetbody" target="attached-targetbody-1"/>
<body name="attached-targetbody-1"/>
<body/>
</body>
</frame>
</body>
</worldbody>
<sensor>
<framepos name="attached-keep-1" objtype="body" objname="attached-body-1"/>
<framepos name="attached-sensor-1" objtype="body" objname="attached-body-1"/>
</sensor>
<tendon>
<fixed name="attached-fixed-1">
@@ -381,9 +363,12 @@ TEST_F(MujocoTest, AttachDifferent) {
</fixed>
</tendon>
<actuator>
<position joint="attached-hinge-1"/>
<position tendon="attached-fixed-1"/>
<position name="attached-hinge-1" joint="attached-hinge-1"/>
<position name="attached-fixed-1" tendon="attached-fixed-1"/>
</actuator>
<contact>
<exclude body1="attached-body-1" body2="attached-targetbody-1"/>
</contact>
</mujoco>)";
// 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<char, 1000> er;
mjtNum tol = 0;
std::string field = "";
static constexpr char xml_parent[] = R"(
<mujoco>
<worldbody>
<body name="sphere">
<freejoint/>
<geom size=".1"/>
<frame name="frame" pos=".1 0 0" euler="0 90 0"/>
</body>
</worldbody>
</mujoco>)";
static constexpr char xml_result[] = R"(
<mujoco>
<worldbody>
<body name="sphere">
<freejoint/>
<geom size=".1"/>
<frame name="frame" pos=".1 0 0" euler="0 90 0"/>
<frame name="pframe">
<frame name="cframe">
<body name="attached-body-1">
<joint type="hinge" name="attached-hinge-1"/>
<geom type="cylinder" size=".1 1 0"/>
<light mode="targetbody" target="attached-targetbody-1"/>
<body name="attached-targetbody-1"/>
<body/>
</body>
</frame>
</frame>
</body>
</worldbody>
<sensor>
<framepos name="attached-sensor-1" objtype="body" objname="attached-body-1"/>
</sensor>
<tendon>
<fixed name="attached-fixed-1">
<joint joint="attached-hinge-1" coef="2"/>
</fixed>
</tendon>
<actuator>
<position name="attached-hinge-1" joint="attached-hinge-1"/>
<position name="attached-fixed-1" tendon="attached-fixed-1"/>
</actuator>
<contact>
<exclude body1="attached-body-1" body2="attached-targetbody-1"/>
</contact>
</mujoco>)";
// 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
test/xml/xml_native_writer_test.cc
+32 -1
View File
@@ -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);
}
}