Commit Graph

518 Commits

Author SHA1 Message Date
Alessio Quaglino 0e3306b62d Add useful flex arrays to MJX API.
This enables using name2id for flexes.

PiperOrigin-RevId: 893923892
Change-Id: I7c2a71791e26fa2280cfa75fd66db328f4599b0b
2026-04-03 00:13:29 -07:00
Taylor Howell e9de329e4e Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 892971818
Change-Id: Ibce0c41294a4fcf2825a246a9533e8927fdc7e82
2026-04-01 09:33:20 -07:00
Taylor Howell 0b54638df4 Import NVIDIA/warp from GitHub.
PiperOrigin-RevId: 892794388
Change-Id: Ic681071dc68bfaecaaab4e131c1c464dca99f9b5
2026-04-01 01:54:01 -07:00
Peter Hawkins 47264877a1 [JAX] Replace jnp.clip(..., a_min=..., a_max=...) with jnp.clip(..., min=..., max=...).
a_min and a_max are deprecated parameter names to jax.numpy.clip.

PiperOrigin-RevId: 891786675
Change-Id: Icc06efe0381c0db06f4aeb16f281a4d01ea37539
2026-03-30 10:28:24 -07:00
Taylor Howell 57bdf89d07 Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 885715192
Change-Id: I13fcb8312c1515bb1bd2d6cdfd0d256e37cb74de
2026-03-18 11:58:25 -07:00
Yuval Tassa 206e245353 Internal change.
PiperOrigin-RevId: 885152209
Change-Id: I56941ba443fe3daad02eff39d6dbcaa9615df41b
2026-03-17 12:29:37 -07:00
Yuval Tassa efae9157a7 Polynomial stiffness and damping https://youtu.be/aKa3ZlEF9_Y
PiperOrigin-RevId: 884607673
Change-Id: If8088dbf37fed1055304778a7eb84dec52cba920
2026-03-16 13:25:25 -07:00
Baruch Tabanpour 7324864366 Add random seed to SensorTest to prevent flakiness.
PiperOrigin-RevId: 884055306
Change-Id: Ia053d2661ed30835e66a145659c1e1b7275ab251
2026-03-15 11:23:44 -07:00
Michael Moss 6d70e66810 Update MuJoCo version from 3.6.0 to 3.7.0 following the 3.6.0 release
PiperOrigin-RevId: 881995135
Change-Id: I0c1a3b4a22c4c7cfada580b49a2cacde64cc84f3
2026-03-11 07:09:24 -07:00
Michael Moss 955d6c0b02 Update MuJoCo version from 3.5.1 to 3.6.0 prior to next release.
The next release will include breaking changes, so this bumps the major
version to reflect that.

PiperOrigin-RevId: 881351265
Change-Id: I970c68a0211651c07fcb8109ebd0b138196cf50a
2026-03-10 04:30:07 -07:00
Sam Haves dade2d8c3e MJX: Use content-hash for np.ndarray pytree metadata fields
Previously mjx dataclasses would return np.ndarray as bytes for jax tracing to hash since they are not inherently hashable. This meant that any tracing that was cached would hold on to a copy of the numpy arrays in the dataclass. Children such as mjx.Model that store large numpy arrays would end up duplicating that data 3+ times in some cases.

This eliminates O(N * array_size) memory duplication across N cached
pytree traces, saving a lot of memory on models with heavy mesh/texture data.

PiperOrigin-RevId: 880985898
Change-Id: I58d4e91fda4112e4c633818f92907d20138e962a
2026-03-09 12:35:16 -07:00
Alessio Quaglino 84f09f3fa6 Add mesh_normalnum to MJX API.
PiperOrigin-RevId: 880957335
Change-Id: Ib3c37afa593c888ec109f520947958a5d47b0be2
2026-03-09 11:34:10 -07:00
Baruch Tabanpour 6bc6c11c49 Add naccdmax to mjx io.
PiperOrigin-RevId: 879780258
Change-Id: I24f9d0b3205577322bebadc3d03796c8dcadde48
2026-03-06 13:02:55 -08:00
Baruch Tabanpour a8d6b0c848 Fix extra dim for depth, matches rgb output shape.
PiperOrigin-RevId: 878098932
Change-Id: Ifcc4ba39dbee1beeea88fda44c3d20b88179705e
2026-03-03 13:26:41 -08:00
Baruch Tabanpour 67b6ab0589 Update mjx-warp docs.
PiperOrigin-RevId: 875966048
Change-Id: I333e8562f3db333f6c64af67cee1f735cb21232d
2026-02-26 17:26:31 -08:00
Baruch Tabanpour 07cda5a3c4 Refactor MJX warp RenderContext, get rid of brittle pytree owner pattern.
PiperOrigin-RevId: 875901011
Change-Id: I2ba0734d9ca001ff91361aab6412fe370919344b
2026-02-26 14:49:32 -08:00
Taylor Howell 52ed96bc3a Move tendon sparsity fields from mjData to mjModel
PiperOrigin-RevId: 875087590
Change-Id: I1a5489d2d2011795ee38b09d547e68acd72e3cc7
2026-02-25 04:32:11 -08:00
Google DeepMind 94acdb33bd Fix Impl.CPP memory leaks due to mjModel/mjData not clearing from pytree.
PiperOrigin-RevId: 874622402
Change-Id: I38a00633dea77d940e4b0422f93e9456ad84daf9
2026-02-24 08:17:06 -08:00
Michael Moss a043df6bdf Update supported Python versions and remove obsolete 3.9 references.
PiperOrigin-RevId: 874063785
Change-Id: Ia4b2956b632aae5f0211fa766b6d413edcba8fe2
2026-02-23 07:29:47 -08:00
Taylor Howell 2bf9539234 MJX put_data with impl='warp'
PiperOrigin-RevId: 874017611
Change-Id: I23a25727d41cd888d90f2ea15b2b534eba5583be
2026-02-23 05:20:07 -08:00
Baruch Tabanpour 0f36b0fc2e Import NVIDIA/warp from GitHub.
PiperOrigin-RevId: 872969263
Change-Id: Ic890cb53c32593c158a1f3926906b492c7f85348
2026-02-20 10:29:13 -08:00
Baruch Tabanpour e5a236774d Add has_side_effect for pmap bvh refit in mjx-warp.
PiperOrigin-RevId: 872934001
Change-Id: If48cab439718b9e783c4b1da4fb1b1a08521a602
2026-02-20 09:21:43 -08:00
Taylor Howell fe50efd41a Sparse tendon Jacobian update for MuJoCo Warp
PiperOrigin-RevId: 872550625
Change-Id: Ifc5afacb048bec36ef6455072b188396db958002
2026-02-19 13:52:25 -08:00
Baruch Tabanpour 62a32386d6 Add pmap support for mjx-warp render.
PiperOrigin-RevId: 872533248
Change-Id: I828303866c31d33d2367dd9fba4eb42028c67540
2026-02-19 13:14:22 -08:00
Taylor Howell 9efe41c0c1 Remove dense code path for tendon Jacobian
PiperOrigin-RevId: 872444531
Change-Id: I6180101abc49469a72aee8bec5726e1e94f142ec
2026-02-19 09:53:49 -08:00
Baruch Tabanpour 3c858f871b Keep render_util arg order consistent between mujoco_warp and mjx-warp.
PiperOrigin-RevId: 872055868
Change-Id: Ic5b4bdedbdb353dc0b1385cf555261a86407cfdb
2026-02-18 14:57:06 -08:00
Baruch Tabanpour 485e2eb028 Fix bug in mjx-warp render vmap.
PiperOrigin-RevId: 871901512
Change-Id: Ifc25e98ffd89284cc6db9248bcca352bb8811e7e
2026-02-18 09:10:34 -08:00
Taylor Howell 8541128853 Update _get_data_into_warp. Fixes #3109
PiperOrigin-RevId: 871790690
Change-Id: I225579f7559ca5c1de92af6606d02ff763f9f4d1
2026-02-18 04:04:29 -08:00
Baruch Tabanpour f6129596a6 Save depth images in visualize_render.py
PiperOrigin-RevId: 871454618
Change-Id: If844e2c4af10e271536c4e6a36e760e8d4d50374
2026-02-17 12:55:39 -08:00
Taylor Howell af802a5386 Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 871338519
Change-Id: I6ef9b4a4c3b623a87f4c43efa798b6757ac534ae
2026-02-17 08:48:20 -08:00
Taylor Howell 66a413419a Reorder Option fields
PiperOrigin-RevId: 870779423
Change-Id: Ie356725810a820e910396abd9daef0f9c9f94abc
2026-02-16 02:18:23 -08:00
Baruch Tabanpour 538c131415 Last bits of mjx-warp render codegen.
PiperOrigin-RevId: 869900622
Change-Id: If86942fddc45b9d71505d7fc4e2345dc120ccab2
2026-02-13 14:55:26 -08:00
Baruch Tabanpour bbfe7fe0e2 Update mjx-warp render with partial codegen.
PiperOrigin-RevId: 869874195
Change-Id: I71a16a6c192873786a09211dc4d326c3a4115ad9
2026-02-13 13:50:07 -08:00
Baruch Tabanpour fb47ff01cd initial implementation of mjx-warp render.
PiperOrigin-RevId: 869822624
Change-Id: Ic20db1dcc2e7287752b273c4ad986066cceccad9
2026-02-13 11:40:55 -08:00
Taylor Howell 5903d4826f Sparse ten_J and ten_J_colind
PiperOrigin-RevId: 869712136
Change-Id: Id2979684aa0c39552cb8453852483cbac55f0ea3
2026-02-13 06:57:22 -08:00
Michael Moss 42958aab81 Update MuJoCo version from 3.5.0 to 3.5.1 following the 3.5.0 release
PiperOrigin-RevId: 869654894
Change-Id: I9cacaf4cbb70ad17faa24c3c30ec38ffe91b5deb
2026-02-13 03:54:25 -08:00
Taylor Howell 881544c0c5 Modify MuJoCo Warp dev version check
PiperOrigin-RevId: 869311325
Change-Id: I2f34bd932d7998d37b96b36e5d561f734567b8d5
2026-02-12 11:39:28 -08:00
Yuval Tassa 6af0d4c823 Replace margin/gap max with sum.
PiperOrigin-RevId: 868842521
Change-Id: Ia4a4ca5c6821c5a7b3ba7b0af7f71d981c2866c1
2026-02-11 14:22:27 -08:00
Google DeepMind 336e1026a4 Sparse ten_J and ten_J_colind memory
PiperOrigin-RevId: 868834376
Change-Id: Iec92d9e09e3134666895e54a78cb02439d74a4de
2026-02-11 14:05:43 -08:00
Taylor Howell d011285380 Sparse ten_J and ten_J_colind memory
PiperOrigin-RevId: 868800229
Change-Id: Id83f9fd40dd9bece06fd99753e4202eea17ca884
2026-02-11 12:45:50 -08:00
Taylor Howell 94e834bf4f Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 868759523
Change-Id: I83a9f014840e66d7775a3684aa2e59ec5a077191
2026-02-11 11:14:24 -08:00
Baruch Tabanpour 210cf86486 Use kernel cache dir per mjx-warp test fixture.
PiperOrigin-RevId: 868259579
Change-Id: I9f1346899591ca603ad2e34949d22b78c464c2ad
2026-02-10 12:10:28 -08:00
Baruch Tabanpour 3c107b3047 Add stage_in and stage_out for MJX-Warp.
PiperOrigin-RevId: 868180029
Change-Id: I0d1b17b6023d32498766b5d1d9e5624bd31cda3d
2026-02-10 09:19:09 -08:00
Yuval Tassa 4fa062c7e5 Update unreleased MuJoCo version from 3.4.1 to 3.5.0.
PiperOrigin-RevId: 868169291
Change-Id: I8041842811f5f664c80116993315f50d549b5575
2026-02-10 08:52:04 -08:00
Google DeepMind b637d7eb76 Call kinematics in case of state mismatch in get_data_into with CPP backend.
PiperOrigin-RevId: 867667583
Change-Id: Ieec5ff86f6768c11459795068abf17f69bd26f3c
2026-02-09 10:29:27 -08:00
Yuval Tassa 6419534bad Add actuator and sensor delays. Fixes #1004
PiperOrigin-RevId: 866478839
Change-Id: Id21a6da0f98454c8fa39ea5af8a5e213d6eae497
2026-02-06 08:47:36 -08:00
Baruch Tabanpour d07f39b4a0 Add graph mode to MJX, in preparation of supporting staged graph capture.
PiperOrigin-RevId: 865380826
Change-Id: I840b48027e1142f87eb745d0d19b0cd3e87904b2
2026-02-04 06:38:13 -08:00
Taylor Howell 4a64017a5f Add documentation for MJX API
PiperOrigin-RevId: 860108133
Change-Id: I85469efcfe49d320ef8b89c9d174f33499a22b31
2026-01-23 07:53:17 -08:00
Taylor Howell bb4df33adf Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 857298931
Change-Id: Icbaddaea8140198c3673124d5994a01d3c92ed69
2026-01-16 14:11:39 -08:00
Alessio Quaglino e56b31e98f Reduce flexedge_J size from nflexedge x nv to the effective sparse size.
PiperOrigin-RevId: 856290141
Change-Id: Ide297d1b6f0e3aa07e3e20bf9ab44b6501097695
2026-01-14 11:18:41 -08:00