Commit Graph

608 Commits

Author SHA1 Message Date
Hana Joo 0aad65868d suppress pyrefly errors in mujoco
Startblock:
  * // Put other blockers before this line to avoid churning.
  * has lgtm
  * is approved
  * and then
  * all comments are resolved
  * and then
PiperOrigin-RevId: 961854282
Change-Id: I9d323a8ef9f36ffe69cc62d449849eac82dc6858
2026-08-09 15:45:36 -07:00
Copybara-Service cb4259e4e4 Merge pull request #3465 from shi-eric:shi-eric/mjx-warp-cache-ffi-callables
PiperOrigin-RevId: 961136122
Change-Id: Ia468812e0557bb4efb7866f7130dab5d1cea673f
2026-08-07 14:59:07 -07:00
Baruch Tabanpour 5e3464f475 Fix #3435. Add token to ensure sequential calls for mjx-warp refit and render.
PiperOrigin-RevId: 960553589
Change-Id: I76caca5a82dd7f96b39e51c5b67b7382ebd1726d
2026-08-06 16:05:47 -07:00
Eric Shi ab4fe36121 Cache MJX-Warp FFI callables
Equivalent JAX retraces rebuilt flattened wrapper functions, causing
Warp to register a new FFI callable for every trace.

Reuse callables with matching structural configurations while keeping
distinct shim functions isolated. This keeps Warp's registry and graph
cache effective without changing callback lifetime.

Document the structural key and lock scope so future changes preserve
the cache's intended boundaries.

Signed-off-by: Eric Shi <ershi@nvidia.com>
2026-08-06 19:49:45 +00:00
Taylor Howell a1d772c9ad Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 960411406
Change-Id: I687b06586881d1915bab1bf8ebab6473e6ef29db
2026-08-06 11:24:34 -07:00
Taylor Howell f848b185d3 update mjx make_data and put_data docstrings for njmax
PiperOrigin-RevId: 958526560
Change-Id: Iaeff307e4d09c9193134fc29689aeedfdcaa9d53
2026-08-03 12:49:29 -07:00
Taylor Howell 382af63dc3 update mjx for naccdmax
PiperOrigin-RevId: 958479874
Change-Id: If0f48f2f372955f2948cb286966f5fce52c3f0ba
2026-08-03 11:20:34 -07:00
Taylor Howell b46d76683d Import Warp v1.15.0
PiperOrigin-RevId: 956678131
Change-Id: I09a672b7335a85a17c95773ffc29eb9272744986
2026-07-30 12:39:05 -07:00
Michael Moss 1db801460b Update MuJoCo version to 3.11.1 following the 3.11.0 release
PiperOrigin-RevId: 954833728
Change-Id: I1c006cdd72072c6821d7b3c28d8ac320de19b8c5
2026-07-27 14:35:26 -07:00
Yuval Tassa 072e963fa0 Add SO3 transmission and native orientation actuator.
https://youtu.be/17XpwnqyCXs

New transmission type mjTRN_SO3: a relative orientation, targeting a ball
joint or a site+refsite pair. It is the first transmission with more than
one force output: its length is the norm of the expmap vector of the
relative rotation and its moment axes are the 3 rows of the
relative rotational Jacobian, without projecting onto per-actuator gears.

New force law mjGAIN_SO3/mjBIAS_SO3: a geodesic PD servo, force =
kp * log(q_current^-1 * q_target) - kv * velocity, exact for arbitrary axis
combinations with a unique equilibrium at every commanded orientation.
Error, moment rows and velocity all live in the child frame (joint or
site): the right-difference error is the gradient of the geodesic
potential in that frame. The parent-frame (left) error is not: driving
child-frame torques with it pumps energy at large angles, settling into
steady-spinning limit cycles (the SO3LargeAngleConvergence test). The
integrator variant stores the 3D orientation setpoint in act (actnum = 3,
re-anchored to a bounded representative at integration time). Exposed in
MJCF as <orientation joint=|site=+refsite= kp kv|dampratio>, or via
<general gaintype="so3" biastype="so3">.

The setpoint input has two charts: an expmap target (3 controls, default)
or a quaternion target (4 controls) -- <orientation input="quat">, the
first actuator with different input and output widths. The signature is
recorded in a new per-actuator field actuator_ctrlspec (mjtCtrlChart),
whose meaning is scoped by the gain type the way gain/bias parameters are;
ctrlnum is derived from it at compile time and remains the layout
authority. An explicit field rather than width inference or a prm slot:
width-as-chart cannot express same-width signatures (upcoming servo input
subsets), and prm slots are the input_mode pattern this stack retires.
The force law normalizes the commanded quaternion, making it scale- and
antipodally-invariant. The all-zero ctrl still maps to the identity via
mju_normalize4, but it is a degenerate point (a nudge of any component
commands a half-turn), so quat inputs reset to the identity quaternion:
new mj_resetCtrl sets neutral ctrl values (zero, except qw = 1), called
by mj_resetData and the viewers' Clear All. The quat chart is
restricted to dyntype 'none': integrating a quaternion setpoint linearly
is not meaningful on the manifold. New mjsActuator.ctrlspec field carries
the signature through the spec and XML round-trip.

Actuator sensors (actuatorpos/vel/frc) now report one value per force
output; dim = 3 on an SO3 actuator.

As the first actuator with nu != nactuator, this commit also makes the
viewers multi-input aware: the control sliders in simulate and studio,
which indexed per-actuator arrays by control index (out of bounds on
this model class), are generated per control and labeled with the
actuator name plus an input suffix ("orient/qw"), via the new
introspection helper mj_actuatorInputName -- the single source of truth
for input names, extended by each new multi-input type (quaternion
components are w-first: qw, qx, qy, qz). Slider ranges now honor a
defined ctrlrange even when ctrllimited is false: range is the UI hint,
limited is the clamp -- wrapped and expmap setpoints are unbounded but
still want finite sliders, while quat components are truly bounded.

The rotational demo model is orientation.xml under
test/engine/testdata/actuation/, upgraded to a three-way contrast:
per-axis wrapped servos vs an expmap-commanded vs a quat-commanded
orientation actuator, on identical checker-textured boxes. It is loaded
by the mixed-axis contrast and input-name tests, and doubles as the
viewer test model (slider groups of 3 independent, 3 grouped, 4 grouped).

PiperOrigin-RevId: 951607063
Change-Id: If235dba8e2f2ca72672e7c62531a27e967c6a373
2026-07-21 11:36:13 -07:00
Michael Moss b942c9f922 Update MuJoCo version to 3.11.0 in preparation for breaking changes release.
PiperOrigin-RevId: 950824721
Change-Id: Icc06b10ea69f5dd84e100eb75233f1598c2cbf88
2026-07-20 07:26:25 -07:00
Hana Joo 78946ca94e Adding type suppressions for pyrefly
PiperOrigin-RevId: 950703735
Change-Id: I6b6d5826a0dd1f4c0a1cbb2a46aee944cd60d884
2026-07-20 02:40:02 -07:00
Yuval Tassa a04b0c5b97 Add boolean flags for gravity compensation and surface velocity in mjModel.
These fields (`flg_gravcomp` and `flg_surfacevel`) replace the fast-path checks originally guarded by `ngravcomp` and (recently) `nsurfacevel`. Since the engine uses these integers only as flags (zero vs non-zero), migrating them to actual booleans makes them writeable from the Python bindings at runtime without violating size/dimension constraints.

The legacy integer field `ngravcomp` is marked as deprecated and will be removed in a future release.

PiperOrigin-RevId: 949779204
Change-Id: Ifab1f026063a4239302e6ad689663b611b59dda8
2026-07-17 15:05:26 -07:00
Yuval Tassa d507e92198 Preparation for MIMO actuators: split actuator counts: nu (inputs), nactuator (objects), nout (outputs).
An actuator now owns a block of consecutive controls
(actuator_ctrladr/ctrlnum, width defined by the actuator type) and a block
of consecutive force outputs (actuator_outadr/outnum, width defined by the
transmission type). Force outputs are the scalars of actuation space: one
force, length, velocity and moment row each. nout = dim(actuator_force) is
derived from transmission types; all current types have width 1, so all
three counts coincide for every existing model and behavior is bit-exact.

Array re-keying: ctrlrange/ctrllimited by nu; forcerange/forcelimited/gear/
acc0/length0/lengthrange and the moment row structure by nout; everything
else per actuator. The mjModel actuator block is re-sorted by size key.

Layout-breaking, not behavior-breaking: saved .mjb files are invalidated
(size list changed) and recompilation is required.

PiperOrigin-RevId: 948351772
Change-Id: Icbc196ffa083cb1eaa6f1a3710869c89d8f62540
2026-07-15 08:30:20 -07:00
Yuval Tassa 23cc383372 internal change
PiperOrigin-RevId: 947585050
Change-Id: Iaf9b60523b06d3cfd513ec661296eb39f4709442
2026-07-14 03:59:35 -07:00
dependabot[bot] bc3f17107e Bump pip from 26.1 to 26.1.2 in /mjx
Bumps [pip](https://github.com/pypa/pip) from 26.1 to 26.1.2.
- [Changelog](https://github.com/pypa/pip/blob/main/NEWS.rst)
- [Commits](https://github.com/pypa/pip/compare/26.1...26.1.2)

---
updated-dependencies:
- dependency-name: pip
  dependency-version: 26.1.2
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
2026-07-10 20:38:40 +00:00
Taylor Howell f6f80cd0a0 Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 945593266
Change-Id: Ic08842caabb92bc32974892ffae98e3295827094
2026-07-10 02:37:52 -07:00
Copybara-Service cb28d0c7d1 Merge pull request #3381 from tkelestemur:tarik/mjx-warp-put-model-nworld
PiperOrigin-RevId: 945578693
Change-Id: I873e4f47551d02f506d3af37040e1014fa0edb61
2026-07-10 02:02:02 -07:00
Tarik Kelestemur 4e5778433b Add batch_sizes model batching for MJX Warp 2026-07-07 18:06:30 -04:00
Taylor Howell 6608f1affa MJX Data.where. Fixes #3377
PiperOrigin-RevId: 943875492
Change-Id: I963e4ad052ce9184a7548a468aa4aea722d10469
2026-07-07 07:15:36 -07:00
Taylor Howell 641c1cf11b update MJX with nvmax
PiperOrigin-RevId: 943836220
Change-Id: I7092b5ed005a48a10812f3280709c50691aa497e
2026-07-07 07:14:47 -07:00
Yuval Tassa 315bcfbf3a Remove mjData.qM
PiperOrigin-RevId: 942520660
Change-Id: I422358e299ca0fcfe4c49cd0ee90c014a4c319d7
2026-07-04 09:38:45 -07:00
Taylor Howell 1f6cf4035c Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 941234486
Change-Id: Ie5d8cc0c9975e66f9a81bb6a86e33f640a5d19e3
2026-07-01 12:27:10 -07:00
Martin Schuck 407be2fc0b Copybara import of the project:
--
83a17d2844770fc2bbff37eda73b82df56076414 by Martin Schuck <martin.schuck@tum.de>:

Fix overflow cast

--
d12211c6665e1d791f77338daf8c62fa2374e3c8 by Martin Schuck <martin.schuck@tum.de>:

Prevent skipping warnings from cached jax functions by clearning the cache before invokation

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/3369 from amacati:fix.overflow_cast d12211c6665e1d791f77338daf8c62fa2374e3c8
PiperOrigin-RevId: 939914544
Change-Id: I0ab6e9ad1e7c45f352d2f49642049ad930955f0b
2026-06-29 10:45:40 -07:00
Taylor Howell 56a182e67c Import Warp v1.14.0
PiperOrigin-RevId: 938464395
Change-Id: Iac99c53566e2e62f6e75f5e0904100919043b463
2026-06-26 02:03:42 -07:00
Michael Moss 5700829122 Update MuJoCo version to 3.10.1 following the 3.10.0 release
PiperOrigin-RevId: 936044969
Change-Id: If0ebcc192419302da01741ebabb2b7dd3d7f7baf
2026-06-22 07:25:02 -07:00
Kevin Zakka b1725b4fc9 Bump pinned pip dependencies to patch Dependabot alerts 2026-06-15 08:53:04 -07:00
Yuval Tassa 60b5cac066 Switch MJX DataJAX from legacy qM format to CSR M format
Switch MJX DataJAX to use the CSR-format `M` as its primary sparse inertia representation, matching MuJoCo C `mjData.M`. This allows both `mjData.qM` and `mjModel.mapM2M` to be deleted in the future.

PiperOrigin-RevId: 932433465
Change-Id: I194ffbfdba0cdd9c718277ee0d5de7ac20f2b06f
2026-06-15 06:52:45 -07:00
Michael Moss 1ae561407a Minor cleanup to avoid unwanted files in mjx packaging.
PiperOrigin-RevId: 931118601
Change-Id: I94adedd0b29370e652140bb91f1488722eb61c55
2026-06-12 06:26:42 -07:00
Taylor Howell 763e713d1b Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 929061622
Change-Id: I53892a819c3727b1f04db943839606ec1f477df1
2026-06-09 02:28:29 -07:00
Copybara-Service 58a8c6ee56 Merge pull request #3323 from mhasek:mhasek/fix_mjwarp_data_sync_mismatch
PiperOrigin-RevId: 929043537
Change-Id: I99c6c67bc02c45dc3a432a585aebac1621757d6d
2026-06-09 01:47:11 -07:00
Copybara-Service a36758c26b Merge pull request #3322 from mhasek:mhasek/add_light_and_mat_prop_to_mjx_model
PiperOrigin-RevId: 928464101
Change-Id: I066caa1b0d45700160909d7452138e557af6b365
2026-06-08 03:42:53 -07:00
Copybara-Service d79520dc87 Merge pull request #3299 from tkelestemur:tarik/mjx-warp-nested-dataclass-methods
PiperOrigin-RevId: 928341876
Change-Id: I0c89bcbbd95a14fadbf33d7785bf212c2c074c99
2026-06-07 22:50:49 -07:00
Mustafa Hasekioglu b4135b571f fix island disable fields dimension mismatch 2026-06-06 16:18:52 -04:00
Mustafa Hasekioglu 994cfa0e4a update mjx model and shims 2026-06-06 14:39:01 -04:00
Taylor Howell 2a4af3246e Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 927227210
Change-Id: Ic72c58524c8e91e6fff44b4cab7f4a0f7de2bb33
2026-06-05 04:56:58 -07:00
btaba 7078f38163 Update generate_warp_types.py 2026-06-04 12:28:38 -07:00
Tarik Kelestemur b4dd8e2d9c Emit TileSet methods in MJX Warp codegen 2026-06-03 12:31:23 -04:00
Tarik Kelestemur 319b3975a0 Merge remote-tracking branch 'origin/main' into tarik/mjx-warp-nested-dataclass-methods
# Conflicts:
#	mjx/mujoco/mjx/warp/types.py
2026-06-03 11:36:53 -04:00
Taylor Howell 5d231c64cd Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 925939529
Change-Id: Idb92c2d50d5be9df7b6e91113b576e9c706847ca
2026-06-03 04:41:10 -07:00
Tarik Kelestemur c80d8fb35d Stabilize MJX Warp types codegen 2026-06-01 10:45:22 -04:00
Tarik Kelestemur f97b9f6e3b TileSet MJX codegen fix 2026-06-01 10:22:54 -04:00
Tarik Kelestemur 9bd8daf53b Minimize generated WARP types diff 2026-05-30 16:29:27 -04:00
Tarik Kelestemur f8a12a827f Merge origin/main into light_active PR 2026-05-30 16:24:27 -04:00
Taylor Howell 50e823e91c Import google-deepmind/mujoco_warp from GitHub.
PiperOrigin-RevId: 923922783
Change-Id: I0b407347a70a0e2024aae7ed5bab1592dc40a074
2026-05-30 07:44:29 -07:00
Taylor Howell b9c7a4b81a Import NVIDIA/warp 1.13.0 from GitHub.
PiperOrigin-RevId: 923878216
Change-Id: If9df5e0672e956eb62874fb6f4f74ba3a02977b3
2026-05-30 04:48:08 -07:00
Google DeepMind 9436972a22 Import NVIDIA/warp 1.13.0 from GitHub.
PiperOrigin-RevId: 922893824
Change-Id: Ib91ffb03257c614c48feda096822f723323be54d
2026-05-28 11:51:39 -07:00
Copybara-Service a22fc2423a Merge pull request #3295 from tkelestemur:tarik/fix-mjx-codegen-python311
PiperOrigin-RevId: 922794565
Change-Id: Ibd42328d739d3882e5c418cf69fc9787572cdb4a
2026-05-28 08:44:28 -07:00
Taylor Howell cae5a608d5 Import NVIDIA/warp 1.13.0 from GitHub.
PiperOrigin-RevId: 922749092
Change-Id: Ide815e4224e64af69e7c3a58102508bc81d22770
2026-05-28 07:00:59 -07:00
Tarik Kelestemur 106f1c5dbb fix mjx codegen and rerun 2026-05-27 17:50:31 -04:00