From 2386353b02c2239fc968e6d6e04b091c6085e038 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Thu, 14 Mar 2024 16:20:04 -0700 Subject: [PATCH] Throw an error for margin/gap with meshes. PiperOrigin-RevId: 615937263 Change-Id: I3289cf1b1fec93385cc78a01e53866427453db59 --- doc/changelog.rst | 1 + doc/mjx.rst | 20 +++++------ mjx/mujoco/mjx/__init__.py | 1 + mjx/mujoco/mjx/_src/collision_driver.py | 46 +++++++++++++++---------- mjx/mujoco/mjx/_src/device.py | 9 +++-- mjx/mujoco/mjx/_src/io.py | 10 ++++-- mjx/mujoco/mjx/_src/io_test.py | 19 ++++++++++ 7 files changed, 73 insertions(+), 33 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index fdfafeef..e32cd0a5 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -11,6 +11,7 @@ MJX 1. Improved performance of SAT for convex collisions. 2. Fixed bug for sphere/capsule-convex deep penetration. 3. Fixed bug where ``mjx.Data`` produced by ``mjx.put_data`` had different treedef than ``mjx.make_data``. +4. Throw an error for margin/gap for convex mesh collisions, since they are not supported. Version 3.1.3 (March 5th, 2024) ----------------------------------- diff --git a/doc/mjx.rst b/doc/mjx.rst index 3d1d42d2..330336ee 100644 --- a/doc/mjx.rst +++ b/doc/mjx.rst @@ -215,6 +215,14 @@ The following features are **in development** and coming soon: * - Category - Feature + * - :ref:`Geom ` + - ``SDF``, ``HFIELD``, ``ELLIPSOID``, ``CYLINDER`` + * - :ref:`Condim ` + - 1, 4, 6 + * - :ref:`Constraint ` + - :ref:`Frictionloss `, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` + * - :ref:`Integrator ` + - ``IMPLICIT``, ``IMPLICITFAST`` * - Dynamics - :ref:`Inverse ` * - :ref:`Transmission ` @@ -227,16 +235,8 @@ The following features are **in development** and coming soon: - ``MUSCLE`` * - :ref:`Tendon Wrapping ` - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER`` - * - :ref:`Geom ` - - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER`` - * - :ref:`Constraint ` - - :ref:`Frictionloss `, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF`` - * - :ref:`Integrator ` - - ``IMPLICIT``, ``IMPLICITFAST`` * - :ref:`Cone ` - ``ELLIPTIC`` - * - :ref:`Condim ` - - 1, 4, 6 * - Fluid Model - :ref:`flEllipsoid` * - :ref:`Tendons ` @@ -258,6 +258,8 @@ The following features are **unsupported**: * - Category - Feature + * - :ref:`margin` and :ref:`gap` + - Unimplemented for collisions with ``Mesh`` :ref:`Geom `. * - :ref:`Transmission ` - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY`` * - :ref:`Actuator Dynamics ` @@ -270,8 +272,6 @@ The following features are **unsupported**: - ``PGS`` * - :ref:`Sensors ` - ``PLUGIN``, ``USER`` - * - :ref:`Geom ` - - ``SDF`` .. _MjxSharpBits: diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py index 472356b0..e21b14fd 100644 --- a/mjx/mujoco/mjx/__init__.py +++ b/mjx/mujoco/mjx/__init__.py @@ -16,6 +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 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py index d492d69b..3b70990c 100644 --- a/mjx/mujoco/mjx/_src/collision_driver.py +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -14,7 +14,7 @@ # ============================================================================== """Collide geometries.""" -from typing import Callable, Dict, Optional, Sequence, Tuple, Union +from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union import jax from jax import numpy as jp @@ -187,6 +187,31 @@ def _dynamic_params( 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]]]]: @@ -276,24 +301,7 @@ def _collide_geoms( if not fn: return Contact.zero() - # 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) + geom1, geom2, params = get_params(m, candidates) g1, g2, in_axes = _pair_info(m, d, geom1, geom2) # Run a crude version of broadphase. diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py index 0d094210..b920b2e6 100644 --- a/mjx/mujoco/mjx/_src/device.py +++ b/mjx/mujoco/mjx/_src/device.py @@ -138,8 +138,7 @@ def _validate(m: mujoco.MjModel): raise NotImplementedError('Tendons are not supported.') # check collision geom types - candidate_set = collision_driver.collision_candidates(m) - for g1, g2, *_ in candidate_set: + 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, @@ -149,6 +148,12 @@ def _validate(m: mujoco.MjModel): 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. diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index ddcd8907..1958829f 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -91,10 +91,16 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model: raise NotImplementedError('only condim=3 is supported') # check collision geom types - for g1, g2, *_ in collision_driver.collision_candidates(m): + 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: - g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) 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 enum_field, enum_type, mj_type in ( (m.actuator_biastype, types.BiasType, mujoco.mjtBias), diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index e3f414df..83147503 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -209,6 +209,25 @@ class ModelIOTest(parameterized.TestCase): """)) + def test_margin_gap_mesh_not_implemented(self): + with self.assertRaises(NotImplementedError): + mjx.put_model(mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + + + """)) + class DataIOTest(parameterized.TestCase): """IO tests for mjx.Data."""