diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
index aff524f6..82e53459 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
@@ -21,7 +21,6 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Model as Model
from mujoco.mjx.third_party.mujoco_warp._src.types import Data as Data
# isort: on
-from ._src import test_util as test_util # used by viewer and testspeed, not meant for public consumption
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import collision as collision
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import nxn_broadphase as nxn_broadphase
from mujoco.mjx.third_party.mujoco_warp._src.collision_driver import sap_broadphase as sap_broadphase
@@ -68,6 +67,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.solver import solve as solve
from mujoco.mjx.third_party.mujoco_warp._src.support import contact_force as contact_force
from mujoco.mjx.third_party.mujoco_warp._src.support import mul_m as mul_m
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate as xfrc_accumulate
+from mujoco.mjx.third_party.mujoco_warp._src.types import BiasType as BiasType
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter as BroadphaseFilter
from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType as BroadphaseType
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType as ConeType
@@ -76,6 +76,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Contact as Contact
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit as DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import DynType as DynType
from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit as EnableBit
+from mujoco.mjx.third_party.mujoco_warp._src.types import GainType as GainType
+from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType as GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType as IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import JointType as JointType
from mujoco.mjx.third_party.mujoco_warp._src.types import Option as Option
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
similarity index 58%
rename from mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py
rename to mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
index 72a16659..d0bc902c 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/test_util.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/benchmark.py
@@ -13,7 +13,7 @@
# limitations under the License.
# ==============================================================================
-"""Utilities for testing."""
+"""Utilities for benchmarking MuJoCo Warp."""
import importlib
import os
@@ -23,183 +23,15 @@ from typing import Callable, Optional, Tuple
import mujoco
import numpy as np
import warp as wp
-from etils import epath
from mujoco.mjx.third_party.mujoco_warp._src import forward
from mujoco.mjx.third_party.mujoco_warp._src import io
from mujoco.mjx.third_party.mujoco_warp._src import warp_util
-from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
-from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
-from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
-from mujoco.mjx.third_party.mujoco_warp._src.types import EnableBit
-from mujoco.mjx.third_party.mujoco_warp._src.types import IntegratorType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
-from mujoco.mjx.third_party.mujoco_warp._src.types import SolverType
from mujoco.mjx.third_party.mujoco_warp._src.util_misc import halton
-def fixture(
- fname: Optional[str] = None,
- xml: Optional[str] = None,
- keyframe: int = -1,
- actuation: bool = True,
- contact: bool = True,
- constraint: bool = True,
- equality: bool = True,
- spring: bool = True,
- damper: bool = True,
- gravity: bool = True,
- clampctrl: bool = True,
- filterparent: bool = True,
- qpos0: bool = False,
- kick: bool = False,
- energy: bool = False,
- eulerdamp: Optional[bool] = None,
- cone: Optional[ConeType] = None,
- integrator: Optional[IntegratorType] = None,
- solver: Optional[SolverType] = None,
- iterations: Optional[int] = None,
- ls_iterations: Optional[int] = None,
- ls_parallel: Optional[bool] = None,
- sparse: Optional[bool] = None,
- broadphase: Optional[BroadphaseType] = None,
- disableflags: Optional[int] = None,
- enableflags: Optional[int] = None,
- applied: bool = False,
- nstep: int = 3,
- seed: int = 42,
- nworld: int = None,
- nconmax: int = None,
- njmax: int = None,
-):
- np.random.seed(seed)
- if fname is not None:
- path = epath.resource_path("mjx") / "third_party/mujoco_warp" / "test_data" / fname
- mjm = mujoco.MjModel.from_xml_path(path.as_posix())
- elif xml is not None:
- mjm = mujoco.MjModel.from_xml_string(xml)
- else:
- raise ValueError("either fname or xml must be provided")
-
- if not actuation:
- mjm.opt.disableflags |= DisableBit.ACTUATION
- if not contact:
- mjm.opt.disableflags |= DisableBit.CONTACT
- if not constraint:
- mjm.opt.disableflags |= DisableBit.CONSTRAINT
- if not equality:
- mjm.opt.disableflags |= DisableBit.EQUALITY
- if not spring:
- mjm.opt.disableflags |= DisableBit.SPRING
- if not damper:
- mjm.opt.disableflags |= DisableBit.DAMPER
- if not gravity:
- mjm.opt.disableflags |= DisableBit.GRAVITY
- if not clampctrl:
- mjm.opt.disableflags |= DisableBit.CLAMPCTRL
- if not eulerdamp:
- mjm.opt.disableflags |= DisableBit.EULERDAMP
- if not filterparent:
- mjm.opt.disableflags |= DisableBit.FILTERPARENT
-
- if energy:
- mjm.opt.enableflags |= EnableBit.ENERGY
-
- if cone is not None:
- mjm.opt.cone = cone
- if integrator is not None:
- mjm.opt.integrator = integrator
- if disableflags is not None:
- mjm.opt.disableflags |= disableflags
- if enableflags is not None:
- mjm.opt.enableflags |= enableflags
- if solver is not None:
- mjm.opt.solver = solver
- if iterations is not None:
- mjm.opt.iterations = iterations
- if ls_iterations is not None:
- mjm.opt.ls_iterations = ls_iterations
- if sparse is not None:
- if sparse:
- mjm.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
- else:
- mjm.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
-
- mjd = mujoco.MjData(mjm)
- if keyframe > -1:
- mujoco.mj_resetDataKeyframe(mjm, mjd, keyframe)
- elif qpos0:
- mjd.qpos[:] = mjm.qpos0
- else:
- # set random qpos, underlying code should gracefully handle un-normalized quats
- mjd.qpos[:] = np.random.random(mjm.nq)
-
- if kick:
- # give the system a little kick to ensure we have non-identity rotations
- mjd.qvel = np.random.uniform(-0.01, 0.01, mjm.nv)
- mjd.ctrl = np.random.uniform(-0.1, 0.1, size=mjm.nu)
- if applied:
- mjd.qfrc_applied = np.random.uniform(-0.1, 0.1, size=mjm.nv)
- mjd.xfrc_applied = np.random.uniform(-0.1, 0.1, size=mjd.xfrc_applied.shape)
- if kick or applied:
- mujoco.mj_step(mjm, mjd, nstep) # let dynamics get state significantly non-zero
-
- if mjm.nmocap:
- mjd.mocap_pos = np.random.random(mjd.mocap_pos.shape)
- mocap_quat = np.random.random(mjd.mocap_quat.shape)
- mjd.mocap_quat = mocap_quat
-
- mujoco.mj_forward(mjm, mjd)
- mjd.qacc_warmstart = mjd.qacc
- m = io.put_model(mjm)
- if ls_parallel is not None:
- m.opt.ls_parallel = ls_parallel
- if broadphase is not None:
- m.opt.broadphase = broadphase
-
- d = io.put_data(mjm, mjd, nworld=nworld, nconmax=nconmax, njmax=njmax)
- return mjm, mjd, m, d
-
-
-def find_keys(model: mujoco.MjModel, keyname_prefix: str) -> list[int]:
- """Finds keyframes that start with keyname_prefix."""
- keys = []
-
- for keyid in range(model.nkey):
- name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_KEY, keyid)
- if name.startswith(keyname_prefix):
- keys.append(keyid)
-
- return keys
-
-
-def make_trajectory(model: mujoco.MjModel, keys: list[int]) -> np.ndarray:
- """Make a ctrl trajectory with linear interpolation."""
- ctrls = []
- prev_ctrl_key = np.zeros(model.nu, dtype=np.float64)
- prev_time, time = 0.0, 0.0
-
- for key in keys:
- ctrl_key, ctrl_time = model.key_ctrl[key], model.key_time[key]
- if not ctrls and ctrl_time != 0.0:
- raise ValueError("first keyframe must have time 0.0")
- elif ctrls and ctrl_time <= prev_time:
- raise ValueError("keyframes must be in time order")
-
- while time < ctrl_time:
- frac = (time - prev_time) / (ctrl_time - prev_time)
- ctrls.append(prev_ctrl_key * (1 - frac) + ctrl_key * frac)
- time += model.opt.timestep
-
- ctrls.append(ctrl_key)
- time += model.opt.timestep
- prev_ctrl_key = ctrl_key
- prev_time = time
-
- return np.array(ctrls)
-
-
def _sum(stack1, stack2):
ret = {}
for k in stack1:
@@ -372,10 +204,10 @@ class BenchmarkSuite:
ctrls = None
if self.replay:
- keys = find_keys(mjm, self.replay)
+ keys = io.find_keys(mjm, self.replay)
if not keys:
raise ValueError(f"Key prefix not find: {self.replay}")
- ctrls = make_trajectory(mjm, keys)
+ ctrls = io.make_trajectory(mjm, keys)
mujoco.mj_resetDataKeyframe(mjm, mjd, keys[0])
if mjm.nkey > 0:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/broadphase_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/broadphase_test.py
index f1555cc5..1568f156 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/broadphase_test.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/broadphase_test.py
@@ -20,16 +20,17 @@ import warp as wp
from absl.testing import absltest
from absl.testing import parameterized
-import mujoco_warp as mjwarp
+import mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp import BroadphaseFilter
+from mujoco.mjx.third_party.mujoco_warp import BroadphaseType
+from mujoco.mjx.third_party.mujoco_warp import DisableBit
+from mujoco.mjx.third_party.mujoco_warp import test_data
from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
-from mujoco.mjx.third_party.mujoco_warp._src import test_util
-from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseFilter
-from mujoco.mjx.third_party.mujoco_warp._src.types import BroadphaseType
def broadphase_caller(m, d):
- if m.opt.broadphase == int(BroadphaseType.NXN):
+ if m.opt.broadphase == BroadphaseType.NXN:
collision_driver.nxn_broadphase(m, d)
else:
collision_driver.sap_broadphase(m, d)
@@ -37,12 +38,12 @@ def broadphase_caller(m, d):
class BroadphaseTest(parameterized.TestCase):
# filter combinations
- plane_sphere = BroadphaseFilter.PLANE.value | BroadphaseFilter.SPHERE.value
- plane_aabb = BroadphaseFilter.PLANE.value | BroadphaseFilter.AABB.value
- plane_obb = BroadphaseFilter.PLANE | BroadphaseFilter.OBB.value
- plane_sphere_aabb = plane_sphere | BroadphaseFilter.AABB.value
- plane_sphere_obb = plane_sphere | BroadphaseFilter.OBB.value
- plane_sphere_aabb_obb = plane_sphere_aabb | BroadphaseFilter.OBB.value
+ plane_sphere = BroadphaseFilter.PLANE | BroadphaseFilter.SPHERE
+ plane_aabb = BroadphaseFilter.PLANE | BroadphaseFilter.AABB
+ plane_obb = BroadphaseFilter.PLANE | BroadphaseFilter.OBB
+ plane_sphere_aabb = plane_sphere | BroadphaseFilter.AABB
+ plane_sphere_obb = plane_sphere | BroadphaseFilter.OBB
+ plane_sphere_aabb_obb = plane_sphere_aabb | BroadphaseFilter.OBB
@parameterized.product(
broadphase=list(BroadphaseType),
@@ -112,7 +113,7 @@ class BroadphaseTest(parameterized.TestCase):
"""
# one world and zero collisions
- mjm, _, m, d0 = test_util.fixture(xml=_XML, keyframe=0)
+ mjm, _, m, d0 = test_data.fixture(xml=_XML, keyframe=0)
m.opt.broadphase = broadphase
m.opt.broadphase_filter = filter
@@ -121,7 +122,7 @@ class BroadphaseTest(parameterized.TestCase):
np.testing.assert_allclose(d0.ncollision.numpy()[0], 0)
# one world and one collision
- _, mjd1, _, d1 = test_util.fixture(xml=_XML, keyframe=1)
+ _, mjd1, _, d1 = test_data.fixture(xml=_XML, keyframe=1)
broadphase_caller(m, d1)
np.testing.assert_allclose(d1.ncollision.numpy()[0], 1)
@@ -129,7 +130,7 @@ class BroadphaseTest(parameterized.TestCase):
np.testing.assert_allclose(d1.collision_pair.numpy()[0][1], 1)
# one world and three collisions
- _, mjd2, _, d2 = test_util.fixture(xml=_XML, keyframe=2)
+ _, mjd2, _, d2 = test_data.fixture(xml=_XML, keyframe=2)
broadphase_caller(m, d2)
ncollision = d2.ncollision.numpy()[0]
@@ -140,7 +141,7 @@ class BroadphaseTest(parameterized.TestCase):
self.assertTrue([d2.collision_pair.numpy()[i][0], d2.collision_pair.numpy()[i][1]] in collision_pairs)
# two worlds and four collisions
- d3 = mjwarp.make_data(mjm, nworld=2, nconmax=512, njmax=512)
+ d3 = mjw.make_data(mjm, nworld=2, nconmax=512, njmax=512)
d3.geom_xpos = wp.array(
np.vstack([np.expand_dims(mjd1.geom_xpos, axis=0), np.expand_dims(mjd2.geom_xpos, axis=0)]),
dtype=wp.vec3,
@@ -162,15 +163,15 @@ class BroadphaseTest(parameterized.TestCase):
self.assertTrue([d3.collision_pair.numpy()[i][0], d3.collision_pair.numpy()[i][1]] in collision_pairs[worldid])
# one world and zero collisions: contype and conaffinity incompatibility
- mjm4, _, m4, d4 = test_util.fixture(xml=_XML, keyframe=1)
+ mjm4, _, m4, d4 = test_data.fixture(xml=_XML, keyframe=1)
mjm4.geom_contype[:3] = 0
- m4 = mjwarp.put_model(mjm4)
+ m4 = mjw.put_model(mjm4)
broadphase_caller(m4, d4)
np.testing.assert_allclose(d4.ncollision.numpy()[0], 0)
# one world and one collision: geomtype ordering
- _, _, _, d5 = test_util.fixture(xml=_XML, keyframe=3)
+ _, _, _, d5 = test_data.fixture(xml=_XML, keyframe=3)
broadphase_caller(m, d5)
np.testing.assert_allclose(d5.ncollision.numpy()[0], 1)
np.testing.assert_allclose(d5.collision_pair.numpy()[0][0], 3)
@@ -195,12 +196,12 @@ class BroadphaseTest(parameterized.TestCase):
"""
- _, _, m, d = test_util.fixture(xml=_MJCF, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_MJCF, keyframe=0)
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], ncollision)
- @parameterized.parameters(True, False)
- def test_broadphase_filterparent(self, filterparent):
+ @parameterized.parameters((0, 0), (DisableBit.FILTERPARENT, 1))
+ def test_broadphase_filterparent(self, disablebit, expected_collisions):
_MJCF = """
@@ -218,16 +219,16 @@ class BroadphaseTest(parameterized.TestCase):
"""
- _, _, m, d = test_util.fixture(xml=_MJCF, filterparent=filterparent, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_MJCF, keyframe=0, overrides={"opt.disableflags": disablebit})
broadphase_caller(m, d)
- self.assertEqual(d.ncollision.numpy()[0], 0 if filterparent else 1)
+ self.assertEqual(d.ncollision.numpy()[0], expected_collisions)
def test_broadphase_filter(self):
- plane = BroadphaseFilter.PLANE.value
- sphere = BroadphaseFilter.SPHERE.value
- aabb = BroadphaseFilter.AABB.value
- obb = BroadphaseFilter.OBB.value
+ plane = BroadphaseFilter.PLANE
+ sphere = BroadphaseFilter.SPHERE
+ aabb = BroadphaseFilter.AABB
+ obb = BroadphaseFilter.OBB
plane_sphere = plane | sphere
plane_aabb = plane | aabb
plane_obb = plane | obb
@@ -260,68 +261,68 @@ class BroadphaseTest(parameterized.TestCase):
"""
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
m.opt.broadphase_filter = plane_sphere
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
m.opt.broadphase_filter = plane_aabb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=0)
m.opt.broadphase_filter = plane_obb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
m.opt.broadphase_filter = plane_sphere
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
# note: collision_driver._plane_filter checks bounding sphere
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
m.opt.broadphase_filter = plane
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 2)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
m.opt.broadphase_filter = plane_sphere
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=1)
m.opt.broadphase_filter = plane_obb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
m.opt.broadphase_filter = plane_sphere
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
m.opt.broadphase_filter = plane_aabb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=2)
m.opt.broadphase_filter = plane_obb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
m.opt.broadphase_filter = plane_sphere
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
m.opt.broadphase_filter = plane_aabb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 1)
- _, _, m, d = test_util.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
+ _, _, m, d = test_data.fixture(xml=_PLANE_CAPSULE_CAPSULE, keyframe=3)
m.opt.broadphase_filter = plane_obb
broadphase_caller(m, d)
self.assertEqual(d.ncollision.numpy()[0], 0)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
index 7d0580ff..77ba312b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
@@ -43,33 +43,33 @@ mat3c = wp.types.matrix(shape=(MULTI_CONTACT_COUNT, 3), dtype=float)
mat63 = wp.types.matrix(shape=(6, 3), dtype=float)
_CONVEX_COLLISION_PAIRS = [
- (GeomType.HFIELD.value, GeomType.SPHERE.value),
- (GeomType.HFIELD.value, GeomType.CAPSULE.value),
- (GeomType.HFIELD.value, GeomType.ELLIPSOID.value),
- (GeomType.HFIELD.value, GeomType.CYLINDER.value),
- (GeomType.HFIELD.value, GeomType.BOX.value),
- (GeomType.HFIELD.value, GeomType.MESH.value),
- (GeomType.SPHERE.value, GeomType.ELLIPSOID.value),
- (GeomType.SPHERE.value, GeomType.MESH.value),
- (GeomType.CAPSULE.value, GeomType.ELLIPSOID.value),
- (GeomType.CAPSULE.value, GeomType.CYLINDER.value),
- (GeomType.CAPSULE.value, GeomType.MESH.value),
- (GeomType.ELLIPSOID.value, GeomType.ELLIPSOID.value),
- (GeomType.ELLIPSOID.value, GeomType.CYLINDER.value),
- (GeomType.ELLIPSOID.value, GeomType.BOX.value),
- (GeomType.ELLIPSOID.value, GeomType.MESH.value),
- (GeomType.CYLINDER.value, GeomType.CYLINDER.value),
- (GeomType.CYLINDER.value, GeomType.BOX.value),
- (GeomType.CYLINDER.value, GeomType.MESH.value),
- (GeomType.BOX.value, GeomType.MESH.value),
- (GeomType.MESH.value, GeomType.MESH.value),
+ (GeomType.HFIELD, GeomType.SPHERE),
+ (GeomType.HFIELD, GeomType.CAPSULE),
+ (GeomType.HFIELD, GeomType.ELLIPSOID),
+ (GeomType.HFIELD, GeomType.CYLINDER),
+ (GeomType.HFIELD, GeomType.BOX),
+ (GeomType.HFIELD, GeomType.MESH),
+ (GeomType.SPHERE, GeomType.ELLIPSOID),
+ (GeomType.SPHERE, GeomType.MESH),
+ (GeomType.CAPSULE, GeomType.ELLIPSOID),
+ (GeomType.CAPSULE, GeomType.CYLINDER),
+ (GeomType.CAPSULE, GeomType.MESH),
+ (GeomType.ELLIPSOID, GeomType.ELLIPSOID),
+ (GeomType.ELLIPSOID, GeomType.CYLINDER),
+ (GeomType.ELLIPSOID, GeomType.BOX),
+ (GeomType.ELLIPSOID, GeomType.MESH),
+ (GeomType.CYLINDER, GeomType.CYLINDER),
+ (GeomType.CYLINDER, GeomType.BOX),
+ (GeomType.CYLINDER, GeomType.MESH),
+ (GeomType.BOX, GeomType.MESH),
+ (GeomType.MESH, GeomType.MESH),
]
def _check_convex_collision_pairs():
prev_idx = -1
for pair in _CONVEX_COLLISION_PAIRS:
- idx = upper_trid_index(len(GeomType), pair[0], pair[1])
+ idx = upper_trid_index(len(GeomType), pair[0].value, pair[1].value)
if pair[1] < pair[0] or idx <= prev_idx:
return False
prev_idx = idx
@@ -84,8 +84,7 @@ def ccd_kernel_builder(
legacy_gjk: bool,
geomtype1: int,
geomtype2: int,
- gjk_iterations: int,
- epa_iterations: int,
+ ccd_iterations: int,
epa_exact_neg_distance: bool,
depth_extension: float,
is_hfield: bool,
@@ -152,7 +151,7 @@ def ccd_kernel_builder(
# TODO(kbayes): remove legacy GJK once multicontact can be enabled
if wp.static(legacy_gjk):
simplex, normal = gjk_legacy(
- gjk_iterations,
+ ccd_iterations,
geom1,
geom2,
geomtype1,
@@ -160,14 +159,14 @@ def ccd_kernel_builder(
)
depth, normal = epa_legacy(
- epa_iterations, geom1, geom2, geomtype1, geomtype2, depth_extension, epa_exact_neg_distance, simplex, normal
+ ccd_iterations, geom1, geom2, geomtype1, geomtype2, depth_extension, epa_exact_neg_distance, simplex, normal
)
dist = -depth
if dist >= 0.0 or depth < -depth_extension:
return 0
- sphere = int(GeomType.SPHERE.value)
- ellipsoid = int(GeomType.ELLIPSOID.value)
+ sphere = GeomType.SPHERE
+ ellipsoid = GeomType.ELLIPSOID
g1 = geoms[0]
g2 = geoms[1]
if geom_type[g1] == sphere or geom_type[g1] == ellipsoid or geom_type[g2] == sphere or geom_type[g2] == ellipsoid:
@@ -177,12 +176,13 @@ def ccd_kernel_builder(
frame = make_frame(normal)
else:
points = mat3c()
+ geom1.margin = margin
+ geom2.margin = margin
dist, ncontact, witness1, witness2 = ccd(
- False,
+ False, # ignored for box-box, multiccd always on
opt_ccd_tolerance[worldid],
0.0,
- gjk_iterations,
- epa_iterations,
+ ccd_iterations,
geom1,
geom2,
geomtype1,
@@ -389,13 +389,13 @@ def ccd_kernel_builder(
geomtype1,
geom1_dataid,
geom_size[worldid, g1],
- mesh_vertadr[geom1_dataid],
- mesh_vertnum[geom1_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom1_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom1_dataid],
- mesh_polyadr[geom1_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
@@ -412,13 +412,13 @@ def ccd_kernel_builder(
geomtype2,
geom2_dataid,
geom_size[worldid, g2],
- mesh_vertadr[geom2_dataid],
- mesh_vertnum[geom2_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom2_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom2_dataid],
- mesh_polyadr[geom2_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
@@ -633,13 +633,11 @@ def convex_narrowphase(m: Model, d: Data):
computations for non-existent pair types.
"""
for geom_pair in _CONVEX_COLLISION_PAIRS:
- g1 = geom_pair[0]
- g2 = geom_pair[1]
+ g1 = geom_pair[0].value
+ g2 = geom_pair[1].value
if m.geom_pair_type_count[upper_trid_index(len(GeomType), g1, g2)]:
wp.launch(
- ccd_kernel_builder(
- m.opt.legacy_gjk, g1, g2, m.opt.gjk_iterations, m.opt.epa_iterations, True, 1e9, g1 == int(GeomType.HFIELD.value)
- ),
+ ccd_kernel_builder(m.opt.legacy_gjk, g1, g2, m.opt.ccd_iterations, True, 1e9, g1 == GeomType.HFIELD),
dim=d.nconmax,
inputs=[
m.opt.ccd_tolerance,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
index 2973c99b..1906baef 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
@@ -233,6 +233,7 @@ def _obb_filter(
return True
+@cache_kernel
def _broadphase_filter(opt_broadphase_filter: int):
@wp.func
def func(
@@ -267,16 +268,16 @@ def _broadphase_filter(opt_broadphase_filter: int):
xmat2 = geom_xmat_in[worldid, geom2]
if rbound1 == 0.0 or rbound2 == 0.0:
- if wp.static(opt_broadphase_filter & int(BroadphaseFilter.PLANE.value)):
+ if wp.static(opt_broadphase_filter & BroadphaseFilter.PLANE):
return _plane_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2, xmat1, xmat2)
else:
- if wp.static(opt_broadphase_filter & int(BroadphaseFilter.SPHERE.value)):
+ if wp.static(opt_broadphase_filter & BroadphaseFilter.SPHERE):
if not _sphere_filter(rbound1, rbound2, margin1, margin2, xpos1, xpos2):
return False
- if wp.static(opt_broadphase_filter & int(BroadphaseFilter.AABB.value)):
+ if wp.static(opt_broadphase_filter & BroadphaseFilter.AABB):
if not _aabb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
- if wp.static(opt_broadphase_filter & int(BroadphaseFilter.OBB.value)):
+ if wp.static(opt_broadphase_filter & BroadphaseFilter.OBB):
if not _obb_filter(center1, center2, size1, size2, margin1, margin2, xpos1, xpos2, xmat1, xmat2):
return False
@@ -533,7 +534,7 @@ def sap_broadphase(m: Model, d: Data):
],
)
- if m.opt.broadphase == int(BroadphaseType.SAP_TILE):
+ if m.opt.broadphase == BroadphaseType.SAP_TILE:
wp.launch_tiled(
kernel=_segmented_sort(m.ngeom),
dim=(d.nworld),
@@ -715,7 +716,7 @@ def collision(m: Model, d: Data):
if d.nconmax == 0 or m.opt.disableflags & (DisableBit.CONSTRAINT | DisableBit.CONTACT):
return
- if m.opt.broadphase == int(BroadphaseType.NXN):
+ if m.opt.broadphase == BroadphaseType.NXN:
nxn_broadphase(m, d)
else:
sap_broadphase(m, d)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py
index 2003c8a5..d3087f84 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver_test.py
@@ -16,33 +16,16 @@
import mujoco
import numpy as np
-import warp as wp
from absl.testing import absltest
from absl.testing import parameterized
-import mujoco_warp as mjwarp
+import mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp import BroadphaseType
+from mujoco.mjx.third_party.mujoco_warp import DisableBit
+from mujoco.mjx.third_party.mujoco_warp import test_data
from mujoco.mjx.third_party.mujoco_warp.test_data.collision_sdf.utils import register_sdf_plugins
-from mujoco.mjx.third_party.mujoco_warp._src import test_util
from mujoco.mjx.third_party.mujoco_warp._src import types
-from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import VolumeData
-from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sample_volume_sdf
-
-
-@wp.kernel
-def sample_sdf_kernel(
- # In:
- points: wp.array(dtype=wp.vec3),
- volume_data: VolumeData,
- # Out:
- results_out: wp.array(dtype=float),
-):
- """Kernel to sample SDF values at given points using Warp volume."""
- tid = wp.tid()
- point = points[tid]
- sdf_value = sample_volume_sdf(point, volume_data)
- results_out[tid] = sdf_value
-
_TOLERANCE = 5e-5
@@ -501,16 +484,16 @@ class CollisionTest(parameterized.TestCase):
@classmethod
def setUpClass(cls):
- register_sdf_plugins(mjwarp)
+ register_sdf_plugins(mjw)
@parameterized.parameters(_SDF_SDF.keys())
def test_sdf_collision(self, fixture):
"""Tests collisions with different geometries."""
- mjm, mjd, m, d = test_util.fixture(xml=self._SDF_SDF[fixture], qpos0=True)
+ mjm, mjd, m, d = test_data.fixture(xml=self._SDF_SDF[fixture])
mujoco.mj_collision(mjm, mjd)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
for i in range(min(mjd.ncon, d.ncon.numpy()[0])):
actual_dist = mjd.contact.dist[i]
actual_pos = mjd.contact.pos[i]
@@ -528,14 +511,14 @@ class CollisionTest(parameterized.TestCase):
@parameterized.parameters(_FIXTURES.keys())
def test_collision(self, fixture):
"""Tests collisions with different geometries."""
- mjm, mjd, m, d = test_util.fixture(xml=self._FIXTURES[fixture], qpos0=True)
+ mjm, mjd, m, d = test_data.fixture(xml=self._FIXTURES[fixture])
# Exempt GJK collisions from exact contact count check
# because GJK generates more contacts
allow_different_contact_count = False
mujoco.mj_collision(mjm, mjd)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertGreater(d.ncon.numpy()[0], 0)
self.assertGreater(mjd.ncon, 0)
@@ -582,16 +565,16 @@ class CollisionTest(parameterized.TestCase):
@parameterized.parameters(_HFIELD_FIXTURES.keys())
def test_hfield_collision(self, fixture):
"""Tests hfield collision with different geometries."""
- mjm, mjd, m, d = test_util.fixture(xml=self._HFIELD_FIXTURES[fixture])
+ mjm, mjd, m, d = test_data.fixture(xml=self._HFIELD_FIXTURES[fixture])
mujoco.mj_collision(mjm, mjd)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(mjd.ncon > 0, d.ncon.numpy()[0] > 0, "If MJ collides, MJW should too")
def test_contact_exclude(self):
"""Tests contact exclude."""
- _, _, m, _ = test_util.fixture(
+ _, _, m, _ = test_data.fixture(
xml="""
@@ -617,11 +600,11 @@ class CollisionTest(parameterized.TestCase):
self.assertEqual(m.nxn_geom_pair.numpy().shape[0], 3)
np.testing.assert_equal(m.nxn_pairid.numpy(), np.array([-2, -1, -1]))
- @parameterized.parameters(list(types.BroadphaseType))
+ @parameterized.parameters(list(BroadphaseType))
def test_contact_pair(self, broadphase):
"""Tests contact pair."""
# no pairs
- _, _, m, _ = test_util.fixture(
+ _, _, m, _ = test_data.fixture(
xml="""
@@ -632,12 +615,12 @@ class CollisionTest(parameterized.TestCase):
""",
- broadphase=broadphase,
+ overrides={"opt.broadphase": broadphase},
)
self.assertTrue((m.nxn_pairid.numpy() == -1).all())
# 1 pair
- _, _, m, d = test_util.fixture(
+ _, _, m, d = test_data.fixture(
xml="""
@@ -654,8 +637,7 @@ class CollisionTest(parameterized.TestCase):
- """,
- qpos0=True,
+ """
)
self.assertTrue((m.nxn_pairid.numpy() == 0).all())
@@ -670,7 +652,7 @@ class CollisionTest(parameterized.TestCase):
):
arr.zero_()
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(d.ncon.numpy()[0], 1)
self.assertEqual(d.contact.includemargin.numpy()[0], -1)
@@ -681,7 +663,7 @@ class CollisionTest(parameterized.TestCase):
np.testing.assert_allclose(d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5]))
# 1 pair: override contype and conaffinity
- _, _, m, d = test_util.fixture(
+ _, _, m, d = test_data.fixture(
xml="""
@@ -698,8 +680,7 @@ class CollisionTest(parameterized.TestCase):
- """,
- qpos0=True,
+ """
)
self.assertTrue((m.nxn_pairid.numpy() == 0).all())
@@ -714,7 +695,7 @@ class CollisionTest(parameterized.TestCase):
):
arr.zero_()
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(d.ncon.numpy()[0], 1)
self.assertEqual(d.contact.includemargin.numpy()[0], -1)
@@ -725,7 +706,7 @@ class CollisionTest(parameterized.TestCase):
np.testing.assert_allclose(d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5]))
# 1 pair: override exclude
- _, _, m, d = test_util.fixture(
+ _, _, m, d = test_data.fixture(
xml="""
@@ -743,8 +724,7 @@ class CollisionTest(parameterized.TestCase):
- """,
- qpos0=True,
+ """
)
self.assertTrue((m.nxn_pairid.numpy() == 0).all())
@@ -759,7 +739,7 @@ class CollisionTest(parameterized.TestCase):
):
arr.zero_()
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(d.ncon.numpy()[0], 1)
self.assertEqual(d.contact.includemargin.numpy()[0], -1)
@@ -770,7 +750,7 @@ class CollisionTest(parameterized.TestCase):
np.testing.assert_allclose(d.contact.solimp.numpy()[0], np.array([0.1, 0.2, 0.3, 0.4, 0.5]))
# 1 pair 1 exclude
- _, _, m, d = test_util.fixture(
+ _, _, m, d = test_data.fixture(
xml="""
@@ -792,8 +772,7 @@ class CollisionTest(parameterized.TestCase):
- """,
- qpos0=True,
+ """
)
np.testing.assert_equal(m.nxn_pairid.numpy(), np.array([-2, -1, 0]))
@@ -808,7 +787,7 @@ class CollisionTest(parameterized.TestCase):
):
arr.zero_()
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(d.ncon.numpy()[0], 2)
self.assertEqual(d.contact.includemargin.numpy()[1], -1)
@@ -818,29 +797,20 @@ class CollisionTest(parameterized.TestCase):
np.testing.assert_allclose(d.contact.solreffriction.numpy()[1], np.array([2.0, 4.0]))
np.testing.assert_allclose(d.contact.solimp.numpy()[1], np.array([0.1, 0.2, 0.3, 0.4, 0.5]))
- @parameterized.parameters(
- (True, True),
- (True, False),
- (False, True),
- (False, False),
- )
+ @parameterized.product(constraint=(0, DisableBit.CONSTRAINT), contact=(0, DisableBit.CONTACT))
def test_collision_disableflags(self, constraint, contact):
"""Tests collision disableflags."""
- mjm, mjd, m, d = test_util.fixture(
- "humanoid/humanoid.xml",
- keyframe=0,
- constraint=constraint,
- contact=contact,
- kick=False,
+ mjm, mjd, m, d = test_data.fixture(
+ "humanoid/humanoid.xml", keyframe=0, overrides={"opt.disableflags": constraint | contact}
)
mujoco.mj_collision(mjm, mjd)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
self.assertEqual(d.ncon.numpy()[0], mjd.ncon)
def test_hfield_maxconpair(self):
- _XML = f"""
+ _XML = """
@@ -858,14 +828,14 @@ class CollisionTest(parameterized.TestCase):
"""
- _, _, m, d = test_util.fixture(xml=_XML, keyframe=0)
+ _, _, m, d = test_data.fixture(xml=_XML, keyframe=0)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
np.testing.assert_equal(d.ncon.numpy()[0], types.MJ_MAXCONPAIR)
def test_min_friction(self):
- _, _, _, d = test_util.fixture(
+ _, _, _, d = test_data.fixture(
xml="""
@@ -891,7 +861,7 @@ class CollisionTest(parameterized.TestCase):
@parameterized.parameters(("1", "1"), ("1", "2"), ("2", "1"))
def test_contact_parameter_mixing(self, priority1, priority2):
- _, mjd, m, d = test_util.fixture(
+ _, mjd, m, d = test_data.fixture(
xml=f"""
@@ -909,7 +879,7 @@ class CollisionTest(parameterized.TestCase):
keyframe=0,
)
- mjwarp.collision(m, d)
+ mjw.collision(m, d)
ncon = d.ncon.numpy()[0]
_assert_eq(ncon, 1, "ncon")
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
index a4345e62..892a7115 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
@@ -90,54 +90,40 @@ class SupportPoint:
@wp.func
def _discrete_geoms(g1: int, g2: int) -> bool:
- return (g1 == int(GeomType.MESH.value) or g1 == int(GeomType.BOX.value) or g1 == int(GeomType.HFIELD.value)) and (
- g2 == int(GeomType.MESH.value) or g2 == int(GeomType.BOX.value) or g2 == int(GeomType.HFIELD.value)
+ return (g1 == GeomType.MESH or g1 == GeomType.BOX or g1 == GeomType.HFIELD) and (
+ g2 == GeomType.MESH or g2 == GeomType.BOX or g2 == GeomType.HFIELD
)
-@wp.func
-def _support_margin(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
- sp = SupportPoint()
- sp.cached_index = -1
- sp.vertex_index = -1
- if geomtype == int(GeomType.SPHERE.value):
- sp.point = geom.pos
- return sp
- elif geomtype == int(GeomType.CAPSULE.value):
- local_dir = wp.transpose(geom.rot) @ dir
- res = wp.vec3()
- res[2] = wp.where(local_dir[2] >= 0, geom.size[1], -geom.size[1])
- sp.point = geom.rot @ res + geom.pos
- return sp
-
-
@wp.func
def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
sp = SupportPoint()
sp.cached_index = -1
sp.vertex_index = -1
+ if geomtype == GeomType.SPHERE:
+ sp.point = geom.pos + (0.5 * geom.margin) * geom.size[0] * dir
+ return sp
+
local_dir = wp.transpose(geom.rot) @ dir
- if geomtype == int(GeomType.SPHERE.value):
- sp.point = geom.pos + geom.size[0] * dir
- elif geomtype == int(GeomType.BOX.value):
+ if geomtype == GeomType.BOX:
tmp = wp.sign(local_dir)
res = wp.cw_mul(tmp, geom.size)
sp.point = geom.rot @ res + geom.pos
sp.vertex_index = wp.where(tmp[0] > 0, 1, 0)
sp.vertex_index += wp.where(tmp[1] > 0, 2, 0)
sp.vertex_index += wp.where(tmp[2] > 0, 4, 0)
- elif geomtype == int(GeomType.CAPSULE.value):
+ elif geomtype == GeomType.CAPSULE:
res = local_dir * geom.size[0]
# add cylinder contribution
res[2] += wp.sign(local_dir[2]) * geom.size[1]
sp.point = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.ELLIPSOID.value):
+ elif geomtype == GeomType.ELLIPSOID:
res = wp.cw_mul(local_dir, geom.size)
res = wp.normalize(res)
# transform to ellipsoid
res = wp.cw_mul(res, geom.size)
sp.point = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.CYLINDER.value):
+ elif geomtype == GeomType.CYLINDER:
res = wp.vec3(0.0, 0.0, 0.0)
# set result in XY plane: support on circle
d = wp.sqrt(local_dir[0] * local_dir[0] + local_dir[1] * local_dir[1])
@@ -148,7 +134,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
# set result in Z direction
res[2] = wp.sign(local_dir[2]) * geom.size[1]
sp.point = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.MESH.value):
+ elif geomtype == GeomType.MESH:
max_dist = float(FLOAT_MIN)
if geom.graphadr == -1 or geom.vertnum < 10:
if geom.index > -1:
@@ -195,7 +181,7 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
sp.point = geom.vert[geom.vertadr + imax]
sp.point = geom.rot @ sp.point + geom.pos
- elif geomtype == int(GeomType.HFIELD.value):
+ elif geomtype == GeomType.HFIELD:
max_dist = float(FLOAT_MIN)
for i in range(6):
vert = geom.hfprism[i]
@@ -205,6 +191,8 @@ def _support(geom: Geom, geomtype: int, dir: wp.vec3) -> SupportPoint:
sp.point = vert
sp.point = geom.rot @ sp.point + geom.pos
+ if geom.margin > 0.0:
+ sp.point += dir * (0.5 * geom.margin)
return sp
@@ -591,10 +579,9 @@ def gjk(
geomtype1: int,
geomtype2: int,
cutoff: float,
- use_margin: bool,
+ is_discrete: bool,
) -> GJKResult:
"""Find distance within a tolerance between two geoms."""
- is_discrete = _discrete_geoms(geomtype1, geomtype2)
cutoff2 = cutoff * cutoff
simplex = mat43()
simplex1 = mat43()
@@ -608,9 +595,6 @@ def gjk(
# set initial guess
x_k = x1_0 - x2_0
- use_margin1 = use_margin and (geomtype1 == int(GeomType.SPHERE.value) or geomtype1 == int(GeomType.CAPSULE.value))
- use_margin2 = use_margin and (geomtype2 == int(GeomType.SPHERE.value) or geomtype2 == int(GeomType.CAPSULE.value))
-
for k in range(gjk_iterations):
xnorm = wp.dot(x_k, x_k)
# TODO(kbayes): determine new constant here
@@ -619,13 +603,13 @@ def gjk(
dir_neg = x_k / wp.sqrt(xnorm)
# compute kth support point in geom1
- sp = wp.where(use_margin1, _support_margin(geom1, geomtype1, -dir_neg), _support(geom1, geomtype1, -dir_neg))
+ sp = _support(geom1, geomtype1, -dir_neg)
simplex1[n] = sp.point
geom1.index = sp.cached_index
simplex_index1[n] = sp.vertex_index
# compute kth support point in geom2
- sp = wp.where(use_margin2, _support_margin(geom2, geomtype2, dir_neg), _support(geom2, geomtype2, dir_neg))
+ sp = _support(geom2, geomtype2, dir_neg)
simplex2[n] = sp.point
geom2.index = sp.cached_index
simplex_index2[n] = sp.vertex_index
@@ -692,8 +676,10 @@ def gjk(
result = GJKResult()
# compute the approximate witness points
- result.x1 = _linear_combine(n, coordinates, simplex1)
- result.x2 = _linear_combine(n, coordinates, simplex2)
+ # if n is zero, then there was an immediate return meaning the initial points
+ # are the witness points
+ result.x1 = wp.where(n == 0, x1_0, _linear_combine(n, coordinates, simplex1))
+ result.x2 = wp.where(n == 0, x2_0, _linear_combine(n, coordinates, simplex2))
result.dist = wp.norm_l2(x_k)
result.dim = n
@@ -1199,10 +1185,17 @@ def _polytope4(
@wp.func
def _epa(
- tolerance: float, epa_iterations: int, pt: Polytope, geom1: Geom, geom2: Geom, geomtype1: int, geomtype2: int
+ # In:
+ tolerance: float,
+ epa_iterations: int,
+ pt: Polytope,
+ geom1: Geom,
+ geom2: Geom,
+ geomtype1: int,
+ geomtype2: int,
+ is_discrete: bool,
) -> Tuple[float, wp.vec3, wp.vec3, int]:
"""Recover penetration data from two geoms in contact given an initial polytope."""
- is_discrete = _discrete_geoms(geomtype1, geomtype2)
upper = FLOAT_MAX
upper2 = FLOAT_MAX
idx = int(-1)
@@ -1808,6 +1801,7 @@ def _plane_intersect(pn: wp.vec3, pd: float, a: wp.vec3, b: wp.vec3) -> Tuple[fl
@wp.func
def _polygon_clip(
# In:
+ prune: bool,
plane_normal: wp.array(dtype=wp.vec3),
plane_dist: wp.array(dtype=float),
face1: wp.array(dtype=wp.vec3),
@@ -1887,7 +1881,7 @@ def _polygon_clip(
if npolygon < 1:
return 0, witness1, witness2
- if npolygon > 4:
+ if prune and npolygon > 4:
quad = _polygon_quad(polygon_out, npolygon)
for i in range(4):
witness2[i] = polygon_out[quad[i]]
@@ -1939,7 +1933,7 @@ def _multicontact(
witness1[0] = x1
witness2[0] = x2
- if geomtype1 == int(GeomType.MESH.value):
+ if geomtype1 == GeomType.MESH:
vert = geom1.vert
polynormal = geom1.mesh_polynormal
polyvertadr = geom1.mesh_polyvertadr
@@ -1948,7 +1942,7 @@ def _multicontact(
polymapadr = geom1.mesh_polymapadr
polymapnum = geom1.mesh_polymapnum
polymap = geom1.mesh_polymap
- elif geomtype2 == int(GeomType.MESH.value):
+ elif geomtype2 == GeomType.MESH:
vert = geom2.vert
polynormal = geom2.mesh_polynormal
polyvertadr = geom2.mesh_polyvertadr
@@ -1966,9 +1960,9 @@ def _multicontact(
dir_neg = -dir
# get all possible face normals for each geom
- if geomtype1 == int(GeomType.BOX.value):
+ if geomtype1 == GeomType.BOX:
nnorms1 = _box_normals(nface1, feature_index1, geom1.rot, dir_neg, n1, idx1)
- elif geomtype1 == int(GeomType.MESH.value):
+ elif geomtype1 == GeomType.MESH:
nnorms1 = _mesh_normals(
nface1,
feature_index1,
@@ -1982,9 +1976,9 @@ def _multicontact(
n1,
idx1,
)
- if geomtype2 == int(GeomType.BOX.value):
+ if geomtype2 == GeomType.BOX:
nnorms2 = _box_normals(nface2, feature_index2, geom2.rot, dir, n2, idx2)
- elif geomtype2 == int(GeomType.MESH.value):
+ elif geomtype2 == GeomType.MESH:
nnorms2 = _mesh_normals(
nface2,
feature_index2,
@@ -2007,11 +2001,11 @@ def _multicontact(
# check if edge-face collision
if nface1 < 3 and nface1 <= nface2:
nnorms1 = 0
- if geomtype1 == int(GeomType.BOX.value):
+ if geomtype1 == GeomType.BOX:
nnorms1 = _box_edge_normals(
nface1, geom1.rot, geom1.pos, geom1.size, feature_vertex1[0], feature_vertex1[1], feature_index1[0], n1, endvert
)
- elif geomtype1 == int(GeomType.MESH.value):
+ elif geomtype1 == GeomType.MESH:
nnorms1 = _mesh_edge_normals(
nface1,
geom1.rot,
@@ -2039,11 +2033,11 @@ def _multicontact(
# check if face-edge collision
elif nface2 < 3:
nnorms2 = 0
- if geomtype2 == int(GeomType.BOX.value):
+ if geomtype2 == GeomType.BOX:
nnorms2 = _box_edge_normals(
nface2, geom2.rot, geom2.pos, geom2.size, feature_vertex2[0], feature_vertex2[1], feature_index2[0], n2, endvert
)
- elif geomtype2 == int(GeomType.MESH.value):
+ elif geomtype2 == GeomType.MESH:
nnorms2 = _mesh_edge_normals(
nface2,
geom2.rot,
@@ -2079,9 +2073,9 @@ def _multicontact(
nface1 = _set_edge(pt.vert1, endvert, face[0], i, face1)
else:
ind = wp.where(is_edge_contact_geom2, idx1[j], idx1[i])
- if geomtype1 == int(GeomType.BOX.value):
+ if geomtype1 == GeomType.BOX:
nface1 = _box_face(geom1.rot, geom1.pos, geom1.size, ind, face1)
- elif geomtype1 == int(GeomType.MESH.value):
+ elif geomtype1 == GeomType.MESH:
nface1 = _mesh_face(
geom1.rot,
geom1.pos,
@@ -2099,9 +2093,9 @@ def _multicontact(
if is_edge_contact_geom2:
nface2 = _set_edge(pt.vert2, endvert, face[0], i, face2)
else:
- if geomtype2 == int(GeomType.BOX.value):
+ if geomtype2 == GeomType.BOX:
nface2 = _box_face(geom2.rot, geom2.pos, geom2.size, idx2[j], face2)
- elif geomtype2 == int(GeomType.MESH.value):
+ elif geomtype2 == GeomType.MESH:
nface2 = _mesh_face(
geom2.rot,
geom2.pos,
@@ -2123,16 +2117,19 @@ def _multicontact(
# face1 is an edge; clip face1 against face2
if is_edge_contact_geom1:
approx_dir = wp.norm_l2(dir) * n2[j]
- return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
+ return _polygon_clip(False, plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
# face2 is an edge; clip face2 against face1
if is_edge_contact_geom2:
approx_dir = -wp.norm_l2(dir) * n1[j]
- return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped)
+ return _polygon_clip(False, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[j], approx_dir, polygon, clipped)
# face-face collision
approx_dir = wp.norm_l2(dir) * n2[j]
- return _polygon_clip(plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped)
+
+ # don't prune box-box collisions (expect up to 8 contacts)
+ prune = not (geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX)
+ return _polygon_clip(prune, plane_normal, plane_dist, face1, nface1, face2, nface2, n1[i], approx_dir, polygon, clipped)
@wp.func
@@ -2153,8 +2150,7 @@ def ccd(
multiccd: bool,
tolerance: float,
cutoff: float,
- gjk_iterations: int,
- epa_iterations: int,
+ ccd_iterations: int,
geom1: Geom,
geom2: Geom,
geomtype1: int,
@@ -2187,19 +2183,30 @@ def ccd(
"""General convex collision detection via GJK/EPA."""
witness1 = mat3c()
witness2 = mat3c()
- margin1 = 0.0
- margin2 = 0.0
+ full_margin1 = 0.0
+ full_margin2 = 0.0
+ size1 = 0.0
+ size2 = 0.0
- if geomtype1 == int(GeomType.SPHERE.value) or geomtype1 == int(GeomType.CAPSULE.value):
- margin1 = geom1.size[0]
+ # determine if the geoms being tested are discrete
+ is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
- if geomtype2 == int(GeomType.SPHERE.value) or geomtype2 == int(GeomType.CAPSULE.value):
- margin2 = geom2.size[0]
+ if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
+ size1 = geom1.size[0]
+ full_margin1 = size1 + 0.5 * geom1.margin
+ geom1.margin = 0.0
+ geom1.size = wp.vec3(0.0, geom1.size[1], geom1.size[2])
+
+ if geomtype2 == GeomType.SPHERE or geomtype2 == GeomType.CAPSULE:
+ size2 = geom2.size[0]
+ full_margin2 = size2 + 0.5 * geom2.margin
+ geom2.margin = 0.0
+ geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
# special handling for sphere and capsule (shrink to point and line respectively)
- if margin1 + margin2 > 0.0:
- cutoff += margin1 + margin2
- result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, True)
+ if size1 + size2 > 0.0:
+ cutoff += full_margin1 + full_margin2
+ result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
# shallow penetration, inflate contact
if result.dist > tolerance:
@@ -2207,15 +2214,19 @@ def ccd(
witness1[0] = result.x1
witness2[0] = result.x2
return result.dist, 1, witness1, witness2
- dist, x1, x2 = _inflate(result.dist, result.x1, result.x2, margin1, margin2)
+ dist, x1, x2 = _inflate(result.dist, result.x1, result.x2, full_margin1, full_margin2)
witness1[0] = x1
witness2[0] = x2
return dist, 1, witness1, witness2
# deep penetration, reset initial conditions and rerun GJK + EPA
- cutoff -= margin1 + margin2
+ geom1.margin = full_margin1 - size1
+ geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
+ geom2.margin = full_margin2 - size2
+ geom2.size = wp.vec3(size2, geom2.size[1], geom2.size[2])
+ cutoff -= full_margin1 + full_margin2
- result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, False)
+ result = gjk(tolerance, ccd_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
# no penetration depth to recover
if result.dist > tolerance or result.dim < 2:
@@ -2305,14 +2316,19 @@ def ccd(
witness2[0] = result.x2
return result.dist, 1, witness1, witness2
- dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2)
+ dist, x1, x2, idx = _epa(tolerance, ccd_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
if idx == -1:
return FLOAT_MAX, 0, witness1, witness2
+ # multiccd is always on for box-box collisions
+ if geomtype1 == GeomType.BOX and geomtype2 == GeomType.BOX:
+ multiccd = True
+
if (
multiccd
- and (geomtype1 == int(GeomType.BOX.value) or (geomtype1 == int(GeomType.MESH.value) and geom1.mesh_polyadr > -1))
- and (geomtype2 == int(GeomType.BOX.value) or (geomtype2 == int(GeomType.MESH.value) and geom2.mesh_polyadr > -1))
+ and (geom1.margin == 0.0 and geom2.margin == 0.0)
+ and (geomtype1 == GeomType.BOX or (geomtype1 == GeomType.MESH and geom1.mesh_polyadr > -1))
+ and (geomtype2 == GeomType.BOX or (geomtype2 == GeomType.MESH and geom2.mesh_polyadr > -1))
):
num, w1, w2 = _multicontact(
polygon,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_legacy.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_legacy.py
index e4b2d28d..df50e4ba 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_legacy.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_legacy.py
@@ -54,23 +54,23 @@ VECI2 = vec6(1, 2, 3, 2, 3, 3)
@wp.func
def _gjk_support_geom(geom: Geom, geomtype: int, dir: wp.vec3):
local_dir = wp.transpose(geom.rot) @ dir
- if geomtype == int(GeomType.SPHERE.value):
+ if geomtype == GeomType.SPHERE:
support_pt = geom.pos + geom.size[0] * dir
- elif geomtype == int(GeomType.BOX.value):
+ elif geomtype == GeomType.BOX:
res = wp.cw_mul(wp.sign(local_dir), geom.size)
support_pt = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.CAPSULE.value):
+ elif geomtype == GeomType.CAPSULE:
res = local_dir * geom.size[0]
# add cylinder contribution
res[2] += wp.sign(local_dir[2]) * geom.size[1]
support_pt = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.ELLIPSOID.value):
+ elif geomtype == GeomType.ELLIPSOID:
res = wp.cw_mul(local_dir, geom.size)
res = wp.normalize(res)
# transform to ellipsoid
res = wp.cw_mul(res, geom.size)
support_pt = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.CYLINDER.value):
+ elif geomtype == GeomType.CYLINDER:
res = wp.vec3(0.0, 0.0, 0.0)
# set result in XY plane: support on circle
d = wp.sqrt(wp.dot(local_dir, local_dir))
@@ -81,7 +81,7 @@ def _gjk_support_geom(geom: Geom, geomtype: int, dir: wp.vec3):
# set result in Z direction
res[2] = wp.sign(local_dir[2]) * geom.size[1]
support_pt = geom.rot @ res + geom.pos
- elif geomtype == int(GeomType.MESH.value):
+ elif geomtype == GeomType.MESH:
max_dist = float(FLOAT_MIN)
if geom.graphadr == -1 or geom.vertnum < 10:
# exhaustive search over all vertices
@@ -117,7 +117,7 @@ def _gjk_support_geom(geom: Geom, geomtype: int, dir: wp.vec3):
support_pt = geom.vert[geom.vertadr + imax]
support_pt = geom.rot @ support_pt + geom.pos
- elif geomtype == int(GeomType.HFIELD.value):
+ elif geomtype == GeomType.HFIELD:
max_dist = float(FLOAT_MIN)
for i in range(6):
vert = geom.hfprism[i]
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py
index 5a3c82de..719ccf27 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk_test.py
@@ -13,23 +13,23 @@
# limitations under the License.
# ==============================================================================
+import numpy as np
import warp as wp
from absl.testing import absltest
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
-from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
+from mujoco.mjx.third_party.mujoco_warp import Data
+from mujoco.mjx.third_party.mujoco_warp import GeomType
+from mujoco.mjx.third_party.mujoco_warp import Model
+from mujoco.mjx.third_party.mujoco_warp import test_data
-from mujoco.mjx.third_party.mujoco_warp._src import test_util
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
-from mujoco.mjx.third_party.mujoco_warp._src.types import Data
-from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
-from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
MAX_ITERATIONS = 20
-def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int, multiccd=False):
+def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int, multiccd=False, margin=0.0):
@nested_kernel(module="unique", enable_backward=False)
def _gjk_kernel(
# Model:
@@ -82,18 +82,17 @@ def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int, multicc
ncon_out: wp.array(dtype=int),
pos_out: wp.array(dtype=wp.vec3),
):
- MESHGEOM = int(GeomType.MESH.value)
-
geom1 = Geom()
geom1.index = -1
geomtype1 = geom_type[gid1]
geom1.pos = geom_xpos_in[0, gid1]
geom1.rot = geom_xmat_in[0, gid1]
geom1.size = geom_size[0, gid1]
+ geom1.margin = margin
geom1.graphadr = -1
geom1.mesh_polyadr = -1
- if geom_dataid[gid1] >= 0 and geom_type[gid1] == MESHGEOM:
+ if geom_dataid[gid1] >= 0 and geom_type[gid1] == GeomType.MESH:
dataid = geom_dataid[gid1]
geom1.vertadr = mesh_vertadr[dataid]
geom1.vertnum = mesh_vertnum[dataid]
@@ -114,10 +113,11 @@ def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int, multicc
geom2.pos = geom_xpos_in[0, gid2]
geom2.rot = geom_xmat_in[0, gid2]
geom2.size = geom_size[0, gid2]
+ geom2.margin = margin
geom2.graphadr = -1
geom2.mesh_polyadr = -1
- if geom_dataid[gid2] >= 0 and geom_type[gid2] == MESHGEOM:
+ if geom_dataid[gid2] >= 0 and geom_type[gid2] == GeomType.MESH:
dataid = geom_dataid[gid2]
geom2.vertadr = mesh_vertadr[dataid]
geom2.vertnum = mesh_vertnum[dataid]
@@ -145,7 +145,6 @@ def _geom_dist(m: Model, d: Data, gid1: int, gid2: int, iterations: int, multicc
1e-6,
1.0e30,
iterations,
- iterations,
geom1,
geom2,
geomtype1,
@@ -268,8 +267,8 @@ class GJKTest(absltest.TestCase):
def test_spheres_distance(self):
"""Test distance between two spheres."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -279,14 +278,16 @@ class GJKTest(absltest.TestCase):
"""
)
- dist, _, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS)
+ dist, _, x1, x2 = _geom_dist(m, d, 0, 1, MAX_ITERATIONS)
self.assertEqual(1.0, dist)
+ self.assertEqual(-0.5, x1[0])
+ self.assertEqual(0.5, x2[0])
def test_spheres_touching(self):
"""Test two touching spheres have zero distance"""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -302,8 +303,8 @@ class GJKTest(absltest.TestCase):
def test_box_mesh_distance(self):
"""Test distance between a mesh and box"""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -347,8 +348,8 @@ class GJKTest(absltest.TestCase):
def test_box_box_contact(self):
"""Test penetration between two boxes."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -358,8 +359,10 @@ class GJKTest(absltest.TestCase):
"""
)
dist, _, x1, x2 = _geom_dist(m, d, 0, 1, MAX_ITERATIONS)
+ diff = x1 - x2
+ normal = diff / np.linalg.norm(diff)
+
self.assertAlmostEqual(-1, dist)
- normal = wp.normalize(x1 - x2)
self.assertAlmostEqual(normal[0], 1)
self.assertAlmostEqual(normal[1], 0)
self.assertAlmostEqual(normal[2], 0)
@@ -367,8 +370,8 @@ class GJKTest(absltest.TestCase):
def test_mesh_mesh_contact(self):
"""Test penetration between two meshes."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -421,8 +424,8 @@ class GJKTest(absltest.TestCase):
def test_box_edge(self):
"""Test box edge."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -436,8 +439,8 @@ class GJKTest(absltest.TestCase):
def test_box_box_ccd(self):
"""Test box box."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -452,8 +455,8 @@ class GJKTest(absltest.TestCase):
def test_mesh_mesh_ccd(self):
"""Test mesh-mesh multiccd."""
- _, _, m, d = test_util.fixture(
- xml=f"""
+ _, _, m, d = test_data.fixture(
+ xml="""
@@ -485,7 +488,28 @@ class GJKTest(absltest.TestCase):
)
_, ncon, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=True)
- self.assertEqual(ncon, 4)
+ self.assertEqual(ncon, 5)
+
+ def test_sphere_mesh_margin(self):
+ """Test sphere-mesh margin."""
+
+ _, _, m, d = test_data.fixture(
+ xml="""
+
+
+
+
+
+
+
+
+
+ """
+ )
+
+ dist, _, _, _ = _geom_dist(m, d, 0, 1, MAX_ITERATIONS, multiccd=False, margin=0.05)
+ self.assertAlmostEqual(dist, -0.001)
if __name__ == "__main__":
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
index 340ca24c..5f15e808 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_primitive.py
@@ -38,6 +38,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import vec5
+from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import kernel as nested_kernel
@@ -59,6 +60,7 @@ class Geom:
rot: wp.mat33
normal: wp.vec3
size: wp.vec3
+ margin: float
hfprism: mat63
vertadr: int
vertnum: int
@@ -84,13 +86,13 @@ def geom(
geom_type: int,
geom_dataid: int,
geom_size: wp.vec3,
- mesh_vertadr: int,
- mesh_vertnum: int,
+ mesh_vertadr: wp.array(dtype=int),
+ mesh_vertnum: wp.array(dtype=int),
mesh_vert: wp.array(dtype=wp.vec3),
- mesh_graphadr: int,
+ mesh_graphadr: wp.array(dtype=int),
mesh_graph: wp.array(dtype=int),
- mesh_polynum: int,
- mesh_polyadr: int,
+ mesh_polynum: wp.array(dtype=int),
+ mesh_polyadr: wp.array(dtype=int),
mesh_polynormal: wp.array(dtype=wp.vec3),
mesh_polyvertadr: wp.array(dtype=int),
mesh_polyvertnum: wp.array(dtype=int),
@@ -109,21 +111,20 @@ def geom(
geom.size = geom_size
geom.normal = wp.vec3(geom_xmat_in[0, 2], geom_xmat_in[1, 2], geom_xmat_in[2, 2]) # plane
- # If geom is MESH, get mesh verts
- if geom_dataid >= 0 and geom_type == int(GeomType.MESH.value):
- geom.vertadr = mesh_vertadr
- geom.vertnum = mesh_vertnum
- geom.graphadr = mesh_graphadr
- geom.mesh_polynum = mesh_polynum
- geom.mesh_polyadr = mesh_polyadr
- else:
- geom.vertadr = -1
- geom.vertnum = -1
- geom.graphadr = -1
- geom.mesh_polynum = -1
- geom.mesh_polyadr = -1
+ if geom_type == GeomType.MESH:
+ if geom_dataid >= 0:
+ geom.vertadr = mesh_vertadr[geom_dataid]
+ geom.vertnum = mesh_vertnum[geom_dataid]
+ geom.graphadr = mesh_graphadr[geom_dataid]
+ geom.mesh_polynum = mesh_polynum[geom_dataid]
+ geom.mesh_polyadr = mesh_polyadr[geom_dataid]
+ else:
+ geom.vertadr = -1
+ geom.vertnum = -1
+ geom.graphadr = -1
+ geom.mesh_polynum = -1
+ geom.mesh_polyadr = -1
- if geom_type == int(GeomType.MESH.value):
geom.vert = mesh_vert
geom.graph = mesh_graph
geom.mesh_polynormal = mesh_polynormal
@@ -135,6 +136,8 @@ def geom(
geom.mesh_polymap = mesh_polymap
geom.index = -1
+ geom.margin = 0.0
+
return geom
@@ -167,56 +170,64 @@ def plane_convex(plane_normal: wp.vec3, plane_pos: wp.vec3, convex: Geom) -> Tup
# exhaustive search over all vertices
if convex.graphadr == -1 or convex.vertnum < 10:
- # Find support points
+ # find support points
max_support = wp.float32(-_HUGE_VAL)
for i in range(convex.vertnum):
support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
max_support = wp.max(support, max_support)
threshold = wp.max(0.0, max_support - 1e-3)
- # Find point a (first support point)
+
+ # find first support point (a)
a_dist = wp.float32(-_HUGE_VAL)
+ a = wp.vec3()
for i in range(convex.vertnum):
- support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
+ vert = convex.vert[convex.vertadr + i]
+ support = wp.dot(plane_pos_local - vert, n)
dist = wp.where(support > threshold, support, -_HUGE_VAL)
if dist > a_dist:
indices[0] = i
a_dist = dist
- a = convex.vert[convex.vertadr + indices[0]]
+ a = vert
- # Find point b (furthest from a)
+ # find point (b) furthest from a
b_dist = wp.float32(-_HUGE_VAL)
+ b = wp.vec3()
for i in range(convex.vertnum):
- support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
+ vert = convex.vert[convex.vertadr + i]
+ support = wp.dot(plane_pos_local - vert, n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
- dist = wp.length_sq(a - convex.vert[convex.vertadr + i]) + dist_mask
+ dist = wp.length_sq(a - vert) + dist_mask
if dist > b_dist:
indices[1] = i
b_dist = dist
- b = convex.vert[convex.vertadr + indices[1]]
+ b = vert
- # Find point c (furthest along axis orthogonal to a-b)
+ # find point (c) furthest along axis orthogonal to a-b
ab = wp.cross(n, a - b)
c_dist = wp.float32(-_HUGE_VAL)
+ c = wp.vec3()
for i in range(convex.vertnum):
- support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
+ vert = convex.vert[convex.vertadr + i]
+ support = wp.dot(plane_pos_local - vert, n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
- ap = a - convex.vert[convex.vertadr + i]
+ ap = a - vert
dist = wp.abs(wp.dot(ap, ab)) + dist_mask
if dist > c_dist:
indices[2] = i
c_dist = dist
- c = convex.vert[convex.vertadr + indices[2]]
+ c = vert
- # Find point d (furthest from other triangle edges)
+ # find point (d) furthest from other triangle edges
ac = wp.cross(n, a - c)
bc = wp.cross(n, b - c)
d_dist = wp.float32(-_HUGE_VAL)
for i in range(convex.vertnum):
- support = wp.dot(plane_pos_local - convex.vert[convex.vertadr + i], n)
+ vert = convex.vert[convex.vertadr + i]
+ support = wp.dot(plane_pos_local - vert, n)
dist_mask = wp.where(support > threshold, 0.0, -_HUGE_VAL)
- ap = a - convex.vert[convex.vertadr + i]
- bp = b - convex.vert[convex.vertadr + i]
+ ap = a - vert
+ bp = b - vert
dist_ap = wp.abs(wp.dot(ap, ac)) + dist_mask
dist_bp = wp.abs(wp.dot(bp, bc)) + dist_mask
if dist_ap + dist_bp > d_dist:
@@ -1406,19 +1417,19 @@ def box_box_wrapper(
_PRIMITIVE_COLLISIONS = {
- (GeomType.PLANE.value, GeomType.SPHERE.value): plane_sphere_wrapper,
- (GeomType.PLANE.value, GeomType.CAPSULE.value): plane_capsule_wrapper,
- (GeomType.PLANE.value, GeomType.ELLIPSOID.value): plane_ellipsoid_wrapper,
- (GeomType.PLANE.value, GeomType.CYLINDER.value): plane_cylinder_wrapper,
- (GeomType.PLANE.value, GeomType.BOX.value): plane_box_wrapper,
- (GeomType.PLANE.value, GeomType.MESH.value): plane_convex_wrapper,
- (GeomType.SPHERE.value, GeomType.SPHERE.value): sphere_sphere_wrapper,
- (GeomType.SPHERE.value, GeomType.CAPSULE.value): sphere_capsule_wrapper,
- (GeomType.SPHERE.value, GeomType.CYLINDER.value): sphere_cylinder_wrapper,
- (GeomType.SPHERE.value, GeomType.BOX.value): sphere_box_wrapper,
- (GeomType.CAPSULE.value, GeomType.CAPSULE.value): capsule_capsule_wrapper,
- (GeomType.CAPSULE.value, GeomType.BOX.value): capsule_box_wrapper,
- (GeomType.BOX.value, GeomType.BOX.value): box_box_wrapper,
+ (GeomType.PLANE, GeomType.SPHERE): plane_sphere_wrapper,
+ (GeomType.PLANE, GeomType.CAPSULE): plane_capsule_wrapper,
+ (GeomType.PLANE, GeomType.ELLIPSOID): plane_ellipsoid_wrapper,
+ (GeomType.PLANE, GeomType.CYLINDER): plane_cylinder_wrapper,
+ (GeomType.PLANE, GeomType.BOX): plane_box_wrapper,
+ (GeomType.PLANE, GeomType.MESH): plane_convex_wrapper,
+ (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere_wrapper,
+ (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule_wrapper,
+ (GeomType.SPHERE, GeomType.CYLINDER): sphere_cylinder_wrapper,
+ (GeomType.SPHERE, GeomType.BOX): sphere_box_wrapper,
+ (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule_wrapper,
+ (GeomType.CAPSULE, GeomType.BOX): capsule_box_wrapper,
+ (GeomType.BOX, GeomType.BOX): box_box_wrapper,
}
@@ -1426,7 +1437,7 @@ _PRIMITIVE_COLLISIONS = {
def _check_primitive_collisions():
prev_idx = -1
for types in _PRIMITIVE_COLLISIONS.keys():
- idx = upper_trid_index(len(GeomType), types[0], types[1])
+ idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
if types[1] < types[0] or idx <= prev_idx:
return False
prev_idx = idx
@@ -1439,13 +1450,7 @@ _primitive_collisions_types = []
_primitive_collisions_func = []
-def _primitive_narrowphase_builder(m: Model):
- for types, func in _PRIMITIVE_COLLISIONS.items():
- idx = upper_trid_index(len(GeomType), types[0], types[1])
- if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types:
- _primitive_collisions_types.append(types)
- _primitive_collisions_func.append(func)
-
+def _create_narrowphase_kernel():
@nested_kernel(module="unique", enable_backward=False)
def _primitive_narrowphase(
# Model:
@@ -1549,13 +1554,13 @@ def _primitive_narrowphase_builder(m: Model):
type1,
geom1_dataid,
geom_size[worldid, g1],
- mesh_vertadr[geom1_dataid],
- mesh_vertnum[geom1_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom1_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom1_dataid],
- mesh_polyadr[geom1_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
@@ -1572,13 +1577,13 @@ def _primitive_narrowphase_builder(m: Model):
type2,
geom2_dataid,
geom_size[worldid, g2],
- mesh_vertadr[geom2_dataid],
- mesh_vertnum[geom2_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom2_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom2_dataid],
- mesh_polyadr[geom2_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
@@ -1625,6 +1630,16 @@ def _primitive_narrowphase_builder(m: Model):
return _primitive_narrowphase
+def _primitive_narrowphase_builder(m: Model):
+ for types, func in _PRIMITIVE_COLLISIONS.items():
+ idx = upper_trid_index(len(GeomType), types[0].value, types[1].value)
+ if m.geom_pair_type_count[idx] and types not in _primitive_collisions_types:
+ _primitive_collisions_types.append(types)
+ _primitive_collisions_func.append(func)
+
+ return _create_narrowphase_kernel()
+
+
@event_scope
def primitive_narrowphase(m: Model, d: Data):
"""Runs collision detection on primitive geom pairs discovered during broadphase.
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
index c08e838d..7427b16b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
@@ -93,11 +93,11 @@ def get_sdf_params(
plugin_index = -1
volume_data = VolumeData()
- if g_type == int(GeomType.SDF.value) and plugin_id != -1:
+ if g_type == GeomType.SDF and plugin_id != -1:
attributes = plugin_attr[plugin_id]
plugin_index = plugin[plugin_id]
- elif g_type == int(GeomType.SDF.value) and mesh_id != -1:
+ elif g_type == GeomType.SDF and mesh_id != -1:
volume_data.center = oct_aabb[mesh_id, 0]
volume_data.half_size = oct_aabb[mesh_id, 1]
volume_data.oct_aabb = oct_aabb
@@ -232,6 +232,7 @@ def find_oct(
rx = vec8f(0.0)
ry = vec8f(0.0)
rz = vec8f(0.0)
+ eps = 1e-6
while niter > 0:
niter -= 1
@@ -243,6 +244,17 @@ def find_oct(
vmin = oct_aabb[node, 0] - oct_aabb[node, 1]
vmax = oct_aabb[node, 0] + oct_aabb[node, 1]
+
+ if (
+ p[0] + eps < vmin[0]
+ or p[0] - eps > vmax[0]
+ or p[1] + eps < vmin[1]
+ or p[1] - eps > vmax[1]
+ or p[2] + eps < vmin[2]
+ or p[2] - eps > vmax[2]
+ ):
+ continue
+
coord = wp.cw_div(p - vmin, vmax - vmin)
# check if the node is a leaf
@@ -270,9 +282,9 @@ def find_oct(
return node, (rx, ry, rz)
# compute which of 8 children to visit next
- x = 1 if coord[0] < 0.5 else 0
- y = 1 if coord[1] < 0.5 else 0
- z = 1 if coord[2] < 0.5 else 0
+ x = 0 if coord[0] < 0.5 else 1
+ y = 0 if coord[1] < 0.5 else 1
+ z = 0 if coord[2] < 0.5 else 1
stack = oct_child[node][4 * z + 2 * y + x]
wp.print("ERROR: Node not found\n")
@@ -345,15 +357,15 @@ def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
@wp.func
def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> float:
- if type == int(GeomType.PLANE.value):
+ if type == GeomType.PLANE:
return p[2]
- elif type == int(GeomType.SPHERE.value):
+ elif type == GeomType.SPHERE:
return sphere(p, attr)
- elif type == int(GeomType.BOX.value):
+ elif type == GeomType.BOX:
return box(p, attr)
- elif type == int(GeomType.ELLIPSOID.value):
+ elif type == GeomType.ELLIPSOID:
return ellipsoid(p, attr)
- elif type == int(GeomType.MESH.value) and mesh_data.valid:
+ elif type == GeomType.MESH and mesh_data.valid:
mesh_data.pnt = p
mesh_data.vec = -wp.normalize(p)
dist = ray_mesh(
@@ -382,7 +394,7 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
-mesh_data.vec,
)
return dist
- elif type == int(GeomType.SDF.value):
+ elif type == GeomType.SDF:
if sdf_type == -1:
return sample_volume_sdf(p, volume_data)
else:
@@ -393,16 +405,16 @@ def sdf(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: Volume
@wp.func
def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: VolumeData, mesh_data: MeshData) -> wp.vec3:
- if type == int(GeomType.PLANE.value):
+ if type == GeomType.PLANE:
grad = wp.vec3(0.0, 0.0, 1.0)
return grad
- elif type == int(GeomType.SPHERE.value):
+ elif type == GeomType.SPHERE:
return grad_sphere(p)
- elif type == int(GeomType.BOX.value):
+ elif type == GeomType.BOX:
return grad_box(p, attr)
- elif type == int(GeomType.ELLIPSOID.value):
+ elif type == GeomType.ELLIPSOID:
return grad_ellipsoid(p, attr)
- elif type == int(GeomType.MESH.value) and mesh_data.valid:
+ elif type == GeomType.MESH and mesh_data.valid:
mesh_data.pnt = p
mesh_data.vec = -wp.normalize(p)
dist = ray_mesh(
@@ -422,7 +434,7 @@ def sdf_grad(type: int, p: wp.vec3, attr: wp.vec3, sdf_type: int, volume_data: V
else:
return wp.vec3(-1.0)
- elif type == int(GeomType.SDF.value):
+ elif type == GeomType.SDF:
if sdf_type == -1:
return sample_volume_grad(p, volume_data)
else:
@@ -448,7 +460,7 @@ def clearance(
mesh_data2: MeshData,
) -> float:
sdf1 = sdf(type1, p1, s1, sdf_type1, volume_data1, mesh_data1)
- sdf2 = sdf(int(GeomType.SDF.value), p2, s2, sdf_type2, volume_data2, mesh_data2)
+ sdf2 = sdf(GeomType.SDF, p2, s2, sdf_type2, volume_data2, mesh_data2)
if sfd_intersection:
return wp.max(sdf1, sdf2)
else:
@@ -471,9 +483,9 @@ def compute_grad(
mesh_data2: MeshData,
) -> wp.vec3:
A = sdf(type1, p1, params.attr1, sdf_type1, volume_data1, mesh_data1)
- B = sdf(int(GeomType.SDF.value), p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
+ B = sdf(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
grad1 = sdf_grad(type1, p1, params.attr1, sdf_type1, volume_data1, mesh_data1)
- grad2 = sdf_grad(int(GeomType.SDF.value), p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
+ grad2 = sdf_grad(GeomType.SDF, p2, params.attr2, sdf_type2, volume_data2, mesh_data2)
grad1_transformed = wp.transpose(params.rel_mat) * grad1
if sfd_intersection:
if A > B:
@@ -594,7 +606,7 @@ def gradient_descent(
grad1 = sdf_grad(type1, x_1, params.attr1, sdf_type1, volume_data1, mesh_data1)
grad1 = wp.transpose(params.rel_mat) * grad1
grad1 = wp.normalize(grad1)
- grad2 = sdf_grad(int(GeomType.SDF.value), x, params.attr2, sdf_type2, volume_data2, mesh_data2)
+ grad2 = sdf_grad(GeomType.SDF, x, params.attr2, sdf_type2, volume_data2, mesh_data2)
grad2 = wp.normalize(grad2)
n = grad1 - grad2
n = wp.normalize(n)
@@ -688,7 +700,7 @@ def _sdf_narrowphase(
geoms = collision_pair_in[contact_tid]
g2 = geoms[1]
type2 = geom_type[g2]
- if type2 != int(GeomType.SDF.value):
+ if type2 != GeomType.SDF:
return
worldid = collision_worldid_in[contact_tid]
_, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -720,13 +732,13 @@ def _sdf_narrowphase(
type1,
geom1_dataid,
geom_size[worldid, g1],
- mesh_vertadr[geom1_dataid],
- mesh_vertnum[geom1_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom1_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom1_dataid],
- mesh_polyadr[geom1_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
@@ -743,13 +755,13 @@ def _sdf_narrowphase(
type2,
geom2_dataid,
geom_size[worldid, g2],
- mesh_vertadr[geom2_dataid],
- mesh_vertnum[geom2_dataid],
+ mesh_vertadr,
+ mesh_vertnum,
mesh_vert,
- mesh_graphadr[geom2_dataid],
+ mesh_graphadr,
mesh_graph,
- mesh_polynum[geom2_dataid],
- mesh_polyadr[geom2_dataid],
+ mesh_polynum,
+ mesh_polyadr,
mesh_polynormal,
mesh_polyvertadr,
mesh_polyvertnum,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
index 201eab65..c1c285df 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
@@ -166,11 +166,10 @@ def _efc_equality_connect(
):
"""Calculates constraint rows for connect equality constraints."""
- worldid, i_eq_connect_adr = wp.tid()
- timestep = opt_timestep[worldid]
- i_eq = eq_connect_adr[i_eq_connect_adr]
+ worldid, eqconnectid = wp.tid()
+ eqid = eq_connect_adr[eqconnectid]
- if not eq_active_in[worldid, i_eq]:
+ if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_connect_out, worldid, 3)
@@ -179,14 +178,14 @@ def _efc_equality_connect(
if efcid + 3 >= njmax_in:
return
- data = eq_data[worldid, i_eq]
+ data = eq_data[worldid, eqid]
anchor1 = wp.vec3f(data[0], data[1], data[2])
anchor2 = wp.vec3f(data[3], data[4], data[5])
- obj1id = eq_obj1id[i_eq]
- obj2id = eq_obj2id[i_eq]
+ obj1id = eq_obj1id[eqid]
+ obj2id = eq_obj2id[eqid]
- if nsite and eq_objtype[i_eq] == wp.static(types.ObjType.SITE.value):
+ if nsite and eq_objtype[eqid] == types.ObjType.SITE:
# body1id stores the index of site_bodyid.
body1id = site_bodyid[obj1id]
body2id = site_bodyid[obj2id]
@@ -235,15 +234,15 @@ def _efc_equality_connect(
invweight = body_invweight0[worldid, body1id][0] + body_invweight0[worldid, body2id][0]
pos_imp = wp.length(pos)
- solref = eq_solref[worldid, i_eq]
- solimp = eq_solimp[worldid, i_eq]
+ solref = eq_solref[worldid, eqid]
+ solimp = eq_solimp[worldid, eqid]
for i in range(3):
efcidi = efcid + i
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcidi,
pos[i],
@@ -254,8 +253,8 @@ def _efc_equality_connect(
0.0,
Jqvel[i],
0.0,
- ConstraintType.EQUALITY.value,
- i_eq,
+ ConstraintType.EQUALITY,
+ eqid,
efc_type_out,
efc_id_out,
efc_pos_out,
@@ -302,10 +301,10 @@ def _efc_equality_joint(
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
):
- worldid, i_eq_joint_adr = wp.tid()
- timestep = opt_timestep[worldid]
- i_eq = eq_jnt_adr[i_eq_joint_adr]
- if not eq_active_in[worldid, i_eq]:
+ worldid, eqjntid = wp.tid()
+ eqid = eq_jnt_adr[eqjntid]
+
+ if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_jnt_out, worldid, 1)
@@ -317,9 +316,9 @@ def _efc_equality_joint(
for i in range(nv):
efc_J_out[worldid, efcid, i] = 0.0
- jntid_1 = eq_obj1id[i_eq]
- jntid_2 = eq_obj2id[i_eq]
- data = eq_data[worldid, i_eq]
+ jntid_1 = eq_obj1id[eqid]
+ jntid_2 = eq_obj2id[eqid]
+ data = eq_data[worldid, eqid]
dofadr1 = jnt_dofadr[jntid_1]
qposadr1 = jnt_qposadr[jntid_1]
efc_J_out[worldid, efcid, dofadr1] = 1.0
@@ -348,19 +347,19 @@ def _efc_equality_joint(
# Update constraint parameters
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
pos,
pos,
invweight,
- eq_solref[worldid, i_eq],
- eq_solimp[worldid, i_eq],
+ eq_solref[worldid, eqid],
+ eq_solimp[worldid, eqid],
0.0,
Jqvel,
0.0,
- ConstraintType.EQUALITY.value,
- i_eq,
+ ConstraintType.EQUALITY,
+ eqid,
efc_type_out,
efc_id_out,
efc_pos_out,
@@ -406,9 +405,8 @@ def _efc_equality_tendon(
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
):
- worldid, tenid = wp.tid()
- timestep = opt_timestep[worldid]
- eqid = eq_ten_adr[tenid]
+ worldid, eqtenid = wp.tid()
+ eqid = eq_ten_adr[eqtenid]
if not eq_active_in[worldid, eqid]:
return
@@ -457,7 +455,7 @@ def _efc_equality_tendon(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
pos,
@@ -468,7 +466,7 @@ def _efc_equality_tendon(
0.0,
Jqvel,
0.0,
- ConstraintType.EQUALITY.value,
+ ConstraintType.EQUALITY,
eqid,
efc_type_out,
efc_id_out,
@@ -509,18 +507,16 @@ def _efc_friction_dof(
efc_frictionloss_out: wp.array2d(dtype=float),
):
worldid, dofid = wp.tid()
- timestep = opt_timestep[worldid]
if dof_frictionloss[worldid, dofid] <= 0.0:
return
+ wp.atomic_add(nf_out, worldid, 1)
efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
return
- wp.atomic_add(nf_out, worldid, 1)
-
for i in range(nv):
efc_J_out[worldid, efcid, i] = 0.0
@@ -529,7 +525,7 @@ def _efc_friction_dof(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
0.0,
@@ -540,7 +536,7 @@ def _efc_friction_dof(
0.0,
Jqvel,
dof_frictionloss[worldid, dofid],
- ConstraintType.FRICTION_DOF.value,
+ ConstraintType.FRICTION_DOF,
dofid,
efc_type_out,
efc_id_out,
@@ -563,6 +559,7 @@ def _efc_friction_tendon(
tendon_frictionloss: wp.array2d(dtype=float),
tendon_invweight0: wp.array2d(dtype=float),
# Data in:
+ njmax_in: int,
qvel_in: wp.array2d(dtype=float),
ten_J_in: wp.array3d(dtype=float),
# In:
@@ -581,14 +578,16 @@ def _efc_friction_tendon(
efc_frictionloss_out: wp.array2d(dtype=float),
):
worldid, tenid = wp.tid()
- timestep = opt_timestep[worldid]
frictionloss = tendon_frictionloss[worldid, tenid]
if frictionloss <= 0.0:
return
- efcid = wp.atomic_add(nefc_out, worldid, 1)
wp.atomic_add(nf_out, worldid, 1)
+ efcid = wp.atomic_add(nefc_out, worldid, 1)
+
+ if efcid >= njmax_in:
+ return
Jqvel = float(0.0)
@@ -600,7 +599,7 @@ def _efc_friction_tendon(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
0.0,
@@ -611,7 +610,7 @@ def _efc_friction_tendon(
0.0,
Jqvel,
frictionloss,
- ConstraintType.FRICTION_TENDON.value,
+ ConstraintType.FRICTION_TENDON,
tenid,
efc_type_out,
efc_id_out,
@@ -668,10 +667,10 @@ def _efc_equality_weld(
efc_aref_out: wp.array2d(dtype=float),
efc_frictionloss_out: wp.array2d(dtype=float),
):
- worldid, i_eq_weld_adr = wp.tid()
- timestep = opt_timestep[worldid]
- i_eq = eq_wld_adr[i_eq_weld_adr]
- if not eq_active_in[worldid, i_eq]:
+ worldid, eqweldid = wp.tid()
+ eqid = eq_wld_adr[eqweldid]
+
+ if not eq_active_in[worldid, eqid]:
return
wp.atomic_add(ne_weld_out, worldid, 6)
@@ -680,12 +679,12 @@ def _efc_equality_weld(
if efcid + 6 >= njmax_in:
return
- is_site = eq_objtype[i_eq] == wp.static(types.ObjType.SITE.value) and nsite > 0
+ is_site = eq_objtype[eqid] == types.ObjType.SITE and nsite > 0
- obj1id = eq_obj1id[i_eq]
- obj2id = eq_obj2id[i_eq]
+ obj1id = eq_obj1id[eqid]
+ obj2id = eq_obj2id[eqid]
- data = eq_data[worldid, i_eq]
+ data = eq_data[worldid, eqid]
anchor1 = wp.vec3(data[0], data[1], data[2])
anchor2 = wp.vec3(data[3], data[4], data[5])
relpose = wp.quat(data[6], data[7], data[8], data[9])
@@ -739,14 +738,14 @@ def _efc_equality_weld(
)
jacdifp = jacp1 - jacp2
- for i in range(wp.static(3)):
+ for i in range(3):
efc_J_out[worldid, efcid + i, dofid] = jacdifp[i]
jacdifr = (jacr1 - jacr2) * torquescale
jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
- for i in range(wp.static(3)):
+ for i in range(3):
efc_J_out[worldid, efcid + 3 + i, dofid] = jacdifr[i]
Jqvelp += jacdifp * qvel_in[worldid, dofid]
@@ -762,8 +761,10 @@ def _efc_equality_weld(
pos_imp = wp.sqrt(wp.length_sq(cpos) + wp.length_sq(crot))
- solref = eq_solref[worldid, i_eq]
- solimp = eq_solimp[worldid, i_eq]
+ solref = eq_solref[worldid, eqid]
+ solimp = eq_solimp[worldid, eqid]
+
+ timestep = opt_timestep[worldid]
for i in range(3):
_update_efc_row(
@@ -779,8 +780,8 @@ def _efc_equality_weld(
0.0,
Jqvelp[i],
0.0,
- ConstraintType.EQUALITY.value,
- i_eq,
+ ConstraintType.EQUALITY,
+ eqid,
efc_type_out,
efc_id_out,
efc_pos_out,
@@ -807,8 +808,8 @@ def _efc_equality_weld(
0.0,
Jqvelr[i],
0.0,
- ConstraintType.EQUALITY.value,
- i_eq,
+ ConstraintType.EQUALITY,
+ eqid,
efc_type_out,
efc_id_out,
efc_pos_out,
@@ -853,7 +854,6 @@ def _efc_limit_slide_hinge(
efc_frictionloss_out: wp.array2d(dtype=float),
):
worldid, jntlimitedid = wp.tid()
- timestep = opt_timestep[worldid]
jntid = jnt_limited_slide_hinge_adr[jntlimitedid]
jntrange = jnt_range[worldid, jntid]
@@ -881,7 +881,7 @@ def _efc_limit_slide_hinge(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
pos,
@@ -892,7 +892,7 @@ def _efc_limit_slide_hinge(
jntmargin,
Jqvel,
0.0,
- ConstraintType.LIMIT_JOINT.value,
+ ConstraintType.LIMIT_JOINT,
jntid,
efc_type_out,
efc_id_out,
@@ -938,7 +938,6 @@ def _efc_limit_ball(
efc_frictionloss_out: wp.array2d(dtype=float),
):
worldid, jntlimitedid = wp.tid()
- timestep = opt_timestep[worldid]
jntid = jnt_limited_ball_adr[jntlimitedid]
qposadr = jnt_qposadr[jntid]
@@ -975,7 +974,7 @@ def _efc_limit_ball(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
pos,
@@ -986,7 +985,7 @@ def _efc_limit_ball(
jntmargin,
Jqvel,
0.0,
- ConstraintType.LIMIT_JOINT.value,
+ ConstraintType.LIMIT_JOINT,
jntid,
efc_type_out,
efc_id_out,
@@ -1036,7 +1035,6 @@ def _efc_limit_tendon(
efc_frictionloss_out: wp.array2d(dtype=float),
):
worldid, tenlimitedid = wp.tid()
- timestep = opt_timestep[worldid]
tenid = tendon_limited_adr[tenlimitedid]
tenrange = tendon_range[worldid, tenid]
@@ -1060,7 +1058,7 @@ def _efc_limit_tendon(
scl = float(dist_min < dist_max) * 2.0 - 1.0
adr = tendon_adr[tenid]
- if wrap_type[adr] == wp.static(types.WrapType.JOINT.value):
+ if wrap_type[adr] == types.WrapType.JOINT:
ten_num = tendon_num[tenid]
for i in range(ten_num):
dofadr = jnt_dofadr[wrap_objid[adr + i]]
@@ -1075,7 +1073,7 @@ def _efc_limit_tendon(
_update_efc_row(
worldid,
- timestep,
+ opt_timestep[worldid],
refsafe_in,
efcid,
pos,
@@ -1086,7 +1084,7 @@ def _efc_limit_tendon(
tenmargin,
Jqvel,
0.0,
- ConstraintType.LIMIT_TENDON.value,
+ ConstraintType.LIMIT_TENDON,
tenid,
efc_type_out,
efc_id_out,
@@ -1159,8 +1157,8 @@ def _efc_contact_pyramidal(
if active:
worldid = worldid_in[conid]
- efcid = wp.atomic_add(nefc_out, worldid, 1)
+ efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
contact_efc_address_out[conid, dimid] = -1
return
@@ -1234,9 +1232,9 @@ def _efc_contact_pyramidal(
Jqvel += J * qvel_in[worldid, i]
if condim == 1:
- efc_type = int(ConstraintType.CONTACT_FRICTIONLESS.value)
+ efc_type = ConstraintType.CONTACT_FRICTIONLESS
else:
- efc_type = int(ConstraintType.CONTACT_PYRAMIDAL.value)
+ efc_type = ConstraintType.CONTACT_PYRAMIDAL
_update_efc_row(
worldid,
@@ -1323,8 +1321,8 @@ def _efc_contact_elliptic(
if active:
worldid = worldid_in[conid]
- efcid = wp.atomic_add(nefc_out, worldid, 1)
+ efcid = wp.atomic_add(nefc_out, worldid, 1)
if efcid >= njmax_in:
contact_efc_address_out[conid, dimid] = -1
return
@@ -1402,9 +1400,9 @@ def _efc_contact_elliptic(
pos_aref = 0.0
if condim == 1:
- efc_type = int(ConstraintType.CONTACT_FRICTIONLESS.value)
+ efc_type = ConstraintType.CONTACT_FRICTIONLESS
else:
- efc_type = int(ConstraintType.CONTACT_ELLIPTIC.value)
+ efc_type = ConstraintType.CONTACT_ELLIPTIC
_update_efc_row(
worldid,
@@ -1454,22 +1452,13 @@ def make_constraint(m: types.Model, d: types.Data):
wp.launch(
zero_constraint_counts,
dim=d.nworld,
- inputs=[
- d.ne,
- d.ne_connect,
- d.ne_weld,
- d.ne_jnt,
- d.ne_ten,
- d.nf,
- d.nl,
- d.nefc,
- ],
+ inputs=[d.ne, d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten, d.nf, d.nl, d.nefc],
)
- if not (m.opt.disableflags & types.DisableBit.CONSTRAINT.value):
+ if not (m.opt.disableflags & types.DisableBit.CONSTRAINT):
refsafe = m.opt.disableflags & types.DisableBit.REFSAFE
- if not (m.opt.disableflags & types.DisableBit.EQUALITY.value):
+ if not (m.opt.disableflags & types.DisableBit.EQUALITY):
wp.launch(
_efc_equality_connect,
dim=(d.nworld, m.eq_connect_adr.size),
@@ -1632,19 +1621,12 @@ def make_constraint(m: types.Model, d: types.Data):
wp.launch(
_num_equality,
- dim=(d.nworld,),
- inputs=[
- d.ne_connect,
- d.ne_weld,
- d.ne_jnt,
- d.ne_ten,
- ],
- outputs=[
- d.ne,
- ],
+ dim=d.nworld,
+ inputs=[d.ne_connect, d.ne_weld, d.ne_jnt, d.ne_ten],
+ outputs=[d.ne],
)
- if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS.value):
+ if not (m.opt.disableflags & types.DisableBit.FRICTIONLOSS):
wp.launch(
_efc_friction_dof,
dim=(d.nworld, m.nv),
@@ -1684,6 +1666,7 @@ def make_constraint(m: types.Model, d: types.Data):
m.tendon_solimp_fri,
m.tendon_frictionloss,
m.tendon_invweight0,
+ d.njmax,
d.qvel,
d.ten_J,
refsafe,
@@ -1704,122 +1687,116 @@ def make_constraint(m: types.Model, d: types.Data):
)
# limit
- if not (m.opt.disableflags & types.DisableBit.LIMIT.value):
- limit_ball = m.jnt_limited_ball_adr.size > 0
- if limit_ball:
- wp.launch(
- _efc_limit_ball,
- dim=(d.nworld, m.jnt_limited_ball_adr.size),
- inputs=[
- m.nv,
- m.opt.timestep,
- m.jnt_qposadr,
- m.jnt_dofadr,
- m.jnt_solref,
- m.jnt_solimp,
- m.jnt_range,
- m.jnt_margin,
- m.jnt_limited_ball_adr,
- m.dof_invweight0,
- d.njmax,
- d.qpos,
- d.qvel,
- refsafe,
- ],
- outputs=[
- d.nl,
- d.nefc,
- d.efc.type,
- d.efc.id,
- d.efc.J,
- d.efc.pos,
- d.efc.margin,
- d.efc.D,
- d.efc.vel,
- d.efc.aref,
- d.efc.frictionloss,
- ],
- )
+ if not (m.opt.disableflags & types.DisableBit.LIMIT):
+ wp.launch(
+ _efc_limit_ball,
+ dim=(d.nworld, m.jnt_limited_ball_adr.size),
+ inputs=[
+ m.nv,
+ m.opt.timestep,
+ m.jnt_qposadr,
+ m.jnt_dofadr,
+ m.jnt_solref,
+ m.jnt_solimp,
+ m.jnt_range,
+ m.jnt_margin,
+ m.jnt_limited_ball_adr,
+ m.dof_invweight0,
+ d.njmax,
+ d.qpos,
+ d.qvel,
+ refsafe,
+ ],
+ outputs=[
+ d.nl,
+ d.nefc,
+ d.efc.type,
+ d.efc.id,
+ d.efc.J,
+ d.efc.pos,
+ d.efc.margin,
+ d.efc.D,
+ d.efc.vel,
+ d.efc.aref,
+ d.efc.frictionloss,
+ ],
+ )
- limit_slide_hinge = m.jnt_limited_slide_hinge_adr.size > 0
- if limit_slide_hinge:
- wp.launch(
- _efc_limit_slide_hinge,
- dim=(d.nworld, m.jnt_limited_slide_hinge_adr.size),
- inputs=[
- m.nv,
- m.opt.timestep,
- m.jnt_qposadr,
- m.jnt_dofadr,
- m.jnt_solref,
- m.jnt_solimp,
- m.jnt_range,
- m.jnt_margin,
- m.jnt_limited_slide_hinge_adr,
- m.dof_invweight0,
- d.njmax,
- d.qpos,
- d.qvel,
- refsafe,
- ],
- outputs=[
- d.nl,
- d.nefc,
- d.efc.type,
- d.efc.id,
- d.efc.J,
- d.efc.pos,
- d.efc.margin,
- d.efc.D,
- d.efc.vel,
- d.efc.aref,
- d.efc.frictionloss,
- ],
- )
+ wp.launch(
+ _efc_limit_slide_hinge,
+ dim=(d.nworld, m.jnt_limited_slide_hinge_adr.size),
+ inputs=[
+ m.nv,
+ m.opt.timestep,
+ m.jnt_qposadr,
+ m.jnt_dofadr,
+ m.jnt_solref,
+ m.jnt_solimp,
+ m.jnt_range,
+ m.jnt_margin,
+ m.jnt_limited_slide_hinge_adr,
+ m.dof_invweight0,
+ d.njmax,
+ d.qpos,
+ d.qvel,
+ refsafe,
+ ],
+ outputs=[
+ d.nl,
+ d.nefc,
+ d.efc.type,
+ d.efc.id,
+ d.efc.J,
+ d.efc.pos,
+ d.efc.margin,
+ d.efc.D,
+ d.efc.vel,
+ d.efc.aref,
+ d.efc.frictionloss,
+ ],
+ )
- limit_tendon = m.tendon_limited_adr.size > 0
- if limit_tendon:
- wp.launch(
- _efc_limit_tendon,
- dim=(d.nworld, m.tendon_limited_adr.size),
- inputs=[
- m.nv,
- m.opt.timestep,
- m.jnt_dofadr,
- m.tendon_adr,
- m.tendon_num,
- m.tendon_limited_adr,
- m.tendon_solref_lim,
- m.tendon_solimp_lim,
- m.tendon_range,
- m.tendon_margin,
- m.tendon_invweight0,
- m.wrap_objid,
- m.wrap_type,
- d.njmax,
- d.qvel,
- d.ten_length,
- d.ten_J,
- refsafe,
- ],
- outputs=[
- d.nl,
- d.nefc,
- d.efc.type,
- d.efc.id,
- d.efc.J,
- d.efc.pos,
- d.efc.margin,
- d.efc.D,
- d.efc.vel,
- d.efc.aref,
- d.efc.frictionloss,
- ],
- )
+ wp.launch(
+ _efc_limit_tendon,
+ dim=(d.nworld, m.tendon_limited_adr.size),
+ inputs=[
+ m.nv,
+ m.opt.timestep,
+ m.jnt_dofadr,
+ m.tendon_adr,
+ m.tendon_num,
+ m.tendon_limited_adr,
+ m.tendon_solref_lim,
+ m.tendon_solimp_lim,
+ m.tendon_range,
+ m.tendon_margin,
+ m.tendon_invweight0,
+ m.wrap_objid,
+ m.wrap_type,
+ d.njmax,
+ d.qvel,
+ d.ten_length,
+ d.ten_J,
+ refsafe,
+ ],
+ outputs=[
+ d.nl,
+ d.nefc,
+ d.efc.type,
+ d.efc.id,
+ d.efc.J,
+ d.efc.pos,
+ d.efc.margin,
+ d.efc.D,
+ d.efc.vel,
+ d.efc.aref,
+ d.efc.frictionloss,
+ ],
+ )
# contact
- if not (m.opt.disableflags & types.DisableBit.CONTACT.value):
- if m.opt.cone == types.ConeType.PYRAMIDAL.value:
+ if not (m.opt.disableflags & types.DisableBit.CONTACT):
+ if m.opt.cone == types.ConeType.PYRAMIDAL:
wp.launch(
_efc_contact_pyramidal,
dim=(d.nconmax, 2 * (m.condim_max - 1) if m.condim_max > 1 else 1),
@@ -1863,7 +1840,7 @@ def make_constraint(m: types.Model, d: types.Data):
d.efc.frictionloss,
],
)
- elif m.opt.cone == types.ConeType.ELLIPTIC.value:
+ elif m.opt.cone == types.ConeType.ELLIPTIC:
wp.launch(
_efc_contact_elliptic,
dim=(d.nconmax, m.condim_max),
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py
index 2b8a3c26..4698c06b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint_test.py
@@ -21,10 +21,9 @@ import warp as wp
from absl.testing import absltest
from absl.testing import parameterized
-import mujoco_warp as mjwarp
-
-from mujoco.mjx.third_party.mujoco_warp._src import test_util
-from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
+import mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp import ConeType
+from mujoco.mjx.third_party.mujoco_warp import test_data
# tolerance for difference between MuJoCo and MJWarp constraint calculations,
# mostly due to float precision
@@ -143,13 +142,13 @@ class ConstraintTest(parameterized.TestCase):
"""
- _, mjd, m, d = test_util.fixture(xml=xml, cone=cone, keyframe=0)
+ _, mjd, m, d = test_data.fixture(xml=xml, keyframe=0, overrides={"opt.cone": cone})
# fill with nan to check whether we are not reading uninitialized values
for arr in (d.efc.J, d.efc.D, d.efc.aref, d.efc.pos, d.efc.margin):
arr.fill_(wp.nan)
- mjwarp.make_constraint(m, d)
+ mjw.make_constraint(m, d)
_assert_eq(d.ncon.numpy()[0], mjd.ncon, "ncon")
_assert_eq(d.efc.J.numpy()[0, : mjd.nefc, :].reshape(-1), mjd.efc_J, "efc_J")
@@ -165,14 +164,14 @@ class ConstraintTest(parameterized.TestCase):
def test_constraints(self, cone):
"""Test constraints."""
for key in range(3):
- _, mjd, m, d = test_util.fixture("constraints.xml", sparse=False, cone=cone, keyframe=key)
+ _, mjd, m, d = test_data.fixture("constraints.xml", keyframe=key, overrides={"opt.cone": cone})
for arr in (d.ne, d.nefc, d.nf, d.nl, d.efc.type):
arr.fill_(-1)
for arr in (d.efc.J, d.efc.D, d.efc.vel, d.efc.aref, d.efc.pos, d.efc.margin):
arr.fill_(wp.nan)
- mjwarp.make_constraint(m, d)
+ mjw.make_constraint(m, d)
_assert_eq(d.ne.numpy()[0], mjd.ne, "ne")
_assert_eq(d.nefc.numpy()[0], mjd.nefc, "nefc")
@@ -183,14 +182,14 @@ class ConstraintTest(parameterized.TestCase):
def test_limit_tendon(self):
"""Test limit tendon constraints."""
for keyframe in range(-1, 1):
- _, mjd, m, d = test_util.fixture("tendon/tendon_limit.xml", sparse=False, keyframe=keyframe)
+ _, mjd, m, d = test_data.fixture("tendon/tendon_limit.xml", keyframe=keyframe)
for arr in (d.nefc, d.nl, d.efc.type):
arr.fill_(-1)
for arr in (d.efc.J, d.efc.D, d.efc.vel, d.efc.aref, d.efc.pos, d.efc.margin):
arr.fill_(wp.nan)
- mjwarp.make_constraint(m, d)
+ mjw.make_constraint(m, d)
_assert_eq(d.nefc.numpy()[0], mjd.nefc, "nefc")
_assert_eq(d.nl.numpy()[0], mjd.nl, "nl")
@@ -199,7 +198,7 @@ class ConstraintTest(parameterized.TestCase):
def test_equality_tendon(self):
"""Test equality tendon constraints."""
- _, mjd, m, d = test_util.fixture(
+ _, mjd, m, d = test_data.fixture(
xml="""