diff --git a/doc/changelog.rst b/doc/changelog.rst index a2f528fc..59bbad37 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -14,12 +14,21 @@ General 2. Added :ref:`timeconst` attribute to the :ref:`position actuator`. When set to a positive value, the actuator is made stateful with :at:`filterexact` dynamics. +MJX +^^^ + +3. Add height-field collision support. Fixes :github:issue:`1491`. +4. Add a pre-compiled field ``mesh_convex`` to ``mjx.Model`` so that mesh properties can be vmapped over. + Fixes :github:issue:`1655`. +5. Fix a bug in convex mesh collisions, where erroneous edge contacts were being created even though face + separating axes were found. Fixes :github:issue:`1695`. + Bug fixes ^^^^^^^^^ -3. Fixed a bug the could cause collisions to be missed when :ref:`fusestatic` is enabled, as is +6. Fixed a bug the could cause collisions to be missed when :ref:`fusestatic` is enabled, as is often the case for URDF imports. Fixes :github:issue:`1069`, :github:issue:`1577`. -4. Fixed a bug that was causing the visualization of SDF iterations to write outside the size of the vector storing +7. Fixed a bug that was causing the visualization of SDF iterations to write outside the size of the vector storing them. Fixes :github:issue:`1539`. Version 3.1.5 (May 7, 2024) diff --git a/doc/mjx.rst b/doc/mjx.rst index 0e3c52f8..361c9cc4 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -196,7 +196,7 @@ The following features are **fully supported** in MJX: * - :ref:`Actuator Bias ` - ``NONE``, ``AFFINE`` * - :ref:`Geom ` - - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` + - ``PLANE``, ``HFIELD``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` are fully implemented. ``ELLIPSOID`` and ``CYLINDER`` are implemented but only collide with other primitives. * - :ref:`Constraint ` - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_FRICTIONLESS``, ``CONTACT_PYRAMIDAL`` * - :ref:`Equality ` @@ -223,7 +223,7 @@ The following features are **in development** and coming soon: * - Category - Feature * - :ref:`Geom ` - - ``SDF``, ``HFIELD``, ``ELLIPSOID``, ``CYLINDER`` + - ``SDF``. Collisions between (``SPHERE``, ``BOX``, ``MESH``, ``HFIELD``) and ``CYLINDER``. Collisions between (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``. * - :ref:`Constraint ` - :ref:`Frictionloss `, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` * - :ref:`Integrator ` diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py index 3c3efd5c..7d2c10dd 100644 --- a/mjx/mujoco/mjx/_src/collision_convex.py +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -14,8 +14,9 @@ # ============================================================================== """Convex collisions.""" +from collections.abc import Callable import functools -from typing import Tuple +from typing import Tuple, Union import jax from jax import numpy as jp @@ -26,16 +27,19 @@ 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.collision_types import HFieldInfo 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 +_GeomInfo = Union[GeomInfo, ConvexInfo] + def collider(ncon: int): """Wraps collision functions for use by collision_driver.""" - def wrapper(func): + def wrapper(collision_fn): def collide( m: Model, d: Data, key: FunctionKey, geom: jax.Array ) -> Collision: @@ -45,18 +49,25 @@ def collider(ncon: int): GeomInfo(d.geom_xpos[g2], d.geom_xmat[g2], m.geom_size[g2]), ] in_axes = [0, 0] + fn = collision_fn 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 + pos=0, mat=0, size=0, face=0, vert=0 ) elif key.types[i] == GeomType.MESH: - infos[i] = mesh.convex(m, key.data_ids[i], infos[i]) + c, cm = infos[i], m.mesh_convex[key.data_ids[i]] + infos[i] = ConvexInfo(**vars(c), **vars(cm)) in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace( - pos=0, mat=0 + pos=0, mat=0, size=0 ) - dist, pos, frame = jax.vmap(func, in_axes=in_axes)(*infos) + elif key.types[i] == GeomType.HFIELD: + hfield_info = mesh.hfield(m, key.data_ids[i]) + infos[i] = hfield_info.replace(pos=infos[i].pos, mat=infos[i].mat) + in_axes[i] = hfield_info.replace(pos=0, mat=0, data=None) + fn = functools.partial(fn, subgrid_size=key.subgrid_size) + dist, pos, frame = jax.vmap(fn, in_axes=in_axes)(*infos) if ncon > 1: return jax.tree_util.tree_map(jp.concatenate, (dist, pos, frame)) return dist, pos, frame @@ -242,9 +253,8 @@ def plane_convex(plane: GeomInfo, convex: ConvexInfo) -> Collision: return dist, pos, frame -@collider(ncon=1) -def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: - """Calculates contact between a sphere and a convex object.""" +def _sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: + """Calculates contact between a sphere and a convex mesh.""" faces = convex.face normals = convex.face_normal @@ -276,7 +286,7 @@ def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: face_normal, ) edge_dist = jax.vmap( - lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm) + lambda plane_pt, plane_norm, pt=pt: (pt - plane_pt).dot(plane_norm) )(edge_p0, side_normals) pt_on_face = jp.all(edge_dist <= 0) # lte to handle degenerate edges @@ -291,7 +301,8 @@ def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: # Get the normal, dist, and contact position. pt_normal, d = math.normalize_with_norm(pt - sphere_pos) - # Ensure normal points towards convex centroid. + # Ensure normal points towards convex centroid. Assume convex centroid is at + # the origin. inside = jp.dot(pt, pt_normal) > 0 sign = jp.where(inside, -1, 1) n = jp.where(pt_on_face | (d < 1e-6), -face_normal, sign * pt_normal) @@ -305,11 +316,17 @@ def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: n = convex.mat @ n pos = convex.mat @ pos + convex.pos + return dist, pos, n + + +@collider(ncon=1) +def sphere_convex(sphere: GeomInfo, convex: ConvexInfo) -> Collision: + """Calculates contact between a sphere and a convex mesh.""" + dist, pos, n = _sphere_convex(sphere, convex) return dist, pos, math.make_frame(n) -@collider(ncon=2) -def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision: +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 @@ -431,6 +448,13 @@ def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision: dist = -jp.where( has_edge_contact, jp.array([edge_penetration, -1]), face_penetration ) + return dist, pos, n + + +@collider(ncon=2) +def capsule_convex(cap: GeomInfo, convex: ConvexInfo) -> Collision: + """Calculates contacts between a capsule and a convex object.""" + dist, pos, n = _capsule_convex(cap, convex) frame = jax.vmap(math.make_frame)(n) return dist, pos, frame @@ -851,6 +875,7 @@ def _sat_gaussmap( incident_face_norm, -best_axis, ) + dist = jp.where(is_face_separating, 1.0, dist) # Handle edge separating axes by checking all edge pairs. a_idx = jp.tile(jp.arange(edges_a.shape[0]), reps=edges_b.shape[0]) @@ -891,8 +916,9 @@ def _sat_gaussmap( best_edge_idx = edge_dist.argmax() best_edge_dist = edge_dist[best_edge_idx] is_edge_contact = jp.where( - dist.max() < 0, best_edge_dist > dist.max() - 1e-6, - (best_edge_dist < 0) & ~jp.isinf(best_edge_dist) + dist.max() < 0.0, + best_edge_dist > dist.max() - 1e-6, + (best_edge_dist < 0) & ~jp.isinf(best_edge_dist), ) is_edge_contact = is_edge_contact & ~is_face_separating normal = jp.where(is_edge_contact, edge_axes[best_edge_idx], normal) @@ -911,9 +937,45 @@ def _sat_gaussmap( return dist, pos, normal -@collider(ncon=4) -def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: - """Calculates contacts between two convex objects.""" +def _box_box(b1: ConvexInfo, b2: ConvexInfo) -> Collision: + """Calculates contacts between two boxes.""" + faces1 = b1.face + faces2 = b2.face + + to_local_pos = b2.mat.T @ (b1.pos - b2.pos) + to_local_mat = b2.mat.T @ b1.mat + + faces1 = to_local_pos + faces1 @ to_local_mat.T + normals1 = b1.face_normal @ to_local_mat.T + normals2 = b2.face_normal + + vertices1 = to_local_pos + b1.vert @ to_local_mat.T + vertices2 = b2.vert + + unique_edges1 = jp.take(vertices1, b1.edge_dir, axis=0) + unique_edges2 = jp.take(vertices2, b2.edge_dir, axis=0) + + # brute-force SAT is more performant for box-box + dist, pos, normal = _sat_bruteforce( + faces1, + faces2, + vertices1, + vertices2, + normals1, + normals2, + unique_edges1, + unique_edges2, + ) + + # Go back to world frame. + pos = b2.pos + pos @ b2.mat.T + n = normal @ b2.mat.T + + return dist, pos, n + + +def _convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: + """Calculates contacts between two convex meshes.""" # 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] @@ -932,6 +994,7 @@ def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: faces1 = c1.face faces2 = c2.face + # convert to c2 frame to_local_pos = c2.mat.T @ (c1.pos - c2.pos) to_local_mat = c2.mat.T @ c1.mat @@ -942,49 +1005,209 @@ def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: vertices1 = to_local_pos + c1.vert @ to_local_mat.T vertices2 = c2.vert - 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) edge_face_normals1 = c1.edge_face_normal @ to_local_mat.T edge_face_normals2 = c2.edge_face_normal - enable_bruteforce = ( - unique_edges1.shape[0] * unique_edges2.shape[0] - < edges1[0].shape[0] * edges2[0].shape[0] + dist, pos, normal = _sat_gaussmap( + to_local_pos, + faces1, + faces2, + vertices1, + vertices2, + normals1, + normals2, + edges1, + edges2, + edge_face_normals1, + edge_face_normals2, ) - if enable_bruteforce: - dist, pos, normal = _sat_bruteforce( - faces1, - faces2, - vertices1, - vertices2, - normals1, - normals2, - unique_edges1, - unique_edges2, - ) - else: - dist, pos, normal = _sat_gaussmap( - to_local_pos, - faces1, - faces2, - vertices1, - vertices2, - normals1, - normals2, - edges1, - edges2, - edge_face_normals1, - edge_face_normals2, - ) # Go back to world frame. 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) + n = normal @ c2.mat.T + n = -n if swapped else n + return dist, pos, n + + +@collider(ncon=4) +def box_box(b1: ConvexInfo, b2: ConvexInfo) -> Collision: + """Calculates contacts between two boxes.""" + dist, pos, n = _box_box(b1, b2) + frame = jax.vmap(math.make_frame)(n) return dist, pos, frame + + +@collider(ncon=4) +def convex_convex(c1: ConvexInfo, c2: ConvexInfo) -> Collision: + """Calculates contacts between two convex objects.""" + dist, pos, n = _convex_convex(c1, c2) + frame = jax.vmap(math.make_frame)(n) + return dist, pos, frame + + +def _hfield_collision( + collider_fn: Callable[[_GeomInfo, _GeomInfo], Collision], + h: HFieldInfo, + obj: _GeomInfo, + obj_rbound: jax.Array, + subgrid_size: Tuple[int, int], +) -> Collision: + """Collides an object with prisms in a height field.""" + # put obj in hfield frame + obj_pos = h.mat.T @ (obj.pos - h.pos) + obj_mat = h.mat.T @ obj.mat + + xmin = obj_pos[0] - obj_rbound + ymin = obj_pos[1] - obj_rbound + cmin = jp.floor((xmin + h.size[0]) / (2 * h.size[0]) * (h.ncol - 1)) + cmin = cmin.astype(int) + rmin = jp.floor((ymin + h.size[1]) / (2 * h.size[1]) * (h.nrow - 1)) + rmin = rmin.astype(int) + + # compute real-valued grid step + dx = 2.0 * h.size[0] / (h.ncol - 1) + dy = 2.0 * h.size[1] / (h.nrow - 1) + + # set zbottom value using base size + bvert = jp.array([0.0, 0.0, -h.size[3]]) + bmask = jp.array([True, True, False]) + + # process all prisms in sub-grid + prisms = [] + for r in range(subgrid_size[1]): + for c in range(subgrid_size[0]): + ri, ci = rmin + r, cmin + c + + # ensure ri, ci are in the bounds of the hfield + ri = jp.clip(ri, 0, h.nrow - 2) + ci = jp.clip(ci, 0, h.ncol - 2) + + p1 = [ + dx * ci - h.size[0], + dy * ri - h.size[1], + h.data[ci, ri] * h.size[2], + ] + p2 = [ + dx * (ci + 1) - h.size[0], + dy * (ri + 1) - h.size[1], + h.data[ci + 1, ri + 1] * h.size[2], + ] + p3 = [ + dx * ci - h.size[0], + dy * (ri + 1) - h.size[1], + h.data[ci, ri + 1] * h.size[2], + ] + top = jp.array([p1, p2, p3]) + bottom = jp.array([p1, p3, p2]) * bmask + bvert + vert = jp.concatenate([bottom, top]) + prisms.append(mesh.hfield_prism(vert)) + + p3 = p2 + p2 = [ + dx * (ci + 1) - h.size[0], + dy * ri - h.size[1], + h.data[ci + 1, ri] * h.size[2], + ] + top = jp.array([p1, p2, p3]) + bottom = jp.array([p1, p3, p2]) * bmask + bvert + vert = jp.concatenate([bottom, top]) + # NB: If the order of verts is updated above, the corresponding + # hfield_prism function must be updated to ensure that all faces have the + # correct winding order. + prisms.append(mesh.hfield_prism(vert)) + + n_prisms = len(prisms) + prisms = jax.tree_util.tree_map(lambda *x: jp.stack(x), *prisms) + dist, pos, n = jax.vmap(collider_fn, in_axes=[None, 0])( + obj.replace(pos=obj_pos, mat=obj_mat), prisms + ) + + dist = dist.flatten() + pos = pos.reshape((-1, 3)) + n = n.reshape((-1, 3)) + n *= -1 # flip the normal since we flipped args in the call to collider_fn + + # Check that we're in the half-space of the hfield norm. If not, pick the top + # face norm. This resolves issues with cracks of doom. + n_repeats = dist.shape[0] // n_prisms + top_norm = jp.repeat(prisms.face_normal[:, 1], n_repeats, axis=0) + cond = jax.vmap(jp.dot, in_axes=[0, None])(n, h.mat[2]) < 1e-6 + n = jp.where(cond[:, None], top_norm, n) + + return dist, pos, n + + +@collider(ncon=4) +def hfield_sphere( + h: HFieldInfo, s: GeomInfo, subgrid_size: Tuple[int, int] +) -> Collision: + """Calculates contacts between a hfield and a sphere.""" + rbound = jp.max(s.size) + dist, pos, n = _hfield_collision(_sphere_convex, h, s, rbound, subgrid_size) + + n_mean = jp.mean(n, axis=0) + mask = dist < jp.minimum(0, dist.min() + 1e-3) + idx = _manifold_points(pos, mask, n_mean) + dist, pos, n = dist[idx], pos[idx], n[idx] + + # zero out non-unique contacts + unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1 + dist = jp.where(unique, dist, 1) + + # back to world frame, _hfield_collision returns collision in hfield frame + pos = jax.vmap(lambda p: h.mat @ p + h.pos)(pos) + n = jax.vmap(lambda n: h.mat @ n)(n) + + return dist, pos, jax.vmap(math.make_frame)(n) + + +@collider(ncon=4) +def hfield_capsule( + h: HFieldInfo, c: GeomInfo, subgrid_size: Tuple[int, int] +) -> Collision: + """Calculates contacts between a hfield and a capsule.""" + rbound = c.size[0] + c.size[1] + dist, pos, n = _hfield_collision(_capsule_convex, h, c, rbound, subgrid_size) + + n_mean = jp.mean(n, axis=0) + mask = dist < jp.minimum(0, dist.min() + 1e-3) + idx = _manifold_points(pos, mask, n_mean) + dist, pos, n = dist[idx], pos[idx], n[idx] + + # zero out non-unique contacts + unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1 + dist = jp.where(unique, dist, 1) + + # back to world frame, _hfield_collision returns collision in hfield frame + pos = jax.vmap(lambda p: h.mat @ p + h.pos)(pos) + n = jax.vmap(lambda n: h.mat @ n)(n) + + return dist, pos, jax.vmap(math.make_frame)(n) + + +@collider(ncon=4) +def hfield_convex( + h: HFieldInfo, c: ConvexInfo, subgrid_size: Tuple[int, int] +) -> Collision: + """Calculates contacts between a hfield and a capsule.""" + rbound = jp.max(c.size) + dist, pos, n = _hfield_collision(_convex_convex, h, c, rbound, subgrid_size) + + n_mean = jp.mean(n, axis=0) + mask = dist < jp.minimum(0, dist.min() + 1e-3) + idx = _manifold_points(pos, mask, n_mean) + dist, pos, n = dist[idx], pos[idx], n[idx] + + # zero out non-unique contacts + unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1 + dist = jp.where(unique, dist, 1) + + # back to world frame, _hfield_collision returns collision in hfield frame + pos = jax.vmap(lambda p: h.mat @ p + h.pos)(pos) + n = jax.vmap(lambda n: h.mat @ n)(n) + + return dist, pos, jax.vmap(math.make_frame)(n) diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index b67168d3..8791fed7 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -45,8 +45,12 @@ from jax import numpy as jp import mujoco from mujoco.mjx._src import support # pylint: disable=g-importing-member +from mujoco.mjx._src.collision_convex import box_box 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 hfield_capsule +from mujoco.mjx._src.collision_convex import hfield_convex +from mujoco.mjx._src.collision_convex import hfield_sphere from mujoco.mjx._src.collision_convex import plane_convex from mujoco.mjx._src.collision_convex import sphere_convex from mujoco.mjx._src.collision_primitive import capsule_capsule @@ -78,6 +82,10 @@ _COLLISION_FUNC = { (GeomType.PLANE, GeomType.ELLIPSOID): plane_ellipsoid, (GeomType.PLANE, GeomType.CYLINDER): plane_cylinder, (GeomType.PLANE, GeomType.MESH): plane_convex, + (GeomType.HFIELD, GeomType.SPHERE): hfield_sphere, + (GeomType.HFIELD, GeomType.CAPSULE): hfield_capsule, + (GeomType.HFIELD, GeomType.BOX): hfield_convex, + (GeomType.HFIELD, GeomType.MESH): hfield_convex, (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere, (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule, (GeomType.SPHERE, GeomType.BOX): sphere_convex, @@ -90,7 +98,7 @@ _COLLISION_FUNC = { (GeomType.ELLIPSOID, GeomType.ELLIPSOID): ellipsoid_ellipsoid, (GeomType.ELLIPSOID, GeomType.CYLINDER): ellipsoid_cylinder, (GeomType.CYLINDER, GeomType.CYLINDER): cylinder_cylinder, - (GeomType.BOX, GeomType.BOX): convex_convex, + (GeomType.BOX, GeomType.BOX): box_box, (GeomType.BOX, GeomType.MESH): convex_convex, (GeomType.MESH, GeomType.MESH): convex_convex, } @@ -210,6 +218,21 @@ def _geom_groups( condim = max(m.geom_condim[g1], m.geom_condim[g2]) key = FunctionKey(types, data_ids, condim) + + if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD: + # add static grid bounds to the grouping key for hfield collisions + geom_rbound_hfield = ( + m.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound + ) + nrow, ncol = m.hfield_nrow[data_ids[0]], m.hfield_ncol[data_ids[0]] + xsize, ysize = m.hfield_size[data_ids[0]][:2] + xtick, ytick = (2 * xsize) / (ncol - 1), (2 * ysize) / (nrow - 1) + xbound = int(np.ceil(2 * geom_rbound_hfield[g2] / xtick)) + 1 + xbound = min(xbound, ncol) + ybound = int(np.ceil(2 * geom_rbound_hfield[g2] / ytick)) + 1 + ybound = min(ybound, nrow) + key = FunctionKey(types, data_ids, condim, (xbound, ybound)) + groups.setdefault(key, []).append((g1, g2, ip)) return groups diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py index 294c6a67..5492ea7f 100644 --- a/mjx/mujoco/mjx/_src/collision_driver_test.py +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -721,6 +721,95 @@ class ConvexTest(absltest.TestCase): self.assertTrue((c.dist > 0).all()) +class HFieldTest(absltest.TestCase): + _HFIELD = """ + + + + + + + + + + + + + + + + + + + + + + + + + + """ + + def test_sphere_hfield(self): + m = mujoco.MjModel.from_xml_string(self._HFIELD) + mx = mjx.put_model(m) + + d = mujoco.MjData(m) + d.qpos[:] = m.keyframe('qpos1').qpos + dx = mjx.put_data(m, d) + + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + # check that all geoms are colliding with the hfield + for geom_id in [1, 2, 3]: + mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1) + c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact) + self.assertTrue((c.dist < 0).any()) + self.assertTrue((c.dist > -1e-3).any()) + # all contact normals are roughly pointing in the right direction + self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all()) + + def test_hfield_outside(self): + """Tests that objects outside of the hfield do not collide.""" + positions = ['2.0 0', '-2.0 0', '0 -2.0', '0 2.0'] + for p in positions: + xml = self._HFIELD.replace('= 0).all()) + + def test_hfield_deep(self): + """Tests that objects with deep penetration do not get stuck.""" + m = mujoco.MjModel.from_xml_string(self._HFIELD) + mx = mjx.put_model(m) + + d = mujoco.MjData(m) + d.qpos[:] = m.keyframe('qpos2').qpos + dx = mjx.put_data(m, d) + + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + # check that all geoms are colliding with the hfield + for geom_id in [1, 2, 3]: + mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1) + c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact) + # all contact normals are in the top half-face of the hfield + self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all()) + + class BodyPairFilterTest(absltest.TestCase): """Tests that certain body pairs get filtered.""" diff --git a/mjx/mujoco/mjx/_src/collision_types.py b/mjx/mujoco/mjx/_src/collision_types.py index 49d77438..e7af8f23 100644 --- a/mjx/mujoco/mjx/_src/collision_types.py +++ b/mjx/mujoco/mjx/_src/collision_types.py @@ -15,31 +15,10 @@ """Collision base types.""" import dataclasses -from typing import Tuple +from typing import Optional, 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 - +from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member +import numpy as np # Collision returned by collision functions: # - distance distance between nearest points; neg: penetration @@ -48,18 +27,53 @@ class ConvexInfo(PyTreeNode): Collision = Tuple[jax.Array, jax.Array, jax.Array] +class GeomInfo(PyTreeNode): + """Geom properties for primitive shapes.""" + + pos: jax.Array + mat: jax.Array + size: jax.Array + + +class ConvexInfo(PyTreeNode): + """Geom properties for convex meshes.""" + + pos: jax.Array + mat: jax.Array + size: jax.Array + vert: jax.Array + face: jax.Array + face_normal: jax.Array + edge: jax.Array + edge_face_normal: jax.Array + edge_dir: Optional[jax.Array] = None + + +class HFieldInfo(PyTreeNode): + """Geom properties for height fields.""" + + pos: jax.Array + mat: jax.Array + size: np.ndarray + nrow: int + ncol: int + data: 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. + 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. + resulting constraint jacobian is determined at compile time. + subgrid_size: the size determines the hfield subgrid to collide with """ types: Tuple[int, int] data_ids: Tuple[int, int] condim: int + subgrid_size: Tuple[int, int] = (-1, -1) diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py index ba513d96..a96bb966 100644 --- a/mjx/mujoco/mjx/_src/dataclasses.py +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -30,7 +30,7 @@ def _jax_in_args(typ) -> bool: return True if dataclasses.is_dataclass(typ): return any(_jax_in_args(f.type) for f in dataclasses.fields(typ)) - if typing.get_origin(typ) in (list, dict, Union, set): + if typing.get_origin(typ) in (tuple, list, dict, Union, set): return any(_jax_in_args(t) for t in typing.get_args(typ)) return False diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 594ca15f..5068373a 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -22,6 +22,7 @@ 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 @@ -68,14 +69,16 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: if m.ntendon: raise NotImplementedError('tendons are not supported') + mesh_geomid = set() 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): + # margin/gap not supported for meshes and height fields + no_margin = {mujoco.mjtGeom.mjGEOM_MESH, mujoco.mjtGeom.mjGEOM_HFIELD} + if no_margin.intersection({t1, t2}): if ip != -1: margin = m.pair_margin[ip] else: @@ -83,6 +86,9 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: if margin.any(): t1, t2 = mujoco.mjtGeom(t1), mujoco.mjtGeom(t2) raise NotImplementedError(f'({t1}, {t2}) margin/gap not implemented.') + for t, g in [(t1, g1), (t2, g2)]: + if t == mujoco.mjtGeom.mjGEOM_MESH: + mesh_geomid.add(g) for enum_field, enum_type, mj_type in ( (m.actuator_biastype, types.BiasType, mujoco.mjtBias), @@ -100,12 +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.') - fields = {f.name: getattr(m, f.name) for f in types.Model.fields()} + mjx_only = {'mesh_convex', 'geom_rbound_hfield'} + mj_field_names = {f.name for f in types.Model.fields()} - mjx_only + fields = {f: getattr(m, f) for f in mj_field_names} + fields['geom_rbound_hfield'] = fields['geom_rbound'] 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) + + # Pre-compile meshes for MJX collisions. + fields['mesh_convex'] = [None] * m.nmesh + for i in mesh_geomid: + dataid = m.geom_dataid[i] + if fields['mesh_convex'][dataid] is None: + fields['mesh_convex'][dataid] = mesh.convex(m, dataid) # pytype: disable=unsupported-operands + fields['mesh_convex'] = tuple(fields['mesh_convex']) + model = types.Model(**{k: copy.copy(v) for k, v in fields.items()}) return jax.device_put(model, device=device) diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py index 2498aae2..31f45571 100644 --- a/mjx/mujoco/mjx/_src/mesh.py +++ b/mjx/mujoco/mjx/_src/mesh.py @@ -16,14 +16,18 @@ import collections import itertools -from typing import Tuple +from typing import Tuple, Union import warnings import jax from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math # pylint: disable=g-importing-member from mujoco.mjx._src.collision_types import ConvexInfo from mujoco.mjx._src.collision_types import GeomInfo +from mujoco.mjx._src.collision_types import HFieldInfo +from mujoco.mjx._src.types import ConvexMesh from mujoco.mjx._src.types import Model # pylint: enable=g-importing-member import numpy as np @@ -141,7 +145,9 @@ def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray: return hull_point_idx -def _merge_coplanar(m: Model, tm: trimesh.Trimesh, meshid: int) -> np.ndarray: +def _merge_coplanar( + m: Union[mujoco.MjModel, Model], tm: trimesh.Trimesh, meshid: int +) -> np.ndarray: """Merges coplanar facets.""" if not tm.facets: return tm.faces.copy() # no facets @@ -221,6 +227,7 @@ def box(info: GeomInfo) -> ConvexInfo: c = ConvexInfo( info.pos, info.mat, + info.size, vert, face, face_normal, @@ -236,22 +243,21 @@ def box(info: GeomInfo) -> ConvexInfo: return c -def convex(m: Model, mesh_id: int, info: GeomInfo) -> ConvexInfo: +def convex(m: Union[mujoco.MjModel, Model], data_id: int) -> ConvexMesh: """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 + data_id: the mesh id to process Returns: - a convex mesh info + a convex mesh """ - vert_beg = m.mesh_vertadr[mesh_id] - vert_end = m.mesh_vertadr[mesh_id + 1] if mesh_id < m.nmesh - 1 else None + vert_beg = m.mesh_vertadr[data_id] + vert_end = m.mesh_vertadr[data_id + 1] if data_id < m.nmesh - 1 else None vert = m.mesh_vert[vert_beg:vert_end] - graphadr = m.mesh_graphadr[mesh_id] + graphadr = m.mesh_graphadr[data_id] graph = m.mesh_graph[graphadr:] graph_idx = 0 @@ -273,21 +279,98 @@ def convex(m: Model, mesh_id: int, info: GeomInfo) -> ConvexInfo: tm_convex = trimesh.Trimesh(vertices=vert, faces=face) vert = np.array(tm_convex.vertices) - face = _merge_coplanar(m, tm_convex, mesh_id) + face = _merge_coplanar(m, tm_convex, data_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, + c = ConvexMesh( vert, face, face_normal, edge, edge_face_normal, - edge_dir, ) return jax.tree_util.tree_map(jp.array, c) + + +def hfield_prism(vert: jax.Array) -> ConvexInfo: + """Builds a hfield prism.""" + # The first 3 vertices define the bottom triangle, and the next 3 vertices + # define the top triangle. The remaining triangles define the side of the + # prism. + face = np.array([ + [0, 1, 2, 0], # bottom + [3, 4, 5, 3], # top + [0, 3, 5, 1], + [0, 2, 4, 3], + [2, 1, 5, 4], + ]) + edges = np.array([ + # bottom + [0, 1], + [1, 2], + [0, 2], + # top + [3, 4], + [3, 5], + [4, 5], + # sides + [0, 3], + [1, 5], + [2, 4], + ]) + edge_face_norm = np.array([ + # bottom + [0, 2], + [0, 4], + [0, 3], + # top + [1, 3], + [1, 2], + [1, 4], + # sides + [2, 3], + [2, 4], + [3, 4], + ]) + + def get_face_norm(face): + # use ccw winding order convention, and avoid using the last vertex + edge0 = face[2, :] - face[1, :] + edge1 = face[0, :] - face[1, :] + return math.normalize(jp.cross(edge0, edge1)) + + centroid = jp.mean(vert, axis=0) + vert = vert - centroid + face = vert[face] + face_norm = jax.vmap(get_face_norm)(face) + + c = ConvexInfo( + centroid, + jp.eye(3, dtype=float), + jp.ones(3), + vert, + face, + face_norm, + edges, + face_norm[edge_face_norm], + None, + ) + + return jax.tree_util.tree_map(jp.array, c) + + +def hfield(m: Union[mujoco.MjModel, Model], data_id: int) -> HFieldInfo: + adr = m.hfield_adr[data_id] + nrow, ncol = m.hfield_nrow[data_id], m.hfield_ncol[data_id] + h = HFieldInfo( + jp.zeros(3, dtype=float), + jp.eye(3, dtype=float), + m.hfield_size[data_id], + nrow, + ncol, + m.hfield_data[adr : adr + nrow * ncol].reshape((ncol, nrow), order='F'), + ) + return h diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index 421126f9..0cd31636 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -197,6 +197,7 @@ def _getnum(m: Union[Model, mujoco.MjModel], obj: mujoco._enums.mjtObj) -> int: mujoco.mjtObj.mjOBJ_SITE: m.nsite, mujoco.mjtObj.mjOBJ_CAMERA: m.ncam, mujoco.mjtObj.mjOBJ_MESH: m.nmesh, + mujoco.mjtObj.mjOBJ_HFIELD: m.nhfield, mujoco.mjtObj.mjOBJ_PAIR: m.npair, mujoco.mjtObj.mjOBJ_EQUALITY: m.neq, mujoco.mjtObj.mjOBJ_ACTUATOR: m.nu, @@ -218,6 +219,7 @@ def _getadr( mujoco.mjtObj.mjOBJ_SITE: m.name_siteadr, mujoco.mjtObj.mjOBJ_CAMERA: m.name_camadr, mujoco.mjtObj.mjOBJ_MESH: m.name_meshadr, + mujoco.mjtObj.mjOBJ_HFIELD: m.name_hfieldadr, mujoco.mjtObj.mjOBJ_PAIR: m.name_pairadr, mujoco.mjtObj.mjOBJ_EQUALITY: m.name_eqadr, mujoco.mjtObj.mjOBJ_ACTUATOR: m.name_actuatoradr, diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index caccbea8..0036e58d 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -15,7 +15,7 @@ """Base types used in MJX.""" import enum - +from typing import Tuple import jax import mujoco from mujoco.mjx._src.dataclasses import PyTreeNode # pylint: disable=g-importing-member @@ -112,6 +112,24 @@ class GeomType(enum.IntEnum): # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE +class ConvexMesh(PyTreeNode): + """Geom properties for convex meshes. + + Attributes: + vert: vertices of the convex mesh + face: faces of the convex mesh + face_normal: normal vectors for the faces + edge: edge indexes for all edges in the convex mesh + edge_face_normal: indexes for face normals adjacent to edges in `edge` + """ + + vert: jax.Array + face: jax.Array + face_normal: jax.Array + edge: jax.Array + edge_face_normal: jax.Array + + class ConeType(enum.IntEnum): """Type of friction cone. @@ -322,6 +340,7 @@ class Model(PyTreeNode): nmesh: number of meshes nmeshvert: number of vertices in all meshes nmeshface: number of triangular faces in all meshes + nhfield: number of heightfields nmat: number of materials npair: number of predefined geom pairs nexclude: number of excluded geom pairs @@ -396,6 +415,7 @@ class Model(PyTreeNode): geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) geom_size: geom-specific size parameters (ngeom, 3) geom_rbound: radius of bounding sphere (ngeom,) + geom_rbound_hfield: static rbound for hfield grid bounds (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) @@ -419,6 +439,12 @@ class Model(PyTreeNode): mesh_vert: vertex positions for all meshes (nmeshvert, 3) mesh_face: vertex face data (nmeshface, 3) mesh_graph: convex graph data (nmeshgraph,) + mesh_convex: pre-compiled convex mesh info for MJX (nmesh,) + hfield_size: (x, y, z_top, z_bottom) (nhfield,) + hfield_nrow: number of rows in grid (nhfield,) + hfield_ncol: number of columns in grid (nhfield,) + hfield_adr: address in hfield_data (nhfield,) + hfield_data: elevation data (nhfielddata,) mat_rgba: rgba (nmat, 4) pair_dim: contact dimensionality (npair,) pair_geom1: id of geom1 (npair,) @@ -488,6 +514,7 @@ class Model(PyTreeNode): nmesh: int nmeshvert: int nmeshface: int + nhfield: int nmat: int npair: int nexclude: int @@ -561,6 +588,7 @@ class Model(PyTreeNode): geom_solimp: jax.Array geom_size: jax.Array geom_rbound: jax.Array + geom_rbound_hfield: np.ndarray geom_pos: jax.Array geom_quat: jax.Array geom_friction: jax.Array @@ -584,6 +612,12 @@ class Model(PyTreeNode): mesh_vert: np.ndarray mesh_face: np.ndarray mesh_graph: np.ndarray + mesh_convex: Tuple[ConvexMesh, ...] + hfield_size: np.ndarray + hfield_nrow: np.ndarray + hfield_ncol: np.ndarray + hfield_adr: np.ndarray + hfield_data: jax.Array mat_rgba: np.ndarray pair_dim: np.ndarray pair_geom1: np.ndarray @@ -632,6 +666,7 @@ class Model(PyTreeNode): name_siteadr: np.ndarray name_camadr: np.ndarray name_meshadr: np.ndarray + name_hfieldadr: np.ndarray name_pairadr: np.ndarray name_eqadr: np.ndarray name_actuatoradr: np.ndarray diff --git a/mjx/mujoco/mjx/test_data/barkour_v0/assets/barkour_v0_mjx.xml b/mjx/mujoco/mjx/test_data/barkour_v0/assets/barkour_v0_mjx.xml index a98db27b..82dd41df 100644 --- a/mjx/mujoco/mjx/test_data/barkour_v0/assets/barkour_v0_mjx.xml +++ b/mjx/mujoco/mjx/test_data/barkour_v0/assets/barkour_v0_mjx.xml @@ -77,6 +77,7 @@ + @@ -101,7 +102,8 @@ - + + diff --git a/mjx/mujoco/mjx/test_data/barkour_v0/assets/hfield_240_280.png b/mjx/mujoco/mjx/test_data/barkour_v0/assets/hfield_240_280.png new file mode 100644 index 00000000..97f72d2a Binary files /dev/null and b/mjx/mujoco/mjx/test_data/barkour_v0/assets/hfield_240_280.png differ