Baruch Tabanpour
a49f3f63fc
Remove disable_jit in test.
...
PiperOrigin-RevId: 614824680
Change-Id: I15f1441dfb1a70b98e5a221c0e15b06572860748
2024-03-11 16:03:40 -07:00
Baruch Tabanpour
22cd96fbbc
Speed up convex-convex collisions.
...
PiperOrigin-RevId: 613450874
Change-Id: Ie66e0a078e2ff51a136b949dd9459847e1fe957d
2024-03-06 21:48:51 -08:00
Baruch Tabanpour
224fbd00c8
Fix egregious bug in broadphase.
...
PiperOrigin-RevId: 612856292
Change-Id: I622ca00509c27f351ce5efefc139b2f2e141c8e9
2024-03-05 08:48:09 -08:00
Baruch Tabanpour
419be4c605
Add crude broadphase impl.
...
PiperOrigin-RevId: 611643683
Change-Id: Id7816dadd7e9f1b8855268b815353c61540c4e0c
2024-02-29 17:12:58 -08:00
Copybara-Service
f9c46e2abb
Merge pull request #1447 from jaro-sevcik:mjx-factor-m-parallel
...
PiperOrigin-RevId: 611348280
Change-Id: I19ef25c732d43bba0b0a8c296b197948c69d5b52
2024-02-28 22:31:01 -08:00
Jaroslav Sevcik
803f0cb0d9
Clean up factor_m
2024-02-28 01:19:12 -08:00
Copybara-Service
a5ea87b0e6
Merge pull request #1446 from jaro-sevcik:mjx-solve-m-parallel
...
PiperOrigin-RevId: 610984251
Change-Id: Ifd69480643df236fe2d6b9f77d1b029d2650acff
2024-02-27 22:21:11 -08:00
Baruch Tabanpour
2e43f18bee
Update barkour notebook.
...
PiperOrigin-RevId: 610906071
Change-Id: Ie67b86888e195bbd8a8e879547e65fbc2c864534
2024-02-27 16:09:20 -08:00
Baruch Tabanpour
2067e2083a
Add cameras to MJX. Fixes #1422
...
PiperOrigin-RevId: 610870732
Change-Id: Iaa39bd00f56046c595dc961281ac23446ca5bdf2
2024-02-27 14:17:21 -08:00
Jaroslav Sevcik
7688e03eea
Remove unnecessary parens
2024-02-27 00:10:21 -08:00
Jaroslav Sevcik
29d6e10375
Update independent elements in solve_m in parallel
2024-02-25 20:27:28 +00:00
Jaroslav Sevcik
13f256dd19
Parallel factor_m
2024-02-25 20:24:54 +00:00
Erik Frey
1b93226fb0
To save memory, avoid loading meshes that will never be used for collisions.
...
PiperOrigin-RevId: 607770456
Change-Id: I34c4a34d1820a9fd88d2ebeaffec9a64637d2614
2024-02-16 12:40:35 -08:00
Erik Frey
6a346c42e8
Adds support for mjOption.impratio to MJX.
...
PiperOrigin-RevId: 607456034
Change-Id: I7f6549a4332fd2bc0331d76a8906b971ce28dbcf
2024-02-15 14:28:36 -08:00
Erik Frey
d561995c03
Use less brittle search to determine whether an MJX field is a data or meta pytree node. Use correct type annotation for geom_convex_* fields in mjx.Model.
...
PiperOrigin-RevId: 607096855
Change-Id: I2e14b9b63e676b4ac3ec2a12aa07316e1d019a79
2024-02-14 13:57:55 -08:00
Baruch Tabanpour
f50a44e193
Allow jax config to choose precision.
...
PiperOrigin-RevId: 606765036
Change-Id: I36977851615700507eb2e83cc6d2d1ff8255be62
2024-02-13 15:24:02 -08:00
Erik Frey
47bb4a828e
Fixes to the MJX Euler integrator and Newton solver, and test improvements to catch future bugs like this.
...
PiperOrigin-RevId: 606389094
Change-Id: I50d0143f338fe29a823c930dba82995ce8a64ce2
2024-02-12 15:46:27 -08:00
Baruch Tabanpour
805d862d5f
Fix manifold point function.
...
PiperOrigin-RevId: 606366838
Change-Id: I8cf031f6f0fd0173777e0a9bb6d11c0c77c617ec
2024-02-12 14:32:24 -08:00
Erik Frey
e77c3cb258
Update MJX viewer to use mjx.get_data_into and fix bad function parameters.
...
PiperOrigin-RevId: 605461208
Change-Id: I7a4f6a1ed92382e522d50224e7b5cd73d8842750
2024-02-08 16:32:27 -08:00
Erik Frey
4933a2c7b6
Improve performance of getting and putting MJX device data by reducing cost of ncon.
...
PiperOrigin-RevId: 605454760
Change-Id: I5c52a84c12ab67d042fe43f6ada284b51a762294
2024-02-08 16:06:58 -08:00
Erik Frey
3e2d62c248
Improve performance of getting and putting MJX device data with large numpy arrays.
...
ndarray.tobytes() is faster than converting arrays to tuples by orders of magnitude.
PiperOrigin-RevId: 605437697
Change-Id: I7d42c10e178b1aec80ff38a82e863e0701c15e62
2024-02-08 15:02:21 -08:00
Erik Frey
ea69f20b43
Unrolled matmul for small matrices in MJX improves steps/sec by 10-20%.
...
PiperOrigin-RevId: 605033188
Change-Id: I9da723d01e1f6511abaf4522b4c8f829a8091772
2024-02-07 10:48:45 -08:00
Copybara-Service
19c41aa965
Merge pull request #1401 from driesmarzougui:main
...
PiperOrigin-RevId: 604771114
Change-Id: I9389828e9b588c22eb24a4074a23283943ca748d
2024-02-06 14:49:35 -08:00
Erik Frey
38cbc9e543
Improve MJX performance on NVIDIA GPU by enabling Triton GEMM.
...
PiperOrigin-RevId: 604691509
Change-Id: Iba11a9e1b79848624a1ef6eb4e61e2093d36e662
2024-02-06 10:31:49 -08:00
Dries Marzougui
ed36956b0d
fix https://github.com/google-deepmind/mujoco/issues/1379
2024-02-05 10:59:08 +01:00
Baruch Tabanpour
92c161a61b
Add ncon to __init__.
...
PiperOrigin-RevId: 603468404
Change-Id: I20c622b8f6f84ecacb9f3cdd26ac9560cf2788be
2024-02-01 13:38:23 -08:00
Baruch Tabanpour
8be966cdf9
Avoid NaN in gradient of where. Fixes #1349 .
...
PiperOrigin-RevId: 602461299
Change-Id: I7c976d3b7262027bdc63967dcc00674fb2802e70
2024-01-29 12:09:29 -08:00
Baruch Tabanpour
508669a9f2
Raise not implemented for frictionloss for now. #1344
...
PiperOrigin-RevId: 602453907
Change-Id: I0ea506b54f96c4fadc0f29f9ed0a21b0b6a7e5d9
2024-01-29 11:45:36 -08:00
Baruch Tabanpour
a1e5ede266
Allow Brax data structures to inherit from MJX data structures.
...
PiperOrigin-RevId: 601290735
Change-Id: I6dad1b6d7e6b3e71acecc36189ab59cc81ab90b4
2024-01-24 18:03:17 -08:00
Erik Frey
e59605135a
Make benchmarks and viewer more accessible.
...
- Drop google_benchmark module dependency, wasn't really using it.
- Move benchmark models into test_data
- Rename benchmark => testspeed, make default output more like MJ testspeed.
- Update viewer to use new io functions and add get_data_into.
- Add "mjx-testspeed" and "mjx-viewer" bin scripts
PiperOrigin-RevId: 601277135
Change-Id: I782f666c58a96ac1a0a1c5ffc8ecf29d9fcd810f
2024-01-24 17:04:02 -08:00
Erik Frey
0a7be1732c
Adds support for explicit dense/sparse mass matrices to MJX.
...
This increases performance, particularly for the Newton solver on TPU.
PiperOrigin-RevId: 600696483
Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff
2024-01-23 00:32:07 -08:00
Erik Frey
5cbaa23388
Remove TEST_FILES from MJX test_util.py.
...
PiperOrigin-RevId: 599886327
Change-Id: I4d166e09db22fffc959dbe34fa3469bd4a1b9f84
2024-01-19 11:15:06 -08:00
Baruch Tabanpour
a02fc405af
Add naive ray mesh implementation.
...
PiperOrigin-RevId: 598993152
Change-Id: I708505bf20a89d5a4840451a11b99771195003e6
2024-01-16 16:29:11 -08:00
Erik Frey
ee922be393
Benchmark improvements.
...
- Use multiple devices if present
- FLAGS more agnostic to diverse platforms
PiperOrigin-RevId: 598990299
Change-Id: Ifbd454ca521a788324f2b62f1c79e7d4b8305a30
2024-01-16 16:17:56 -08:00
Erik Frey
85ab6183bc
Add ray filtering based on geom transparency and filter params.
...
PiperOrigin-RevId: 597316447
Change-Id: If9619320f0024c486bb7abe2c63e9a1ba1a631f0
2024-01-10 11:59:11 -08:00
Erik Frey
feb92bf535
MJX ray for planes, spheres, capsules, and boxes.
...
PiperOrigin-RevId: 596004453
Change-Id: I6c15268f7ec6afc1aadd370a78ef225ab9c5fd04
2024-01-05 08:42:17 -08:00
Baruch Tabanpour
8ce2c92021
Add refsite transmission.
...
PiperOrigin-RevId: 595239860
Change-Id: Id1ca8fe2ffab4b55013e8987e92eb99847278dc4
2024-01-02 16:45:50 -08:00
Baruch Tabanpour
77b4132c47
Add site transmission to MJX.
...
PiperOrigin-RevId: 592373042
Change-Id: I54d941f4ee6f74404fe2aea7253f5b426d097e23
2023-12-19 16:19:17 -08:00
Baruch Tabanpour
80f50c943c
Add filterexact to MJX.
...
PiperOrigin-RevId: 592026712
Change-Id: Ia73d67f36db5a4a1786a2532b5daa9e9a36f714c
2023-12-18 15:29:18 -08:00
Erik Frey
0915d69c3f
Correctly transform site_xmat in MJX get_data/put_data.
...
PiperOrigin-RevId: 590750273
Change-Id: I836e1f5d20725ed9b93bf86fe3e8e61954a56b47
2023-12-13 16:34:42 -08:00
Erik Frey
b3ccf67ebf
Cleanup MJX tests and migrate them to put_model/put_data.
...
PiperOrigin-RevId: 588304729
Change-Id: I58075383e6eed64ae00cea305a568eba6465b0b0
2023-12-05 23:08:05 -08:00
Erik Frey
67fa7c1dc2
Add put_model, put_data, and get_data to MJX.
...
These io functions replace device_put, device_get_into. These new functions correctly copy intermediate fields such as efc_J, which the old functions could not.
PiperOrigin-RevId: 586712273
Change-Id: Ia3a4104af17db2ebd41b41a2ca501c5de9ac5d05
2023-11-30 10:14:46 -08:00
Erik Frey
2f37147fc0
Guard against NaNs in MJX solver when no cost is imposed. Fixes #1225 .
...
PiperOrigin-RevId: 586020958
Change-Id: I82f398b1538d2c3801a35c80c862e49be3b110a2
2023-11-28 10:12:17 -08:00
Baruch Tabanpour
35c90844ae
Add sites to MJX.
...
PiperOrigin-RevId: 585691582
Change-Id: I6cc594daf436231f91534f137ff259e7b39811d9
2023-11-27 10:28:50 -08:00
Baruch Tabanpour
4d2a176a8a
Modify name of mjx integration tests.
...
PiperOrigin-RevId: 583108693
Change-Id: If972c376181ca98201ff61ff06d200c82d214e2b
2023-11-16 11:24:07 -08:00
Erik Frey
c8146372cf
Simplify pendula test data. Fix bug in `scan.body_tree` that led to incorrect smooth dynamics for some kinematic tree layouts.
...
PiperOrigin-RevId: 582125477
Change-Id: I47505ceb4be4c88052bc670f401b8c45ff320bbd
2023-11-13 16:48:49 -08:00
Erik Frey
3c0a56c1e5
Adds support for Newton solver to MJX.
...
PiperOrigin-RevId: 578738033
Change-Id: Id0e68c7b961380cb1ed88b40d99f202c3fa8e228
2023-11-01 22:05:10 -07:00
Baruch Tabanpour
45878b7eef
Raise an exception in MJX if given a model with tendons.
...
Fixes #1149
PiperOrigin-RevId: 578278430
Change-Id: Ie633975b830379ca554438307a6e674675eaf8e7
2023-10-31 13:09:27 -07:00
Erik Frey
7069976557
Implement joint equality constraint. Streamline constraint functions.
...
PiperOrigin-RevId: 577905570
Change-Id: I70580d21e3c612606b967cdd84ad22794c672397
2023-10-30 11:38:15 -07:00
Yuval Tassa
86d9c84e59
Fix bug in implementation of torquescale
...
PiperOrigin-RevId: 577466789
Change-Id: I3c8e5dd9b4af9a629e777f7ccad8417075b0046c
2023-10-28 08:25:25 -07:00