Merge main branch from mujoco

.github/workflows/build.yml

@@ -24,8 +24,24 @@ jobs:
       fail-fast: false
       matrix:
         include:
+          - os: ubuntu-24.04
+            label: "ubuntu-24.04-gcc-14"
+            cmake_args: >-
+              -G Ninja
+              -DCMAKE_C_COMPILER:STRING=gcc-14
+              -DCMAKE_CXX_COMPILER:STRING=g++-14
+              -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
+            tmpdir: "/tmp"
+          - os: ubuntu-24.04
+            label: "ubuntu-24.04-gcc-13"
+            cmake_args: >-
+              -G Ninja
+              -DCMAKE_C_COMPILER:STRING=gcc-13
+              -DCMAKE_CXX_COMPILER:STRING=g++-13
+              -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
+            tmpdir: "/tmp"
           - os: ubuntu-22.04
-            additional_label: "-gcc-12"
+            label: "ubuntu-22.04-gcc-12"
             cmake_args: >-
               -G Ninja
               -DCMAKE_C_COMPILER:STRING=gcc-12
@@ -33,7 +49,7 @@ jobs:
               -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
             tmpdir: "/tmp"
           - os: ubuntu-22.04
-            additional_label: "-gcc-11"
+            label: "ubuntu-22.04-gcc-11"
             cmake_args: >-
               -G Ninja
               -DCMAKE_C_COMPILER:STRING=gcc-11
@@ -41,23 +57,47 @@ jobs:
               -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
             tmpdir: "/tmp"
           - os: ubuntu-22.04
-            additional_label: "-gcc-10"
+            label: "ubuntu-22.04-gcc-10"
             cmake_args: >-
               -G Ninja
               -DCMAKE_C_COMPILER:STRING=gcc-10
               -DCMAKE_CXX_COMPILER:STRING=g++-10
               -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
             tmpdir: "/tmp"
-          - os: ubuntu-22.04
+          - os: ubuntu-24.04
-            additional_label: "-gcc-9"
+            label: "ubuntu-24.04-clang-18"
             cmake_args: >-
               -G Ninja
-              -DCMAKE_C_COMPILER:STRING=gcc-9
+              -DCMAKE_C_COMPILER:STRING=clang-18
-              -DCMAKE_CXX_COMPILER:STRING=g++-9
+              -DCMAKE_CXX_COMPILER:STRING=clang++-18
-              -DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
+              -DMUJOCO_HARDEN:BOOL=ON
+            tmpdir: "/tmp"
+          - os: ubuntu-24.04
+            label: "ubuntu-24.04-clang-17"
+            cmake_args: >-
+              -G Ninja
+              -DCMAKE_C_COMPILER:STRING=clang-17
+              -DCMAKE_CXX_COMPILER:STRING=clang++-17
+              -DMUJOCO_HARDEN:BOOL=ON
+            tmpdir: "/tmp"
+          - os: ubuntu-24.04
+            label: "ubuntu-24.04-clang-16"
+            cmake_args: >-
+              -G Ninja
+              -DCMAKE_C_COMPILER:STRING=clang-16
+              -DCMAKE_CXX_COMPILER:STRING=clang++-16
+              -DMUJOCO_HARDEN:BOOL=ON
             tmpdir: "/tmp"
           - os: ubuntu-22.04
-            additional_label: "-clang-14"
+            label: "ubuntu-22.04-clang-15"
+            cmake_args: >-
+              -G Ninja
+              -DCMAKE_C_COMPILER:STRING=clang-15
+              -DCMAKE_CXX_COMPILER:STRING=clang++-15
+              -DMUJOCO_HARDEN:BOOL=ON
+            tmpdir: "/tmp"
+          - os: ubuntu-22.04
+            label: "ubuntu-22.04-clang-14"
             cmake_args: >-
               -G Ninja
               -DCMAKE_C_COMPILER:STRING=clang-14
@@ -65,49 +105,33 @@ jobs:
               -DMUJOCO_HARDEN:BOOL=ON
             tmpdir: "/tmp"
           - os: ubuntu-22.04
-            additional_label: "-clang-13"
+            label: "ubuntu-22.04-clang-13"
             cmake_args: >-
               -G Ninja
               -DCMAKE_C_COMPILER:STRING=clang-13
               -DCMAKE_CXX_COMPILER:STRING=clang++-13
               -DMUJOCO_HARDEN:BOOL=ON
             tmpdir: "/tmp"
-          - os: ubuntu-20.04
+          - os: macos-15
-            additional_label: "-clang-12"
+            label: "macos-15-arm64"
-            cmake_args: >-
-              -G Ninja
-              -DCMAKE_C_COMPILER:STRING=clang-12
-              -DCMAKE_CXX_COMPILER:STRING=clang++-12
-              -DMUJOCO_HARDEN:BOOL=ON
-            tmpdir: "/tmp"
-          - os: ubuntu-20.04
-            additional_label: "-clang-11"
-            cmake_args: >-
-              -G Ninja
-              -DCMAKE_C_COMPILER:STRING=clang-11
-              -DCMAKE_CXX_COMPILER:STRING=clang++-11
-              -DMUJOCO_HARDEN:BOOL=ON
-            tmpdir: "/tmp"
-          - os: ubuntu-20.04
-            additional_label: "-clang-10"
-            cmake_args: >-
-              -G Ninja
-              -DCMAKE_C_COMPILER:STRING=clang-10
-              -DCMAKE_CXX_COMPILER:STRING=clang++-10
-              -DMUJOCO_HARDEN:BOOL=ON
-            tmpdir: "/tmp"
-          - os: macos-13
             cmake_args: >-
               -G Ninja
               -DMUJOCO_HARDEN:BOOL=ON
             tmpdir: "/tmp"
-          - os: windows-2022
+          - os: macos-15-large
+            label: "macos-15-x86_64"
             cmake_args: >-
-              -DCMAKE_SYSTEM_VERSION="10.0.22621.0"
+              -G Ninja
+              -DMUJOCO_HARDEN:BOOL=ON
+            tmpdir: "/tmp"
+          - os: windows-2025
+            label: "windows-2025"
+            cmake_args: >-
+              -DCMAKE_SYSTEM_VERSION="10.0.26100.0"
             cmake_build_args: "-- -m"
             tmpdir: "C:/Temp"
 
-    name: "${{ matrix.os }}${{ matrix.additional_label }}"
+    name: "${{ matrix.label }}"
     runs-on: ${{ matrix.os }}
 
     steps:
@@ -179,7 +203,7 @@ jobs:
            cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
-           cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
+           cp lib/libsdf_plugin.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
     - name: Copy plugins (Windows)
       if: ${{ runner.os == 'Windows' }}
       working-directory: build
@@ -273,7 +297,7 @@ jobs:
         CHATMSG_AUTHOR_NAME: ${{ github.event.head_commit.author.name }}
         CHATMSG_AUTHOR_EMAIL: ${{ github.event.head_commit.author.email }}
         CHATMSG_COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
-        CHATMSG_JOB_ID: ${{ matrix.os }}${{ matrix.additional_label }}
+        CHATMSG_JOB_ID: ${{ matrix.label }}
       if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API_URL != '' }}
       run: |
         CHATMSG="$(cat <<-'EOF' | python3
@@ -288,7 +312,7 @@ jobs:
             email=env('CHATMSG_AUTHOR_EMAIL'),
             msg=env('CHATMSG_COMMIT_MESSAGE').replace('```', '')
         )
-        text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name}<{email}>\n\n{msg}```'.format(**data)
+        text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name} <{email}>\n\n{msg}```'.format(**data)
         print(json.dumps({'text' : text}))
         EOF
         )" &&

CMakeLists.txt

@@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco
-  VERSION 3.3.1
+  VERSION 3.3.3
   DESCRIPTION "MuJoCo Physics Simulator"
   HOMEPAGE_URL "https://mujoco.org"
 )
@@ -42,6 +42,14 @@ option(MUJOCO_BUILD_EXAMPLES "Build samples for MuJoCo" ON)
 option(MUJOCO_BUILD_SIMULATE "Build simulate library for MuJoCo" ON)
 option(MUJOCO_BUILD_TESTS "Build tests for MuJoCo" ON)
 option(MUJOCO_TEST_PYTHON_UTIL "Build and test utility libraries for Python bindings" ON)
+option(MUJOCO_BUILD_USD_PLUGINS "Build OpenUSD plugins" OFF)
+
+# USD libs to compile against.
+set(MUJOCO_USD_ALLOWED_TARGET_VALUES "USD" "Houdini")
+set(MUJOCO_USD_TARGET "USD" CACHE STRING "Select the USD target for the project.")
+set_property(CACHE MUJOCO_USD_TARGET
+  PROPERTY STRINGS ${MUJOCO_USD_ALLOWED_TARGET_VALUES}
+)
 
 if(APPLE AND (MUJOCO_BUILD_EXAMPLES OR MUJOCO_BUILD_SIMULATE))
   enable_language(OBJC)
@@ -181,6 +189,10 @@ if(MUJOCO_BUILD_EXAMPLES)
   add_subdirectory(sample)
 endif()
 
+if(MUJOCO_BUILD_USD_PLUGINS)
+  add_subdirectory(src/experimental/usd/plugins)
+endif()
+
 if(BUILD_TESTING AND MUJOCO_BUILD_TESTS)
   enable_testing()
   add_subdirectory(test)

README.md

@@ -53,6 +53,8 @@ running on Google Colab:
 
  - The **introductory** tutorial teaches MuJoCo basics:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)
+ - The **Model Editing** tutorial shows how to create and edit models procedurally:
+   [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb)
  - The **rollout** tutorial shows how to use the multithreaded `rollout` module:
    [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb)
  - The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a

cmake/MujocoDependencies.cmake

@@ -39,12 +39,12 @@ set(MUJOCO_DEP_VERSION_qhull
     CACHE STRING "Version of `qhull` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_Eigen3
-    66f7f51b7e069d0a03a21157fa60b24aece69aeb
+    464c1d097891a1462ab28bf8bb763c1683883892
     CACHE STRING "Version of `Eigen3` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_abseil
-    9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0
+    d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1
     CACHE STRING "Version of `abseil` to be fetched."
 )
 

dist/mujoco.rc

@@ -1,6 +1,6 @@
 1 VERSIONINFO
-FILEVERSION 3,3,1,0
+FILEVERSION 3,3,3,0
-PRODUCTVERSION 3,3,1,0
+PRODUCTVERSION 3,3,3,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -9,9 +9,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "3.3.1"
+      VALUE "ProductVersion", "3.3.3"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "3.3.1"
+      VALUE "FileVersion", "3.3.3"
       VALUE "InternalName", "mujoco.dll"
       VALUE "OriginalFilename", "mujoco.dll"
       VALUE "CompanyName", "Google DeepMind"

dist/simulate.rc

@@ -1,8 +1,8 @@
 MUJOCO ICON "mujoco.ico"
 
 1 VERSIONINFO
-FILEVERSION 3,3,1,0
+FILEVERSION 3,3,3,0
-PRODUCTVERSION 3,3,1,0
+PRODUCTVERSION 3,3,3,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -11,9 +11,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "3.3.1"
+      VALUE "ProductVersion", "3.3.3"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "3.3.1"
+      VALUE "FileVersion", "3.3.3"
       VALUE "InternalName", "simulate.exe"
       VALUE "OriginalFilename", "simulate.exe"
       VALUE "CompanyName", "Google DeepMind"

doc/APIreference/APIglobals.rst

@@ -448,7 +448,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - 200
      - The number of iterations where solver statistics can be stored in ``mjData.solver``. This array is used
        to store diagnostic information about each iteration of the constraint solver.
-       The actual number of iterations is given by ``mjData.solver_iter``.
+       The actual number of iterations is given by ``mjData.solver_niter``.
    * - ``mjNISLAND``
      - 20
      - The number of islands for which solver statistics can be stored in ``mjData.solver``. This array is
@@ -517,7 +517,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - Maximum number of UI rectangles.
        Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
    * - ``mjVERSION_HEADER``
-     - 331
+     - 333
      - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software
        version, so 210 corresponds to version 2.1. Defined in  mujoco.h. The API function :ref:`mj_version` returns a
        number with the same meaning but for the compiled library.

doc/APIreference/APItypes.rst

@@ -165,6 +165,17 @@ values are used in ``m->cam_mode`` and ``m->light_mode``.
 .. mujoco-include:: mjtCamLight
 
 
+.. _mjtLightType:
+
+mjtLightType
+~~~~~~~~~~~~
+
+The type of a light source describing how its position, orientation and other properties will interact with the
+objects in the scene. These values are used in ``m->light_type``.
+
+.. mujoco-include:: mjtLightType
+
+
 .. _mjtTexture:
 
 mjtTexture
@@ -357,6 +368,26 @@ last argument to :ref:`mj_local2global`.
 .. mujoco-include:: mjtSameFrame
 
 
+.. _mjtFlexSelf:
+
+mjtFlexSelf
+~~~~~~~~~~~~
+
+Types of flex self-collisions midphase.
+
+.. mujoco-include:: mjtFlexSelf
+
+
+.. _mjtSDFType:
+
+mjtSDFType
+~~~~~~~~~~~
+
+Formulas used to combine SDFs when calling mjc_distance and mjc_gradient.
+
+.. mujoco-include:: mjtSDFType
+
+
 .. _tyDataEnums:
 
 Data
@@ -883,7 +914,7 @@ mjSolverStat
 
 This is the data structure holding information about one solver iteration. ``mjData.solver`` is a preallocated array
 of mjSolverStat data structures, one for each iteration of the solver, up to a maximum of mjNSOLVER. The actual number
-of solver iterations is given by ``mjData.solver_iter``.
+of solver iterations is given by ``mjData.solver_niter``.
 
 .. mujoco-include:: mjSolverStat
 
@@ -967,17 +998,6 @@ This structure contains everything needed to render the 3D scene in OpenGL.
 .. mujoco-include:: mjvScene
 
 
-.. _mjvSceneState:
-
-mjvSceneState
-~~~~~~~~~~~~~
-
-This structure contains the portions of :ref:`mjModel` and :ref:`mjData` that are required for
-various ``mjv_*`` functions.
-
-.. mujoco-include:: mjvSceneState
-
-
 .. _mjvFigure:
 
 mjvFigure

doc/APIreference/functions.rst

@@ -51,6 +51,15 @@ Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multi
 :ref:`mjModel` instance that takes the edits into account.
 If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`.
 
+.. _mj_copyBack:
+
+`mj_copyBack <#mj_copyBack>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_copyBack
+
+Copy real-valued arrays from model to spec, returns 1 on success.
+
 .. _mj_recompile:
 
 `mj_recompile <#mj_recompile>`__
@@ -459,7 +468,7 @@ Multiply vector by (inertia matrix)^(1/2).
 .. mujoco-include:: mj_addM
 
 Add inertia matrix to destination matrix.
-Destination can be sparse uncompressed, or dense when all int* are NULL
+Destination can be sparse or dense when all int* are NULL.
 
 .. _mj_applyFT:
 
@@ -748,8 +757,7 @@ Compare forward and inverse dynamics, save results in fwdinv.
 Sub components
 ^^^^^^^^^^^^^^
 
-These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
+These are sub-components of the simulation pipeline, called internally from the components above.
-that the user will need to call them.
 
 .. _mj_sensorPos:
 
@@ -886,6 +894,18 @@ Compute actuator transmission lengths and moments.
 
 Run composite rigid body inertia algorithm (CRB).
 
+.. _mj_makeM:
+
+`mj_makeM <#mj_makeM>`__
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_makeM
+
+Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due
+to :ref:`tendon armature<tendon-spatial-armature>`. The joint-space inertia matrix is stored in both ``mjData.qM`` and
+``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row,
+respectively).
+
 .. _mj_factorM:
 
 `mj_factorM <#mj_factorM>`__
@@ -1364,6 +1384,15 @@ If the model buffer is unallocated the initial configuration will not be set.
 Copy mjData.
 m is only required to contain the size fields from MJMODEL_INTS.
 
+.. _mjv_copyData:
+
+`mjv_copyData <#mjv_copyData>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjv_copyData
+
+Copy mjData, skip large arrays not required for visualization.
+
 .. _mj_resetData:
 
 `mj_resetData <#mj_resetData>`__
@@ -1995,15 +2024,6 @@ Rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_
 
 Move camera with mouse; action is mjtMouse.
 
-.. _mjv_moveCameraFromState:
-
-`mjv_moveCameraFromState <#mjv_moveCameraFromState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_moveCameraFromState
-
-Move camera with mouse given a scene state; action is mjtMouse.
-
 .. _mjv_movePerturb:
 
 `mjv_movePerturb <#mjv_movePerturb>`__
@@ -2013,15 +2033,6 @@ Move camera with mouse given a scene state; action is mjtMouse.
 
 Move perturb object with mouse; action is mjtMouse.
 
-.. _mjv_movePerturbFromState:
-
-`mjv_movePerturbFromState <#mjv_movePerturbFromState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_movePerturbFromState
-
-Move perturb object with mouse given a scene state; action is mjtMouse.
-
 .. _mjv_moveModel:
 
 `mjv_moveModel <#mjv_moveModel>`__
@@ -2164,15 +2175,6 @@ Free abstract scene.
 
 Update entire scene given model state.
 
-.. _mjv_updateSceneFromState:
-
-`mjv_updateSceneFromState <#mjv_updateSceneFromState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_updateSceneFromState
-
-Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings.
-
 .. _mjv_copyModel:
 
 `mjv_copyModel <#mjv_copyModel>`__
@@ -2182,42 +2184,6 @@ Update entire scene from a scene state, return the number of new mjWARN_VGEOMFUL
 
 Copy mjModel, skip large arrays not required for abstract visualization.
 
-.. _mjv_defaultSceneState:
-
-`mjv_defaultSceneState <#mjv_defaultSceneState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_defaultSceneState
-
-Set default scene state.
-
-.. _mjv_makeSceneState:
-
-`mjv_makeSceneState <#mjv_makeSceneState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_makeSceneState
-
-Allocate resources and initialize a scene state object.
-
-.. _mjv_freeSceneState:
-
-`mjv_freeSceneState <#mjv_freeSceneState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_freeSceneState
-
-Free scene state.
-
-.. _mjv_updateSceneState:
-
-`mjv_updateSceneState <#mjv_updateSceneState>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjv_updateSceneState
-
-Update a scene state from model and data.
-
 .. _mjv_addGeoms:
 
 `mjv_addGeoms <#mjv_addGeoms>`__
@@ -3808,41 +3774,14 @@ Free all pointers with ``mju_free()``.
 
 Attachment
 ^^^^^^^^^^
-.. _mjs_attachBody:
+.. _mjs_attach:
 
-`mjs_attachBody <#mjs_attachBody>`__
+`mjs_attach <#mjs_attach>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. mujoco-include:: mjs_attachBody
+.. mujoco-include:: mjs_attach
 
-Attach child body to a parent frame, return the attached body if success or NULL otherwise.
+Attach child to a parent, return the attached element if success or NULL otherwise.
-
-.. _mjs_attachFrame:
-
-`mjs_attachFrame <#mjs_attachFrame>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjs_attachFrame
-
-Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
-
-.. _mjs_attachToSite:
-
-`mjs_attachToSite <#mjs_attachToSite>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjs_attachToSite
-
-Attach child body to a parent site, return the attached body if success or NULL otherwise.
-
-.. _mjs_attachFrameToSite:
-
-`mjs_attachFrameToSite <#mjs_attachFrameToSite>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mjs_attachFrameToSite
-
-Attach child frame to a parent site, return the attached frame if success or NULL otherwise.
 
 .. _mjs_detachBody:
 
@@ -4428,6 +4367,15 @@ Get string contents.
 
 Get double array contents and optionally its size.
 
+.. _mjs_getPluginAttributes:
+
+`mjs_getPluginAttributes <#mjs_getPluginAttributes>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_getPluginAttributes
+
+Get plugin attributes.
+
 .. _SpecUtilities:
 
 Spec utilities
@@ -4448,7 +4396,7 @@ Set element's default.
 
 .. mujoco-include:: mjs_setFrame
 
-Set element's enclosing frame.
+Set element's enclosing frame, return 0 on success.
 
 .. _mjs_resolveOrientation:
 
@@ -4477,6 +4425,17 @@ Transform body into a frame.
 
 Set user payload, overriding the existing value for the specified key if present.
 
+.. _mjs_setUserValueWithCleanup:
+
+`mjs_setUserValueWithCleanup <#mjs_setUserValueWithCleanup>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_setUserValueWithCleanup
+
+Set user payload, overriding the existing value for the specified key if
+present. This version differs from mjs_setUserValue in that it takes a
+cleanup function that will be called when the user payload is deleted.
+
 .. _mjs_getUserValue:
 
 `mjs_getUserValue <#mjs_getUserValue>`__

doc/APIreference/functions_override.rst

@@ -113,8 +113,14 @@ Integrates the simulation state using an implicit-in-velocity integrator (either
 
 .. _Subcomponents:
 
-These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
+These are sub-components of the simulation pipeline, called internally from the components above.
-that the user will need to call them.
+
+.. _mj_makeM:
+
+Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due
+to :ref:`tendon armature<tendon-spatial-armature>`. The joint-space inertia matrix is stored in both ``mjData.qM`` and
+``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row,
+respectively).
 
 .. _mj_factorM:
 

doc/XMLreference.rst

@@ -175,11 +175,10 @@ replicating 200 times, suffixes will be ``000, 001, ...`` etc). All referencing
 and namespaced appropriately. Detailed examples of models using replicate can be found in the
 `model/replicate/ <https://github.com/google-deepmind/mujoco/tree/main/model/replicate>`__ directory.
 
-There is a caveat concerning :ref:`keyframes<keyframe>` when using replicate. Since :ref:`mjs_attachFrame` is used to
+There is a caveat concerning :ref:`keyframes<keyframe>` when using replicate. Since :ref:`mjs_attach` is used to
 self-attach multiple times the enclosed kinematic tree, if this tree contains further :ref:`attach<body-attach>`
 elements, keyframes will not be replicated nor namespaced by :ref:`replicate<replicate>`, but they will be attached and
-namespaced once by the innermost call of :ref:`mjs_attachFrame` or :ref:`mjs_attachBody`. See the limitations discussed
+namespaced once by the innermost call of :ref:`mjs_attach`. See the limitations discussed in :ref:`attach<body-attach>`.
-in :ref:`attach<body-attach>`.
 
 .. _replicate-count:
 
@@ -833,6 +832,10 @@ has any effect. The settings here are global and apply to the entire model.
    necessary to adjust this attribute and the geom-specific groups so as to exclude world geoms from the inertial
    computation.
 
+.. _compiler-saveinertial:
+
+:at:`saveinertial`: :at-val:`[false, true], "false"`
+   If set to "true", the compiler will save explicit :ref:`inertial <body-inertial>` clauses for all bodies.
 
 .. _compiler-lengthrange:
 
@@ -1461,16 +1464,19 @@ still be specified here but this functionality is now deprecated and will be rem
 :el-prefix:`asset/` |-| **texture** (*)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-| This element creates a texture asset, which is then referenced from a :ref:`material <asset-material>` asset, which is
+   This element creates a texture asset, which is then referenced from a :ref:`material <asset-material>` asset, which
-  finally referenced from a model element that needs to be textured. MuJoCo provides access to the texture mapping
+   is finally referenced from a model element that needs to be textured.
-  mechanism in OpenGL. Texture coordinates are generated automatically in GL_OBJECT_PLANE mode, using either 2D or cube
+
-  mapping. MIP maps are always enabled in GL_LINEAR_MIPMAP_LINEAR mode. The texture color is combined with the object
+   The texture data can be loaded from files or can be generated by the compiler as a procedural texture. Because
-  color in GL_MODULATE mode. The texture data can be loaded from PNG files, with provisions for loading cube and skybox
+   different texture types require different parameters, only a subset of the attributes below are used for any given
-  textures. Alternatively the data can be generated by the compiler as a procedural texture. Because different texture
+   texture. Provisions are provided for loading cube and skybox textures from individual image files.
-  types require different parameters, only a subset of the attributes below are used for any given texture.
+
-| A second file format is supported for loading textures, in addition to PNG. If the file name extension is
+   Currently, three file formats are supported for loading textures: PNG, KTX, and a custom MuJoCo texture format. The
-  different from .png or .PNG, or if the ``content_type`` attribute is set to ``image/vnd.mujoco.texture``, then MuJoCo
+   loader will use the extension of the file name to determine which format to use, defaulting to the custom format if
-  assumes that the texture is in this format. This is a custom binary file format, containing the following data:
+   the extension is not recognized. Alternatively, the content_type attribute can be used to specify the format
+   explicitly. Only ``image/png``, ``image/ktx``, or ``image/vnd.mujoco.texture`` are supported.
+
+   The custom MuJoCo format is assumed to be a binary file containing the following data:
 
 .. code:: Text
 
@@ -1523,13 +1529,19 @@ still be specified here but this functionality is now deprecated and will be rem
    with the texrepeat attribute of :ref:`material <asset-material>`. The data can be loaded from a single file or
    created procedurally.
 
+.. _asset-texture-colorspace:
+
+:at:`colorspace`: :at-val:`[auto, linear, sRGB], "auto"`
+   This attribute determines the color space of the texture. The default value ``auto`` means that the color space will
+   be determined from the image file itself.  If no color space is defined in the file, then ``linear`` is assumed.
+
 .. _asset-texture-content_type:
 
 :at:`content_type`: :at-val:`string, optional`
    If the file attribute is specified, then this sets the
    `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`_ (formerly known as MIME types) of the
-   file to be loaded. Any filename extensions will be ignored.  Currently ``image/png`` and ``image/vnd.mujoco.texture``
+   file to be loaded. Any filename extensions will be ignored.  Currently ``image/png``, ``image/ktx``, and
-   are supported.
+   ``image/vnd.mujoco.texture`` are supported.
 
 .. _asset-texture-file:
 
@@ -1777,11 +1789,11 @@ properties are grouped together.
 .. _asset-material-rgba:
 
 :at:`rgba`: :at-val:`real(4), "1 1 1 1"`
-   Color and transparency of the material. All components should be in the range [0 1]. Note that textures are applied
+   Color and transparency of the material. All components should be in the range [0 1]. Note that the texture color (if
-   in GL_MODULATE mode, meaning that the texture color and the color specified here are multiplied component-wise. Thus
+   assigned) and the color specified here are multiplied component-wise. Thus the default value of "1 1 1 1" has the
-   the default value of "1 1 1 1" has the effect of leaving the texture unchanged. When the material is applied to a
+   effect of leaving the texture unchanged. When the material is applied to a model element which defines its own local
-   model element which defines its own local rgba attribute, the local definition has precedence. Note that this "local"
+   rgba attribute, the local definition has precedence. Note that this "local" definition could in fact come from a
-   definition could in fact come from a defaults class. The remaining material properties always apply.
+   defaults class. The remaining material properties always apply.
 
 .. _material-layer:
 
@@ -2151,14 +2163,28 @@ rotations as unit quaternions.
    corresponding to mjModel.qpos_spring is also used to compute the spring reference lengths of all tendons, stored in
    mjModel.tendon_lengthspring. This is because :ref:`tendons <tendon>` can also have springs.
 
+.. image:: images/XMLreference/armature.gif
+   :width: 40%
+   :align: right
+   :class: only-light
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml
+.. image:: images/XMLreference/armature_dark.gif
+   :width: 40%
+   :align: right
+   :class: only-dark
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml
+
 .. _body-joint-armature:
 
 :at:`armature`: :at-val:`real, "0"`
    Additional inertia associated with movement of the joint that is not due to body mass. This added inertia is usually
    due to a rotor (a.k.a `armature <https://en.wikipedia.org/wiki/Armature_(electrical)>`__) spinning faster than the
-   joint itself due to a geared transmission; in this case the added inertia is known as "reflected inertia" and its
+   joint itself due to a geared transmission. In the illustration, we compare (*left*) a 2-dof system with an armature
-   value is the rotational inertia of the spinning element multiplied by the square of the gear ratio. The value applies
+   body (purple box), coupled with a gear ratio of :math:`3` to the pendulum using a :ref:`joint
-   to all degrees of freedom created by this joint.
+   equality<equality-joint>` constraint, and (*right*) a simple 1-dof pendulum with an equivalent :at:`armature`.
+   Because the gear ratio appears twice, multiplying both forces and lengths, the effect is known as "reflected
+   inertia" and the equivalent value is the inertia of the spinning body multiplied by the *square of the gear ratio*,
+   in this case :math:`9=3^2`. The value applies to all degrees of freedom created by this joint.
 
    Besides increasing the realism of joints with geared transmission, positive :at:`armature` significantly improves
    simulation stability, even for small values, and is a recommended possible fix when encountering stability issues.
@@ -2849,11 +2875,15 @@ and the +Y axis points up. Thus the frame position and orientation are the key a
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 This element creates a light, which moves with the body where it is defined. To create a fixed light, define it in the
-world body. The lights created here are in addition to the default headlight which is always defined and is adjusted via
+world body. The lights created here are in addition to the headlight which is always defined and is configured via the
-the :ref:`visual <visual>` element. MuJoCo relies on the standard lighting model in OpenGL (fixed functionality)
+:ref:`visual <visual>` element. Lights shine along the direction specified by the dir attribute. They do not have
-augmented with shadow mapping. The effects of lights are additive, thus adding a light always makes the scene brighter.
+a full spatial frame with three orthogonal axes.
-The maximum number of lights that can be active simultaneously is 8, counting the headlight. The light is shining along
+
-the direction specified by the dir attribute. It does not have a full spatial frame with three orthogonal axes.
+By default, MuJoCo uses the standard OpenGL (fixed functional) Phong lighting model for its rendering, with augmented
+with shadow mapping. (See the OpenGL documentation for more information, including details about various attributes.)
+
+MJCF also supports alternative lighting models (e.g. physically-based rendering) by providing additional attributes.
+Attributes may be applied or ignored depending on the lighting model being used.
 
 
 .. _body-light-name:
@@ -2879,10 +2909,17 @@ the direction specified by the dir attribute. It does not have a full spatial fr
    This is identical to the target attribute of :ref:`camera <body-camera>` above. It specifies which body should be
    targeted in "targetbody" and "targetbodycom" modes.
 
+.. _body-light-type:
+
+:at:`type`: :at-val:`[spot, directional, point, image], "spot"`
+   Determines the type of light. Note that some light types may not be supported by some renderers (e.g. only spot and
+   directional lights are supported by the default native renderer).
+
 .. _body-light-directional:
 
 :at:`directional`: :at-val:`[false, true], "false"`
-   The light is directional if this attribute is "true", otherwise it is a spotlight.
+   This is a deprecated legacy attribute. Please use :ref:`light <body-light-type>` type instead. If set to "true", and
+   no type is specified, this will change the light type to be directional.
 
 .. _body-light-castshadow:
 
@@ -2902,12 +2939,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr
    these clipping planes bound the cone or box shadow volume in the light direction. As a result, some shadows
    (especially those very close to the light) may be clipped.
 
-.. _body-light-bulbradius:
-
-:at:`radius`: :at-val:`real, "0.02"`
-   Radius of the light, affects shadow softness. This attribute has no effect in MuJoCo's native renderer, but it can be
-   useful when rendering scenes with an external renderer.
-
 .. _body-light-active:
 
 :at:`active`: :at-val:`[false, true], "true"`
@@ -2924,11 +2955,50 @@ the direction specified by the dir attribute. It does not have a full spatial fr
 :at:`dir`: :at-val:`real(3), "0 0 -1"`
    Direction of the light.
 
+.. _body-light-diffuse:
+
+:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"`
+   The color of the light. For the Phong (default) lighting model, this defines the diffuse color of
+   the light.
+
+.. _body-light-texture:
+
+:at:`texture`: :at-val:`string, optional`
+   The texture to use for image-based lighting. This is unused by the default Phong lighting model.
+
+.. _body-light-intensity:
+
+:at:`intensity`: :at-val:`real, "0.0"`
+   The intensity of the light source, measured in candela, used for physically-based lighting models.
+   This is unused by the default Phong lighting model.
+
+.. _body-light-ambient:
+
+:at:`ambient`: :at-val:`real(3), "0 0 0"`
+   The ambient color of the light, used by the default Phong lighting model.
+
+.. _body-light-specular:
+
+:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"`
+   The specular color of the light, used by the default Phong lighting model.
+
+.. _body-light-range:
+
+:at:`range`: :at-val:`real, "10.0"`
+   The effective range of the light. Objects further than this distance from the light position
+   will not be illuminated by this light. This only applies to spotlights.
+
+.. _body-light-bulbradius:
+
+:at:`bulbradius`: :at-val:`real, "0.02"`
+   The radius of the light source which can affect shadow softness depending on the
+   renderer. This only applies to spotlights.
+
 .. _body-light-attenuation:
 
 :at:`attenuation`: :at-val:`real(3), "1 0 0"`
-   These are the constant, linear and quadratic attenuation coefficients in OpenGL. The default corresponds to no
+   These are the constant, linear and quadratic attenuation coefficients for Phong lighting.
-   attenuation. See the OpenGL documentation for more information on this and all other OpenGL-related properties.
+   The default corresponds to no attenuation.
 
 .. _body-light-cutoff:
 
@@ -2940,22 +3010,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr
 :at:`exponent`: :at-val:`real, "10"`
    Exponent for spotlights. This setting controls the softness of the spotlight cutoff.
 
-.. _body-light-ambient:
-
-:at:`ambient`: :at-val:`real(3), "0 0 0"`
-   The ambient color of the light.
-
-.. _body-light-diffuse:
-
-:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"`
-   The diffuse color of the light.
-
-.. _body-light-specular:
-
-:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"`
-   The specular color of the light.
-
-
 .. _body-composite:
 
 :el-prefix:`body/` |-| **composite** (*)
@@ -3003,8 +3057,13 @@ cable, which produces an inextensible chain of bodies connected with ball joints
 .. _body-composite-offset:
 
 :at:`offset`: :at-val:`real(3), "0 0 0"`
-   It specifies a 3D offset from the center of the parent body to the center of the grid of elements. The offset is
+   It specifies a 3D offset from the center of the parent body to the center of the first body of the cable. The offset
-   expressed in the local coordinate frame of the parent body.
+   is expressed in the local coordinate frame of the parent body.
+
+.. _body-composite-quat:
+
+:at:`quat`: :at-val:`real(4), "1 0 0 0"`
+   It specifies a quaternion that rotates the first body frame. The quaternion is expressed in the parent body frame.
 
 .. _body-composite-vertex:
 
@@ -3532,6 +3591,7 @@ saving the XML:
 
 .. _flexcomp-contact-internal:
 .. _flexcomp-contact-selfcollide:
+.. _flexcomp-contact-vertcollide:
 .. _flexcomp-contact-activelayers:
 .. _flexcomp-contact-contype:
 .. _flexcomp-contact-conaffinity:
@@ -3545,8 +3605,8 @@ saving the XML:
 .. _flexcomp-contact-gap:
 
 .. |body/flexcomp/contact attrib list| replace::
-   :at:`internal`, :at:`selfcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`, :at:`condim`,
+   :at:`internal`, :at:`selfcollide`, :at:`vertcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`,
-   :at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap`
+   :at:`condim`, :at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap`
 
 |body/flexcomp/contact attrib list|
    Same as in :ref:`flex/contact<flex-contact>`. All attributes are passed through to the automatically-generated flex.
@@ -3588,9 +3648,10 @@ element is used to adjust the properties of all edges in the flex.
 .. _flexcomp-elasticity-poisson:
 .. _flexcomp-elasticity-damping:
 .. _flexcomp-elasticity-thickness:
+.. _flexcomp-elasticity-elastic2d:
 
 .. |body/flexcomp/elasticity attrib list| replace::
-   :at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness`
+   :at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness`, :at:`elastic2d`
 
 |body/flexcomp/elasticity attrib list|
    Same as in :ref:`flex/elasticity<flex-elasticity>`.
@@ -3698,8 +3759,9 @@ all attachments will appear in the saved XML file.
 
 .. _body-attach-body:
 
-:at:`body`: :at-val:`string, required`
+:at:`body`: :at-val:`string, optional`
-   Name of the body in the sub-model to attach here. The body and its subtree will be attached.
+   Name of the body in the sub-model to attach here. The body and its subtree will be attached. If this attribute is not
+   specified, the contents of the world body will be attached in a new :ref:`frame<body-frame>`.
 
 .. _body-attach-prefix:
 
@@ -4030,6 +4092,11 @@ stress-strain relationship.. See also :ref:`deformable <CDeformable>` objects.
    This thickness can be set equal to 2 times the :ref:`radius <deformable-flex-radius>` in order to match the geometry,
    but is exposed separately since the radius might be constrained by considerations related to collision detection.
 
+.. _flex-elasticity-elastic2d:
+
+:at:`elastic2d`: :at-val:`[none, bend, stretch, both], "none"`
+   Elastic contribution to passive forces of 2D flexes. "none": none, "bend": bending only, "stretch": stretching only,
+   "both": bending and stretching.
 
 .. _flex-contact:
 
@@ -4063,6 +4130,13 @@ extensions specific to flexes.
    **sap** in 1D and 2D, and **bvh** in 3D. Which strategy performs better depends on the specifics of the model. The
    automatic setting is just a simple rule which we have found to perform well in general.
 
+.. _flex-contact-vertcollide:
+
+:at:`vertcollide`: :at-val:`[true, false], "false"`
+   Enables or disables vertex collisions. if **true**, spherical geoms are added at the vertices of flex, with radius
+   equal to the radius of the flex. These geoms can collide with other geoms and are not visible by default. If
+   **false**, no additional geoms are added.
+
 .. _flex-contact-activelayers:
 
 :at:`activelayers`: :at-val:`int(1), "1"`
@@ -4607,12 +4681,28 @@ length X, as in the clip on the right of `this example model
    solver. If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler <compiler>`, length limits
    will be enabled if range is defined.
 
+.. _tendon-spatial-actuatorfrclimited:
+
+:at:`actuatorfrclimited`: :at-val:`[false, true, auto], "auto"`
+   This attribute specifies whether actuator forces acting on the tendon should be clamped. See :ref:`CForceRange` for
+   details. This attribute interacts with the :ref:`actuatorfrcrange<tendon-spatial-actuatorfrcrange>` attribute. If
+   this attribute is "false", actuator force clamping is disabled. If it is "true", actuator force clamping is enabled.
+   If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler <compiler>`, actuator force clamping will
+   be enabled if :at:`actuatorfrcrange` is defined.
+
 .. _tendon-spatial-range:
 
 :at:`range`: :at-val:`real(2), "0 0"`
    Range of allowed tendon lengths. Setting this attribute without specifying :at:`limited` is an error, unless
    :at:`autolimits` is set in :ref:`compiler <compiler>`.
 
+.. _tendon-spatial-actuatorfrcrange:
+
+:at:`actuatorfrcrange`: :at-val:`real(2), "0 0"`
+   Range for clamping total actuator forces acting on this tendon. See :ref:`CForceRange` for details. The compiler
+   expects the lower bound to be nonpositive and the upper bound to be nonnegative. |br| Setting this attribute without
+   specifying :at:`actuatorfrclimited` is an error if :at:`compiler-autolimits` is "false".
+
 .. _tendon-spatial-solreflimit:
 
 .. _tendon-spatial-solimplimit:
@@ -4683,6 +4773,31 @@ length X, as in the clip on the right of `this example model
    joint damping which is integrated implicitly by the Euler method, tendon damping is not integrated implicitly, thus
    joint damping should be used if possible.
 
+.. image:: images/XMLreference/tendon_armature.gif
+   :width: 30%
+   :align: right
+   :class: only-light
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml
+.. image:: images/XMLreference/tendon_armature_dark.gif
+   :width: 30%
+   :align: right
+   :class: only-dark
+   :target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml
+
+.. _tendon-spatial-armature:
+
+:at:`armature`: :at-val:`real, "0"`
+   Inertia associated with changes in tendon length. Setting this attribute to a positive value :math:`m` adds a kinetic
+   energy term :math:`\frac{1}{2}mv^2`, where :math:`v` is the tendon velocity. Tendon inertia is most valuable
+   when modeling the :ref:`armature<body-joint-armature>` inertia in a linear actuator which contains a spinning element
+   or the inertial motion of a fluid in a linear hydraulic actuator. In the illustration, we compare (*left*) a 3-dof
+   system with a "tendon" implemented with a rotational joint and a slider joint with
+   :ref:`armature<body-joint-armature>`, attached to the world with a :ref:`connect<equality-connect>` constraint and
+   (*right*) an equivalent 1-dof model with an armature-bearing tendon. Like joint :ref:`armature<body-joint-armature>`,
+   this added inertia is only associated with changes in tendon length, and would not affect the dynamics of a moving
+   fixed-length tendon. Because the tendon Jacobian :math:`J` is position-dependent, tendon armature leads to an
+   additional bias-force term :math:`c = m J \dot{J}^T \dot{q}`.
+
 .. _tendon-spatial-user:
 
 :at:`user`: :at-val:`real(nuser_tendon), "0 0 ..."`
@@ -4772,8 +4887,12 @@ as above.
 
 .. _tendon-fixed-limited:
 
+.. _tendon-fixed-actuatorfrclimited:
+
 .. _tendon-fixed-range:
 
+.. _tendon-fixed-actuatorfrcrange:
+
 .. _tendon-fixed-solreflimit:
 
 .. _tendon-fixed-solimplimit:
@@ -4792,6 +4911,8 @@ as above.
 
 .. _tendon-fixed-damping:
 
+.. _tendon-fixed-armature:
+
 .. _tendon-fixed-user:
 
 .. |tendon/fixed attrib list| replace::
@@ -6335,6 +6456,33 @@ joint or when a single actuator acts on multiple joints. See :ref:`CForceRange`
    The joint where actuator forces will be sensed. The sensor output is copied from ``mjData.qfrc_actuator``.
 
 
+.. _sensor-tendonactuatorfrc:
+
+:el-prefix:`sensor/` |-| **tendonactuatorfrc** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This element creates an actuator force sensor, measured at a tendon. The quantity being sensed is the total force
+contributed by all actuators to a single tendon. This type of sensor is important when multiple actuators act on a
+single tendon. See :ref:`CForceRange` for details.
+
+
+.. _sensor-tendonactuatorfrc-name:
+
+.. _sensor-tendonactuatorfrc-noise:
+
+.. _sensor-tendonactuatorfrc-cutoff:
+
+.. _sensor-tendonactuatorfrc-user:
+
+:at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
+   See :ref:`CSensor`.
+
+.. _sensor-tendonactuatorfrc-tendon:
+
+:at:`tendon`: :at-val:`string, required`
+   The tendon where actuator forces will be sensed.
+
+
 .. _sensor-ballquat:
 
 :el-prefix:`sensor/` |-| **ballquat** (*)
@@ -8152,7 +8300,7 @@ if omitted.
 
 .. _default-light-dir:
 
-.. _default-light-bulbradius:
+.. _default-light-type:
 
 .. _default-light-directional:
 
@@ -8160,18 +8308,24 @@ if omitted.
 
 .. _default-light-active:
 
+.. _default-light-diffuse:
+
+.. _default-light-intensity:
+
+.. _default-light-ambient:
+
+.. _default-light-specular:
+
+.. _default-light-bulbradius:
+
+.. _default-light-range:
+
 .. _default-light-attenuation:
 
 .. _default-light-cutoff:
 
 .. _default-light-exponent:
 
-.. _default-light-ambient:
-
-.. _default-light-diffuse:
-
-.. _default-light-specular:
-
 .. _default-light-mode:
 
 :el-prefix:`default/` |-| **light** (?)
@@ -8227,8 +8381,12 @@ if omitted.
 
 .. _default-tendon-limited:
 
+.. _default-tendon-actuatorfrclimited:
+
 .. _default-tendon-range:
 
+.. _default-tendon-actuatorfrcrange:
+
 .. _default-tendon-solreflimit:
 
 .. _default-tendon-solimplimit:

doc/XMLschema.rst

@@ -54,7 +54,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`discardvisual<compiler-discardvisual>`                    | :ref:`usethread<compiler-usethread>`                            | :ref:`fusestatic<compiler-fusestatic>`                          | :ref:`inertiafromgeom<compiler-inertiafromgeom>`                | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`inertiagrouprange<compiler-inertiagrouprange>`            | :ref:`assetdir<compiler-assetdir>`                              | :ref:`alignfree<compiler-alignfree>`                            |                                                                 | |
+|                                    |    |    | :ref:`inertiagrouprange<compiler-inertiagrouprange>`            | :ref:`saveinertial<compiler-saveinertial>`                      | :ref:`assetdir<compiler-assetdir>`                              | :ref:`alignfree<compiler-alignfree>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| compiler |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -153,17 +153,17 @@
 | :ref:`texture                      | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <asset-texture>`                   |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`name<asset-texture-name>`                                 | :ref:`type<asset-texture-type>`                                 | :ref:`content_type<asset-texture-content_type>`                 | :ref:`file<asset-texture-file>`                                 | |
+|                                    |    |    | :ref:`name<asset-texture-name>`                                 | :ref:`type<asset-texture-type>`                                 | :ref:`colorspace<asset-texture-colorspace>`                     | :ref:`content_type<asset-texture-content_type>`                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`gridsize<asset-texture-gridsize>`                         | :ref:`gridlayout<asset-texture-gridlayout>`                     | :ref:`fileright<asset-texture-fileright>`                       | :ref:`fileleft<asset-texture-fileleft>`                         | |
+|                                    |    |    | :ref:`file<asset-texture-file>`                                 | :ref:`gridsize<asset-texture-gridsize>`                         | :ref:`gridlayout<asset-texture-gridlayout>`                     | :ref:`fileright<asset-texture-fileright>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`fileup<asset-texture-fileup>`                             | :ref:`filedown<asset-texture-filedown>`                         | :ref:`filefront<asset-texture-filefront>`                       | :ref:`fileback<asset-texture-fileback>`                         | |
+|                                    |    |    | :ref:`fileleft<asset-texture-fileleft>`                         | :ref:`fileup<asset-texture-fileup>`                             | :ref:`filedown<asset-texture-filedown>`                         | :ref:`filefront<asset-texture-filefront>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`builtin<asset-texture-builtin>`                           | :ref:`rgb1<asset-texture-rgb1>`                                 | :ref:`rgb2<asset-texture-rgb2>`                                 | :ref:`mark<asset-texture-mark>`                                 | |
+|                                    |    |    | :ref:`fileback<asset-texture-fileback>`                         | :ref:`builtin<asset-texture-builtin>`                           | :ref:`rgb1<asset-texture-rgb1>`                                 | :ref:`rgb2<asset-texture-rgb2>`                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`markrgb<asset-texture-markrgb>`                           | :ref:`random<asset-texture-random>`                             | :ref:`width<asset-texture-width>`                               | :ref:`height<asset-texture-height>`                             | |
+|                                    |    |    | :ref:`mark<asset-texture-mark>`                                 | :ref:`markrgb<asset-texture-markrgb>`                           | :ref:`random<asset-texture-random>`                             | :ref:`width<asset-texture-width>`                               | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`hflip<asset-texture-hflip>`                               | :ref:`vflip<asset-texture-vflip>`                               | :ref:`nchannel<asset-texture-nchannel>`                         |                                                                 | |
+|                                    |    |    | :ref:`height<asset-texture-height>`                             | :ref:`hflip<asset-texture-hflip>`                               | :ref:`vflip<asset-texture-vflip>`                               | :ref:`nchannel<asset-texture-nchannel>`                         | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| asset |br| |_| |L|             |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -313,13 +313,15 @@
 | :ref:`light                        | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <body-light>`                      |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`name<body-light-name>`                                    | :ref:`class<body-light-class>`                                  | :ref:`directional<body-light-directional>`                      | :ref:`castshadow<body-light-castshadow>`                        | |
+|                                    |    |    | :ref:`name<body-light-name>`                                    | :ref:`class<body-light-class>`                                  | :ref:`directional<body-light-directional>`                      | :ref:`type<body-light-type>`                                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`active<body-light-active>`                                | :ref:`pos<body-light-pos>`                                      | :ref:`dir<body-light-dir>`                                      | :ref:`bulbradius<body-light-bulbradius>`                        | |
+|                                    |    |    | :ref:`castshadow<body-light-castshadow>`                        | :ref:`active<body-light-active>`                                | :ref:`pos<body-light-pos>`                                      | :ref:`dir<body-light-dir>`                                      | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`attenuation<body-light-attenuation>`                      | :ref:`cutoff<body-light-cutoff>`                                | :ref:`exponent<body-light-exponent>`                            | :ref:`ambient<body-light-ambient>`                              | |
+|                                    |    |    | :ref:`bulbradius<body-light-bulbradius>`                        | :ref:`intensity<body-light-intensity>`                          | :ref:`range<body-light-range>`                                  | :ref:`attenuation<body-light-attenuation>`                      | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`diffuse<body-light-diffuse>`                              | :ref:`specular<body-light-specular>`                            | :ref:`mode<body-light-mode>`                                    | :ref:`target<body-light-target>`                                | |
+|                                    |    |    | :ref:`cutoff<body-light-cutoff>`                                | :ref:`exponent<body-light-exponent>`                            | :ref:`ambient<body-light-ambient>`                              | :ref:`diffuse<body-light-diffuse>`                              | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`specular<body-light-specular>`                            | :ref:`mode<body-light-mode>`                                    | :ref:`target<body-light-target>`                                | :ref:`texture<body-light-texture>`                              | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| body |br| |_| |L|              |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -344,6 +346,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`vertex<body-composite-vertex>`                            | :ref:`initial<body-composite-initial>`                          | :ref:`curve<body-composite-curve>`                              | :ref:`size<body-composite-size>`                                | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`quat<body-composite-quat>`                                |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| composite |br| |_2| |L|       |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`joint                        | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -439,6 +443,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`young<flexcomp-elasticity-young>`                         | :ref:`poisson<flexcomp-elasticity-poisson>`                     | :ref:`damping<flexcomp-elasticity-damping>`                     | :ref:`thickness<flexcomp-elasticity-thickness>`                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`elastic2d<flexcomp-elasticity-elastic2d>`                 |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`contact                      | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -450,7 +456,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`margin<flexcomp-contact-margin>`                          | :ref:`gap<flexcomp-contact-gap>`                                | :ref:`internal<flexcomp-contact-internal>`                      | :ref:`selfcollide<flexcomp-contact-selfcollide>`                | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`activelayers<flexcomp-contact-activelayers>`              |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`activelayers<flexcomp-contact-activelayers>`              | :ref:`vertcollide<flexcomp-contact-vertcollide>`                |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -500,7 +506,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`margin<flex-contact-margin>`                              | :ref:`gap<flex-contact-gap>`                                    | :ref:`internal<flex-contact-internal>`                          | :ref:`selfcollide<flex-contact-selfcollide>`                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`activelayers<flex-contact-activelayers>`                  |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`activelayers<flex-contact-activelayers>`                  | :ref:`vertcollide<flex-contact-vertcollide>`                    |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| flex |br| |_2| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -516,6 +522,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`young<flex-elasticity-young>`                             | :ref:`poisson<flex-elasticity-poisson>`                         | :ref:`damping<flex-elasticity-damping>`                         | :ref:`thickness<flex-elasticity-thickness>`                     | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`elastic2d<flex-elasticity-elastic2d>`                     |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| deformable |br| |_| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`skin                         | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -619,13 +627,15 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<tendon-spatial-name>`                                | :ref:`class<tendon-spatial-class>`                              | :ref:`group<tendon-spatial-group>`                              | :ref:`limited<tendon-spatial-limited>`                          | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`range<tendon-spatial-range>`                              | :ref:`solreflimit<tendon-spatial-solreflimit>`                  | :ref:`solimplimit<tendon-spatial-solimplimit>`                  | :ref:`solreffriction<tendon-spatial-solreffriction>`            | |
+|                                    |    |    | :ref:`actuatorfrclimited<tendon-spatial-actuatorfrclimited>`    | :ref:`range<tendon-spatial-range>`                              | :ref:`actuatorfrcrange<tendon-spatial-actuatorfrcrange>`        | :ref:`solreflimit<tendon-spatial-solreflimit>`                  | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`solimpfriction<tendon-spatial-solimpfriction>`            | :ref:`frictionloss<tendon-spatial-frictionloss>`                | :ref:`springlength<tendon-spatial-springlength>`                | :ref:`width<tendon-spatial-width>`                              | |
+|                                    |    |    | :ref:`solimplimit<tendon-spatial-solimplimit>`                  | :ref:`solreffriction<tendon-spatial-solreffriction>`            | :ref:`solimpfriction<tendon-spatial-solimpfriction>`            | :ref:`frictionloss<tendon-spatial-frictionloss>`                | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`material<tendon-spatial-material>`                        | :ref:`margin<tendon-spatial-margin>`                            | :ref:`stiffness<tendon-spatial-stiffness>`                      | :ref:`damping<tendon-spatial-damping>`                          | |
+|                                    |    |    | :ref:`springlength<tendon-spatial-springlength>`                | :ref:`width<tendon-spatial-width>`                              | :ref:`material<tendon-spatial-material>`                        | :ref:`margin<tendon-spatial-margin>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`rgba<tendon-spatial-rgba>`                                | :ref:`user<tendon-spatial-user>`                                |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`stiffness<tendon-spatial-stiffness>`                      | :ref:`damping<tendon-spatial-damping>`                          | :ref:`armature<tendon-spatial-armature>`                        | :ref:`rgba<tendon-spatial-rgba>`                                | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`user<tendon-spatial-user>`                                |                                                                 |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| spatial |br| |_2| |L|         |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -655,11 +665,13 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<tendon-fixed-name>`                                  | :ref:`class<tendon-fixed-class>`                                | :ref:`group<tendon-fixed-group>`                                | :ref:`limited<tendon-fixed-limited>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`range<tendon-fixed-range>`                                | :ref:`solreflimit<tendon-fixed-solreflimit>`                    | :ref:`solimplimit<tendon-fixed-solimplimit>`                    | :ref:`solreffriction<tendon-fixed-solreffriction>`              | |
+|                                    |    |    | :ref:`actuatorfrclimited<tendon-fixed-actuatorfrclimited>`      | :ref:`range<tendon-fixed-range>`                                | :ref:`actuatorfrcrange<tendon-fixed-actuatorfrcrange>`          | :ref:`solreflimit<tendon-fixed-solreflimit>`                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`solimpfriction<tendon-fixed-solimpfriction>`              | :ref:`frictionloss<tendon-fixed-frictionloss>`                  | :ref:`springlength<tendon-fixed-springlength>`                  | :ref:`margin<tendon-fixed-margin>`                              | |
+|                                    |    |    | :ref:`solimplimit<tendon-fixed-solimplimit>`                    | :ref:`solreffriction<tendon-fixed-solreffriction>`              | :ref:`solimpfriction<tendon-fixed-solimpfriction>`              | :ref:`frictionloss<tendon-fixed-frictionloss>`                  | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`stiffness<tendon-fixed-stiffness>`                        | :ref:`damping<tendon-fixed-damping>`                            | :ref:`user<tendon-fixed-user>`                                  |                                                                 | |
+|                                    |    |    | :ref:`springlength<tendon-fixed-springlength>`                  | :ref:`margin<tendon-fixed-margin>`                              | :ref:`stiffness<tendon-fixed-stiffness>`                        | :ref:`damping<tendon-fixed-damping>`                            | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`armature<tendon-fixed-armature>`                          | :ref:`user<tendon-fixed-user>`                                  |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| fixed |br| |_2| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -1002,6 +1014,15 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`tendonactuatorfrc            | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <sensor-tendonactuatorfrc>`        |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<sensor-tendonactuatorfrc-name>`                      | :ref:`tendon<sensor-tendonactuatorfrc-tendon>`                  | :ref:`cutoff<sensor-tendonactuatorfrc-cutoff>`                  | :ref:`noise<sensor-tendonactuatorfrc-noise>`                    | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`user<sensor-tendonactuatorfrc-user>`                      |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| sensor |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`ballquat                     | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <sensor-ballquat>`                 |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
@@ -1442,13 +1463,13 @@
 | :ref:`light                        | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
 | <default-light>`                   |    |                                                                                                                                                                                                                                                                              |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`pos<default-light-pos>`                                   | :ref:`dir<default-light-dir>`                                   | :ref:`bulbradius<default-light-bulbradius>`                     | :ref:`directional<default-light-directional>`                   | |
+|                                    |    |    | :ref:`pos<default-light-pos>`                                   | :ref:`dir<default-light-dir>`                                   | :ref:`bulbradius<default-light-bulbradius>`                     | :ref:`intensity<default-light-intensity>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`castshadow<default-light-castshadow>`                     | :ref:`active<default-light-active>`                             | :ref:`attenuation<default-light-attenuation>`                   | :ref:`cutoff<default-light-cutoff>`                             | |
+|                                    |    |    | :ref:`range<default-light-range>`                               | :ref:`directional<default-light-directional>`                   | :ref:`type<default-light-type>`                                 | :ref:`castshadow<default-light-castshadow>`                     | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`exponent<default-light-exponent>`                         | :ref:`ambient<default-light-ambient>`                           | :ref:`diffuse<default-light-diffuse>`                           | :ref:`specular<default-light-specular>`                         | |
+|                                    |    |    | :ref:`active<default-light-active>`                             | :ref:`attenuation<default-light-attenuation>`                   | :ref:`cutoff<default-light-cutoff>`                             | :ref:`exponent<default-light-exponent>`                         | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`mode<default-light-mode>`                                 |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`ambient<default-light-ambient>`                           | :ref:`diffuse<default-light-diffuse>`                           | :ref:`specular<default-light-specular>`                         | :ref:`mode<default-light-mode>`                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| default |br| |_| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |

doc/changelog.rst

@@ -2,26 +2,100 @@
 Changelog
 =========
 
-Upcoming version (not yet released)
+Upcoming version (not yet release)
------------------------------------
+----------------------------------
+
+General
+^^^^^^^
+- Refactored island implementation so that island data is memory-contiguous. This speeds up island processing in the
+  solver and clears the way for the addition of the Newton and PGS solvers (currently only CG is supported).
+- Removed the :at:`shell` plugin. This is now supported by :ref:`flexcomp<body-flexcomp>` and is active depending on
+  the :ref:`elastic2d<flexcomp-elasticity-elastic2d>` attribute (off by default).
+- Replaced the :ref:`directional<body-light-directional>` (boolean) field for lights with a
+  :ref:`type<body-light-type>` field (of type :ref:`mjtLightType<mjtLightType>`) to allow for additional lighting
+  types.
+- Added new sub-component :ref:`mj_makeM` which combines the :ref:`mj_crb` call with additional logic to support the
+  introduction in 3.3.1 of :ref:`tendon armature<tendon-spatial-armature>`. In addition to the traditional
+  ``mjData.qM``, :ref:`mj_makeM` also computes ``mjData.M``, a CSR representation of the same matrix.
+- Added a new function :ref:`mj_copyBack` to copy real-valued arrays in an mjModel to a compatible mjSpec.
+
+Simulate
+^^^^^^^^
+- The struct ``mjv_sceneState`` has been removed. This struct was used for partial synchronization of ``mjModel`` and
+  ``mjData`` when the Python viewer is used in passive mode. This functionality is now provided by :ref:`mjv_copyModel`
+  and :ref:`mjv_copyData`, which don't copy arrays which are not required for visualization.
+
+.. image:: images/changelog/procedural_terrain_generation.png
+   :width: 25%
+   :align: right
+
+Python bindings
+^^^^^^^^^^^^^^^
+
+- Added examples of procedural terrain generation to the Model Editing tutorial: |mjspec_colab|
+
+Version 3.3.2 (April 28, 2025)
+------------------------------
+
+MJX
+^^^
+1. Added inverse dynamics.
+2. Added tendon actuator force sensor.
+3. Fix :github:issue:`2606` such that ``make_data`` copies over ``mocap_pos`` and ``mocap_quat``
+   from ``body_pos`` and ``body_quat``.
+
+Version 3.3.1 (Apr 9, 2025)
+----------------------------
 
 .. admonition:: Breaking API changes
    :class: attention
 
-   - The default value of the flag for toggling :ref:`internal flex contacts<flex-contact-internal>` was changed from
+   1. The default value of the flag for toggling :ref:`internal flex contacts<flex-contact-internal>` was changed from
       "true" to "false". This feature has proven to be counterintuitive for users.
+   2. All of the attach functions (``mjs_attachBody``, ``mjs_attachFrame``, ``mjs_attachToSite``,
+      ``mjs_attachFrameToSite``) have been removed and replaced by a single function :ref:`mjs_attach`.
+
+General
+^^^^^^^
+3. Added :ref:`tendon armature<tendon-spatial-armature>`: inertia associated with changes in tendon length.
+4. Added the :ref:`compiler/saveinertial<compiler-saveinertial>` flag, writing explicit inertial clauses for all
+   bodies when saving to XML.
+5. Added :ref:`orientation<body-composite-quat>` attribute to :ref:`composite<body-composite>`. Moreover, allow the
+   composite to be the direct child of a frame.
+6. Added :ref:`tendon actuator force limits<tendon-spatial-actuatorfrclimited>` and
+   :ref:`tendon actuator force sensor<sensor-tendonactuatorfrc>`.
+
+MJX
+^^^
+7. Added tendon actuator force limits.
 
 Bug fixes
 ^^^^^^^^^
-- :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
+8. :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
    which the Jacobian is computed, now fixed.
-- Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
+9. Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
    mjSpec to a frame or a site.
+10. Fixed a bug that caused shadow rendering to flicker on platforms (e.g., MacOS) that do not support
+    ARB_clip_control. Fixed in collaboration with :github:user:`aftersomemath`.
+
+Python bindings
+^^^^^^^^^^^^^^^
+
+.. youtube:: LbANnKMDOHg
+   :aspect: 16:7
+   :align: right
+   :width: 240px
+
+11. Added examples of procedural model creation to the Model Editing tutorial: |mjspec_colab|
+12. Added support for nameless :ref:`mjSpec` objects in the ``bind`` method, see the corresponding
+    :ref:`section<PyMJCF>` in the documentation.
+
+.. |mjspec_colab| image:: https://colab.research.google.com/assets/colab-badge.svg
+                  :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb
 
 Version 3.3.0 (Feb 26, 2025)
 ----------------------------
 
-
 Feature promotion
 ^^^^^^^^^^^^^^^^^
 .. youtube:: qJFbx-FR7Bc

doc/computation/index.rst

@@ -1692,7 +1692,7 @@ The stages below compute quantities that depend on the generalized positions ``m
 4. Compute quantities related to :ref:`flex<deformable-flex>` objects: :ref:`mj_flex`
 5. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
    tendons: :ref:`mj_tendon`
-6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb`
+6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_makeM`
 7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM`
 8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection:
    :ref:`mj_collision`

doc/includes/references.h

@@ -154,7 +154,6 @@ struct mjData_ {
 
   // solver statistics
   mjSolverStat  solver[mjNISLAND*mjNSOLVER];  // solver statistics per island, per iteration
-  int           solver_nisland;               // number of islands processed by solver
   int           solver_niter[mjNISLAND];      // number of solver iterations, per island
   int           solver_nnz[mjNISLAND];        // number of nonzeros in Hessian or efc_AR, per island
   mjtNum        solver_fwdinv[2];             // forward-inverse comparison: qfrc, efc
@@ -172,6 +171,7 @@ struct mjData_ {
   int     nJ;                // number of non-zeros in constraint Jacobian
   int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
+  int     nidof;             // number of dofs in all islands
 
   // global properties
   mjtNum  time;              // simulation time
@@ -267,12 +267,13 @@ struct mjData_ {
   int*    moment_colind;     // column indices in sparse Jacobian                (nJmom x 1)
   mjtNum* actuator_moment;   // actuator moments                                 (nJmom x 1)
 
-  // computed by mj_fwdPosition/mj_crb
+  // computed by mj_fwdPosition/mj_makeM
   mjtNum* crb;               // com-based composite inertia and mass             (nbody x 10)
-  mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
+  mjtNum* qM;                // inertia (sparse)                                 (nM x 1)
+  mjtNum* M;                 // reduced inertia (compressed sparse row)          (nC x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -305,27 +306,23 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
   int*    B_rownnz;          // body-dof: non-zeros in each row                  (nbody x 1)
   int*    B_rowadr;          // body-dof: address of each row in B_colind        (nbody x 1)
   int*    B_colind;          // body-dof: column indices of non-zeros            (nB x 1)
-  int*    M_rownnz;          // inertia: non-zeros in each row                   (nv x 1)
+  int*    M_rownnz;          // reduced inertia: non-zeros in each row           (nv x 1)
-  int*    M_rowadr;          // inertia: address of each row in M_colind         (nv x 1)
+  int*    M_rowadr;          // reduced inertia: address of each row in M_colind (nv x 1)
-  int*    M_colind;          // inertia: column indices of non-zeros             (nM x 1)
+  int*    M_colind;          // reduced inertia: column indices of non-zeros     (nC x 1)
-  int*    mapM2M;            // index mapping from M (legacy) to M (CSR)         (nM x 1)
+  int*    mapM2M;            // index mapping from qM to M                       (nC x 1)
-  int*    C_rownnz;          // reduced dof-dof: non-zeros in each row           (nv x 1)
+  int*    D_rownnz;          // full inertia: non-zeros in each row              (nv x 1)
-  int*    C_rowadr;          // reduced dof-dof: address of each row in C_colind (nv x 1)
+  int*    D_rowadr;          // full inertia: address of each row in D_colind    (nv x 1)
-  int*    C_colind;          // reduced dof-dof: column indices of non-zeros     (nC x 1)
+  int*    D_diag;            // full inertia: index of diagonal element          (nv x 1)
-  int*    mapM2C;            // index mapping from M to C                        (nC x 1)
+  int*    D_colind;          // full inertia: column indices of non-zeros        (nD x 1)
-  int*    D_rownnz;          // dof-dof: non-zeros in each row                   (nv x 1)
+  int*    mapM2D;            // index mapping from qM to D                       (nD x 1)
-  int*    D_rowadr;          // dof-dof: address of each row in D_colind         (nv x 1)
+  int*    mapD2M;            // index mapping from D to qM                       (nM x 1)
-  int*    D_diag;            // dof-dof: index of diagonal element               (nv x 1)
-  int*    D_colind;          // dof-dof: column indices of non-zeros             (nD x 1)
-  int*    mapM2D;            // index mapping from M to D                        (nD x 1)
-  int*    mapD2M;            // index mapping from D to M                        (nM x 1)
 
   // computed by mj_implicit/mj_derivative
   mjtNum* qDeriv;            // d (passive + actuator - bias) / d qvel           (nD x 1)
@@ -347,8 +344,8 @@ struct mjData_ {
   mjtNum* qfrc_constraint;   // constraint force                                 (nv x 1)
 
   // computed by mj_inverse
-  mjtNum* qfrc_inverse;      // net external force; should equal:                (nv x 1)
+  mjtNum* qfrc_inverse;      // net external force; should equal:
-                             // qfrc_applied + J'*xfrc_applied + qfrc_actuator
+                             // qfrc_applied + J'*xfrc_applied + qfrc_actuator   (nv x 1)
 
   // computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format
   mjtNum* cacc;              // com-based acceleration                           (nbody x 6)
@@ -382,16 +379,50 @@ struct mjData_ {
   mjtNum* efc_R;             // inverse constraint mass                          (nefc x 1)
   int*    tendon_efcadr;     // first efc address involving tendon; -1: none     (ntendon x 1)
 
-  // computed by mj_island
+  // computed by mj_island (island dof structure)
   int*    dof_island;        // island id of this dof; -1: none                  (nv x 1)
-  int*    island_dofnum;     // number of dofs in island                         (nisland x 1)
+  int*    island_nv;         // number of dofs in this island                    (nisland x 1)
-  int*    island_dofadr;     // start address in island_dofind                   (nisland x 1)
+  int*    island_idofadr;    // island start address in idof vector              (nisland x 1)
-  int*    island_dofind;     // island dof indices; -1: none                     (nv x 1)
+  int*    island_dofadr;     // island start address in dof vector               (nisland x 1)
-  int*    dof_islandind;     // dof island indices; -1: none                     (nv x 1)
+  int*    map_dof2idof;      // map from dof to idof                             (nv x 1)
+  int*    map_idof2dof;      // map from idof to dof;  >= nidof: unconstrained   (nv x 1)
+
+  // computed by mj_island (dofs sorted by island)
+  mjtNum* ifrc_smooth;       // net unconstrained force                          (nidof x 1)
+  mjtNum* iacc_smooth;       // unconstrained acceleration                       (nidof x 1)
+  int*    iM_rownnz;         // inertia: non-zeros in each row                   (nidof x 1)
+  int*    iM_rowadr;         // inertia: address of each row in iM_colind        (nidof x 1)
+  int*    iM_colind;         // inertia: column indices of non-zeros             (nC x 1)
+  mjtNum* iM;                // total inertia (sparse)                           (nC x 1)
+  mjtNum* iLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
+  mjtNum* iLDiagInv;         // 1/diag(D)                                        (nidof x 1)
+  mjtNum* iacc;              // acceleration                                     (nidof x 1)
+
+  // computed by mj_island (island constraint structure)
   int*    efc_island;        // island id of this constraint                     (nefc x 1)
-  int*    island_efcnum;     // number of constraints in island                  (nisland x 1)
+  int*    island_ne;         // number of equality constraints in island         (nisland x 1)
-  int*    island_efcadr;     // start address in island_efcind                   (nisland x 1)
+  int*    island_nf;         // number of friction constraints in island         (nisland x 1)
-  int*    island_efcind;     // island constraint indices                        (nefc x 1)
+  int*    island_nefc;       // number of constraints in island                  (nisland x 1)
+  int*    island_iefcadr;    // start address in iefc vector                     (nisland x 1)
+  int*    map_efc2iefc;      // map from efc to iefc                             (nefc x 1)
+  int*    map_iefc2efc;      // map from iefc to efc                             (nefc x 1)
+
+  // computed by mj_island (constraints sorted by island)
+  int*    iefc_type;         // constraint type (mjtConstraint)                  (nefc x 1)
+  int*    iefc_id;           // id of object of specified type                   (nefc x 1)
+  int*    iefc_J_rownnz;     // number of non-zeros in constraint Jacobian row   (nefc x 1)
+  int*    iefc_J_rowadr;     // row start address in colind array                (nefc x 1)
+  int*    iefc_J_rowsuper;   // number of subsequent rows in supernode           (nefc x 1)
+  int*    iefc_J_colind;     // column indices in constraint Jacobian            (nJ x 1)
+  int*    iefc_JT_rownnz;    // number of non-zeros in constraint Jacobian row T (nidof x 1)
+  int*    iefc_JT_rowadr;    // row start address in colind array              T (nidof x 1)
+  int*    iefc_JT_rowsuper;  // number of subsequent rows in supernode         T (nidof x 1)
+  int*    iefc_JT_colind;    // column indices in constraint Jacobian          T (nJ x 1)
+  mjtNum* iefc_J;            // constraint Jacobian                              (nJ x 1)
+  mjtNum* iefc_JT;           // constraint Jacobian transposed                   (nJ x 1)
+  mjtNum* iefc_frictionloss; // frictionloss (friction)                          (nefc x 1)
+  mjtNum* iefc_D;            // constraint mass                                  (nefc x 1)
+  mjtNum* iefc_R;            // inverse constraint mass                          (nefc x 1)
 
   // computed by mj_projectConstraint (PGS solver)
   int*    efc_AR_rownnz;     // number of non-zeros in AR                        (nefc x 1)
@@ -409,11 +440,18 @@ struct mjData_ {
 
   // computed by mj_fwdConstraint/mj_inverse
   mjtNum* efc_b;             // linear cost term: J*qacc_smooth - aref           (nefc x 1)
-  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
+  mjtNum* iefc_aref;         // reference pseudo-acceleration                    (nefc x 1)
+  int*    iefc_state;        // constraint state (mjtConstraintState)            (nefc x 1)
+  mjtNum* iefc_force;        // constraint force in constraint space             (nefc x 1)
   int*    efc_state;         // constraint state (mjtConstraintState)            (nefc x 1)
+  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
+  mjtNum* ifrc_constraint;   // constraint force                                 (nidof x 1)
 
   // thread pool pointer
   uintptr_t threadpool;
+
+  // compilation signature
+  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
 typedef struct mjData_ mjData;
 typedef enum mjtDisableBit_ {     // disable default feature bitflags
@@ -488,6 +526,12 @@ typedef enum mjtCamLight_ {       // tracking mode for camera and light
   mjCAMLIGHT_TARGETBODY,          // pos fixed in body, rot tracks target body
   mjCAMLIGHT_TARGETBODYCOM        // pos fixed in body, rot tracks target subtree com
 } mjtCamLight;
+typedef enum mjtLightType_ {      // type of light
+  mjLIGHT_SPOT    = 0,            // spot
+  mjLIGHT_DIRECTIONAL,            // directional
+  mjLIGHT_POINT,                  // point
+  mjLIGHT_IMAGE,                  // image-based
+} mjtLightType;
 typedef enum mjtTexture_ {        // type of texture
   mjTEXTURE_2D        = 0,        // 2d texture, suitable for planes and hfields
   mjTEXTURE_CUBE,                 // cube texture, suitable for all other geom types
@@ -506,6 +550,11 @@ typedef enum mjtTextureRole_ {    // role of texture map in rendering
   mjTEXROLE_ORM,                  // occlusion, roughness, metallic
   mjNTEXROLE
 } mjtTextureRole;
+typedef enum mjtColorSpace_ {     // type of color space encoding
+  mjCOLORSPACE_AUTO   = 0,        // attempts to autodetect color space, defaults to linear
+  mjCOLORSPACE_LINEAR,            // linear color space
+  mjCOLORSPACE_SRGB               // standard RGB color space
+} mjtColorSpace;
 typedef enum mjtIntegrator_ {     // integrator mode
   mjINT_EULER         = 0,        // semi-implicit Euler
   mjINT_RK4,                      // 4th-order Runge Kutta
@@ -604,7 +653,8 @@ typedef enum mjtObj_ {            // type of MujoCo object
 
   // meta elements, do not appear in mjModel
   mjOBJ_FRAME         = 100,      // frame
-  mjOBJ_DEFAULT                   // default
+  mjOBJ_DEFAULT,                  // default
+  mjOBJ_MODEL                     // entire model
 
 } mjtObj;
 typedef enum mjtConstraint_ {     // type of constraint
@@ -645,6 +695,7 @@ typedef enum mjtSensor_ {         // type of sensor
   mjSENS_ACTUATORVEL,             // scalar actuator velocity
   mjSENS_ACTUATORFRC,             // scalar actuator force
   mjSENS_JOINTACTFRC,             // scalar actuator force, measured at the joint
+  mjSENS_TENDONACTFRC,            // scalar actuator force, measured at the tendon
 
   // sensors related to ball joints
   mjSENS_BALLQUAT,                // 4D ball joint quaternion
@@ -722,6 +773,12 @@ typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
   mjFLEXSELF_SAP,                 // use SAP in midphase
   mjFLEXSELF_AUTO                 // choose between BVH and SAP automatically
 } mjtFlexSelf;
+typedef enum mjtSDFType_ {        // signed distance function (SDF) type
+  mjSDFTYPE_SINGLE     = 0,       // single SDF
+  mjSDFTYPE_INTERSECTION,         // max(A, B)
+  mjSDFTYPE_MIDSURFACE,           // A - B
+  mjSDFTYPE_COLLISION,            // A + B + abs(max(A, B))
+} mjtSDFType;
 struct mjLROpt_ {                 // options for mj_setLengthRange()
   // flags
   int mode;                       // which actuators to process (mjtLRMode)
@@ -1128,9 +1185,12 @@ struct mjModel_ {
   int*      light_mode;           // light tracking mode (mjtCamLight)        (nlight x 1)
   int*      light_bodyid;         // id of light's body                       (nlight x 1)
   int*      light_targetbodyid;   // id of targeted body; -1: none            (nlight x 1)
-  mjtByte*  light_directional;    // directional light                        (nlight x 1)
+  int*      light_type;           // spot, directional, etc. (mjtLightType)   (nlight x 1)
+  int*      light_texid;          // texture id for image lights              (nlight x 1)
   mjtByte*  light_castshadow;     // does light cast shadows                  (nlight x 1)
   float*    light_bulbradius;     // light radius for soft shadows            (nlight x 1)
+  float*    light_intensity;      // intensity, in candela                    (nlight x 1)
+  float*    light_range;          // range of effectiveness                   (nlight x 1)
   mjtByte*  light_active;         // is light on                              (nlight x 1)
   mjtNum*   light_pos;            // position rel. to body frame              (nlight x 3)
   mjtNum*   light_dir;            // direction rel. to body frame             (nlight x 3)
@@ -1182,6 +1242,7 @@ struct mjModel_ {
   int*      flex_nodebodyid;      // node body ids                            (nflexnode x 1)
   int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
   int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
+  int*      flex_edgeflap;        // adjacent vertex ids (dim=2 only)         (nflexedge x 2)
   int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
   int*      flex_elemtexcoord;    // element texture coordinates (dim+1)      (nflexelemdata x 1)
   int*      flex_elemedge;        // element edge ids                         (nflexelemedge x 1)
@@ -1196,6 +1257,7 @@ struct mjModel_ {
   mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
   mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
   mjtNum*   flex_stiffness;       // finite element stiffness matrix          (nflexelem x 21)
+  mjtNum*   flex_bending;         // bending stiffness                        (nflexedge x 16)
   mjtNum*   flex_damping;         // Rayleigh's damping coefficient           (nflex x 1)
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
@@ -1276,6 +1338,7 @@ struct mjModel_ {
 
   // textures
   int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
+  int*      tex_colorspace;       // texture colorspace (mjtColorSpace)       (ntex x 1)
   int*      tex_height;           // number of rows in texture image          (ntex x 1)
   int*      tex_width;            // number of columns in texture image       (ntex x 1)
   int*      tex_nchannel;         // number of channels in texture image      (ntex x 1)
@@ -1326,15 +1389,18 @@ struct mjModel_ {
   int*      tendon_matid;         // material id for rendering                (ntendon x 1)
   int*      tendon_group;         // group for visibility                     (ntendon x 1)
   mjtByte*  tendon_limited;       // does tendon have length limits           (ntendon x 1)
+  mjtByte*  tendon_actfrclimited; // does tendon have actuator force limits   (ntendon x 1)
   mjtNum*   tendon_width;         // width for rendering                      (ntendon x 1)
   mjtNum*   tendon_solref_lim;    // constraint solver reference: limit       (ntendon x mjNREF)
   mjtNum*   tendon_solimp_lim;    // constraint solver impedance: limit       (ntendon x mjNIMP)
   mjtNum*   tendon_solref_fri;    // constraint solver reference: friction    (ntendon x mjNREF)
   mjtNum*   tendon_solimp_fri;    // constraint solver impedance: friction    (ntendon x mjNIMP)
   mjtNum*   tendon_range;         // tendon length limits                     (ntendon x 2)
+  mjtNum*   tendon_actfrcrange;   // range of total actuator force            (ntendon x 2)
   mjtNum*   tendon_margin;        // min distance for limit detection         (ntendon x 1)
   mjtNum*   tendon_stiffness;     // stiffness coefficient                    (ntendon x 1)
   mjtNum*   tendon_damping;       // damping coefficient                      (ntendon x 1)
+  mjtNum*   tendon_armature;      // inertia associated with tendon velocity  (ntendon x 1)
   mjtNum*   tendon_frictionloss;  // loss due to friction                     (ntendon x 1)
   mjtNum*   tendon_lengthspring;  // spring resting length range              (ntendon x 2)
   mjtNum*   tendon_length0;       // tendon length in qpos0                   (ntendon x 1)
@@ -1451,6 +1517,9 @@ struct mjModel_ {
 
   // paths
   char*     paths;                // paths to assets, 0-terminated            (npaths x 1)
+
+  // compilation signature
+  uint64_t  signature;            // also held by the mjSpec that compiled this model
 };
 typedef struct mjModel_ mjModel;
 struct mjResource_ {
@@ -1536,6 +1605,15 @@ struct mjpPlugin_ {
   void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
 };
 typedef struct mjpPlugin_ mjpPlugin;
+struct mjSDF_ {
+  const mjpPlugin** plugin;
+  int* id;
+  mjtSDFType type;
+  mjtNum* relpos;
+  mjtNum* relmat;
+  mjtGeom* geomtype;
+};
+typedef struct mjSDF_ mjSDF;
 typedef enum mjtGridPos_ {        // grid position for overlay
   mjGRID_TOPLEFT      = 0,        // top left
   mjGRID_TOPRIGHT,                // top right
@@ -1713,6 +1791,7 @@ typedef enum mjtOrientation_ {     // type of orientation specifier
 } mjtOrientation;
 typedef struct mjsElement_ {       // element type, do not modify
   mjtObj elemtype;                 // element type
+  uint64_t signature;              // compilation signature
 } mjsElement;
 typedef struct mjsCompiler_ {      // compiler options
   mjtByte autolimits;              // infer "limited" attribute based on range
@@ -1728,6 +1807,7 @@ typedef struct mjsCompiler_ {      // compiler options
   mjtByte fusestatic;              // fuse static bodies with parent
   int inertiafromgeom;             // use geom inertias (mjtInertiaFromGeom)
   int inertiagrouprange[2];        // range of geom groups used to compute inertia
+  mjtByte saveinertial;            // save explicit inertial clause for all bodies to XML
   int alignfree;                   // align free joints with inertial frame
   mjLROpt LRopt;                   // options for lengthrange computation
 } mjsCompiler;
@@ -1963,9 +2043,12 @@ typedef struct mjsLight_ {         // light specification
 
   // intrinsics
   mjtByte active;                  // is light active
-  mjtByte directional;             // is light directional or spot
+  mjtLightType type;               // type of light
+  mjString* texture;               // texture name for image lights
   mjtByte castshadow;              // does light cast shadows
-  double bulbradius;               // bulb radius, for soft shadows
+  float bulbradius;                // bulb radius, for soft shadows
+  float intensity;                 // intensity, in candelas
+  float range;                     // range of effectiveness
   float attenuation[3];            // OpenGL attenuation (quadratic model)
   float cutoff;                    // OpenGL cutoff
   float exponent;                  // OpenGL exponent
@@ -1997,7 +2080,8 @@ typedef struct mjsFlex_ {          // flex specification
   double radius;                   // radius around primitive element
   mjtByte internal;                // enable internal collisions
   mjtByte flatskin;                // render flex skin with flat shading
-  int selfcollide;                 // mode for flex self colllision
+  int selfcollide;                 // mode for flex self collision
+  int vertcollide;                 // mode for vertex collision
   int activelayers;                // number of active element layers in 3D
   int group;                       // group for visualizatioh
   double edgestiffness;            // edge stiffness
@@ -2008,6 +2092,7 @@ typedef struct mjsFlex_ {          // flex specification
   double poisson;                  // Poisson's ratio
   double damping;                  // Rayleigh's damping
   double thickness;                // thickness (2D only)
+  int elastic2d;                   // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
 
   // mesh properties
   mjStringVec* nodebody;           // node body names
@@ -2079,6 +2164,7 @@ typedef struct mjsTexture_ {       // texture specification
   mjsElement* element;             // element type
   mjString* name;                  // name
   mjtTexture type;                 // texture type
+  mjtColorSpace colorspace;        // colorspace
 
   // method 1: builtin
   int builtin;                     // builtin type (mjtBuiltin)
@@ -2165,17 +2251,20 @@ typedef struct mjsTendon_ {        // tendon specification
   mjsElement* element;             // element type
   mjString* name;                  // name
 
-  // stiffness, damping, friction
+  // stiffness, damping, friction, armature
   double stiffness;                // stiffness coefficient
   double springlength[2];          // spring resting length; {-1, -1}: use qpos_spring
   double damping;                  // damping coefficient
   double frictionloss;             // friction loss
   mjtNum solref_friction[mjNREF];  // solver reference: tendon friction
   mjtNum solimp_friction[mjNIMP];  // solver impedance: tendon friction
+  double armature;                 // inertia associated with tendon velocity
 
   // length range
   int limited;                     // does tendon have limits (mjtLimited)
+  int actfrclimited;               // does tendon have actuator force limits
   double range[2];                 // length limits
+  double actfrcrange[2];           // actuator force limits
   double margin;                   // margin value for tendon limit detection
   mjtNum solref_limit[mjNREF];     // solver reference: tendon limits
   mjtNum solimp_limit[mjNIMP];     // solver impedance: tendon limits
@@ -2728,7 +2817,7 @@ typedef struct mjvGLCamera_ mjvGLCamera;
 struct mjvGeom_ {                 // abstract geom
   // type info
   int      type;                  // geom type (mjtGeom)
-  int      dataid;                // mesh, hfield or plane id; -1: none
+  int      dataid;                // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull)
   int      objtype;               // mujoco object type; mjOBJ_UNKNOWN for decor
   int      objid;                 // mujoco object id; -1 for decor
   int      category;              // visual category
@@ -2759,6 +2848,8 @@ typedef struct mjvGeom_ mjvGeom;
 struct mjvLight_ {                // OpenGL light
   float    pos[3];                // position rel. to body frame
   float    dir[3];                // direction rel. to body frame
+  int      type;                  // type (mjtLightType)
+  int      texid;                 // texture id for image lights
   float    attenuation[3];        // OpenGL attenuation (quadratic model)
   float    cutoff;                // OpenGL cutoff
   float    exponent;              // OpenGL exponent
@@ -2766,9 +2857,10 @@ struct mjvLight_ {                // OpenGL light
   float    diffuse[3];            // diffuse rgb (alpha=1)
   float    specular[3];           // specular rgb (alpha=1)
   mjtByte  headlight;             // headlight
-  mjtByte  directional;           // directional light
   mjtByte  castshadow;            // does light cast shadows
   float    bulbradius;            // bulb radius for soft shadows
+  float    intensity;             // intensity, in candelas
+  float    range;                 // range of effectiveness
 };
 typedef struct mjvLight_ mjvLight;
 struct mjvOption_ {                  // abstract visualization options
@@ -2889,287 +2981,6 @@ struct mjvFigure_ {               // abstract 2D figure passed to OpenGL rendere
   float   yaxisdata[2];           // range of y-axis in data units
 };
 typedef struct mjvFigure_ mjvFigure;
-struct mjvSceneState_ {
-  int nbuffer;                     // size of the buffer in bytes
-  void* buffer;                    // heap-allocated memory for all arrays in this struct
-  int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
-
-  // fields in mjModel that are necessary to re-render a scene
-  struct {
-    int nv;
-    int nu;
-    int na;
-    int nbody;
-    int nbvh;
-    int nbvhstatic;
-    int njnt;
-    int ngeom;
-    int nsite;
-    int ncam;
-    int nlight;
-    int nmesh;
-    int nskin;
-    int nflex;
-    int nflexvert;
-    int nflextexcoord;
-    int nskinvert;
-    int nskinface;
-    int nskinbone;
-    int nskinbonevert;
-    int nmat;
-    int neq;
-    int ntendon;
-    int ntree;
-    int nwrap;
-    int nsensor;
-    int nnames;
-    int npaths;
-    int nsensordata;
-    int narena;
-
-    mjOption opt;
-    mjVisual vis;
-    mjStatistic stat;
-
-    int* body_parentid;
-    int* body_rootid;
-    int* body_weldid;
-    int* body_mocapid;
-    int* body_jntnum;
-    int* body_jntadr;
-    int* body_dofnum;
-    int* body_dofadr;
-    int* body_geomnum;
-    int* body_geomadr;
-    mjtNum* body_iquat;
-    mjtNum* body_mass;
-    mjtNum* body_inertia;
-    int* body_bvhadr;
-    int* body_bvhnum;
-
-    int* bvh_depth;
-    int* bvh_child;
-    int* bvh_nodeid;
-    mjtNum* bvh_aabb;
-
-    int* jnt_type;
-    int* jnt_bodyid;
-    int* jnt_group;
-
-    int* geom_type;
-    int* geom_bodyid;
-    int* geom_contype;
-    int* geom_conaffinity;
-    int* geom_dataid;
-    int* geom_matid;
-    int* geom_group;
-    mjtNum* geom_size;
-    mjtNum* geom_aabb;
-    mjtNum* geom_rbound;
-    float* geom_rgba;
-
-    int* site_type;
-    int* site_bodyid;
-    int* site_matid;
-    int* site_group;
-    mjtNum* site_size;
-    float* site_rgba;
-
-    int* cam_orthographic;
-    mjtNum* cam_fovy;
-    mjtNum* cam_ipd;
-    int* cam_resolution;
-    float* cam_sensorsize;
-    float* cam_intrinsic;
-
-    mjtByte* light_directional;
-    mjtByte* light_castshadow;
-    float* light_bulbradius;
-    mjtByte* light_active;
-    float* light_attenuation;
-    float* light_cutoff;
-    float* light_exponent;
-    float* light_ambient;
-    float* light_diffuse;
-    float* light_specular;
-
-    mjtByte* flex_flatskin;
-    int* flex_dim;
-    int* flex_matid;
-    int* flex_group;
-    int* flex_interp;
-    int* flex_nodeadr;
-    int* flex_nodenum;
-    int* flex_nodebodyid;
-    int* flex_vertadr;
-    int* flex_vertnum;
-    int* flex_elem;
-    int* flex_elemtexcoord;
-    int* flex_elemlayer;
-    int* flex_elemadr;
-    int* flex_elemnum;
-    int* flex_elemdataadr;
-    int* flex_shell;
-    int* flex_shellnum;
-    int* flex_shelldataadr;
-    int* flex_texcoordadr;
-    int* flex_bvhadr;
-    int* flex_bvhnum;
-    mjtByte* flex_centered;
-    mjtNum* flex_node;
-    mjtNum* flex_radius;
-    float* flex_rgba;
-    float* flex_texcoord;
-
-    int* hfield_pathadr;
-
-    int* mesh_bvhadr;
-    int* mesh_bvhnum;
-    int* mesh_texcoordadr;
-    int* mesh_graphadr;
-    int* mesh_pathadr;
-
-    int* skin_matid;
-    int* skin_group;
-    float* skin_rgba;
-    float* skin_inflate;
-    int* skin_vertadr;
-    int* skin_vertnum;
-    int* skin_texcoordadr;
-    int* skin_faceadr;
-    int* skin_facenum;
-    int* skin_boneadr;
-    int* skin_bonenum;
-    float* skin_vert;
-    int* skin_face;
-    int* skin_bonevertadr;
-    int* skin_bonevertnum;
-    float* skin_bonebindpos;
-    float* skin_bonebindquat;
-    int* skin_bonebodyid;
-    int* skin_bonevertid;
-    float* skin_bonevertweight;
-    int* skin_pathadr;
-
-    int* tex_pathadr;
-
-    int* mat_texid;
-    mjtByte* mat_texuniform;
-    float* mat_texrepeat;
-    float* mat_emission;
-    float* mat_specular;
-    float* mat_shininess;
-    float* mat_reflectance;
-    float* mat_metallic;
-    float* mat_roughness;
-    float* mat_rgba;
-
-    int* eq_type;
-    int* eq_obj1id;
-    int* eq_obj2id;
-    int* eq_objtype;
-    mjtNum* eq_data;
-
-    int* tendon_num;
-    int* tendon_matid;
-    int* tendon_group;
-    mjtByte* tendon_limited;
-    mjtNum* tendon_width;
-    mjtNum* tendon_range;
-    mjtNum* tendon_stiffness;
-    mjtNum* tendon_damping;
-    mjtNum* tendon_frictionloss;
-    mjtNum* tendon_lengthspring;
-    float* tendon_rgba;
-
-    int* actuator_trntype;
-    int* actuator_dyntype;
-    int* actuator_trnid;
-    int* actuator_actadr;
-    int* actuator_actnum;
-    int* actuator_group;
-    mjtByte* actuator_ctrllimited;
-    mjtByte* actuator_actlimited;
-    mjtNum* actuator_ctrlrange;
-    mjtNum* actuator_actrange;
-    mjtNum* actuator_cranklength;
-
-    int* sensor_type;
-    int* sensor_objid;
-    int* sensor_adr;
-
-    int* name_bodyadr;
-    int* name_jntadr;
-    int* name_geomadr;
-    int* name_siteadr;
-    int* name_camadr;
-    int* name_lightadr;
-    int* name_eqadr;
-    int* name_tendonadr;
-    int* name_actuatoradr;
-    char* names;
-    char* paths;
-  } model;
-
-  // fields in mjData that are necessary to re-render a scene
-  struct {
-    mjWarningStat warning[mjNWARNING];
-
-    int nefc;
-    int ncon;
-    int nisland;
-
-    mjtNum time;
-
-    mjtNum* act;
-
-    mjtNum* ctrl;
-    mjtNum* xfrc_applied;
-    mjtByte* eq_active;
-
-    mjtNum* sensordata;
-
-    mjtNum* xpos;
-    mjtNum* xquat;
-    mjtNum* xmat;
-    mjtNum* xipos;
-    mjtNum* ximat;
-    mjtNum* xanchor;
-    mjtNum* xaxis;
-    mjtNum* geom_xpos;
-    mjtNum* geom_xmat;
-    mjtNum* site_xpos;
-    mjtNum* site_xmat;
-    mjtNum* cam_xpos;
-    mjtNum* cam_xmat;
-    mjtNum* light_xpos;
-    mjtNum* light_xdir;
-
-    mjtNum* subtree_com;
-
-    int* ten_wrapadr;
-    int* ten_wrapnum;
-    int* wrap_obj;
-    mjtNum* ten_length;
-    mjtNum* wrap_xpos;
-
-    mjtNum* bvh_aabb_dyn;
-    mjtByte* bvh_active;
-    int* island_dofadr;
-    int* island_dofind;
-    int* dof_island;
-    int* efc_island;
-    int* tendon_efcadr;
-
-    mjtNum* flexvert_xpos;
-
-    mjContact* contact;
-    mjtNum* efc_force;
-    void* arena;
-  } data;
-};
-typedef struct mjvSceneState_ mjvSceneState;
 
 //----------------------------- MJAPI FUNCTIONS --------------------------------
 void mj_defaultVFS(mjVFS* vfs);
@@ -3181,6 +2992,7 @@ mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int err
 mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
 mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz);
 mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
+int mj_copyBack(mjSpec* s, const mjModel* m);
 int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
 int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz);
 void mj_freeLastXML(void);
@@ -3204,6 +3016,7 @@ void mj_deleteModel(mjModel* m);
 int mj_sizeModel(const mjModel* m);
 mjData* mj_makeData(const mjModel* m);
 mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
+mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src);
 void mj_resetData(const mjModel* m, mjData* d);
 void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
 void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
@@ -3259,6 +3072,7 @@ void mj_flex(const mjModel* m, mjData* d);
 void mj_tendon(const mjModel* m, mjData* d);
 void mj_transmission(const mjModel* m, mjData* d);
 void mj_crb(const mjModel* m, mjData* d);
+void mj_makeM(const mjModel* m, mjData* d);
 void mj_factorM(const mjModel* m, mjData* d);
 void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
 void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
@@ -3357,14 +3171,8 @@ mjtNum mjv_frustumHeight(const mjvScene* scn);
 void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum forward[3]);
 void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
                     const mjvScene* scn, mjvCamera* cam);
-void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action,
-                             mjtNum reldx, mjtNum reldy,
-                             const mjvScene* scn, mjvCamera* cam);
 void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
                      mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
-void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action,
-                              mjtNum reldx, mjtNum reldy,
-                              const mjvScene* scn, mjvPerturb* pert);
 void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
                    const mjtNum roomup[3], mjvScene* scn);
 void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPerturb* pert);
@@ -3387,16 +3195,7 @@ void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom);
 void mjv_freeScene(mjvScene* scn);
 void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
                      const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
-int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt,
-                             const mjvPerturb* pert, mjvCamera* cam, int catmask,
-                             mjvScene* scn);
 void mjv_copyModel(mjModel* dest, const mjModel* src);
-void mjv_defaultSceneState(mjvSceneState* scnstate);
-void mjv_makeSceneState(const mjModel* m, const mjData* d,
-                        mjvSceneState* scnstate, int maxgeom);
-void mjv_freeSceneState(mjvSceneState* scnstate);
-void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
-                          mjvSceneState* scnstate);
 void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
                   const mjvPerturb* pert, int catmask, mjvScene* scn);
 void mjv_makeLights(const mjModel* m, const mjData* d, mjvScene* scn);
@@ -3580,6 +3379,10 @@ void mju_insertionSortInt(int* list, int n);
 mjtNum mju_Halton(int index, int base);
 char* mju_strncpy(char *dst, const char *src, int n);
 mjtNum mju_sigmoid(mjtNum x);
+const mjpPlugin* mjc_getSDF(const mjModel* m, int id);
+mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]);
+void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3],
+                  const mjtNum x[3]);
 void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
                       mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
 void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
@@ -3605,13 +3408,7 @@ void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
 void mju_threadPoolDestroy(mjThreadPool* thread_pool);
 void mju_defaultTask(mjTask* task);
 void mju_taskJoin(mjTask* task);
-mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
+mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
-                        const char* prefix, const char* suffix);
-mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
-                          const char* prefix, const char* suffix);
-mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
-                          const char* prefix, const char* suffix);
-mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child,
                        const char* prefix, const char* suffix);
 int mjs_detachBody(mjSpec* s, mjsBody* b);
 int mjs_detachDefault(mjSpec* s, mjsDefault* d);
@@ -3675,12 +3472,16 @@ void mjs_setDouble(mjDoubleVec* dest, const double* array, int size);
 void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes);
 const char* mjs_getString(const mjString* source);
 const double* mjs_getDouble(const mjDoubleVec* source, int* size);
+const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
 void mjs_setDefault(mjsElement* element, const mjsDefault* def);
-void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
+int mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
                                    const mjsOrientation* orientation);
 mjsFrame* mjs_bodyToFrame(mjsBody** body);
 void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
+void mjs_setUserValueWithCleanup(mjsElement* element, const char* key,
+                                 const void* data,
+                                 void (*cleanup)(const void*));
 const void* mjs_getUserValue(mjsElement* element, const char* key);
 void mjs_deleteUserValue(mjsElement* element, const char* key);
 void mjs_defaultSpec(mjSpec* spec);

doc/mjx.rst

@@ -235,6 +235,8 @@ The following features are **fully supported** in MJX:
      - 1, 3, 4, 6 (1 is not supported with ``ELLIPTIC``)
    * - :ref:`Solver <mjtSolver>`
      - ``CG``, ``NEWTON``
+   * - Dynamics
+     - :ref:`Inverse <mj_inverse>`
    * - Fluid Model
      - :ref:`flInertia`
    * - :ref:`Tendons <tendon>`
@@ -244,7 +246,7 @@ The following features are **fully supported** in MJX:
        ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``,
        ``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``TENDONVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``,
        ``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``,
-       ``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``
+       ``ACTUATORFRC``, ``JOINTACTFRC``, ``TENDONACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``
        (``ACCELEROMETER``, ``FORCE``, ``TORQUE`` not supported with connect or weld equality constraints)
 
 The following features are **in development** and coming soon:
@@ -262,8 +264,6 @@ The following features are **in development** and coming soon:
        (``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``.
    * - :ref:`Integrator <mjtIntegrator>`
      - ``IMPLICIT``
-   * - Dynamics
-     - :ref:`Inverse <mj_inverse>`
    * - Fluid Model
      - :ref:`flEllipsoid`
    * - :ref:`Sensors <mjtSensor>`

doc/modeling.rst

@@ -724,7 +724,10 @@ Force clamping at joint input with :ref:`joint/actuatorfrcrange<body-joint-actua
   :ref:`jointactuatorfrc<sensor-jointactuatorfrc>` sensor to report the total actuator force acting on a joint.
   The standard :ref:`actuatorfrc<sensor-actuatorfrc>` sensor will continue to report the pre-clamped actuator force.
 
-The three clamping options above are non-exclusive and can be combined as required.
+Force clamping at tendon input with :ref:`tendon/actuatorfrcrange<tendon-spatial-actuatorfrcrange>`:
+  This tendon attribute clamps input forces from all actuators acting on the tendon.
+
+The clamping options above are non-exclusive and can be combined as required.
 
 .. _CLengthRange:
 

doc/programming/modeledit.rst

@@ -111,9 +111,10 @@ This framework introduces a powerful new feature: attaching and detaching model
 to power the :ref:`attach<body-attach>` an :ref:`replicate<replicate>` meta-elements in MJCF. Attachment allows the user
 to move or copy a subtree from one model into another, while also copying or moving related referenced assets and
 referencing elements from outside the kinematic tree (e.g., actuators and sensors). Similarly, detaching a subtree will
-remove all associated elements from the model. The default behavior is to move during attach. The user can select to
+remove all associated elements from the model. The default behavior is to move the child into the parent while
-instead copy by passing the corresponding flag to ``mjs_setDeepCopy``. This flag is temporary set to true while parsing
+attaching, so subsequent changes to the child will also change the parent. Alternatively, the user can choose to make an
-XMLs. It is possible to :ref:`attach a body to a frame<mjs_attachBody>`:
+entirely new copy during attach using :ref:`mjs_setDeepCopy`. This flag is temporarily set to true while parsing XMLs.
+It is possible to :ref:`attach a body to a frame<mjs_attach>`:
 
 .. code-block:: C
 
@@ -121,29 +122,29 @@ XMLs. It is possible to :ref:`attach a body to a frame<mjs_attachBody>`:
    mjSpec* child = mj_makeSpec();
    parent->compiler.degree = 0;
    child->compiler.degree = 1;
-   mjsFrame* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL);
+   mjsElement* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL)->element;
-   mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL);
+   mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element;
-   mjsBody* attached_body_1 = mjs_attachBody(frame, body, "attached-", "-1");
+   mjsBody* attached_body_1 = mjs_asBody(mjs_attach(frame, body, "attached-", "-1"));
 
-or :ref:`attach a body to a site<mjs_attachToSite>`:
+or :ref:`attach a body to a site<mjs_attach>`:
 
 .. code-block:: C
 
    mjSpec* parent = mj_makeSpec();
    mjSpec* child = mj_makeSpec();
-   mjsSite* site = mjs_addSite(mjs_findBody(parent, "world"), NULL);
+   mjsElement* site = mjs_addSite(mjs_findBody(parent, "world"), NULL)->element;
-   mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL);
+   mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element;
-   mjsBody* attached_body_2 = mjs_attachToSite(site, body, "attached-", "-2");
+   mjsBody* attached_body_2 = mjs_asBody(mjs_attach(site, body, "attached-", "-2"));
 
-or :ref:`attach a frame to a body<mjs_attachFrame>`:
+or :ref:`attach a frame to a body<mjs_attach>`:
 
 .. code-block:: C
 
    mjSpec* parent = mj_makeSpec();
    mjSpec* child = mj_makeSpec();
-   mjsBody* body = mjs_addBody(mjs_findBody(parent, "world"), NULL);
+   mjsElement* body = mjs_addBody(mjs_findBody(parent, "world"), NULL)->element;
-   mjsFrame* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL);
+   mjsElement* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL)->element;
-   mjsFrame* attached_frame = mjs_attachFrame(body, frame, "attached-", "-1");
+   mjsFrame* attached_frame = mjs_asFrame(mjs_attach(body, frame, "attached-", "-1"));
 
 Note that in the above examples, the parent and child models have different values for ``compiler.degree``,
 corresponding to the :ref:`compiler/angle<compiler-angle>` attribute, specifying the units in which angles are

doc/programming/simulation.rst

@@ -807,7 +807,7 @@ to implement high-resolution timers in C without bringing in additional dependen
 does not need timing, and in that case there is no reason to call timing functions.
 
 One part of the simulation pipeline that needs to be monitored closely is the iterative constraint solver. The
-simplest diagnostic here is ``mjData.solver_iter`` which shows how many iterations the solver took on the last call to
+simplest diagnostic here is ``mjData.solver_niter`` which shows how many iterations the solver took on the last call to
 mj_step or ``mj_forward``. Note that the solver has tolerance parameters for early termination, so this number is
 usually smaller than the maximum number of iterations allowed. The array ``mjData.solver`` contains one
 :ref:`mjSolverStat` data structure per iteration of the constraint solver, with information about the constraint state
@@ -913,7 +913,7 @@ elliptic, depending on which solver is selected in ``mjModel.opt``. The function
 can be used to determine which friction cone type is used. For pyramidal cones, the interpretation of the contact force
 (whose address we computed above) is non-trivial, because the components are forces along redundant non-orthogonal axes
 corresponding to the edges of the pyramid. The function :ref:`mj_contactForce` can be
-used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D toque.
+used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D torque.
 The torque component will be zero when :at:`condim` is 1 or 3, and non-zero otherwise. This force and torque are
 expressed in the contact frame given by mjContact.frame. Unlike all other matrices in mjData, this matrix is stored in
 transposed form. Normally a 3-by-3 matrix corresponding to a coordinate frame would have the frame axes along the

doc/python.rst

@@ -601,6 +601,14 @@ Lists of all elements in a spec can be accessed using named properties, using th
 ``equalities``, ``tendons``, ``actuators``, ``skins``, ``textures``, ``texts``, ``tuples``, ``flexes``, ``hfields``,
 ``keys``, ``numerics``, ``excludes``, ``sensors``, ``plugins``.
 
+Element removal
+^^^^^^^^^^^^^^^
+For elements that can have children (bodies and defaults), the methods ``spec.detach_body(body)`` and
+``spec.detach_default(def)`` remove, respectively, ``body`` and ``def`` from the spec, together with all of their
+children. When detaching body subtrees, all elements which reference elements in the subtree, will also be removed. For
+all other elements, the method ``delete()`` removes the corresponding element from the spec, e.g.
+``spec.geom('my_geom').delete()`` will remove the geom named "my_geom" and all of the elements that reference it.
+
 Tree traversal
 ^^^^^^^^^^^^^^
 Traversal of the kinematic tree is aided by the following methods which return tree-related lists of elements:
@@ -620,8 +628,16 @@ Parent:
   The parent body of a given element -- including bodies and frames -- can be accessed via the ``parent`` property.
   For example, the parent of a site can be accessed via ``site.parent``.
 
-Relationship to ``PyMJCF``
+Serialization
---------------------------
+^^^^^^^^^^^^^
+The ``MjSpec`` object can be serialized with all of its assets using the function ``spec.to_zip(file)``, where ``file``
+can be either a path to a file or a file object. In order to load the spec from a zip file, use ``spec =
+MjSpec.from_zip(file)``, where ``file`` is a path to a zip file or a zip file object.
+
+.. _PyMJCF:
+
+Relationship to ``PyMJCF`` and ``bind``
+---------------------------------------
 
 `dm_control <https://github.com/google-deepmind/dm_control/tree/main>`__'s
 `PyMJCF <https://github.com/google-deepmind/dm_control/blob/main/dm_control/mjcf/README.md>`__ module provides similar
@@ -637,9 +653,8 @@ includes a reimplementation of the ``PyMJCF`` example in the ``dm_control``
 
 ``PyMJCF`` provides a notion of "binding", giving access to :ref:`mjModel` and :ref:`mjData` values via a helper class.
 In the native API, the helper class is not needed, so it is possible to directly bind an ``mjs`` object to
-:ref:`mjModel` and :ref:`mjData`. This requires the objects to have a non-empty name. For example, say we have multiple
+:ref:`mjModel` and :ref:`mjData`. For example, say we have multiple geoms containing the string "torso" in their name.
-geoms containing the string "torso" in their name. We want to get their Cartesian positions in the XY plane from
+We want to get their Cartesian positions in the XY plane from ``mjData``. This can be done as follows:
-``mjData``. This can be done as follows:
 
 .. code-block:: python
 
@@ -647,6 +662,9 @@ geoms containing the string "torso" in their name. We want to get their Cartesia
    pos_x = [torso.xpos[0] for torso in torsos]
    pos_y = [torso.xpos[1] for torso in torsos]
 
+Using the ``bind`` method requires the :ref:`mjModel` and :ref:`mjData` to be compiled from the :`ref:`mjSpec`. If
+objects are added or removed from the :ref:`mjSpec` since the last compilation, an error is raised.
+
 Notes
 -----
 

doc/requirements.txt

@@ -1,6 +1,6 @@
 Sphinx==5.3.0
 furo==2022.9.29
-sphinxcontrib-bibtex==2.6.1
+sphinxcontrib-bibtex==2.6.3
 sphinxcontrib-katex==0.9.4
 sphinxcontrib-youtube==1.2.0
 sphinx-copybutton==0.5.2
@@ -10,7 +10,7 @@ sphinx-toolbox==3.8.2
 nbsphinx==0.9.1
 pandoc==1.1.0
 pygments==2.15.0
-jq==1.4.1
+jq==1.8.0
 Jinja2~=3.0
 wheel
 # see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe

doc/unity.rst

@@ -37,14 +37,14 @@ _____
 
 The MuJoCo app needs to be run at least once before the native library can be used, in order to register the library as
 a trusted binary. Then, copy the dynamic library file from
-``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.3.1.dylib`` (it can be
+``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.3.3.dylib`` (it can be
 found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``.
 
 Linux
 _____
 
 Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from
-``~/.mujoco/mujoco-3.3.1/lib/libmujoco.so.3.3.1`` and rename it as ``libmujoco.so``.
+``~/.mujoco/mujoco-3.3.3/lib/libmujoco.so.3.3.3`` and rename it as ``libmujoco.so``.
 
 Windows
 _______

include/mujoco/mjdata.h

@@ -182,7 +182,6 @@ struct mjData_ {
 
   // solver statistics
   mjSolverStat  solver[mjNISLAND*mjNSOLVER];  // solver statistics per island, per iteration
-  int           solver_nisland;               // number of islands processed by solver
   int           solver_niter[mjNISLAND];      // number of solver iterations, per island
   int           solver_nnz[mjNISLAND];        // number of nonzeros in Hessian or efc_AR, per island
   mjtNum        solver_fwdinv[2];             // forward-inverse comparison: qfrc, efc
@@ -200,6 +199,7 @@ struct mjData_ {
   int     nJ;                // number of non-zeros in constraint Jacobian
   int     nA;                // number of non-zeros in constraint inverse inertia matrix
   int     nisland;           // number of detected constraint islands
+  int     nidof;             // number of dofs in all islands
 
   // global properties
   mjtNum  time;              // simulation time
@@ -295,12 +295,13 @@ struct mjData_ {
   int*    moment_colind;     // column indices in sparse Jacobian                (nJmom x 1)
   mjtNum* actuator_moment;   // actuator moments                                 (nJmom x 1)
 
-  // computed by mj_fwdPosition/mj_crb
+  // computed by mj_fwdPosition/mj_makeM
   mjtNum* crb;               // com-based composite inertia and mass             (nbody x 10)
-  mjtNum* qM;                // total inertia (sparse)                           (nM x 1)
+  mjtNum* qM;                // inertia (sparse)                                 (nM x 1)
+  mjtNum* M;                 // reduced inertia (compressed sparse row)          (nC x 1)
 
   // computed by mj_fwdPosition/mj_factorM
-  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nM x 1)
+  mjtNum* qLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collisionTree
@@ -333,27 +334,23 @@ struct mjData_ {
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nC x 1)
   mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
   int*    B_rownnz;          // body-dof: non-zeros in each row                  (nbody x 1)
   int*    B_rowadr;          // body-dof: address of each row in B_colind        (nbody x 1)
   int*    B_colind;          // body-dof: column indices of non-zeros            (nB x 1)
-  int*    M_rownnz;          // inertia: non-zeros in each row                   (nv x 1)
+  int*    M_rownnz;          // reduced inertia: non-zeros in each row           (nv x 1)
-  int*    M_rowadr;          // inertia: address of each row in M_colind         (nv x 1)
+  int*    M_rowadr;          // reduced inertia: address of each row in M_colind (nv x 1)
-  int*    M_colind;          // inertia: column indices of non-zeros             (nM x 1)
+  int*    M_colind;          // reduced inertia: column indices of non-zeros     (nC x 1)
-  int*    mapM2M;            // index mapping from M (legacy) to M (CSR)         (nM x 1)
+  int*    mapM2M;            // index mapping from qM to M                       (nC x 1)
-  int*    C_rownnz;          // reduced dof-dof: non-zeros in each row           (nv x 1)
+  int*    D_rownnz;          // full inertia: non-zeros in each row              (nv x 1)
-  int*    C_rowadr;          // reduced dof-dof: address of each row in C_colind (nv x 1)
+  int*    D_rowadr;          // full inertia: address of each row in D_colind    (nv x 1)
-  int*    C_colind;          // reduced dof-dof: column indices of non-zeros     (nC x 1)
+  int*    D_diag;            // full inertia: index of diagonal element          (nv x 1)
-  int*    mapM2C;            // index mapping from M to C                        (nC x 1)
+  int*    D_colind;          // full inertia: column indices of non-zeros        (nD x 1)
-  int*    D_rownnz;          // dof-dof: non-zeros in each row                   (nv x 1)
+  int*    mapM2D;            // index mapping from qM to D                       (nD x 1)
-  int*    D_rowadr;          // dof-dof: address of each row in D_colind         (nv x 1)
+  int*    mapD2M;            // index mapping from D to qM                       (nM x 1)
-  int*    D_diag;            // dof-dof: index of diagonal element               (nv x 1)
-  int*    D_colind;          // dof-dof: column indices of non-zeros             (nD x 1)
-  int*    mapM2D;            // index mapping from M to D                        (nD x 1)
-  int*    mapD2M;            // index mapping from D to M                        (nM x 1)
 
   // computed by mj_implicit/mj_derivative
   mjtNum* qDeriv;            // d (passive + actuator - bias) / d qvel           (nD x 1)
@@ -375,8 +372,8 @@ struct mjData_ {
   mjtNum* qfrc_constraint;   // constraint force                                 (nv x 1)
 
   // computed by mj_inverse
-  mjtNum* qfrc_inverse;      // net external force; should equal:                (nv x 1)
+  mjtNum* qfrc_inverse;      // net external force; should equal:
-                             // qfrc_applied + J'*xfrc_applied + qfrc_actuator
+                             // qfrc_applied + J'*xfrc_applied + qfrc_actuator   (nv x 1)
 
   // computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format
   mjtNum* cacc;              // com-based acceleration                           (nbody x 6)
@@ -410,16 +407,50 @@ struct mjData_ {
   mjtNum* efc_R;             // inverse constraint mass                          (nefc x 1)
   int*    tendon_efcadr;     // first efc address involving tendon; -1: none     (ntendon x 1)
 
-  // computed by mj_island
+  // computed by mj_island (island dof structure)
   int*    dof_island;        // island id of this dof; -1: none                  (nv x 1)
-  int*    island_dofnum;     // number of dofs in island                         (nisland x 1)
+  int*    island_nv;         // number of dofs in this island                    (nisland x 1)
-  int*    island_dofadr;     // start address in island_dofind                   (nisland x 1)
+  int*    island_idofadr;    // island start address in idof vector              (nisland x 1)
-  int*    island_dofind;     // island dof indices; -1: none                     (nv x 1)
+  int*    island_dofadr;     // island start address in dof vector               (nisland x 1)
-  int*    dof_islandind;     // dof island indices; -1: none                     (nv x 1)
+  int*    map_dof2idof;      // map from dof to idof                             (nv x 1)
+  int*    map_idof2dof;      // map from idof to dof;  >= nidof: unconstrained   (nv x 1)
+
+  // computed by mj_island (dofs sorted by island)
+  mjtNum* ifrc_smooth;       // net unconstrained force                          (nidof x 1)
+  mjtNum* iacc_smooth;       // unconstrained acceleration                       (nidof x 1)
+  int*    iM_rownnz;         // inertia: non-zeros in each row                   (nidof x 1)
+  int*    iM_rowadr;         // inertia: address of each row in iM_colind        (nidof x 1)
+  int*    iM_colind;         // inertia: column indices of non-zeros             (nC x 1)
+  mjtNum* iM;                // total inertia (sparse)                           (nC x 1)
+  mjtNum* iLD;               // L'*D*L factorization of M (sparse)               (nC x 1)
+  mjtNum* iLDiagInv;         // 1/diag(D)                                        (nidof x 1)
+  mjtNum* iacc;              // acceleration                                     (nidof x 1)
+
+  // computed by mj_island (island constraint structure)
   int*    efc_island;        // island id of this constraint                     (nefc x 1)
-  int*    island_efcnum;     // number of constraints in island                  (nisland x 1)
+  int*    island_ne;         // number of equality constraints in island         (nisland x 1)
-  int*    island_efcadr;     // start address in island_efcind                   (nisland x 1)
+  int*    island_nf;         // number of friction constraints in island         (nisland x 1)
-  int*    island_efcind;     // island constraint indices                        (nefc x 1)
+  int*    island_nefc;       // number of constraints in island                  (nisland x 1)
+  int*    island_iefcadr;    // start address in iefc vector                     (nisland x 1)
+  int*    map_efc2iefc;      // map from efc to iefc                             (nefc x 1)
+  int*    map_iefc2efc;      // map from iefc to efc                             (nefc x 1)
+
+  // computed by mj_island (constraints sorted by island)
+  int*    iefc_type;         // constraint type (mjtConstraint)                  (nefc x 1)
+  int*    iefc_id;           // id of object of specified type                   (nefc x 1)
+  int*    iefc_J_rownnz;     // number of non-zeros in constraint Jacobian row   (nefc x 1)
+  int*    iefc_J_rowadr;     // row start address in colind array                (nefc x 1)
+  int*    iefc_J_rowsuper;   // number of subsequent rows in supernode           (nefc x 1)
+  int*    iefc_J_colind;     // column indices in constraint Jacobian            (nJ x 1)
+  int*    iefc_JT_rownnz;    // number of non-zeros in constraint Jacobian row T (nidof x 1)
+  int*    iefc_JT_rowadr;    // row start address in colind array              T (nidof x 1)
+  int*    iefc_JT_rowsuper;  // number of subsequent rows in supernode         T (nidof x 1)
+  int*    iefc_JT_colind;    // column indices in constraint Jacobian          T (nJ x 1)
+  mjtNum* iefc_J;            // constraint Jacobian                              (nJ x 1)
+  mjtNum* iefc_JT;           // constraint Jacobian transposed                   (nJ x 1)
+  mjtNum* iefc_frictionloss; // frictionloss (friction)                          (nefc x 1)
+  mjtNum* iefc_D;            // constraint mass                                  (nefc x 1)
+  mjtNum* iefc_R;            // inverse constraint mass                          (nefc x 1)
 
   // computed by mj_projectConstraint (PGS solver)
   int*    efc_AR_rownnz;     // number of non-zeros in AR                        (nefc x 1)
@@ -437,11 +468,18 @@ struct mjData_ {
 
   // computed by mj_fwdConstraint/mj_inverse
   mjtNum* efc_b;             // linear cost term: J*qacc_smooth - aref           (nefc x 1)
-  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
+  mjtNum* iefc_aref;         // reference pseudo-acceleration                    (nefc x 1)
+  int*    iefc_state;        // constraint state (mjtConstraintState)            (nefc x 1)
+  mjtNum* iefc_force;        // constraint force in constraint space             (nefc x 1)
   int*    efc_state;         // constraint state (mjtConstraintState)            (nefc x 1)
+  mjtNum* efc_force;         // constraint force in constraint space             (nefc x 1)
+  mjtNum* ifrc_constraint;   // constraint force                                 (nidof x 1)
 
   // thread pool pointer
   uintptr_t threadpool;
+
+  // compilation signature
+  uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
 typedef struct mjData_ mjData;
 

include/mujoco/mjmodel.h

@@ -129,6 +129,14 @@ typedef enum mjtCamLight_ {       // tracking mode for camera and light
 } mjtCamLight;
 
 
+typedef enum mjtLightType_ {      // type of light
+  mjLIGHT_SPOT    = 0,            // spot
+  mjLIGHT_DIRECTIONAL,            // directional
+  mjLIGHT_POINT,                  // point
+  mjLIGHT_IMAGE,                  // image-based
+} mjtLightType;
+
+
 typedef enum mjtTexture_ {        // type of texture
   mjTEXTURE_2D        = 0,        // 2d texture, suitable for planes and hfields
   mjTEXTURE_CUBE,                 // cube texture, suitable for all other geom types
@@ -151,6 +159,13 @@ typedef enum mjtTextureRole_ {    // role of texture map in rendering
 } mjtTextureRole;
 
 
+typedef enum mjtColorSpace_ {     // type of color space encoding
+  mjCOLORSPACE_AUTO   = 0,        // attempts to autodetect color space, defaults to linear
+  mjCOLORSPACE_LINEAR,            // linear color space
+  mjCOLORSPACE_SRGB               // standard RGB color space
+} mjtColorSpace;
+
+
 typedef enum mjtIntegrator_ {     // integrator mode
   mjINT_EULER         = 0,        // semi-implicit Euler
   mjINT_RK4,                      // 4th-order Runge Kutta
@@ -269,7 +284,8 @@ typedef enum mjtObj_ {            // type of MujoCo object
 
   // meta elements, do not appear in mjModel
   mjOBJ_FRAME         = 100,      // frame
-  mjOBJ_DEFAULT                   // default
+  mjOBJ_DEFAULT,                  // default
+  mjOBJ_MODEL                     // entire model
 
 } mjtObj;
 
@@ -316,6 +332,7 @@ typedef enum mjtSensor_ {         // type of sensor
   mjSENS_ACTUATORVEL,             // scalar actuator velocity
   mjSENS_ACTUATORFRC,             // scalar actuator force
   mjSENS_JOINTACTFRC,             // scalar actuator force, measured at the joint
+  mjSENS_TENDONACTFRC,            // scalar actuator force, measured at the tendon
 
   // sensors related to ball joints
   mjSENS_BALLQUAT,                // 4D ball joint quaternion
@@ -405,6 +422,14 @@ typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
 } mjtFlexSelf;
 
 
+typedef enum mjtSDFType_ {        // signed distance function (SDF) type
+  mjSDFTYPE_SINGLE     = 0,       // single SDF
+  mjSDFTYPE_INTERSECTION,         // max(A, B)
+  mjSDFTYPE_MIDSURFACE,           // A - B
+  mjSDFTYPE_COLLISION,            // A + B + abs(max(A, B))
+} mjtSDFType;
+
+
 //---------------------------------- mjLROpt -------------------------------------------------------
 
 struct mjLROpt_ {                 // options for mj_setLengthRange()
@@ -832,9 +857,12 @@ struct mjModel_ {
   int*      light_mode;           // light tracking mode (mjtCamLight)        (nlight x 1)
   int*      light_bodyid;         // id of light's body                       (nlight x 1)
   int*      light_targetbodyid;   // id of targeted body; -1: none            (nlight x 1)
-  mjtByte*  light_directional;    // directional light                        (nlight x 1)
+  int*      light_type;           // spot, directional, etc. (mjtLightType)   (nlight x 1)
+  int*      light_texid;          // texture id for image lights              (nlight x 1)
   mjtByte*  light_castshadow;     // does light cast shadows                  (nlight x 1)
   float*    light_bulbradius;     // light radius for soft shadows            (nlight x 1)
+  float*    light_intensity;      // intensity, in candela                    (nlight x 1)
+  float*    light_range;          // range of effectiveness                   (nlight x 1)
   mjtByte*  light_active;         // is light on                              (nlight x 1)
   mjtNum*   light_pos;            // position rel. to body frame              (nlight x 3)
   mjtNum*   light_dir;            // direction rel. to body frame             (nlight x 3)
@@ -886,6 +914,7 @@ struct mjModel_ {
   int*      flex_nodebodyid;      // node body ids                            (nflexnode x 1)
   int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
   int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
+  int*      flex_edgeflap;        // adjacent vertex ids (dim=2 only)         (nflexedge x 2)
   int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
   int*      flex_elemtexcoord;    // element texture coordinates (dim+1)      (nflexelemdata x 1)
   int*      flex_elemedge;        // element edge ids                         (nflexelemedge x 1)
@@ -900,6 +929,7 @@ struct mjModel_ {
   mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
   mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
   mjtNum*   flex_stiffness;       // finite element stiffness matrix          (nflexelem x 21)
+  mjtNum*   flex_bending;         // bending stiffness                        (nflexedge x 16)
   mjtNum*   flex_damping;         // Rayleigh's damping coefficient           (nflex x 1)
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
@@ -980,6 +1010,7 @@ struct mjModel_ {
 
   // textures
   int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
+  int*      tex_colorspace;       // texture colorspace (mjtColorSpace)       (ntex x 1)
   int*      tex_height;           // number of rows in texture image          (ntex x 1)
   int*      tex_width;            // number of columns in texture image       (ntex x 1)
   int*      tex_nchannel;         // number of channels in texture image      (ntex x 1)
@@ -1030,15 +1061,18 @@ struct mjModel_ {
   int*      tendon_matid;         // material id for rendering                (ntendon x 1)
   int*      tendon_group;         // group for visibility                     (ntendon x 1)
   mjtByte*  tendon_limited;       // does tendon have length limits           (ntendon x 1)
+  mjtByte*  tendon_actfrclimited; // does tendon have actuator force limits   (ntendon x 1)
   mjtNum*   tendon_width;         // width for rendering                      (ntendon x 1)
   mjtNum*   tendon_solref_lim;    // constraint solver reference: limit       (ntendon x mjNREF)
   mjtNum*   tendon_solimp_lim;    // constraint solver impedance: limit       (ntendon x mjNIMP)
   mjtNum*   tendon_solref_fri;    // constraint solver reference: friction    (ntendon x mjNREF)
   mjtNum*   tendon_solimp_fri;    // constraint solver impedance: friction    (ntendon x mjNIMP)
   mjtNum*   tendon_range;         // tendon length limits                     (ntendon x 2)
+  mjtNum*   tendon_actfrcrange;   // range of total actuator force            (ntendon x 2)
   mjtNum*   tendon_margin;        // min distance for limit detection         (ntendon x 1)
   mjtNum*   tendon_stiffness;     // stiffness coefficient                    (ntendon x 1)
   mjtNum*   tendon_damping;       // damping coefficient                      (ntendon x 1)
+  mjtNum*   tendon_armature;      // inertia associated with tendon velocity  (ntendon x 1)
   mjtNum*   tendon_frictionloss;  // loss due to friction                     (ntendon x 1)
   mjtNum*   tendon_lengthspring;  // spring resting length range              (ntendon x 2)
   mjtNum*   tendon_length0;       // tendon length in qpos0                   (ntendon x 1)
@@ -1155,6 +1189,9 @@ struct mjModel_ {
 
   // paths
   char*     paths;                // paths to assets, 0-terminated            (npaths x 1)
+
+  // compilation signature
+  uint64_t  signature;            // also held by the mjSpec that compiled this model
 };
 typedef struct mjModel_ mjModel;
 

include/mujoco/mjplugin.h

@@ -135,6 +135,16 @@ struct mjpPlugin_ {
 };
 typedef struct mjpPlugin_ mjpPlugin;
 
+struct mjSDF_ {
+  const mjpPlugin** plugin;
+  int* id;
+  mjtSDFType type;
+  mjtNum* relpos;
+  mjtNum* relmat;
+  mjtGeom* geomtype;
+};
+typedef struct mjSDF_ mjSDF;
+
 #if defined(__has_attribute)
 
   #if __has_attribute(constructor)

include/mujoco/mjspec.h

@@ -23,6 +23,7 @@
 // this is a C-API
 #ifdef __cplusplus
 #include <cstddef>
+#include <cstdint>
 #include <string>
 #include <vector>
 
@@ -119,6 +120,7 @@ typedef enum mjtOrientation_ {     // type of orientation specifier
 
 typedef struct mjsElement_ {       // element type, do not modify
   mjtObj elemtype;                 // element type
+  uint64_t signature;              // compilation signature
 } mjsElement;
 
 
@@ -136,6 +138,7 @@ typedef struct mjsCompiler_ {      // compiler options
   mjtByte fusestatic;              // fuse static bodies with parent
   int inertiafromgeom;             // use geom inertias (mjtInertiaFromGeom)
   int inertiagrouprange[2];        // range of geom groups used to compute inertia
+  mjtByte saveinertial;            // save explicit inertial clause for all bodies to XML
   int alignfree;                   // align free joints with inertial frame
   mjLROpt LRopt;                   // options for lengthrange computation
 } mjsCompiler;
@@ -391,9 +394,12 @@ typedef struct mjsLight_ {         // light specification
 
   // intrinsics
   mjtByte active;                  // is light active
-  mjtByte directional;             // is light directional or spot
+  mjtLightType type;               // type of light
+  mjString* texture;               // texture name for image lights
   mjtByte castshadow;              // does light cast shadows
-  double bulbradius;               // bulb radius, for soft shadows
+  float bulbradius;                // bulb radius, for soft shadows
+  float intensity;                 // intensity, in candelas
+  float range;                     // range of effectiveness
   float attenuation[3];            // OpenGL attenuation (quadratic model)
   float cutoff;                    // OpenGL cutoff
   float exponent;                  // OpenGL exponent
@@ -427,7 +433,8 @@ typedef struct mjsFlex_ {          // flex specification
   double radius;                   // radius around primitive element
   mjtByte internal;                // enable internal collisions
   mjtByte flatskin;                // render flex skin with flat shading
-  int selfcollide;                 // mode for flex self colllision
+  int selfcollide;                 // mode for flex self collision
+  int vertcollide;                 // mode for vertex collision
   int activelayers;                // number of active element layers in 3D
   int group;                       // group for visualizatioh
   double edgestiffness;            // edge stiffness
@@ -438,6 +445,7 @@ typedef struct mjsFlex_ {          // flex specification
   double poisson;                  // Poisson's ratio
   double damping;                  // Rayleigh's damping
   double thickness;                // thickness (2D only)
+  int elastic2d;                   // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
 
   // mesh properties
   mjStringVec* nodebody;           // node body names
@@ -518,6 +526,7 @@ typedef struct mjsTexture_ {       // texture specification
   mjsElement* element;             // element type
   mjString* name;                  // name
   mjtTexture type;                 // texture type
+  mjtColorSpace colorspace;        // colorspace
 
   // method 1: builtin
   int builtin;                     // builtin type (mjtBuiltin)
@@ -614,17 +623,20 @@ typedef struct mjsTendon_ {        // tendon specification
   mjsElement* element;             // element type
   mjString* name;                  // name
 
-  // stiffness, damping, friction
+  // stiffness, damping, friction, armature
   double stiffness;                // stiffness coefficient
   double springlength[2];          // spring resting length; {-1, -1}: use qpos_spring
   double damping;                  // damping coefficient
   double frictionloss;             // friction loss
   mjtNum solref_friction[mjNREF];  // solver reference: tendon friction
   mjtNum solimp_friction[mjNIMP];  // solver impedance: tendon friction
+  double armature;                 // inertia associated with tendon velocity
 
   // length range
   int limited;                     // does tendon have limits (mjtLimited)
+  int actfrclimited;               // does tendon have actuator force limits
   double range[2];                 // length limits
+  double actfrcrange[2];           // actuator force limits
   double margin;                   // margin value for tendon limit detection
   mjtNum solref_limit[mjNREF];     // solver reference: tendon limits
   mjtNum solimp_limit[mjNIMP];     // solver impedance: tendon limits

include/mujoco/mjvisualize.h

@@ -226,7 +226,7 @@ typedef struct mjvGLCamera_ mjvGLCamera;
 struct mjvGeom_ {                 // abstract geom
   // type info
   int      type;                  // geom type (mjtGeom)
-  int      dataid;                // mesh, hfield or plane id; -1: none
+  int      dataid;                // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull)
   int      objtype;               // mujoco object type; mjOBJ_UNKNOWN for decor
   int      objid;                 // mujoco object id; -1 for decor
   int      category;              // visual category
@@ -261,6 +261,8 @@ typedef struct mjvGeom_ mjvGeom;
 struct mjvLight_ {                // OpenGL light
   float    pos[3];                // position rel. to body frame
   float    dir[3];                // direction rel. to body frame
+  int      type;                  // type (mjtLightType)
+  int      texid;                 // texture id for image lights
   float    attenuation[3];        // OpenGL attenuation (quadratic model)
   float    cutoff;                // OpenGL cutoff
   float    exponent;              // OpenGL exponent
@@ -268,9 +270,10 @@ struct mjvLight_ {                // OpenGL light
   float    diffuse[3];            // diffuse rgb (alpha=1)
   float    specular[3];           // specular rgb (alpha=1)
   mjtByte  headlight;             // headlight
-  mjtByte  directional;           // directional light
   mjtByte  castshadow;            // does light cast shadows
   float    bulbradius;            // bulb radius for soft shadows
+  float    intensity;             // intensity, in candelas
+  float    range;                 // range of effectiveness
 };
 typedef struct mjvLight_ mjvLight;
 
@@ -404,289 +407,4 @@ struct mjvFigure_ {               // abstract 2D figure passed to OpenGL rendere
 };
 typedef struct mjvFigure_ mjvFigure;
 
-
-//---------------------------------- mjvSceneState -------------------------------------------------
-
-struct mjvSceneState_ {
-  int nbuffer;                     // size of the buffer in bytes
-  void* buffer;                    // heap-allocated memory for all arrays in this struct
-  int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
-
-  // fields in mjModel that are necessary to re-render a scene
-  struct {
-    int nv;
-    int nu;
-    int na;
-    int nbody;
-    int nbvh;
-    int nbvhstatic;
-    int njnt;
-    int ngeom;
-    int nsite;
-    int ncam;
-    int nlight;
-    int nmesh;
-    int nskin;
-    int nflex;
-    int nflexvert;
-    int nflextexcoord;
-    int nskinvert;
-    int nskinface;
-    int nskinbone;
-    int nskinbonevert;
-    int nmat;
-    int neq;
-    int ntendon;
-    int ntree;
-    int nwrap;
-    int nsensor;
-    int nnames;
-    int npaths;
-    int nsensordata;
-    int narena;
-
-    mjOption opt;
-    mjVisual vis;
-    mjStatistic stat;
-
-    int* body_parentid;
-    int* body_rootid;
-    int* body_weldid;
-    int* body_mocapid;
-    int* body_jntnum;
-    int* body_jntadr;
-    int* body_dofnum;
-    int* body_dofadr;
-    int* body_geomnum;
-    int* body_geomadr;
-    mjtNum* body_iquat;
-    mjtNum* body_mass;
-    mjtNum* body_inertia;
-    int* body_bvhadr;
-    int* body_bvhnum;
-
-    int* bvh_depth;
-    int* bvh_child;
-    int* bvh_nodeid;
-    mjtNum* bvh_aabb;
-
-    int* jnt_type;
-    int* jnt_bodyid;
-    int* jnt_group;
-
-    int* geom_type;
-    int* geom_bodyid;
-    int* geom_contype;
-    int* geom_conaffinity;
-    int* geom_dataid;
-    int* geom_matid;
-    int* geom_group;
-    mjtNum* geom_size;
-    mjtNum* geom_aabb;
-    mjtNum* geom_rbound;
-    float* geom_rgba;
-
-    int* site_type;
-    int* site_bodyid;
-    int* site_matid;
-    int* site_group;
-    mjtNum* site_size;
-    float* site_rgba;
-
-    int* cam_orthographic;
-    mjtNum* cam_fovy;
-    mjtNum* cam_ipd;
-    int* cam_resolution;
-    float* cam_sensorsize;
-    float* cam_intrinsic;
-
-    mjtByte* light_directional;
-    mjtByte* light_castshadow;
-    float* light_bulbradius;
-    mjtByte* light_active;
-    float* light_attenuation;
-    float* light_cutoff;
-    float* light_exponent;
-    float* light_ambient;
-    float* light_diffuse;
-    float* light_specular;
-
-    mjtByte* flex_flatskin;
-    int* flex_dim;
-    int* flex_matid;
-    int* flex_group;
-    int* flex_interp;
-    int* flex_nodeadr;
-    int* flex_nodenum;
-    int* flex_nodebodyid;
-    int* flex_vertadr;
-    int* flex_vertnum;
-    int* flex_elem;
-    int* flex_elemtexcoord;
-    int* flex_elemlayer;
-    int* flex_elemadr;
-    int* flex_elemnum;
-    int* flex_elemdataadr;
-    int* flex_shell;
-    int* flex_shellnum;
-    int* flex_shelldataadr;
-    int* flex_texcoordadr;
-    int* flex_bvhadr;
-    int* flex_bvhnum;
-    mjtByte* flex_centered;
-    mjtNum* flex_node;
-    mjtNum* flex_radius;
-    float* flex_rgba;
-    float* flex_texcoord;
-
-    int* hfield_pathadr;
-
-    int* mesh_bvhadr;
-    int* mesh_bvhnum;
-    int* mesh_texcoordadr;
-    int* mesh_graphadr;
-    int* mesh_pathadr;
-
-    int* skin_matid;
-    int* skin_group;
-    float* skin_rgba;
-    float* skin_inflate;
-    int* skin_vertadr;
-    int* skin_vertnum;
-    int* skin_texcoordadr;
-    int* skin_faceadr;
-    int* skin_facenum;
-    int* skin_boneadr;
-    int* skin_bonenum;
-    float* skin_vert;
-    int* skin_face;
-    int* skin_bonevertadr;
-    int* skin_bonevertnum;
-    float* skin_bonebindpos;
-    float* skin_bonebindquat;
-    int* skin_bonebodyid;
-    int* skin_bonevertid;
-    float* skin_bonevertweight;
-    int* skin_pathadr;
-
-    int* tex_pathadr;
-
-    int* mat_texid;
-    mjtByte* mat_texuniform;
-    float* mat_texrepeat;
-    float* mat_emission;
-    float* mat_specular;
-    float* mat_shininess;
-    float* mat_reflectance;
-    float* mat_metallic;
-    float* mat_roughness;
-    float* mat_rgba;
-
-    int* eq_type;
-    int* eq_obj1id;
-    int* eq_obj2id;
-    int* eq_objtype;
-    mjtNum* eq_data;
-
-    int* tendon_num;
-    int* tendon_matid;
-    int* tendon_group;
-    mjtByte* tendon_limited;
-    mjtNum* tendon_width;
-    mjtNum* tendon_range;
-    mjtNum* tendon_stiffness;
-    mjtNum* tendon_damping;
-    mjtNum* tendon_frictionloss;
-    mjtNum* tendon_lengthspring;
-    float* tendon_rgba;
-
-    int* actuator_trntype;
-    int* actuator_dyntype;
-    int* actuator_trnid;
-    int* actuator_actadr;
-    int* actuator_actnum;
-    int* actuator_group;
-    mjtByte* actuator_ctrllimited;
-    mjtByte* actuator_actlimited;
-    mjtNum* actuator_ctrlrange;
-    mjtNum* actuator_actrange;
-    mjtNum* actuator_cranklength;
-
-    int* sensor_type;
-    int* sensor_objid;
-    int* sensor_adr;
-
-    int* name_bodyadr;
-    int* name_jntadr;
-    int* name_geomadr;
-    int* name_siteadr;
-    int* name_camadr;
-    int* name_lightadr;
-    int* name_eqadr;
-    int* name_tendonadr;
-    int* name_actuatoradr;
-    char* names;
-    char* paths;
-  } model;
-
-  // fields in mjData that are necessary to re-render a scene
-  struct {
-    mjWarningStat warning[mjNWARNING];
-
-    int nefc;
-    int ncon;
-    int nisland;
-
-    mjtNum time;
-
-    mjtNum* act;
-
-    mjtNum* ctrl;
-    mjtNum* xfrc_applied;
-    mjtByte* eq_active;
-
-    mjtNum* sensordata;
-
-    mjtNum* xpos;
-    mjtNum* xquat;
-    mjtNum* xmat;
-    mjtNum* xipos;
-    mjtNum* ximat;
-    mjtNum* xanchor;
-    mjtNum* xaxis;
-    mjtNum* geom_xpos;
-    mjtNum* geom_xmat;
-    mjtNum* site_xpos;
-    mjtNum* site_xmat;
-    mjtNum* cam_xpos;
-    mjtNum* cam_xmat;
-    mjtNum* light_xpos;
-    mjtNum* light_xdir;
-
-    mjtNum* subtree_com;
-
-    int* ten_wrapadr;
-    int* ten_wrapnum;
-    int* wrap_obj;
-    mjtNum* ten_length;
-    mjtNum* wrap_xpos;
-
-    mjtNum* bvh_aabb_dyn;
-    mjtByte* bvh_active;
-    int* island_dofadr;
-    int* island_dofind;
-    int* dof_island;
-    int* efc_island;
-    int* tendon_efcadr;
-
-    mjtNum* flexvert_xpos;
-
-    mjContact* contact;
-    mjtNum* efc_force;
-    void* arena;
-  } data;
-};
-typedef struct mjvSceneState_ mjvSceneState;
-
 #endif  // MUJOCO_MJVISUALIZE_H_

include/mujoco/mujoco.h

@@ -16,7 +16,7 @@
 #define MUJOCO_MUJOCO_H_
 
 // header version; should match the library version as returned by mj_version()
-#define mjVERSION_HEADER 331
+#define mjVERSION_HEADER 333
 
 // needed to define size_t, fabs and log10
 #include <stdlib.h>
@@ -107,6 +107,9 @@ MJAPI mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error,
 // Compile spec to model.
 MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
 
+// Copy real-valued arrays from model to spec, returns 1 on success.
+MJAPI int mj_copyBack(mjSpec* s, const mjModel* m);
+
 // Recompile spec to model, preserving the state, return 0 on success.
 MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
 
@@ -187,6 +190,9 @@ MJAPI mjData* mj_makeData(const mjModel* m);
 // m is only required to contain the size fields from MJMODEL_INTS.
 MJAPI mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
 
+// Copy mjData, skip large arrays not required for visualization.
+MJAPI mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src);
+
 // Reset data to defaults.
 MJAPI void mj_resetData(const mjModel* m, mjData* d);
 
@@ -361,6 +367,9 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
 // Run composite rigid body inertia algorithm (CRB).
 MJAPI void mj_crb(const mjModel* m, mjData* d);
 
+// Make inertia matrix.
+MJAPI void mj_makeM(const mjModel* m, mjData* d);
+
 // Compute sparse L'*D*L factorizaton of inertia matrix.
 MJAPI void mj_factorM(const mjModel* m, mjData* d);
 
@@ -485,7 +494,7 @@ MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
 MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
 
 // Add inertia matrix to destination matrix.
-// Destination can be sparse uncompressed, or dense when all int* are NULL
+// Destination can be sparse or dense when all int* are NULL.
 MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind);
 
 // Apply Cartesian force and torque (outside xfrc_applied mechanism).
@@ -622,20 +631,10 @@ MJAPI void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum fo
 MJAPI void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
                           const mjvScene* scn, mjvCamera* cam);
 
-// Move camera with mouse given a scene state; action is mjtMouse.
-MJAPI void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action,
-                                   mjtNum reldx, mjtNum reldy,
-                                   const mjvScene* scn, mjvCamera* cam);
-
 // Move perturb object with mouse; action is mjtMouse.
 MJAPI void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
                            mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
 
-// Move perturb object with mouse given a scene state; action is mjtMouse.
-MJAPI void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action,
-                                    mjtNum reldx, mjtNum reldy,
-                                    const mjvScene* scn, mjvPerturb* pert);
-
 // Move model with mouse; action is mjtMouse.
 MJAPI void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
                          const mjtNum roomup[3], mjvScene* scn);
@@ -692,28 +691,9 @@ MJAPI void mjv_freeScene(mjvScene* scn);
 MJAPI void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
                            const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
 
-// Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings.
-MJAPI int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt,
-                                   const mjvPerturb* pert, mjvCamera* cam, int catmask,
-                                   mjvScene* scn);
-
 // Copy mjModel, skip large arrays not required for abstract visualization.
 MJAPI void mjv_copyModel(mjModel* dest, const mjModel* src);
 
-// Set default scene state.
-MJAPI void mjv_defaultSceneState(mjvSceneState* scnstate);
-
-// Allocate resources and initialize a scene state object.
-MJAPI void mjv_makeSceneState(const mjModel* m, const mjData* d,
-                              mjvSceneState* scnstate, int maxgeom);
-
-// Free scene state.
-MJAPI void mjv_freeSceneState(mjvSceneState* scnstate);
-
-// Update a scene state from model and data.
-MJAPI void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
-                                mjvSceneState* scnstate);
-
 // Add geoms from selected categories.
 MJAPI void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
                         const mjvPerturb* pert, int catmask, mjvScene* scn);
@@ -1308,6 +1288,19 @@ MJAPI char* mju_strncpy(char *dst, const char *src, int n);
 MJAPI mjtNum mju_sigmoid(mjtNum x);
 
 
+//---------------------------------- Signed Distance Function --------------------------------------
+
+// get sdf from geom id
+MJAPI const mjpPlugin* mjc_getSDF(const mjModel* m, int id);
+
+// signed distance function
+MJAPI mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]);
+
+// gradient of sdf
+MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3],
+                        const mjtNum x[3]);
+
+
 //---------------------------------- Derivatives ---------------------------------------------------
 
 // Finite differenced transition matrices (control theory notation)
@@ -1414,20 +1407,8 @@ MJAPI void mju_taskJoin(mjTask* task);
 
 //---------------------------------- Attachment ----------------------------------------------------
 
-// Attach child body to a parent frame, return the attached body if success or NULL otherwise.
+// Attach child to a parent, return the attached element if success or NULL otherwise.
-MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
+MJAPI mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
-                              const char* prefix, const char* suffix);
-
-// Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
-MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
-                                const char* prefix, const char* suffix);
-
-// Attach child body to a parent site, return the attached body if success or NULL otherwise.
-MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
-                                const char* prefix, const char* suffix);
-
-// Attach child frame to a parent site, return the attached frame if success or NULL otherwise.
-MJAPI mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child,
                              const char* prefix, const char* suffix);
 
 // Delete body and descendants from mjSpec, remove all references, return 0 on success.
@@ -1634,14 +1615,17 @@ MJAPI const char* mjs_getString(const mjString* source);
 // Get double array contents and optionally its size.
 MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size);
 
+// Get plugin attributes.
+MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
+
 
 //---------------------------------- Spec utilities ------------------------------------------------
 
 // Set element's default.
 MJAPI void mjs_setDefault(mjsElement* element, const mjsDefault* def);
 
-// Set element's enclosing frame.
+// Set element's enclosing frame, return 0 on success.
-MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
+MJAPI int mjs_setFrame(mjsElement* dest, mjsFrame* frame);
 
 // Resolve alternative orientations to quat, return error if any.
 MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
@@ -1653,6 +1637,13 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
 // Set user payload, overriding the existing value for the specified key if present.
 MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
 
+// Set user payload, overriding the existing value for the specified key if
+// present. This version differs from mjs_setUserValue in that it takes a
+// cleanup function that will be called when the user payload is deleted.
+MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key,
+                                       const void* data,
+                                       void (*cleanup)(const void*));
+
 // Return user payload or NULL if none found.
 MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key);
 

mjx/mujoco/mjx/__init__.py

@@ -17,6 +17,7 @@
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
 from mujoco.mjx._src.constraint import make_constraint
+from mujoco.mjx._src.derivative import deriv_smooth_vel
 from mujoco.mjx._src.forward import euler
 from mujoco.mjx._src.forward import forward
 from mujoco.mjx._src.forward import fwd_acceleration
@@ -26,6 +27,7 @@ from mujoco.mjx._src.forward import fwd_velocity
 from mujoco.mjx._src.forward import implicit
 from mujoco.mjx._src.forward import rungekutta4
 from mujoco.mjx._src.forward import step
+from mujoco.mjx._src.inverse import inverse
 from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import get_data_into
 from mujoco.mjx._src.io import make_data

mjx/mujoco/mjx/_src/collision_convex.py

@@ -28,8 +28,10 @@ from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.collision_types import HFieldInfo
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 
 _GeomInfo = Union[GeomInfo, ConvexInfo]
@@ -42,6 +44,9 @@ def collider(ncon: int):
     def collide(
         m: Model, d: Data, key: FunctionKey, geom: jax.Array
     ) -> Collision:
+      if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+        raise ValueError('collider requires JAX backend implementation.')
+
       g1, g2 = geom.T
       infos = [
           GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
@@ -56,7 +61,7 @@ def collider(ncon: int):
               pos=0, mat=0, size=0, face=0, vert=0
           )
         elif key.types[i] == GeomType.MESH:
-          c, cm = infos[i], m.mesh_convex[key.data_ids[i]]
+          c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]]
           infos[i] = ConvexInfo(**vars(c), **vars(cm))
           in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
               pos=0, mat=0, size=0

mjx/mujoco/mjx/_src/collision_driver.py

@@ -71,9 +71,11 @@ from mujoco.mjx._src.collision_sdf import sphere_ellipsoid
 from mujoco.mjx._src.collision_types import FunctionKey
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -227,7 +229,7 @@ def _geom_groups(
     if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD:
       # add static grid bounds to the grouping key for hfield collisions
       geom_rbound_hfield = (
-          m.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound
+          m._impl.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound  # pytype: disable=attribute-error
       )
       nrow, ncol = m.hfield_nrow[data_ids[0]], m.hfield_ncol[data_ids[0]]
       xsize, ysize = m.hfield_size[data_ids[0]][:2]
@@ -323,11 +325,11 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
         solref=solref,
         solreffriction=solreffriction,
         solimp=solimp,
-        dim=d.contact.dim,
+        dim=d._impl.contact.dim,  # pytype: disable=attribute-error
         geom1=jp.array(geom[:, 0]),
         geom2=jp.array(geom[:, 1]),
         geom=jp.array(geom[:, :2]),
-        efc_address=d.contact.efc_address,
+        efc_address=d._impl.contact.efc_address,  # pytype: disable=attribute-error
     )
 
   return groups
@@ -361,8 +363,12 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
   condim_counts = {}
   for k, v in group_counts.items():
+    if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
+      ncon = m.opt.sdf_initpoints
+    else:
       func = _COLLISION_FUNC[k.types]
-    num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v  # pytype: disable=attribute-error
+      ncon = func.ncon  # pytype: disable=attribute-error
+    num_contacts = condim_counts.get(k.condim, 0) + ncon * v
     if max_contact_points > -1:
       num_contacts = min(max_contact_points, num_contacts)
     condim_counts[k.condim] = num_contacts
@@ -374,7 +380,10 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
 
 def collision(m: Model, d: Data) -> Data:
   """Collides geometries."""
-  if d.ncon == 0:
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('collision requires JAX backend implementation.')
+
+  if d._impl.ncon == 0:  # pytype: disable=attribute-error
     return d
 
   max_geom_pairs = _numeric(m, 'max_geom_pairs')
@@ -424,4 +433,4 @@ def collision(m: Model, d: Data) -> Data:
   contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
   contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
 
-  return d.replace(contact=contact)
+  return d.tree_replace({'_impl.contact': contact})

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -119,7 +119,7 @@ class SphereCollisionTest(parameterized.TestCase):
   def test_sphere(self, name, mjcf):
     d, dx = _collide(mjcf)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-5)
 
   _SPHERE_CONVEX = """
     <mujoco>
@@ -143,7 +143,7 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     self.assertEmpty(d.contact.dist)
-    self.assertGreater(dx.contact.dist, 0)
+    self.assertGreater(dx._impl.contact.dist, 0)
 
     # face contact
     xml = self._SPHERE_CONVEX.replace(
@@ -151,38 +151,38 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'face', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'face', 1e-4)
 
     # deep face contact
     xml = self._SPHERE_CONVEX.replace(
         '<body pos="0.52 0 0.52">', '<body pos="0.48 0 0.47">'
     )
     d, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     self.assertTrue((d.contact.dist < 0).all())
-    np.testing.assert_allclose(dx.contact.dist, [-0.07], atol=1e-5)
+    np.testing.assert_allclose(dx._impl.contact.dist, [-0.07], atol=1e-5)
-    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+        dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
     )
 
   def test_sphere_convex_edge(self):
     # edge contact
     d, dx = _collide(self._SPHERE_CONVEX)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'edge', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'edge', 1e-4)
 
     # deep edge penetration
     xml = self._SPHERE_CONVEX.replace(
         '<body pos="0.52 0 0.52">', '<body pos="0.49 0 0.49">'
     )
     d, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     self.assertTrue((d.contact.dist < 0).all())
-    np.testing.assert_allclose(dx.contact.dist, [-0.06], atol=1e-5)
+    np.testing.assert_allclose(dx._impl.contact.dist, [-0.06], atol=1e-5)
-    np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
+    np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
+        dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
     )
 
     # vertex contact
@@ -191,7 +191,7 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex', 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, 'vertex', 1e-4)
 
     # sphere center on vertex
     xml = self._SPHERE_CONVEX.replace(
@@ -199,7 +199,9 @@ class SphereCollisionTest(parameterized.TestCase):
     )
     d, dx = _collide(xml)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'vertex_center', 1e-4)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'vertex_center', 1e-4
+      )
 
 
 class EllipsoidCollisionTest(parameterized.TestCase):
@@ -219,10 +221,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_plane_ellipsoid(self):
     """Tests ellipsoid plane contact."""
     d, dx = _collide(self._ELLIPSOID_PLANE)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
       )
 
   _ELLIPSOID_ELLIPSOID = """
@@ -242,10 +244,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_ellipsoid_ellipsoid(self):
     """Tests ellipsoid ellipsoid contact."""
     d, dx = _collide(self._ELLIPSOID_ELLIPSOID)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
       )
 
   _ELLIPSOID_SPHERE = """
@@ -266,10 +268,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     """Tests ellipsoid capsule contact."""
     d, dx = _collide(self._ELLIPSOID_SPHERE)
     d.contact.pos[0][2] = 0.03  # MJX finds the point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
       )
 
   _ELLIPSOID_CAPSULE = """
@@ -289,10 +291,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
   def test_capsule_ellipsoid(self):
     """Tests ellipsoid capsule contact."""
     d, dx = _collide(self._ELLIPSOID_CAPSULE)
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
       )
 
   _ELLIPSOID_CYLINDER = """
@@ -313,10 +315,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
     """Tests ellipsoid cylinder contact."""
     d, dx = _collide(self._ELLIPSOID_CYLINDER)
     d.contact.pos[0][2] = 0.04  # MJX finds the deepest point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
+          dx._impl.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
       )
 
 
@@ -357,7 +359,7 @@ class CapsuleCollisionTest(parameterized.TestCase):
   def test_capsule(self, name, mjcf):
     d, dx = _collide(mjcf)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4)
+      _assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-4)
 
   _PARALLEL_CAP = """
     <mujoco>
@@ -378,14 +380,14 @@ class CapsuleCollisionTest(parameterized.TestCase):
     """Tests that two parallel capsules are colliding at the midpoint."""
     _, dx = _collide(self._PARALLEL_CAP)
 
-    np.testing.assert_allclose(dx.contact.dist, -0.05)
+    np.testing.assert_allclose(dx._impl.contact.dist, -0.05)
     np.testing.assert_allclose(
-        dx.contact.pos[0],
+        dx._impl.contact.pos[0],
         np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]),
         atol=1e-5,
     )
     np.testing.assert_allclose(
-        dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
+        dx._impl.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
     )
 
   _CAP_BOX = """
@@ -408,30 +410,32 @@ class CapsuleCollisionTest(parameterized.TestCase):
     d, dx = _collide(self._CAP_BOX)
 
     # sort positions for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
     d.contact.frame[:] = d.contact.frame[idx]
     d.contact.dist[:] = d.contact.dist[idx]
 
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'capsule_convex', 1e-4
+      )
 
   def test_capsule_convex_face_deep(self):
     """Tests deep face penetration."""
     xml = self._CAP_BOX.replace('<body pos="0 0 0.54">', '<body pos="0 0 0.4">')
 
     _, dx = _collide(xml)
-    self.assertTrue((dx.contact.dist < 0).all())
+    self.assertTrue((dx._impl.contact.dist < 0).all())
     np.testing.assert_array_almost_equal(
-        dx.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
+        dx._impl.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.dist, np.array([-0.15, -0.15])
+        dx._impl.contact.dist, np.array([-0.15, -0.15])
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
+        dx._impl.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
     )
 
   _CAP_EDGE_BOX = """
@@ -453,11 +457,11 @@ class CapsuleCollisionTest(parameterized.TestCase):
     """Tests edge contact."""
     d, dx = _collide(self._CAP_EDGE_BOX)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
@@ -469,13 +473,17 @@ class CapsuleCollisionTest(parameterized.TestCase):
     )
     _, dx = _collide(xml)
 
-    np.testing.assert_array_equal(dx.contact.dist < 0, np.array([True, False]))
+    np.testing.assert_array_equal(
-    np.testing.assert_array_almost_equal(dx.contact.dist[0], np.array([-0.13]))
+        dx._impl.contact.dist < 0, np.array([True, False])
-    np.testing.assert_array_almost_equal(
-        dx.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
     )
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
+        dx._impl.contact.dist[0], np.array([-0.13])
+    )
+    np.testing.assert_array_almost_equal(
+        dx._impl.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
+    )
+    np.testing.assert_array_almost_equal(
+        dx._impl.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
     )
 
   def test_capsule_convex_edge_shallow_tip(self):
@@ -489,16 +497,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
     xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.58 0 0.55"')
     d, dx = _collide(xml)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
+        dx._impl.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
     )
 
     # the capsule sphere is outside the edge voronoi region, so there is a
@@ -510,16 +518,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
     xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.5 0 0.52"')
     d, dx = _collide(xml)
 
-    c = dx.contact
+    c = dx._impl.contact
     self.assertEqual(c.pos.shape[0], 2)
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4)
     np.testing.assert_array_almost_equal(
-        dx.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
+        dx._impl.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
     )
 
 
@@ -542,17 +550,17 @@ class CylinderTest(absltest.TestCase):
     d, dx = _collide(self._CYLINDER_PLANE)
 
     # cylinder is lying flat
-    np.testing.assert_array_less(dx.contact.dist[:2], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
-    np.testing.assert_array_less(-dx.contact.dist[2:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
 
     # sort position for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
 
     # extract the contact points with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:2], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:2], dx._impl.contact)
     c = c.replace(dim=c.dim[:2], efc_address=c.efc_address[:2])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'cylinder_plane', 1e-5)
@@ -564,9 +572,11 @@ class CylinderTest(absltest.TestCase):
     xml = xml.replace('pos="0 0 0.04"', 'pos="0 0 0.095"')
     d, dx = _collide(xml)
 
-    np.testing.assert_array_less(dx.contact.dist, 0)
+    np.testing.assert_array_less(dx._impl.contact.dist, 0)
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'cylinder_plane', 1e-5)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'cylinder_plane', 1e-5
+      )
 
   _SPHERE_CYLINDER = """
     <mujoco>
@@ -586,10 +596,10 @@ class CylinderTest(absltest.TestCase):
     """Tests sphere cylinder contact."""
     d, dx = _collide(self._SPHERE_CYLINDER)
     d.contact.pos[0][2] = 0.05  # MJX finds the deepest point on the surface
-    self.assertLess(dx.contact.dist[0], 0)
+    self.assertLess(dx._impl.contact.dist[0], 0)
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(
-          dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
+          dx._impl.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
       )
 
 
@@ -612,11 +622,11 @@ class ConvexTest(absltest.TestCase):
     """Tests box collision with a plane."""
     d, dx = _collide(self._BOX_PLANE)
 
-    np.testing.assert_array_less(dx.contact.dist[:2], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
-    np.testing.assert_array_less(-dx.contact.dist[2:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
     # extract the contact points with penetration
     c = jax.tree_util.tree_map(
-        lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact
+        lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx._impl.contact
     )
     c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]])
     for field in dataclasses.fields(Contact):
@@ -638,16 +648,18 @@ class ConvexTest(absltest.TestCase):
     """Tests box collision with a plane."""
     d, dx = _collide(self._FLAT_BOX_PLANE)
 
-    np.testing.assert_array_less(dx.contact.dist, 0)
+    np.testing.assert_array_less(dx._impl.contact.dist, 0)
 
     # sort positions for comparison
-    idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+    idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
-    dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+    dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
     idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
     d.contact.pos[:] = d.contact.pos[idx]
 
     for field in dataclasses.fields(Contact):
-      _assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5)
+      _assert_attr_eq(
+          dx._impl.contact, d.contact, field.name, 'flat_box_plane', 1e-5
+      )
 
   _BOX_BOX = """
     <mujoco>
@@ -669,7 +681,7 @@ class ConvexTest(absltest.TestCase):
   def test_box_box(self):
     """Tests a face contact for a box-box collision."""
     d, dx = _collide(self._BOX_BOX, keyframe=0)
-    c = dx.contact
+    c = dx._impl.contact
 
     self.assertEqual(c.pos.shape[0], 4)
     np.testing.assert_array_less(c.dist, 0)
@@ -701,10 +713,10 @@ class ConvexTest(absltest.TestCase):
     d, dx = _collide(self._BOX_BOX_EDGE)
 
     # Only one contact point.
-    np.testing.assert_array_less(dx.contact.dist[:1], 0)
+    np.testing.assert_array_less(dx._impl.contact.dist[:1], 0)
-    np.testing.assert_array_less(-dx.contact.dist[1:], 0)
+    np.testing.assert_array_less(-dx._impl.contact.dist[1:], 0)
     # extract the contact point with penetration
-    c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
+    c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
     c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
@@ -736,7 +748,7 @@ class ConvexTest(absltest.TestCase):
         ).read_bytes(),
     }
     _, dx = _collide(self._CONVEX_CONVEX, assets=assets)
-    c = dx.contact
+    c = dx._impl.contact
 
     # Only one contact point for an edge contact.
     self.assertLess(c.dist[0], 0)
@@ -766,7 +778,7 @@ class ConvexTest(absltest.TestCase):
   def test_convex_convex_edge(self):
     """Tests convex-convex collisions with edge contact via _sat_gaussmap."""
     _, dx = _collide(self._CONVEX_CONVEX_THIN)
-    c = dx.contact
+    c = dx._impl.contact
 
     # Only one contact point for an edge contact.
     self.assertLess(c.dist[0], 0)
@@ -781,7 +793,7 @@ class ConvexTest(absltest.TestCase):
             'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'
         )
     )
-    c = dx.contact
+    c = dx._impl.contact
     self.assertTrue((c.dist > 0).all())
 
 
@@ -837,8 +849,8 @@ class HFieldTest(absltest.TestCase):
 
     # check that all geoms are colliding with the hfield
     for geom_id in [1, 2, 3]:
-      mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
+      mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
-      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
+      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
       self.assertTrue((c.dist < 0).any())
       self.assertTrue((c.dist > -1e-3).any())
       # all contact normals are roughly pointing in the right direction
@@ -850,7 +862,7 @@ class HFieldTest(absltest.TestCase):
     for p in positions:
       xml = self._HFIELD.replace('<body pos="0 0', f'<body pos="{p}')
       _, dx = _collide(xml)
-      self.assertTrue((dx.contact.dist >= 0).all())
+      self.assertTrue((dx._impl.contact.dist >= 0).all())
 
   def test_hfield_deep(self):
     """Tests that objects with deep penetration do not get stuck."""
@@ -868,8 +880,8 @@ class HFieldTest(absltest.TestCase):
 
     # check that all geoms are colliding with the hfield
     for geom_id in [1, 2, 3]:
-      mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
+      mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
-      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
+      c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
       # all contact normals are in the top half-face of the hfield
       self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all())
 
@@ -892,8 +904,8 @@ class BodyPairFilterTest(absltest.TestCase):
   def test_filter_self_collision(self):
     """Tests that self collisions get filtered."""
     d, dx = _collide(self._SELF_COLLISION)
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 0)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 0)
 
   _PARENT_CHILD = """
     <mujoco>
@@ -923,8 +935,8 @@ class BodyPairFilterTest(absltest.TestCase):
     dx = kinematics_jit_fn(mx, dx)
     dx = collision_jit_fn(mx, dx)
 
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 0)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 0)
 
   def test_disable_filter_parent_child(self):
     """Tests that filterparent flag disables parent-child filtering."""
@@ -941,8 +953,8 @@ class BodyPairFilterTest(absltest.TestCase):
     dx = collision_jit_fn(mx, dx)
 
     # one collision between parent-child spheres
-    self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
+    self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
-    self.assertEqual(dx.contact.pos.shape[0], 1)
+    self.assertEqual(dx._impl.contact.pos.shape[0], 1)
 
 
 class DimTest(parameterized.TestCase):
@@ -1027,8 +1039,8 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.contact.dist.shape, (3,))
+    self.assertEqual(dx_all._impl.contact.dist.shape, (3,))
-    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
+    self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
 
   _CAPSULES_MAX_PAIR = """
     <mujoco>
@@ -1073,9 +1085,9 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.contact.dist.shape, (6,))
+    self.assertEqual(dx_all._impl.contact.dist.shape, (6,))
-    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
+    self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
-    self.assertTrue((dx_top_k.contact.dist < 0).all())
+    self.assertTrue((dx_top_k._impl.contact.dist < 0).all())
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/collision_sdf.py

@@ -32,7 +32,9 @@ from mujoco.mjx._src.collision_types import Collision
 from mujoco.mjx._src.collision_types import GeomInfo
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 # pylint: enable=g-importing-member
 
 # the SDF function takes position in, and returns a distance or objective

mjx/mujoco/mjx/_src/constraint.py

@@ -28,10 +28,12 @@ from mujoco.mjx._src.types import ConeType
 from mujoco.mjx._src.types import ConstraintType
 from mujoco.mjx._src.types import Contact
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import ObjType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -281,6 +283,10 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for tendon equality constraints."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_equality_tendon requires JAX backend implementation.'
+    )
 
   eq_id = np.nonzero(m.eq_type == EqType.TENDON)[0]
 
@@ -314,9 +320,9 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
     return jax.tree_util.tree_map(lambda x: x * active, efc)
 
   inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
-  jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id]
+  jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id]
-  pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id]
+  pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id]
-  pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id]
+  pos2 = d._impl.ten_length[obj2id] - m.tendon_length0[obj2id]
   invweight = inv1 + inv2 * (obj2id > -1)
 
   return rows(
@@ -326,8 +332,11 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for dof frictionloss."""
-  dof_id = np.nonzero(m.dof_hasfrictionloss)[0]
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
-  tendon_id = np.nonzero(m.tendon_hasfrictionloss)[0]
+    raise ValueError('_efc_friction requires JAX backend implementation.')
+
+  dof_id = np.nonzero(m._impl.dof_hasfrictionloss)[0]
+  tendon_id = np.nonzero(m._impl.tendon_hasfrictionloss)[0]
 
   size = dof_id.size + tendon_id.size
   if (m.opt.disableflags & DisableBit.FRICTIONLOSS) or (size == 0):
@@ -337,7 +346,7 @@ def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
   args_dof += (m.dof_solimp,)
   args_dof = jax.tree_util.tree_map(lambda x: x[dof_id], args_dof)
 
-  args_ten = (d.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
+  args_ten = (d._impl.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
   args_ten += (m.tendon_solref_fri, m.tendon_solimp_fri)
   args_ten = jax.tree_util.tree_map(lambda x: x[tendon_id], args_ten)
 
@@ -415,6 +424,9 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for tendon limits."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_efc_limit_tendon requires JAX backend implementation.')
+
   tendon_id = np.nonzero(m.tendon_limited)[0]
 
   if (m.opt.disableflags & DisableBit.LIMIT) or tendon_id.size == 0:
@@ -423,8 +435,8 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
   length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map(
       lambda x: x[tendon_id],
       (
-          d.ten_length,
+          d._impl.ten_length,
-          d.ten_J,
+          d._impl.ten_J,
           m.tendon_range,
           m.tendon_margin,
           m.tendon_invweight0,
@@ -446,8 +458,12 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
 
 def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for frictionless contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_frictionless requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == 1)[0]
+  con_id = np.nonzero(d._impl.contact.dim == 1)[0]
 
   if con_id.size == 0:
     return None
@@ -473,15 +489,19 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
 
   return rows(contact)
 
 
 def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
   """Calculates constraint rows for frictional pyramidal contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_pyramidal requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == condim)[0]
+  con_id = np.nonzero(d._impl.contact.dim == condim)[0]
 
   if con_id.size == 0:
     return None
@@ -518,15 +538,19 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
   # concatenate to drop row grouping
   return jax.tree_util.tree_map(jp.concatenate, rows(contact))
 
 
 def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
   """Calculates constraint rows for frictional elliptic contacts."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError(
+        '_efc_contact_elliptic requires JAX backend implementation.'
+    )
 
-  con_id = np.nonzero(d.contact.dim == condim)[0]
+  con_id = np.nonzero(d._impl.contact.dim == condim)[0]
 
   if con_id.size == 0:
     return None
@@ -563,7 +587,7 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
         jp.zeros_like(pos),
     )
 
-  contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
+  contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
   # concatenate to drop row grouping
   return jax.tree_util.tree_map(jp.concatenate, rows(contact))
 
@@ -602,14 +626,14 @@ def make_efc_type(
 
   if not m.opt.disableflags & DisableBit.FRICTIONLOSS:
     nf_dof = (
-        m.dof_hasfrictionloss.sum()
+        m._impl.dof_hasfrictionloss.sum()
-        if isinstance(m, Model)
+        if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
         else (m.dof_frictionloss > 0).sum()
     )
     efc_types += [ConstraintType.FRICTION_DOF] * nf_dof
     nf_tendon = (
-        m.tendon_hasfrictionloss.sum()
+        m._impl.tendon_hasfrictionloss.sum()
-        if isinstance(m, Model)
+        if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
         else (m.tendon_frictionloss > 0).sum()
     )
     efc_types += [ConstraintType.FRICTION_TENDON] * nf_tendon
@@ -683,10 +707,14 @@ def make_constraint(m: Model, d: Data) -> Data:
 
   if not efcs:
     z = jp.empty(0)
-    d = d.replace(efc_J=jp.empty((0, m.nv)))
+    d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})
-    d = d.replace(
+    d = d.tree_replace({
-        efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z, efc_margin=z
+        '_impl.efc_D': z,
-    )
+        '_impl.efc_aref': z,
+        '_impl.efc_frictionloss': z,
+        '_impl.efc_pos': z,
+        '_impl.efc_margin': z,
+    })
     return d
 
   efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -699,9 +727,13 @@ def make_constraint(m: Model, d: Data) -> Data:
     return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
 
   aref, r, pos, margin, frictionloss = fn(efc)
-  d = d.replace(
+  d = d.tree_replace({
-      efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin
+      '_impl.efc_J': efc.J,
-  )
+      '_impl.efc_D': 1 / r,
-  d = d.replace(efc_frictionloss=frictionloss)
+      '_impl.efc_aref': aref,
+      '_impl.efc_pos': pos,
+      '_impl.efc_margin': margin,
+  })
+  d = d.tree_replace({'_impl.efc_frictionloss': frictionloss})
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -69,16 +69,16 @@ class ConstraintTest(parameterized.TestCase):
 
       order = test_util.efc_order(m, d, dx)
       d_efc_j = d.efc_J.reshape((-1, m.nv))
-      _assert_eq(d_efc_j, dx.efc_J[order][: d.nefc], 'efc_J')
+      _assert_eq(d_efc_j, dx._impl.efc_J[order][: d.nefc], 'efc_J')
-      _assert_eq(0, dx.efc_J[order][d.nefc :], 'efc_J')
+      _assert_eq(0, dx._impl.efc_J[order][d.nefc :], 'efc_J')
-      _assert_eq(d.efc_aref, dx.efc_aref[order][: d.nefc], 'efc_aref')
+      _assert_eq(d.efc_aref, dx._impl.efc_aref[order][: d.nefc], 'efc_aref')
-      _assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
+      _assert_eq(0, dx._impl.efc_aref[order][d.nefc :], 'efc_aref')
-      _assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
+      _assert_eq(d.efc_D, dx._impl.efc_D[order][: d.nefc], 'efc_D')
-      _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
+      _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
-      _assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos')
+      _assert_eq(dx._impl.efc_pos[order][d.nefc :], 0, 'efc_pos')
       _assert_eq(
           d.efc_frictionloss,
-          dx.efc_frictionloss[order][: d.nefc],
+          dx._impl.efc_frictionloss[order][: d.nefc],
           'efc_frictionloss',
       )
 
@@ -104,7 +104,7 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 0)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 0)
+    self.assertEqual(dx._impl.efc_J.shape[0], 0)
 
   def test_disable_equality(self):
     m = test_util.load_test_file('constraints.xml')
@@ -115,7 +115,9 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 180)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 187)  # only joint/tendon limit, contact
+    self.assertEqual(
+        dx._impl.efc_J.shape[0], 187
+    )  # only joint/tendon limit, contact
 
   def test_disable_contact(self):
     m = test_util.load_test_file('constraints.xml')
@@ -126,7 +128,9 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 0)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 27)  # only equality, joint/tendon limit
+    self.assertEqual(
+        dx._impl.efc_J.shape[0], 27
+    )  # only equality, joint/tendon limit
 
   def test_disable_frictionloss(self):
     m = test_util.load_test_file('constraints.xml')
@@ -137,7 +141,7 @@ class ConstraintTest(parameterized.TestCase):
     self.assertEqual(nl, 5)
     self.assertEqual(nc, 180)
     dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    self.assertEqual(dx.efc_J.shape[0], 205)
+    self.assertEqual(dx._impl.efc_J.shape[0], 205)
 
   def test_margin(self):
     """Test margin."""
@@ -163,8 +167,8 @@ class ConstraintTest(parameterized.TestCase):
     dx = mjx.make_constraint(mx, dx)
 
     order = test_util.efc_order(m, d, dx)
-    _assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
+    _assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
-    _assert_eq(d.efc_margin, dx.efc_margin[order][: d.nefc], 'efc_margin')
+    _assert_eq(d.efc_margin, dx._impl.efc_margin[order][: d.nefc], 'efc_margin')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/dataclasses.py

@@ -18,6 +18,7 @@ import copy
 import dataclasses
 import typing
 from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
+import warnings
 import jax
 import numpy as np
 

mjx/mujoco/mjx/_src/derivative.py

@@ -0,0 +1,61 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Derivative functions."""
+
+from typing import Optional
+
+import jax
+from jax import numpy as jp
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import BiasType
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import DynType
+from mujoco.mjx._src.types import GainType
+from mujoco.mjx._src.types import Model
+
+
+def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
+  """Analytical derivative of smooth forces w.r.t velocities."""
+
+  qderiv = None
+
+  # qDeriv += d qfrc_actuator / d qvel
+  if not m.opt.disableflags & DisableBit.ACTUATION:
+    affine_bias = m.actuator_biastype == BiasType.AFFINE
+    bias_vel = m.actuator_biasprm[:, 2] * affine_bias
+    affine_gain = m.actuator_gaintype == GainType.AFFINE
+    gain_vel = m.actuator_gainprm[:, 2] * affine_gain
+    ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
+    vel = bias_vel + gain_vel * ctrl
+    qderiv = d._impl.actuator_moment.T @ jax.vmap(jp.multiply)(
+        d._impl.actuator_moment, vel
+    )
+
+  # qDeriv += d qfrc_passive / d qvel
+  if not m.opt.disableflags & DisableBit.PASSIVE:
+    if qderiv is None:
+      qderiv = -jp.diag(m.dof_damping)
+    else:
+      qderiv -= jp.diag(m.dof_damping)
+    if m.ntendon:
+      qderiv -= d._impl.ten_J.T @ jp.diag(m.tendon_damping) @ d._impl.ten_J
+    # TODO(robotics-simulation): fluid drag model
+    if m.opt.has_fluid_params:  # pytype: disable=attribute-error
+      raise NotImplementedError('fluid drag not supported for implicitfast')
+
+  # TODO(team): rne derivative
+
+  return qderiv

mjx/mujoco/mjx/_src/forward.py

@@ -22,6 +22,7 @@ from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
+from mujoco.mjx._src import derivative
 from mujoco.mjx._src import math
 from mujoco.mjx._src import passive
 from mujoco.mjx._src import scan
@@ -32,12 +33,15 @@ from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import BiasType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import DynType
 from mujoco.mjx._src.types import GainType
 from mujoco.mjx._src.types import IntegratorType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
+from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -79,10 +83,13 @@ def fwd_position(m: Model, d: Data) -> Data:
 @named_scope
 def fwd_velocity(m: Model, d: Data) -> Data:
   """Velocity-dependent computations."""
-  d = d.replace(
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
-      actuator_velocity=d.actuator_moment @ d.qvel,
+    raise ValueError('fwd_velocity requires JAX backend implementation.')
-      ten_velocity=d.ten_J @ d.qvel,
+
-  )
+  d = d.tree_replace({
+      '_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
+      '_impl.ten_velocity': d._impl.ten_J @ d.qvel,
+  })
   d = smooth.com_vel(m, d)
   d = passive.passive(m, d)
   d = smooth.rne(m, d)
@@ -171,13 +178,41 @@ def fwd_actuation(m: Model, d: Data) -> Data:
       m.actuator_gainprm,
       m.actuator_biastype,
       m.actuator_biasprm,
-      d.actuator_length,
+      d._impl.actuator_length,
-      d.actuator_velocity,
+      d._impl.actuator_velocity,
       ctrl_act,
       jp.array(m.actuator_lengthrange),
       jp.array(m.actuator_acc0),
       group_by='u',
   )
+
+  # tendon total force clamping
+  if np.any(m.tendon_actfrclimited):
+    (tendon_actfrclimited_id,) = np.nonzero(m.tendon_actfrclimited)
+    actuator_tendon = m.actuator_trntype == TrnType.TENDON
+
+    force_mask = [
+        actuator_tendon & (m.actuator_trnid[:, 0] == tendon_id)
+        for tendon_id in tendon_actfrclimited_id
+    ]
+    force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask])
+    force_mat = np.array(force_mask)[:, force_ids]
+    tendon_total_force = force_mat @ force[force_ids]
+
+    force_scaling = jp.where(
+        tendon_total_force < m.tendon_actfrcrange[tendon_actfrclimited_id, 0],
+        m.tendon_actfrcrange[tendon_actfrclimited_id, 0] / tendon_total_force,
+        1,
+    )
+    force_scaling = jp.where(
+        tendon_total_force > m.tendon_actfrcrange[tendon_actfrclimited_id, 1],
+        m.tendon_actfrcrange[tendon_actfrclimited_id, 1] / tendon_total_force,
+        force_scaling,
+    )
+
+    tendon_forces = force[force_ids] * (force_mat.T @ force_scaling)
+    force = force.at[force_ids].set(tendon_forces)
+
   forcerange = jp.where(
       m.actuator_forcelimited[:, None],
       m.actuator_forcerange,
@@ -185,7 +220,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
   )
   force = jp.clip(force, forcerange[:, 0], forcerange[:, 1])
 
-  qfrc_actuator = d.actuator_moment.T @ force
+  qfrc_actuator = d._impl.actuator_moment.T @ force
 
   if m.ngravcomp:
     # actuator-level gravity compensation, skip if added as passive force
@@ -301,13 +336,17 @@ def _advance(
 @named_scope
 def euler(m: Model, d: Data) -> Data:
   """Euler integrator, semi-implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('euler requires JAX backend implementation.')
+
   # integrate damping implicitly
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     if support.is_sparse(m):
-      dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
     else:
-      dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping))
+      qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
+    dh = d.tree_replace({'_impl.qM': qM})
     dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
@@ -317,7 +356,7 @@ def euler(m: Model, d: Data) -> Data:
 @named_scope
 def rungekutta4(m: Model, d: Data) -> Data:
   """Runge-Kutta explicit order 4 integrator."""
-  d_t0 = d
+  d0 = d
   # pylint: disable=invalid-name
   A, B = _RK4_A, _RK4_B
   C = jp.tril(A).sum(axis=0)  # C(i) = sum_j A(i,j)
@@ -338,9 +377,9 @@ def rungekutta4(m: Model, d: Data) -> Data:
         lambda k: a * k, (kqvel, d.qacc, d.act_dot)
     )
     # get intermediate RK solutions
-    kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel)
+    kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d0.qpos, dqvel)
-    kact = d_t0.act + dact_dot * m.opt.timestep
+    kact = d0.act + dact_dot * m.opt.timestep
-    kqvel = d_t0.qvel + dqacc * m.opt.timestep
+    kqvel = d0.qvel + dqacc * m.opt.timestep
     d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t)
     d = forward(m, d)
 
@@ -352,44 +391,25 @@ def rungekutta4(m: Model, d: Data) -> Data:
 
   abt = jp.vstack([jp.diag(A), B[1:4], T]).T
   out, _ = jax.lax.scan(f, (qvel, qacc, act_dot, kqvel, d), abt, unroll=3)
-  qvel, qacc, act_dot, *_ = out
+  qvel, qacc, act_dot, _, d1 = out
 
-  d = _advance(m, d_t0, act_dot, qacc, qvel)
+  d = d1.replace(qpos=d0.qpos, qvel=d0.qvel, act=d0.act, time=d0.time)
+  d = _advance(m, d, act_dot, qacc, qvel)
   return d
 
 
 @named_scope
 def implicit(m: Model, d: Data) -> Data:
   """Integrates fully implicit in velocity."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('implicit requires JAX backend implementation.')
 
-  qderiv = None
+  qderiv = derivative.deriv_smooth_vel(m, d)
-
-  # qDeriv += d qfrc_actuator / d qvel
-  if not m.opt.disableflags & DisableBit.ACTUATION:
-    affine_bias = m.actuator_biastype == BiasType.AFFINE
-    bias_vel = m.actuator_biasprm[:, 2] * affine_bias
-    affine_gain = m.actuator_gaintype == GainType.AFFINE
-    gain_vel = m.actuator_gainprm[:, 2] * affine_gain
-    ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
-    vel = bias_vel + gain_vel * ctrl
-    qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment
-
-  # qDeriv += d qfrc_passive / d qvel
-  if not m.opt.disableflags & DisableBit.PASSIVE:
-    if qderiv is None:
-      qderiv = -jp.diag(m.dof_damping)
-    else:
-      qderiv -= jp.diag(m.dof_damping)
-    if m.ntendon:
-      qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J
-    # TODO(robotics-simulation): fluid drag model
-    if m.opt.has_fluid_params:
-      raise NotImplementedError('fluid drag not supported for implicitfast')
 
   qacc = d.qacc
   if qderiv is not None:
     # TODO(robotics-simulation): use smooth.factor_m / solve_m here:
-    qm = support.full_m(m, d) if support.is_sparse(m) else d.qM
+    qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
     qm -= m.opt.timestep * qderiv
     qh, _ = jax.scipy.linalg.cho_factor(qm)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
@@ -401,6 +421,9 @@ def implicit(m: Model, d: Data) -> Data:
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('forward requires JAX backend implementation.')
+
   d = fwd_position(m, d)
   d = sensor.sensor_pos(m, d)
   d = fwd_velocity(m, d)
@@ -409,7 +432,7 @@ def forward(m: Model, d: Data) -> Data:
   d = fwd_acceleration(m, d)
   d = sensor.sensor_acc(m, d)
 
-  if d.efc_J.size == 0:
+  if d._impl.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)
     return d
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
+from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
@@ -130,6 +131,7 @@ class ForwardTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'qpos')
     _assert_attr_eq(d, dx, 'act')
     _assert_attr_eq(d, dx, 'time')
+    _assert_attr_eq(d, dx, 'xpos')
 
   def test_eulerdamp(self):
     m = test_util.load_test_file('pendula.xml')
@@ -195,6 +197,27 @@ class ActuatorTest(parameterized.TestCase):
     dx = jax.jit(mjx.euler)(mx, dx)
     _assert_attr_eq(d, dx, 'act')
 
+  def test_tendon_force_clamp(self):
+    m = test_util.load_test_file('actuator/tendon_force_clamp.xml')
+    d = mujoco.MjData(m)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = dx.replace(ctrl=jp.array([1.0, 1.0, 1.0, -4.0, 1.0, -20.0, 5.0, -5.0]))
+    dx = mjx.forward(mx, dx)
+
+    _assert_eq(
+        dx.actuator_force,
+        jp.array([1.0, 1.0, 1.0, -4.0 / 3.0, 1.0 / 3.0, -10.0, 5.0, -5.0]),
+        'actuator_force',
+    )
+
+    _assert_eq(
+        dx.sensordata,
+        jp.array([3.0, -1.0, -10.0, 0.0]),
+        'sensordata',
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/inverse.py

@@ -0,0 +1,106 @@
+# Copyright 2025 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Inverse dynamics functions."""
+
+from jax import numpy as jp
+from mujoco.mjx._src import derivative
+from mujoco.mjx._src import forward
+from mujoco.mjx._src import sensor
+from mujoco.mjx._src import smooth
+from mujoco.mjx._src import solver
+from mujoco.mjx._src import support
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import EnableBit
+from mujoco.mjx._src.types import IntegratorType
+from mujoco.mjx._src.types import Model
+
+
+def discrete_acc(m: Model, d: Data) -> Data:
+  """Convert discrete-time qacc to continuous-time qacc."""
+
+  if m.opt.integrator == IntegratorType.RK4:
+    raise RuntimeError(
+        'discrete inverse dynamics is not supported by RK4 integrator'
+    )
+  elif m.opt.integrator == IntegratorType.EULER:
+    dsbl_eulerdamp = m.opt.disableflags & DisableBit.EULERDAMP
+    no_dof_damping = (m.dof_damping == 0).all()
+    if dsbl_eulerdamp or no_dof_damping:
+      return d
+
+    # set qfrc = (M + h*diag(B)) * qacc
+    qfrc = support.mul_m(m, d, d.qacc)
+    qfrc += m.opt.timestep * m.dof_damping * d.qacc
+  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+    qm = support.full_m(m, d)
+
+    # compute analytical derivative qDeriv; skip rne derivative
+    qderiv = derivative.deriv_smooth_vel(m, d)
+    if qderiv is not None:
+      # M = M - dt*qDeriv
+      qm -= m.opt.timestep * qderiv
+
+    # set qfrc = (M - dt*qDeriv) * qacc
+    qfrc = qm @ d.qacc
+  else:
+    raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
+
+  # solve for qacc: qfrc = M * qacc
+  qacc = smooth.solve_m(m, d, qfrc)
+
+  return d.replace(qacc=qacc)
+
+
+def inv_constraint(m: Model, d: Data) -> Data:
+  """Inverse constraint solver."""
+
+  # no constraints
+  if d._impl.efc_J.size == 0:  # pytype: disable=attribute-error
+    return d.replace(qfrc_constraint=jp.zeros(m.nv))
+
+  # update
+  ctx = solver.Context.create(m, d, grad=False)
+
+  return d.tree_replace({
+      'qfrc_constraint': ctx.qfrc_constraint,
+      '_impl.efc_force': ctx.efc_force,
+  })
+
+
+def inverse(m: Model, d: Data) -> Data:
+  """Inverse dynamics."""
+  d = forward.fwd_position(m, d)
+  d = sensor.sensor_pos(m, d)
+  d = forward.fwd_velocity(m, d)
+  d = sensor.sensor_vel(m, d)
+
+  qacc = d.qacc
+  if m.opt.enableflags & EnableBit.INVDISCRETE:
+    d = discrete_acc(m, d)
+
+  d = inv_constraint(m, d)
+  d = smooth.rne(m, d, flg_acc=True)
+  d = sensor.sensor_acc(m, d)
+
+  qfrc_inverse = (
+      d.qfrc_bias + m.dof_armature * d.qacc - d.qfrc_passive - d.qfrc_constraint
+  )
+
+  if m.opt.enableflags & EnableBit.INVDISCRETE:
+    return d.replace(qfrc_inverse=qfrc_inverse, qacc=qacc)
+  else:
+    return d.replace(qfrc_inverse=qfrc_inverse)

mjx/mujoco/mjx/_src/inverse_test.py

@@ -0,0 +1,114 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for inverse dynamics functions."""
+from absl.testing import absltest
+from absl.testing import parameterized
+from jax import numpy as jp
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import support
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX calculations - mostly
+# due to float precision
+_TOLERANCE = 1e-5
+
+
+def _assert_eq(a, b, name, tol=_TOLERANCE):
+  tol = tol * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class InverseTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      (mujoco.mjtIntegrator.mjINT_EULER, False, False),
+      (mujoco.mjtIntegrator.mjINT_EULER, False, True),
+      (mujoco.mjtIntegrator.mjINT_EULER, True, False),
+      (mujoco.mjtIntegrator.mjINT_EULER, True, True),
+      (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, False, False),
+      (mujoco.mjtIntegrator.mjINT_IMPLICITFAST, True, False),
+  )
+  def test_forward_inverse_match(self, integrator, invdiscrete, eulerdamp):
+    m = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <option timestep=".005" gravity="-1 -1 -10"/>
+        <worldbody>
+          <geom type="plane" size="10 10 .001"/>
+          <body pos="0 0 1">
+            <geom type="sphere" size=".1" pos=".1 .2 .3"/>
+            <joint name="jnt1" type="hinge" axis="0 1 0" stiffness=".25" damping=".125"/>
+            <body pos="0 0 1">
+              <geom type="sphere" size=".1" pos=".1 .2 .3"/>
+              <joint name="jnt2" type="hinge" axis="0 1 0" stiffness=".6" damping=".3"/>
+            </body>
+          </body>
+        </worldbody>
+        <actuator>
+          <motor joint="jnt1"/>
+        </actuator>
+        <equality>
+          <joint joint1="jnt1" joint2="jnt2"/>
+        </equality>
+      </mujoco>
+    """)
+    m.opt.integrator = integrator
+    if invdiscrete:
+      m.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_INVDISCRETE
+    if not eulerdamp:
+      m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_EULERDAMP
+
+    d = mujoco.MjData(m)
+    d.qvel = np.random.uniform(low=-0.01, high=0.01, size=d.qvel.shape)
+    d.ctrl = np.random.uniform(low=-0.01, high=0.01, size=d.ctrl.shape)
+    d.qfrc_applied = np.random.uniform(
+        low=-0.01, high=0.01, size=d.qfrc_applied.shape
+    )
+    d.xfrc_applied = np.random.uniform(
+        low=-0.01, high=0.01, size=d.xfrc_applied.shape
+    )
+    mujoco.mj_step(m, d, 100)
+
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    dx_next = mjx.step(mx, dx)
+    qacc_fd = (dx_next.qvel - dx.qvel) / mx.opt.timestep
+
+    dx = mjx.forward(mx, dx)
+
+    if invdiscrete:
+      dx = dx.replace(qacc=qacc_fd)
+
+    dxinv = mjx.inverse(mx, dx)
+
+    fwdinv0 = jp.linalg.norm(
+        dxinv.qfrc_constraint - dx.qfrc_constraint, ord=np.inf
+    )
+    fwdinv1 = jp.linalg.norm(
+        dxinv.qfrc_inverse
+        - (
+            dx.qfrc_applied + dx.qfrc_actuator + support.xfrc_accumulate(mx, dx)
+        ),
+        ord=np.inf,
+    )
+
+    self.assertLess(fwdinv0, 1.0e-3)
+    self.assertLess(fwdinv1, 1.0e-3)
+    _assert_eq(dxinv.qacc, dx.qacc, 'qacc')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/io_test.py

@@ -14,14 +14,19 @@
 # ==============================================================================
 """Tests for io functions."""
 
+import os
+from unittest import mock
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import io as mjx_io
 from mujoco.mjx._src import test_util
+
 # pylint: disable=g-importing-member
+from mujoco.mjx._src.types import BackendImpl
 from mujoco.mjx._src.types import ConeType
 # pylint: enable=g-importing-member
 import numpy as np
@@ -91,14 +96,29 @@ _MULTIPLE_CONSTRAINTS = """
   </mujoco>
 """
 
+_SIMPLE_BODY = """
+  <mujoco>
+    <worldbody>
+      <!-- nC < nM because this body has an inertia-aligned free joint -->
+      <body name="simplebody">
+        <freejoint/>
+        <geom type="sphere" size="0.01"/>
+      </body>
+    </worldbody>
+  </mujoco>
+"""
+
 
 class ModelIOTest(parameterized.TestCase):
   """IO tests for mjx.Model."""
 
-  @parameterized.parameters(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS)
+  @parameterized.product(
-  def test_put_model(self, xml):
+      xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
+      backend_impl=('jax', 'c'),
+  )
+  def test_put_model(self, xml, backend_impl):
     m = mujoco.MjModel.from_xml_string(xml)
-    mx = mjx.put_model(m)
+    mx = mjx.put_model(m, backend_impl=backend_impl)
 
     def assert_not_weak_type(x):
       if isinstance(x, jax.Array):
@@ -120,8 +140,14 @@ class ModelIOTest(parameterized.TestCase):
     self.assertEqual(mx.nM, m.nM)
     self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
 
+    if backend_impl == 'jax':
       # fields restricted to MuJoCo should not be populated
-    self.assertEqual(mx.bvh_aabb.shape, (0,))
+      self.assertFalse(hasattr(mx, 'bvh_aabb'))
+    elif backend_impl == 'c':
+      # Options specific to C are populated.
+      self.assertEqual(mx.opt.apirate, m.opt.apirate)
+      # Fields private to C backend impl are populated.
+      self.assertTrue(hasattr(mx._impl, 'bvh_aabb'))
 
     np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
     np.testing.assert_allclose(mx.geom_type, m.geom_type)
@@ -150,7 +176,8 @@ class ModelIOTest(parameterized.TestCase):
     m = mjx.put_model(
         mujoco.MjModel.from_xml_string(
             '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
-        )
+        ),
+        backend_impl='jax',
     )
     self.assertTrue(m.opt.has_fluid_params)
 
@@ -174,7 +201,8 @@ class ModelIOTest(parameterized.TestCase):
 
   def test_margin_gap_mesh_not_implemented(self):
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string("""
         <mujoco>
           <asset>
             <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 1"/>
@@ -189,27 +217,32 @@ class ModelIOTest(parameterized.TestCase):
               <geom size="0.05"/>
             </body>
           </worldbody>
-        </mujoco>"""))
+        </mujoco>"""),
+          backend_impl='jax',
+      )
 
   def test_implicitfast_fluid_not_implemented(self):
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string("""
         <mujoco>
           <option viscosity="3.0" integrator="implicitfast"/>
           <worldbody/>
-        </mujoco>"""))
+        </mujoco>"""),
+          backend_impl='jax',
+      )
 
   def test_wrap_inside(self):
     m = test_util.load_test_file('tendon/wrap_sidesite.xml')
-    mx0 = mjx.put_model(m)
+    mx0 = mjx.put_model(m, backend_impl='jax')
     np.testing.assert_equal(
-        mx0.is_wrap_inside,
+        mx0._impl.is_wrap_inside,
         np.array([1, 0, 1, 0, 1, 1, 0]),
     )
     m.site_pos[2] = m.site_pos[1]
-    mx1 = mjx.put_model(m)
+    mx1 = mjx.put_model(m, backend_impl='jax')
     np.testing.assert_equal(
-        mx1.is_wrap_inside,
+        mx1._impl.is_wrap_inside,
         np.array([0, 0, 1, 0, 1, 0, 0]),
     )
 
@@ -217,11 +250,11 @@ class ModelIOTest(parameterized.TestCase):
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""
 
-  def test_make_data(self):
+  @parameterized.parameters('jax', 'c')
+  def test_make_data(self, backend_impl: str):
     """Test that make_data returns the correct shapes."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
-    d = mjx.make_data(m)
+    d = mjx.make_data(m, backend_impl=backend_impl)
 
     nq = 22
     nbody = 5
@@ -230,7 +263,7 @@ class DataIOTest(parameterized.TestCase):
     nv = 19
     nefc = 185
 
-    self.assertEqual(d.nefc, nefc)
+    self.assertEqual(d._impl.nefc, nefc)
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
     self.assertEqual(d.act.shape, (0,))
@@ -251,57 +284,68 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.geom_xpos.shape, (6, 3))
     self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
     self.assertEqual(d.subtree_com.shape, (nbody, 3))
-    self.assertEqual(d.cdof.shape, (nv, 6))
+    self.assertEqual(d._impl.cdof.shape, (nv, 6))
-    self.assertEqual(d.cinert.shape, (nbody, 10))
+    self.assertEqual(d._impl.cinert.shape, (nbody, 10))
-    self.assertEqual(d.crb.shape, (nbody, 10))
+    self.assertEqual(d._impl.crb.shape, (nbody, 10))
-    self.assertEqual(d.actuator_length.shape, (1,))
+    self.assertEqual(d._impl.actuator_length.shape, (1,))
-    self.assertEqual(d.actuator_moment.shape, (1, nv))
+    self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
-    self.assertEqual(d.qM.shape, (nv, nv))
+    self.assertEqual(d._impl.contact.dist.shape, (ncon,))
-    self.assertEqual(d.qLD.shape, (nv, nv))
+    self.assertEqual(d._impl.contact.pos.shape, (ncon, 3))
-    self.assertEqual(d.qLDiagInv.shape, (0,))
+    self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3))
-    self.assertEqual(d.contact.dist.shape, (ncon,))
+    self.assertEqual(d._impl.contact.solref.shape, (ncon, 2))
-    self.assertEqual(d.contact.pos.shape, (ncon, 3))
+    self.assertEqual(d._impl.contact.solimp.shape, (ncon, 5))
-    self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
+    self.assertEqual(d._impl.contact.geom1.shape, (ncon,))
-    self.assertEqual(d.contact.solref.shape, (ncon, 2))
+    self.assertEqual(d._impl.contact.geom2.shape, (ncon,))
-    self.assertEqual(d.contact.solimp.shape, (ncon, 5))
+    self.assertEqual(d._impl.efc_J.shape, (nefc, nv))
-    self.assertEqual(d.contact.geom1.shape, (ncon,))
+    self.assertEqual(d._impl.efc_frictionloss.shape, (nefc,))
-    self.assertEqual(d.contact.geom2.shape, (ncon,))
+    self.assertEqual(d._impl.efc_D.shape, (nefc,))
-    self.assertEqual(d.efc_J.shape, (nefc, nv))
+    self.assertEqual(d._impl.actuator_velocity.shape, (1,))
-    self.assertEqual(d.efc_frictionloss.shape, (nefc,))
-    self.assertEqual(d.efc_D.shape, (nefc,))
-    self.assertEqual(d.actuator_velocity.shape, (1,))
     self.assertEqual(d.cvel.shape, (nbody, 6))
-    self.assertEqual(d.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d._impl.cdof_dot.shape, (nv, 6))
     self.assertEqual(d.qfrc_bias.shape, (nv,))
     self.assertEqual(d.qfrc_passive.shape, (nv,))
-    self.assertEqual(d.efc_aref.shape, (nefc,))
+    self.assertEqual(d._impl.efc_aref.shape, (nefc,))
     self.assertEqual(d.qfrc_actuator.shape, (nv,))
     self.assertEqual(d.qfrc_smooth.shape, (nv,))
     self.assertEqual(d.qacc_smooth.shape, (nv,))
     self.assertEqual(d.qfrc_constraint.shape, (nv,))
     self.assertEqual(d.qfrc_inverse.shape, (nv,))
-    self.assertEqual(d.efc_force.shape, (nefc,))
+    self.assertEqual(d._impl.efc_force.shape, (nefc,))
+
+    if backend_impl == 'jax':
+      self.assertEqual(d._impl.qM.shape, (nv, nv))
+      self.assertEqual(d._impl.qLD.shape, (nv, nv))
+      self.assertEqual(d._impl.qLDiagInv.shape, (0,))
+    elif backend_impl == 'c':
+      self.assertEqual(d._impl.qM.shape, (nm,))
+      self.assertEqual(d._impl.qLD.shape, (nm,))
+      self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
     # test sparse
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
-    d = mjx.make_data(m)
+    d = mjx.make_data(m, backend_impl=backend_impl)
-    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d._impl.qM.shape, (nm,))
-    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d._impl.qLD.shape, (nm,))
-    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
-  def test_put_data(self):
+    if backend_impl == 'c':
+      # check C specific fields
+      self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3))
+      self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,))
+
+  @parameterized.parameters('jax', 'c')
+  def test_put_data(self, backend_impl: str):
     """Test that put_data puts the correct data for dense and sparse."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
 
     # check a few fields
     np.testing.assert_allclose(dx.qpos, d.qpos)
     np.testing.assert_allclose(dx.xpos, d.xpos)
     np.testing.assert_allclose(dx.cvel, d.cvel)
-    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot)
 
     # check that there are no weak types
     self.assertFalse(
@@ -312,21 +356,26 @@ class DataIOTest(parameterized.TestCase):
         )
     )
 
+    if backend_impl == 'jax':
       # check that qM is transformed properly
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mj_fullM(m, qm, d.qM)
       np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
+    elif backend_impl == 'c':
+      np.testing.assert_allclose(dx._impl.qM, d.qM)
+      np.testing.assert_allclose(dx._impl.qLD, d.qLD)
+      np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv)
 
     # 4 contacts, 2 for each capsule against the plane
-    self.assertEqual(dx.contact.dist.shape, (4,))
+    self.assertEqual(dx._impl.contact.dist.shape, (4,))
     self.assertEqual(d.ncon, 1)  # however only 1 contact in this step
-    np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
+    np.testing.assert_allclose(dx._impl.contact.dist[0], d.contact.dist[0])
-    self.assertTrue((dx.contact.dist[1:] > 0).all())
+    self.assertTrue((dx._impl.contact.dist[1:] > 0).all())
-    self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
+    self.assertEqual(dx._impl.contact.frame.shape, (4, 3, 3))
     np.testing.assert_allclose(
-        dx.contact.frame[0].reshape(9), d.contact.frame[0]
+        dx._impl.contact.frame[0].reshape(9), d.contact.frame[0]
     )
-    np.testing.assert_allclose(dx.contact.frame[1:], 0)
+    np.testing.assert_allclose(dx._impl.contact.frame[1:], 0)
 
     # xmat, ximat, geom_xmat are all shape transformed
     self.assertEqual(dx.xmat.shape, (3, 3, 3))
@@ -339,57 +388,68 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
 
     # tendon data is correct
-    np.testing.assert_allclose(dx.ten_length, d.ten_length)
+    np.testing.assert_allclose(dx._impl.ten_length, d.ten_length)
-    np.testing.assert_equal(dx.ten_wrapadr, np.zeros((1,)))
+    np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
-    np.testing.assert_equal(dx.ten_wrapnum, np.zeros((1,)))
+    np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
-    np.testing.assert_equal(dx.wrap_obj, np.zeros((2, 2)))
+    np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
-    np.testing.assert_equal(dx.wrap_xpos, np.zeros((2, 6)))
+    np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6)))
 
     # efc_ are also shape transformed and padded
-    self.assertEqual(dx.efc_J.shape, (45, 8))  # nefc, nv
+    self.assertEqual(dx._impl.efc_J.shape, (45, 8))  # nefc, nv
     d_efc_j = d.efc_J.reshape((-1, 8))
-    np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3])  # connect eq
+    np.testing.assert_allclose(dx._impl.efc_J[:3], d_efc_j[:3])  # connect eq
-    np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3])  # one active limit
+    np.testing.assert_allclose(
-    np.testing.assert_allclose(dx.efc_J[4], 0)  # one inactive limit
+        dx._impl.efc_J[3], d_efc_j[3]
-    np.testing.assert_allclose(dx.efc_J[5:15], d_efc_j[4:14])  # contact
+    )  # one active limit
-    np.testing.assert_allclose(dx.efc_J[15:], 0)  # no contact
+    np.testing.assert_allclose(dx._impl.efc_J[4], 0)  # one inactive limit
+    np.testing.assert_allclose(dx._impl.efc_J[5:15], d_efc_j[4:14])  # contact
+    np.testing.assert_allclose(dx._impl.efc_J[15:], 0)  # no contact
 
     # check another efc_ too
-    self.assertEqual(dx.efc_aref.shape, (45,))  # nefc
+    self.assertEqual(dx._impl.efc_aref.shape, (45,))  # nefc
-    np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
+    np.testing.assert_allclose(dx._impl.efc_aref[:3], d.efc_aref[:3])
-    np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
+    np.testing.assert_allclose(dx._impl.efc_aref[3], d.efc_aref[3])
-    np.testing.assert_allclose(dx.efc_aref[4], 0)
+    np.testing.assert_allclose(dx._impl.efc_aref[4], 0)
-    np.testing.assert_allclose(dx.efc_aref[5:15], d.efc_aref[4:14])
+    np.testing.assert_allclose(dx._impl.efc_aref[5:15], d.efc_aref[4:14])
-    np.testing.assert_allclose(dx.efc_aref[15:], 0)
+    np.testing.assert_allclose(dx._impl.efc_aref[15:], 0)
 
     # check sparse transform is correct
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_sparse = mjx.put_data(m, d)
+    dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl)
-    np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8)
 
     # check sparse mass matrices are correct
-    np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
-    np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8)
-    np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
+    np.testing.assert_allclose(
+        dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8
+    )
 
     # check dense mass matrices are correct
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_from_dense = mjx.put_data(m, d)
+    dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl)
+    if backend_impl == 'jax':
       qm = np.zeros((m.nv, m.nv))
       mujoco.mj_fullM(m, qm, d.qM)
-    np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
+      np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
+    elif backend_impl == 'c':
+      np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
 
-  def test_get_data(self):
+  @parameterized.parameters(
+      ('jax', False), ('jax', True), ('c', False), ('c', True)
+  )
+  def test_get_data(self, backend_impl: str, sparse: bool):
     """Test that get_data makes correct MjData."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    if sparse:
+      m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     d_2: mujoco.MjData = mjx.get_data(m, dx)
 
     # check a few fields
@@ -398,6 +458,8 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(d_2.cvel, d.cvel)
     np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
     np.testing.assert_allclose(d_2.qM, d.qM)
+    np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6)
+    np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6)
 
     # only 1 contact active
     self.assertEqual(d_2.contact.dist.shape, (1,))
@@ -418,12 +480,45 @@ class DataIOTest(parameterized.TestCase):
 
     # efc_* are also shape transformed and filtered
     self.assertEqual(d_2.nefc, 14)
+    if sparse:
+      efc_j = np.zeros((d.nefc, m.nv))
+      mujoco.mju_sparse2dense(
+          efc_j,
+          d.efc_J,
+          d.efc_J_rownnz,
+          d.efc_J_rowadr,
+          d.efc_J_colind,
+      )
+      efc_j2 = np.zeros((d_2.nefc, m.nv))
+      mujoco.mju_sparse2dense(
+          efc_j2,
+          d_2.efc_J,
+          d_2.efc_J_rownnz,
+          d_2.efc_J_rowadr,
+          d_2.efc_J_colind,
+      )
+      np.testing.assert_allclose(efc_j, efc_j2)
+    else:
       self.assertEqual(d_2.efc_J.shape, (112,))  # nefc * nv
       np.testing.assert_allclose(d_2.efc_J, d.efc_J)
     self.assertEqual(d_2.efc_aref.shape, (14,))  # nefc
     np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
     np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
 
+    if backend_impl == 'c':
+      # check fields specific to the C implementation
+      np.testing.assert_allclose(d_2.bvh_active, d.bvh_active)
+
+  def test_get_data_simplebody(self):
+    """Test that get_data works with simple bodies where nC < nM."""
+    m = mujoco.MjModel.from_xml_string(_SIMPLE_BODY)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2: mujoco.MjData = mjx.get_data(m, dx)
+    np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6)
+    np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6)
+
   def test_get_data_runs(self):
     xml = """
       <mujoco>
@@ -444,13 +539,14 @@ class DataIOTest(parameterized.TestCase):
     dx = mjx.put_data(m, d)
     mjx.get_data(m, dx)
 
-  def test_get_data_batched(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_batched(self, backend_impl):
     """Test that get_data makes correct List[MjData] for batched Data."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     # second data in batch has contact dist > 0, disables contact
     dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
     ds = mjx.get_data(m, dx_b)
@@ -460,13 +556,14 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(ds[0].ncon, 1)
     self.assertEqual(ds[1].ncon, 0)
 
-  def test_get_data_into(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_into(self, backend_impl):
     """Test that get_data_into correctly populates an MjData."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     d_2 = mujoco.MjData(m)
     mjx.get_data_into(d_2, m, dx)
 
@@ -482,32 +579,33 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d_2.contact.frame.shape, (1, 9))
     np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
 
-  def test_get_data_into_wrong_shape(self):
+  @parameterized.parameters('jax', 'c')
+  def test_get_data_into_wrong_shape(self, backend_impl):
     """Tests that get_data_into throwsif input and output shapes don't match."""
 
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
     m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
     d_2 = mujoco.MjData(m_2)
     with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'):
       mjx.get_data_into(d_2, m, dx)
 
-  def test_make_matches_put(self):
+  @parameterized.parameters('jax', 'c')
+  def test_make_matches_put(self, backend_impl):
     """Test that make_data produces a pytree that matches put_data."""
-
     m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx = mjx.put_data(m, d)
+    dx = mjx.put_data(m, d, backend_impl=backend_impl)
 
     step_fn = lambda d: d.replace(time=d.time + 1)
     step_fn_jit = jax.jit(step_fn).lower(dx).compile()
 
     # placing an MjData onto device should yield the same treedef mjx.Data as
     # calling make_data.  they should be interchangeable for jax functions:
-    step_fn_jit(mjx.make_data(m))
+    step_fn_jit(mjx.make_data(m, backend_impl=backend_impl))
 
   def test_contact_elliptic_condim1(self):
     """Test that condim=1 with ConeType.ELLIPTIC is not implemented."""
@@ -556,7 +654,327 @@ class DataIOTest(parameterized.TestCase):
         </mujoco>
       """)
     with self.assertRaises(NotImplementedError):
-      mjx.put_model(m)
+      mjx.put_model(m, backend_impl='jax')
+
+
+class FullCompatTest(parameterized.TestCase):
+  """Tests for the _full_compat flag."""
+
+  def test_full_compat_deprecated(self):
+    """Tests that _full_compat is deprecated."""
+    xml = """
+      <mujoco>
+        <worldbody>
+          <body name="box">
+            <joint name="slide1" type="slide" axis="1 0 0" />
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+          </body>
+        </worldbody>
+        <actuator>
+          <motor joint="slide1"/>
+        </actuator>
+      </mujoco>
+    """
+    m = mujoco.MjModel.from_xml_string(xml)
+    with self.assertWarns(DeprecationWarning):
+      out = mjx_io.put_model(m, _full_compat=True)
+      self.assertEqual(out.backend_impl, BackendImpl.C)
+    with self.assertWarns(DeprecationWarning):
+      out = mjx_io.make_data(m, _full_compat=True)
+      self.assertEqual(out.backend_impl, BackendImpl.C)
+
+
+# Test cases for `_resolve_backend_impl_and_device` where the device is
+# specified by the user and the device is available.
+_DEVICE_TEST_CASES = [
+    # Arguments use the following format:
+    # (device_type_str, backend_impl_str,
+    #  (expected_device, expected_backend_impl)))
+    # No backend specified.
+    ('cpu', None, ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
+    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    # JAX backend specified.
+    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    # WARP backend specified.
+    ('cpu', 'warp', ('cpu', 'error')),
+    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('tpu', 'warp', ('tpu', 'error')),
+    # C backend specified.
+    ('cpu', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', 'c', ('cpu', 'error')),
+    ('gpu-nvidia', 'c', ('cpu', 'error')),
+    ('tpu', 'c', ('tpu', 'error')),
+]
+
+# Test cases for `_resolve_backend_impl_and_device` where the user does NOT
+# specify a device. We mock the JAX default device.
+_DEFAULT_DEVICE_TEST_CASES = [
+    # Arguments use the following format:
+    # (jax.default_device, backend_impl_str,
+    #  (expected_device, expected_backend_impl))
+    # No backend impl specified.
+    ('cpu', None, ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
+    ('tpu', None, ('tpu', BackendImpl.JAX)),
+    # JAX backend impl specified.
+    ('cpu', 'jax', ('cpu', BackendImpl.JAX)),
+    ('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
+    ('tpu', 'jax', ('tpu', BackendImpl.JAX)),
+    # WARP backend impl specified.
+    ('cpu', 'warp', ('cpu', 'error')),
+    ('gpu-notnvidia', 'warp', ('cpu', 'error')),
+    ('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
+    ('tpu', 'warp', ('tpu', 'error')),
+    # C backend impl specified, CPU should always be available.
+    ('cpu', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)),
+    ('gpu-nvidia', 'c', ('cpu', BackendImpl.C)),
+    ('tpu', 'c', ('cpu', BackendImpl.C)),
+]
+
+
+class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
+  """Tests for the _resolve_backend_impl_and_device function."""
+
+  def setUp(self):
+    super().setUp()
+
+    # Create mock devices
+    self.mock_cpu = mock.Mock(spec=jax.Device)
+    self.mock_cpu.platform = 'cpu'
+    self.mock_cpu.device_kind = 'Mock CPU'
+    self.mock_cpu.id = 0
+
+    self.mock_nvidia_gpu = mock.Mock(spec=jax.Device)
+    self.mock_nvidia_gpu.platform = 'gpu'
+    self.mock_nvidia_gpu.device_kind = 'NVIDIA Mocked GPU'
+    self.mock_nvidia_gpu.id = 0
+
+    self.mock_other_gpu = mock.Mock(spec=jax.Device)
+    self.mock_other_gpu.platform = 'gpu'
+    self.mock_other_gpu.device_kind = 'Other Mocked GPU'
+    self.mock_other_gpu.id = 1
+
+    self.mock_tpu = mock.Mock(spec=jax.Device)
+    self.mock_tpu.platform = 'tpu'
+    self.mock_tpu.device_kind = 'Mock TPU'
+    self.mock_tpu.id = 0
+
+    # Patch jax.devices for the entire test class using enter_context
+    self.mock_jax_devices = self.enter_context(mock.patch('jax.devices'))
+    self.mock_default_backend = self.enter_context(
+        mock.patch('jax.default_backend')
+    )
+
+  @parameterized.named_parameters(
+      (f'{str(args[0])}_{str(args[1])}', *args) for args in _DEVICE_TEST_CASES
+  )
+  @mock.patch.dict(
+      os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
+  )
+  def test_resolve_with_device(
+      self,
+      device_type_str,
+      backend_impl_str,
+      expected,
+  ):
+    """Tests various combinations of device and backend impls."""
+    input_device = {
+        'cpu': self.mock_cpu,
+        'gpu-nvidia': self.mock_nvidia_gpu,
+        'gpu-notnvidia': self.mock_other_gpu,
+        'tpu': self.mock_tpu,
+    }[device_type_str]
+
+    def devices_side_effect(backend=None):
+      # assume the user-specified device is always available
+      if backend == 'cpu':
+        return [self.mock_cpu]
+      elif backend == 'gpu':
+        if 'nvidia' in device_type_str:
+          return [self.mock_nvidia_gpu]
+        return [self.mock_other_gpu]
+      elif backend == 'tpu':
+        return [self.mock_tpu]
+      elif backend == 'cuda':
+        return [self.mock_nvidia_gpu]
+
+      raise AssertionError('Should not be called.')
+
+    self.mock_jax_devices.side_effect = devices_side_effect
+
+    expected_device, expected_backend_impl = expected
+    if expected_backend_impl == 'error':
+      with self.assertRaises(AssertionError):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=input_device
+        )
+      return
+
+    actual_backend_impl, actual_device = (
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=input_device
+        )
+    )
+
+    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertIsNotNone(actual_device)
+    self.assertEqual(actual_device.platform, expected_device)
+
+  @parameterized.named_parameters(
+      (f'{str(args[0])}_{str(args[1])}', *args)
+      for args in _DEFAULT_DEVICE_TEST_CASES
+  )
+  @mock.patch.dict(
+      os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
+  )
+  def test_resolve_without_device(
+      self,
+      default_device_str,
+      backend_impl_str,
+      expected,
+  ):
+    """Tests various combinations of jax.default_device and backend impls."""
+    default_devices = {
+        'cpu': [self.mock_cpu],
+        'gpu-nvidia': [self.mock_nvidia_gpu, self.mock_cpu],
+        'gpu-notnvidia': [self.mock_other_gpu, self.mock_cpu],
+        'tpu': [self.mock_tpu, self.mock_cpu],
+    }[default_device_str]
+
+    def devices_side_effect(backend=None):
+      if backend == 'cpu':
+        return [self.mock_cpu]  # CPU is always available
+      if backend == 'gpu' and default_device_str == 'gpu-notnvidia':
+        return [self.mock_other_gpu]
+      if backend == 'gpu' and default_device_str == 'gpu-nvidia':
+        return [self.mock_nvidia_gpu]
+      if backend == 'cuda' and default_device_str == 'gpu-nvidia':
+        return [self.mock_nvidia_gpu]
+      if backend == 'tpu' and default_device_str == 'tpu':
+        return [self.mock_tpu]
+      if backend is None:
+        return default_devices
+      if backend == 'cuda':
+        raise RuntimeError('cuda backend not supported')
+      raise AssertionError('jax.devices error')
+
+    self.mock_jax_devices.side_effect = devices_side_effect
+    default_device_side_effect_str = {
+        'cpu': 'cpu',
+        'gpu-nvidia': 'gpu',
+        'gpu-notnvidia': 'gpu',
+        'tpu': 'tpu',
+    }[default_device_str]
+    self.mock_default_backend.side_effect = (
+        lambda: default_device_side_effect_str
+    )
+
+    expected_device, expected_backend_impl = expected
+    if (
+        expected_backend_impl == 'error'
+        and default_device_str != 'gpu-nvidia'
+        and backend_impl_str == 'warp'
+    ):
+      with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+      return
+
+    if expected_backend_impl == 'error':
+      with self.assertRaises(AssertionError):
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+      return
+
+    actual_backend_impl, actual_device = (
+        mjx_io._resolve_backend_impl_and_device(
+            backend_impl=backend_impl_str, device=None
+        )
+    )
+
+    self.assertEqual(actual_backend_impl, expected_backend_impl)
+    self.assertIsNotNone(actual_device)
+    self.assertEqual(actual_device.platform, expected_device)
+
+  @mock.patch.dict(os.environ, {'MJX_WARP_ENABLED': 'false'})
+  def test_resolve_warp_disabled(self):
+    """Tests behavior when MJX_WARP_ENABLED is false."""
+    self.mock_jax_devices.side_effect = lambda backend=None: (
+        [self.mock_nvidia_gpu, self.mock_cpu]
+        if backend is None
+        else ([self.mock_nvidia_gpu] if backend == 'gpu' else [self.mock_cpu])
+    )
+    self.mock_default_backend.side_effect = lambda: 'gpu'
+
+    # Default to JAX instead of WARP on NVIDIA GPU.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'gpu')
+
+    # Specifying an NVIDIA GPU should still choose JAX.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=self.mock_nvidia_gpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'gpu')
+
+    # Requesting warp explicitly should fail since it is disabled.
+    with self.assertRaises(AssertionError):
+      mjx_io._resolve_backend_impl_and_device(
+          backend_impl='warp', device=self.mock_nvidia_gpu
+      )
+    with self.assertRaises(AssertionError):
+      mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None)
+
+  @mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'})
+  def test_resolve_c_disabled(self):
+    """Tests behavior when MJX_C_DEFAULT_ENABLED is false."""
+    # Users expect that CPU defaults to the JAX impl. But in the future, it will
+    # default to the C backend implementation. This test checks that
+    # MJX_C_DEFAULT_ENABLED=false defaults to the old behavior, until the
+    # migration to MJEP-15 is complete.
+    self.mock_jax_devices.side_effect = lambda backend=None: ([self.mock_cpu])
+    self.mock_default_backend.side_effect = lambda: 'cpu'
+
+    # Default to JAX instead of C on CPU.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'cpu')
+
+    # Specifing CPU should still choose JAX.
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl=None, device=self.mock_cpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.JAX)
+    self.assertEqual(device.platform, 'cpu')
+
+    # Specifying C should choose C!
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl='c', device=None
+    )
+    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(device.platform, 'cpu')
+
+    backend_impl, device = mjx_io._resolve_backend_impl_and_device(
+        backend_impl='c', device=self.mock_cpu
+    )
+    self.assertEqual(backend_impl, BackendImpl.C)
+    self.assertEqual(device.platform, 'cpu')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/passive.py

@@ -23,9 +23,12 @@ from mujoco.mjx._src import scan
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
+from mujoco.mjx._src.types import OptionJAX
 # pylint: enable=g-importing-member
 
 
@@ -72,11 +75,11 @@ def _spring_damper(m: Model, d: Data) -> jax.Array:
   qfrc -= m.dof_damping * d.qvel
 
   # tendon-level spring-dampers
-  below, above = m.tendon_lengthspring.T - d.ten_length
+  below, above = m.tendon_lengthspring.T - d._impl.ten_length
   frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
   frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
-  frc_damper = -m.tendon_damping * d.ten_velocity
+  frc_damper = -m.tendon_damping * d._impl.ten_velocity
-  qfrc += d.ten_J.T @ (frc_spring + frc_damper)
+  qfrc += d._impl.ten_J.T @ (frc_spring + frc_damper)
 
   return qfrc
 
@@ -113,6 +116,9 @@ def _fluid(m: Model, d: Data) -> jax.Array:
 
 def passive(m: Model, d: Data) -> Data:
   """Adds all passive forces."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('passive requires JAX backend implementation.')
+
   if m.opt.disableflags & DisableBit.PASSIVE:
     return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv))
 
@@ -124,7 +130,7 @@ def passive(m: Model, d: Data) -> Data:
     # add gravcomp unless added via actuators
     qfrc_passive += qfrc_gravcomp * (1 - m.jnt_actgravcomp[m.dof_jntid])
 
-  if m.opt.has_fluid_params:
+  if m.opt.has_fluid_params:  # pytype: disable=attribute-error
     qfrc_passive += _fluid(m, d)
 
   d = d.replace(qfrc_passive=qfrc_passive, qfrc_gravcomp=qfrc_gravcomp)

mjx/mujoco/mjx/_src/scan.py

@@ -132,6 +132,8 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
 
 def _check_input(m: Model, args: Any, in_types: str) -> None:
   """Checks that scan input has the right shape."""
+  if m.nv == 0:
+    raise ValueError('Scan across Model with zero DoFs unsupported.')
   size = {
       'b': m.nbody,
       'j': m.njnt,

mjx/mujoco/mjx/_src/scan_test.py

@@ -61,10 +61,7 @@ class ScanTest(absltest.TestCase):
       return body_id + 1
 
     b_in = jp.array([1])
-    b_expect = jp.array([2])
+    self.assertRaises(ValueError, scan.flat, m, fn, 'b', 'b', b_in)
-    b_out = scan.flat(m, fn, 'b', 'b', b_in)
-
-    np.testing.assert_equal(np.array(b_out), np.array(b_expect))
 
   def test_flat_joints(self):
     """Tests scanning over bodies with joints of different types."""

mjx/mujoco/mjx/_src/sensor.py

@@ -22,16 +22,20 @@ from mujoco.mjx._src import math
 from mujoco.mjx._src import ray
 from mujoco.mjx._src import smooth
 from mujoco.mjx._src import support
+from mujoco.mjx._src.types import BackendImpl
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import ObjType
 from mujoco.mjx._src.types import SensorType
+from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
 import numpy as np
 
 
-def apply_cutoff(
+def _apply_cutoff(
     sensor: jax.Array, cutoff: jax.Array, data_type: int
 ) -> jax.Array:
   """Clip sensor to cutoff value."""
@@ -50,6 +54,8 @@ def apply_cutoff(
 
 def sensor_pos(m: Model, d: Data) -> Data:
   """Compute position-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_pos requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -162,15 +168,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
         sensor, _ = jax.vmap(
             ray.ray, in_axes=(None, None, 0, 0, None, None, None)
         )(m, d, site_xpos, site_mat, (), True, sid)
-        sensors.append(apply_cutoff(sensor, cutoffs, data_type[0]))
+        sensors.append(_apply_cutoff(sensor, cutoffs, data_type[0]))
         adrs.append(adr[idxs])
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.JOINTPOS:
       sensor = d.qpos[m.jnt_qposadr[objid]]
     elif sensor_type == SensorType.TENDONPOS:
-      sensor = d.ten_length[objid]
+      sensor = d._impl.ten_length[objid]
     elif sensor_type == SensorType.ACTUATORPOS:
-      sensor = d.actuator_length[objid]
+      sensor = d._impl.actuator_length[objid]
     elif sensor_type == SensorType.BALLQUAT:
       jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
       quat = d.qpos[jnt_qposadr]
@@ -193,7 +199,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
         cutofft = cutoff[idxt]
         sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt)
         adrt = adr[idxt, None] + np.arange(3)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type in frame_axis:
@@ -213,7 +219,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
         cutofft = cutoff[idxt]
         sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt)
         adrt = adr[idxt, None] + np.arange(3)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.FRAMEQUAT:
@@ -254,7 +260,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
             )
         )(quat, refquat, refidt)
         adrt = adr[idxt, None] + np.arange(4)[None]
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.SUBTREECOM:
@@ -266,7 +272,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:
@@ -281,6 +287,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
 
 def sensor_vel(m: Model, d: Data) -> Data:
   """Compute velocity-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_vel requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -332,9 +340,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.JOINTVEL:
       sensor = d.qvel[m.jnt_dofadr[objid]]
     elif sensor_type == SensorType.TENDONVEL:
-      sensor = d.ten_velocity[objid]
+      sensor = d._impl.ten_velocity[objid]
     elif sensor_type == SensorType.ACTUATORVEL:
-      sensor = d.actuator_velocity[objid]
+      sensor = d._impl.actuator_velocity[objid]
     elif sensor_type == SensorType.BALLANGVEL:
       jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
       sensor = d.qvel[jnt_dotadr]
@@ -392,20 +400,20 @@ def sensor_vel(m: Model, d: Data) -> Data:
 
         adrt = adr[idxt, None] + np.arange(3)[None]
 
-        sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
+        sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
         adrs.append(adrt.reshape(-1))
       continue  # avoid adding to sensors/adrs list a second time
     elif sensor_type == SensorType.SUBTREELINVEL:
-      sensor = d.subtree_linvel[objid]
+      sensor = d._impl.subtree_linvel[objid]
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.SUBTREEANGMOM:
-      sensor = d.subtree_angmom[objid]
+      sensor = d._impl.subtree_angmom[objid]
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     else:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:
@@ -420,6 +428,8 @@ def sensor_vel(m: Model, d: Data) -> Data:
 
 def sensor_acc(m: Model, d: Data) -> Data:
   """Compute acceleration/force-dependent sensors values."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('sensor_acc requires JAX backend implementation.')
 
   if m.opt.disableflags & DisableBit.SENSOR:
     return d
@@ -441,6 +451,8 @@ def sensor_acc(m: Model, d: Data) -> Data:
       SensorType.ACCELEROMETER,
       SensorType.FORCE,
       SensorType.TORQUE,
+      SensorType.FRAMELINACC,
+      SensorType.FRAMEANGACC,
   }:
     d = smooth.rne_postconstraint(m, d)
 
@@ -457,14 +469,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       # compute contact forces
       forces = []
       condim_ids = []
-      for dim in set(d.contact.dim):
+      for dim in set(d._impl.contact.dim):
         force, condim_id = support.contact_force_dim(m, d, dim)
         forces.append(force)
         condim_ids.append(condim_id)
       forces = jp.concatenate(forces)[np.argsort(np.concatenate(condim_ids))]
 
       # get bodies of contact geoms
-      conbody = jp.array(m.geom_bodyid)[d.contact.geom]
+      conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom]
 
       # get site information
       site_bodyid = m.site_bodyid[objid]
@@ -474,12 +486,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       site_type = m.site_type[objid]
       conbody0 = site_bodyid[:, None] == conbody[:, 0]
       conbody1 = site_bodyid[:, None] == conbody[:, 1]
-      contacts = (d.contact.efc_address >= 0)[None] & (conbody0 | conbody1)
+      contacts = (d._impl.contact.efc_address >= 0)[None] & (
+          conbody0 | conbody1
+      )
 
       # compute conray, flip if second body
       conray = jax.vmap(
           lambda frame, force: math.normalize(frame[0] * force[0])
-      )(d.contact.frame, forces)
+      )(d._impl.contact.frame, forces)
       conray = jp.where(conbody1[..., None], -conray, conray)
 
       # compute distance, mapping over sites and contacts
@@ -501,7 +515,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
             site_xpos[dist_id_site],
             site_xmat[dist_id_site],
             st,
-            d.contact.pos,
+            d._impl.contact.pos,
             conray[dist_id_site],
         )
         dist.append(jp.where(jp.isinf(dist_site), 0, dist_site))
@@ -523,14 +537,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
       bodyid = m.site_bodyid[objid]
       rot = d.site_xmat[objid]
       cvel = d.cvel[bodyid]
-      cacc = d.cacc[bodyid]
+      cacc = d._impl.cacc[bodyid]
       dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]]
 
       sensor = _accelerometer(cvel, cacc, dif, rot)
       adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
     elif sensor_type == SensorType.FORCE:
       bodyid = m.site_bodyid[objid]
-      cfrc_int = d.cfrc_int[bodyid]
+      cfrc_int = d._impl.cfrc_int[bodyid]
       site_xmat = d.site_xmat[objid]
       sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
           site_xmat, cfrc_int[:, 3:]
@@ -539,7 +553,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
     elif sensor_type == SensorType.TORQUE:
       bodyid = m.site_bodyid[objid]
       rootid = m.body_rootid[bodyid]
-      cfrc_int = d.cfrc_int[bodyid]
+      cfrc_int = d._impl.cfrc_int[bodyid]
       site_xmat = d.site_xmat[objid]
       dif = d.site_xpos[objid] - d.subtree_com[rootid]
       sensor = jax.vmap(
@@ -550,6 +564,15 @@ def sensor_acc(m: Model, d: Data) -> Data:
       sensor = d.actuator_force[objid]
     elif sensor_type == SensorType.JOINTACTFRC:
       sensor = d.qfrc_actuator[m.jnt_dofadr[objid]]
+    elif sensor_type == SensorType.TENDONACTFRC:
+      force_mask = [
+          (m.actuator_trntype == TrnType.TENDON)
+          & (m.actuator_trnid[:, 0] == tendon_id)
+          for tendon_id in objid
+      ]
+      force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask])
+      force_mat = np.array(force_mask)[:, force_ids]
+      sensor = force_mat @ d.actuator_force[force_ids]
     elif sensor_type in (SensorType.FRAMELINACC, SensorType.FRAMEANGACC):
       objtype = m.sensor_objtype[idx]
 
@@ -559,7 +582,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
         pos, bodyid = objtype_data[ot]
         pos = pos[objidt]
         bodyid = bodyid[objidt]
-        cacc = d.cacc[bodyid]
+        cacc = d._impl.cacc[bodyid]
 
         if sensor_type == SensorType.FRAMELINACC:
 
@@ -589,7 +612,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
       # TODO(taylorhowell): raise error after adding sensor check to io.py
       continue  # unsupported sensor type
 
-    sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
+    sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
     adrs.append(adr)
 
   if not adrs:

mjx/mujoco/mjx/_src/sensor_test.py

@@ -64,14 +64,14 @@ class SensorTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
 
     mx = mjx.put_model(m)
-    dx = mjx.put_data(m, d).replace(
+    dx = mjx.put_data(m, d).tree_replace({
-        sensordata=jp.zeros_like(d.sensordata),
+        'sensordata': jp.zeros_like(d.sensordata),
-        subtree_linvel=jp.zeros_like(d.subtree_linvel),
+        '_impl.subtree_linvel': jp.zeros_like(d.subtree_linvel),
-        subtree_angmom=jp.zeros_like(d.subtree_angmom),
+        '_impl.subtree_angmom': jp.zeros_like(d.subtree_angmom),
-        cacc=jp.zeros_like(d.cacc),
+        '_impl.cacc': jp.zeros_like(d.cacc),
-        cfrc_int=jp.zeros_like(d.cfrc_int),
+        '_impl.cfrc_int': jp.zeros_like(d.cfrc_int),
-        cfrc_ext=jp.zeros_like(d.cfrc_ext),
+        '_impl.cfrc_ext': jp.zeros_like(d.cfrc_ext),
-    )
+    })
     dx = jax.jit(mjx.sensor_pos)(mx, dx)
     dx = jax.jit(mjx.sensor_vel)(mx, dx)
     dx = jax.jit(mjx.sensor_acc)(mx, dx)

mjx/mujoco/mjx/_src/smooth.py

@@ -23,10 +23,12 @@ from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import CamLightType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import EqType
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import TrnType
 from mujoco.mjx._src.types import WrapType
 # pylint: enable=g-importing-member
@@ -35,7 +37,6 @@ import numpy as np
 
 def kinematics(m: Model, d: Data) -> Data:
   """Converts position/velocity from generalized coordinates to maximal."""
-
   def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
     # calculate joint anchors, axes, body pos and quat in global frame
     # also normalize qpos while we're at it
@@ -131,6 +132,8 @@ def kinematics(m: Model, d: Data) -> Data:
 
 def com_pos(m: Model, d: Data) -> Data:
   """Maps inertias and motion dofs to global frame centered at subtree-CoM."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('com_pos requires JAX backend implementation.')
 
   # calculate center of mass of each subtree
   def subtree_sum(carry, xipos, body_mass):
@@ -162,7 +165,7 @@ def com_pos(m: Model, d: Data) -> Data:
   root_com = subtree_com[m.body_rootid]
   offset = d.xipos - root_com
   cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
-  d = d.replace(cinert=cinert)
+  d = d.tree_replace({'_impl.cinert': cinert})
 
   # map motion dofs to global frame centered at subtree_com
   def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis):
@@ -201,13 +204,16 @@ def com_pos(m: Model, d: Data) -> Data:
       d.xanchor,
       d.xaxis,
   )
-  d = d.replace(cdof=cdof)
+  d = d.tree_replace({'_impl.cdof': cdof})
 
   return d
 
 
 def camlight(m: Model, d: Data) -> Data:
   """Computes camera and light positions and orientations."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('camlight requires JAX backend implementation.')
+
   if m.ncam == 0:
     return d.replace(cam_xpos=jp.zeros((0, 3)), cam_xmat=jp.zeros((0, 3, 3)))
 
@@ -278,32 +284,33 @@ def camlight(m: Model, d: Data) -> Data:
 
 def crb(m: Model, d: Data) -> Data:
   """Runs composite rigid body inertia algorithm."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('crb requires JAX backend implementation.')
 
   def crb_fn(crb_child, crb_body):
     if crb_child is not None:
       crb_body += crb_child
     return crb_body
 
-  crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d.cinert, reverse=True)
+  crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True)
   crb_body = crb_body.at[0].set(0.0)
-  d = d.replace(crb=crb_body)
+  d = d.tree_replace({'_impl.crb': crb_body})
 
   crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
-  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
+  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof)
-  qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
+  qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature)
-  d = d.replace(qM=qm)
+  d = d.tree_replace({'_impl.qM': qm})
-  if support.is_sparse(m) and d._qM_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qM_sparse=qm)
-
   return d
 
 
 def factor_m(m: Model, d: Data) -> Data:
   """Gets factorizaton of inertia-like matrix M, assumed spd."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('factor_m requires JAX backend implementation.')
 
   if not support.is_sparse(m):
-    qh, _ = jax.scipy.linalg.cho_factor(d.qM)
+    qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM)
-    d = d.replace(qLD=qh)
+    d = d.tree_replace({'_impl.qLD': qh})
     return d
 
   # build up indices for where we will do backwards updates over qLD
@@ -325,7 +332,7 @@ def factor_m(m: Model, d: Data) -> Data:
           (out_beg, out_end, madr_d, madr_ij)
       )
 
-  qld = d.qM
+  qld = d._impl.qM
 
   for _, updates in sorted(updates.items(), reverse=True):
     # combine the updates into one update batch (per depth level)
@@ -353,20 +360,17 @@ def factor_m(m: Model, d: Data) -> Data:
   qld_diag = qld[m.dof_Madr]
   qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
 
-  d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag)
+  d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
-  if d._qLD_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qLD_sparse=d.qLD)
-  if d._qLDiagInv_sparse.size > 0:  # pylint: disable=protected-access
-    d = d.replace(_qLDiagInv_sparse=d.qLDiagInv)
-
   return d
 
 
 def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   """Computes sparse backsubstitution:  x = inv(L'*D*L)*y ."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('solve_m requires JAX backend implementation.')
 
   if not support.is_sparse(m):
-    return jax.scipy.linalg.cho_solve((d.qLD, False), x)
+    return jax.scipy.linalg.cho_solve((d._impl.qLD, False), x)
 
   depth = []
   for i in range(m.nv):
@@ -385,21 +389,23 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   # x <- inv(L') * x
   for _, vals in sorted(updates_j.items(), reverse=True):
     j, madr_ij, i = np.array(vals).T
-    x = x.at[j].add(-d.qLD[madr_ij] * x[i])
+    x = x.at[j].add(-d._impl.qLD[madr_ij] * x[i])
 
   # x <- inv(D) * x
-  x = x * d.qLDiagInv
+  x = x * d._impl.qLDiagInv
 
   # x <- inv(L) * x
   for _, vals in sorted(updates_i.items()):
     i, madr_ij, j = np.array(vals).T
-    x = x.at[i].add(-d.qLD[madr_ij] * x[j])
+    x = x.at[i].add(-d._impl.qLD[madr_ij] * x[j])
 
   return x
 
 
 def com_vel(m: Model, d: Data) -> Data:
   """Computes cvel, cdof_dot."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('com_vel requires JAX backend implementation.')
 
   # forward scan down tree: accumulate link center of mass velocity
   def fn(parent, jnt_typs, cdof, qvel):
@@ -431,17 +437,19 @@ def com_vel(m: Model, d: Data) -> Data:
       'jvv',
       'bv',
       m.jnt_type,
-      d.cdof,
+      d._impl.cdof,
       d.qvel,
   )
 
-  d = d.replace(cvel=cvel, cdof_dot=cdof_dot)
+  d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot})
 
   return d
 
 
 def subtree_vel(m: Model, d: Data) -> Data:
   """Subtree linear velocity and angular momentum."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('subtree_vel requires JAX backend implementation.')
 
   # bodywise quantities
   def _forward(cvel, xipos, ximat, subtree_com_root, mass, inertia):
@@ -529,14 +537,22 @@ def subtree_vel(m: Model, d: Data) -> Data:
       reverse=True,
   )
 
-  return d.replace(subtree_linvel=subtree_linvel, subtree_angmom=subtree_angmom)
+  return d.tree_replace({
+      '_impl.subtree_linvel': subtree_linvel,
+      '_impl.subtree_angmom': subtree_angmom,
+  })
 
 
-def rne(m: Model, d: Data) -> Data:
+def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
-  """Computes inverse dynamics using the recursive Newton-Euler algorithm."""
+  """Computes inverse dynamics using the recursive Newton-Euler algorithm.
+
+  flg_acc=False removes inertial term.
+  """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('rne requires JAX backend implementation.')
 
   # forward scan over tree: accumulate link center of mass acceleration
-  def cacc_fn(cacc, cdof_dot, qvel):
+  def cacc_fn(cacc, cdof_dot, qvel, cdof, qacc):
     if cacc is None:
       if m.opt.disableflags & DisableBit.GRAVITY:
         cacc = jp.zeros((6,))
@@ -545,9 +561,15 @@ def rne(m: Model, d: Data) -> Data:
 
     cacc += jp.sum(jax.vmap(jp.multiply)(cdof_dot, qvel), axis=0)
 
+    # cacc += cdof * qacc
+    if flg_acc:
+      cacc += jp.sum(jax.vmap(jp.multiply)(cdof, qacc), axis=0)
+
     return cacc
 
-  cacc = scan.body_tree(m, cacc_fn, 'vv', 'b', d.cdof_dot, d.qvel)
+  cacc = scan.body_tree(
+      m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc
+  )
 
   def frc(cinert, cacc, cvel):
     frc = math.inert_mul(cinert, cacc)
@@ -555,7 +577,7 @@ def rne(m: Model, d: Data) -> Data:
 
     return frc
 
-  loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel)
+  loc_cfrc = jax.vmap(frc)(d._impl.cinert, cacc, d.cvel)
 
   # backward scan up tree: accumulate body forces
   def cfrc_fn(cfrc_child, cfrc):
@@ -564,7 +586,7 @@ def rne(m: Model, d: Data) -> Data:
     return cfrc
 
   cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True)
-  qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)])
+  qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)])
 
   d = d.replace(qfrc_bias=qfrc_bias)
 
@@ -573,6 +595,8 @@ def rne(m: Model, d: Data) -> Data:
 
 def rne_postconstraint(m: Model, d: Data) -> Data:
   """RNE with complete data: compute cacc, cfrc_ext, cfrc_int."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('rne_postconstraint requires JAX backend implementation.')
 
   def _transform_force(frc, offset):
     force, torque = jp.split(frc, 2)
@@ -593,7 +617,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
   # compute contact forces for each condim
   forces = []
   condim_idx = []
-  for dim in set(d.contact.dim):
+  for dim in set(d._impl.contact.dim):
     force, idx = support.contact_force_dim(m, d, dim)
     forces.append(force)
     condim_idx.append(idx)
@@ -620,10 +644,10 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
       )
 
     condim_idx = jp.concatenate(condim_idx)
-    frame = d.contact.frame[condim_idx]
+    frame = d._impl.contact.frame[condim_idx]
-    pos = d.contact.pos[condim_idx]
+    pos = d._impl.contact.pos[condim_idx]
-    id1 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 0]]
+    id1 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 0]]
-    id2 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 1]]
+    id2 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 1]]
     com1 = d.subtree_com[jp.array(m.body_rootid)][id1]
     com2 = d.subtree_com[jp.array(m.body_rootid)][id2]
 
@@ -659,8 +683,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
     )
 
     # cacc = cacc_parent + cdofdot * qvel + cdof * qacc
-    cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
+    cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel)
-    cacc_acc = d.cdof.T @ (mask * d.qacc)
+    cacc_acc = d._impl.cdof.T @ (mask * d.qacc)
     cacc = cacc_parent + cacc_vel + cacc_acc
 
     # cfrc_body = cinert * cacc + cvel x (cinert * cvel)
@@ -678,7 +702,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
       'bbbbb',
       'bb',
       cfrc_ext,
-      d.cinert,
+      d._impl.cinert,
       d.cvel,
       jp.array(m.body_dofadr),
       jp.array(m.body_dofnum),
@@ -695,11 +719,18 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
   )
 
   # update data
-  return d.replace(cacc=cacc, cfrc_int=cfrc_int, cfrc_ext=cfrc_ext)
+  return d.tree_replace({
+      '_impl.cacc': cacc,
+      '_impl.cfrc_int': cfrc_int,
+      '_impl.cfrc_ext': cfrc_ext,
+  })
 
 
 def tendon(m: Model, d: Data) -> Data:
   """Computes tendon lengths and moments."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('tendon requires JAX backend implementation.')
+
   if not m.ntendon:
     return d
 
@@ -827,8 +858,8 @@ def tendon(m: Model, d: Data) -> Data:
   # wrap inside
   # TODO(taylorhowell): check that is_wrap_inside is consistent with
   # site and geom relative positions
-  (wrap_inside_id,) = np.nonzero(m.is_wrap_inside)
+  (wrap_inside_id,) = np.nonzero(m._impl.is_wrap_inside)
-  (wrap_outside_id,) = np.nonzero(~m.is_wrap_inside)
+  (wrap_outside_id,) = np.nonzero(~m._impl.is_wrap_inside)
 
   # compute geom wrap length and connect points (if wrap occurs)
   v_wrap = jax.vmap(
@@ -844,9 +875,9 @@ def tendon(m: Model, d: Data) -> Data:
       has_sidesite[wrap_inside_id],
       is_sphere[wrap_inside_id],
       True,
-      m.wrap_inside_maxiter,
+      m._impl.wrap_inside_maxiter,
-      m.wrap_inside_tolerance,
+      m._impl.wrap_inside_tolerance,
-      m.wrap_inside_z_init,
+      m._impl.wrap_inside_z_init,
   )
 
   lengths_outside, pnt0_outside, pnt1_outside = v_wrap(
@@ -859,9 +890,9 @@ def tendon(m: Model, d: Data) -> Data:
       has_sidesite[wrap_outside_id],
       is_sphere[wrap_outside_id],
       False,
-      m.wrap_inside_maxiter,
+      m._impl.wrap_inside_maxiter,
-      m.wrap_inside_tolerance,
+      m._impl.wrap_inside_tolerance,
-      m.wrap_inside_z_init,
+      m._impl.wrap_inside_z_init,
   )
 
   wrap_id = np.argsort(np.concatenate([wrap_inside_id, wrap_outside_id]))
@@ -943,12 +974,14 @@ def tendon(m: Model, d: Data) -> Data:
   )
 
   # assemble length and moment
-  ten_length = jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt)
+  ten_length = (
+      jp.zeros_like(d._impl.ten_length).at[tendon_id_jnt].set(length_jnt)
+  )
   ten_length = ten_length.at[tendon_id_site].add(length_site)
   ten_length = ten_length.at[tendon_id_geom].add(length_geom)
 
   ten_moment = (
-      jp.zeros_like(d.ten_J)
+      jp.zeros_like(d._impl.ten_J)
       .at[adr_moment_jnt, dofadr_moment_jnt]
       .set(moment_jnt)
   )
@@ -1011,14 +1044,14 @@ def tendon(m: Model, d: Data) -> Data:
       [wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)]
   ).reshape((m.nwrap, 2))
 
-  return d.replace(
+  return d.tree_replace({
-      ten_length=ten_length,
+      '_impl.ten_length': ten_length,
-      ten_J=ten_moment,
+      '_impl.ten_J': ten_moment,
-      ten_wrapadr=jp.array(ten_wrapadr, dtype=int),
+      '_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
-      ten_wrapnum=jp.array(ten_wrapnum, dtype=int),
+      '_impl.ten_wrapnum': jp.array(ten_wrapnum, dtype=int),
-      wrap_xpos=wrap_xpos,
+      '_impl.wrap_xpos': wrap_xpos,
-      wrap_obj=jp.array(wrap_obj, dtype=int),
+      '_impl.wrap_obj': jp.array(wrap_obj, dtype=int),
-  )
+  })
 
 
 def _site_dof_mask(m: Model) -> np.ndarray:
@@ -1052,6 +1085,9 @@ def _site_dof_mask(m: Model) -> np.ndarray:
 
 def transmission(m: Model, d: Data) -> Data:
   """Computes actuator/transmission lengths and moments."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('transmission requires JAX backend implementation.')
+
   if not m.nu:
     return d
 
@@ -1111,8 +1147,8 @@ def transmission(m: Model, d: Data) -> Data:
       wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
       moment = jac @ wrench
     elif trntype == TrnType.TENDON:
-      length = d.ten_length[trnid[0]] * gear[:1]
+      length = d._impl.ten_length[trnid[0]] * gear[:1]
-      moment = d.ten_J[trnid[0]] * gear[0]
+      moment = d._impl.ten_J[trnid[0]] * gear[0]
     else:
       raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
 
@@ -1144,5 +1180,7 @@ def transmission(m: Model, d: Data) -> Data:
   length = length.reshape((m.nu,))
   moment = moment.reshape((m.nu, m.nv))
 
-  d = d.replace(actuator_length=length, actuator_moment=moment)
+  d = d.tree_replace(
+      {'_impl.actuator_length': length, '_impl.actuator_moment': moment}
+  )
   return d

mjx/mujoco/mjx/_src/smooth_test.py

@@ -20,7 +20,7 @@ import jax
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
-from mujoco.mjx._src.types import ConeType
+from mujoco.mjx._src.types import ConeType  # pylint: disable=g-importing-member
 import numpy as np
 
 # tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -78,33 +78,37 @@ class SmoothTest(absltest.TestCase):
     # com_pos
     dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'subtree_com')
-    _assert_attr_eq(d, dx, 'cinert')
+    _assert_attr_eq(d, dx._impl, 'cinert')
-    _assert_attr_eq(d, dx, 'cdof')
+    _assert_attr_eq(d, dx._impl, 'cdof')
     # camlight
     dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cam_xpos')
     _assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat')
     # crb
     dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'crb')
+    _assert_attr_eq(d, dx._impl, 'crb')
-    _assert_attr_eq(d, dx, 'qM')
+    _assert_attr_eq(d, dx._impl, 'qM')
-    _assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
     qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
-    for i in range(m.nM):
+    for i in range(m.nC):
       qLDLegacy[d.mapM2M[i]] = d.qLD[i]
-    _assert_eq(qLDLegacy, dx.qLD, 'qLD')
+    _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
-    _assert_attr_eq(d, dx, 'qLDiagInv')
+    _assert_attr_eq(d, dx._impl, 'qLDiagInv')
-    _assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
-    _assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
     # com_vel
     dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'cvel')
-    _assert_attr_eq(d, dx, 'cdof_dot')
+    _assert_attr_eq(d, dx._impl, 'cdof_dot')
     # rne
     dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx, 'qfrc_bias')
+    # rne (flg_acc=True)
+    qfrc_bias = np.zeros(m.nv)
+    mujoco.mj_rne(m, d, 1, qfrc_bias)
+    dx = jax.jit(mjx.rne, static_argnums=(2,))(
+        mx, mjx.put_data(m, d), flg_acc=True
+    )
+    _assert_eq(dx.qfrc_bias, qfrc_bias, 'qfrc_bias')
 
     # set dense jacobian for tendon:
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
@@ -115,11 +119,11 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_forward(m, d)
     # tendon
     dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
-    _assert_attr_eq(d, dx, 'ten_J')
+    _assert_attr_eq(d, dx._impl, 'ten_J')
-    _assert_attr_eq(d, dx, 'ten_length')
+    _assert_attr_eq(d, dx._impl, 'ten_length')
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx._impl, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
@@ -130,7 +134,7 @@ class SmoothTest(absltest.TestCase):
         d.moment_rowadr,
         d.moment_colind,
     )
-    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
+    _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
 
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""
@@ -190,7 +194,7 @@ class SmoothTest(absltest.TestCase):
 
     mujoco.mj_transmission(m, d)
     dx = jax.jit(mjx.transmission)(mx, dx)
-    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx._impl, 'actuator_length')
 
     # convert sparse actuator_moment to dense representation
     moment = np.zeros((m.nu, m.nv))
@@ -201,7 +205,7 @@ class SmoothTest(absltest.TestCase):
         d.moment_rowadr,
         d.moment_colind,
     )
-    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
+    _assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
 
   def test_subtree_vel(self):
     """Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel."""
@@ -219,8 +223,8 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_subtreeVel(m, d)
     dx = jax.jit(mjx.subtree_vel)(mx, dx)
 
-    _assert_attr_eq(d, dx, 'subtree_linvel')
+    _assert_attr_eq(d, dx._impl, 'subtree_linvel')
-    _assert_attr_eq(d, dx, 'subtree_angmom')
+    _assert_attr_eq(d, dx._impl, 'subtree_angmom')
 
 
 class RnePostConstraintTest(parameterized.TestCase):
@@ -271,9 +275,9 @@ class RnePostConstraintTest(parameterized.TestCase):
     mujoco.mj_rnePostConstraint(m, d)
     dx = jax.jit(mjx.rne_postconstraint)(mx, dx)
 
-    _assert_eq(d.cacc, dx.cacc, 'cacc')
+    _assert_eq(d.cacc, dx._impl.cacc, 'cacc')
-    _assert_eq(d.cfrc_ext, dx.cfrc_ext, 'cfrc_ext')
+    _assert_eq(d.cfrc_ext, dx._impl.cfrc_ext, 'cfrc_ext')
-    _assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int')
+    _assert_eq(d.cfrc_int, dx._impl.cfrc_int, 'cfrc_int')
 
 
 class TendonTest(parameterized.TestCase):
@@ -305,12 +309,12 @@ class TendonTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
     dx = jax.jit(mjx.forward)(mx, dx)
 
-    _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
+    _assert_eq(d.ten_length, dx._impl.ten_length, 'ten_length')
-    _assert_eq(d.ten_J, dx.ten_J, 'ten_J')
+    _assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
-    _assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum')
+    _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
-    _assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr')
+    _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
-    _assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj')
+    _assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
-    _assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos')
+    _assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/solver.py

@@ -24,13 +24,15 @@ from mujoco.mjx._src import support
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import ConeType
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DataJAX
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import ModelJAX
 from mujoco.mjx._src.types import SolverType
 # pylint: enable=g-importing-member
 
 
-class _Context(PyTreeNode):
+class Context(PyTreeNode):
   """Data updated during each solver iteration.
 
   Attributes:
@@ -72,25 +74,30 @@ class _Context(PyTreeNode):
   h: jax.Array
 
   @classmethod
-  def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context':
+  def create(cls, m: Model, d: Data, grad: bool = True) -> 'Context':
-    jaref = d.efc_J @ d.qacc - d.efc_aref
+    if not isinstance(d._impl, DataJAX):
+      raise ValueError(
+          'Constraint context requires JAX backend implementation.'
+      )
+
+    jaref = d._impl.efc_J @ d.qacc - d._impl.efc_aref
     # TODO(robotics-team): determine nv at which sparse mul is faster
     ma = support.mul_m(m, d, d.qacc)
     nv_0 = jp.zeros(m.nv)
     fri = 0.0
     if m.opt.cone == ConeType.ELLIPTIC:
-      friction = d.contact.friction[d.contact.dim > 1]
+      friction = d._impl.contact.friction[d._impl.contact.dim > 1]
-      dim = d.contact.dim[d.contact.dim > 1]
+      dim = d._impl.contact.dim[d._impl.contact.dim > 1]
       mu = friction[:, 0] / jp.sqrt(m.opt.impratio)
       fri = jp.concatenate((mu[:, None], friction), axis=1)
       for condim in (3, 4, 6):
         fri = fri.at[dim == condim, condim:].set(0)
 
-    ctx = _Context(
+    ctx = Context(
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
-        efc_force=d.efc_force,
+        efc_force=d._impl.efc_force,
         Ma=ma,
         grad=nv_0,
         Mgrad=nv_0,
@@ -133,7 +140,7 @@ class _LSPoint(PyTreeNode):
       cls,
       m: Model,
       d: Data,
-      ctx: _Context,
+      ctx: Context,
       alpha: jax.Array,
       jv: jax.Array,
       quad: jax.Array,
@@ -145,18 +152,20 @@ class _LSPoint(PyTreeNode):
   ) -> '_LSPoint':
     """Creates a linesearch point with first and second derivatives."""
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
+    if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+      raise ValueError('LSPoint requires JAX backend implementation.')
 
     cost, deriv_0, deriv_1 = 0.0, 0.0, 0.0
     quad_total = quad_gauss
     x = ctx.Jaref + alpha * jv
-    active = (x < 0).at[: d.ne + d.nf].set(True)
+    active = (x < 0).at[: d._impl.ne + d._impl.nf].set(True)
 
-    dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+    dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
     if (dof_fl.any() or ten_fl.any()) and not (
         m.opt.disableflags & DisableBit.FRICTIONLOSS
     ):
-      f = d.efc_frictionloss
+      f = d._impl.efc_frictionloss
-      r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+      r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
       rf, z = r * f, jp.zeros_like(f)
       linear_neg = (x <= -rf)[:, None]
       linear_pos = (x >= rf)[:, None]
@@ -174,13 +183,13 @@ class _LSPoint(PyTreeNode):
       middle_zone = (tsqr > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
 
       # quadratic cost for equality, friction, limits, frictionless contacts
-      dim1 = d.contact.efc_address[d.contact.dim == 1]
+      dim1 = d._impl.contact.efc_address[d._impl.contact.dim == 1]
-      nefl = d.ne + d.nf + d.nl
+      nefl = d._impl.ne + d._impl.nf + d._impl.nl
       active = active.at[nefl:].set(False).at[dim1].set(active[dim1])
       quad_efld = jax.vmap(jp.multiply)(quad, active)
       quad_total += jp.sum(quad_efld, axis=0)
       # elliptic bottom zone: quadratic cost
-      efc_elliptic = d.contact.efc_address[d.contact.dim > 1]
+      efc_elliptic = d._impl.contact.efc_address[d._impl.contact.dim > 1]
       quad_c = jax.vmap(jp.multiply)(quad[efc_elliptic], bottom_zone)
       quad_total += jp.sum(quad_c, axis=0)
       # elliptic middle zone
@@ -241,7 +250,7 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
   return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]
 
 
-def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
+def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
   """Updates constraint force and resulting cost given last solver iteration.
 
   Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
@@ -254,17 +263,20 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   Returns:
     context with new constraint force and costs
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_update_constraint requires JAX backend implementation.')
+
   # ne constraints are always active, nf are conditionally active, others are
   # non-negative constraints.
-  active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
+  active = (ctx.Jaref < 0).at[: d._impl.ne + d._impl.nf].set(True)
 
-  floss_force, floss_cost = jp.zeros(d.nefc), 0.0
+  floss_force, floss_cost = jp.zeros(d._impl.nefc), 0.0
-  dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
+  dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
   if (dof_fl.any() or ten_fl.any()) and not (
       m.opt.disableflags & DisableBit.FRICTIONLOSS
   ):
-    f = d.efc_frictionloss
+    f = d._impl.efc_frictionloss
-    r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
+    r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
     linear_neg = (ctx.Jaref <= -r * f) * (f > 0)
     linear_pos = (ctx.Jaref >= r * f) * (f > 0)
     active = active & ~linear_neg & ~linear_pos
@@ -274,13 +286,13 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
     floss_cost = floss_cost.sum()
 
   if m.opt.cone == ConeType.PYRAMIDAL:
-    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
+    efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
-    cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
+    cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
     dm, u, h = 0.0, 0.0, 0.0
   elif m.opt.cone == ConeType.ELLIPTIC:
-    friction = d.contact.friction[d.contact.dim > 1]
+    friction = d._impl.contact.friction[d._impl.contact.dim > 1]
-    efc_address = d.contact.efc_address[d.contact.dim > 1]
+    efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
-    dim = d.contact.dim[d.contact.dim > 1]
+    dim = d._impl.contact.dim[d._impl.contact.dim > 1]
     # to prevent out of range append zeros to ctx.Jaref
     slice_fn = jax.vmap(
         lambda x: jax.lax.dynamic_slice(
@@ -297,12 +309,12 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
       adr_i.extend(range(addr, addr + condim))
       adr_j.extend([i] * condim)
     active = active.at[jp.array(adr_i)].set(bottom_zone[jp.array(adr_j)])
-    efc_force = d.efc_D * -ctx.Jaref * active + floss_force
+    efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
-    cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
+    cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
 
     # middle zone: cone
     middle_zone = (t > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
-    dm = d.efc_D[efc_address] / jp.maximum(
+    dm = d._impl.efc_D[efc_address] / jp.maximum(
         mu * mu * (1 + mu * mu), mujoco.mjMINVAL
     )
     nmt = n - mu * t
@@ -339,7 +351,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   else:
     raise NotImplementedError(f'unsupported cone type: {m.opt.cone}')
 
-  qfrc_constraint = d.efc_J.T @ efc_force
+  qfrc_constraint = d._impl.efc_J.T @ efc_force
   gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
   ctx = ctx.replace(
       qfrc_constraint=qfrc_constraint,
@@ -356,7 +368,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   return ctx
 
 
-def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
+def _update_gradient(m: Model, d: Data, ctx: Context) -> Context:
   """Updates grad and M / grad given latest solver iteration.
 
   Corresponds to CGupdateGradient in mujoco/src/engine/engine_solver.c
@@ -371,6 +383,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
   Raises:
     NotImplementedError: for unsupported solver type
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_update_gradient requires JAX backend implementation.')
 
   grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint
 
@@ -378,16 +392,16 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
     mgrad = smooth.solve_m(m, d, grad)
   elif m.opt.solver == SolverType.NEWTON:
     if m.opt.cone == ConeType.ELLIPTIC:
-      cm = jp.diag(d.efc_D * ctx.active)
+      cm = jp.diag(d._impl.efc_D * ctx.active)
-      efc_address = d.contact.efc_address[d.contact.dim > 1]
+      efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
-      dim = d.contact.dim[d.contact.dim > 1]
+      dim = d._impl.contact.dim[d._impl.contact.dim > 1]
       # set efc of cone H along diagonal
       for i, (condim, addr) in enumerate(zip(dim, efc_address)):
         h_cone = ctx.h[i, :condim, :condim]
         cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone)
-      h = d.efc_J.T @ cm @ d.efc_J
+      h = d._impl.efc_J.T @ cm @ d._impl.efc_J
     else:
-      h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J
+      h = (d._impl.efc_J.T * d._impl.efc_D * ctx.active) @ d._impl.efc_J
     h = support.full_m(m, d) + h
     h_ = jax.scipy.linalg.cho_factor(h)
     mgrad = jax.scipy.linalg.cho_solve(h_, grad)
@@ -403,7 +417,7 @@ def _rescale(m: Model, value: jax.Array) -> jax.Array:
   return value / (m.stat.meaninertia * max(1, m.nv))
 
 
-def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
+def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
   """Performs a zoom linesearch to find optimal search step size.
 
   Args:
@@ -414,12 +428,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   Returns:
     updated context with new qacc, Ma, Jaref
   """
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('_lineasearch requires JAX backend implementation.')
+
   smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv)
   gtol = m.opt.tolerance * m.opt.ls_tolerance * smag
 
   # compute Mv, Jv
   mv = support.mul_m(m, d, ctx.search)
-  jv = d.efc_J @ ctx.search
+  jv = d._impl.efc_J @ ctx.search
 
   # prepare quadratics
   quad_gauss = jp.stack((
@@ -428,13 +445,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
       0.5 * jp.dot(ctx.search, mv),
   ))
   quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
-  quad = (quad * d.efc_D).T
+  quad = (quad * d._impl.efc_D).T
   uu, v0, uv, vv = 0.0, 0.0, 0.0, 0.0
   if m.opt.cone == ConeType.ELLIPTIC:
-    mask = d.contact.dim > 1
+    mask = d._impl.contact.dim > 1
     # complete vector quadratic (for bottom zone)
     efc_con, efc_fri = [], []
-    for condim, addr in zip(d.contact.dim[mask], d.contact.efc_address[mask]):
+    for condim, addr in zip(
+        d._impl.contact.dim[mask], d._impl.contact.efc_address[mask]
+    ):
       efc_con.extend([addr] * (condim - 1))
       efc_fri.extend(range(addr + 1, addr + condim))
     quad = quad.at[jp.array(efc_con)].add(quad[jp.array(efc_fri)])
@@ -446,7 +465,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
             jp.concatenate((jv, jp.zeros(3))), (x,), (6,)
         )
     )
-    efc_elliptic = d.contact.efc_address[mask]
+    efc_elliptic = d._impl.contact.efc_address[mask]
     v = jv_fn(efc_elliptic) * ctx.fri
     uu = jp.sum(ctx.u[:, 1:] * ctx.u[:, 1:], axis=1)
     v0 = v[:, 0]
@@ -529,7 +548,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
 def solve(m: Model, d: Data) -> Data:
   """Finds forces that satisfy constraints using conjugate gradient descent."""
 
-  def cond(ctx: _Context) -> jax.Array:
+  def cond(ctx: Context) -> jax.Array:
     improvement = _rescale(m, ctx.prev_cost - ctx.cost)
     gradient = _rescale(m, math.norm(ctx.grad))
 
@@ -539,7 +558,7 @@ def solve(m: Model, d: Data) -> Data:
 
     return ~done
 
-  def body(ctx: _Context) -> _Context:
+  def body(ctx: Context) -> Context:
     ctx = _linesearch(m, d, ctx)
     prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad  # pylint: disable=invalid-name
     ctx = _update_constraint(m, d, ctx)
@@ -560,22 +579,22 @@ def solve(m: Model, d: Data) -> Data:
   # warmstart:
   qacc = d.qacc_smooth
   if not m.opt.disableflags & DisableBit.WARMSTART:
-    warm = _Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
+    warm = Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
-    smth = _Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False)
+    smth = Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False)
     qacc = jp.where(warm.cost < smth.cost, d.qacc_warmstart, d.qacc_smooth)
   d = d.replace(qacc=qacc)
 
-  ctx = _Context.create(m, d)
+  ctx = Context.create(m, d)
   if m.opt.iterations == 1:
     ctx = body(ctx)
   else:
     ctx = jax.lax.while_loop(cond, body, ctx)
 
-  d = d.replace(
+  d = d.tree_replace({
-      qacc_warmstart=ctx.qacc,
+      'qacc_warmstart': ctx.qacc,
-      qacc=ctx.qacc,
+      'qfrc_constraint': ctx.qfrc_constraint,
-      qfrc_constraint=ctx.qfrc_constraint,
+      'qacc': ctx.qacc,
-      efc_force=ctx.efc_force,
+      '_impl.efc_force': ctx.efc_force,
-  )
+  })
 
   return d

mjx/mujoco/mjx/_src/solver_test.py

@@ -74,7 +74,7 @@ class SolverTest(parameterized.TestCase):
 
       # compare costs
       mj_cost = cost(d.qacc)
-      ctx = solver._Context.create(mjx.put_model(m), mjx.put_data(m, d))
+      ctx = solver.Context.create(mjx.put_model(m), mjx.put_data(m, d))
       mjx_cost = ctx.cost - ctx.gauss
       _assert_eq(mj_cost, mjx_cost, 'cost')
 
@@ -87,8 +87,8 @@ class SolverTest(parameterized.TestCase):
 
       # MJX finds very similar solutions with the newton solver
       if solver_ == mujoco.mjtSolver.mjSOL_NEWTON:
-        nnz = dx.efc_J.any(axis=1)
+        nnz = dx._impl.efc_J.any(axis=1)
-        _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+        _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
         _assert_attr_eq(d, dx, 'qfrc_constraint')
         _assert_attr_eq(d, dx, 'qacc')
 
@@ -108,9 +108,9 @@ class SolverTest(parameterized.TestCase):
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
     # even without warmstart, newton converges quickly
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-4)
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force', tol=2e-4)
 
   def test_sparse(self):
     """Test solver works with sparse mass matrices."""
@@ -130,8 +130,8 @@ class SolverTest(parameterized.TestCase):
 
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
 
   def test_quad_frictionloss(self):
     """Test a case with quadratic frictionloss constraints."""
@@ -144,8 +144,8 @@ class SolverTest(parameterized.TestCase):
 
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
-    nnz = dx.efc_J.any(axis=1)
+    nnz = dx._impl.efc_J.any(axis=1)
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+    _assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
 
   # TODO(taylorhowell): condim=1 with ConeType.ELLIPTIC
   @parameterized.product(condim=(3, 4, 6), cone=tuple(ConeType))

mjx/mujoco/mjx/_src/support.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 """Engine support functions."""
+
 from collections.abc import Iterable, Sequence
 from typing import Optional, Tuple, Union
 
@@ -92,7 +93,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
   """Reconstitute dense mass matrix from qM."""
 
   if not is_sparse(m):
-    return d.qM
+    return d._impl.qM  # pytype: disable=attribute-error
 
   ij = []
   for i in range(m.nv):
@@ -103,7 +104,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
 
   i, j = (jp.array(x) for x in zip(*ij))
 
-  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM)  # pytype: disable=attribute-error
 
   # also set upper triangular
   mat = mat + jp.tril(mat, -1).T
@@ -115,9 +116,9 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
   """Multiply vector by inertia matrix."""
 
   if not is_sparse(m):
-    return d.qM @ vec
+    return d._impl.qM @ vec  # pytype: disable=attribute-error
 
-  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
+  diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec  # pytype: disable=attribute-error
 
   is_, js, madr_ijs = [], [], []
   for i in range(m.nv):
@@ -131,8 +132,8 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
 
   i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
 
-  out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j])
+  out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j])  # pytype: disable=attribute-error
-  out = out.at[j].add(d.qM[madr_ij] * vec[i])
+  out = out.at[j].add(d._impl.qM[madr_ij] * vec[i])  # pytype: disable=attribute-error
 
   return out
 
@@ -147,9 +148,9 @@ def jac(
   mask = mask[jp.array(m.dof_bodyid)] > 0
 
   offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
-  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof)
+  jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof)  # pytype: disable=attribute-error
   jacp = jax.vmap(jp.multiply)(jacp, mask)
-  jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask)
+  jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask)  # pytype: disable=attribute-error
 
   return jacp, jacr
 
@@ -295,73 +296,56 @@ class BindModel(object):
     self.prefix = ''
     ids = []
     for spec in specs:
-      if not spec.name:
+      if model.signature != spec.signature:
-        raise KeyError(f'cannot bind spec with empty name')
+        raise ValueError(
+            'mjSpec signature does not match mjx.Model signature:'
+            f' {spec.signature} != {model.signature}'
+        )
+      elif spec.id < 0:
+        raise KeyError(f'invalid id: {spec.id}')
       elif isinstance(spec, mujoco.MjsBody):
         self.prefix = 'body_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
       elif isinstance(spec, mujoco.MjsJoint):
         self.prefix = 'jnt_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
       elif isinstance(spec, mujoco.MjsGeom):
         self.prefix = 'geom_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
       elif isinstance(spec, mujoco.MjsSite):
         self.prefix = 'site_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
       elif isinstance(spec, mujoco.MjsLight):
         self.prefix = 'light_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
       elif isinstance(spec, mujoco.MjsCamera):
         self.prefix = 'cam_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
       elif isinstance(spec, mujoco.MjsMesh):
         self.prefix = 'mesh_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_MESH, spec.name)
       elif isinstance(spec, mujoco.MjsHField):
         self.prefix = 'hfield_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, spec.name)
       elif isinstance(spec, mujoco.MjsPair):
         self.prefix = 'pair_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_PAIR, spec.name)
       elif isinstance(spec, mujoco.MjsTendon):
         self.prefix = 'tendon_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
       elif isinstance(spec, mujoco.MjsActuator):
         self.prefix = 'actuator_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
       elif isinstance(spec, mujoco.MjsSensor):
         self.prefix = 'sensor_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
       elif isinstance(spec, mujoco.MjsNumeric):
         self.prefix = 'numeric_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_NUMERIC, spec.name)
       elif isinstance(spec, mujoco.MjsText):
         self.prefix = 'text_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TEXT, spec.name)
       elif isinstance(spec, mujoco.MjsTuple):
         self.prefix = 'tuple_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TUPLE, spec.name)
       elif isinstance(spec, mujoco.MjsKey):
         self.prefix = 'key_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_KEY, spec.name)
       elif isinstance(spec, mujoco.MjsEquality):
         self.prefix = 'eq_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
       elif isinstance(spec, mujoco.MjsExclude):
         self.prefix = 'exclude_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EXCLUDE, spec.name)
       elif isinstance(spec, mujoco.MjsSkin):
         self.prefix = 'skin_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SKIN, spec.name)
       elif isinstance(spec, mujoco.MjsMaterial):
         self.prefix = 'material_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, spec.name)
       else:
         raise ValueError('invalid spec type')
-      if idx < 0:
+      ids.append(spec.id)
-        raise KeyError(f'invalid name: {spec.name}')  # pytype: disable=attribute-error
-      ids.append(idx)
     if len(ids) == 1:
       self.id = ids[0]
     else:
@@ -402,42 +386,36 @@ class BindData(object):
     self.prefix = ''
     ids = []
     for spec in specs:
-      if not spec.name:
+      if model.signature != spec.signature:
-        raise KeyError(f'cannot bind spec with empty name')
+        raise ValueError(
+            'mjSpec signature does not match mjx.Model signature:'
+            f' {spec.signature} != {model.signature}'
+        )
+      if spec.id < 0:
+        raise KeyError(f'invalid id: {spec.id}')
       elif isinstance(spec, mujoco.MjsBody):
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
+        pass
       elif isinstance(spec, mujoco.MjsJoint):
         self.prefix = 'jnt_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
       elif isinstance(spec, mujoco.MjsGeom):
         self.prefix = 'geom_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
       elif isinstance(spec, mujoco.MjsSite):
         self.prefix = 'site_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
       elif isinstance(spec, mujoco.MjsLight):
         self.prefix = 'light_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
       elif isinstance(spec, mujoco.MjsCamera):
         self.prefix = 'cam_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
       elif isinstance(spec, mujoco.MjsTendon):
         self.prefix = 'ten_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
       elif isinstance(spec, mujoco.MjsActuator):
         self.prefix = 'actuator_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
       elif isinstance(spec, mujoco.MjsSensor):
         self.prefix = 'sensor_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
       elif isinstance(spec, mujoco.MjsEquality):
         self.prefix = 'eq_'
-        idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
       else:
         raise ValueError('invalid spec type')
-      if idx < 0:
+      ids.append(spec.id)
-        raise KeyError(f'invalid name: {spec.name}')  # pytype: disable=attribute-error
-      ids.append(idx)
     if len(ids) == 1:
       self.id = ids[0]
     else:
@@ -455,11 +433,13 @@ class BindData(object):
         return name
       else:
         raise AttributeError('ctrl is not available for this type')
-    if name == 'qpos' or name == 'qvel' or name == 'qacc':
+    if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
       if self.prefix == 'jnt_':
         return name
       else:
-        raise AttributeError('qpos, qvel, qacc are not available for this type')
+        raise AttributeError(
+            'qpos, qvel, qacc, qfrc are not available for this type'
+        )
     else:
       return self.prefix + name
 
@@ -473,7 +453,9 @@ class BindData(object):
     return var[..., idx, :]
 
   def __getattr__(self, name: str):
-    if name in ('sensordata', 'qpos', 'qvel', 'qacc'):
+    if name in ('sensordata', 'qpos', 'qvel', 'qacc') or (
+        name.startswith('qfrc_')
+    ):
       adr = num = 0
       if name == 'sensordata':
         adr = self.model.sensor_adr[self.id]
@@ -482,7 +464,7 @@ class BindData(object):
         adr = self.model.jnt_qposadr[self.id]
         typ = self.model.jnt_type[self.id]
         num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
-      elif name == 'qvel' or name == 'qacc':
+      elif name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
         adr = self.model.jnt_dofadr[self.id]
         typ = self.model.jnt_type[self.id]
         num = sum((typ == jt) * jt.dof_width() for jt in JointType)
@@ -495,6 +477,8 @@ class BindData(object):
         return self._slice(self.__getname(name), slice(adr, adr + num))
       else:
         return self._slice(self.__getname(name), adr)
+    elif name in ('mocap_pos', 'mocap_quat'):
+      return self._slice(self.__getname(name), self.model.body_mocapid[self.id])
     return self._slice(self.__getname(name), self.id)
 
   def set(self, name: str, value: jax.Array) -> Data:
@@ -507,7 +491,7 @@ class BindData(object):
       iter(value)
     except TypeError:
       value = [value]
-    if name in ('qpos', 'qvel', 'qacc'):
+    if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
       adr = num = 0
       if name == 'qpos':
         adr = self.model.jnt_qposadr[self.id]
@@ -517,16 +501,23 @@ class BindData(object):
         adr = self.model.jnt_dofadr[self.id]
         typ = self.model.jnt_type[self.id]
         num = sum((typ == jt) * jt.dof_width() for jt in JointType)
+      elif name == 'mocap_pos':
+        adr = self.model.body_mocapid[self.id] * 3
+        num = np.ones_like(self.id, dtype=int) * 3
+      elif name == 'mocap_quat':
+        adr = self.model.body_mocapid[self.id] * 4
+        num = np.ones_like(self.id, dtype=int) * 4
       if not isinstance(self.id, list):
         adr = [adr]
         num = [num]
     elif isinstance(self.id, list):
-      adr = self.id * dim
+      adr = (np.array(self.id) * dim).tolist()
       num = [dim for _ in range(len(self.id))]
     else:
       adr = [self.id * dim]
       num = [dim]
     i = 0
+    value = jax.numpy.array(value).flatten()
     for a, n in zip(adr, num):
       shape = array.shape
       array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
@@ -571,20 +562,20 @@ def contact_force(
     m: Model, d: Data, contact_id: int, to_world_frame: bool = False
 ) -> jax.Array:
   """Extract 6D force:torque for one contact, in contact frame by default."""
-  efc_address = d.contact.efc_address[contact_id]
+  efc_address = d._impl.contact.efc_address[contact_id]  # pytype: disable=attribute-error
-  condim = d.contact.dim[contact_id]
+  condim = d._impl.contact.dim[contact_id]  # pytype: disable=attribute-error
   if m.opt.cone == ConeType.PYRAMIDAL:
     force = _decode_pyramid(
-        d.efc_force[efc_address:], d.contact.friction[contact_id], condim
+        d._impl.efc_force[efc_address:], d._impl.contact.friction[contact_id], condim  # pytype: disable=attribute-error
     )
   elif m.opt.cone == ConeType.ELLIPTIC:
-    force = d.efc_force[efc_address : efc_address + condim]
+    force = d._impl.efc_force[efc_address : efc_address + condim]  # pytype: disable=attribute-error
     force = jp.concatenate([force, jp.zeros((6 - condim))])
   else:
     raise ValueError(f'Unknown cone type: {m.opt.cone}')
 
   if to_world_frame:
-    force = force.reshape((-1, 3)) @ d.contact.frame[contact_id]
+    force = force.reshape((-1, 3)) @ d._impl.contact.frame[contact_id]  # pytype: disable=attribute-error
     force = force.reshape(-1)
 
   return force * (efc_address >= 0)
@@ -595,21 +586,21 @@ def contact_force_dim(
 ) -> Tuple[jax.Array, np.ndarray]:
   """Extract 6D force:torque for contacts with dimension dim."""
   # valid contact and condim indices
-  idx_dim = (d.contact.efc_address >= 0) & (d.contact.dim == dim)
+  idx_dim = (d._impl.contact.efc_address >= 0) & (d._impl.contact.dim == dim)  # pytype: disable=attribute-error
 
   # contact force from efc
   if m.opt.cone == ConeType.PYRAMIDAL:
     efc_address = (
-        d.contact.efc_address[idx_dim, None]
+        d._impl.contact.efc_address[idx_dim, None]  # pytype: disable=attribute-error
         + np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None]
     )
-    efc_force = d.efc_force[efc_address]
+    efc_force = d._impl.efc_force[efc_address]  # pytype: disable=attribute-error
     force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
-        efc_force, d.contact.friction[idx_dim], dim
+        efc_force, d._impl.contact.friction[idx_dim], dim  # pytype: disable=attribute-error
     )
   elif m.opt.cone == ConeType.ELLIPTIC:
-    efc_address = d.contact.efc_address[idx_dim, None] + np.arange(dim)[None]
+    efc_address = d._impl.contact.efc_address[idx_dim, None] + np.arange(dim)[None]  # pytype: disable=attribute-error
-    force = d.efc_force[efc_address]
+    force = d._impl.efc_force[efc_address]  # pytype: disable=attribute-error
     force = jp.hstack([force, jp.zeros((force.shape[0], 6 - dim))])
   else:
     raise ValueError(f'Unknown cone type: {m.opt.cone}.')

mjx/mujoco/mjx/_src/support_test.py

@@ -173,18 +173,20 @@ class SupportTest(parameterized.TestCase):
             <joint axis="0 0 1" type="slide" name="joint3"/>
             <geom size="7 8 9" type="box" name="geom3"/>
           </body>
+          <body pos="100 110 120" name="body4" mocap="true"/>
+          <body pos="130 140 150" name="body5" mocap="true"/>
         </worldbody>
 
         <actuator>
-          <motor name="actuator1" joint="joint1"/>
+          <motor joint="joint1"/>
-          <motor name="actuator2" joint="joint2"/>
+          <motor joint="joint2"/>
-          <motor name="actuator3" joint="joint3"/>
+          <motor joint="joint3"/>
         </actuator>
 
         <sensor>
-          <framepos name="sensor1" objtype="body" objname="body1"/>
+          <framepos objtype="body" objname="body1"/>
-          <framepos name="sensor2" objtype="body" objname="body2"/>
+          <framepos objtype="body" objname="body2"/>
-          <framepos name="sensor3" objtype="body" objname="body3"/>
+          <framepos objtype="body" objname="body3"/>
         </sensor>
     </mujoco>
     """
@@ -247,6 +249,11 @@ class SupportTest(parameterized.TestCase):
           dx.bind(mx, s.joints[i]).qacc,
           d.qacc[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6
       )
+      np.testing.assert_array_almost_equal(
+          dx.bind(mx, s.joints[i]).qfrc_actuator,
+          d.qfrc_actuator[m.jnt_dofadr[i] : m.jnt_dofadr[i] + dofnum[i]],
+          decimal=6,
+      )
 
     np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
     for i in range(m.nu):
@@ -308,30 +315,49 @@ class SupportTest(parameterized.TestCase):
       np.testing.assert_array_equal(
           dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0]
       )
+    dx10 = dx.bind(mx, s.bodies[0:2]).set(
+        'xfrc_applied',
+        np.array([np.array([1, 1, 1, 1, 1, 1]), np.array([2, 2, 2, 2, 2, 2])]),
+    )
+    np.testing.assert_array_equal(
+        dx10.bind(mx, s.bodies[0]).xfrc_applied, [1, 1, 1, 1, 1, 1]
+    )
+    np.testing.assert_array_equal(
+        dx10.bind(mx, s.bodies[1]).xfrc_applied, [2, 2, 2, 2, 2, 2]
+    )
+    dx11 = dx.bind(mx, s.bodies[-1]).set(
+        'mocap_pos',
+        [1, 2, 3],
+    )
+    np.testing.assert_array_equal(
+        dx11.bind(mx, s.bodies[-2]).mocap_pos, [100, 110, 120]
+    )
+    np.testing.assert_array_equal(
+        dx11.bind(mx, s.bodies[-1]).mocap_pos, [1, 2, 3]
+    )
 
-    # test invalid name
+    # test attribute and type mismatches
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='ctrl is not available for this type'
+        AttributeError, 'ctrl is not available for this type'
     ):
       print(dx.bind(mx, s.geoms).ctrl)
-    with self.assertRaises(
+    with self.assertRaises(KeyError):
-        KeyError, msg='actuator_actuator_ctrl'
-    ):
       print(dx.bind(mx, s.actuators).actuator_ctrl)
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='actuator_actuator_ctrl'
+        AttributeError,
+        "'Data' object has no attribute 'actuator_actuator_ctrl'",
     ):
       print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
-    with self.assertRaises(
+    with self.assertRaisesRegex(
-        AttributeError, msg='qpos, qvel, qacc are not available for this type'
+        AttributeError, 'qpos, qvel, qacc, qfrc are not available for this type'
     ):
       print(dx.bind(mx, s.geoms).qpos)
-    with self.assertRaises(KeyError, msg='invalid name: invalid_actuator_name'):
+
-      s.actuators[0].name = 'invalid_actuator_name'
+    # test that modified names do not raise an error
-      print(dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3]))
+    s.actuators[0].name = 'modified_actuator_name'
-    with self.assertRaises(KeyError, msg='invalid name: invalid_geom_name'):
+    np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
-      s.geoms[0].name = 'invalid_geom_name'
+    s.geoms[0].name = 'modified_geom_name'
-      print(mx.bind(s.geoms).pos)
+    np.testing.assert_array_equal(mx.bind(s.geoms[0]).pos, m.geom_pos[0, :])
 
     # test batched data
     batch_size = 16
@@ -343,12 +369,15 @@ class SupportTest(parameterized.TestCase):
           vdx.bind(mx, s.bodies[i]).xpos, [d.xpos[i, :]] * batch_size
       )
 
-    # test emtpy name
+    # test that adding a body requires recompilation
     s.worldbody.add_body()
-    m = s.compile()
+    with self.assertRaises(ValueError) as e:
-    mx = mjx.put_model(m)
-    with self.assertRaises(KeyError, msg='cannot bind spec with empty name'):
       mx.bind(s.bodies)
+    self.assertEqual(
+        str(e.exception),
+        'mjSpec signature does not match mjx.Model signature:'
+        ' 15297169659434471387 != 2785811613804955188',
+    )
 
   _CONTACTS = """
     <mujoco>
@@ -383,7 +412,7 @@ class SupportTest(parameterized.TestCase):
 
     # map MJX contacts to MJ ones
     def _find(g):
-      val = (g == dx.contact.geom).sum(axis=1)
+      val = (g == dx._impl.contact.geom).sum(axis=1)
       return np.where(val == 2)[0][0]
 
     contact_id_map = {i: _find(d.contact.geom[i]) for i in range(d.ncon)}
@@ -397,7 +426,7 @@ class SupportTest(parameterized.TestCase):
       np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
 
       # check for zeros after first condim elements
-      condim = dx.contact.dim[j]
+      condim = dx._impl.contact.dim[j]
       if condim < 6:
         np.testing.assert_allclose(force[condim:], 0, rtol=1e-5, atol=1e-5)
 
@@ -410,8 +439,8 @@ class SupportTest(parameterized.TestCase):
           ),
       )(mx, dx, j, True)
       # back to contact frame
-      force = force.at[:3].set(dx.contact.frame[j] @ force[:3])
+      force = force.at[:3].set(dx._impl.contact.frame[j] @ force[:3])
-      force = force.at[3:].set(dx.contact.frame[j] @ force[3:])
+      force = force.at[3:].set(dx._impl.contact.frame[j] @ force[3:])
       np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
 
   def test_wrap_inside(self):

mjx/mujoco/mjx/_src/test_util.py

@@ -106,20 +106,20 @@ def benchmark(
 
 
 def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray:
-  """Returns a sort order such that dx.efc_*[order][:d.nefc] == d.efc_*."""
+  """Returns a sort order such that dx.efc_*[order][:d._impl.nefc] == d.efc_*."""  # pytype: disable=attribute-error
   # reorder efc rows to skip inactive constraints and match contact order
-  efl = dx.ne + dx.nf + dx.nl
+  efl = dx._impl.ne + dx._impl.nf + dx._impl.nl  # pytype: disable=attribute-error
   order = np.arange(efl)
-  order[(dx.efc_J[:efl] == 0).all(axis=1)] = 2**16  # move empty rows to end
+  order[(dx._impl.efc_J[:efl] == 0).all(axis=1)] = 2**16  # move empty rows to end  # pytype: disable=attribute-error
-  for i in range(dx.ncon):
+  for i in range(dx._impl.ncon):  # pytype: disable=attribute-error
-    num_rows = dx.contact.dim[i]
+    num_rows = dx._impl.contact.dim[i]  # pytype: disable=attribute-error
-    if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
+    if dx._impl.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:  # pytype: disable=attribute-error
-      num_rows = (dx.contact.dim[i] - 1) * 2
+      num_rows = (dx._impl.contact.dim[i] - 1) * 2  # pytype: disable=attribute-error
-    if dx.contact.dist[i] > 0:  # move empty contacts to end
+    if dx._impl.contact.dist[i] > 0:  # move empty contacts to end  # pytype: disable=attribute-error
       order = np.append(order, np.repeat(2**16, num_rows))
       continue
-    contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1)
+    contact_match = (d.contact.geom == dx._impl.contact.geom[i]).all(axis=-1)  # pytype: disable=attribute-error
-    contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1)
+    contact_match &= (d.contact.pos == dx._impl.contact.pos[i]).all(axis=-1)  # pytype: disable=attribute-error
     assert contact_match.any(), f'contact {i} not found'
     contact_id = np.nonzero(contact_match)[0][0]
     order = np.append(order, np.repeat(efl + contact_id, num_rows))

mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml

@@ -0,0 +1,55 @@
+<mujoco model="fixed_site">
+  <worldbody>
+    <body>
+      <joint name="joint0" type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <site name="site0" pos="0.25 0 0.1" size="0.025"/>
+      <body pos="0.5 0 0">
+        <joint name="joint1" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <site name="site1" pos="0.25 0 0.1" size="0.025"/>
+        <body pos="0.5 0 0">
+          <joint name="joint2" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+          <site name="site2" pos="0.25 0 0.1" size="0.025"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <tendon>
+    <spatial name="spatial0" width="0.0125">
+      <site site="site0"/>
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+    <spatial name="spatial1" width="0.0125" actuatorfrclimited="true" actuatorfrcrange="-1 1">
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+    <fixed name="fixed0" actuatorfrclimited="true" actuatorfrcrange="-10 10">
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint1" coef=".2"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+    <fixed name="fixed1" actuatorfrclimited="true" actuatorfrcrange="0 1">
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+  </tendon>
+  <actuator>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial0"/>
+    <motor tendon="spatial1"/>
+    <motor tendon="spatial1"/>
+    <motor tendon="fixed0"/>
+    <motor tendon="fixed1"/>
+    <motor tendon="fixed1"/>
+  </actuator>
+  <sensor>
+    <tendonactuatorfrc tendon="spatial0"/>
+    <tendonactuatorfrc tendon="spatial1"/>
+    <tendonactuatorfrc tendon="fixed0"/>
+    <tendonactuatorfrc tendon="fixed1"/>
+  </sensor>
+</mujoco>

mjx/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-mjx"
-version = "3.3.1"
+version = "3.3.3"
 authors = [
     {name = "Google DeepMind", email = "mujoco@deepmind.com"},
 ]
@@ -30,7 +30,7 @@ dependencies = [
     "etils[epath]",
     "jax",
     "jaxlib",
-    "mujoco>=3.3.1.dev0",
+    "mujoco>=3.3.3.dev0",
     "scipy",
     "trimesh",
 ]
@@ -41,9 +41,9 @@ mjx-viewer = "mujoco.mjx.viewer:main"
 
 [project.urls]
 Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
-Documentation = "https://mujoco.readthedocs.io/en/3.3.1"
+Documentation = "https://mujoco.readthedocs.io/en/3.3.3"
 Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
-Changelog = "https://mujoco.readthedocs.io/en/3.3.1/changelog.html"
+Changelog = "https://mujoco.readthedocs.io/en/3.3.3/changelog.html"
 
 [tool.isort]
 force_single_line = true

mjx/requirements.txt

@@ -76,8 +76,9 @@ scipy==1.13.1; python_version == '3.9' \
     --hash=sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004 \
     --hash=sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24 \
     --hash=sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5
-setuptools==75.5.0 \
+setuptools==78.1.1 \
-    --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
+    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
+    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
 trimesh==4.5.2 \
     --hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73
 wheel==0.45.0 \

model/flex/bunny.xml

@@ -31,7 +31,7 @@
     </body>
     <flexcomp type="mesh" file="bunny.obj" pos="0 0 .1" dim="2" euler="90 0 0"
               radius=".001" rgba="0 .7 .7 1" mass=".05" name="softbody" dof="trilinear">
-      <elasticity young="1e3" poisson="0.1" damping="0.001"/>
+      <elasticity young="1e3" poisson="0.1" damping="0.001" elastic2d="stretch"/>
       <contact selfcollide="none" internal="false"/>
     </flexcomp>
   </worldbody>

model/flex/bunny_with_uv.xml

@@ -37,7 +37,7 @@
     </body>
     <flexcomp type="mesh" file="bunny_with_uv.obj" pos="0 0 .1" dim="2" euler="90 0 0"
               radius=".001" material="matsponge" mass=".05" name="softbody" dof="trilinear">
-      <elasticity young="1e3" poisson="0.1" damping="0.001"/>
+      <elasticity young="1e3" poisson="0.1" damping="0.001" elastic2d="stretch"/>
       <contact selfcollide="none" internal="false"/>
     </flexcomp>
   </worldbody>

model/flex/flag.xml

@@ -26,10 +26,6 @@
   <option wind="5 5 0" density="10" solver="CG" tolerance="1e-6">
   </option>
 
-  <extension>
-    <plugin plugin="mujoco.elasticity.shell"/>
-  </extension>
-
   <worldbody>
     <geom name="floor" type="plane" size="0 0 .1"/>
     <light diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
@@ -37,12 +33,7 @@
       <flexcomp type="grid" count="9 19 1" spacing=".05 .05 .05" mass="10"
                 name="flag" radius="0.001">
         <edge equality="true" damping="0.001"/>
-        <plugin plugin="mujoco.elasticity.shell">
+        <elasticity poisson="0" thickness="1e-2" young="3e6" elastic2d="none"/>
-          <config key="poisson" value="0"/>
-          <config key="thickness" value="1e-2"/>
-          <!--Units are in Pa (SI)-->
-          <config key="young" value="3e6"/>
-        </plugin>
       </flexcomp>
     </body>
   </worldbody>

model/flex/pancake.xml

@@ -16,10 +16,6 @@
 <mujoco model="Pancake">
   <include file="scene.xml"/>
 
-  <extension>
-    <plugin plugin="mujoco.elasticity.shell"/>
-  </extension>
-
   <option timestep="0.001" solver="CG" tolerance="1e-6" integrator="implicitfast"/>
 
   <size memory="10M"/>
@@ -39,12 +35,7 @@
               radius=".01" mass=".5" name="plate" dim="2">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
       <edge equality="true" damping="10"/>
-      <plugin plugin="mujoco.elasticity.shell">
+      <elasticity poisson="0" thickness="8e-3" young="3e5" elastic2d="bend"/>
-        <config key="poisson" value="0"/>
-        <config key="thickness" value="8e-3"/>
-        <!--Units are in Pa (SI)-->
-        <config key="young" value="3e5"/>
-      </plugin>
     </flexcomp>
   </worldbody>
 </mujoco>

model/flex/plate.xml

@@ -16,10 +16,6 @@
 <mujoco model="Plate">
   <include file="scene.xml"/>
 
-  <extension>
-    <plugin plugin="mujoco.elasticity.shell"/>
-  </extension>
-
   <option solver="CG" tolerance="1e-6" integrator="implicitfast"/>
 
   <size memory="10M"/>
@@ -42,12 +38,7 @@
               radius=".001" mass="100" name="plate">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
       <edge equality="true" damping="0.001"/>
-      <plugin plugin="mujoco.elasticity.shell">
+      <elasticity poisson="0" thickness="1e-2" young="3e7" elastic2d="bend"/>
-        <config key="poisson" value="0"/>
-        <config key="thickness" value="1e-2"/>
-        <!--Units are in Pa (SI)-->
-        <config key="young" value="3e7"/>
-      </plugin>
     </flexcomp>
   </worldbody>
 </mujoco>

model/flex/poncho.xml

@@ -19,10 +19,6 @@
     <flag energy="enable"/>
   </option>
 
-  <extension>
-    <plugin plugin="mujoco.elasticity.shell"/>
-  </extension>
-
   <visual>
     <map force="0.1" zfar="30"/>
     <rgba haze="0.15 0.25 0.35 1"/>
@@ -1418,13 +1414,8 @@
                 398 399 418
                 398 376 378">
       <edge equality="true" damping="0.1"/>
+      <elasticity young="3e5" poisson="0" thickness="8e-3" elastic2d="bend"/>
       <contact solref="0.003"/>
-      <plugin plugin="mujoco.elasticity.shell">
-        <config key="poisson" value="0"/>
-        <config key="thickness" value="8e-3"/>
-        <!--Units are in Pa (SI)-->
-        <config key="young" value="3e5"/>
-      </plugin>
     </flexcomp>
 
   </worldbody>

model/flex/trampoline.xml

@@ -39,7 +39,7 @@
               radius=".001" mass="10" name="plate" dim="2">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
       <edge equality="false" damping="10"/>
-      <elasticity young="3e5" poisson="0" thickness="1e-2"/>
+      <elasticity young="3e5" poisson="0" thickness="1e-2" elastic2d="stretch"/>
       <pin id="0 15 240 255"/>
     </flexcomp>
   </worldbody>

model/plugin/sensor/touch_stress.xml

@@ -0,0 +1,64 @@
+<mujoco model="touchtest">
+  <compiler autolimits="true"/>
+
+  <extension>
+    <plugin plugin="mujoco.sdf.gear">
+      <instance name="gear">
+        <config key="alpha" value="0"/>
+      </instance>
+    </plugin>
+    <plugin plugin="mujoco.sensor.touch_stress">
+      <instance name="touch_stress">
+        <config key="size" value="37 37"/>
+        <config key="fov" value="45 45"/>
+        <config key="gamma" value="0"/>
+        <config key="nchannel" value="3"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <asset>
+    <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4"
+             width="300" height="300" mark="edge" markrgb=".2 .3 .4"/>
+    <material name="grid" texture="grid" texrepeat="3 1" texuniform="true"/>
+    <mesh name="gear">
+      <plugin instance="gear"/>
+    </mesh>
+  </asset>
+
+  <visual>
+      <headlight ambient=".7 .7 .7" diffuse=".2 .2 .2" specular="0.1 0.1 0.1"/>
+      <map znear="0.01"/>
+      <scale contactwidth=".02" contactheight=".5"/>
+  </visual>
+
+  <default>
+    <geom friction="0.4" solimp="0 0.95 0.02"/>
+  </default>
+
+  <statistic center="0 0 1" extent="1" meansize=".1"/>
+
+  <worldbody>
+    <light pos="1 0 .3" dir="-1 0 -.3"/>
+    <light pos="-1 0 .3" dir="1 0 -.3"/>
+    <geom name="floor" pos="0 0 -0.01" type="plane" size="3 3 .01"/>
+
+    <geom type="sdf" name="gear" mesh="gear" rgba="0.4 0.4 0.4 1" pos="-1.3 0 0">
+      <plugin instance="gear"/>
+    </geom>
+
+    <body name="ball" pos="0 0 1">
+      <joint name="x" type="slide" axis="1 0 0" damping="1"/>
+      <joint name="y" type="slide" axis="0 1 0" damping="1"/>
+      <joint name="z" type="slide" axis="0 0 1"/>
+      <joint name="rx" axis="1 0 0" springdamper="0.2 1"/>
+      <joint name="ry" axis="0 1 0" springdamper="0.2 1"/>
+      <geom type="sphere" size=".3" mass="0.1" rgba=".5 .5 .5 .3"/>
+      <site name="touch"/>
+    </body>
+  </worldbody>
+
+  <sensor>
+    <plugin instance="touch_stress" objtype="site" objname="touch"/>
+  </sensor>
+</mujoco>

plugin/elasticity/CMakeLists.txt

@@ -22,8 +22,6 @@ set(MUJOCO_ELASTICITY_SRCS
     elasticity.cc
     elasticity.h
     register.cc
-    shell.cc
-    shell.h
 )
 
 add_library(elasticity SHARED)

plugin/elasticity/README.md

@@ -18,15 +18,3 @@ Parameters:
    - `bend` [Pa]: bending stiffness.
    - `flat` [bool]: if true, the stress-equilibrium configuration is that of a straight cable; if false or unspecified, it is the configuration defined in the XML.
    - `vmax` [N/m^2]: If greater than zero, the cable is colored using mechanical stresses; the value represent the maximum stress in the color scale.
-
-### Shell
-
-Implemented in [shell.cc](shell.cc).
-
-The shell plugin discretizes an inextensible 2D continuum. It is intended to simulate the bending of plates where the stretching is negligible compared to other deformation modes.
-
-Parameters:
-
-   - `young` [Pa]: Young's modulus.
-   - `poisson` [Pa]: Poisson's ratio; if 0, then the material only opposed shear deformations; if near 0.5, then the material is nearly incompressible (rubber-like).
-   - `thickness` [m]: shell thickness, used to scale the bending stiffness.

plugin/elasticity/register.cc

@@ -14,13 +14,11 @@
 
 #include <mujoco/mjplugin.h>
 #include "cable.h"
-#include "shell.h"
 
 namespace mujoco::plugin::elasticity {
 
 mjPLUGIN_LIB_INIT {
   Cable::RegisterPlugin();
-  Shell::RegisterPlugin();
 }
 
 }  // namespace mujoco::plugin::elasticity

plugin/elasticity/shell.cc

@@ -1,271 +0,0 @@
-// Copyright 2023 DeepMind Technologies Limited
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include <algorithm>
-#include <cassert>
-#include <cmath>
-#include <cstdint>
-#include <cstdlib>
-#include <optional>
-#include <string>
-#include <unordered_map>
-#include <utility>
-#include <vector>
-
-#include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
-#include <mujoco/mujoco.h>
-#include "elasticity.h"
-#include "shell.h"
-
-
-namespace mujoco::plugin::elasticity {
-namespace {
-
-// local tetrahedron numbering
-constexpr int kNumEdges = Stencil2D::kNumEdges;
-constexpr int kNumVerts = Stencil2D::kNumVerts;
-constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}};
-
-// cotangent between two edges
-mjtNum cot(mjtNum* x, int v0, int v1, int v2) {
-  mjtNum normal[3];
-  mjtNum edge1[3];
-  mjtNum edge2[3];
-
-  mju_sub3(edge1, x+3*v1, x+3*v0);
-  mju_sub3(edge2, x+3*v2, x+3*v0);
-  mju_cross(normal, edge1, edge2);
-
-  return mju_dot3(edge1, edge2) / mju_norm3(normal);
-}
-
-// area of a triangle
-mjtNum ComputeVolume(const mjtNum* x, const int v[kNumVerts]) {
-  mjtNum normal[3];
-  mjtNum edge1[3];
-  mjtNum edge2[3];
-
-  mju_sub3(edge1, x+3*v[1], x+3*v[0]);
-  mju_sub3(edge2, x+3*v[2], x+3*v[0]);
-  mju_cross(normal, edge1, edge2);
-
-  return mju_norm3(normal) / 2;
-}
-
-}  // namespace
-
-// factory function
-std::optional<Shell> Shell::Create(const mjModel* m, mjData* d, int instance) {
-    if (CheckAttr("face", m, instance) &&
-        CheckAttr("edge", m, instance) &&
-        CheckAttr("poisson", m, instance) &&
-        CheckAttr("young", m, instance) &&
-        CheckAttr("thickness", m, instance)) {
-      mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr);
-      mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
-      mjtNum thick =
-          strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
-      std::vector<int> face, edge;
-      String2Vector(mj_getPluginConfig(m, instance, "face"), face);
-      String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
-      return Shell(m, d, instance, nu, E, thick, face, edge);
-    } else {
-      mju_warning("Invalid parameter specification in shell plugin");
-      return std::nullopt;
-    }
-}
-
-// create map from triangles to vertices and edges and from edges to vertices
-void Shell::CreateStencils(const std::vector<int>& simplex,
-                           const std::vector<int>& edgeidx) {
-  // populate stencil
-  nt = simplex.size() / kNumVerts;
-  elements.resize(nt);
-  for (int t = 0; t < nt; t++) {
-    for (int v = 0; v < kNumVerts; v++) {
-      elements[t].vertices[v] = simplex[kNumVerts*t+v];
-    }
-  }
-
-  // map from edge vertices to their index in `edges` vector
-  std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
-
-  // loop over all triangles
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
-
-    // compute edges to vertices map for fast computations
-    for (int e = 0; e < kNumEdges; e++) {
-      auto pair = std::pair(
-        std::min(v[edge[e][0]], v[edge[e][1]]),
-        std::max(v[edge[e][0]], v[edge[e][1]])
-      );
-
-      // if edge is already present in the vector only store its index
-      auto [it, inserted] = edge_indices.insert({pair, ne});
-
-      if (inserted) {
-        StencilFlap flap;
-        flap.vertices[0] = v[edge[e][0]];
-        flap.vertices[1] = v[edge[e][1]];
-        flap.vertices[2] = v[(edge[e][1]+1) % 3];
-        flap.vertices[3] = -1;
-        flaps.push_back(flap);
-        elements[t].edges[e] = ne++;
-      } else {
-        elements[t].edges[e] = it->second;
-        flaps[it->second].vertices[3] = v[(edge[e][1]+1) % 3];
-      }
-
-      if (!edgeidx.empty()) {
-        assert(elements[t].edges[e] == edgeidx[kNumEdges*t+e]);
-      }
-    }
-  }
-}
-
-// plugin constructor
-Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-             mjtNum thick, const std::vector<int>& face,
-             const std::vector<int>& edgeidx)
-    : thickness(thick) {
-  // count plugin bodies
-  nv = ne = 0;
-  for (int i = 1; i < m->nbody; i++) {
-    if (m->body_plugin[i] == instance) {
-      if (!nv++) {
-        i0 = i;
-      }
-    }
-  }
-
-  // generate triangles from the vertices
-  CreateStencils(face, edgeidx);
-
-  // material parameters
-  mjtNum mu = E / (2*(1+nu));
-
-  // loop over all triangles
-  for (int t = 0; t < nt; t++) {
-    int* v = elements[t].vertices;
-    for (int i = 0; i < kNumVerts; i++) {
-      if (m->body_plugin[i0+v[i]] != instance) {
-        mju_error("This body does not have the requested plugin instance");
-      }
-    }
-  }
-
-  // allocate array
-  position.assign(nv*3, 0);
-  bending.assign(ne*16, 0);
-
-  // store previous positions
-  mju_copy(position.data(), m->body_pos+3*i0, 3*nv);
-
-  // assemble bending Hessian
-  for (int e = 0; e < ne; e++)  {
-    int* v = flaps[e].vertices;
-    int vadj[3] = {v[1], v[0], v[3]};
-
-    if (v[3]== -1) {
-      // skip boundary edges
-      continue;
-    }
-
-    // cotangent operator from Wardetzky at al., "Discrete Quadratic Curvature
-    // Energies", https://cims.nyu.edu/gcl/papers/wardetzky2007dqb.pdf
-
-    mjtNum a01 = cot(m->body_pos+3*i0, v[0], v[1], v[2]);
-    mjtNum a02 = cot(m->body_pos+3*i0, v[0], v[3], v[1]);
-    mjtNum a03 = cot(m->body_pos+3*i0, v[1], v[2], v[0]);
-    mjtNum a04 = cot(m->body_pos+3*i0, v[1], v[0], v[3]);
-    mjtNum c[4] = {a03 + a04, a01 + a02, -(a01 + a03), -(a02 + a04)};
-    mjtNum volume = ComputeVolume(m->body_pos+3*i0, v) +
-                    ComputeVolume(m->body_pos+3*i0, vadj);
-
-    for (int v1 = 0; v1 < StencilFlap::kNumVerts; v1++) {
-      for (int v2 = 0; v2 < StencilFlap::kNumVerts; v2++) {
-        bending[16 * e + 4 * v1 + v2] +=
-            1.5 * c[v1] * c[v2] / volume * mu * pow(thickness, 3) / 12;
-      }
-    }
-  }
-}
-
-void Shell::Compute(const mjModel* m, mjData* d, int instance) {
-  for (int e = 0; e < ne; e++) {
-    int* v = flaps[e].vertices;
-    mjtNum force[12] = {0};
-    if (v[3] == -1) {
-      // skip boundary edges
-      continue;
-    }
-    for (int i = 0; i < StencilFlap::kNumVerts; i++) {
-      for (int j = 0; j < StencilFlap::kNumVerts; j++) {
-        for (int x = 0; x < 3; x++) {
-          force[3*i+x] += bending[16*e+4*i+j] * d->xpos[3*(i0+v[j])+x];
-        }
-      }
-    }
-
-    // update stored positions
-    mju_copy(position.data(), d->xpos+3*i0, 3*nv);
-
-    // insert into global force
-    for (int i = 0; i < StencilFlap::kNumVerts; i++) {
-      for (int x = 0; x < 3; x++) {
-        d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
-      }
-    }
-  }
-}
-
-
-
-void Shell::RegisterPlugin() {
-  mjpPlugin plugin;
-  mjp_defaultPlugin(&plugin);
-
-  plugin.name = "mujoco.elasticity.shell";
-  plugin.capabilityflags |= mjPLUGIN_PASSIVE;
-
-  const char* attributes[] = {"face",    "edge",    "young",
-                              "poisson", "thickness", "damping"};
-  plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
-  plugin.attributes = attributes;
-  plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
-
-  plugin.init = +[](const mjModel* m, mjData* d, int instance) {
-    auto elasticity_or_null = Shell::Create(m, d, instance);
-    if (!elasticity_or_null.has_value()) {
-      return -1;
-    }
-    d->plugin_data[instance] = reinterpret_cast<uintptr_t>(
-        new Shell(std::move(*elasticity_or_null)));
-    return 0;
-  };
-  plugin.destroy = +[](mjData* d, int instance) {
-    delete reinterpret_cast<Shell*>(d->plugin_data[instance]);
-    d->plugin_data[instance] = 0;
-  };
-  plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) {
-    auto* elasticity = reinterpret_cast<Shell*>(d->plugin_data[instance]);
-    elasticity->Compute(m, d, instance);
-  };
-
-  mjp_registerPlugin(&plugin);
-}
-
-}  // namespace mujoco::plugin::elasticity

plugin/elasticity/shell.h

@@ -1,74 +0,0 @@
-// Copyright 2023 DeepMind Technologies Limited
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef MUJOCO_PLUGIN_ELASTICITY_SHELL_H_
-#define MUJOCO_PLUGIN_ELASTICITY_SHELL_H_
-
-#include <optional>
-#include <vector>
-
-#include <mujoco/mjdata.h>
-#include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
-#include "elasticity.h"
-
-
-namespace mujoco::plugin::elasticity {
-
-struct StencilFlap {
-  static constexpr int kNumVerts = 4;
-  int vertices[kNumVerts];
-};
-
-class Shell {
- public:
-  // Returns a new Shell instance or nullopt on failure.
-  static std::optional<Shell> Create(const mjModel* m, mjData* d,
-                                     int instance);
-  Shell(Shell&&) = default;
-
-  Shell& operator=(Shell&& other) = default;
-
-  void Compute(const mjModel* m, mjData* d, int instance);
-
-  static void RegisterPlugin();
-
-  int i0;  // index of first body
-  int nc;  // number of quads in the grid
-  int nv;  // number of vertices (bodies) in the Shell
-  int nt;  // number of area elements (triangles)
-  int ne;  // number of edges in the Shell
-
-  // connectivity info for mapping tetrahedra to edges and vertices
-  std::vector<Stencil2D> elements;          // triangles              (nt x 6)
-  std::vector<StencilFlap> flaps;           // adjacent triangles     (ne x 4)
-
-  // precomputed quantities
-  std::vector<mjtNum> position;             // previous-step positions (nv x 3)
-  std::vector<mjtNum> bending;              // bending Hessian         (ne x 16)
-
-  mjtNum thickness;
-
- private:
-  Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
-        mjtNum thick, const std::vector<int>& face,
-        const std::vector<int>& edgeidx);
-
-  void CreateStencils(const std::vector<int>& simplex,
-                      const std::vector<int>& edgeidx);
-};
-
-}  // namespace mujoco::plugin::elasticity
-
-#endif  // MUJOCO_PLUGIN_ELASTICITY_SHELL_H_

plugin/sdf/CMakeLists.txt

@@ -34,19 +34,19 @@ set(MUJOCO_SDF_SRCS
     torus.h
 )
 
-add_library(sdf SHARED)
+add_library(sdf_plugin SHARED)
-target_sources(sdf PRIVATE ${MUJOCO_SDF_SRCS})
+target_sources(sdf_plugin PRIVATE ${MUJOCO_SDF_SRCS})
-target_include_directories(sdf PRIVATE ${MUJOCO_SDF_INCLUDE})
+target_include_directories(sdf_plugin PRIVATE ${MUJOCO_SDF_INCLUDE})
-target_link_libraries(sdf PRIVATE mujoco SdfLib)
+target_link_libraries(sdf_plugin PRIVATE mujoco SdfLib)
 target_compile_options(
-  sdf
+  sdf_plugin
   PRIVATE ${AVX_COMPILE_OPTIONS}
           ${MUJOCO_MACOS_COMPILE_OPTIONS}
           ${EXTRA_COMPILE_OPTIONS}
           ${MUJOCO_CXX_FLAGS}
 )
 target_link_options(
-  sdf
+  sdf_plugin
   PRIVATE
   ${MUJOCO_MACOS_LINK_OPTIONS}
   ${EXTRA_LINK_OPTIONS}

plugin/sensor/CMakeLists.txt

@@ -18,9 +18,11 @@ set(MUJOCO_SENSOR_INCLUDE
 )
 
 set(MUJOCO_SENSOR_SRCS
-    sensor.cc
+    register.cc
     touch_grid.h
     touch_grid.cc
+    touch_stress.h
+    touch_stress.cc
 )
 
 add_library(sensor SHARED)

plugin/sensor/README.md

@@ -9,6 +9,9 @@ plugins](https://mujoco.readthedocs.io/en/latest/programming/extension.html#engi
   - [Illustration of foveal deformation](#illustration-of-foveal-deformation)
   - [Illustration combining resolution, fields-of-view and foveal deformation](#illustration-combining-resolution-fields-of-view-and-foveal-deformation)
 
+- [Touch Stress](#touch-stress)
+  - [Example model with analytical SDF](#example-model-with-analytical-sdf)
+
 ## [Touch Grid](touch_grid.h)
 
 This sensor aggregates contact forces into "taxels": a rectangular array of pixel-like elements.
@@ -23,7 +26,7 @@ The output of the sensor is a stack of 1 to 6 "touch images" corresponding to fo
 and torques in the frame of the sensor. Forces and torques are in the in [z, x,
 y] order, corresponding to the ordering in contact frames: [normal, tangent,
 tangent] and [torsional, rolling, rolling]. Each "taxel" corresponds to an angular bin
-in spherical coordinates, and aggregates all the forces occuring inside this bin, which occur
+in spherical coordinates, and aggregates all the forces occurring inside this bin, which occur
 between the body in which the sensor's site is defined and any other body.
 
 The sensor is parametrized by 6 numbers:
@@ -84,3 +87,64 @@ See [touch_grid.xml](../../model/plugin/sensor/touch_grid.xml) to play with the
 ### Illustration combining resolution, fields-of-view and foveal deformation
 
 [![touch grid illustration](https://img.youtube.com/vi/YScjmR8LwQI/0.jpg)](https://www.youtube.com/watch?v=YScjmR8LwQI)
+
+## [Touch Stress](touch_stress.h)
+
+This sensor is based on similar concepts and parametrization as the `touch_grid`,
+while overcoming some of its limitations. In particular, the `touch_grid` can
+only provide sparse information, depending on the number of contact points
+generated. The `touch_stress` sensor can instead generate a high-resolution
+touch image. In order to do this, it requires a signed distance function (SDF)
+of the object that is in contact with the sensor. This is handled internally for
+primitives or it must be declared explicitly in the model using SDF plugins.
+
+There is one important difference with respect to the `touch_grid`: in this case,
+the force is computed in the local taxel frame and not in the frame of the sensor.
+This allows for a more intuitive interpretation of normal and tangential stresses,
+as shown in the images below.
+
+Note that in this case, the absolute values of the stresses reported by the
+sensor are unrelated to the contact forces. They are purely based on geometric
+and kinematic considerations, i.e. the SDF for the normal stress and the sliding
+velocity for the tangential contributions.
+
+### Example model with analytical SDF
+
+```xml
+<extension>
+  <plugin plugin="mujoco.sdf.gear">
+    <instance name="gear">
+      <config key="alpha" value="0"/>
+    </instance>
+  </plugin>
+  <plugin plugin="mujoco.sensor.touch_stress">
+    <instance name="touch_stress">
+      <config key="size" value="37 37"/>
+      <config key="fov" value="45 45"/>
+      <config key="gamma" value="0"/>
+      <config key="nchannel" value="3"/>
+    </instance>
+  </plugin>
+  ...
+  <asset>
+    <mesh name="gear">
+      <plugin instance="gear"/>
+    </mesh>
+  </asset>
+  ...
+  <worldbody>
+    <geom type="sdf" name="gear" mesh="gear">
+      <plugin instance="gear"/>
+    </geom>
+  ...
+  <sensor>
+    <plugin instance="touch_stress" objtype="site" objname="touch_site"/>
+  </sensor>
+</extension>
+```
+The images below show a static sphere over a gear described by an analytic SDF
+and the same sphere dragged along the x and y axes.
+
+<img src="images/normal.png" style="width: 300px;"/>
+<img src="images/tangential1.png" style="width: 300px;"/>
+<img src="images/tangential2.png" style="width: 300px;"/>

plugin/sensor/register.cc

@@ -14,9 +14,13 @@
 
 #include <mujoco/mjplugin.h>
 #include "touch_grid.h"
+#include "touch_stress.h"
 
 namespace mujoco::plugin::sensor {
 
-mjPLUGIN_LIB_INIT { TouchGrid::RegisterPlugin(); }
+mjPLUGIN_LIB_INIT {
+  TouchGrid::RegisterPlugin();
+  TouchStress::RegisterPlugin();
+}
 
 }  // namespace mujoco::plugin::sensor

plugin/sensor/touch_stress.cc

@@ -0,0 +1,557 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "touch_stress.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#include <sstream>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjplugin.h>
+#include <mujoco/mjtnum.h>
+#include <mujoco/mjvisualize.h>
+#include <mujoco/mujoco.h>
+
+namespace mujoco::plugin::sensor {
+
+namespace {
+
+// Checks that a plugin config attribute exists.
+bool CheckAttr(const std::string& input) {
+  char* end;
+  std::string value = input;
+  value.erase(std::remove_if(value.begin(), value.end(), isspace), value.end());
+  strtod(value.c_str(), &end);
+  return end == value.data() + value.size();
+}
+
+// Converts a string into a numeric vector
+template <typename T>
+void ReadVector(std::vector<T>& output, const std::string& input) {
+  std::stringstream ss(input);
+  std::string item;
+  char delim = ' ';
+  while (getline(ss, item, delim)) {
+    CheckAttr(item);
+    output.push_back(strtod(item.c_str(), nullptr));
+  }
+}
+
+// Evenly spaced numbers over a specified interval.
+void LinSpace(mjtNum lower, mjtNum upper, int n, mjtNum array[]) {
+  mjtNum increment = n > 1 ? (upper - lower) / (n - 1) : 0;
+  for (int i = 0; i < n; ++i) {
+    *array = lower;
+    ++array;
+    lower += increment;
+  }
+}
+
+// Parametrized linear/quintic interpolated nonlinearity.
+mjtNum Fovea(mjtNum x, mjtNum gamma) {
+  // Quick return.
+  if (!gamma) return x;
+
+  // Foveal deformation.
+  mjtNum g = mjMAX(0, mjMIN(1, gamma));
+  return g*mju_pow(x, 5) + (1 - g)*x;
+}
+
+// Make bin edges.
+void BinEdges(mjtNum* x_edges, mjtNum* y_edges, int size[2], mjtNum fov[2],
+              mjtNum gamma) {
+  // Make unit bin edges.
+  LinSpace(-1, 1, size[0] + 1, x_edges);
+  LinSpace(-1, 1, size[1] + 1, y_edges);
+
+  // Apply foveal deformation.
+  for (int i = 0; i < size[0] + 1; i++) {
+    x_edges[i] = Fovea(x_edges[i], gamma);
+  }
+  for (int i = 0; i < size[1] + 1; i++) {
+    y_edges[i] = Fovea(y_edges[i], gamma);
+  }
+
+  // Scale by field-of-view.
+  mju_scl(x_edges, x_edges, fov[0]*mjPI / 180, size[0] + 1);
+  mju_scl(y_edges, y_edges, fov[1]*mjPI / 180, size[1] + 1);
+}
+
+// Permute 3-vector from 0,1,2 to 2,0,1.
+static void xyz2zxy(mjtNum* x) {
+  mjtNum z = x[2];
+  x[2] = x[1];
+  x[1] = x[0];
+  x[0] = z;
+}
+
+// Transform spherical (azimuth, elevation, radius) to Cartesian (x,y,z).
+void SphericalToCartesian(const mjtNum aer[3], mjtNum xyz[3]) {
+  mjtNum a = aer[0], e = aer[1], r = aer[2];
+  xyz[0] = r * mju_cos(e) * mju_sin(a);
+  xyz[1] = r * mju_sin(e);
+  xyz[2] = -r * mju_cos(e) * mju_cos(a);
+}
+
+// Tangent frame in Cartesian coordinates.
+void TangentFrame(const mjtNum aer[3], mjtNum mat[9]) {
+  mjtNum a = aer[0], e = aer[1], r = aer[2];
+  mjtNum ta[3] = {r * mju_cos(e) * mju_cos(a), 0,
+                  r * mju_cos(e) * mju_sin(a)};
+  mjtNum te[3] = {-r * mju_sin(e) * mju_sin(a), r * mju_cos(e),
+                   r * mju_sin(e) * mju_cos(a)};
+  mju_normalize3(ta);
+  mju_normalize3(te);
+  mju_copy3(mat, ta);
+  mju_copy3(mat+3, te);
+  mju_cross(mat+6, te, ta);
+}
+
+}  // namespace
+
+// Creates a TouchStress instance if all config attributes are defined and
+// within their allowed bounds.
+TouchStress* TouchStress::Create(const mjModel* m, mjData* d,
+                                                 int instance) {
+  if (CheckAttr(std::string(mj_getPluginConfig(m, instance, "gamma"))) &&
+      CheckAttr(std::string(mj_getPluginConfig(m, instance, "nchannel")))) {
+    // nchannel
+    int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr);
+    if (!nchannel) nchannel = 1;
+    if (nchannel < 1 || nchannel > 3) {
+      mju_error("nchannel must be between 1 and 3");
+      return nullptr;
+    }
+
+    // size
+    std::vector<int> size;
+    std::string size_str = std::string(mj_getPluginConfig(m, instance, "size"));
+    ReadVector(size, size_str.c_str());
+    if (size.size()!= 2) {
+      mju_error("Both horizontal and vertical resolutions must be specified");
+      return nullptr;
+    }
+    if (size[0] <= 0 || size[1] <= 0) {
+      mju_error("Horizontal and vertical resolutions must be positive");
+      return nullptr;
+    }
+
+    // field of view
+    std::vector<mjtNum> fov;
+    std::string fov_str = std::string(mj_getPluginConfig(m, instance, "fov"));
+    ReadVector(fov, fov_str.c_str());
+    if (fov.size()!= 2) {
+      mju_error(
+          "Both horizontal and vertical fields of view must be specified");
+      return nullptr;
+    }
+    if (fov[0] <= 0 || fov[0] > 180) {
+      mju_error("`fov[0]` must be a float between (0, 180] degrees");
+      return nullptr;
+    }
+    if (fov[1] <= 0 || fov[1] > 90) {
+      mju_error("`fov[1]` must be a float between (0, 90] degrees");
+      return nullptr;
+    }
+
+    // gamma
+    mjtNum gamma = strtod(mj_getPluginConfig(m, instance, "gamma"), nullptr);
+    if (gamma < 0 || gamma > 1) {
+      mju_error("`gamma` must be a nonnegative float between [0, 1]");
+      return nullptr;
+    }
+
+    return new TouchStress(m, d, instance, nchannel, size.data(), fov.data(),
+                         gamma);
+  } else {
+    mju_error("Invalid or missing parameters in touch_grid sensor plugin");
+    return nullptr;
+  }
+}
+
+TouchStress::TouchStress(const mjModel* m, mjData* d, int instance,
+                         int nchannel, int size[2], mjtNum fov[2], mjtNum gamma)
+    : nchannel_(nchannel),
+      size_{size[0], size[1]},
+      fov_{fov[0], fov[1]},
+      gamma_(gamma) {
+  // Make sure sensor is attached to a site.
+  for (int i = 0; i < m->nsensor; ++i) {
+    if (m->sensor_type[i] == mjSENS_PLUGIN && m->sensor_plugin[i] == instance) {
+      if (m->sensor_objtype[i] != mjOBJ_SITE) {
+        mju_error("Touch Grid sensor must be attached to a site");
+      }
+    }
+  }
+
+  // Get sensor id.
+  for (id_ = 0; id_ < m->nsensor; ++id_) {
+    if (m->sensor_type[id_] == mjSENS_PLUGIN &&
+        m->sensor_plugin[id_] == instance) {
+      break;
+    }
+  }
+
+  // Get parent weld id.
+  int site_id = m->sensor_objid[id_];
+  int parent_body = m->body_weldid[m->site_bodyid[site_id]];
+  parent_weld_ = m->body_weldid[parent_body];
+
+  // Get geom id.
+  if (m->body_geomnum[parent_body] != 1) {
+    mju_error("Touch sensor must be attached to a body with exactly one geom");
+  }
+  geom_id_ = m->body_geomadr[parent_body];
+
+  // Create bin edges.
+  x_edges_.assign(size[0] + 1, 0);
+  y_edges_.assign(size[1] + 1, 0);
+  BinEdges(x_edges_.data(), y_edges_.data(), size_, fov_, gamma_);
+  dist_.resize(size[0]*size[1], 0);
+  pos_.resize(3*size[0]*size[1], 0);
+  mat_.resize(9*size[0]*size[1], 0);
+
+  // Precompute spherical coordinates.
+  for (int i = 0; i < size[0]; i++) {
+    for (int j = 0; j < size[1]; j++) {
+      mjtNum aer[3];
+      aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]);
+      aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]);
+      aer[2] = m->geom_size[3*geom_id_];
+      SphericalToCartesian(aer, pos_.data() + 3 * (i * size[1] + j));
+      dist_[i*size[1]+j] = mju_abs(aer[2]);
+      TangentFrame(aer, mat_.data() + 9 * (i * size[1] + j));
+    }
+  }
+}
+
+void TouchStress::Reset(const mjModel* m, int instance) {}
+
+void TouchStress::Compute(const mjModel* m, mjData* d, int instance) {
+  mj_markStack(d);
+
+  // Clear sensordata and distance matrix.
+  mjtNum* sensordata = d->sensordata + m->sensor_adr[id_];
+  mju_zero(sensordata, m->sensor_dim[id_]);
+
+  // Get site id.
+  int site_id = m->sensor_objid[id_];
+
+  // Count contacts and get contact geom ids.
+  std::unordered_set<int> contact_geom_ids;
+  for (int i = 0; i < d->ncon; i++) {
+    int body1 = m->body_weldid[m->geom_bodyid[d->contact[i].geom1]];
+    int body2 = m->body_weldid[m->geom_bodyid[d->contact[i].geom2]];
+    if (body1 == parent_weld_) {
+      contact_geom_ids.insert(d->contact[i].geom2);
+    }
+    if (body2 == parent_weld_) {
+      contact_geom_ids.insert(d->contact[i].geom1);
+    }
+  }
+
+  // No contacts, return.
+  if (contact_geom_ids.empty()) {
+    mj_freeStack(d);
+    return;
+  }
+
+  // All of the quadrature points are contact points.
+  int ncon = size_[0]*size_[1];
+
+  // Get site frame.
+  mjtNum* site_pos = d->site_xpos + 3*site_id;
+  mjtNum* site_mat = d->site_xmat + 9*site_id;
+
+  // Allocate contact forces and positions.
+  mjtNum* forces = mj_stackAllocNum(d, ncon*3);
+  mjtNum* forcesT = mj_stackAllocNum(d, ncon*3);
+
+  // Iterate over colliding geoms.
+  for (auto geom : contact_geom_ids) {
+    int body = m->geom_bodyid[geom];
+
+    // Get sdf plugin of the geoms.
+    int sdf_instance[2] = {-1, geom_id_};
+    mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE};
+    const mjpPlugin* sdf_ptr[2] = {NULL, NULL};
+    if (m->geom_type[geom] == mjGEOM_SDF) {
+      sdf_instance[0] = m->geom_plugin[geom];
+      sdf_ptr[0] = mjc_getSDF(m, geom);
+    } else {
+      sdf_instance[0] = geom;
+      geomtype[0] = (mjtGeom)m->geom_type[geom];
+    }
+
+    // Set SDF parameters.
+    mjSDF geom_sdf;
+    geom_sdf.id = &sdf_instance[0];
+    geom_sdf.type = mjSDFTYPE_SINGLE;
+    geom_sdf.plugin = &sdf_ptr[0];
+    geom_sdf.geomtype = &geomtype[0];
+
+    mjSDF sensor_sdf;
+    sensor_sdf.id = &sdf_instance[1];
+    sensor_sdf.type = mjSDFTYPE_SINGLE;
+    sensor_sdf.plugin = &sdf_ptr[1];
+    sensor_sdf.geomtype = &geomtype[1];
+
+    // Get forces and positions in spherical coordinates.
+    int node = 0;
+    for (int j = 0; j < size_[1]; j++) {
+      for (int i = 0; i < size_[0]; i++) {
+        // Position in site frame.
+        mjtNum* pos = pos_.data() + 3*(i*size_[1] + j);
+        mjtNum* mat = mat_.data() + 9*(i*size_[1] + j);
+
+        // Position in global frame.
+        mjtNum xpos[3];
+        mju_mulMatVec3(xpos, site_mat, pos);
+        mju_addTo3(xpos, site_pos);
+
+        // Position in other geom frame.
+        mjtNum lpos[3], tmp[3];
+        mju_sub3(tmp, xpos, d->geom_xpos + 3*geom);
+        mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp);
+
+        // Add mesh position if needed.
+        if (m->geom_type[geom] == mjGEOM_MESH ||
+            m->geom_type[geom] == mjGEOM_SDF) {
+          mjtNum mesh_mat[9];
+          mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]);
+          mju_mulMatVec3(lpos, mesh_mat, lpos);
+          mju_addTo3(lpos, m->mesh_pos + 3 * m->geom_dataid[geom]);
+        }
+
+        // Compute distance.
+        mjtNum depth = mju_min(mjc_distance(m, d, &geom_sdf, lpos), 0);
+        if (depth == 0) {
+          mju_zero3(forces + 3*node);
+          node++;
+          continue;
+        }
+
+        // Get velocity in global frame.
+        mjtNum vel_sensor[6], vel_other[6], vel_rel[3];
+        mju_transformSpatial(
+            vel_sensor, d->cvel + 6 * parent_weld_, 0, xpos,
+            d->subtree_com + 3 * m->body_rootid[parent_weld_], NULL);
+        mju_transformSpatial(
+            vel_other, d->cvel + 6 * body, 0, d->geom_xpos + 3 * geom,
+            d->subtree_com + 3 * m->body_rootid[body], NULL);
+        mju_sub3(vel_rel, vel_sensor+3, vel_other+3);
+
+        // Get contact force/torque, rotate into node frame.
+        mjtNum tmp_force[3], normal[3];
+        mjtNum kMaxDepth = 0.05;
+        mjtNum pressure = 1 / (kMaxDepth - depth) - 1 / kMaxDepth;
+        mjc_gradient(m, d, &sensor_sdf, normal, pos);
+        mju_scl3(tmp_force, normal, pressure);
+        mju_mulMatTVec3(forces + 3*node, mat, tmp_force);
+        forces[3*node+0] = mju_abs(mju_dot3(vel_rel, mat + 0));
+        forces[3*node+1] = mju_abs(mju_dot3(vel_rel, mat + 3));
+
+        // Permute forces from x,y,z to z,x,y (normal, tangent, tangent)
+        xyz2zxy(forces + 3*node);
+        node++;
+      }
+    }
+
+    // Transpose forces.
+    mju_transpose(forcesT, forces, ncon, 3);
+
+    // Compute sensor output.
+    for (int c = 0; c < nchannel_; c++) {
+      if (!mju_isZero(forcesT + c*ncon, ncon)) {
+        mju_addTo(sensordata + c*ncon, forcesT + c*ncon, size_[0]*size_[1]);
+      }
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+// Thickness of taxel-visualization boxes relative to contact distance.
+static const mjtNum kRelativeThickness = 0.02;
+
+void TouchStress::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
+                             mjvScene* scn, int instance) {
+  mj_markStack(d);
+
+  // Get sensor data.
+  mjtNum* sensordata = d->sensordata + m->sensor_adr[id_];
+
+  // Get maximum absolute normal force.
+  mjtNum maxval = 0;
+  int frame = size_[0]*size_[1];
+  for (int j=0; j < frame; j++) {
+    maxval = mju_max(maxval, mju_abs(sensordata[j]));
+  }
+
+  // If no normal force readings, quick return.
+  if (!maxval) {
+    mj_freeStack(d);
+    return;
+  }
+
+  // Get site id and frame.
+  int site_id = m->sensor_objid[id_];
+  mjtNum* site_pos = d->site_xpos + 3*site_id;
+  mjtNum* site_mat = d->site_xmat + 9*site_id;
+  mjtNum site_quat[4];
+  mju_mat2Quat(site_quat, site_mat);
+
+  // Draw geoms.
+  for (int i=0; i < size_[0]; i++) {
+    for (int j=0; j < size_[1]; j++) {
+      mjtNum dist = dist_[i*size_[1]+j];
+      if (!dist) {
+        continue;
+      }
+      if (scn->ngeom >= scn->maxgeom) {
+        mj_warning(d, mjWARN_VGEOMFULL, scn->maxgeom);
+        mj_freeStack(d);
+        return;
+      } else {
+        // size
+        mjtNum size[3];
+        size[0] = dist*0.5*(x_edges_[i+1]-x_edges_[i]);
+        size[1] = dist*0.5*(y_edges_[j+1]-y_edges_[j]);
+        size[2] = dist*kRelativeThickness;
+
+        // position
+        mjtNum pos[3];
+        mjtNum aer[3];
+        aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]);
+        aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]);
+        aer[2] = dist*(1-kRelativeThickness);
+        SphericalToCartesian(aer, pos);
+        mju_mulMatVec3(pos, site_mat, pos);
+        mju_addTo3(pos, site_pos);
+
+        // orientation
+        mjtNum a_quat[4];
+        mjtNum site_y[3] = {-site_mat[1], -site_mat[4], -site_mat[7]};
+        mju_axisAngle2Quat(a_quat, site_y, aer[0]);
+        mjtNum e_quat[4];
+        mjtNum site_x[3] = {site_mat[0], site_mat[3], site_mat[6]};
+        mju_axisAngle2Quat(e_quat, site_x, aer[1]);
+        mjtNum quat[4];
+        mju_mulQuat(quat, e_quat, site_quat);
+        mju_mulQuat(quat, a_quat, quat);
+        mjtNum mat[9];
+        mju_quat2Mat(mat, quat);
+
+        // color
+        float rgba[4] = {1, 1, 1, 1.0};
+        for (int k=0; k < mjMIN(nchannel_, 3); k++) {
+          rgba[k] = mju_abs(sensordata[k*frame + j*size_[0] + i]) / maxval;
+        }
+
+        // draw box geom
+        mjvGeom* thisgeom = scn->geoms + scn->ngeom;
+        mjv_initGeom(thisgeom, mjGEOM_BOX, size, pos, mat, rgba);
+        thisgeom->objtype = mjOBJ_UNKNOWN;
+        thisgeom->objid = id_;
+        thisgeom->category = mjCAT_DECOR;
+        thisgeom->segid = scn->ngeom;
+        scn->ngeom++;
+      }
+    }
+  }
+
+  mj_freeStack(d);
+}
+
+
+void TouchStress::RegisterPlugin() {
+  mjpPlugin plugin;
+  mjp_defaultPlugin(&plugin);
+
+  plugin.name = "mujoco.sensor.touch_stress";
+  plugin.capabilityflags |= mjPLUGIN_SENSOR;
+
+  // Parameterized by 4 attributes.
+  const char* attributes[] = {"nchannel", "size", "fov", "gamma"};
+  plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
+  plugin.attributes = attributes;
+
+  // Stateless.
+  plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
+
+  // Sensor dimension = nchannel * size[0] * size[1]
+  plugin.nsensordata = +[](const mjModel* m, int instance, int sensor_id) {
+    int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr);
+    if (!nchannel) nchannel = 1;
+    std::vector<int> size;
+    std::string size_str = std::string(mj_getPluginConfig(m, instance, "size"));
+    ReadVector(size, size_str.c_str());
+    return nchannel * size[0] * size[1];
+  };
+
+  // Can only run after forces have been computed.
+  plugin.needstage = mjSTAGE_ACC;
+
+  // Initialization callback.
+  plugin.init = +[](const mjModel* m, mjData* d, int instance) {
+    auto* TouchStress = TouchStress::Create(m, d, instance);
+    if (!TouchStress) {
+      return -1;
+    }
+    d->plugin_data[instance] = reinterpret_cast<uintptr_t>(TouchStress);
+    return 0;
+  };
+
+  // Destruction callback.
+  plugin.destroy = +[](mjData* d, int instance) {
+    delete reinterpret_cast<TouchStress*>(d->plugin_data[instance]);
+    d->plugin_data[instance] = 0;
+  };
+
+  // Reset callback.
+  plugin.reset = +[](const mjModel* m, mjtNum* plugin_state, void* plugin_data,
+                     int instance) {
+    auto* TouchStress = reinterpret_cast<class TouchStress*>(plugin_data);
+    TouchStress->Reset(m, instance);
+  };
+
+  // Compute callback.
+  plugin.compute =
+      +[](const mjModel* m, mjData* d, int instance, int capability_bit) {
+        auto* TouchStress =
+            reinterpret_cast<class TouchStress*>(d->plugin_data[instance]);
+        TouchStress->Compute(m, d, instance);
+      };
+
+  // Visualization callback.
+  plugin.visualize = +[](const mjModel* m, mjData* d, const mjvOption* opt,
+                         mjvScene* scn, int instance) {
+    auto* TouchStress =
+        reinterpret_cast<class TouchStress*>(d->plugin_data[instance]);
+    TouchStress->Visualize(m, d, opt, scn, instance);
+  };
+
+  // Register the plugin.
+  mjp_registerPlugin(&plugin);
+}
+
+}  // namespace mujoco::plugin::sensor

plugin/sensor/touch_stress.h

@@ -0,0 +1,80 @@
+// Copyright 2025 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_
+#define MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_
+
+#include <vector>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
+#include <mujoco/mjvisualize.h>
+
+namespace mujoco::plugin::sensor {
+
+// A touch grid sensor is associated with a site and senses contact stresses
+// between the site's parent body and all other bodies. The site's
+// frame determines the orientation of the sensor with the same convention used
+// for cameras and lights: the sensor points in the frame's negative-Z
+// direction, so the X and Y axes correspond to horizontal and vertical
+// directions respectively.
+//
+// The output of the sensor is a stack of 3 "touch images" corresponding to
+// forces in the local frame of the taxels. Forces are in the [z, x, y] order,
+// corresponding to the ordering in contact frames: [normal, tangent, tangent].
+//
+// The sensor has 6 parameters:
+//  1. (int) Number of channels [1-3]. Defaults to 1.
+//  2. (int) Horizontal resolution.
+//  3. (int) Vertical resolution.
+//  4. (float) Horizontal field-of-view (fov_x), in degrees.
+//  5. (float) Vertical field-of-view (fov_y), in degrees.
+//  6. (float) Foveal deformation. Defaults to 0.
+class TouchStress {
+ public:
+  static TouchStress* Create(const mjModel* m, mjData* d, int instance);
+  TouchStress(TouchStress&&) = default;
+  ~TouchStress() = default;
+
+  void Reset(const mjModel* m, int instance);
+  void Compute(const mjModel* m, mjData* d, int instance);
+  void Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
+                 mjvScene* scn, int instance);
+
+  static void RegisterPlugin();
+
+  int nchannel_;           // number of channels (1-3)
+  int size_[2];            // horizontal and vertical resolution
+  mjtNum fov_[2];          // horizontal and vertical field of view, in degrees
+  mjtNum gamma_;           // foveal deformation
+
+ private:
+  TouchStress(const mjModel* m, mjData* d, int instance, int nchannel,
+               int* size, mjtNum* fov_x, mjtNum gamma);
+
+  std::vector<mjtNum> x_edges_;
+  std::vector<mjtNum> y_edges_;
+  std::vector<mjtNum> dist_;
+  std::vector<mjtNum> pos_;
+  std::vector<mjtNum> mat_;
+
+  int id_;
+  int parent_weld_;
+  int geom_id_;
+};
+
+}  // namespace mujoco::plugin::sensor
+
+#endif  // MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_

python/build_requirements.txt

@@ -57,8 +57,9 @@ PyOpenGL==3.1.7 \
     --hash=sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a
 pytest==8.3.3 \
     --hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2
-setuptools==75.5.0 \
+setuptools==78.1.1 \
-    --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
+    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
+    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
 wheel==0.45.0 \
     --hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7
 

python/make_sdist_requirements.txt

@@ -4,8 +4,9 @@ build==1.2.2.post1 \
     --hash=sha256:1d61c0887fa860c01971625baae8bdd338e517b836a2f70dd1f7aa3a6b2fc5b5
 pip==24.3.1 \
     --hash=sha256:3790624780082365f47549d032f3770eeb2b1e8bd1f7b2e02dace1afa361b4ed
-setuptools==75.5.0 \
+setuptools==78.1.1 \
-    --hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
+    --hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
+    --hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
 
 # Transitive dependencies of build
 colorama==0.4.6; platform_system == 'Windows' \

python/mujoco/CMakeLists.txt

@@ -84,7 +84,7 @@ if(NOT TARGET mujoco)
     if(MUJOCO_FRAMEWORK)
       message("MuJoCo framework is at ${MUJOCO_FRAMEWORK}/mujoco.framework")
       set(MUJOCO_LIBRARY
-          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.1.dylib
+          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.3.dylib
       )
       target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
     endif()
@@ -92,7 +92,7 @@ if(NOT TARGET mujoco)
 
   if(NOT MUJOCO_FRAMEWORK)
     find_library(
-      MUJOCO_LIBRARY mujoco mujoco.3.3.1 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
+      MUJOCO_LIBRARY mujoco mujoco.3.3.3 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
     )
     find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
     message("MuJoCo is at ${MUJOCO_LIBRARY}")
@@ -140,7 +140,7 @@ findorfetch(
   GIT_REPO
   https://github.com/abseil/abseil-cpp
   GIT_TAG
-  9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0
+  d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1
   TARGETS
   ${MUJOCO_PYTHON_ABSL_TARGETS}
   EXCLUDE_FROM_ALL
@@ -173,7 +173,7 @@ findorfetch(
   GIT_REPO
   https://gitlab.com/libeigen/eigen
   GIT_TAG
-  66f7f51b7e069d0a03a21157fa60b24aece69aeb
+  464c1d097891a1462ab28bf8bb763c1683883892
   TARGETS
   Eigen3::Eigen
   EXCLUDE_FROM_ALL
@@ -278,6 +278,25 @@ target_link_libraries(
             raw
 )
 
+add_library(structs_wrappers STATIC
+  structs_wrappers.cc
+  serialization.h
+  indexers.cc
+)
+target_include_directories(structs_wrappers PRIVATE ${Python3_INCLUDE_DIRS})
+target_link_libraries(
+  structs_wrappers
+  PRIVATE absl::flat_hash_map
+          absl::span
+          crossplatform
+          mujoco
+          raw
+          structs_header
+          pybind11::headers
+          Eigen3::Eigen
+)
+
+
 add_library(functions_header INTERFACE)
 target_sources(functions_header INTERFACE functions.h)
 set_target_properties(functions_header PROPERTIES PUBLIC_HEADER functions.h)
@@ -403,6 +422,7 @@ target_link_libraries(
           func_wrap
           function_traits
           structs_header
+          structs_wrappers
 )
 
 if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/specs.cc.inc)
@@ -417,6 +437,8 @@ endif()
 mujoco_pybind11_module(
   _specs
   specs.cc
+  specs_wrapper.h
+  specs_wrapper.cc
   specs.cc.inc
 )
 target_link_libraries(

python/mujoco/bindings_test.py

@@ -623,10 +623,12 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     ncon = 13
     nefc = 17
-    mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc)
+    nj = 21
+    mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc, nJ=nj)
 
     self.assertLen(self.data.contact, ncon)
     self.assertEqual(self.data.efc_id.shape, (nefc,))
+    self.assertEqual(self.data.efc_J.shape, (nj,))
     self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
 
     expected_error = 'insufficient arena memory available'
@@ -1083,6 +1085,20 @@ Euler integrator, semi-implicit in velocity.
     ):
       mujoco.mj_forward(self.model, self.data)
 
+  def test_timer_installed_by_default(self):
+    timer_step = mujoco.mjtTimer.mjTIMER_STEP
+    self.assertEqual(self.data.timer[timer_step].number, 0)
+    self.assertEqual(self.data.timer[timer_step].duration, 0.0)
+
+    mujoco.mj_step(self.model, self.data)
+    self.assertEqual(self.data.timer[timer_step].number, 1)
+    duration_1 = self.data.timer[timer_step].duration
+    self.assertGreater(duration_1, 0.0)
+
+    mujoco.mj_step(self.model, self.data, 5)
+    self.assertEqual(self.data.timer[timer_step].number, 6)
+    self.assertGreater(self.data.timer[timer_step].duration, duration_1)
+
   def test_mjcb_time(self):
 
     class CallCounter:

python/mujoco/codegen/generate_spec_bindings.py

@@ -84,6 +84,7 @@ def _value_binding_code(
   fulltype = fulltype.replace('mjOption', 'raw::MjOption')
   fulltype = fulltype.replace('mjVisual', 'raw::MjVisual')
   fulltype = fulltype.replace('mjStatistic', 'raw::MjStatistic')
+  element = '.element' if fullvarname == 'plugin' else ''
 
   def_property_args = (
       f'"{varname}"',
@@ -91,7 +92,7 @@ def _value_binding_code(
         return self.{fullvarname};
       }}""",
       f"""[]({rawclassname}& self, {fulltype} {varname}) {{
-        self.{fullvarname} = {varname};
+        self.{fullvarname}{element} = {varname}{element};
       }}""",
   )
 
@@ -613,12 +614,44 @@ def generate_find() -> None:
     print(code)
 
 
+def generate_signature() -> None:
+  """Generate signature functions."""
+  for key, _, _, _, _ in SPECS:
+    elem = key.removeprefix('mjs')
+    titlecase = 'Mjs' + elem
+    code = f"""\n
+      {key}.def_property_readonly("signature",
+      [](raw::{titlecase}& self) -> uint64_t {{
+        return mjs_getSpec(self.element)->element->signature;
+      }});
+    """
+    print(code)
+
+
+def generate_id() -> None:
+  """Generate id functions."""
+  for key, _, _, _, _ in SPECS:
+    if key == 'mjsPlugin':
+      continue
+    elem = key.removeprefix('mjs')
+    titlecase = 'Mjs' + elem
+    code = f"""\n
+      {key}.def_property_readonly("id",
+      [](raw::{titlecase}& self) -> int {{
+        return mjs_getId(self.element);
+      }});
+    """
+    print(code)
+
+
 def main(argv: Sequence[str]) -> None:
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
   generate()
   generate_add()
   generate_find()
+  generate_signature()
+  generate_id()
 
 
 if __name__ == '__main__':

python/mujoco/functions.cc

@@ -211,6 +211,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_tendon>(pymodule);
   Def<traits::mj_transmission>(pymodule);
   Def<traits::mj_crb>(pymodule);
+  Def<traits::mj_makeM>(pymodule);
   Def<traits::mj_factorM>(pymodule);
   DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM, "n")(
       pymodule,
@@ -1467,10 +1468,10 @@ PYBIND11_MODULE(_functions, pymodule) {
 
   pymodule.def(
       "_realloc_con_efc",
-      [](MjDataWrapper& d, int ncon, int nefc) {
+      [](MjDataWrapper& d, int ncon, int nefc, int nJ) {
         raw::MjData* data = d.get();
 
-        auto cleanup = [](raw::MjData* data) {
+        auto cleanup = [](raw::MjData* data, int nJ) {
 #ifdef ADDRESS_SANITIZER
         ASAN_POISON_MEMORY_REGION(
             static_cast<char*>(data->arena),
@@ -1479,6 +1480,7 @@ PYBIND11_MODULE(_functions, pymodule) {
           data->parena = 0;
           data->ncon = 0;
           data->nefc = 0;
+          if (nJ > -1) data->nJ = 0;
           data->contact = static_cast<raw::MjContact*>(data->arena);
 #define X(type, name, nr, nc) data->name = nullptr;
           MJDATA_ARENA_POINTERS_SOLVER
@@ -1486,14 +1488,15 @@ PYBIND11_MODULE(_functions, pymodule) {
 #undef X
         };
 
-        cleanup(data);
+        cleanup(data, nJ);
         data->ncon = ncon;
         data->nefc = nefc;
+        if (nJ > -1) data->nJ = nJ;
         data->contact =
             static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
                 data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
         if (!data->contact) {
-          cleanup(data);
+          cleanup(data, nJ);
           throw FatalError("insufficient arena memory available");
         }
 
@@ -1505,7 +1508,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
       data, sizeof(type) * (nr) * (nc), alignof(type)));                  \
   if (!data->name) {                                                      \
-    cleanup(data);                                                        \
+    cleanup(data, nJ);                                                        \
     throw FatalError("insufficient arena memory available");              \
   }
 
@@ -1519,7 +1522,7 @@ PYBIND11_MODULE(_functions, pymodule) {
 #undef MJ_M
 #define MJ_M(x) x
       },
-      py::arg("d"), py::arg("ncon"), py::arg("nefc"),
+      py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
       py::call_guard<py::gil_scoped_release>());
 }  // PYBIND11_MODULE NOLINT(readability/fn_size)
 }  // namespace

python/mujoco/indexer_xmacro.h

@@ -152,7 +152,7 @@
   X( int,     light_, mode,          nlight,  1 ) \
   X( int,     light_, bodyid,        nlight,  1 ) \
   X( int,     light_, targetbodyid,  nlight,  1 ) \
-  X( mjtByte, light_, directional,   nlight,  1 ) \
+  X( int,     light_, type,          nlight,  1 ) \
   X( mjtByte, light_, castshadow,    nlight,  1 ) \
   X( mjtByte, light_, active,        nlight,  1 ) \
   X( mjtNum,  light_, pos,           nlight,  3 ) \

python/mujoco/introspect/codegen/generate_functions.py

@@ -99,8 +99,13 @@ class MjFunctionVisitor:
       return ''.join(strings)
 
   def visit(self, node: ClangJsonNode) -> None:
-    if (node.get('kind') == 'FunctionDecl' and
+    # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated
-        node.get('name', '').startswith('mj')):
+    # objects and doesn't need a python wrapper.
+    if (
+        node.get('kind') == 'FunctionDecl'
+        and node.get('name', '').startswith('mj')
+        and node.get('name', '') != 'mjs_setUserValueWithCleanup'
+    ):
       func_decl = self._make_function(node)
       self._functions[func_decl.name] = func_decl
 

python/mujoco/introspect/codegen/generate_structs.py

@@ -47,7 +47,7 @@ _EXCLUDED = (
     'mjResource_',
 )
 
-_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s\s+\((.+) x (.+)\)\Z')
+_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s+\(([^\(\)]+) x ([^\(\)]+)\)\Z')
 
 
 def traverse(node, visitor):
@@ -93,16 +93,19 @@ class MjStructVisitor:
     # No valid normalization, just parse the declname.
     return type_parsing.parse_type(declname)
 
-  def _make_comment(self, node: ClangJsonNode) -> str:
+  def _make_comment(self, node: ClangJsonNode, strip: bool = True) -> str:
     """Makes a comment string from a Clang AST FullComment node."""
     kind = node.get('kind')
     if kind == 'TextComment':
-      return node['text'].replace('\N{NO-BREAK SPACE}', ' ').strip()
+      retval = node['text'].replace('\N{NO-BREAK SPACE}', ' ')
     else:
       strings = []
       for child in node['inner']:
-        strings.append(self._make_comment(child))
+        strings.append(self._make_comment(child, strip=False))
-      return ''.join(strings).strip()
+      retval = ''.join(strings)
+    if strip:
+      retval = retval.strip()
+    return retval
 
   def _make_field(
       self, node: ClangJsonNode

python/mujoco/introspect/enums.py

@@ -111,6 +111,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjCAMLIGHT_TARGETBODYCOM', 4),
          ]),
      )),
+    ('mjtLightType',
+     EnumDecl(
+         name='mjtLightType',
+         declname='enum mjtLightType_',
+         values=dict([
+             ('mjLIGHT_SPOT', 0),
+             ('mjLIGHT_DIRECTIONAL', 1),
+             ('mjLIGHT_POINT', 2),
+             ('mjLIGHT_IMAGE', 3),
+         ]),
+     )),
     ('mjtTexture',
      EnumDecl(
          name='mjtTexture',
@@ -139,6 +150,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjNTEXROLE', 10),
          ]),
      )),
+    ('mjtColorSpace',
+     EnumDecl(
+         name='mjtColorSpace',
+         declname='enum mjtColorSpace_',
+         values=dict([
+             ('mjCOLORSPACE_AUTO', 0),
+             ('mjCOLORSPACE_LINEAR', 1),
+             ('mjCOLORSPACE_SRGB', 2),
+         ]),
+     )),
     ('mjtIntegrator',
      EnumDecl(
          name='mjtIntegrator',
@@ -288,6 +309,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjNOBJECT', 26),
              ('mjOBJ_FRAME', 100),
              ('mjOBJ_DEFAULT', 101),
+             ('mjOBJ_MODEL', 102),
          ]),
      )),
     ('mjtConstraint',
@@ -339,34 +361,35 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjSENS_ACTUATORVEL', 14),
              ('mjSENS_ACTUATORFRC', 15),
              ('mjSENS_JOINTACTFRC', 16),
-             ('mjSENS_BALLQUAT', 17),
+             ('mjSENS_TENDONACTFRC', 17),
-             ('mjSENS_BALLANGVEL', 18),
+             ('mjSENS_BALLQUAT', 18),
-             ('mjSENS_JOINTLIMITPOS', 19),
+             ('mjSENS_BALLANGVEL', 19),
-             ('mjSENS_JOINTLIMITVEL', 20),
+             ('mjSENS_JOINTLIMITPOS', 20),
-             ('mjSENS_JOINTLIMITFRC', 21),
+             ('mjSENS_JOINTLIMITVEL', 21),
-             ('mjSENS_TENDONLIMITPOS', 22),
+             ('mjSENS_JOINTLIMITFRC', 22),
-             ('mjSENS_TENDONLIMITVEL', 23),
+             ('mjSENS_TENDONLIMITPOS', 23),
-             ('mjSENS_TENDONLIMITFRC', 24),
+             ('mjSENS_TENDONLIMITVEL', 24),
-             ('mjSENS_FRAMEPOS', 25),
+             ('mjSENS_TENDONLIMITFRC', 25),
-             ('mjSENS_FRAMEQUAT', 26),
+             ('mjSENS_FRAMEPOS', 26),
-             ('mjSENS_FRAMEXAXIS', 27),
+             ('mjSENS_FRAMEQUAT', 27),
-             ('mjSENS_FRAMEYAXIS', 28),
+             ('mjSENS_FRAMEXAXIS', 28),
-             ('mjSENS_FRAMEZAXIS', 29),
+             ('mjSENS_FRAMEYAXIS', 29),
-             ('mjSENS_FRAMELINVEL', 30),
+             ('mjSENS_FRAMEZAXIS', 30),
-             ('mjSENS_FRAMEANGVEL', 31),
+             ('mjSENS_FRAMELINVEL', 31),
-             ('mjSENS_FRAMELINACC', 32),
+             ('mjSENS_FRAMEANGVEL', 32),
-             ('mjSENS_FRAMEANGACC', 33),
+             ('mjSENS_FRAMELINACC', 33),
-             ('mjSENS_SUBTREECOM', 34),
+             ('mjSENS_FRAMEANGACC', 34),
-             ('mjSENS_SUBTREELINVEL', 35),
+             ('mjSENS_SUBTREECOM', 35),
-             ('mjSENS_SUBTREEANGMOM', 36),
+             ('mjSENS_SUBTREELINVEL', 36),
-             ('mjSENS_GEOMDIST', 37),
+             ('mjSENS_SUBTREEANGMOM', 37),
-             ('mjSENS_GEOMNORMAL', 38),
+             ('mjSENS_GEOMDIST', 38),
-             ('mjSENS_GEOMFROMTO', 39),
+             ('mjSENS_GEOMNORMAL', 39),
-             ('mjSENS_E_POTENTIAL', 40),
+             ('mjSENS_GEOMFROMTO', 40),
-             ('mjSENS_E_KINETIC', 41),
+             ('mjSENS_E_POTENTIAL', 41),
-             ('mjSENS_CLOCK', 42),
+             ('mjSENS_E_KINETIC', 42),
-             ('mjSENS_PLUGIN', 43),
+             ('mjSENS_CLOCK', 43),
-             ('mjSENS_USER', 44),
+             ('mjSENS_PLUGIN', 44),
+             ('mjSENS_USER', 45),
          ]),
      )),
     ('mjtStage',
@@ -426,6 +449,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjFLEXSELF_AUTO', 4),
          ]),
      )),
+    ('mjtSDFType',
+     EnumDecl(
+         name='mjtSDFType',
+         declname='enum mjtSDFType_',
+         values=dict([
+             ('mjSDFTYPE_SINGLE', 0),
+             ('mjSDFTYPE_INTERSECTION', 1),
+             ('mjSDFTYPE_MIDSURFACE', 2),
+             ('mjSDFTYPE_COLLISION', 3),
+         ]),
+     )),
     ('mjtTaskStatus',
      EnumDecl(
          name='mjtTaskStatus',