From fd9b989bf860df0c368a644d691e69f49c4e1d2e Mon Sep 17 00:00:00 2001 From: Taylor Howell Date: Tue, 7 Oct 2025 09:48:21 +0100 Subject: [PATCH] Import google-deepmind/mujoco_warp from GitHub. PiperOrigin-RevId: 816099854 Change-Id: I6302a97c2f5d4594a08ebf4bc53a180b804360e6 --- .../mjx/third_party/mujoco_warp/__init__.py | 4 +- .../_src/{test_util.py => benchmark.py} | 174 +- .../mujoco_warp/_src/broadphase_test.py | 81 +- .../mujoco_warp/_src/collision_convex.py | 88 +- .../mujoco_warp/_src/collision_driver.py | 13 +- .../mujoco_warp/_src/collision_driver_test.py | 104 +- .../mujoco_warp/_src/collision_gjk.py | 160 +- .../mujoco_warp/_src/collision_gjk_legacy.py | 14 +- .../mujoco_warp/_src/collision_gjk_test.py | 98 +- .../mujoco_warp/_src/collision_primitive.py | 147 +- .../mujoco_warp/_src/collision_sdf.py | 76 +- .../mujoco_warp/_src/constraint.py | 391 +- .../mujoco_warp/_src/constraint_test.py | 23 +- .../mujoco_warp/_src/derivative.py | 32 +- .../third_party/mujoco_warp/_src/forward.py | 52 +- .../mujoco_warp/_src/forward_test.py | 183 +- .../mujoco_warp/_src/inverse_test.py | 64 +- .../mjx/third_party/mujoco_warp/_src/io.py | 249 +- .../mujoco_warp/_src/io_jax_test.py | 32 +- .../third_party/mujoco_warp/_src/io_test.py | 101 +- .../third_party/mujoco_warp/_src/jax_test.py | 19 +- .../third_party/mujoco_warp/_src/passive.py | 81 +- .../mujoco_warp/_src/passive_test.py | 41 +- .../mjx/third_party/mujoco_warp/_src/ray.py | 18 +- .../third_party/mujoco_warp/_src/ray_test.py | 66 +- .../third_party/mujoco_warp/_src/sensor.py | 595 ++- .../mujoco_warp/_src/sensor_test.py | 171 +- .../third_party/mujoco_warp/_src/smooth.py | 103 +- .../mujoco_warp/_src/smooth_test.py | 137 +- .../third_party/mujoco_warp/_src/solver.py | 40 +- .../mujoco_warp/_src/solver_test.py | 147 +- .../third_party/mujoco_warp/_src/support.py | 6 +- .../mujoco_warp/_src/support_test.py | 17 +- .../mjx/third_party/mujoco_warp/_src/types.py | 36 +- .../mujoco_warp/_src/unroll_test.py | 62 + .../third_party/mujoco_warp/_src/util_misc.py | 18 +- .../mujoco_warp/test_data/__init__.py | 104 + ...8426c7d466d59b03a7f048c5304d28205b3069.stl | Bin 0 -> 23384 bytes ...129cf77510d68c89cf2ddf67bd36dacef387b3.stl | Bin 0 -> 25684 bytes ...5f08a7e5e7488eb120fc62c3e6458e6be0f956.stl | Bin 0 -> 6484 bytes ...1ca8ed83291edde7bc9ab61d7d4828b37b0343.stl | Bin 0 -> 40084 bytes ...c8ad369e839057b4e10539978b87e39b5f81bf.stl | Bin 0 -> 40084 bytes ...e2ad6e912937b9c5f0ebde8e23f34ff982455a.stl | Bin 0 -> 40084 bytes ...e7ca5723c6e8bedf8a1138046f5419c7440b72.stl | Bin 0 -> 84084 bytes ...10c72c0b4ae36b2e4f79e29fe1fe8f3a9a678a.stl | Bin 0 -> 84084 bytes ...df49f763bf3784925b76eab16d8c28641a4529.stl | Bin 0 -> 40084 bytes ...8fe51a40113c78c2172ae79d2d0b8a5d9002d1.obj | 1338 +++++++ ...c239941aa3757c1a8c08c0e7ac9ce503ec8d05.obj | 570 +++ ...90f600631515ffcf851e485d44658c713792e2.obj | 1134 ++++++ ...78c968f22f8fb86ec096bfa4739e8fa1eeceaf.obj | 141 + ...d922b53e782e2add994cbafff464e869755941.obj | 411 ++ ...e56b18a79a8daec67f71cf3dfb02c50659863c.obj | 375 ++ ...8e1afc2bcaa23895712349e76fcd00eb0bd33f.obj | 345 ++ ...2d5fffa71966e308e7f2993748005d70a874c9.obj | 405 ++ ...a0de5070f55e2bb145b0122ecba1848e40cae6.obj | 420 ++ ...27db3114dd0fe49bd66244990f88c818c756b3.obj | 402 ++ ...ce18dc0d773030ae0aeb24475d7789b68e858e.obj | 327 ++ ...a6f38e04f915c4ac7c0a5134a0b680881190b8.obj | 435 +++ ...8f4a2bfbfe1a245ef5339cf88ba2d67f6047e3.obj | 882 +++++ ...44d4c9c075207d4618aa4990a635cce9621dd7.obj | 576 +++ ...987c61a726a3ac722b52fb4eee10763febd48d.obj | 603 +++ ...8a14acef9d0fb5ceb5658cdb1e61bbec00b058.obj | 972 +++++ ...77d2ce50ff654661aea7888c3851962f049c9c.obj | 489 +++ ...36653cef4a18fe2bc8967e34ca58404b73fa4d.obj | 567 +++ ...93d308c471712f4b27bd7ac69e433f779b31a6.obj | 903 +++++ ...3fcdc88dfbd4473e87852461b45193c22f7a5d.obj | 537 +++ ...4b94344d8aa949eb4e59ce7cfc72e6c6421751.obj | 705 ++++ ...39e444107a6087da79e6278eafc574c027bcd0.obj | 2712 +++++++++++++ ...56bb587dfb53ac2dc9a1b879f6ee4621776c47.obj | 3042 +++++++++++++++ ...d1a01665521786bae65a8926f051b0bfdbdbd6.obj | 2736 +++++++++++++ ...3ea4e7a3d15bd82c74f43117462d5098777802.obj | 3015 ++++++++++++++ ...5956077c32ec52183ba6330c26f2388b662199.obj | 1051 +++++ ...922fffdad3044d90075c80dcb505681d487cb5.obj | 141 + ...797f8472c190ed84c6122055337442976fd74e.obj | 2591 ++++++++++++ ...948b4fccbd63e4c4835435b6e32c87f0cd1b5b.obj | 585 +++ ...08605fa032bbd3c3458a0823fc3adb4e6095d1.obj | 618 +++ ...5221cb64e522621919eddad8b1cbb56367ee40.obj | 225 ++ ...ef67903a31437ad5b27e2f0fca54596d68a026.obj | 285 ++ ...4e24db22fafda8618c5c5c784debc3c74d4298.obj | 279 ++ ...c5cc87f26f9079dec574f7683dfd9e950cd10e.obj | 258 ++ ...8dfdca17412c2cacd2d3558f53110d9f9df0f2.obj | 264 ++ ...6e65dc60499d1b3ebabafd73eef5615cefc236.obj | 306 ++ ...04b26d016115d98af353973731e730d3a6dc0a.obj | 231 ++ ...37344b83d744900ac4fc2ca3397c2da1dba5cc.obj | 267 ++ ...4714646f4c886da3b1472b959a9ada8f08333b.obj | 258 ++ ...94f002a719c9606862e1fb1227f9110825bb79.obj | 594 +++ ...933c58c0e8ad8b2de88ff6af7f4293ad96929c.obj | 615 +++ ...65a7ef409ce22edd43097a0569168883c3eb5c.obj | 399 ++ ...3b05250829877106a511befd9cf6fa8823675a.obj | 282 ++ ...a5679233100eec0a3295abb54b1b7b9b4b0839.obj | 225 ++ ...74e4da93bd05a4065813d1d2efe3cd2102dfe7.obj | 300 ++ ...40e96d3a62e6fac0428a167cc13de2f120c57a.obj | 267 ++ ...df8358c6c7f7dd89929ae63b66ef778bd8746c.obj | 219 ++ ...923025c8be39433f5b8ef1fac9d74f5ba3556b.obj | 1443 +++++++ ...5abc551d7a44ee3d4d97301c152cb05af78767.obj | 522 +++ ...400f78b8859f07b401002cd19ffaa952da7468.obj | 537 +++ ...318237a2ff73b64244d87128194ea897852e93.obj | 543 +++ ...7c5e4cd695911b0efeafbf97c62e1a82418c67.obj | 516 +++ ...9963e0fb42ed6852e78415fdc74faa645347ee.obj | 387 ++ ...cda3d287d1cfcabe91c62bc3aa506569f518b7.obj | 522 +++ ...b9408de472c6c15fcfe0ddb302b65c5bf9e960.obj | 567 +++ ...82299967302621d8e6ed94305028a9b209dfd2.obj | 498 +++ ...c5d4471387579448d36399cc44ef392e8ef355.obj | 863 ++++ ...ce5ed902c7b7a8a0c407f48118fff68016dfec.obj | 462 +++ ...98452198415d0e7d49f87079b6eeca643a356f.obj | 459 +++ ...c4c0aeef953ee56e3531c0f2ae64b2dc3655ce.obj | 249 ++ ...7a70b9746eb429fd6fb1555649ae813a4cee43.obj | 231 ++ ...6a9bb9b03dc60fb32bb6ca9f7fddbd172f08fa.obj | 165 + ...0bb87696882db99697715207acc1d326b62947.obj | 207 + ...351b3484dabfb2999f03923638186f714e52ae.obj | 234 ++ ...ede1030270611430302fee1fc06cf8f6a8c101.obj | 195 + ...05d61a021487a768f29be67bf681bf19bbbd36.obj | 192 + ...f2ecf0799db93fd1d265370bb939e0f33c5c18.obj | 201 + ...412f7d3891d34d1425bf5829dcdb9576b7fee5.obj | 267 ++ ...4cedd984c14af625555605dec4650b7286d5fa.obj | 462 +++ ...1ea9d6f39011787da8fbb18f65ed5932cd23db.obj | 459 +++ ...78c38f2458a8e06326875928b64ea676c3a083.obj | 300 ++ ...19cc650173b632bc1329ddf4c2ddcf9337d5cd.obj | 234 ++ ...7124ea953dd835721ae4893a0d480c6202e9fc.obj | 195 + ...2ef675dbe9332f84d2ec2bd523ca31c5fc9c21.obj | 189 + ...f5b867c535a7b91a54511fd9e36eee82e9b943.obj | 174 + ...b5946fc3874b7ebfd1ac0f0a2bac3407b6d1ad.obj | 192 + ...5ee4e74c9ce3b86d1b7233bfccb4f6e11876c9.obj | 1946 +++++++++ ...e74c9ce3b86d1b7233bfccb4f6e11876c9.obj.mtl | 13 + ...44eed200b016b11d4936140ec35102f9e8d4ab.obj | 3459 ++++++++++++++++ ...ef9250fff030b5c0cc3df5565b7763875b1f11.obj | 906 +++++ ...95d51ef7bde4cdca2697b29a7e58076df9cb5f.obj | 1946 +++++++++ ...1ef7bde4cdca2697b29a7e58076df9cb5f.obj.mtl | 13 + ...9901813433e123bb23132256974b4cbb27f450.obj | 3471 ++++++++++++++++ ...f2f4ceb88b73a331d804f2c4658746126af80f.obj | 903 +++++ ...8a04b6db2eff1ca910caab9aa265d77782cc6d.obj | 1039 +++++ ...b903ddd2b86d74a023dc3caa2977ddb0414ae8.obj | 141 + ...204bd53e09bed87fb790d0204e6bc9aa92673b.obj | 1942 +++++++++ ...d53e09bed87fb790d0204e6bc9aa92673b.obj.mtl | 13 + ...8f82718a43061e3f7d46041e964fe213a45c18.obj | 3303 ++++++++++++++++ ...c4706030166e8c1b0834f343d6812fdc8623e0.obj | 816 ++++ ...ca1674b22089a5c0558949aa539d95f5fe9039.obj | 1942 +++++++++ ...74b22089a5c0558949aa539d95f5fe9039.obj.mtl | 13 + ...bf7ffccaf2bf740bd3134f32fa2c221678fa0f.obj | 3474 +++++++++++++++++ ...9061bdc92ac3a5e4ee6614811f4d022863ac30.obj | 891 +++++ ...c36bb29d030fdd3ee6d3150c46429a2c01fb6a.obj | 1039 +++++ ...dd02393237429b33a2296feda0a0f2e3887e90.obj | 141 + ...c614ac600458b2d5a77e9abf8ac3930d0b3f01.obj | 1198 ++++++ ...ac600458b2d5a77e9abf8ac3930d0b3f01.obj.mtl | 13 + ...fd8fe1964def48aa4fe42635f9542bc6015573.obj | 411 ++ ...3900e1c3bad73bd2b2a016cf4ae951dc16a798.obj | 354 ++ ...c34f530a06ad54b80a9006aaf215bbac9509fa.obj | 447 +++ ...2322ca1c81b79a05bff62846ac490b8d1ac3d4.obj | 435 +++ ...c2a944b721ef784d23454f3252f30bae943bd2.obj | 393 ++ ...80b71c4d74bf20ac97b1f0875e34a789547c53.obj | 399 ++ ...99b40aab7dab531fef77ff92bdec5a3410b9bc.obj | 384 ++ ...8b7651225ffadfd716665f99970af90a86e52f.obj | 276 ++ ...ede03f8a622ddcc01fc62430a1d6c807741421.obj | 327 ++ ...e7a1ee44f3e24dd5a8973509e8efd2b0169c8c.obj | 1184 ++++++ ...ee44f3e24dd5a8973509e8efd2b0169c8c.obj.mtl | 13 + ...dd649f5629943594425caa55f1e4d5a4e93e0d.obj | 753 ++++ ...9b4cd611d64e52b1529afd3668ba7bed221919.obj | 627 +++ ...fcfa80779b37d0b35006ff6ff99d1e2d77e774.obj | 612 +++ ...8912a3efd6064071faf75a14023617d7a29750.obj | 795 ++++ ...d16325a58a37ef9fa2bbdaf15b6d20f798e33e.obj | 753 ++++ ...44cc1653ede725016118d7cdd0eceb1af55583.obj | 672 ++++ ...fef0d9372bf84dec2f832426530128338bbfb2.obj | 654 ++++ ...b4279ab8f42b3a35de1b776130701108a50bc3.obj | 720 ++++ ...2d72dd026c36e9694154a5d8aba6193d1c1878.obj | 675 ++++ ...38c6eae9bf1f3ae55798660678516b5820fc46.obj | 417 ++ ...6d2b3c89ad237434b655c8490aa663274cb58f.obj | 372 ++ ...c9e799427a1745776a8110a2b9aeb34cb5d3b3.obj | 378 ++ ...e866ca632c952bcd7e8fb4d08b1c02ca70230f.obj | 405 ++ ...b9a772794eb9b36235a69486bb86d42e381ad2.obj | 384 ++ ...6a8bb29df4bfb7dcc3af405310880df1109f0b.obj | 327 ++ ...5d98f451863b8794bb95c618796c8c752c6518.obj | 423 ++ ...466ad73a6129bf39af3c092662bc00b560a663.obj | 219 ++ ...396c1ac75f494900bf8db473bcc89d6412342d.obj | 300 ++ ...66afde19240856451b0fc8f5e31590e7e998f0.obj | 141 + ...df90098757a980b68bc9b546bddce3257c7263.stl | Bin 0 -> 25684 bytes .../mujoco_warp/test_data/aloha_pot/scene.xml | 635 +++ ...055adb309111efb32fb9fc0417883865cf9bf2.obj | 414 ++ ...aaf6399ee429de60273fe61f29a7a7fcd20440.obj | 370 ++ ...72d9cc47a6f6792f1b4a7fd4ac6e6862d06105.obj | 190 + ...11efddeb54716caf91218e0beee848e0083aad.stl | Bin 0 -> 49984 bytes ...ee3af162e071b59ae35b1b7279fe257c375ffb.stl | Bin 0 -> 63884 bytes ...b3692a22ad52ba9b3c439033c552c51222ed28.stl | Bin 0 -> 51384 bytes ...4b69653fd0cc02e83bb7d65ba3e7ec1604b71c.stl | Bin 0 -> 49584 bytes ...55d345ad2564df22bcbd24557f5a69d362ef68.stl | Bin 0 -> 49984 bytes ...facdc0498ab0119fb164f573de994a27c7401e.stl | Bin 0 -> 70784 bytes ...f4b35b626deb7894668277f953fc4143ad9898.stl | Bin 0 -> 25684 bytes ...8e61f159d0594d74deb8fabda682e9971dc2b5.stl | Bin 0 -> 25634 bytes ...d29471843d7ca6e0e8a970bdc9256a822c8b15.stl | Bin 0 -> 25684 bytes ...982864f8591ef106a4b771fc7502c724ef3d7c.stl | Bin 0 -> 25684 bytes ...349941da5624f43b137273d462f62ab9605c36.stl | Bin 0 -> 25684 bytes ...c7acdce0b7e4d90ec0ec56b4fcd8a7ff0cf2e7.stl | Bin 0 -> 25684 bytes .../test_data/collision_sdf/utils.py | 24 +- .../mjx/third_party/mujoco_warp/viewer.py | 53 +- 193 files changed, 92473 insertions(+), 1599 deletions(-) rename mjx/mujoco/mjx/third_party/mujoco_warp/_src/{test_util.py => benchmark.py} (58%) create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/_src/unroll_test.py create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/__init__.py create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/angled_extrusion-2b8426c7d466d59b03a7f048c5304d28205b3069.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/corner_bracket-c4129cf77510d68c89cf2ddf67bd36dacef387b3.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/d405_solid-a35f08a7e5e7488eb120fc62c3e6458e6be0f956.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_1000-421ca8ed83291edde7bc9ab61d7d4828b37b0343.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_1220-a3c8ad369e839057b4e10539978b87e39b5f81bf.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_150-62e2ad6e912937b9c5f0ebde8e23f34ff982455a.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_2040_1000-3ae7ca5723c6e8bedf8a1138046f5419c7440b72.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_2040_880-df10c72c0b4ae36b2e4f79e29fe1fe8f3a9a678a.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/extrusion_600-87df49f763bf3784925b76eab16d8c28641a4529.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-10_col_collision_0-208fe51a40113c78c2172ae79d2d0b8a5d9002d1.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-10_col_collision_1-e1c239941aa3757c1a8c08c0e7ac9ce503ec8d05.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-10_col_collision_2-e190f600631515ffcf851e485d44658c713792e2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-11_col_collision_0-a478c968f22f8fb86ec096bfa4739e8fa1eeceaf.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_0-0cd922b53e782e2add994cbafff464e869755941.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_1-88e56b18a79a8daec67f71cf3dfb02c50659863c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_2-ec8e1afc2bcaa23895712349e76fcd00eb0bd33f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_3-722d5fffa71966e308e7f2993748005d70a874c9.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_4-61a0de5070f55e2bb145b0122ecba1848e40cae6.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_5-f927db3114dd0fe49bd66244990f88c818c756b3.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_6-cfce18dc0d773030ae0aeb24475d7789b68e858e.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-13_col_collision_7-d3a6f38e04f915c4ac7c0a5134a0b680881190b8.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_0-778f4a2bfbfe1a245ef5339cf88ba2d67f6047e3.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_1-ee44d4c9c075207d4618aa4990a635cce9621dd7.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_2-74987c61a726a3ac722b52fb4eee10763febd48d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_3-928a14acef9d0fb5ceb5658cdb1e61bbec00b058.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_4-8977d2ce50ff654661aea7888c3851962f049c9c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_5-c336653cef4a18fe2bc8967e34ca58404b73fa4d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_6-6493d308c471712f4b27bd7ac69e433f779b31a6.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_7-403fcdc88dfbd4473e87852461b45193c22f7a5d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-14_col_collision_8-634b94344d8aa949eb4e59ce7cfc72e6c6421751.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-17_col_collision_0-b439e444107a6087da79e6278eafc574c027bcd0.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-17_col_collision_1-ca56bb587dfb53ac2dc9a1b879f6ee4621776c47.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-17_col_collision_2-98d1a01665521786bae65a8926f051b0bfdbdbd6.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-17_col_collision_3-213ea4e7a3d15bd82c74f43117462d5098777802.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-19_0-7a5956077c32ec52183ba6330c26f2388b662199.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-19_col_collision_0-ea922fffdad3044d90075c80dcb505681d487cb5.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_0-51797f8472c190ed84c6122055337442976fd74e.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_0-c8948b4fccbd63e4c4835435b6e32c87f0cd1b5b.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_1-2708605fa032bbd3c3458a0823fc3adb4e6095d1.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_10-225221cb64e522621919eddad8b1cbb56367ee40.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_11-e3ef67903a31437ad5b27e2f0fca54596d68a026.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_12-b44e24db22fafda8618c5c5c784debc3c74d4298.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_13-cbc5cc87f26f9079dec574f7683dfd9e950cd10e.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_14-8e8dfdca17412c2cacd2d3558f53110d9f9df0f2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_15-3a6e65dc60499d1b3ebabafd73eef5615cefc236.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_16-3104b26d016115d98af353973731e730d3a6dc0a.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_17-df37344b83d744900ac4fc2ca3397c2da1dba5cc.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_18-e24714646f4c886da3b1472b959a9ada8f08333b.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_2-3b94f002a719c9606862e1fb1227f9110825bb79.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_3-e4933c58c0e8ad8b2de88ff6af7f4293ad96929c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_4-5465a7ef409ce22edd43097a0569168883c3eb5c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_5-d03b05250829877106a511befd9cf6fa8823675a.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_6-12a5679233100eec0a3295abb54b1b7b9b4b0839.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_7-2374e4da93bd05a4065813d1d2efe3cd2102dfe7.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_8-d140e96d3a62e6fac0428a167cc13de2f120c57a.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-21_col_collision_9-6edf8358c6c7f7dd89929ae63b66ef778bd8746c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_0-b6923025c8be39433f5b8ef1fac9d74f5ba3556b.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_0-db5abc551d7a44ee3d4d97301c152cb05af78767.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_1-96400f78b8859f07b401002cd19ffaa952da7468.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_2-8a318237a2ff73b64244d87128194ea897852e93.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_3-df7c5e4cd695911b0efeafbf97c62e1a82418c67.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_4-dc9963e0fb42ed6852e78415fdc74faa645347ee.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_5-6bcda3d287d1cfcabe91c62bc3aa506569f518b7.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_6-60b9408de472c6c15fcfe0ddb302b65c5bf9e960.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-22_col_collision_7-9382299967302621d8e6ed94305028a9b209dfd2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_0-34c5d4471387579448d36399cc44ef392e8ef355.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_0-12ce5ed902c7b7a8a0c407f48118fff68016dfec.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_1-d498452198415d0e7d49f87079b6eeca643a356f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_10-73c4c0aeef953ee56e3531c0f2ae64b2dc3655ce.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_11-8b7a70b9746eb429fd6fb1555649ae813a4cee43.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_12-486a9bb9b03dc60fb32bb6ca9f7fddbd172f08fa.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_13-920bb87696882db99697715207acc1d326b62947.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_14-f1351b3484dabfb2999f03923638186f714e52ae.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_15-f8ede1030270611430302fee1fc06cf8f6a8c101.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_16-d605d61a021487a768f29be67bf681bf19bbbd36.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_17-c3f2ecf0799db93fd1d265370bb939e0f33c5c18.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_18-bb412f7d3891d34d1425bf5829dcdb9576b7fee5.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_2-6f4cedd984c14af625555605dec4650b7286d5fa.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_3-b71ea9d6f39011787da8fbb18f65ed5932cd23db.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_4-ea78c38f2458a8e06326875928b64ea676c3a083.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_5-d519cc650173b632bc1329ddf4c2ddcf9337d5cd.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_6-5a7124ea953dd835721ae4893a0d480c6202e9fc.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_7-552ef675dbe9332f84d2ec2bd523ca31c5fc9c21.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_8-a4f5b867c535a7b91a54511fd9e36eee82e9b943.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-23_col_collision_9-c3b5946fc3874b7ebfd1ac0f0a2bac3407b6d1ad.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-25_0-b85ee4e74c9ce3b86d1b7233bfccb4f6e11876c9.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-25_0-b85ee4e74c9ce3b86d1b7233bfccb4f6e11876c9.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-25_col_collision_0-0f44eed200b016b11d4936140ec35102f9e8d4ab.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-25_col_collision_1-16ef9250fff030b5c0cc3df5565b7763875b1f11.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-26_0-0d95d51ef7bde4cdca2697b29a7e58076df9cb5f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-26_0-0d95d51ef7bde4cdca2697b29a7e58076df9cb5f.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-26_col_collision_0-259901813433e123bb23132256974b4cbb27f450.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-26_col_collision_1-87f2f4ceb88b73a331d804f2c4658746126af80f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-27_0-078a04b6db2eff1ca910caab9aa265d77782cc6d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-27_col_collision_0-c5b903ddd2b86d74a023dc3caa2977ddb0414ae8.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-29_0-59204bd53e09bed87fb790d0204e6bc9aa92673b.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-29_0-59204bd53e09bed87fb790d0204e6bc9aa92673b.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-29_col_collision_0-898f82718a43061e3f7d46041e964fe213a45c18.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-29_col_collision_1-c2c4706030166e8c1b0834f343d6812fdc8623e0.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-30_0-2cca1674b22089a5c0558949aa539d95f5fe9039.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-30_0-2cca1674b22089a5c0558949aa539d95f5fe9039.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-30_col_collision_0-18bf7ffccaf2bf740bd3134f32fa2c221678fa0f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-30_col_collision_1-249061bdc92ac3a5e4ee6614811f4d022863ac30.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-31_0-fac36bb29d030fdd3ee6d3150c46429a2c01fb6a.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-31_col_collision_0-d7dd02393237429b33a2296feda0a0f2e3887e90.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_0-98c614ac600458b2d5a77e9abf8ac3930d0b3f01.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_0-98c614ac600458b2d5a77e9abf8ac3930d0b3f01.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_0-03fd8fe1964def48aa4fe42635f9542bc6015573.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_1-543900e1c3bad73bd2b2a016cf4ae951dc16a798.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_2-6cc34f530a06ad54b80a9006aaf215bbac9509fa.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_3-e12322ca1c81b79a05bff62846ac490b8d1ac3d4.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_4-edc2a944b721ef784d23454f3252f30bae943bd2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_5-7b80b71c4d74bf20ac97b1f0875e34a789547c53.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_6-6a99b40aab7dab531fef77ff92bdec5a3410b9bc.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_7-2a8b7651225ffadfd716665f99970af90a86e52f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-5_col_collision_8-f3ede03f8a622ddcc01fc62430a1d6c807741421.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_0-15e7a1ee44f3e24dd5a8973509e8efd2b0169c8c.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_0-15e7a1ee44f3e24dd5a8973509e8efd2b0169c8c.obj.mtl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_0-91dd649f5629943594425caa55f1e4d5a4e93e0d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_1-d99b4cd611d64e52b1529afd3668ba7bed221919.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_2-a2fcfa80779b37d0b35006ff6ff99d1e2d77e774.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_3-118912a3efd6064071faf75a14023617d7a29750.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_4-a5d16325a58a37ef9fa2bbdaf15b6d20f798e33e.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_5-da44cc1653ede725016118d7cdd0eceb1af55583.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_6-ebfef0d9372bf84dec2f832426530128338bbfb2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-6_col_collision_7-dbb4279ab8f42b3a35de1b776130701108a50bc3.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_0-d32d72dd026c36e9694154a5d8aba6193d1c1878.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_0-3638c6eae9bf1f3ae55798660678516b5820fc46.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_1-616d2b3c89ad237434b655c8490aa663274cb58f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_2-25c9e799427a1745776a8110a2b9aeb34cb5d3b3.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_3-27e866ca632c952bcd7e8fb4d08b1c02ca70230f.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_4-06b9a772794eb9b36235a69486bb86d42e381ad2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_5-4b6a8bb29df4bfb7dcc3af405310880df1109f0b.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_6-505d98f451863b8794bb95c618796c8c752c6518.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_7-2e466ad73a6129bf39af3c092662bc00b560a663.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-7_col_collision_8-7f396c1ac75f494900bf8db473bcc89d6412342d.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/original-8_col_collision_0-3e66afde19240856451b0fc8f5e31590e7e998f0.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/overhead_mount-2edf90098757a980b68bc9b546bddce3257c7263.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/scene.xml create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/table_mesh_mapped-6c055adb309111efb32fb9fc0417883865cf9bf2.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/tablelegs-d0aaf6399ee429de60273fe61f29a7a7fcd20440.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/tabletop-2972d9cc47a6f6792f1b4a7fd4ac6e6862d06105.obj create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_1_base-8911efddeb54716caf91218e0beee848e0083aad.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_2_shoulder-5dee3af162e071b59ae35b1b7279fe257c375ffb.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_3_upper_arm-f0b3692a22ad52ba9b3c439033c552c51222ed28.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_4_upper_forearm-6d4b69653fd0cc02e83bb7d65ba3e7ec1604b71c.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_5_lower_forearm-8555d345ad2564df22bcbd24557f5a69d362ef68.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_6_wrist-b3facdc0498ab0119fb164f573de994a27c7401e.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_7_gripper_bar-c5f4b35b626deb7894668277f953fc4143ad9898.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_7_gripper_prop-fb8e61f159d0594d74deb8fabda682e9971dc2b5.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_7_gripper_wrist_mount-61d29471843d7ca6e0e8a970bdc9256a822c8b15.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_8_custom_finger_left-bf982864f8591ef106a4b771fc7502c724ef3d7c.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/vx300s_8_custom_finger_right-41349941da5624f43b137273d462f62ab9605c36.stl create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/test_data/aloha_pot/wormseye_mount-4ac7acdce0b7e4d90ec0ec56b4fcd8a7ff0cf2e7.stl 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="""