Skip to content

combine tests for nxn_broadphase and sap_broadphase #219

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

combine tests for nxn_broadphase and sap_broadphase #219

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6ba8e0
combine tests for nxn_broadphase and sap_broadphase
thowell May 6, 2025
5203c7c
Merge remote-tracking branch 'upstream/main' into broadphase_test
thowell May 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
145 changes: 34 additions & 111 deletions mujoco_warp/_src/broadphase_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import numpy as np
import warp as wp
from absl.testing import absltest
from absl.testing import parameterized

import mujoco_warp as mjwarp

Expand All @@ -26,91 +27,12 @@
from . import test_util


class BroadphaseTest(absltest.TestCase):
def test_sap_broadphase(self):
"""Tests sap_broadphase."""
class BroadphaseTest(parameterized.TestCase):
@parameterized.parameters(collision_driver.nxn_broadphase, collision_driver.sap_broadphase)
def test_broadphase(self, broadphase):
"""Tests collision broadphase algorithms."""

_SAP_MODEL = """
<mujoco>
<worldbody>
<geom size="40 40 40" type="plane"/> <!- (0) intersects with nothing -->
<body pos="0 0 0.7">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (1) intersects with 2, 6, 7 -->
</body>
<body pos="0.1 0 0.7">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (2) intersects with 1, 6, 7 -->
</body>
<body pos="1.8 0 0.7">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (3) intersects with 4 -->
</body>
<body pos="1.6 0 0.7">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (4) intersects with 3 -->
</body>
<body pos="0 0 1.8">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (5) intersects with 7 -->
<geom size="0.5 0.5 0.5" type="box" pos="0 0 -1"/> <!- (6) intersects with 2, 1, 7 -->
</body>
<body pos="0 0.5 1.2">
<freejoint/>
<geom size="0.5 0.5 0.5" type="box"/> <!- (7) intersects with 5, 6 -->
</body>
</worldbody>
</mujoco>
"""

collision_pair = [
(0, 1),
(0, 2),
(0, 3),
(0, 4),
(0, 6),
(1, 2),
(1, 4),
(1, 5),
(1, 6),
(1, 7),
(2, 3),
(2, 4),
(2, 5),
(2, 6),
(2, 7),
(3, 4),
(4, 6),
(5, 7),
(6, 7),
]

_, _, m, d = test_util.fixture(xml=_SAP_MODEL)

mjwarp.sap_broadphase(m, d)

ncollision = d.ncollision.numpy()[0]
np.testing.assert_equal(ncollision, len(collision_pair), "ncollision")

for i in range(ncollision):
pair = d.collision_pair.numpy()[i]
if pair[0] > pair[1]:
pair_tuple = (int(pair[1]), int(pair[0]))
else:
pair_tuple = (int(pair[0]), int(pair[1]))

np.testing.assert_equal(
pair_tuple in collision_pair,
True,
f"geom pair {pair_tuple} not found in brute force results",
)

# TODO(team): test DisableBit.FILTERPARENT

def test_nxn_broadphase(self):
"""Tests nxn_broadphase."""

_NXN_MODEL = """
_XML = """
<mujoco>
<worldbody>
<body>
Expand Down Expand Up @@ -169,59 +91,60 @@ def test_nxn_broadphase(self):
</keyframe>
</mujoco>
"""

# one world and zero collisions
mjm, _, m, d0 = test_util.fixture(xml=_NXN_MODEL, keyframe=0)
collision_driver.nxn_broadphase(m, d0)
mjm, _, m, d0 = test_util.fixture(xml=_XML, keyframe=0)
broadphase(m, d0)
np.testing.assert_allclose(d0.ncollision.numpy()[0], 0)

# one world and one collision
_, mjd1, _, d1 = test_util.fixture(xml=_NXN_MODEL, keyframe=1)
collision_driver.nxn_broadphase(m, d1)
_, mjd1, _, d1 = test_util.fixture(xml=_XML, keyframe=1)
broadphase(m, d1)

np.testing.assert_allclose(d1.ncollision.numpy()[0], 1)
np.testing.assert_allclose(d1.collision_pair.numpy()[0][0], 0)
np.testing.assert_allclose(d1.collision_pair.numpy()[0][1], 1)

# one world and three collisions
_, mjd2, _, d2 = test_util.fixture(xml=_NXN_MODEL, keyframe=2)
collision_driver.nxn_broadphase(m, d2)
np.testing.assert_allclose(d2.ncollision.numpy()[0], 3)
np.testing.assert_allclose(d2.collision_pair.numpy()[0][0], 0)
np.testing.assert_allclose(d2.collision_pair.numpy()[0][1], 1)
np.testing.assert_allclose(d2.collision_pair.numpy()[1][0], 0)
np.testing.assert_allclose(d2.collision_pair.numpy()[1][1], 2)
np.testing.assert_allclose(d2.collision_pair.numpy()[2][0], 1)
np.testing.assert_allclose(d2.collision_pair.numpy()[2][1], 2)
_, mjd2, _, d2 = test_util.fixture(xml=_XML, keyframe=2)
broadphase(m, d2)

ncollision = d2.ncollision.numpy()[0]
np.testing.assert_allclose(ncollision, 3)

collision_pairs = [[0, 1], [0, 2], [1, 2]]
for i in range(ncollision):
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.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,
)
broadphase(m, d3)

ncollision = d3.ncollision.numpy()[0]
np.testing.assert_allclose(ncollision, 4)

collision_driver.nxn_broadphase(m, d3)
np.testing.assert_allclose(d3.ncollision.numpy()[0], 4)
np.testing.assert_allclose(d3.collision_pair.numpy()[0][0], 0)
np.testing.assert_allclose(d3.collision_pair.numpy()[0][1], 1)
np.testing.assert_allclose(d3.collision_pair.numpy()[1][0], 0)
np.testing.assert_allclose(d3.collision_pair.numpy()[1][1], 1)
np.testing.assert_allclose(d3.collision_pair.numpy()[2][0], 0)
np.testing.assert_allclose(d3.collision_pair.numpy()[2][1], 2)
np.testing.assert_allclose(d3.collision_pair.numpy()[3][0], 1)
np.testing.assert_allclose(d3.collision_pair.numpy()[3][1], 2)
collision_pairs = [[[0, 1]], [[0, 1], [0, 2], [1, 2]]]
worldids = [0, 1, 1, 1]
for i in range(ncollision):
worldid = d3.collision_worldid.numpy()[i]
self.assertTrue(worldid == worldids[i])
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, _, _, d4 = test_util.fixture(xml=_NXN_MODEL, keyframe=1)
mjm4, _, m4, d4 = test_util.fixture(xml=_XML, keyframe=1)
mjm4.geom_contype[:3] = 0
m4 = mjwarp.put_model(mjm4)

collision_driver.nxn_broadphase(m4, d4)
broadphase(m4, d4)
np.testing.assert_allclose(d4.ncollision.numpy()[0], 0)

# one world and one collision: geomtype ordering
_, _, _, d5 = test_util.fixture(xml=_NXN_MODEL, keyframe=3)
collision_driver.nxn_broadphase(m, d5)
_, _, _, d5 = test_util.fixture(xml=_XML, keyframe=3)
broadphase(m, d5)
np.testing.assert_allclose(d5.ncollision.numpy()[0], 1)
np.testing.assert_allclose(d5.collision_pair.numpy()[0][0], 3)
np.testing.assert_allclose(d5.collision_pair.numpy()[0][1], 2)
Expand Down
8 changes: 5 additions & 3 deletions mujoco_warp/_src/collision_driver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,10 +220,12 @@ def _sap_broadphase(
# get geom indices and swap if necessary
geom1 = sap_sort_index_in[worldid, i]
geom2 = sap_sort_index_in[worldid, j]

# find linear index of (geom1, geom2) in upper triangular nxn_pairid
if geom2 < geom1:
tmp = geom1
geom1 = geom2
geom2 = tmp
idx = _upper_tri_index(ngeom, geom2, geom1)
else:
idx = _upper_tri_index(ngeom, geom1, geom2)

# find linear index of (geom1, geom2) in upper triangular nxn_pairid
idx = _upper_tri_index(ngeom, geom1, geom2)
Expand Down
Loading