From 90da796468159ca56bb85174e14d9255d28a643f Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 23 Apr 2026 08:24:37 +0000
Subject: [PATCH 01/61] Bump pytest from 8.3.3 to 9.0.3 in /mjx

Bumps [pytest](https://github.com/pytest-dev/pytest) from 8.3.3 to 9.0.3.
- [Release notes](https://github.com/pytest-dev/pytest/releases)
- [Changelog](https://github.com/pytest-dev/pytest/blob/main/CHANGELOG.rst)
- [Commits](https://github.com/pytest-dev/pytest/compare/8.3.3...9.0.3)

---
updated-dependencies:
- dependency-name: pytest
  dependency-version: 9.0.3
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 mjx/requirements.txt | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/mjx/requirements.txt b/mjx/requirements.txt
index 835dbfa8..5706eb31 100644
--- a/mjx/requirements.txt
+++ b/mjx/requirements.txt
@@ -67,8 +67,9 @@ jaxlib==0.4.38; python_version >= '3.10' and python_version <= '3.13' and sys_pl
 pip==25.2 \
     --hash=sha256:578283f006390f85bb6282dffb876454593d637f5d1be494b5202ce4877e71f2 \
     --hash=sha256:6d67a2b4e7f14d8b31b8b52648866fa717f45a1eb70e83002f4331d07e953717
-pytest==8.3.3 \
-    --hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2
+pytest==9.0.3 \
+    --hash=sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9 \
+    --hash=sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c
 pytest-xdist==3.6.1 \
     --hash=sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7
 scipy==1.17.0; python_version >= '3.13' \

From 6b898395f14847171ee7bba533b9889cba22cf01 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 6 May 2026 15:52:40 +0000
Subject: [PATCH 02/61] Bump pip from 25.2 to 26.1 in /mjx

Bumps [pip](https://github.com/pypa/pip) from 25.2 to 26.1.
- [Changelog](https://github.com/pypa/pip/blob/main/NEWS.rst)
- [Commits](https://github.com/pypa/pip/compare/25.2...26.1)

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

Signed-off-by: dependabot[bot] <support@github.com>
---
 mjx/requirements.txt | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mjx/requirements.txt b/mjx/requirements.txt
index 835dbfa8..2958ce65 100644
--- a/mjx/requirements.txt
+++ b/mjx/requirements.txt
@@ -64,9 +64,9 @@ jaxlib==0.4.38; python_version >= '3.10' and python_version <= '3.13' and sys_pl
     --hash=sha256:248cca3771ebf24b070f49701364ceada33e6139445b06c782cca5ac5ad92bf4 \
     --hash=sha256:2ce77ba8cda9259a4bca97afc1c722e4291a6c463a63f8d372c6edc85117d625 \
     --hash=sha256:4103db0b3a38a5dc132741237453c24d8547290a22079ba1b577d6c88c95300a
-pip==25.2 \
-    --hash=sha256:578283f006390f85bb6282dffb876454593d637f5d1be494b5202ce4877e71f2 \
-    --hash=sha256:6d67a2b4e7f14d8b31b8b52648866fa717f45a1eb70e83002f4331d07e953717
+pip==26.1 \
+    --hash=sha256:4e8486d821d814b77319acb7b9e8bf5a4ee7590a643e7cb21029f209be8573c1 \
+    --hash=sha256:81e13ebcca3ffa8cc85e4deff5c27e1ee26dea0aa7fc2f294a073ac208806ff3
 pytest==8.3.3 \
     --hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2
 pytest-xdist==3.6.1 \

From 7b83f96e49eeb871c6b0c743ced95103ad80713b Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 20 May 2026 07:47:55 -0700
Subject: [PATCH 03/61] Prevent visited docs links from changing color.

PiperOrigin-RevId: 918445755
Change-Id: I4884c7f349ef39fab2678ba68eae422e6140b0a2
---
 doc/conf.py                 | 2 ++
 doc/css/theme_overrides.css | 5 -----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/doc/conf.py b/doc/conf.py
index c3cdfd37..b7ad83b1 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -168,6 +168,7 @@ SHARED_CSS_VARIABLES = {
 # font-stack--monospace used in code blocks, Inconsolata fits in 100 chars.
 html_theme_options = {
     'light_css_variables': {
+        'color-brand-visited': 'var(--color-brand-content)',
         'font-stack--monospace': 'Inconsolata,Consolas,ui-monospace,monospace',
         'at-color': '#830b2b',
         'at-val-color': '#bc103e',
@@ -182,6 +183,7 @@ html_theme_options = {
         'wy-nav-side-background-color': '#0053d6',
     },
     'dark_css_variables': {
+        'color-brand-visited': 'var(--color-brand-content)',
         'at-color': '#ffaab7',
         'at-val-color': '#ff95a6',
         'body-color': '#14234b',
diff --git a/doc/css/theme_overrides.css b/doc/css/theme_overrides.css
index a7318056..3c888606 100644
--- a/doc/css/theme_overrides.css
+++ b/doc/css/theme_overrides.css
@@ -197,11 +197,6 @@ h4 {
   padding: 6px;
 }
 
-/* Don't change color of visited links.
-.rst-content a.reference:visited {
-  color: var(--rst-content-a-color);
-} */
-
 .rst-content code, .rst-content tt {
   font-size: 90%;
   padding: inherit;

From 6aa738e75b6a2d364dda221fd03d14042f173a4f Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 20 May 2026 08:05:32 -0700
Subject: [PATCH 04/61] Add Expand/Collapse All buttons to XML schema
 documentation.

PiperOrigin-RevId: 918454244
Change-Id: Idcd2cec39ccaf90fd2d6265aaa364dd6bfb79254
---
 doc/XMLschema.rst | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst
index 6c9bb2a8..26eb579d 100755
--- a/doc/XMLschema.rst
+++ b/doc/XMLschema.rst
@@ -1,3 +1,10 @@
+.. raw:: html
+
+   <div class="schema-controls" style="margin-bottom: 1em;">
+     <button onclick="document.querySelectorAll('.schema-small details').forEach(d => d.open = true)" class="sd-btn sd-btn-outline-primary sd-btn-sm">Expand All</button>
+     <button onclick="document.querySelectorAll('.schema-small details details').forEach(d => d.open = false)" class="sd-btn sd-btn-outline-secondary sd-btn-sm" style="margin-left: 0.5em;">Collapse All</button>
+   </div>
+
 .. dropdown:: :ref:`mujoco<mujoco>` :octicon:`star`
    :open:
 

From 57d2c493163ad157c44fb2e54b385fcb2fb45322 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Wed, 20 May 2026 08:53:51 -0700
Subject: [PATCH 05/61] Track size/scale separately for renderables.

PiperOrigin-RevId: 918476985
Change-Id: I414b1abe2f9b3a3a7549ab5195971127d5598c36
---
 .../filament/compat/scene_geom_util.cc        | 32 +++++++++----------
 .../filament/filament/renderable.cc           | 23 ++++++++++---
 .../filament/filament/renderable.h            | 26 +++++++++++++--
 src/experimental/filament/filament_util.h     | 17 +---------
 .../filament/render_context_filament.cc       | 10 ++++--
 .../filament/render_context_filament.h        | 12 ++++---
 src/experimental/platform/ux/imgui_bridge.cc  |  4 +--
 7 files changed, 73 insertions(+), 51 deletions(-)

diff --git a/src/experimental/filament/compat/scene_geom_util.cc b/src/experimental/filament/compat/scene_geom_util.cc
index 0dc9d1f0..0d542ac2 100644
--- a/src/experimental/filament/compat/scene_geom_util.cc
+++ b/src/experimental/filament/compat/scene_geom_util.cc
@@ -57,29 +57,21 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
   std::memcpy(position, &geom.pos, 3 * sizeof(float));
   float rotation[9];
   std::memcpy(rotation, &geom.mat, 9 * sizeof(float));
-  float size[3];
-  std::memcpy(size, &geom.size, 3 * sizeof(float));
 
   const mjtGeom geom_type = (mjtGeom)geom.type;
   switch (geom_type) {
     case mjGEOM_MESH:
     case mjGEOM_SDF:
       mjrf_setRenderableMesh(renderable, model_objs->GetMesh(geom.dataid), 0, 0);
-      // Ignore size for meshes.
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_HFIELD:
       mjrf_setRenderableMesh(renderable, model_objs->GetHeightField(geom.dataid), 0, 0);
-      // Ignore size for meshes.
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_PLANE: {
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
 
+      float size[3];
+      std::memcpy(size, &geom.size, 3 * sizeof(float));
       const bool is_infinite = !(size[0] > 0 && size[1] > 0);
       if (is_infinite) {
         // Infinite planes are scaled to match the tile size used by
@@ -90,40 +82,52 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       }
       // Planes only define an xy size, so set the z-dimension to 1.0f.
       size[2] = 1.0f;
+      mjrf_setRenderableSize(renderable, size);
       break;
     }
     case mjGEOM_SPHERE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ELLIPSOID:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_BOX:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_CAPSULE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_CYLINDER:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW1:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_ARROW2:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_LINE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_LINEBOX:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_TRIANGLE:
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
+      mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_FLEX:
       mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
@@ -133,9 +137,6 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       rotation[0] = 1.f;
       rotation[4] = 1.f;
       rotation[8] = 1.f;
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_SKIN:
       mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
@@ -145,9 +146,6 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       rotation[0] = 1.f;
       rotation[4] = 1.f;
       rotation[8] = 1.f;
-      size[0] = 1.f;
-      size[1] = 1.f;
-      size[2] = 1.f;
       break;
     case mjGEOM_NONE:
     case mjGEOM_LABEL:
@@ -158,7 +156,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       break;
   }
 
-  mjrf_setRenderableTransform(renderable, position, rotation, size);
+  mjrf_setRenderableTransform(renderable, position, rotation);
 }
 
 static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
diff --git a/src/experimental/filament/filament/renderable.cc b/src/experimental/filament/filament/renderable.cc
index 535110cb..f181b82c 100644
--- a/src/experimental/filament/filament/renderable.cc
+++ b/src/experimental/filament/filament/renderable.cc
@@ -22,13 +22,13 @@
 #include <filament/RenderableManager.h>
 #include <filament/Scene.h>
 #include <filament/TransformManager.h>
+#include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/builtins.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
@@ -40,6 +40,7 @@ namespace mujoco {
 using filament::math::float2;
 using filament::math::float3;
 using filament::math::float4;
+using filament::math::mat3f;
 using filament::math::mat4f;
 
 // An arbitrary scale factor for arrows.
@@ -149,21 +150,33 @@ void Renderable::InitPartEntity(Part& part) {
   }
 }
 
-void Renderable::SetTransform(const Trs& trs) {
+void Renderable::SetTransform(const float3& position, const mat3f& rotation) {
+  trs_.translation = position;
+  trs_.rotation = rotation;
+  UpdateTransform();
+}
+
+void Renderable::SetSize(const float3& size) {
+  trs_.size = size;
+  UpdateTransform();
+}
+
+void Renderable::UpdateTransform() {
   if (parts_.empty()) {
-    transform_ = trs.ToTransform();
+    transform_ = trs_.ToTransform();
     return;
   }
 
   filament::TransformManager& tm = GetEngine()->getTransformManager();
   if (get_transform_fn_) {
     for (int i = 0; i < parts_.size(); ++i) {
-      const mat4f& transform = get_transform_fn_(i, trs);
+      const mat4f transform = get_transform_fn_(i, trs_);
       tm.setTransform(tm.getInstance(parts_[i].entity), transform);
     }
   } else {
+    const mat4f transform = trs_.ToTransform();
     for (Part& part : parts_) {
-      tm.setTransform(tm.getInstance(part.entity), trs.ToTransform());
+      tm.setTransform(tm.getInstance(part.entity), transform);
     }
   }
   transform_ = tm.getTransform(tm.getInstance(parts_[0].entity));
diff --git a/src/experimental/filament/filament/renderable.h b/src/experimental/filament/filament/renderable.h
index c41c83d3..ab0cd3b2 100644
--- a/src/experimental/filament/filament/renderable.h
+++ b/src/experimental/filament/filament/renderable.h
@@ -21,10 +21,11 @@
 
 #include <filament/Engine.h>
 #include <filament/Scene.h>
+#include <math/mat3.h>
 #include <math/mat4.h>
+#include <math/vec3.h>
 #include <utils/Entity.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
 #include "experimental/filament/render_context_filament.h"
@@ -53,8 +54,14 @@ class Renderable : public mjrRenderable {
   // Sets the mesh of this renderable to a built-in mesh based on the geom type.
   void SetGeomMesh(mjtGeom type, int nstack, int nslice, int nquad);
 
-  // Sets the transform of this renderable.
-  void SetTransform(const Trs& trs);
+  // Sets the position and rotation of this renderable.
+  void SetTransform(const filament::math::float3& position,
+                    const filament::math::mat3f& rotation);
+
+  // Sets the size of this renderable. Note: this is effectively the same as a
+  // scale for most renderables. However, for e.g. capsules, the spherical ends
+  // are not scaled and remain fixed in size.
+  void SetSize(const filament::math::float3& size);
 
   // Returns the current transform of this renderable.
   const filament::math::mat4f& GetTransform() const;
@@ -114,12 +121,24 @@ class Renderable : public mjrRenderable {
     int elem_count = 0;
   };
 
+  // A tuple of translation, rotation, and size.
+  struct Trs {
+    filament::math::float3 translation{0.0f, 0.0f, 0.0f};
+    filament::math::mat3f rotation;
+    filament::math::float3 size{1.0f, 1.0f, 1.0f};
+    filament::math::mat4f ToTransform() const {
+      return filament::math::mat4f(rotation, translation) *
+            filament::math::mat4f::scaling(size);
+    }
+  };
+
   // When composing a multi-part renderable, each Entity will have its own
   // transform offset based on the transform of the Renderable itself.
   using GetTransformFn = std::function<filament::math::mat4f(int, const Trs&)>;
 
   void AppendMesh(const Mesh* mesh);
   void InitPartEntity(Part& part);
+  void UpdateTransform();
 
   void AssignMaterial(mjrDrawMode mode,
                       ObjectManager::MaterialType material_type);
@@ -136,6 +155,7 @@ class Renderable : public mjrRenderable {
   std::vector<Part> parts_;
   filament::math::mat4f transform_;
   GetTransformFn get_transform_fn_;
+  Trs trs_;
   bool wireframe_ = false;
 };
 
diff --git a/src/experimental/filament/filament_util.h b/src/experimental/filament/filament_util.h
index 39c4c640..c8398522 100644
--- a/src/experimental/filament/filament_util.h
+++ b/src/experimental/filament/filament_util.h
@@ -15,6 +15,7 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_UTIL_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_UTIL_H_
 
+#include <string_view>
 #include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
@@ -57,22 +58,6 @@ inline filament::math::mat3f ReadMat3(const T* arr, int index = 0) {
   // clang-format on
 }
 
-// A tuple of translation, rotation, and size.
-struct Trs {
-  filament::math::float3 translation{0.0f, 0.0f, 0.0f};
-  filament::math::mat3f rotation;
-  // Note: this is _slightly_ different than scale. For example, for capsules,
-  // the size determines the length of the tube and the radius of the domes,
-  // but the shape remains a capsule.
-  filament::math::float3 size{1.0f, 1.0f, 1.0f};
-
-  // Converts the TRS to a transform matrix.
-  filament::math::mat4f ToTransform() const {
-    return filament::math::mat4f(rotation, translation) *
-           filament::math::mat4f::scaling(size);
-  }
-};
-
 // Calculates a reflection matrix for a plane defined by its transform.
 filament::math::mat4 ToReflectionMatrix(const filament::math::mat4& xform);
 
diff --git a/src/experimental/filament/render_context_filament.cc b/src/experimental/filament/render_context_filament.cc
index b3e3e2f7..f122110b 100644
--- a/src/experimental/filament/render_context_filament.cc
+++ b/src/experimental/filament/render_context_filament.cc
@@ -247,14 +247,18 @@ void mjrf_setRenderableMaterial(mjrRenderable* renderable,
 
 void mjrf_setRenderableTransform(mjrRenderable* renderable,
                                  const float position[3],
-                                 const float rotation[9], const float size[3]) {
+                                 const float rotation[9]) {
   const filament::math::float3 fposition{position[0], position[1], position[2]};
-  const filament::math::float3 fsize{size[0], size[1], size[2]};
   const filament::math::mat3f frotation{rotation[0], rotation[3], rotation[6],
                                         rotation[1], rotation[4], rotation[7],
                                         rotation[2], rotation[5], rotation[8]};
   mujoco::Renderable::downcast(renderable)
-      ->SetTransform({fposition, frotation, fsize});
+      ->SetTransform(fposition, frotation);
+}
+
+void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]) {
+  const filament::math::float3 fsize{size[0], size[1], size[2]};
+  mujoco::Renderable::downcast(renderable)->SetSize(fsize);
 }
 
 void mjrf_setRenderableLayerMask(mjrRenderable* renderable,
diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h
index e7043bec..023cc62c 100644
--- a/src/experimental/filament/render_context_filament.h
+++ b/src/experimental/filament/render_context_filament.h
@@ -695,12 +695,16 @@ void mjrf_setRenderableGeomMesh(mjrRenderable* renderable, mjtGeom type,
 void mjrf_setRenderableMaterial(mjrRenderable* renderable,
                                 const mjrMaterial* material);
 
-// Sets the transform (position, rotation, and size) of the renderable. Note
-// that `size` is not the same as `scale`. For example, the z-size of a capsule
-// only scales the tubular-portion of its geometry, but not the spherical caps.
+// Sets the transform position and rotation of the renderable.
 void mjrf_setRenderableTransform(mjrRenderable* renderable,
                                  const float position[3],
-                                 const float rotation[9], const float size[3]);
+                                 const float rotation[9]);
+
+// Sets the size of the renderable. Note that, for most renderables, this is
+// equivalent to setting the scale. However, for some geom-based renderables,
+// the size scale is not applied uniformly (e.g. the spherical ends of a
+// capsule are scaled such that they always remain spherical).
+void mjrf_setRenderableSize(mjrRenderable* renderable, const float size[3]);
 
 // Sets whether the renderable casts shadows or not.
 void mjrf_setRenderableCastShadows(mjrRenderable* renderable,
diff --git a/src/experimental/platform/ux/imgui_bridge.cc b/src/experimental/platform/ux/imgui_bridge.cc
index ef21e4f0..91147fa4 100644
--- a/src/experimental/platform/ux/imgui_bridge.cc
+++ b/src/experimental/platform/ux/imgui_bridge.cc
@@ -270,10 +270,8 @@ void ImguiBridge::Update() {
       }
       mjrf_setRenderableMaterial(renderable.get(), &material);
 
-      const float position[] = {0, 0, 0};
-      const float rotation[] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
       const float size[] = {scale.x, scale.y, 1.0f};
-      mjrf_setRenderableTransform(renderable.get(), position, rotation, size);
+      mjrf_setRenderableSize(renderable.get(), size);
 
       index_offset += command.ElemCount;
       ++renderable_index;

From fa8ac7a516cbb0c177ae0a6621669e38aaf75743 Mon Sep 17 00:00:00 2001
From: Sam Haves <shaves@google.com>
Date: Wed, 20 May 2026 09:32:00 -0700
Subject: [PATCH 06/61] Implement mjz encoder.

PiperOrigin-RevId: 918493808
Change-Id: I150d182ff27801d829537a8001a0b3a04c1a0499
---
 CMakeLists.txt                                |    2 +-
 src/xml/mjz/CMakeLists.txt                    |    1 +
 src/xml/mjz/mjz_encoder.cc                    |  421 +++++++
 test/xml/mjz/CMakeLists.txt                   |    4 +
 test/xml/mjz/mjz_encoder_test.cc              | 1105 +++++++++++++++++
 .../asset_collision_test/child_a/child.xml    |   10 +
 .../asset_collision_test/child_a/mesh.obj     |   20 +
 .../asset_collision_test/child_b/child.xml    |   10 +
 .../asset_collision_test/child_b/mesh.obj     |    8 +
 test/xml/mjz/testdata/disk_mesh/box.obj       |   20 +
 test/xml/mjz/testdata/disk_mesh/model.xml     |   10 +
 test/xml/mjz/testdata/include_test/child.xml  |    7 +
 test/xml/mjz/testdata/include_test/parent.xml |    8 +
 .../mjz/testdata/meshdir_test/assets/box.obj  |   20 +
 test/xml/mjz/testdata/meshdir_test/model.xml  |   11 +
 test/xml/mjz/testdata/mixed_test/disk_box.obj |   20 +
 .../mjz/testdata/subdir_mesh/meshes/box.obj   |   20 +
 test/xml/mjz/testdata/subdir_mesh/model.xml   |   10 +
 .../mjz/testdata/transitive_test/child.xml    |    8 +
 .../testdata/transitive_test/grandchild.xml   |    7 +
 .../mjz/testdata/transitive_test/parent.xml   |    8 +
 21 files changed, 1729 insertions(+), 1 deletion(-)
 create mode 100644 src/xml/mjz/mjz_encoder.cc
 create mode 100644 test/xml/mjz/mjz_encoder_test.cc
 create mode 100644 test/xml/mjz/testdata/asset_collision_test/child_a/child.xml
 create mode 100644 test/xml/mjz/testdata/asset_collision_test/child_a/mesh.obj
 create mode 100644 test/xml/mjz/testdata/asset_collision_test/child_b/child.xml
 create mode 100644 test/xml/mjz/testdata/asset_collision_test/child_b/mesh.obj
 create mode 100644 test/xml/mjz/testdata/disk_mesh/box.obj
 create mode 100644 test/xml/mjz/testdata/disk_mesh/model.xml
 create mode 100644 test/xml/mjz/testdata/include_test/child.xml
 create mode 100644 test/xml/mjz/testdata/include_test/parent.xml
 create mode 100644 test/xml/mjz/testdata/meshdir_test/assets/box.obj
 create mode 100644 test/xml/mjz/testdata/meshdir_test/model.xml
 create mode 100644 test/xml/mjz/testdata/mixed_test/disk_box.obj
 create mode 100644 test/xml/mjz/testdata/subdir_mesh/meshes/box.obj
 create mode 100644 test/xml/mjz/testdata/subdir_mesh/model.xml
 create mode 100644 test/xml/mjz/testdata/transitive_test/child.xml
 create mode 100644 test/xml/mjz/testdata/transitive_test/grandchild.xml
 create mode 100644 test/xml/mjz/testdata/transitive_test/parent.xml

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9157b7c4..2ed0692e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -157,6 +157,7 @@ add_subdirectory(plugin/stl_decoder)
 add_subdirectory(src/engine)
 add_subdirectory(src/user)
 add_subdirectory(src/xml)
+add_subdirectory(src/xml/mjz)
 add_subdirectory(src/thread)
 if(NOT EMSCRIPTEN AND NOT MUJOCO_USE_FILAMENT_MJR_COMPAT)
   add_subdirectory(src/render/classic)
@@ -256,7 +257,6 @@ if(MUJOCO_BUILD_SIMULATE)
 endif()
 
 if(MUJOCO_BUILD_STUDIO)
-  add_subdirectory(src/xml/mjz)
   add_subdirectory(src/experimental/platform)
   add_subdirectory(src/experimental/studio)
 endif()
diff --git a/src/xml/mjz/CMakeLists.txt b/src/xml/mjz/CMakeLists.txt
index a44e096b..ab92623b 100644
--- a/src/xml/mjz/CMakeLists.txt
+++ b/src/xml/mjz/CMakeLists.txt
@@ -14,6 +14,7 @@
 
 set(MUJOCO_MJZ_SRCS
     mjz_decoder.cc
+    mjz_encoder.cc
 )
 
 target_sources(mujoco PRIVATE ${MUJOCO_MJZ_SRCS})
diff --git a/src/xml/mjz/mjz_encoder.cc b/src/xml/mjz/mjz_encoder.cc
new file mode 100644
index 00000000..b209fb04
--- /dev/null
+++ b/src/xml/mjz/mjz_encoder.cc
@@ -0,0 +1,421 @@
+// Copyright 2026 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 <cctype>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <filesystem>
+#include <functional>
+#include <stack>
+#include <string>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+// Disable unused function warnings for miniz.
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wunused-function"
+#endif
+#include <miniz.h>
+#if defined(__GNUC__) || defined(__clang__)
+#pragma GCC diagnostic pop
+#endif
+
+#include <mujoco/mujoco.h>
+#include "tinyxml2.h"
+
+namespace {
+
+namespace fs = std::filesystem;
+
+// asset file to be packed into the archive
+struct AssetEntry {
+  fs::path archive_path;
+  // modelfiledir of the spec that originally owned the asset
+  fs::path source_dir;
+  fs::path disk_path;
+};
+
+// Key for identifying an asset element that needs its file path rewritten.
+// Using (tag, file, name) allows disambiguation of unnamed elements and
+// individual cubemap face files on the same texture element.
+struct RewriteKey {
+  std::string_view tag;
+  std::string_view file;
+  std::string_view name;
+
+  bool operator==(const RewriteKey&) const = default;
+};
+
+struct RewriteKeyHash {
+  std::size_t operator()(const RewriteKey& k) const {
+    const std::size_t h1 = std::hash<std::string_view>{}(k.tag);
+    const std::size_t h2 = std::hash<std::string_view>{}(k.file);
+    const std::size_t h3 = std::hash<std::string_view>{}(k.name);
+    std::size_t seed = h1;
+    seed ^= h2 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+    seed ^= h3 + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+    return seed;
+  }
+};
+
+using RewriteMap = std::unordered_map<RewriteKey, std::string, RewriteKeyHash>;
+
+fs::path SanitizePath(const fs::path& path) {
+  // replace characters that are not alphanumeric, '/', '.', '_', or '-'
+  // with '_'
+  std::string result = path.string();
+  for (char& c : result) {
+    if (!std::isalnum(static_cast<unsigned char>(c)) && c != '/' && c != '.' &&
+        c != '_' && c != '-') {
+      c = '_';
+    }
+  }
+  return result;
+}
+
+// Remove leading ".." path components from a normalized path.
+std::string RemoveLeadingDotDot(const fs::path& p) {
+  fs::path result;
+  bool skipping = true;
+  for (const auto& component : p) {
+    if (skipping && component == "..") {
+      continue;
+    }
+    skipping = false;
+    result /= component;
+  }
+  return result.string();
+}
+
+// Apply file-attribute rewrites to the serialized XML.
+// For each element, all attributes are checked against the rewrite map keyed
+// by (tag_name, attribute_value, element_name). This handles both regular
+// "file" attributes and cubemap face attributes (fileright, fileleft, etc.).
+void ApplyRewrites(std::string& xml, const RewriteMap& rewrites) {
+  if (rewrites.empty()) return;
+
+  tinyxml2::XMLDocument doc;
+  if (doc.Parse(xml.c_str()) != tinyxml2::XML_SUCCESS) return;
+
+  // Walk all elements in the document.
+  std::stack<tinyxml2::XMLElement*> stack;
+  if (auto* root = doc.RootElement()) stack.push(root);
+  while (!stack.empty()) {
+    auto* elem = stack.top();
+    stack.pop();
+
+    const char* tag_name = elem->Value();
+    const char* name_attr = elem->Attribute("name");
+
+    // Check file-related attributes for paths that need rewriting.
+    // This covers "file" and cubemap face attributes: fileright, fileleft, etc.
+    for (const tinyxml2::XMLAttribute* attr = elem->FirstAttribute(); attr;
+         attr = attr->Next()) {
+      std::string_view attr_name{attr->Name()};
+      if (!attr_name.starts_with("file")) continue;
+
+      const char* attr_value = attr->Value();
+      if (!attr_value) continue;
+
+      RewriteKey key{tag_name, attr_value, name_attr ? name_attr : ""};
+      auto it = rewrites.find(key);
+      if (it != rewrites.end()) {
+        elem->SetAttribute(attr->Name(), it->second.c_str());
+      }
+    }
+
+    for (auto* child = elem->FirstChildElement(); child;
+         child = child->NextSiblingElement()) {
+      stack.push(child);
+    }
+  }
+
+  tinyxml2::XMLPrinter printer;
+  doc.Print(&printer);
+  xml = printer.CStr();
+}
+
+// collect all referenced asset files from the spec
+//
+// returns a map from unique archive entry names to AssetEntry structs
+// and populates xml_rewrites with XML file attributes that need to be updated.
+std::unordered_map<std::string, AssetEntry> CollectAssets(
+    const mjSpec* spec, RewriteMap& xml_rewrites) {
+  const mjString* root_meshdir = spec->compiler.meshdir;
+  const mjString* root_texturedir = spec->compiler.texturedir;
+
+  struct PathHash {
+    std::size_t operator()(const fs::path& p) const {
+      return fs::hash_value(p);
+    }
+  };
+  // maps full disk paths to their archive paths
+  std::unordered_map<fs::path, std::string, PathHash> archived_paths;
+  // all asset entries, keyed by the final archive path
+  std::unordered_map<std::string, AssetEntry> archive_entries;
+
+  auto process = [&](mjsElement* elem, const mjString* raw_file,
+                     bool use_meshdir, std::string_view tag_name) {
+    // skip empty files
+    if (!raw_file || raw_file->empty()) return;
+
+    const mjString* elem_name = mjs_getName(elem);
+    const mjSpec* owning_spec = mjs_getOriginSpec(elem);
+
+    mjsCompiler* comp = mjs_getCompiler(elem);
+    if (!comp) {
+      mju_error(
+          "MJZ encoder: no compiler for element '%s', this should never "
+          "happen.",
+          elem_name->c_str());
+    }
+
+    auto rewrite_key = RewriteKey{tag_name, *raw_file, *elem_name};
+
+    const fs::path owning_spec_dir{
+        owning_spec->modelfiledir ? *owning_spec->modelfiledir : ""};
+    const fs::path raw_path{*raw_file};
+    const fs::path prefix_dir{use_meshdir ? *comp->meshdir : *comp->texturedir};
+    const fs::path root_dir =
+        fs::path(use_meshdir ? *root_meshdir : *root_texturedir);
+
+    const char* ch = std::strchr(raw_file->c_str(), ':');
+    const bool is_uri = ch != nullptr;
+
+    // full path of the asset relative to its spec, or absolute/URI if the
+    // raw path is not relative.
+    fs::path full_spec_path = raw_path;
+    // If the raw path is relative AND it does not have a valid resource
+    // provider URI, then we need to prefix it with the specs
+    // meshdir/texturedir.
+    if (!is_uri && raw_path.is_relative()) {
+      full_spec_path = prefix_dir / raw_path;
+    }
+
+    // full path of the asset regardless of spec, URI and absolute paths
+    // remain unchanged but relative paths are prefixed with modelfiledir
+    fs::path full_path = full_spec_path;
+    if (!is_uri && full_spec_path.is_relative()) {
+      full_path = owning_spec_dir / full_spec_path;
+    }
+
+    // sanitize the path to remove any URI schemes or other non-path characters
+    // if the file path has a URI scheme (e.g. "http://foo/mesh.stl"), strip it
+    // so the archive entry uses a concrete path and decoding the MJZ won't try
+    // to invoke a resource provider.
+    // First remove any leading ".." path components.
+    // my_provider:a/../b/c_$.obj -> my_provider_a/../b/c__.obj
+    const fs::path sanitized = SanitizePath(raw_path);
+    // my_provider_a/../b/c__.obj -> my_provider_a/b/c__.obj
+    const fs::path normalized = sanitized.lexically_normal();
+    // ../../b/c__.obj -> b/c__.obj
+    // Also remove any leading '/'.
+    const fs::path localized = RemoveLeadingDotDot(normalized.relative_path());
+    // path relative to the root XML in the archive
+    fs::path archive_path = root_dir / localized;
+
+    // If this file was already archived, we may still need to add a rewrite
+    // if raw_file was sanitized or collision-renamed for the first occurrence.
+    if (auto it = archived_paths.find(full_path); it != archived_paths.end()) {
+      if (archive_path != it->second || localized != raw_path) {
+        xml_rewrites[rewrite_key] = it->second;
+      }
+      return;
+    }
+
+    if (localized != raw_path) {
+      xml_rewrites[rewrite_key] = archive_path.string();
+    }
+
+    // Collision renaming: if this archive path is already in use,
+    // try again with an incremented suffix.
+    fs::path parent = archive_path.parent_path();
+    fs::path stem = archive_path.stem();
+    fs::path extension = archive_path.extension();
+    for (int i = 0; archive_entries.contains(archive_path.string()); ++i) {
+      std::string new_name =
+          stem.string() + "_" + std::to_string(i) + extension.string();
+      archive_path = parent / new_name;
+      xml_rewrites[rewrite_key] = archive_path.string();
+    }
+
+    archived_paths[full_path] = archive_path.string();
+    archive_entries[archive_path.string()] =
+        AssetEntry{archive_path, owning_spec_dir, full_spec_path};
+  };
+
+  // Meshes, heightfields, and skins use meshdir.
+  {
+    const mjsMesh* mesh = mjs_asMesh(mjs_firstElement(spec, mjOBJ_MESH));
+    while (mesh != nullptr) {
+      if (!mesh->file) continue;
+      process(mesh->element, mesh->file, /*use_meshdir=*/true, "mesh");
+      mesh = mjs_asMesh(mjs_nextElement(spec, mesh->element));
+    }
+  }
+
+  {
+    const mjsHField* hf = mjs_asHField(mjs_firstElement(spec, mjOBJ_HFIELD));
+    while (hf != nullptr) {
+      if (!hf->file) continue;
+      process(hf->element, hf->file, /*use_meshdir=*/true, "hfield");
+      hf = mjs_asHField(mjs_nextElement(spec, hf->element));
+    }
+  }
+
+  {
+    const mjsSkin* skin = mjs_asSkin(mjs_firstElement(spec, mjOBJ_SKIN));
+    while (skin != nullptr) {
+      if (!skin->file) continue;
+      process(skin->element, skin->file, /*use_meshdir=*/true, "skin");
+      skin = mjs_asSkin(mjs_nextElement(spec, skin->element));
+    }
+  }
+
+  // Textures use texturedir.
+  {
+    const mjsTexture* tex =
+        mjs_asTexture(mjs_firstElement(spec, mjOBJ_TEXTURE));
+    while (tex != nullptr) {
+      if (!tex->file) continue;
+      process(tex->element, tex->file, /*use_meshdir=*/false, "texture");
+      if (tex->cubefiles) {
+        for (const mjString& file : *tex->cubefiles) {
+          process(tex->element, &file, /*use_meshdir=*/false, "texture");
+        }
+      }
+      tex = mjs_asTexture(mjs_nextElement(spec, tex->element));
+    }
+  }
+
+  return archive_entries;
+}
+
+int MjzEncode(const mjSpec* spec, const mjModel* model, const mjVFS* vfs,
+              mjResource* resource) {
+  if (!spec || !resource) {
+    return -1;
+  }
+
+  const fs::path archive_path(resource->name);
+  const std::string stem = archive_path.stem().string();
+
+  // Collect assets and compute archive paths + XML rewrite list.
+  RewriteMap xml_rewrites;
+  std::unordered_map<std::string, AssetEntry> assets =
+      CollectAssets(spec, xml_rewrites);
+
+  // Serialize spec to XML.
+  char error[1024] = {0};
+  int xml_sz = 1024 * 1024;
+  std::vector<char> xml_buf(xml_sz);
+  int result =
+      mj_saveXMLString(spec, xml_buf.data(), xml_sz, error, sizeof(error));
+  // mj_saveXMLString returns 0 on success, -1 on failure, or a positive value
+  // indicating the required buffer size when the buffer is too small.
+  if (result > 0) {
+    xml_sz = result + 1;
+    xml_buf.resize(xml_sz);
+    result =
+        mj_saveXMLString(spec, xml_buf.data(), xml_sz, error, sizeof(error));
+  }
+  if (result != 0) {
+    mju_warning("MJZ encoder: failed to serialize spec to XML: %s", error);
+    return -1;
+  }
+
+  // Post-process the XML: apply file attribute rewrites for collision-renamed
+  // assets. Only colliding entries get rewritten; non-colliding models are
+  // unchanged.
+  std::string xml_str(xml_buf.data());
+  ApplyRewrites(xml_str, xml_rewrites);
+
+  // NOLINTBEGIN(misc-include-cleaner) - Disable clang-tidy for miniz.
+  // Initialize zip writer.
+  mz_zip_archive zip;
+  std::memset(&zip, 0, sizeof(zip));
+  if (!mz_zip_writer_init_heap(&zip, 0, 0)) {
+    mju_warning("MJZ encoder: failed to init zip writer");
+    return -1;
+  }
+
+  // Add XML to archive.
+  const std::string xml_name = stem + ".xml";
+  if (!mz_zip_writer_add_mem(&zip, xml_name.c_str(), xml_str.data(),
+                             xml_str.size(), MZ_DEFAULT_COMPRESSION)) {
+    mju_warning("MJZ encoder: failed to add XML to archive");
+    mz_zip_writer_end(&zip);
+    return -1;
+  }
+
+  // Pack each asset file into the archive using the unique archive paths.
+  for (const auto& [archive_entry, entry] : assets) {
+    mjResource* res = mju_openResource(entry.source_dir.string().c_str(),
+                                       entry.disk_path.string().c_str(), vfs,
+                                       error, sizeof(error));
+    if (!res) {
+      mju_warning("MJZ encoder: failed to open resource '%s' (dir='%s'): %s",
+                  entry.disk_path.c_str(), entry.source_dir.c_str(), error);
+      continue;
+    }
+
+    const void* buf = nullptr;
+    int nbytes = mju_readResource(res, &buf);
+    if (nbytes > 0 && buf) {
+      if (!mz_zip_writer_add_mem(&zip, archive_entry.c_str(), buf, nbytes,
+                                 MZ_DEFAULT_COMPRESSION)) {
+        mju_warning("MJZ encoder: failed to write '%s' to archive",
+                    archive_entry.c_str());
+      }
+    }
+    mju_closeResource(res);
+  }
+
+  // Finalize archive.
+  void* archive_buf = nullptr;
+  size_t archive_size = 0;
+  if (!mz_zip_writer_finalize_heap_archive(&zip, &archive_buf, &archive_size)) {
+    mju_warning("MJZ encoder: failed to finalize archive");
+    mz_zip_writer_end(&zip);
+    return -1;
+  }
+  mz_zip_writer_end(&zip);
+  // NOLINTEND(misc-include-cleaner) - End clang-tidy disable for miniz.
+
+  resource->data = archive_buf;
+
+  return static_cast<int>(archive_size);
+}
+
+void MjzCloseResource(mjResource* resource) {
+  if (resource && resource->data) {
+    std::free(resource->data);
+    resource->data = nullptr;
+  }
+}
+}  // namespace
+
+mjPLUGIN_LIB_INIT(mjz_encoder) {
+  mjpEncoder encoder;
+  mjp_defaultEncoder(&encoder);
+  encoder.content_type = "application/zip";
+  encoder.extension = ".mjz|.zip";
+  encoder.encode = MjzEncode;
+  encoder.close_resource = MjzCloseResource;
+  mjp_registerEncoder(&encoder);
+}
diff --git a/test/xml/mjz/CMakeLists.txt b/test/xml/mjz/CMakeLists.txt
index 721be4fc..718cf9df 100644
--- a/test/xml/mjz/CMakeLists.txt
+++ b/test/xml/mjz/CMakeLists.txt
@@ -17,3 +17,7 @@ mujoco_test(mjz_api_test)
 mujoco_test(
   mjz_decoder_test
 )
+
+mujoco_test(
+  mjz_encoder_test
+)
diff --git a/test/xml/mjz/mjz_encoder_test.cc b/test/xml/mjz/mjz_encoder_test.cc
new file mode 100644
index 00000000..c9cd1ff1
--- /dev/null
+++ b/test/xml/mjz/mjz_encoder_test.cc
@@ -0,0 +1,1105 @@
+// Copyright 2026 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 the MJZ encoder plugin.
+
+#include <array>
+#include <cctype>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <filesystem>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <absl/strings/match.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjplugin.h>
+#include <mujoco/mjspec.h>
+#include <mujoco/mujoco.h>
+#include "src/xml/xml_numeric_format.h"
+#include "test/fixture.h"
+
+namespace mujoco {
+
+namespace fs = std::filesystem;
+using ::testing::NotNull;
+
+static mjSpec* MakeSimpleSpec() {
+  mjSpec* s = mj_makeSpec();
+  mjsBody* world = mjs_findBody(s, "world");
+  mjsBody* body = mjs_addBody(world, nullptr);
+  mjsGeom* geom = mjs_addGeom(body, nullptr);
+  geom->size[0] = 1.0;
+  geom->size[1] = 1.0;
+  geom->size[2] = 1.0;
+  return s;
+}
+
+static std::string SanitizePathForTestName(const std::string& path) {
+  fs::path p(path);
+  std::string name = p.stem().string();
+  std::string full_name = p.string();
+  size_t pos = full_name.find("third_party/mujoco");
+  if (pos != std::string::npos) {
+    full_name = full_name.substr(pos);
+  }
+  std::string sanitized;
+  for (char c : full_name) {
+    if (std::isalnum(c)) {
+      sanitized += c;
+    } else {
+      sanitized += '_';
+    }
+  }
+  return sanitized;
+}
+
+std::vector<std::string> GetWriteReadTestModels() {
+  std::vector<std::string> models;
+  std::string ext(".xml");
+  for (const auto& path : {GetTestDataFilePath("."), GetModelPath(".")}) {
+    for (const auto& p : std::filesystem::recursive_directory_iterator(path)) {
+      if (p.path().extension() == ext) {
+        std::string xml = p.path().string();
+        if (  // if file is meant to fail, skip it
+            absl::StrContains(xml, "malformed_") ||
+            absl::StrContains(xml, "_fail") ||
+            // exclude files that are too slow to load
+            absl::StrContains(xml, "cow") || absl::StrContains(xml, "gmsh_") ||
+            absl::StrContains(xml, "shark_") ||
+            absl::StrContains(xml, "perf") ||
+            // exclude files that fail the comparison test
+            absl::StrContains(xml, "rfcamera") ||
+            absl::StrContains(xml, "tactile") ||
+            absl::StrContains(xml, "makemesh") ||
+            absl::StrContains(xml, "many_dependencies") ||
+            absl::StrContains(xml, "usd") ||
+            absl::StrContains(xml, "torus_maxhull") ||
+            absl::StrContains(xml, "fitmesh_") ||
+            absl::StrContains(xml, "lengthrange") ||
+            absl::StrContains(xml, "hfield_xml") ||
+            absl::StrContains(xml, "fromto_convex") ||
+            absl::StrContains(xml, "cube_skin") ||
+            absl::StrContains(xml, "cube_3x3x3") ||
+            // exclude files that fail since we do not save pinned flex nodes
+            absl::StrContains(xml, "gripper_trilinear") ||
+            absl::StrContains(xml, "strain") ||
+            // exclude mjz testdata with VFS files
+            absl::StrContains(xml, "mixed_test")) {
+          continue;
+        }
+        models.push_back(xml);
+      }
+    }
+  }
+  return models;
+}
+
+using MjzEncoderTest = MujocoTest;
+
+TEST_F(MjzEncoderTest, EncoderIsRegistered) {
+  const mjpEncoder* enc = mjp_findEncoder("model.mjz", nullptr);
+  ASSERT_THAT(enc, testing::NotNull());
+  ASSERT_THAT(enc->encode, testing::NotNull());
+}
+
+TEST_F(MjzEncoderTest, EncodeReturnsPositiveByteCount) {
+  mjSpec* spec = MakeSimpleSpec();
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  const mjpEncoder* enc = mjp_findEncoder("model.mjz", nullptr);
+  ASSERT_THAT(enc, testing::NotNull());
+
+  mjResource resource = {};
+  resource.name = const_cast<char*>("model.mjz");
+
+  int nbytes = enc->encode(spec, model, nullptr, &resource);
+  EXPECT_GT(nbytes, 0);
+  EXPECT_THAT(resource.data, testing::NotNull());
+
+  std::free(resource.data);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, RoundTrip) {
+  mjSpec* spec = MakeSimpleSpec();
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  const std::string tmp_path =
+      testing::TempDir() + "/mjz_encoder_roundtrip.mjz";
+
+  char error[1024] = {0};
+  int nbytes = mj_encode(spec, model, tmp_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mjSpec* decoded =
+      mj_parse(tmp_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull());
+
+  EXPECT_EQ(decoded_model->nbody, model->nbody);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+  EXPECT_EQ(decoded_model->njnt, model->njnt);
+
+  std::remove(tmp_path.c_str());
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripFromXmlString) {
+  const char* xml = R"(
+    <mujoco model="test">
+      <worldbody>
+        <body name="body1" pos="0 0 1">
+          <joint type="hinge"/>
+          <geom type="sphere" size="0.1"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  const std::string tmp_path = testing::TempDir() + "/xml_string.mjz";
+  int nbytes = mj_encode(spec, model, tmp_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mjSpec* decoded =
+      mj_parse(tmp_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull());
+
+  EXPECT_EQ(decoded_model->nbody, model->nbody);
+  EXPECT_EQ(decoded_model->njnt, model->njnt);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  std::remove(tmp_path.c_str());
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripPreservesModelCounts) {
+  const char* xml = R"(
+    <mujoco model="counts_test">
+      <worldbody>
+        <body name="b1" pos="0 0 1">
+          <joint type="hinge"/>
+          <geom type="sphere" size="0.1"/>
+          <body name="b2" pos="0 0 2">
+            <joint type="slide"/>
+            <geom type="box" size="0.1 0.1 0.1"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  const std::string tmp_path = testing::TempDir() + "/counts.mjz";
+  int nbytes = mj_encode(spec, model, tmp_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mjSpec* decoded =
+      mj_parse(tmp_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull());
+
+  EXPECT_EQ(decoded_model->nbody, model->nbody);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+  EXPECT_EQ(decoded_model->njnt, model->njnt);
+  EXPECT_EQ(decoded_model->nq, model->nq);
+  EXPECT_EQ(decoded_model->nv, model->nv);
+
+  std::remove(tmp_path.c_str());
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripWithMeshFromVfs) {
+  const char* xml = R"(
+    <mujoco model="vfs_mesh">
+      <asset>
+        <mesh name="box" file="box.obj"/>
+      </asset>
+      <worldbody>
+        <body name="b1">
+          <geom type="mesh" mesh="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mj_addFileVFS(
+      &vfs, nullptr,
+      GetTestDataFilePath("xml/mjz/testdata/disk_mesh/box.obj").c_str());
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, &vfs, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+  auto* mesh = mjs_asMesh(mjs_firstElement(spec, mjOBJ_MESH));
+  ASSERT_THAT(mesh, testing::NotNull());
+  ASSERT_THAT(mesh->file, testing::NotNull());
+
+  mjModel* model = mj_compile(spec, &vfs);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(spec);
+
+  const std::string tmp_path = testing::TempDir() + "/vfs_mesh.mjz";
+  int nbytes = mj_encode(spec, model, tmp_path.c_str(), nullptr, &vfs, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Reset VFS
+  mj_deleteVFS(&vfs);
+  mj_defaultVFS(&vfs);
+
+  mjSpec* decoded =
+      mj_parse(tmp_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  std::remove(tmp_path.c_str());
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripWithMeshFromDisk) {
+  std::string xml_path =
+      GetTestDataFilePath("xml/mjz/testdata/disk_mesh/model.xml");
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "disk_mesh_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, RoundTripWithMeshInSubdirectory) {
+  std::string xml_path =
+      GetTestDataFilePath("xml/mjz/testdata/subdir_mesh/model.xml");
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "subdir_mesh_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, RoundTripWithMeshdir) {
+  std::string xml_path =
+      GetTestDataFilePath("xml/mjz/testdata/meshdir_test/model.xml");
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "meshdir_test_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, RoundTripWithInclude) {
+  std::string xml_path =
+      GetTestDataFilePath("xml/mjz/testdata/include_test/parent.xml");
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "include_test_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull());
+
+  EXPECT_EQ(decoded_model->nbody, model->nbody);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  fs::remove_all(tmpdir);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripTransitiveInclude) {
+  std::string xml_path =
+      GetTestDataFilePath("xml/mjz/testdata/transitive_test/parent.xml");
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "transitive_test_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull());
+
+  EXPECT_EQ(decoded_model->nbody, model->nbody);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  fs::remove_all(tmpdir);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, RoundTripMixedDiskAndVfs) {
+  static const char* xml = R"(
+    <mujoco model="mixed_test">
+      <asset>
+        <mesh name="disk_box" />
+        <mesh name="vfs_box" file="xml/mjz/testdata/disk_mesh/box.obj"/>
+      </asset>
+      <worldbody>
+        <body name="b1">
+          <geom type="mesh" mesh="disk_box"/>
+        </body>
+        <body name="b2">
+          <geom type="mesh" mesh="vfs_box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+
+  )";
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+  mj_addFileVFS(
+      &vfs, nullptr,
+      GetTestDataFilePath("xml/mjz/testdata/disk_mesh/box.obj").c_str());
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, &vfs, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  // Fix the on disk mesh path
+  mjsMesh* mesh = mjs_asMesh(mjs_findElement(spec, mjOBJ_MESH, "disk_box"));
+  ASSERT_THAT(mesh, testing::NotNull());
+  mjs_setString(
+      mesh->file,
+      GetTestDataFilePath("xml/mjz/testdata/mixed_test/disk_box.obj").c_str());
+
+  mjModel* model = mj_compile(spec, &vfs);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(spec);
+
+  fs::path tmpdir = fs::path(testing::TempDir()) / "mixed_test_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, &vfs, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+  EXPECT_EQ(decoded_model->ngeom, model->ngeom);
+
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+  mj_deleteVFS(&vfs);
+}
+
+TEST_F(MjzEncoderTest, AttachedSpecsWithCollidingAssetNames) {
+  std::string child_a_xml = GetTestDataFilePath(
+      "xml/mjz/testdata/asset_collision_test/child_a/child.xml");
+  std::string child_b_xml = GetTestDataFilePath(
+      "xml/mjz/testdata/asset_collision_test/child_b/child.xml");
+
+  char error[1024] = {0};
+
+  // Parse both children from their respective directories.
+  mjSpec* child_a =
+      mj_parse(child_a_xml.c_str(), nullptr, nullptr, error, sizeof(error));
+  ASSERT_THAT(child_a, testing::NotNull()) << error;
+  mjSpec* child_b =
+      mj_parse(child_b_xml.c_str(), nullptr, nullptr, error, sizeof(error));
+  ASSERT_THAT(child_b, testing::NotNull()) << error;
+
+  // Create parent spec and attach both children.
+  mjSpec* parent = mj_makeSpec();
+  ASSERT_THAT(parent, testing::NotNull());
+  mjsBody* world = mjs_findBody(parent, "world");
+
+  // Attach child_a with prefix "a_".
+  mjsFrame* frame_a = mjs_addFrame(world, nullptr);
+  mjs_attach(frame_a->element, mjs_findBody(child_a, "body")->element, "a_",
+             "");
+
+  // Attach child_b with prefix "b_".
+  mjsFrame* frame_b = mjs_addFrame(world, nullptr);
+  mjs_attach(frame_b->element, mjs_findBody(child_b, "body")->element, "b_",
+             "");
+
+  // Compile the combined model.
+  mjModel* model = mj_compile(parent, nullptr);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(parent);
+  EXPECT_EQ(model->nmesh, 2) << "Expected two distinct meshes";
+
+  // Record the original vertex counts for both meshes.
+  int mesh0_nvert = model->mesh_vertnum[0];
+  int mesh1_nvert = model->mesh_vertnum[1];
+  EXPECT_NE(mesh0_nvert, mesh1_nvert)
+      << "The two meshes should have different vertex counts";
+
+  // Encode to MJZ.
+  fs::path tmpdir = fs::path(testing::TempDir()) / "asset_collision_test_out";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+  int nbytes = mj_encode(parent, model, out_path.c_str(), nullptr, nullptr,
+                         error, sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Decode from MJZ.
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  // Verify both meshes survived with correct, distinct vertex counts.
+  EXPECT_EQ(decoded_model->nmesh, 2) << "Decoded model should have two meshes";
+  EXPECT_EQ(decoded_model->mesh_vertnum[0], mesh0_nvert)
+      << "First mesh vertex count mismatch after roundtrip";
+  EXPECT_EQ(decoded_model->mesh_vertnum[1], mesh1_nvert)
+      << "Second mesh vertex count mismatch after roundtrip";
+
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(parent);
+  mj_deleteSpec(child_a);
+  mj_deleteSpec(child_b);
+}
+
+// --- Fake resource provider for URI stripping tests ---
+
+const char* kBoxObj = R"(
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
+)";
+
+// Static STL data served by the fake provider.
+static std::vector<char> g_fake_data;
+
+int FakeUriOpen(mjResource* resource) {
+  if (g_fake_data.empty()) return 0;
+  resource->data = g_fake_data.data();
+  return 1;
+}
+
+int FakeUriRead(mjResource* resource, const void** buffer) {
+  if (!resource->data) return -1;
+  *buffer = resource->data;
+  return static_cast<int>(g_fake_data.size());
+}
+
+void FakeUriClose(mjResource* resource) { resource->data = nullptr; }
+
+// Register the fake provider once (idempotent due to AppendIfUnique).
+static void EnsureFakeProviderRegistered() {
+  static bool registered = false;
+  if (registered) return;
+  mjpResourceProvider provider = {};
+  mjp_defaultResourceProvider(&provider);
+  provider.prefix = "testmjzenc";
+  provider.open = FakeUriOpen;
+  provider.read = FakeUriRead;
+  provider.close = FakeUriClose;
+  mjp_registerResourceProvider(&provider);
+  registered = true;
+}
+
+TEST_F(MjzEncoderTest, RoundTripStripsUriPrefix) {
+  EnsureFakeProviderRegistered();
+
+  // Set up the fake STL data that the provider will serve.
+  g_fake_data.assign(kBoxObj, kBoxObj + std::strlen(kBoxObj));
+
+  const char* xml = R"(
+    <mujoco model="uri_test">
+      <asset>
+        <mesh name="uri_box" file="testmjzenc:meshes/box.obj"/>
+      </asset>
+      <worldbody>
+        <body name="body">
+          <geom type="mesh" mesh="uri_box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  // Compile.
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(spec);
+  EXPECT_EQ(model->nmesh, 1);
+
+  // Encode to MJZ.
+  fs::path tmpdir = fs::path(testing::TempDir()) / "uri_strip_test";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Decode from MJZ.
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  // Verify the decoded mesh file attribute does NOT contain the URI prefix.
+  auto* decoded_mesh = mjs_asMesh(mjs_firstElement(decoded, mjOBJ_MESH));
+  ASSERT_THAT(decoded_mesh, testing::NotNull());
+  const char* decoded_file = mjs_getString(decoded_mesh->file);
+  ASSERT_THAT(decoded_file, testing::NotNull());
+  EXPECT_FALSE(absl::StrContains(decoded_file, ":"))
+      << "Decoded mesh file should not contain ':': " << decoded_file;
+  EXPECT_THAT(decoded_file, testing::HasSubstr("box.obj"));
+
+  // Verify the model compiles correctly from the decoded spec.
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+
+  // Cleanup.
+  g_fake_data.clear();
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, RoundTripStripsUriWithSpecialChars) {
+  EnsureFakeProviderRegistered();
+
+  // Set up the fake STL data.
+  g_fake_data.assign(kBoxObj, kBoxObj + std::strlen(kBoxObj));
+
+  const char* xml = R"(
+    <mujoco model="uri_special_test">
+      <asset>
+        <mesh name="uri_box" file="testmjzenc://store/abc+123%20v2/box.obj"/>
+      </asset>
+      <worldbody>
+        <body name="body">
+          <geom type="mesh" mesh="uri_box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull());
+
+  // Encode to MJZ.
+  fs::path tmpdir = fs::path(testing::TempDir()) / "uri_special_test";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Decode and verify no URI prefix and no special chars in the file attr.
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  // Check the decoded mesh file attribute.
+  auto* decoded_mesh = mjs_asMesh(mjs_firstElement(decoded, mjOBJ_MESH));
+  ASSERT_THAT(decoded_mesh, testing::NotNull());
+  const char* decoded_file = mjs_getString(decoded_mesh->file);
+  ASSERT_THAT(decoded_file, testing::NotNull());
+  EXPECT_FALSE(absl::StrContains(decoded_file, "testmjzenc://"))
+      << "Decoded mesh file should not contain the URI prefix: "
+      << decoded_file;
+  EXPECT_FALSE(absl::StrContains(decoded_file, "+"))
+      << "Decoded mesh file should not contain '+': " << decoded_file;
+  EXPECT_FALSE(absl::StrContains(decoded_file, "%"))
+      << "Decoded mesh file should not contain '%': " << decoded_file;
+  EXPECT_THAT(decoded_file, testing::HasSubstr(".obj"));
+
+  // Verify roundtrip compiles.
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+
+  g_fake_data.clear();
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+class MjzEncoderParameterizedTest
+    : public MujocoTest,
+      public ::testing::WithParamInterface<std::string> {};
+
+TEST_P(MjzEncoderParameterizedTest, WriteReadCompare) {
+  mujoco::FullFloatPrecision increase_precision;
+
+  std::string xml = GetParam();
+  std::array<char, 1000> error;
+
+  mjSpec* s = mj_parseXML(xml.c_str(), nullptr, error.data(), error.size());
+  ASSERT_THAT(s, NotNull()) << error.data();
+
+  mjModel* m = mj_compile(s, nullptr);
+  ASSERT_THAT(m, NotNull()) << mjs_getError(s);
+
+  const std::string tmp_path = testing::TempDir() + "/mjz_roundtrip.mjz";
+  int nbytes = mj_encode(s, m, tmp_path.c_str(), nullptr, nullptr, error.data(),
+                         error.size());
+  ASSERT_GT(nbytes, 0) << error.data();
+  mj_deleteSpec(s);
+
+  mjVFS vfs;
+  mj_defaultVFS(&vfs);
+
+  mjSpec* stemp =
+      mj_parse(tmp_path.c_str(), nullptr, &vfs, error.data(), error.size());
+  ASSERT_THAT(stemp, NotNull()) << error.data();
+
+  mjModel* mtemp = mj_compile(stemp, &vfs);
+  ASSERT_THAT(mtemp, NotNull()) << mjs_getError(stemp);
+  mj_deleteSpec(stemp);
+
+  mjtNum tol = 0;
+  if (absl::StrContains(xml, "belt.xml") ||
+      absl::StrContains(xml, "cable.xml")) {
+    tol = 1e-13;
+  }
+
+  // Make paths identical to avoid failure in CompareModel due to localization.
+  // For example, we might localize "../../y" to "y", which changes the paths.
+  char* old_m_paths = m->paths;
+  char* old_mtemp_paths = mtemp->paths;
+  int old_m_npaths = m->npaths;
+  int old_mtemp_npaths = mtemp->npaths;
+
+  std::string dummy = "ignored_path";
+  int dummy_len = dummy.length() + 1;
+
+  m->paths = (char*)mju_malloc(dummy_len);
+  mtemp->paths = (char*)mju_malloc(dummy_len);
+
+  std::memcpy(m->paths, dummy.c_str(), dummy_len);
+  std::memcpy(mtemp->paths, dummy.c_str(), dummy_len);
+
+  m->npaths = dummy_len;
+  mtemp->npaths = dummy_len;
+
+  std::string field = "";
+  mjtNum result = CompareModel(m, mtemp, field);
+
+  // Restore old pointers and sizes so mj_deleteModel can free them correctly!
+  mju_free(m->paths);
+  mju_free(mtemp->paths);
+  m->paths = old_m_paths;
+  mtemp->paths = old_mtemp_paths;
+  m->npaths = old_m_npaths;
+  mtemp->npaths = old_mtemp_npaths;
+
+  EXPECT_LE(result, tol) << "Loaded and saved models are different!\n"
+                         << "Affected file: " << xml << '\n'
+                         << "Different field: " << field << '\n';
+
+  mj_deleteVFS(&vfs);
+  mj_deleteModel(mtemp);
+  mj_deleteModel(m);
+  std::remove(tmp_path.c_str());
+}
+
+INSTANTIATE_TEST_SUITE_P(MjzEncoderParameterizedTests,
+                         MjzEncoderParameterizedTest,
+                         ::testing::ValuesIn(GetWriteReadTestModels()),
+                         [](const ::testing::TestParamInfo<std::string>& info) {
+                           return SanitizePathForTestName(info.param);
+                         });
+
+TEST_F(MjzEncoderTest, DuplicateFileReferencesRewrite) {
+  EnsureFakeProviderRegistered();
+
+  // Set up the fake data that the provider will serve.
+  g_fake_data.assign(kBoxObj, kBoxObj + std::strlen(kBoxObj));
+
+  // Two meshes referencing the same URI-prefixed file.
+  const char* xml = R"(
+    <mujoco model="dup_file_test">
+      <asset>
+        <mesh name="box_a" file="testmjzenc:meshes/box.obj"/>
+        <mesh name="box_b" file="testmjzenc:meshes/box.obj"/>
+      </asset>
+      <worldbody>
+        <body name="a">
+          <geom type="mesh" mesh="box_a"/>
+        </body>
+        <body name="b">
+          <geom type="mesh" mesh="box_b"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(spec);
+
+  // Encode to MJZ.
+  fs::path tmpdir = fs::path(testing::TempDir()) / "dup_file_test";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Decode from MJZ.
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  // Both decoded meshes should have sanitized file paths (no URI prefix).
+  auto* mesh_a = mjs_asMesh(mjs_firstElement(decoded, mjOBJ_MESH));
+  ASSERT_THAT(mesh_a, testing::NotNull());
+  const char* file_a = mjs_getString(mesh_a->file);
+  ASSERT_THAT(file_a, testing::NotNull());
+  EXPECT_FALSE(absl::StrContains(file_a, ":"))
+      << "First mesh file should not contain ':': " << file_a;
+
+  auto* mesh_b = mjs_asMesh(mjs_nextElement(decoded, mesh_a->element));
+  ASSERT_THAT(mesh_b, testing::NotNull());
+  const char* file_b = mjs_getString(mesh_b->file);
+  ASSERT_THAT(file_b, testing::NotNull());
+  EXPECT_FALSE(absl::StrContains(file_b, ":"))
+      << "Second mesh file should not contain ':': " << file_b;
+
+  // Both should compile correctly.
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+  EXPECT_EQ(decoded_model->nmesh, model->nmesh);
+
+  // Cleanup.
+  g_fake_data.clear();
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+TEST_F(MjzEncoderTest, CubefileUriRewrite) {
+  EnsureFakeProviderRegistered();
+
+  // The fake provider serves arbitrary data; cubefiles need valid PNG data to
+  // compile textures, but the encoder just packs bytes. We verify archive
+  // contents and XML rewriting without requiring valid image data.
+  // Use a minimal valid 1x1 white PNG for each face.
+  static const unsigned char kTinyPng[] = {
+      0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
+      0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
+      0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xde, 0x00, 0x00, 0x00,
+      0x0c, 0x49, 0x44, 0x41, 0x54, 0x78, 0x9c, 0x63, 0xf8, 0xff, 0xff, 0x3f,
+      0x00, 0x05, 0xfe, 0x02, 0xfe, 0x0d, 0xef, 0x46, 0xb8, 0x00, 0x00, 0x00,
+      0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82,
+  };
+  g_fake_data.assign(
+      reinterpret_cast<const char*>(kTinyPng),
+      reinterpret_cast<const char*>(kTinyPng) + sizeof(kTinyPng));
+
+  const char* xml = R"(
+    <mujoco model="cubefile_test">
+      <asset>
+        <texture name="sky" type="cube"
+                 fileright="testmjzenc:faces/right.png"
+                 fileleft="testmjzenc:faces/left.png"
+                 fileup="testmjzenc:faces/up.png"
+                 filedown="testmjzenc:faces/down.png"
+                 filefront="testmjzenc:faces/front.png"
+                 fileback="testmjzenc:faces/back.png"/>
+        <material name="sky_mat" texture="sky"/>
+      </asset>
+      <worldbody>
+        <body name="body">
+          <geom type="sphere" size="0.1" material="sky_mat"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  )";
+
+  char error[1024] = {0};
+  mjSpec* spec = mj_parseXMLString(xml, nullptr, error, sizeof(error));
+  ASSERT_THAT(spec, testing::NotNull()) << error;
+
+  mjModel* model = mj_compile(spec, nullptr);
+  ASSERT_THAT(model, testing::NotNull()) << mjs_getError(spec);
+
+  // Encode to MJZ.
+  fs::path tmpdir = fs::path(testing::TempDir()) / "cubefile_test";
+  fs::create_directories(tmpdir);
+  const std::string out_path = (tmpdir / "model.mjz").string();
+  int nbytes = mj_encode(spec, model, out_path.c_str(), nullptr, nullptr, error,
+                         sizeof(error));
+  ASSERT_GT(nbytes, 0) << error;
+
+  // Decode from MJZ.
+  mjVFS decode_vfs;
+  mj_defaultVFS(&decode_vfs);
+  mjSpec* decoded =
+      mj_parse(out_path.c_str(), nullptr, &decode_vfs, error, sizeof(error));
+  ASSERT_THAT(decoded, testing::NotNull()) << error;
+
+  // Verify the decoded texture's cubefile paths don't contain URI prefix.
+  auto* decoded_tex = mjs_asTexture(mjs_firstElement(decoded, mjOBJ_TEXTURE));
+  ASSERT_THAT(decoded_tex, testing::NotNull());
+  ASSERT_THAT(decoded_tex->cubefiles, testing::NotNull());
+  for (const mjString& cubefile : *decoded_tex->cubefiles) {
+    EXPECT_FALSE(absl::StrContains(cubefile, ":"))
+        << "Cubefile should not contain ':': " << cubefile;
+    EXPECT_THAT(std::string(cubefile), testing::HasSubstr(".png"));
+  }
+
+  // Verify the model compiles correctly from the decoded spec.
+  mjModel* decoded_model = mj_compile(decoded, &decode_vfs);
+  ASSERT_THAT(decoded_model, testing::NotNull()) << mjs_getError(decoded);
+
+  // Cleanup.
+  g_fake_data.clear();
+  fs::remove_all(tmpdir);
+  mj_deleteVFS(&decode_vfs);
+  mj_deleteModel(decoded_model);
+  mj_deleteSpec(decoded);
+  mj_deleteModel(model);
+  mj_deleteSpec(spec);
+}
+
+}  // namespace mujoco
diff --git a/test/xml/mjz/testdata/asset_collision_test/child_a/child.xml b/test/xml/mjz/testdata/asset_collision_test/child_a/child.xml
new file mode 100644
index 00000000..60850688
--- /dev/null
+++ b/test/xml/mjz/testdata/asset_collision_test/child_a/child.xml
@@ -0,0 +1,10 @@
+<mujoco model="child_a">
+  <asset>
+    <mesh name="m" file="mesh.obj"/>
+  </asset>
+  <worldbody>
+    <body name="body">
+      <geom type="mesh" mesh="m"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/asset_collision_test/child_a/mesh.obj b/test/xml/mjz/testdata/asset_collision_test/child_a/mesh.obj
new file mode 100644
index 00000000..b4fbc738
--- /dev/null
+++ b/test/xml/mjz/testdata/asset_collision_test/child_a/mesh.obj
@@ -0,0 +1,20 @@
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
diff --git a/test/xml/mjz/testdata/asset_collision_test/child_b/child.xml b/test/xml/mjz/testdata/asset_collision_test/child_b/child.xml
new file mode 100644
index 00000000..1886f514
--- /dev/null
+++ b/test/xml/mjz/testdata/asset_collision_test/child_b/child.xml
@@ -0,0 +1,10 @@
+<mujoco model="child_b">
+  <asset>
+    <mesh name="m" file="mesh.obj"/>
+  </asset>
+  <worldbody>
+    <body name="body">
+      <geom type="mesh" mesh="m"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/asset_collision_test/child_b/mesh.obj b/test/xml/mjz/testdata/asset_collision_test/child_b/mesh.obj
new file mode 100644
index 00000000..dfc38fb3
--- /dev/null
+++ b/test/xml/mjz/testdata/asset_collision_test/child_b/mesh.obj
@@ -0,0 +1,8 @@
+v 0 0 0
+v 1 0 0
+v 0.5 0.866 0
+v 0.5 0.289 0.816
+f 1 2 3
+f 1 2 4
+f 2 3 4
+f 1 3 4
diff --git a/test/xml/mjz/testdata/disk_mesh/box.obj b/test/xml/mjz/testdata/disk_mesh/box.obj
new file mode 100644
index 00000000..b4fbc738
--- /dev/null
+++ b/test/xml/mjz/testdata/disk_mesh/box.obj
@@ -0,0 +1,20 @@
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
diff --git a/test/xml/mjz/testdata/disk_mesh/model.xml b/test/xml/mjz/testdata/disk_mesh/model.xml
new file mode 100644
index 00000000..e12ba5b8
--- /dev/null
+++ b/test/xml/mjz/testdata/disk_mesh/model.xml
@@ -0,0 +1,10 @@
+<mujoco model="disk_mesh">
+  <asset>
+    <mesh name="box" file="box.obj"/>
+  </asset>
+  <worldbody>
+    <body name="b1">
+      <geom type="mesh" mesh="box"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/include_test/child.xml b/test/xml/mjz/testdata/include_test/child.xml
new file mode 100644
index 00000000..5bd5d122
--- /dev/null
+++ b/test/xml/mjz/testdata/include_test/child.xml
@@ -0,0 +1,7 @@
+<mujoco>
+  <worldbody>
+    <body name="child_body" pos="1 0 0">
+      <geom type="sphere" size="0.05"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/include_test/parent.xml b/test/xml/mjz/testdata/include_test/parent.xml
new file mode 100644
index 00000000..22a0b82f
--- /dev/null
+++ b/test/xml/mjz/testdata/include_test/parent.xml
@@ -0,0 +1,8 @@
+<mujoco model="include_test">
+  <include file="child.xml"/>
+  <worldbody>
+    <body name="parent_body" pos="0 0 1">
+      <geom type="sphere" size="0.1"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/meshdir_test/assets/box.obj b/test/xml/mjz/testdata/meshdir_test/assets/box.obj
new file mode 100644
index 00000000..b4fbc738
--- /dev/null
+++ b/test/xml/mjz/testdata/meshdir_test/assets/box.obj
@@ -0,0 +1,20 @@
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
diff --git a/test/xml/mjz/testdata/meshdir_test/model.xml b/test/xml/mjz/testdata/meshdir_test/model.xml
new file mode 100644
index 00000000..c3defdba
--- /dev/null
+++ b/test/xml/mjz/testdata/meshdir_test/model.xml
@@ -0,0 +1,11 @@
+<mujoco model="meshdir_test">
+  <compiler meshdir="assets"/>
+  <asset>
+    <mesh name="box" file="box.obj"/>
+  </asset>
+  <worldbody>
+    <body name="b1">
+      <geom type="mesh" mesh="box"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/mixed_test/disk_box.obj b/test/xml/mjz/testdata/mixed_test/disk_box.obj
new file mode 100644
index 00000000..b4fbc738
--- /dev/null
+++ b/test/xml/mjz/testdata/mixed_test/disk_box.obj
@@ -0,0 +1,20 @@
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
diff --git a/test/xml/mjz/testdata/subdir_mesh/meshes/box.obj b/test/xml/mjz/testdata/subdir_mesh/meshes/box.obj
new file mode 100644
index 00000000..b4fbc738
--- /dev/null
+++ b/test/xml/mjz/testdata/subdir_mesh/meshes/box.obj
@@ -0,0 +1,20 @@
+v 0 0 0
+v 1 0 0
+v 1 1 0
+v 0 1 0
+v 0 0 1
+v 1 0 1
+v 1 1 1
+v 0 1 1
+f 1 2 3
+f 1 3 4
+f 5 6 7
+f 5 7 8
+f 1 2 6
+f 1 6 5
+f 2 3 7
+f 2 7 6
+f 3 4 8
+f 3 8 7
+f 4 1 5
+f 4 5 8
diff --git a/test/xml/mjz/testdata/subdir_mesh/model.xml b/test/xml/mjz/testdata/subdir_mesh/model.xml
new file mode 100644
index 00000000..e3e001df
--- /dev/null
+++ b/test/xml/mjz/testdata/subdir_mesh/model.xml
@@ -0,0 +1,10 @@
+<mujoco model="subdir_mesh">
+  <asset>
+    <mesh name="box" file="meshes/box.obj"/>
+  </asset>
+  <worldbody>
+    <body name="b1">
+      <geom type="mesh" mesh="box"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/transitive_test/child.xml b/test/xml/mjz/testdata/transitive_test/child.xml
new file mode 100644
index 00000000..23724726
--- /dev/null
+++ b/test/xml/mjz/testdata/transitive_test/child.xml
@@ -0,0 +1,8 @@
+<mujoco>
+  <include file="grandchild.xml"/>
+  <worldbody>
+    <body name="child_body" pos="1 0 0">
+      <geom type="sphere" size="0.05"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/transitive_test/grandchild.xml b/test/xml/mjz/testdata/transitive_test/grandchild.xml
new file mode 100644
index 00000000..7f209770
--- /dev/null
+++ b/test/xml/mjz/testdata/transitive_test/grandchild.xml
@@ -0,0 +1,7 @@
+<mujoco>
+  <worldbody>
+    <body name="gc_body" pos="2 0 0">
+      <geom type="sphere" size="0.02"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/test/xml/mjz/testdata/transitive_test/parent.xml b/test/xml/mjz/testdata/transitive_test/parent.xml
new file mode 100644
index 00000000..3780230f
--- /dev/null
+++ b/test/xml/mjz/testdata/transitive_test/parent.xml
@@ -0,0 +1,8 @@
+<mujoco model="transitive_test">
+  <include file="child.xml"/>
+  <worldbody>
+    <body name="parent_body" pos="0 0 1">
+      <geom type="sphere" size="0.1"/>
+    </body>
+  </worldbody>
+</mujoco>

From 53b3137a12e86cbf9580dcf264a53652c550ab22 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Wed, 20 May 2026 10:00:30 -0700
Subject: [PATCH 07/61] Separate Light management into a separate class.

PiperOrigin-RevId: 918508291
Change-Id: I9cadc1b612006a0450f8a2fd9088d56cc2e610a3
---
 src/experimental/filament/CMakeLists.txt      |   2 +
 .../filament/compat/light_manager.cc          | 213 ++++++++++++++++++
 .../filament/compat/light_manager.h           |  56 +++++
 .../filament/compat/scene_bridge.cc           | 192 ++--------------
 .../filament/compat/scene_bridge.h            |  13 +-
 5 files changed, 291 insertions(+), 185 deletions(-)
 create mode 100644 src/experimental/filament/compat/light_manager.cc
 create mode 100644 src/experimental/filament/compat/light_manager.h

diff --git a/src/experimental/filament/CMakeLists.txt b/src/experimental/filament/CMakeLists.txt
index 99639d24..7f5aac1d 100644
--- a/src/experimental/filament/CMakeLists.txt
+++ b/src/experimental/filament/CMakeLists.txt
@@ -51,6 +51,8 @@ target_sources(${MUJOCO_FILAMENT_TARGET_NAME}
         filament/scene_view.h
         filament/texture.cc
         filament/texture.h
+        compat/light_manager.cc
+        compat/light_manager.h
         compat/model_objects.cc
         compat/model_objects.h
         compat/scene_bridge.cc
diff --git a/src/experimental/filament/compat/light_manager.cc b/src/experimental/filament/compat/light_manager.cc
new file mode 100644
index 00000000..7ee1973a
--- /dev/null
+++ b/src/experimental/filament/compat/light_manager.cc
@@ -0,0 +1,213 @@
+// 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.
+
+#include "experimental/filament/compat/light_manager.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <math/mat4.h>
+#include <math/mathfwd.h>
+#include <math/vec3.h>
+#include <math/vec4.h>
+#include <mujoco/mujoco.h>
+#include "experimental/filament/compat/model_objects.h"
+#include "experimental/filament/filament_util.h"
+#include "experimental/filament/render_context_filament_cpp.h"
+#include "experimental/filament/render_context_filament.h"
+
+namespace mujoco {
+
+using filament::math::float3;
+using filament::math::float4;
+using filament::math::mat3;
+using filament::math::mat4;
+
+static UniquePtr<mjrTexture> CreateFallbackIndirectLightTexture(
+    mjrfContext* ctx) {
+  const std::string filename = ResolveFilamentAssetPath("ibl.ktx");
+  mjResource* resource =
+      mju_openResource("", filename.c_str(), nullptr, nullptr, 0);
+  if (!resource) {
+    mju_error("Failed to open resource: %s", filename.c_str());
+  }
+  const void* bytes = nullptr;
+  const int nbytes = mju_readResource(resource, &bytes);
+  if (bytes == nullptr || nbytes <= 0) {
+    mju_error("Failed to read resource: %s", filename.c_str());
+  }
+
+  mjrTextureConfig config;
+  mjr_defaultTextureConfig(&config);
+  config.width = 1;
+  config.height = 1;
+  config.sampler_type = mjTEXTURE_CUBE;
+  config.format = mjPIXEL_FORMAT_KTX;
+  config.color_space = mjCOLORSPACE_AUTO;
+
+  auto texture = CreateTexture(ctx, config);
+
+  mjrTextureData payload;
+  mjr_defaultTextureData(&payload);
+  payload.bytes = bytes;
+  payload.nbytes = nbytes;
+  payload.release_callback = +[](void* user_data) {
+    mju_closeResource((mjResource*)user_data);
+  };
+  payload.user_data = resource;
+
+  mjrf_setTextureData(texture.get(), &payload);
+  return texture;
+}
+
+LightManager::LightManager(mjrfContext* ctx, mjrScene* scene,
+                           ModelObjects* model_objects)
+    : ctx_(ctx), scene_(scene) {
+  const mjModel* model = model_objects->GetModel();
+  default_shadow_map_size_ = ReadElement(
+      model, "filament.shadows.map_size", default_shadow_map_size_);
+  default_vsm_blur_width_ = ReadElement(
+      model, "filament.shadows.vsm_blur_width", default_vsm_blur_width_);
+  fallback_head_light_intensity_ =
+      ReadElement(model, "filament.fallback.head_light_intensity",
+                  fallback_head_light_intensity_);
+  fallback_scene_light_intensity_ =
+      ReadElement(model, "filament.fallback.scene_light_intensity",
+                  fallback_scene_light_intensity_);
+  fallback_environment_light_intensity_ =
+      ReadElement(model, "filament.fallback.environment_light_intensity",
+                  fallback_environment_light_intensity_);
+  Prepare(model_objects);
+}
+
+LightManager::~LightManager() {
+  for (auto& iter : lights_) {
+    mjrf_removeLightFromScene(scene_, iter.get());
+  }
+  lights_.clear();
+  if (fallback_ibl_) {
+    mjrf_removeLightFromScene(scene_, fallback_ibl_.get());
+  }
+  fallback_ibl_.reset();
+}
+
+void LightManager::Prepare(ModelObjects* model_objects) {
+  const mjModel* model = model_objects->GetModel();
+
+  bool has_image_based_light = false;
+  float total_light_intensity = 0.0f;
+  for (int i = 0; i < model->nlight; ++i) {
+    total_light_intensity += model->light_intensity[i];
+
+    if (model->light_type[i] == mjLIGHT_IMAGE) {
+      mjrLightParams params;
+      mjr_defaultLightParams(&params);
+      params.type = mjLIGHT_IMAGE;
+      params.texture = model_objects->GetTexture(model->light_texid[i]);
+      params.intensity = model->light_intensity[i];
+      auto light_obj = CreateLight(ctx_, params);
+      mjrf_addLightToScene(scene_, light_obj.get());
+      lights_.emplace_back(std::move(light_obj));
+      has_image_based_light = true;
+    } else {
+      mjrLightParams params;
+      mjr_defaultLightParams(&params);
+      params.color[0] = model->light_diffuse[0];
+      params.color[1] = model->light_diffuse[1];
+      params.color[2] = model->light_diffuse[2];
+      params.type = (mjtLightType)model->light_type[i];
+      params.cast_shadows = model->light_castshadow[i];
+      params.bulb_radius = model->light_bulbradius[i];
+      params.range = model->light_range[i];
+      params.intensity = model->light_intensity[i];
+      params.shadow_map_size = default_shadow_map_size_;
+      params.vsm_blur_width = default_vsm_blur_width_;
+      if (params.type == mjLIGHT_SPOT) {
+        params.spot_cone_angle = model->light_cutoff[i];
+      }
+
+      auto light_obj = CreateLight(ctx_, params);
+      mjrf_addLightToScene(scene_, light_obj.get());
+      lights_.emplace_back(std::move(light_obj));
+    }
+  }
+
+  // Add a placeholder (black) headlight as our last light. Going forward, we'll
+  // assume lights_.back() is always the headlight.
+  {
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    // We break with the spec here slightly and use a spot light for the head
+    // light instead of a directional params. This is because filament only
+    // supports a single directional light, and we'd rather allow a scene
+    // light to be that directional params. It's also a bit odd for a
+    // directional light to move with the camera.
+    params.type = mjLIGHT_SPOT;
+    params.cast_shadows = 0;
+    params.intensity = 0.0f;
+    params.spot_cone_angle = 90.0f;
+    auto light_obj = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, light_obj.get());
+    lights_.emplace_back(std::move(light_obj));
+  }
+
+  if (!has_image_based_light && total_light_intensity > 0.0f) {
+    // Create a black indirect light to ensure that the skybox is
+    // oriented to respect mujoco's Z-up convention.
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    params.type = mjLIGHT_IMAGE;
+    params.intensity = 10.0f;
+    fallback_ibl_ = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, fallback_ibl_.get());
+  }
+
+  // There are no "physical" lights in the scene which means we're likely
+  // dealing with a "classic renderer" scene. In this case, let's add a
+  // default environment light and set the light intensity ourselves.
+  if (total_light_intensity == 0.0f) {
+    // Create a fallback environment light.
+    fallback_ibl_texture_ = CreateFallbackIndirectLightTexture(ctx_);
+
+    mjrLightParams params;
+    mjr_defaultLightParams(&params);
+    params.type = mjLIGHT_IMAGE;
+    params.texture = fallback_ibl_texture_.get();
+    params.intensity = fallback_environment_light_intensity_;
+    fallback_ibl_ = CreateLight(ctx_, params);
+    mjrf_addLightToScene(scene_, fallback_ibl_.get());
+
+    // Distribute the fallback scene light intensity among the lights.
+    const float intensity = fallback_scene_light_intensity_ / lights_.size();
+    for (auto& light : lights_) {
+      if (light) {
+        const bool is_headlight = (light == lights_.back());
+        mjrf_setLightIntensity(light.get(),
+                               is_headlight ? fallback_head_light_intensity_
+                                            : intensity);
+      }
+    }
+  }
+
+  mjrf_setSceneSkybox(scene_, model_objects->GetSkyboxTexture());
+}
+
+mjrLight* LightManager::GetLight(int index) {
+  if (index < 0 || index >= lights_.size()) {
+    return nullptr;
+  }
+  return lights_[index].get();
+}
+}  // namespace mujoco
diff --git a/src/experimental/filament/compat/light_manager.h b/src/experimental/filament/compat/light_manager.h
new file mode 100644
index 00000000..1f81b144
--- /dev/null
+++ b/src/experimental/filament/compat/light_manager.h
@@ -0,0 +1,56 @@
+// 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_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_
+#define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_
+
+#include <vector>
+
+#include "experimental/filament/compat/model_objects.h"
+#include "experimental/filament/render_context_filament.h"
+#include "experimental/filament/render_context_filament_cpp.h"
+
+namespace mujoco {
+
+// Manages Light entities for an mjrScene using data from an mjvScene.
+class LightManager {
+ public:
+  LightManager(mjrfContext* ctx, mjrScene* scene, ModelObjects* model_objects);
+  ~LightManager();
+
+  // Returns the light with the given index in the mjModel. Note that an extra
+  // headlight is assigned of the index `nlight`.
+  mjrLight* GetLight(int index);
+
+  LightManager(const LightManager&) = delete;
+  LightManager& operator=(const LightManager&) = delete;
+
+ private:
+  void Prepare(ModelObjects* model_objects);
+
+  mjrfContext* ctx_ = nullptr;
+  mjrScene* scene_ = nullptr;
+  UniquePtr<mjrLight> fallback_ibl_{nullptr, nullptr};
+  UniquePtr<mjrTexture> fallback_ibl_texture_{nullptr, nullptr};
+  std::vector<UniquePtr<mjrLight>> lights_;
+  int default_shadow_map_size_ = 2048;
+  float default_vsm_blur_width_ = 0.0f;
+  float fallback_head_light_intensity_ = 0.f;
+  float fallback_scene_light_intensity_ = 80'000.f;
+  float fallback_environment_light_intensity_ = 5'000.f;
+};
+
+}  // namespace mujoco
+
+#endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_COMPAT_LIGHT_MANAGER_H_
diff --git a/src/experimental/filament/compat/scene_bridge.cc b/src/experimental/filament/compat/scene_bridge.cc
index ba2057b9..61a24d2f 100644
--- a/src/experimental/filament/compat/scene_bridge.cc
+++ b/src/experimental/filament/compat/scene_bridge.cc
@@ -16,8 +16,6 @@
 
 #include <memory>
 #include <optional>
-#include <string>
-#include <string_view>
 #include <utility>
 
 #include <math/TMatHelpers.h>
@@ -27,6 +25,7 @@
 #include <math/vec3.h>
 #include <math/vec4.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/compat/light_manager.h"
 #include "experimental/filament/compat/model_objects.h"
 #include "experimental/filament/compat/scene_geom_util.h"
 #include "experimental/filament/filament_util.h"
@@ -40,43 +39,6 @@ using filament::math::float4;
 using filament::math::mat3;
 using filament::math::mat4;
 
-static UniquePtr<mjrTexture> CreateFallbackIndirectLightTexture(
-    mjrfContext* ctx) {
-  const std::string filename = ResolveFilamentAssetPath("ibl.ktx");
-  mjResource* resource =
-      mju_openResource("", filename.c_str(), nullptr, nullptr, 0);
-  if (!resource) {
-    mju_error("Failed to open resource: %s", filename.c_str());
-  }
-  const void* bytes = nullptr;
-  const int nbytes = mju_readResource(resource, &bytes);
-  if (bytes == nullptr || nbytes <= 0) {
-    mju_error("Failed to read resource: %s", filename.c_str());
-  }
-
-  mjrTextureConfig config;
-  mjr_defaultTextureConfig(&config);
-  config.width = 1;
-  config.height = 1;
-  config.sampler_type = mjTEXTURE_CUBE;
-  config.format = mjPIXEL_FORMAT_KTX;
-  config.color_space = mjCOLORSPACE_AUTO;
-
-  auto texture = CreateTexture(ctx, config);
-
-  mjrTextureData payload;
-  mjr_defaultTextureData(&payload);
-  payload.bytes = bytes;
-  payload.nbytes = nbytes;
-  payload.release_callback = +[](void* user_data) {
-    mju_closeResource((mjResource*)user_data);
-  };
-  payload.user_data = resource;
-
-  mjrf_setTextureData(texture.get(), &payload);
-  return texture;
-}
-
 SceneBridge::SceneBridge(mjrfContext* ctx, const mjModel* model)
     : ctx_(ctx) {
   mjrSceneParams params;
@@ -92,32 +54,12 @@ SceneBridge::SceneBridge(mjrfContext* ctx, const mjModel* model)
                                  filament::math::float4(0, 0, 0, 1));
   mjrf_setClearColor(ctx_, &clear_color[0]);
 
-  default_shadow_map_size_ = ReadElement(
-      model, "filament.shadows.map_size", default_shadow_map_size_);
-  default_vsm_blur_width_ = ReadElement(
-      model, "filament.shadows.vsm_blur_width", default_vsm_blur_width_);
-  fallback_head_light_intensity_ =
-      ReadElement(model, "filament.fallback.head_light_intensity",
-                  fallback_head_light_intensity_);
-  fallback_scene_light_intensity_ =
-      ReadElement(model, "filament.fallback.scene_light_intensity",
-                  fallback_scene_light_intensity_);
-  fallback_environment_light_intensity_ =
-      ReadElement(model, "filament.fallback.environment_light_intensity",
-                  fallback_environment_light_intensity_);
-
-  PrepareLights();
+  light_manager_ =
+      std::make_unique<LightManager>(ctx_, scene_.get(), model_objects_.get());
 }
 
 SceneBridge::~SceneBridge() {
-  for (auto& iter : lights_) {
-    mjrf_removeLightFromScene(scene_.get(), iter.get());
-  }
-  lights_.clear();
-  if (fallback_ibl_) {
-    mjrf_removeLightFromScene(scene_.get(), fallback_ibl_.get());
-  }
-  fallback_ibl_.reset();
+  light_manager_.reset();
   for (auto& iter : renderables_) {
     mjrf_removeRenderableFromScene(scene_.get(), iter.get());
   }
@@ -132,107 +74,6 @@ std::optional<float3> SceneBridge::ClipFromWorld(const float3& pos) const{
   return clip_pos.xyz / clip_pos.w;
 }
 
-void SceneBridge::PrepareLights() {
-  const mjModel* model = model_objects_->GetModel();
-
-  bool has_image_based_light = false;
-  float total_light_intensity = 0.0f;
-  for (int i = 0; i < model->nlight; ++i) {
-    total_light_intensity += model->light_intensity[i];
-
-    if (model->light_type[i] == mjLIGHT_IMAGE) {
-      mjrLightParams params;
-      mjr_defaultLightParams(&params);
-      params.type = mjLIGHT_IMAGE;
-      params.texture = model_objects_->GetTexture(model->light_texid[i]);
-      params.intensity = model->light_intensity[i];
-      auto light_obj = CreateLight(ctx_, params);
-      mjrf_addLightToScene(scene_.get(), light_obj.get());
-      lights_.emplace_back(std::move(light_obj));
-      has_image_based_light = true;
-    } else {
-      mjrLightParams params;
-      mjr_defaultLightParams(&params);
-      params.color[0] = model->light_diffuse[0];
-      params.color[1] = model->light_diffuse[1];
-      params.color[2] = model->light_diffuse[2];
-      params.type = (mjtLightType)model->light_type[i];
-      params.cast_shadows = model->light_castshadow[i];
-      params.bulb_radius = model->light_bulbradius[i];
-      params.range = model->light_range[i];
-      params.intensity = model->light_intensity[i];
-      params.shadow_map_size = default_shadow_map_size_;
-      params.vsm_blur_width = default_vsm_blur_width_;
-      if (params.type == mjLIGHT_SPOT) {
-        params.spot_cone_angle = model->light_cutoff[i];
-      }
-
-      auto light_obj = CreateLight(ctx_, params);
-      mjrf_addLightToScene(scene_.get(), light_obj.get());
-      lights_.emplace_back(std::move(light_obj));
-    }
-  }
-
-  // Add a placeholder (black) headlight as our last light. Going forward, we'll
-  // assume lights_.back() is always the headlight.
-  {
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    // We break with the spec here slightly and use a spot light for the head
-    // light instead of a directional params. This is because filament only
-    // supports a single directional light, and we'd rather allow a scene
-    // light to be that directional params. It's also a bit odd for a
-    // directional light to move with the camera.
-    params.type = mjLIGHT_SPOT;
-    params.cast_shadows = 0;
-    params.intensity = 0.0f;
-    params.spot_cone_angle = 90.0f;
-    auto light_obj = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), light_obj.get());
-    lights_.emplace_back(std::move(light_obj));
-  }
-
-  if (!has_image_based_light && total_light_intensity > 0.0f) {
-    // Create a black indirect light to ensure that the skybox is
-    // oriented to respect mujoco's Z-up convention.
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    params.type = mjLIGHT_IMAGE;
-    params.intensity = 10.0f;
-    fallback_ibl_ = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), fallback_ibl_.get());
-  }
-
-  // There are no "physical" lights in the scene which means we're likely
-  // dealing with a "classic renderer" scene. In this case, let's add a
-  // default environment light and set the light intensity ourselves.
-  if (total_light_intensity == 0.0f) {
-    // Create a fallback environment light.
-    fallback_ibl_texture_ = CreateFallbackIndirectLightTexture(ctx_);
-
-    mjrLightParams params;
-    mjr_defaultLightParams(&params);
-    params.type = mjLIGHT_IMAGE;
-    params.texture = fallback_ibl_texture_.get();
-    params.intensity = fallback_environment_light_intensity_;
-    fallback_ibl_ = CreateLight(ctx_, params);
-    mjrf_addLightToScene(scene_.get(), fallback_ibl_.get());
-
-    // Distribute the fallback scene light intensity among the lights.
-    const float intensity = fallback_scene_light_intensity_ / lights_.size();
-    for (auto& light : lights_) {
-      if (light) {
-        const bool is_headlight = (light == lights_.back());
-        mjrf_setLightIntensity(light.get(),
-                               is_headlight ? fallback_head_light_intensity_
-                                            : intensity);
-      }
-    }
-  }
-
-  mjrf_setSceneSkybox(scene_.get(), model_objects_->GetSkyboxTexture());
-}
-
 mat4 CalculateClipFromWorld(const mjrRect& viewport, const mjrCamera& cam) {
   const float3 cam_pos(cam.pos[0], cam.pos[1], cam.pos[2]);
   const float3 cam_fwd(cam.forward[0], cam.forward[1], cam.forward[2]);
@@ -295,10 +136,16 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
     renderables_.push_back(std::move(renderable));
   }
 
+  mjrLight* headlight = nullptr;
   bool headlight_enabled = false;
   for (int i = 0; i < scene->nlight; ++i) {
     const mjvLight& scene_light = scene->lights[i];
     if (scene_light.id < 0 && scene_light.headlight) {
+      // The headlight, if it exists, is assigned the id `scene->nlight`.
+      headlight = light_manager_->GetLight(scene->nlight);
+      if (!headlight) {
+        continue;
+      }
       // We position the headlight slightly behind the camera to avoid some
       // odd clipping issues.
       headlight_enabled = true;
@@ -306,24 +153,21 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
       headpos[1] -= gazedir[1] * 0.05f;
       headpos[2] -= gazedir[2] * 0.05f;
 
-      // The headlight is always the "back" light.
-      UniquePtr<mjrLight>& light = lights_.back();
-      mjrf_setLightColor(light.get(), scene_light.diffuse);
-      mjrf_setLightTransform(light.get(), headpos, gazedir);
+      mjrf_setLightColor(headlight, scene_light.diffuse);
+      mjrf_setLightTransform(headlight, headpos, gazedir);
       continue;
-    } else if (scene_light.id < lights_.size() - 1) {
-      UniquePtr<mjrLight>& light = lights_[scene_light.id];
-      if (light) {
-        mjrf_setLightColor(light.get(), scene_light.diffuse);
-        mjrf_setLightTransform(light.get(), scene_light.pos, scene_light.dir);
-      }
+    } else if (mjrLight* light = light_manager_->GetLight(scene_light.id)) {
+      mjrf_setLightColor(light, scene_light.diffuse);
+      mjrf_setLightTransform(light, scene_light.pos, scene_light.dir);
     } else {
       mju_error("Unexpected light id: %d", scene_light.id);
     }
   }
 
   // Enable/disable the headlight based on whether or not it's in the scene.
-  mjrf_setLightEnabled(lights_.back().get(), headlight_enabled);
+  if (headlight) {
+    mjrf_setLightEnabled(headlight, headlight_enabled);
+  }
 }
 
 void SceneBridge::UploadMesh(const mjModel* model, int id) {
diff --git a/src/experimental/filament/compat/scene_bridge.h b/src/experimental/filament/compat/scene_bridge.h
index 348001d0..6323863b 100644
--- a/src/experimental/filament/compat/scene_bridge.h
+++ b/src/experimental/filament/compat/scene_bridge.h
@@ -18,13 +18,13 @@
 #include <functional>
 #include <memory>
 #include <optional>
-#include <string_view>
 #include <vector>
 
 #include <math/mat4.h>
 #include <math/vec3.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/compat/light_manager.h"
 #include "experimental/filament/compat/model_objects.h"
 #include "experimental/filament/render_context_filament.h"
 #include "experimental/filament/render_context_filament_cpp.h"
@@ -57,8 +57,6 @@ class SceneBridge {
   SceneBridge& operator=(const SceneBridge&) = delete;
 
  private:
-  void PrepareLights();
-
   // Converts a point in world space to clip space, eg. in the range [-1,-1, 0]
   // to [1, 1, 1]. Returns std::nullopt if the point is behind the camera.
   std::optional<filament::math::float3> ClipFromWorld(
@@ -66,19 +64,12 @@ class SceneBridge {
 
   mjrfContext* ctx_ = nullptr;
   std::unique_ptr<ModelObjects> model_objects_;
+  std::unique_ptr<LightManager> light_manager_;
   mjrCamera camera_;
   DrawTextAtFn draw_text_callback_;
   UniquePtr<mjrScene> scene_{nullptr, nullptr};
-  UniquePtr<mjrLight> fallback_ibl_{nullptr, nullptr};
-  UniquePtr<mjrTexture> fallback_ibl_texture_{nullptr, nullptr};
-  std::vector<UniquePtr<mjrLight>> lights_;
   std::vector<UniquePtr<mjrRenderable>> renderables_;
   filament::math::mat4 clip_from_world_;
-  int default_shadow_map_size_ = 2048;
-  float default_vsm_blur_width_ = 0.0f;
-  float fallback_head_light_intensity_ = 0.f;
-  float fallback_scene_light_intensity_ = 80'000.f;
-  float fallback_environment_light_intensity_ = 5'000.f;
 };
 
 }  // namespace mujoco

From 7174d33f0872d7a0d1df0932ca3fa35aaca301a1 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Wed, 20 May 2026 10:47:11 -0700
Subject: [PATCH 08/61] Introduce `mjPreContact`, a minimal struct passed into
 the collision functions.

PiperOrigin-RevId: 918533795
Change-Id: I2b5af05c1479b25d5c2cfdc690321a26fce6ede6
---
 doc/APIreference/APItypes.rst            |  11 +-
 doc/changelog.rst                        |   3 +
 doc/includes/references.h                |   9 +-
 include/mujoco/mjdata.h                  |  14 +-
 python/mujoco/introspect/structs.py      |  36 ++++
 python/mujoco/raw.h                      |   1 +
 python/mujoco/structs.cc                 |  26 +++
 python/mujoco/structs.h                  |  22 +++
 python/mujoco/structs_wrappers.cc        |  20 ++
 src/engine/engine_collision_box.c        | 136 +++++++-------
 src/engine/engine_collision_convex.c     | 224 ++++++++++++-----------
 src/engine/engine_collision_convex.h     |   8 +-
 src/engine/engine_collision_driver.c     |  81 +++++---
 src/engine/engine_collision_primitive.c  | 163 +++++++++++------
 src/engine/engine_collision_primitive.h  |  67 +++----
 src/engine/engine_collision_sdf.c        |  86 +++++----
 src/engine/engine_collision_sdf.h        |   6 +-
 src/engine/engine_support.c              |   6 +-
 src/engine/engine_util_spatial.c         |   2 +-
 test/engine/engine_collision_box_test.cc |  73 ++++----
 unity/Runtime/Bindings/MjBindings.cs     |   8 +
 wasm/codegen/generated/bindings.cc       | 108 ++++++++---
 wasm/codegen/generators/constants.py     |   1 +
 wasm/codegen/generators/structs.py       |   2 +-
 wasm/codegen/tests/generators_test.py    |   2 +-
 25 files changed, 709 insertions(+), 406 deletions(-)

diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst
index 6fe1c6c5..dd34f411 100644
--- a/doc/APIreference/APItypes.rst
+++ b/doc/APIreference/APItypes.rst
@@ -936,6 +936,15 @@ MJCF element :ref:`statistic <statistic>`. One instance of it is embedded in mjM
 .. mujoco-include:: mjStatistic
 
 
+.. _mjPreContact:
+
+mjPreContact
+~~~~~~~~~~~~
+
+This is the data structure holding information about one contact filled out by a narrowphase collision detector.
+
+.. mujoco-include:: mjPreContact
+
 .. _mjContact:
 
 mjContact
@@ -1780,7 +1789,7 @@ mjfCollision
 .. code-block:: C
 
    typedef int (*mjfCollision)(const mjModel* m, const mjData* d,
-                               mjContact* con, int g1, int g2, mjtNum margin);
+                               mjPreContact* con, int g1, int g2, mjtNum margin);
 
 This is the function type of the callbacks in the collision table :ref:`mjCOLLISIONFUNC`.
 
diff --git a/doc/changelog.rst b/doc/changelog.rst
index ea44b45a..29b713cc 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -70,6 +70,9 @@ General
      Negative ``margin`` values are now permitted (corresponding to ``gap > margin`` under the old semantics). The
      constraint ``margin + gap >= 0`` should be maintained to ensure valid collision detection.
 
+   - The :ref:`mjfCollision` functions now populate the :ref:`mjPreContact` struct instead of the :ref:`mjContact`
+     struct. The :ref:`mjPreContact` only contains the necessary fields needed for the narrowphase collision detection.
+
    - MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
      ``naconmax``.
 
diff --git a/doc/includes/references.h b/doc/includes/references.h
index cc3d6b49..7d929ee3 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -101,8 +101,15 @@ typedef enum mjtSleepState_ {       // sleep state of an object
   mjS_ASLEEP = 0,                   // object is asleep
   mjS_AWAKE  = 1                    // object is awake
 } mjtSleepState;
+struct mjPreContact_ {             // contact parameters set by narrowphase collision functions
+  mjtNum dist;
+  mjtNum pos[3];
+  mjtNum normal[3];                // contact normal of the collision
+  mjtNum tangent[3];               // first tangent direction
+};
+typedef struct mjPreContact_ mjPreContact;
 struct mjContact_ {                // result of collision detection functions
-  // contact parameters set by near-phase collision function
+  // contact parameters set by narrowphase collision function
   mjtNum  dist;                    // distance between nearest points; neg: penetration
   mjtNum  pos[3];                  // position of contact point: midpoint between geoms
   mjtNum  frame[9];                // normal is in [0-2], points from geom[0] to geom[1]
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
index 3709896a..8a22911c 100644
--- a/include/mujoco/mjdata.h
+++ b/include/mujoco/mjdata.h
@@ -121,10 +121,18 @@ typedef enum mjtSleepState_ {       // sleep state of an object
 } mjtSleepState;
 
 
-//---------------------------------- mjContact -----------------------------------------------------
+//------------------------------------- Contact ----------------------------------------------------
+
+struct mjPreContact_ {             // contact parameters set by narrowphase collision functions
+  mjtNum dist;
+  mjtNum pos[3];
+  mjtNum normal[3];                // contact normal of the collision
+  mjtNum tangent[3];               // first tangent direction
+};
+typedef struct mjPreContact_ mjPreContact;
 
 struct mjContact_ {                // result of collision detection functions
-  // contact parameters set by near-phase collision function
+  // contact parameters set by narrowphase collision function
   mjtNum  dist;                    // distance between nearest points; neg: penetration
   mjtNum  pos[3];                  // position of contact point: midpoint between geoms
   mjtNum  frame[9];                // normal is in [0-2], points from geom[0] to geom[1]
@@ -533,7 +541,7 @@ typedef mjtNum (*mjfTime)(void);
 typedef mjtNum (*mjfAct)(const mjModel* m, const mjData* d, int id);
 
 // collision detection
-typedef int (*mjfCollision)(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+typedef int (*mjfCollision)(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                             mjtNum margin);
 
 #endif  // MUJOCO_MJDATA_H_
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
index 0b924bf7..443b9f12 100644
--- a/python/mujoco/introspect/structs.py
+++ b/python/mujoco/introspect/structs.py
@@ -5140,6 +5140,42 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
+    ('mjPreContact',
+     StructDecl(
+         name='mjPreContact',
+         declname='struct mjPreContact_',
+         fields=(
+             StructFieldDecl(
+                 name='dist',
+                 type=ValueType(name='mjtNum'),
+                 doc='',
+             ),
+             StructFieldDecl(
+                 name='pos',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(3,),
+                 ),
+                 doc='',
+             ),
+             StructFieldDecl(
+                 name='normal',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(3,),
+                 ),
+                 doc='contact normal of the collision',
+             ),
+             StructFieldDecl(
+                 name='tangent',
+                 type=ArrayType(
+                     inner_type=ValueType(name='mjtNum'),
+                     extents=(3,),
+                 ),
+                 doc='first tangent direction',
+             ),
+         ),
+     )),
     ('mjContact',
      StructDecl(
          name='mjContact',
diff --git a/python/mujoco/raw.h b/python/mujoco/raw.h
index 7776e88c..542f0292 100644
--- a/python/mujoco/raw.h
+++ b/python/mujoco/raw.h
@@ -25,6 +25,7 @@
 // to them under the "raw" namespace.
 namespace mujoco::raw {
 
+using MjPreContact = ::mjPreContact;
 using MjContact = ::mjContact;
 using MjData = ::mjData;
 using MjLROpt = ::mjLROpt;
diff --git a/python/mujoco/structs.cc b/python/mujoco/structs.cc
index 79840742..a07f89b1 100644
--- a/python/mujoco/structs.cc
+++ b/python/mujoco/structs.cc
@@ -615,6 +615,32 @@ This is useful for example when the MJB is not available as a file on disk.)"));
   X(int, nupdate);
 #undef X
 
+  // ==================== MJPRECONTACT ==========================================
+  py::class_<MjPreContactWrapper> mjPreContact(m, "MjPreContact");
+  mjPreContact.def(py::init<>());
+  mjPreContact.def("__copy__", [](const MjPreContactWrapper& self) {
+    return MjPreContactWrapper(self);
+  });
+  mjPreContact.def("__deepcopy__", [](const MjPreContactWrapper& self, py::dict) {
+    return MjPreContactWrapper(self);
+  });
+  DefineStructFunctions(mjPreContact);
+
+#define X(var)                                                           \
+  mjPreContact.def_property(                                             \
+      #var, [](const MjPreContactWrapper& c) { return c.get()->var; },   \
+      [](MjPreContactWrapper& c, decltype(raw::MjPreContact::var) rhs) { \
+        c.get()->var = rhs;                                              \
+      })
+  X(dist);
+#undef X
+
+#define X(var) DefinePyArray(mjPreContact, #var, &MjPreContactWrapper::var)
+  X(pos);
+  X(normal);
+  X(tangent);
+#undef X
+
   // ==================== MJCONTACT ============================================
   py::class_<MjContactWrapper> mjContact(m, "MjContact");
   mjContact.def(py::init<>());
diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h
index 03080618..60bcf558 100644
--- a/python/mujoco/structs.h
+++ b/python/mujoco/structs.h
@@ -547,6 +547,27 @@ template <>
 struct enable_if_mj_struct<raw::MjModel> { using type = void; };
 
 // ==================== MJCONTACT ==============================================
+template <>
+class MjWrapper<raw::MjPreContact> : public WrapperBase<raw::MjPreContact> {
+ public:
+  MjWrapper();
+  MjWrapper(const MjWrapper&);
+  MjWrapper(MjWrapper&&) = default;
+  MjWrapper(raw::MjPreContact* ptr, pybind11::handle owner);
+  ~MjWrapper() = default;
+
+  #define X(var)                                                     \
+    py_array_or_tuple_t<                                             \
+        std::remove_all_extents_t<decltype(raw::MjPreContact::var)>> \
+        var
+  X(pos);
+  X(normal);
+  X(tangent);
+  #undef X
+};
+
+using MjPreContactWrapper = MjWrapper<raw::MjPreContact>;
+
 template <>
 class MjWrapper<raw::MjContact> : public WrapperBase<raw::MjContact> {
  public:
@@ -961,6 +982,7 @@ using _impl::MjSolverStatWrapper;
 using _impl::MjModelWrapper;
 using _impl::MjDataWrapper;
 using _impl::MjContactWrapper;
+using _impl::MjPreContactWrapper;
 using _impl::MjvPerturbWrapper;
 using _impl::MjvCameraWrapper;
 using _impl::MjvGLCameraWrapper;
diff --git a/python/mujoco/structs_wrappers.cc b/python/mujoco/structs_wrappers.cc
index 1a35c0dc..c5c8b6ca 100644
--- a/python/mujoco/structs_wrappers.cc
+++ b/python/mujoco/structs_wrappers.cc
@@ -535,6 +535,26 @@ std::unique_ptr<MjModelWrapper> MjModelWrapper::Deserialize(
   return std::unique_ptr<MjModelWrapper>(new MjModelWrapper(model));
 }
 
+// ==================== MJPRECONTACT ==========================================
+#define X(var) var(InitPyArray(ptr_->var, owner_))
+MjPreContactWrapper::MjWrapper()
+    : WrapperBase(new raw::MjPreContact{}),
+      X(pos),
+      X(normal),
+      X(tangent) {}
+
+MjPreContactWrapper::MjWrapper(raw::MjPreContact* ptr, py::handle owner)
+    : WrapperBase(ptr, owner),
+      X(pos),
+      X(normal),
+      X(tangent) {}
+#undef X
+
+MjPreContactWrapper::MjWrapper(const MjPreContactWrapper& other)
+    : MjPreContactWrapper() {
+  *this->ptr_ = *other.ptr_;
+}
+
 // ==================== MJCONTACT ==============================================
 #define X(var) var(InitPyArray(ptr_->var, owner_))
 MjContactWrapper::MjWrapper()
diff --git a/src/engine/engine_collision_box.c b/src/engine/engine_collision_box.c
index e30b7e16..6cf09568 100644
--- a/src/engine/engine_collision_box.c
+++ b/src/engine/engine_collision_box.c
@@ -20,24 +20,18 @@
 #include "engine/engine_util_misc.h"
 
 // hard-clamp vector to range [-limit(i), +limit(i)]
-static void mju_clampVec(mjtNum* vec, const mjtNum* limit, int n)
-{
-  int i;
-
-  // loop over active limits
-  for (i = 0; i < n; i++) {
+static void mju_clampVec(mjtNum* vec, const mjtNum* limit, int n) {
+  for (int i = 0; i < n; i++) {
+    // loop over active limits
     if (limit[i] > 0) {
-      if (vec[i] < -limit[i])
-        vec[i] = -limit[i];
-      else if (vec[i] > limit[i])
-        vec[i] = limit[i];
+      vec[i] = mju_clip(vec[i], -limit[i], limit[i]);
     }
   }
 }
 
 
 // raw sphere : box
-int mjraw_SphereBox(mjContact* con, mjtNum margin,
+int mjraw_SphereBox(mjPreContact* con, mjtNum margin,
                     const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
                     const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
   int i, k;
@@ -74,29 +68,32 @@ int mjraw_SphereBox(mjContact* con, mjtNum margin,
 
     mji_copy3(pos, center);
     mji_addToScl3(pos, nearest, (size1[0] - closest) / 2);
-    mji_mulMatVec3(con[0].frame, mat2, nearest);
+    mji_mulMatVec3(con[0].normal, mat2, nearest);
     dist = -closest;
   } else {
     mji_addToScl3(deepest, tmp, size1[0]);
     mju_zero3(pos);
     mji_addToScl3(pos, clamped, 0.5);
     mji_addToScl3(pos, deepest, 0.5);
-    mji_mulMatVec3(con[0].frame, mat2, tmp);
+    mji_mulMatVec3(con[0].normal, mat2, tmp);
   }
 
   mji_mulMatVec3(tmp, mat2, pos);
   mji_add3(con[0].pos, tmp, pos2);
   con[0].dist = dist - size1[0];
-  mju_zero3(con[0].frame + 3);
-
+  mji_zero3(con[0].tangent);
   return 1;
 }
 
 
 // sphere : box
-int mjc_SphereBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO;
-
+int mjc_SphereBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1  = d->geom_xpos + 3*g1;
+  const mjtNum* mat1  = d->geom_xmat + 9*g1;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* pos2  = d->geom_xpos + 3*g2;
+  const mjtNum* mat2  = d->geom_xmat + 9*g2;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_SphereBox(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
@@ -115,9 +112,10 @@ int mjc_SphereBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, m
 */
 
 // raw capsule : box
-int mjraw_CapsuleBox(mjContact* con, mjtNum margin,
+int mjraw_CapsuleBox(mjPreContact* con, mjtNum margin,
                      const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
-                     const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
+                     const mjtNum* pos2, const mjtNum* mat2,
+                     const mjtNum* size2) {
   mjtNum tmp1[3], tmp2[3], tmp3[3], halfaxis[3], axis[3], dif[3];
   mjtNum pos[3];          // position of capsule in box-local frame
 
@@ -148,7 +146,7 @@ int mjraw_CapsuleBox(mjContact* con, mjtNum margin,
   secondpos = -4;  // initialize to no 2nd contact (valid values are between -1 and 1)
 
   mji_sub3(tmp1, pos1, pos2);       // bring capsule to box-local frame (center's box is at (0,0,0))
-  mji_mulMatTVec3(pos, mat2, tmp1);  // and axis parralel to world
+  mji_mulMatTVec3(pos, mat2, tmp1);  // and axis parallel to world
 
   tmp1[0] = mat1[2];  // capsule's axis
   tmp1[1] = mat1[5];
@@ -591,21 +589,25 @@ skip:
 
 
 // capsule : box
-int mjc_CapsuleBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_CapsuleBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1  = d->geom_xpos + 3*g1;
+  const mjtNum* pos2  = d->geom_xpos + 3*g2;
+  const mjtNum* mat1  = d->geom_xmat + 9*g1;
+  const mjtNum* mat2  = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_CapsuleBox(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
 
 // internal box : box
 static inline
-int _boxbox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2, mjtNum margin)
-{
+int _boxbox(const mjModel* M, const mjData* D, mjPreContact* con, int g1, int g2, mjtNum margin) {
   const mjtNum* pos1 = D->geom_xpos + 3 * g1;
-  const mjtNum* mat1 = D->geom_xmat + 9 * g1;
-  const mjtNum* size1 = M->geom_size + 3 * g1;
   const mjtNum* pos2 = D->geom_xpos + 3 * g2;
+  const mjtNum* mat1 = D->geom_xmat + 9 * g1;
   const mjtNum* mat2 = D->geom_xmat + 9 * g2;
+  const mjtNum* size1 = M->geom_size + 3 * g1;
   const mjtNum* size2 = M->geom_size + 3 * g2;
 
   mjtNum pos12[3], pos21[3], rot[9], rott[9], rotabs[9], rottabs[9], tmp1[3], tmp2[3], plen1[3],
@@ -883,8 +885,10 @@ int _boxbox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2, m
           if (mju_abs(c2) > ss[1 - q])
             continue;
 
-          mji_copy3(points[n], lines[i]);
-          mji_addToScl3(points[n++], lines[i] + 3, c1);
+          if (n < mjMAXCONPAIR) {
+            mji_copy3(points[n], lines[i]);
+            mji_addToScl3(points[n++], lines[i] + 3, c1);
+          }
         }
       }
     }
@@ -911,10 +915,12 @@ int _boxbox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2, m
       if (u <= 0 || v <= 0 || u >= 1 || v >= 1)
         continue;
 
-      points[n][0] = llx;
-      points[n][1] = lly;
-      points[n][2] = (pts[0][2] + u * pts[1][2] + v * pts[2][2]);
-      n++;
+      if (n < mjMAXCONPAIR) {
+        points[n][0] = llx;
+        points[n][1] = lly;
+        points[n][2] = (pts[0][2] + u * pts[1][2] + v * pts[2][2]);
+        n++;
+      }
     }
   }
 
@@ -929,7 +935,9 @@ int _boxbox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2, m
       if (tmp1[1] <= -ly || tmp1[1] >= ly)
         continue;
 
-    mji_copy3(points[n++], tmp1);
+    if (n < mjMAXCONPAIR) {
+      mji_copy3(points[n++], tmp1);
+    }
   }
 
 
@@ -955,22 +963,23 @@ int _boxbox(const mjModel* M, const mjData* D, mjContact* con, int g1, int g2, m
   tmp2[1] = (q2 ? -1 : 1) * r[5];
   tmp2[2] = (q2 ? -1 : 1) * r[8];
 
-  mji_copy3(con[0].frame, tmp2);
-  mju_zero3(con[0].frame + 3);
+  mji_copy3(con[0].normal, tmp2);
+  mji_zero3(con[0].tangent);
 
 
 
 
-  for (i = 0; i < n; i++)
-  {
+  for (i = 0; i < n; i++) {
     con[i].dist = 2 * points[i][2];
     points[i][2] += hz;
 
     mji_mulMatVec3(tmp2, r, points[i]);
     mji_add3(con[i].pos, tmp2, p);
 
-    if (i)
-      mji_copy6(con[i].frame, con[0].frame);
+    if (i) {
+      mji_copy3(con[i].normal, con[0].normal);
+      mji_zero3(con[i].tangent);
+    }
   }
 
 
@@ -1315,8 +1324,8 @@ edgeedge:
 
   mji_mulMatVec3(tmp1, r, rnorm);
 
-  mji_scl3(con[0].frame, tmp1, in ? -1 : 1);
-  mju_zero3(con[0].frame + 3);
+  mji_scl3(con[0].normal, tmp1, in ? -1 : 1);
+  mji_zero3(con[0].tangent);
 
 
   for (i = 0; i < n; i++) {
@@ -1327,7 +1336,8 @@ edgeedge:
 
     mji_add3(con[i].pos, tmp2, pos1);
 
-    mji_copy6(con[i].frame, con[0].frame);
+    mji_copy3(con[i].normal, con[0].normal);
+    mji_zero3(con[i].tangent);
   }
 
   return n;
@@ -1338,13 +1348,13 @@ edgeedge:
 
 
 // box : box
-int mjc_BoxBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  int num = _boxbox(m, d, con, g1, g2, margin);
+int mjc_BoxBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  mjPreContact tmp[mjMAXCONPAIR];
+  int num = _boxbox(m, d, tmp, g1, g2, margin);
+
+  // -1: bad, 0: good
+  int dupe[mjMAXCONPAIR] = {0};
 
-  // use dim field to mark: -1: bad, 0: good
-  for (int i=0; i < num; i++) {
-    con[i].dim = 0;
-  }
 
   // get box info
   const mjtNum* pos1 =  d->geom_xpos + 3 * g1;
@@ -1364,28 +1374,28 @@ int mjc_BoxBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtN
     static mjtNum kRemoveRatio = 1.01;
 
     // is the contact outside: 1, inside: -1, within the removal width: 0
-    int out1 = mju_outsideBox(con[i].pos, pos1, mat1, sz1, kRemoveRatio);
-    int out2 = mju_outsideBox(con[i].pos, pos2, mat2, sz2, kRemoveRatio);
+    int out1 = mju_outsideBox(tmp[i].pos, pos1, mat1, sz1, kRemoveRatio);
+    int out2 = mju_outsideBox(tmp[i].pos, pos2, mat2, sz2, kRemoveRatio);
 
     // mark as bad if outside one box and not inside the other box
     if ((out1 == 1 && out2 != -1) || (out2 == 1 && out1 != -1)) {
-      con[i].dim = -1;
+      dupe[i] = -1;
     }
   }
 
   // find duplicates
   for (int i=0; i < num-1; i++) {
-    if (con[i].dim == -1) {
+    if (dupe[i] == -1) {
       continue;  // already marked bad: skip
     }
     for (int j=i+1; j < num; j++) {
-      if (con[j].dim == -1) {
+      if (dupe[j] == -1) {
         continue;  // already marked bad: skip
       }
-      if (con[i].pos[0] == con[j].pos[0] &&
-          con[i].pos[1] == con[j].pos[1] &&
-          con[i].pos[2] == con[j].pos[2]) {
-        con[i].dim = -1;
+      if (tmp[i].pos[0] == tmp[j].pos[0] &&
+          tmp[i].pos[1] == tmp[j].pos[1] &&
+          tmp[i].pos[2] == tmp[j].pos[2]) {
+        dupe[i] = -1;
         break;
       }
     }
@@ -1394,14 +1404,8 @@ int mjc_BoxBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtN
   // consolidate good
   int i = 0;
   for (int j=0; j < num; j++) {
-    // good: maybe copy
-    if (con[j].dim == 0) {
-      // different: copy
-      if (i < j) {
-        con[i] = con[j];
-      }
-
-      // advance either way
+    if (dupe[j] == 0) {
+      con[i] = tmp[j];
       i++;
     }
   }
diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c
index 46f4cf09..2c52a789 100644
--- a/src/engine/engine_collision_convex.c
+++ b/src/engine/engine_collision_convex.c
@@ -41,7 +41,7 @@ static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
 
 // wrapper around libccd; returns number of collisions found
 static inline int _libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* obj2,
-                                  mjContact* con, mjtNum margin) {
+                                  mjPreContact* con, mjtNum margin) {
   ccd_t ccd;
   CCD_INIT(&ccd);
   ccd.mpr_tolerance = m->opt.ccd_tolerance;
@@ -64,10 +64,10 @@ static inline int _libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* ob
     if (ccdVec3Eq(&ccd_dir, ccd_vec3_origin)) {
       return 0;
     }
-    con->dist = margin - ccd_depth;
-    mji_copy3(con->frame, ccd_dir.v);
-    mji_copy3(con->pos, ccd_pos.v);
-    mji_zero3(con->frame + 3);
+    con[0].dist = margin - ccd_depth;
+    mji_copy3(con[0].normal, ccd_dir.v);
+    mji_copy3(con[0].pos, ccd_pos.v);
+    mji_zero3(con[0].tangent);
     return 1;
   }
   return 0;
@@ -76,7 +76,7 @@ static inline int _libccd_wrapper(const mjModel* m, mjCCDObj* obj1, mjCCDObj* ob
 
 // find penetration info between two geoms; returns number of collisions found
 static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj* obj2,
-                           mjContact* con, int ncon, mjtNum margin) {
+                           mjPreContact* con, int ncon, mjtNum margin) {
   if (mjDISABLED(mjDSBL_NATIVECCD)) {
     return _libccd_wrapper(m, obj1, obj2, con, margin);
   }
@@ -97,14 +97,14 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
   if ((dist = mjc_ccd(&config, &status, obj1, obj2)) < 0) {
     mj_freeStack(d);
     int nwitness = status.nx;
-    for (int i = 0; i < nwitness; i++, con++) {
-      con->dist = margin + dist;
-      con->pos[0] = 0.5*(status.x1[3*i + 0] + status.x2[3*i + 0]);
-      con->pos[1] = 0.5*(status.x1[3*i + 1] + status.x2[3*i + 1]);
-      con->pos[2] = 0.5*(status.x1[3*i + 2] + status.x2[3*i + 2]);
-      mji_sub3(con->frame, status.x1 + 3*i, status.x2 + 3*i);
-      mju_normalize3(con->frame);
-      mji_zero3(con->frame + 3);
+    for (int i = 0; i < nwitness; i++) {
+      con[i].dist = margin + dist;
+      con[i].pos[0] = 0.5*(status.x1[3*i + 0] + status.x2[3*i + 0]);
+      con[i].pos[1] = 0.5*(status.x1[3*i + 1] + status.x2[3*i + 1]);
+      con[i].pos[2] = 0.5*(status.x1[3*i + 2] + status.x2[3*i + 2]);
+      mji_sub3(con[i].normal, status.x1 + 3*i, status.x2 + 3*i);
+      mju_normalize3(con[i].normal);
+      mji_zero3(con[i].tangent);
     }
     return nwitness;
   }
@@ -780,9 +780,10 @@ static void mjc_setCCDObjFlex(mjCCDObj* obj, int flex, int elem, int vert) {
 
 
 // compare new contact to previous contacts, return 1 if it is far from all of them
-static int mjc_isDistinctContact(mjContact* con, int ncon, mjtNum tolerance) {
+static int mjc_isDistinctContact(const mjPreContact* con, int ncon, mjtNum tolerance) {
+  const mjtNum* last_pos = con[ncon - 1].pos;
   for (int i=0; i < ncon-1; i++) {
-    if (mju_dist3(con[i].pos, con[ncon - 1].pos) <= tolerance) {
+    if (mju_dist3(con[i].pos, last_pos) <= tolerance) {
       return 0;
     }
   }
@@ -838,7 +839,7 @@ static int maxContacts(const mjModel* m, const mjCCDObj* obj1, const mjCCDObj* o
 
 
 // multi-point convex-convex collision, using libccd
-int mjc_Convex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
+int mjc_Convex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
   // init ccd objects
   mjCCDObj obj1, obj2;
   mjc_initCCDObj(&obj1, m, d, g1, margin);
@@ -857,7 +858,7 @@ int mjc_Convex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtN
   }
 
   // look for additional contacts
-  if (ncon == 1 && !mjDISABLED(mjDSBL_MULTICCD)  // TODO(tassa) leave as bitflag or make geom attribute (?)
+  if (ncon == 1 && !mjDISABLED(mjDSBL_MULTICCD)
       && m->geom_type[g1] != mjGEOM_ELLIPSOID && m->geom_type[g1] != mjGEOM_SPHERE
       && m->geom_type[g2] != mjGEOM_ELLIPSOID && m->geom_type[g2] != mjGEOM_SPHERE) {
     // multiCCD parameters
@@ -866,7 +867,8 @@ int mjc_Convex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtN
 
     // complete frame of initial contact
     mjtNum frame[9];
-    mji_copy9(frame, con[0].frame);
+    mji_copy3(frame, con[0].normal);
+    mju_zero(frame+3, 6);
     mju_makeFrame(frame);
 
     // tolerance for determining if newly found contacts are distinct
@@ -926,7 +928,7 @@ const int maxplanemesh = 3;
 const mjtNum tolplanemesh = 0.3;
 
 // add one plane-mesh contact
-static int addplanemesh(mjContact* con, const float vertex[3],
+static int addplanemesh(mjPreContact* con, const float vertex[3],
                         const mjtNum pos1[3], const mjtNum normal1[3],
                         const mjtNum pos2[3], const mjtNum mat2[9],
                         const mjtNum first[3], mjtNum rbound) {
@@ -945,44 +947,47 @@ static int addplanemesh(mjContact* con, const float vertex[3],
   mji_sub3(dif, pnt, pos1);
 
   // set distance
-  con->dist = mju_dot3(normal1, dif);
+  con[0].dist = mju_dot3(normal1, dif);
 
   // set position
-  mji_copy3(con->pos, pnt);
-  mji_addToScl3(con->pos, normal1, -0.5*con->dist);
+  mji_copy3(con[0].pos, pnt);
+  mji_addToScl3(con[0].pos, normal1, -0.5*con[0].dist);
 
   // set frame
-  mji_copy3(con->frame, normal1);
-  mju_zero3(con->frame+3);
+  mji_copy3(con[0].normal, normal1);
+  mji_zero3(con[0].tangent);
 
   return 1;
 }
 
 
 // plane-convex collision, using libccd
-int mjc_PlaneConvex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
-  mjtNum dist, dif[3], normal[3] = {mat1[2], mat1[5], mat1[8]};
-  ccd_vec3_t dir, vec;
+int mjc_PlaneConvex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1  = d->geom_xpos + 3*g1;
+  const mjtNum* mat1  = d->geom_xmat + 9*g1;
+  const mjtNum* pos2  = d->geom_xpos + 3*g2;
+  const mjtNum* mat2  = d->geom_xmat + 9*g2;
+
+  mjtNum dif[3], normal[3] = {mat1[2], mat1[5], mat1[8]};
+  ccd_vec3_t ccd_dir, ccd_vec;
   mjCCDObj obj;
   mjc_initCCDObj(&obj, m, d, g2, 0);
   // get support point in -normal direction
-  ccdVec3Set(&dir, -mat1[2], -mat1[5], -mat1[8]);
-  mjccd_support(&obj, &dir, &vec);
+  ccdVec3Set(&ccd_dir, -mat1[2], -mat1[5], -mat1[8]);
+  mjccd_support(&obj, &ccd_dir, &ccd_vec);
 
   // compute normal distance, return if too far
-  mji_sub3(dif, vec.v, pos1);
-  dist = mju_dot3(normal, dif);
-  if (dist > margin) {
+  mji_sub3(dif, ccd_vec.v, pos1);
+  con[0].dist = mju_dot3(normal, dif);
+  if (con[0].dist > margin) {
     return 0;
   }
 
   // fill in contact data
-  con->dist = dist;
-  mji_copy3(con->pos, vec.v);
-  mji_addToScl3(con->pos, normal, -0.5*dist);
-  mji_copy3(con->frame, normal);
-  mju_zero3(con->frame+3);
+  mji_copy3(con[0].pos, ccd_vec.v);
+  mji_addToScl3(con[0].pos, normal, -0.5*con[0].dist);
+  mji_copy3(con[0].normal, normal);
+  mji_zero3(con[0].tangent);
 
   //--------------- add all/connected vertices below margin
   float* vertdata;
@@ -1001,7 +1006,7 @@ int mjc_PlaneConvex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
 
   // express dir in geom local frame
   mjtNum locdir[3];
-  mju_mulMatTVec3(locdir, d->geom_xmat+9*g, dir.v);
+  mju_mulMatTVec3(locdir, d->geom_xmat+9*g, ccd_dir.v);
 
   // inclusion threshold along locdir, relative to geom2 center
   mji_sub3(dif, pos2, pos1);
@@ -1097,7 +1102,8 @@ static inline void addPrismVert(mjCCDObj* obj, int r, int c, int i, mjtNum dx, m
 
 
 // entry point for heightfield collisions
-int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
+int mjc_ConvexHField(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                     mjtNum margin) {
   // hfield frame
   const mjtNum* pos1 = d->geom_xpos + 3*g1;
   const mjtNum* mat1 = d->geom_xmat + 9*g1;
@@ -1116,16 +1122,16 @@ int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2
   // try early return using box-sphere test
 
   // express geom2 pos in hfield frame
-  mjtNum pos[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]};
-  mju_mulMatTVec3(pos, mat1, pos);
+  mjtNum local_pos[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]};
+  mju_mulMatTVec3(local_pos, mat1, local_pos);
 
   // sphere radius is geom2 rbound + margin
   mjtNum radius = m->geom_rbound[g2] + margin;
 
   // box-sphere test
-  if ((size0 < pos[0] - radius) || (-size0 > pos[0] + radius) ||
-      (size1 < pos[1] - radius) || (-size1 > pos[1] + radius) ||
-      (size2 < pos[2] - radius) || (-size3 > pos[2] + radius)) {
+  if ((size0 < local_pos[0] - radius) || (-size0 > local_pos[0] + radius) ||
+      (size1 < local_pos[1] - radius) || (-size1 > local_pos[1] + radius) ||
+      (size2 < local_pos[2] - radius) || (-size3 > local_pos[2] + radius)) {
     return 0;
   }
 
@@ -1142,42 +1148,42 @@ int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2
   mji_mulMatTMat3(mat, mat1, mat2);
 
   mji_copy9(obj2.mat, mat);
-  mji_copy3(obj2.pos, pos);
+  mji_copy3(obj2.pos, local_pos);
 
-  mjtNum dir[3] = {0, 0, 0}, res[3];
+  mjtNum local_dir[3] = {0, 0, 0}, res[3];
 
   // get support point in +X
-  dir[0] = 1;
-  obj2.support(res, &obj2, dir);
+  local_dir[0] = 1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum xmax = res[0];
 
   // get support point in -X
-  dir[0] = -1;
-  obj2.support(res, &obj2, dir);
+  local_dir[0] = -1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum xmin = res[0];
 
-  dir[0] = 0;
+  local_dir[0] = 0;
 
   // get support point in +Y
-  dir[1] = 1;
-  obj2.support(res, &obj2, dir);
+  local_dir[1] = 1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum ymax = res[1];
 
   // get support point in -Y
-  dir[1] = -1;
-  obj2.support(res, &obj2, dir);
+  local_dir[1] = -1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum ymin = res[1];
 
-  dir[1] = 0;
+  local_dir[1] = 0;
 
   // get support point in +Z
-  dir[2] = 1;
-  obj2.support(res, &obj2, dir);
+  local_dir[2] = 1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum zmax = res[2];
 
   // get support point in -Z
-  dir[2] = -1;
-  obj2.support(res, &obj2, dir);
+  local_dir[2] = -1;
+  obj2.support(res, &obj2, local_dir);
   mjtNum zmin = res[2];
 
   // AABB box-box test
@@ -1226,10 +1232,10 @@ int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2
         // run penetration function, save contact
         if (mjc_penetration(m, d, &obj1, &obj2, con + ncon, 1, 0.0)) {
           // transform to global coordinates
-          mji_copy3(dir, con[ncon].frame);
-          mji_copy3(pos, con[ncon].pos);
-          mji_mulMatVec3(con[ncon].frame, mat1, dir);
-          mji_mulMatVec3(con[ncon].pos, mat1, pos);
+          mji_copy3(local_dir, con[ncon].normal);
+          mji_copy3(local_pos, con[ncon].pos);
+          mji_mulMatVec3(con[ncon].normal, mat1, local_dir);
+          mji_mulMatVec3(con[ncon].pos, mat1, local_pos);
           mji_addTo3(con[ncon].pos, pos1);
 
           // force out of all loops if max contacts reached
@@ -1247,7 +1253,7 @@ int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2
   if (mjDISABLED(mjDSBL_NATIVECCD)) {
     // fix contact normals
     for (int i=0; i < ncon; i++) {
-      mjc_fixNormal(m, d, con+i, g1, g2);
+      mjc_fixNormal(m, d, con + i, g1, g2);
     }
   }
 
@@ -1362,7 +1368,7 @@ static int mjc_ellipsoidOutside(mjtNum nrm[3], const mjtNum pos[3], const mjtNum
 
 
 // fix normals if required
-void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, int g2) {
+void mjc_fixNormal(const mjModel* m, const mjData* d, mjPreContact* con, int g1, int g2) {
   mjtNum dst1, dst2;
 
   // get geom ids and types
@@ -1389,8 +1395,8 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
 
   // init normals
   mjtNum normal[2][3] = {
-    {con->frame[0], con->frame[1], con->frame[2]},
-    {-con->frame[0], -con->frame[1], -con->frame[2]}
+    {con->normal[0], con->normal[1], con->normal[2]},
+    {-con->normal[0], -con->normal[1], -con->normal[2]}
   };
 
   // process geoms in type range
@@ -1402,27 +1408,27 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
       mjtNum* size = m->geom_size + 3*gid[i];
 
       // map contact point and normal to local frame
-      mjtNum dif[3], pos[3], nrm[3];
+      mjtNum dif[3], pos1[3], nrm[3];
       mju_sub3(dif, con->pos, d->geom_xpos+3*gid[i]);
-      mju_mulMatTVec3(pos, mat, dif);
+      mju_mulMatTVec3(pos1, mat, dif);
       mju_mulMatTVec3(nrm, mat, normal[i]);
 
       // process according to type
       switch (type[i]) {
       case mjGEOM_SPHERE:
-        mji_copy3(nrm, pos);
+        mji_copy3(nrm, pos1);
         processed[i] = 1;
         break;
 
       case mjGEOM_CAPSULE:
         // Z: bottom cap
-        if (pos[2] < -size[1]) {
-          nrm[2] = pos[2]+size[1];
+        if (pos1[2] < -size[1]) {
+          nrm[2] = pos1[2]+size[1];
         }
 
         // Z: top cap
-        else if (pos[2] > size[1]) {
-          nrm[2] = pos[2]-size[1];
+        else if (pos1[2] > size[1]) {
+          nrm[2] = pos1[2]-size[1];
         }
 
         // Z: cylinder
@@ -1431,8 +1437,8 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
         }
 
         // copy XY
-        nrm[0] = pos[0];
-        nrm[1] = pos[1];
+        nrm[0] = pos1[0];
+        nrm[1] = pos1[1];
         processed[i] = 1;
         break;
 
@@ -1443,27 +1449,27 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
         }
 
         // compute elliptic distance^2
-        dst1 = pos[0]*pos[0]/(size[0]*size[0]) +
-               pos[1]*pos[1]/(size[1]*size[1]) +
-               pos[2]*pos[2]/(size[2]*size[2]);
+        dst1 = pos1[0]*pos1[0]/(size[0]*size[0]) +
+               pos1[1]*pos1[1]/(size[1]*size[1]) +
+               pos1[2]*pos1[2]/(size[2]*size[2]);
 
         // dispatch to inside or outside solver
         if (dst1 <= 1) {
-          processed[i] = mjc_ellipsoidInside(nrm, pos, size);
+          processed[i] = mjc_ellipsoidInside(nrm, pos1, size);
         } else {
-          processed[i] = mjc_ellipsoidOutside(nrm, pos, size);
+          processed[i] = mjc_ellipsoidOutside(nrm, pos1, size);
         }
         break;
 
       case mjGEOM_CYLINDER:
         // skip if within 5% length of flat wall
-        if (mju_abs(pos[2]) > 0.95*size[1]) {
+        if (mju_abs(pos1[2]) > 0.95*size[1]) {
           break;
         }
 
         // compute distances to flat and round wall
-        dst1 = mju_abs(size[1]-mju_abs(pos[2]));
-        dst2 = mju_abs(size[0]-mju_norm(pos, 2));
+        dst1 = mju_abs(size[1]-mju_abs(pos1[2]));
+        dst2 = mju_abs(size[0]-mju_norm(pos1, 2));
 
         // require 4x closer to round than flat wall
         if (dst1 < 0.25*dst2) {
@@ -1471,8 +1477,8 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
         }
 
         // set normal for round wall
-        nrm[0] = pos[0];
-        nrm[1] = pos[1];
+        nrm[0] = pos1[0];
+        nrm[1] = pos1[1];
         nrm[2] = 0;
         processed[i] = 1;
         break;
@@ -1492,23 +1498,18 @@ void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, in
 
   // both processed: average
   if (processed[0] && processed[1]) {
-    mji_sub3(con->frame, normal[0], normal[1]);
-    mju_normalize3(con->frame);
+    mji_sub3(con->normal, normal[0], normal[1]);
+    mju_normalize3(con->normal);
   }
 
   // first processed: copy
   else if (processed[0]) {
-    mji_copy3(con->frame, normal[0]);
+    mji_copy3(con->normal, normal[0]);
   }
 
   // second processed: copy reverse
   else if (processed[1]) {
-    mji_scl3(con->frame, normal[1], -1);
-  }
-
-  // clear second frame axis if processed, just in case
-  if (processed[0] || processed[1]) {
-    mju_zero3(con->frame+3);
+    mji_scl3(con->normal, normal[1], -1);
   }
 }
 
@@ -1525,7 +1526,8 @@ int mjc_ConvexElem(const mjModel* m, mjData* d, mjContact* con, int g1, int f1,
   mjc_setCCDObjFlex(&obj2, f2, e2, -1);
 
   // find contacts
-  int ncon = mjc_penetration(m, d, &obj1, &obj2, con, 1, margin);
+  mjPreContact precon[8];
+  int ncon = mjc_penetration(m, d, &obj1, &obj2, precon, 1, margin);
 
   // fix normals for 2D flex
   if (ncon && !mjDISABLED(mjDSBL_NATIVECCD)) {
@@ -1536,11 +1538,19 @@ int mjc_ConvexElem(const mjModel* m, mjData* d, mjContact* con, int g1, int f1,
 
     if (isflex2d) {
       for (int i = 0; i < ncon; i++) {
-        mjc_fixNormal(m, d, con + i, g1, -1);
+        mjc_fixNormal(m, d, precon + i, g1, -1);
       }
     }
   }
 
+  // copy to con
+  for (int i = 0; i < ncon; i++) {
+    con[i].dist = precon[i].dist;
+    mji_copy3(con[i].pos, precon[i].pos);
+    mji_copy3(con[i].frame, precon[i].normal);
+    mju_zero3(con[i].frame + 3);
+  }
+
   return ncon;
 }
 
@@ -1659,14 +1669,14 @@ int mjc_HFieldElem(const mjModel* m, mjData* d, mjContact* con, int g, int f, in
           }
 
           // run ccd, save contact
-          if (mjc_penetration(m, d, &obj1, &obj2, con + cnt, 1, 0.0)) {
+          mjContact* ccon = con + cnt;
+          mjPreContact precon;
+          if (mjc_penetration(m, d, &obj1, &obj2, &precon, 1, 0.0)) {
             // transform to global coordinates
-            mjtNum dir[3], pos[3];
-            mji_copy3(dir, con[cnt].frame);
-            mji_copy3(pos, con[cnt].pos);
-            mji_mulMatVec3(con[cnt].frame, hmat, dir);
-            mji_mulMatVec3(con[cnt].pos, hmat, pos);
-            mji_addTo3(con[cnt].pos, hpos);
+            mji_zero3(ccon->frame + 3);
+            mji_mulMatVec3(ccon->frame, hmat, precon.normal);
+            mji_mulMatVec3(ccon->pos, hmat, precon.pos);
+            mji_addTo3(ccon->pos, hpos);
 
             // count, stop if max number reached
             cnt++;
diff --git a/src/engine/engine_collision_convex.h b/src/engine/engine_collision_convex.h
index d9212fbd..88fda2ea 100644
--- a/src/engine/engine_collision_convex.h
+++ b/src/engine/engine_collision_convex.h
@@ -109,9 +109,9 @@ void mjc_pointSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]);
 void mjc_lineSupport(mjtNum res[3], mjCCDObj* obj, const mjtNum dir[3]);
 
 // pairwise geom collision functions using ccd
-int mjc_PlaneConvex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
-int mjc_ConvexHField(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
-MJAPI int mjc_Convex(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_PlaneConvex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
+int mjc_ConvexHField(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
+MJAPI int mjc_Convex(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
 
 // geom-elem or elem-elem or vert-elem collision function using ccd
 int mjc_ConvexElem(const mjModel* m, mjData* d, mjContact* con, int g1, int f1, int e1, int v1,
@@ -122,7 +122,7 @@ int mjc_HFieldElem(const mjModel* m, mjData* d, mjContact* con, int g, int f, in
                    mjtNum margin);
 
 // fix contact frame normal
-void mjc_fixNormal(const mjModel* m, const mjData* d, mjContact* con, int g1, int g2);
+void mjc_fixNormal(const mjModel* m, const mjData* d, mjPreContact* con, int g1, int g2);
 
 #ifdef __cplusplus
 }
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index 5c31b16c..a598602f 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -1745,20 +1745,11 @@ void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2) {
     return;
   }
 
-  // allocate mjContact[mjMAXCONPAIR] on the arena
-  mjContact* con =
-    (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * mjMAXCONPAIR, _Alignof(mjContact));
-  if (!con) {
-    mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
-    return;
-  }
+
 
   // call collision detector to generate contacts
-  num = collisionFunc(m, d, con, g1, g2, margin + gap);
-
-  // check contacts
-  if (!num) {
-    resetArena(d);
+  mjPreContact precon[mjMAXCONPAIR];
+  if (!(num = collisionFunc(m, d, precon, g1, g2, margin + gap))) {
     return;
   }
 
@@ -1767,6 +1758,12 @@ void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2) {
     mjERROR("too many contacts returned by collision function");
   }
 
+  mjContact* con = (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * num, _Alignof(mjContact));
+  if (!con) {
+    mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
+    return;
+  }
+
   // set condim, solref, solimp, friction: dynamic
   if (ipair < 0) {
     mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
@@ -1787,7 +1784,12 @@ void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2) {
 
   // add contacts returned by collision detector
   for (int i=0; i < num; i++) {
-    // set contact ids
+    // set contact parameters
+    con[i].dist = precon[i].dist;
+    mji_copy3(con[i].pos, precon[i].pos);
+    mji_copy3(con[i].frame + 0, precon[i].normal);
+    mji_copy3(con[i].frame + 3, precon[i].tangent);
+
     con[i].geom[0] = g1;
     con[i].geom[1] = g2;
     con[i].flex[0] = -1;
@@ -2109,23 +2111,32 @@ void mj_collideGeomElem(const mjModel* m, mjData* d, int g, int f, int e) {
                    pos, mat, size);
 
     // call raw primitive for corresponding geom type
-    if (type == mjGEOM_SPHERE) {
-      num = mjraw_SphereCapsule(con, margin + gap,
-                                d->geom_xpos+3*g, d->geom_xmat+9*g, m->geom_size+3*g,
-                                pos, mat, size);
+    mjPreContact precon[2];
+    switch (type) {
+      case mjGEOM_SPHERE:
+        num = mjraw_SphereCapsule(precon, margin + gap, d->geom_xpos+3*g, d->geom_xmat+9*g,
+                                  m->geom_size+3*g, pos, mat, size);
+        break;
+      case mjGEOM_CAPSULE:
+        num = mjraw_CapsuleCapsule(precon, margin + gap, d->geom_xpos+3*g, d->geom_xmat+9*g,
+                                   m->geom_size+3*g, pos, mat, size);
+        break;
+      case mjGEOM_BOX:
+        num = mjraw_CapsuleBox(precon, margin + gap, pos, mat, size, d->geom_xpos+3*g,
+                               d->geom_xmat+9*g, m->geom_size+3*g);
+        break;
+      default:
+        num = 0;
     }
-    else if (type == mjGEOM_CAPSULE) {
-      num = mjraw_CapsuleCapsule(con, margin + gap,
-                                 d->geom_xpos+3*g, d->geom_xmat+9*g, m->geom_size+3*g,
-                                 pos, mat, size);
-    }
-    else {
-      num = mjraw_CapsuleBox(con, margin + gap,
-                             pos, mat, size,
-                             d->geom_xpos+3*g, d->geom_xmat+9*g, m->geom_size+3*g);
+
+    for (int i=0; i < num; i++) {
+      con[i].dist = precon[i].dist;
+      mju_copy3(con[i].pos, precon[i].pos);
+      mju_copy3(con[i].frame + 0, precon[i].normal);
+      mju_copy3(con[i].frame + 3, precon[i].tangent);
 
       // reverse contact normals, since box geom is second
-      for (int i=0; i < num; i++) {
+      if (type == mjGEOM_BOX) {
         mju_scl3(con[i].frame, con[i].frame, -1);
       }
     }
@@ -2256,7 +2267,14 @@ void mj_collideElems(const mjModel* m, mjData* d, int f1, int e1, int f2, int e2
                    pos2, mat2, size2);
 
     // raw primitive
-    num = mjraw_CapsuleCapsule(con, margin + gap, pos1, mat1, size1, pos2, mat2, size2);
+    mjPreContact precon[mjMAXCONPAIR];
+    num = mjraw_CapsuleCapsule(precon, margin + gap, pos1, mat1, size1, pos2, mat2, size2);
+    for (int i=0; i < num; i++) {
+      con[i].dist = precon[i].dist;
+      mju_copy3(con[i].pos, precon[i].pos);
+      mju_copy3(con[i].frame + 0, precon[i].normal);
+      mju_copy3(con[i].frame + 3, precon[i].tangent);
+    }
   }
 
   // general convex collision
@@ -2332,7 +2350,12 @@ void mj_collideElemVert(const mjModel* m, mjData* d, int f, int e, int v) {
     mjtNum pos[3], mat[9], size[2];
     mjtNum I[9] = {1, 0, 0, 0, 1, 0, 0, 0, 1};
     mj_makeCapsule(m, d, f, edata, pos, mat, size);
-    num = mjraw_SphereCapsule(con, 0, vert, I, &radius, pos, mat, size);
+    mjPreContact precon;
+    num = mjraw_SphereCapsule(&precon, 0, vert, I, &radius, pos, mat, size);
+    con->dist = precon.dist;
+    mju_copy3(con->pos, precon.pos);
+    mju_copy3(con->frame + 0, precon.normal);
+    mju_copy3(con->frame + 3, precon.tangent);
   }
 
   // sphere : triangle
diff --git a/src/engine/engine_collision_primitive.c b/src/engine/engine_collision_primitive.c
index 145dfc50..a8823008 100644
--- a/src/engine/engine_collision_primitive.c
+++ b/src/engine/engine_collision_primitive.c
@@ -25,70 +25,87 @@
 //--------------------------- plane collisions -----------------------------------------------------
 
 // raw plane : sphere
-static int mjraw_PlaneSphere(mjContact* con, mjtNum margin,
+static int mjraw_PlaneSphere(mjPreContact* con, mjtNum margin,
                              const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
                              const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
   // set normal
-  con[0].frame[0] = mat1[2];
-  con[0].frame[1] = mat1[5];
-  con[0].frame[2] = mat1[8];
+  con[0].normal[0] = mat1[2];
+  con[0].normal[1] = mat1[5];
+  con[0].normal[2] = mat1[8];
 
   // compute distance, return if too large
   mjtNum tmp[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]};
-  mjtNum cdist = mju_dot3(tmp, con[0].frame);
+  mjtNum cdist = mju_dot3(tmp, con[0].normal);
   if (cdist > margin + size2[0]) {
     return 0;
   }
 
   // depth and position
   con[0].dist = cdist - size2[0];
-  mji_scl3(tmp, con[0].frame, -con[0].dist / 2 - size2[0]);
+  mji_scl3(tmp, con[0].normal, -con[0].dist / 2 - size2[0]);
   mji_add3(con[0].pos, pos2, tmp);
-
-  mju_zero3(con[0].frame+3);
+  mji_zero3(con[0].tangent);
   return 1;
 }
 
 
 // plane : sphere
-int mjc_PlaneSphere(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_PlaneSphere(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                    mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_PlaneSphere(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
 
 // plane : capsule
-int mjc_PlaneCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_PlaneCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                     mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
 
   // get capsule axis, segment = scaled axis
   mjtNum axis[3] = {mat2[2], mat2[5], mat2[8]};
   mjtNum segment[3] = {size2[1]*axis[0], size2[1]*axis[1], size2[1]*axis[2]};
 
   // get point 1, do sphere-plane test
-  mjtNum pos[3];
-  mju_add3(pos, pos2, segment);
-  int n1 = mjraw_PlaneSphere(con, margin, pos1, mat1, size1, pos, mat2, size2);
+  mjtNum endpoint[3];
+  mju_add3(endpoint, pos2, segment);
+  int n1 = mjraw_PlaneSphere(con, margin, pos1, mat1, size1, endpoint, mat2, size2);
 
   // get point 2, do sphere-plane test
-  mju_sub3(pos, pos2, segment);
-  int n2 = mjraw_PlaneSphere(con+n1, margin, pos1, mat1, size1, pos, mat2, size2);
+  mju_sub3(endpoint, pos2, segment);
+  int n2 = mjraw_PlaneSphere(con + n1, margin, pos1, mat1, size1, endpoint, mat2, size2);
 
   // align contact frames with capsule axis
   if (n1) {
-    mji_copy3(con->frame + 3, axis);
+    mji_copy3(con[0].tangent, axis);
   }
   if (n2) {
-    mji_copy3((con + n1)->frame + 3, axis);
+    mji_copy3(con[n1].tangent, axis);
   }
 
-  return n1+n2;
+  return n1 + n2;
 }
 
 
 // plane : cylinder
-int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                      mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size2 = m->geom_size + 3*g2;
+
   mjtNum normal[3] = {mat1[2], mat1[5], mat1[8]};
   mjtNum axis[3] = {mat2[2], mat2[5], mat2[8]};
 
@@ -137,8 +154,8 @@ int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g
     mji_add3(con[cnt].pos, pos2, vec);
     mji_addTo3(con[cnt].pos, axis);
     mji_addToScl3(con[cnt].pos, normal, -con[cnt].dist * 0.5);
-    mji_copy3(con[cnt].frame, normal);
-    mju_zero3(con[cnt].frame+3);
+    mji_copy3(con[cnt].normal, normal);
+    mji_zero3(con[cnt].tangent);
     cnt++;
   } else {
     return 0;  // nearest point is above margin: no contacts
@@ -150,8 +167,8 @@ int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g
     mji_add3(con[cnt].pos, pos2, vec);
     mji_subFrom3(con[cnt].pos, axis);
     mji_addToScl3(con[cnt].pos, normal, -con[cnt].dist * 0.5);
-    mji_copy3(con[cnt].frame, normal);
-    mju_zero3(con[cnt].frame+3);
+    mji_copy3(con[cnt].normal, normal);
+    mji_zero3(con[cnt].tangent);
     cnt++;
   }
 
@@ -170,8 +187,8 @@ int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g
     mji_addTo3(con[cnt].pos, axis);
     mji_addToScl3(con[cnt].pos, vec, -0.5);
     mji_addToScl3(con[cnt].pos, normal, -con[cnt].dist * 0.5);
-    mji_copy3(con[cnt].frame, normal);
-    mju_zero3(con[cnt].frame+3);
+    mji_copy3(con[cnt].normal, normal);
+    mji_zero3(con[cnt].tangent);
     cnt++;
 
     // add point B
@@ -180,8 +197,8 @@ int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g
     mji_addTo3(con[cnt].pos, axis);
     mji_addToScl3(con[cnt].pos, vec, -0.5);
     mji_addToScl3(con[cnt].pos, normal, -con[cnt].dist * 0.5);
-    mji_copy3(con[cnt].frame, normal);
-    mju_zero3(con[cnt].frame+3);
+    mji_copy3(con[cnt].normal, normal);
+    mji_zero3(con[cnt].tangent);
     cnt++;
   }
 
@@ -190,8 +207,12 @@ int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g
 
 
 // plane : box
-int mjc_PlaneBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_PlaneBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size2 = m->geom_size + 3*g2;
 
   // get normal, difference between centers, normal distance
   mjtNum norm[3] = {mat1[2], mat1[5], mat1[8]};
@@ -219,11 +240,11 @@ int mjc_PlaneBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mj
 
     // construct contact
     con[cnt].dist = dist + ldist;
-    mji_copy3(con[cnt].frame, norm);
-    mju_zero3(con[cnt].frame+3);
+    mji_copy3(con[cnt].normal, norm);
     mji_addTo3(corner, pos2);
     mji_scl3(vec, norm, -con[cnt].dist / 2);
     mji_add3(con[cnt].pos, corner, vec);
+    mji_zero3(con[cnt].tangent);
 
     // count; max is 4
     if (++cnt >= 4) {
@@ -238,7 +259,7 @@ int mjc_PlaneBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mj
 //--------------------------- sphere and capsule collisions ----------------------------------------
 
 // sphere : sphere (actual implementation, can be called with modified parameters)
-static int mjraw_SphereSphere(mjContact* con, mjtNum margin,
+static int mjraw_SphereSphere(mjPreContact* con, mjtNum margin,
                               const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
                               const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
   // check bounding spheres (this is called from other functions)
@@ -251,36 +272,45 @@ static int mjraw_SphereSphere(mjContact* con, mjtNum margin,
 
   // depth and normal
   con[0].dist = mju_sqrt(cdist_sqr) - size1[0] - size2[0];
-  mji_sub3(con[0].frame, pos2, pos1);
-  mjtNum len = mju_normalize3(con[0].frame);
+  mji_sub3(con[0].normal, pos2, pos1);
+  mjtNum len = mju_normalize3(con[0].normal);
 
   // if centers are the same, norm = cross-product of z axes
   //  if z axes are parallel, norm = [1;0;0]
   if (len < mjMINVAL) {
     mjtNum axis1[3] = {mat1[2], mat1[5], mat1[8]};
     mjtNum axis2[3] = {mat2[2], mat2[5], mat2[8]};
-    mji_cross(con[0].frame, axis1, axis2);
-    mju_normalize3(con[0].frame);
+    mji_cross(con[0].normal, axis1, axis2);
+    mju_normalize3(con[0].normal);
   }
 
   // position
-  mji_scl3(con[0].pos, con[0].frame, size1[0] + con[0].dist / 2);
+  mji_scl3(con[0].pos, con[0].normal, size1[0] + con[0].dist / 2);
   mji_addTo3(con[0].pos, pos1);
 
-  mju_zero3(con[0].frame+3);
+  // axis
+  mji_zero3(con[0].tangent);
+
   return 1;
 }
 
 
 // sphere : sphere
-int mjc_SphereSphere(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_SphereSphere(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                     mjtNum margin) {
+  const mjtNum* pos1  = d->geom_xpos + 3*g1;
+  const mjtNum* mat1  = d->geom_xmat + 9*g1;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* pos2  = d->geom_xpos + 3*g2;
+  const mjtNum* mat2  = d->geom_xmat + 9*g2;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_SphereSphere(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
 
+
 // raw sphere : capsule
-int mjraw_SphereCapsule(mjContact* con, mjtNum margin,
+int mjraw_SphereCapsule(mjPreContact* con, mjtNum margin,
                         const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
                         const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
   // get capsule length and axis
@@ -299,16 +329,27 @@ int mjraw_SphereCapsule(mjContact* con, mjtNum margin,
 
 
 // sphere : capsule
-int mjc_SphereCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO
+int mjc_SphereCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                      mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_SphereCapsule(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
 
 // sphere : cylinder
-int mjc_SphereCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+int mjc_SphereCylinder(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                        mjtNum margin) {
-  mjGETINFO
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
 
   // get cylinder sizes and axis
   mjtNum radius = size2[0];
@@ -359,9 +400,10 @@ int mjc_SphereCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int
       mat_cap = flipmat;
     }
     int ncon = mjraw_PlaneSphere(con, margin, pos_cap, mat_cap, size2, pos1, mat1, size1);
+
     if (ncon) {
-      // flip frame normal (because mjGEOM_PLANE < mjGEOM_SPHERE < mjGEOM_CYLINDER)
-      mju_scl3(con->frame, con->frame, -1);
+      // flip direction normal (because mjGEOM_PLANE < mjGEOM_SPHERE < mjGEOM_CYLINDER)
+      mju_scl3(con[0].normal, con[0].normal, -1);
     }
     return ncon;
   }
@@ -380,7 +422,7 @@ int mjc_SphereCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int
 
 
 // raw capsule : capsule
-int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
+int mjraw_CapsuleCapsule(mjPreContact* con, mjtNum margin,
                          const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
                          const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2) {
   // get capsule axes (scaled) and center difference
@@ -444,7 +486,7 @@ int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
     x2 = mju_clip((v + mb) / mc, -1, 1);
     mji_scl3(vec2, axis2, x2);
     mji_addTo3(vec2, pos2);
-    int n2 = mjraw_SphereSphere(con+n1, margin, vec1, mat1, size1, vec2, mat2, size2);
+    int n2 = mjraw_SphereSphere(con + n1, margin, vec1, mat1, size1, vec2, mat2, size2);
 
     // return if two contacts already found
     if (n1+n2 >= 2) {
@@ -456,7 +498,7 @@ int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
     mjtNum x1 = mju_clip((u - mb) / ma, -1, 1);
     mji_scl3(vec1, axis1, x1);
     mji_addTo3(vec1, pos1);
-    int n3 = mjraw_SphereSphere(con+n1+n2, margin, vec1, mat1, size1, vec2, mat2, size2);
+    int n3 = mjraw_SphereSphere(con + n1 + n2, margin, vec1, mat1, size1, vec2, mat2, size2);
 
     // return if two contacts already found
     if (n1+n2+n3 >= 2) {
@@ -468,7 +510,7 @@ int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
     x1 = mju_clip((u + mb) / ma, -1, 1);
     mji_scl3(vec1, axis1, x1);
     mji_addTo3(vec1, pos1);
-    int n4 = mjraw_SphereSphere(con+n1+n2+n3, margin, vec1, mat1, size1, vec2, mat2, size2);
+    int n4 = mjraw_SphereSphere(con + n1 + n2 + n3, margin, vec1, mat1, size1, vec2, mat2, size2);
 
     return n1+n2+n3+n4;
   }
@@ -476,9 +518,14 @@ int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
 
 
 // capsule : capsule
-int mjc_CapsuleCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
-                       mjtNum margin) {
-  mjGETINFO
+int mjc_CapsuleCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                        mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_size + 3*g1;
+  const mjtNum* size2 = m->geom_size + 3*g2;
   return mjraw_CapsuleCapsule(con, margin, pos1, mat1, size1, pos2, mat2, size2);
 }
 
@@ -601,6 +648,7 @@ int mjraw_SphereTriangle(mjContact* con, mjtNum margin,
   return 1;
 }
 
+
 // box : triangle with radius
 int mjraw_BoxTriangle(mjContact* con, mjtNum margin, const mjtNum* pos,
                       const mjtNum* mat, const mjtNum* size, const mjtNum* t1,
@@ -693,6 +741,7 @@ int mjraw_BoxTriangle(mjContact* con, mjtNum margin, const mjtNum* pos,
   return cnt;
 }
 
+
 // capsule : triangle with radius
 int mjraw_CapsuleTriangle(mjContact* con, mjtNum margin, const mjtNum* pos,
                           const mjtNum* mat, const mjtNum* size,
diff --git a/src/engine/engine_collision_primitive.h b/src/engine/engine_collision_primitive.h
index 806e8742..b598c0e4 100644
--- a/src/engine/engine_collision_primitive.h
+++ b/src/engine/engine_collision_primitive.h
@@ -18,68 +18,57 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-
-// define and extract geom info
-#define mjGETINFO \
-    const mjtNum* pos1  = d->geom_xpos + 3*g1; \
-    const mjtNum* mat1  = d->geom_xmat + 9*g1; \
-    const mjtNum* size1 = m->geom_size + 3*g1; \
-    const mjtNum* pos2  = d->geom_xpos + 3*g2; \
-    const mjtNum* mat2  = d->geom_xmat + 9*g2; \
-    const mjtNum* size2 = m->geom_size + 3*g2; \
-    (void) size1; (void) size2;  // size1 and size2 are sometimes unused
+#include <mujoco/mjtnum.h>
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 // raw collision functions (called by mjc_XXX)
-int mjraw_SphereCapsule (mjContact* con, mjtNum margin,
-                         const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
-                         const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2);
-int mjraw_CapsuleCapsule(mjContact* con, mjtNum margin,
-                         const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
-                         const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2);
-int mjraw_CapsuleBox    (mjContact* con, mjtNum margin,
-                         const mjtNum* pos1, const mjtNum* mat1, const mjtNum* size1,
-                         const mjtNum* pos2, const mjtNum* mat2, const mjtNum* size2);
-int mjraw_SphereTriangle(mjContact* con, mjtNum margin,
-                         const mjtNum* s, mjtNum rs,
+int mjraw_SphereCapsule(mjPreContact* con, mjtNum margin, const mjtNum* pos1, const mjtNum* mat1,
+                        const mjtNum* size1, const mjtNum* pos2, const mjtNum* mat2,
+                        const mjtNum* size2);
+int mjraw_CapsuleCapsule(mjPreContact* con, mjtNum margin, const mjtNum* pos1, const mjtNum* mat1,
+                         const mjtNum* size1, const mjtNum* pos2, const mjtNum* mat2,
+                         const mjtNum* size2);
+int mjraw_CapsuleBox(mjPreContact* con, mjtNum margin, const mjtNum* pos1, const mjtNum* mat1,
+                     const mjtNum* size1, const mjtNum* pos2, const mjtNum* mat2,
+                     const mjtNum* size2);
+int mjraw_SphereTriangle(mjContact* con, mjtNum margin, const mjtNum* s, mjtNum rs,
                          const mjtNum* t1, const mjtNum* t2, const mjtNum* t3, mjtNum rt);
-int mjraw_BoxTriangle(mjContact* con, mjtNum margin, const mjtNum* pos,
-                      const mjtNum* mat, const mjtNum* size, const mjtNum* t1,
-                      const mjtNum* t2, const mjtNum* t3, mjtNum rt);
-int mjraw_CapsuleTriangle(mjContact* con, mjtNum margin, const mjtNum* pos,
-                          const mjtNum* mat, const mjtNum* size,
-                          const mjtNum* t1, const mjtNum* t2, const mjtNum* t3,
+int mjraw_BoxTriangle(mjContact* con, mjtNum margin, const mjtNum* pos, const mjtNum* mat,
+                      const mjtNum* size, const mjtNum* t1, const mjtNum* t2, const mjtNum* t3,
+                      mjtNum rt);
+int mjraw_CapsuleTriangle(mjContact* con, mjtNum margin, const mjtNum* pos, const mjtNum* mat,
+                          const mjtNum* size, const mjtNum* t1, const mjtNum* t2, const mjtNum* t3,
                           mjtNum rt);
 
 // plane collisions
-MJAPI int mjc_PlaneSphere(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_PlaneSphere(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                           mjtNum margin);
-MJAPI int mjc_PlaneCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_PlaneCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                            mjtNum margin);
-MJAPI int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_PlaneCylinder(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                             mjtNum margin);
-MJAPI int mjc_PlaneBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_PlaneBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                        mjtNum margin);
 
 // sphere and capsule collisions
-MJAPI int mjc_SphereSphere(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
-                           mjtNum margin);
-MJAPI int mjc_SphereCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_SphereSphere(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                             mjtNum margin);
-MJAPI int mjc_SphereCylinder(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_SphereCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
+                            mjtNum margin);
+MJAPI int mjc_SphereCylinder(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                              mjtNum margin);
-MJAPI int mjc_CapsuleCapsule(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_CapsuleCapsule(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                              mjtNum margin);
 
 // box collisions: from engine_collision_box.c
-MJAPI int mjc_CapsuleBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_CapsuleBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                          mjtNum margin);
-MJAPI int mjc_SphereBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_SphereBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                         mjtNum margin);
-MJAPI int mjc_BoxBox(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
+MJAPI int mjc_BoxBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
                      mjtNum margin);
 
 #ifdef __cplusplus
diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c
index cd074841..c5d093c0 100644
--- a/src/engine/engine_collision_sdf.c
+++ b/src/engine/engine_collision_sdf.c
@@ -21,7 +21,6 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include <mujoco/mjtnum.h>
-#include "engine/engine_collision_primitive.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
@@ -543,10 +542,10 @@ static int isknown(const mjtNum* points, const mjtNum x[3], int cnt) {
 // adds candidate point to result
 // flipNormal: 0 = normal points INTO SDF (for mesh-SDF where SDF is g2)
 //             1 = normal points OUT of SDF (for flex-SDF where SDF is g1)
-static int addContact(mjtNum* points, mjContact* con, const mjtNum x[3],
-                      const mjtNum pos2[3], const mjtNum quat2[4], mjtNum dist,
-                      int cnt, const mjModel* m, const mjSDF* s, const mjData* d,
-                      int flipNormal) {
+static int addPreContact(mjtNum* points, mjPreContact* con, const mjtNum x[3],
+                         const mjtNum pos2[3], const mjtNum quat2[4], mjtNum dist,
+                         int cnt, const mjModel* m, const mjSDF* s, const mjData* d,
+                         int flipNormal) {
   // check if there is a collision
   if (dist > 0 || isknown(points, x, cnt)) {
     return cnt;
@@ -570,18 +569,37 @@ static int addContact(mjtNum* points, mjContact* con, const mjtNum x[3],
   }
 
   // construct contact
-  con[cnt].dist = dist;
-  mju_rotVecQuat(con[cnt].frame, norm, quat2);
-  mju_zero3(con[cnt].frame+3);
-  mju_makeFrame(con[cnt].frame);
-  mju_scl3(vec, con[cnt].frame, -con[cnt].dist/2);
-  mju_rotVecQuat(con[cnt].pos, x, quat2);
-  mju_addTo3(con[cnt].pos, pos2);
-  mju_addTo3(con[cnt].pos, vec);
+  con->dist = dist;
+  mju_rotVecQuat(con->normal, norm, quat2);
+  mju_scl3(vec, con->normal, - 0.5 * dist);
+  mju_rotVecQuat(con->pos, x, quat2);
+  mju_zero3(con->tangent);
+  mju_addTo3(con->pos, pos2);
+  mju_addTo3(con->pos, vec);
 
   return cnt+1;
 }
 
+
+// adds candidate point to result
+// flipNormal: 0 = normal points INTO SDF (for mesh-SDF where SDF is g2)
+//             1 = normal points OUT of SDF (for flex-SDF where SDF is g1)
+static int addContact(mjtNum* points, mjContact* con, const mjtNum x[3],
+                      const mjtNum pos2[3], const mjtNum quat2[4], mjtNum dist,
+                      int cnt, const mjModel* m, const mjSDF* s, const mjData* d,
+                      int flipNormal) {
+  mjPreContact precon;
+  int ncon = addPreContact(points, &precon, x, pos2, quat2, dist, cnt, m, s, d,
+                           flipNormal);
+  if (ncon > cnt) {
+    con[cnt].dist = precon.dist;
+    mju_copy3(con[cnt].pos, precon.pos);
+    mju_copy3(con[cnt].frame, precon.normal);
+    mju_zero3(con[cnt].frame + 3);
+  }
+  return ncon;
+}
+
 // finds minimum of Frank-Wolfe objective
 static mjtNum stepFrankWolfe(mjtNum x[3], const mjtNum* corners, int ncorners,
                              const mjModel* m, const mjSDF* sdf, const mjData* d) {
@@ -952,18 +970,21 @@ static int meshFaceCallback(int face_id, int node, void* ctx) {
 //------------------------------ collision functions -----------------------------------------------
 
 // collision between a height field and a signed distance field
-int mjc_HFieldSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
+int mjc_HFieldSDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
   mju_warning("HField vs SDF collision not yet supported!");
   return 0;
 }
 
 
 // collision between a mesh and a signed distance field
-int mjc_MeshSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO;
+int mjc_MeshSDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1 = d->geom_xpos + 3*g1;
+  const mjtNum* mat1 = d->geom_xmat + 9*g1;
+  const mjtNum* pos2 = d->geom_xpos + 3*g2;
+  const mjtNum* mat2 = d->geom_xmat + 9*g2;
 
   mjtNum offset[3], rotation[9];
-  mjtNum points[3*mjMAXCONPAIR], dist[mjMAXCONPAIR], candidate[3*mjMAXCONPAIR];
+  mjtNum points[3*mjMAXCONPAIR], dist2[mjMAXCONPAIR], candidate[3*mjMAXCONPAIR];
   int cnt = 0, ncandidate = 0;
 
   // get sdf plugin
@@ -996,7 +1017,7 @@ int mjc_MeshSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjt
   ctx.faceadr = m->mesh_faceadr[m->geom_dataid[g1]];
   ctx.nstartpts = mju_max(1, m->opt.sdf_initpoints);
   ctx.candidate = candidate;
-  ctx.dist = dist;
+  ctx.dist = dist2;
   ctx.ncandidate = &ncandidate;
 
   // BVH traversal for mesh faces
@@ -1014,34 +1035,37 @@ int mjc_MeshSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjt
   // if few candidates, add them all directly
   if (ncandidate <= mjMAXCONPAIR) {
     for (int i = 0; i < ncandidate; i++) {
-      cnt = addContact(points, con, candidate + 3*i, pos2, sdf_quat,
-                       dist[i], cnt, m, &sdf, d, 0);
+      cnt = addPreContact(points, con + cnt, candidate + 3*i, pos2, sdf_quat,
+                          dist2[i], cnt, m, &sdf, d, 0);
     }
     return cnt;
   }
 
   // farthest point sampling (FPS) for spatial coverage
   int selected_indices[mjMAXCONPAIR];
-  int nselected = selectFPS(candidate, dist, ncandidate, selected_indices, mjMAXCONPAIR);
+  int nselected = selectFPS(candidate, dist2, ncandidate, selected_indices, mjMAXCONPAIR);
 
   // add selected contacts
   for (int i = 0; i < nselected; i++) {
     int idx = selected_indices[i];
-    cnt = addContact(points, con, candidate + 3*idx, pos2, sdf_quat,
-                     dist[idx], cnt, m, &sdf, d, 0);
+    cnt = addPreContact(points, con + cnt, candidate + 3*idx, pos2, sdf_quat,
+                        dist2[idx], cnt, m, &sdf, d, 0);
   }
 
   return cnt;
 }
 
 // collision between two SDFs
-int mjc_SDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin) {
-  mjGETINFO;
-  size1 = m->geom_aabb + 6*g1;
-  size2 = m->geom_aabb + 6*g2;
+int mjc_SDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin) {
+  const mjtNum* pos1  = d->geom_xpos + 3*g1;
+  const mjtNum* mat1  = d->geom_xmat + 9*g1;
+  const mjtNum* pos2  = d->geom_xpos + 3*g2;
+  const mjtNum* mat2  = d->geom_xmat + 9*g2;
+  const mjtNum* size1 = m->geom_aabb + 6*g1;
+  const mjtNum* size2 = m->geom_aabb + 6*g2;
 
   int cnt = 0;
-  mjtNum x[3], y[3], dist, vec1[3], vec2[3];
+  mjtNum x[3], y[3], dist2, vec1[3], vec2[3];
   mjtNum aabb1[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
   mjtNum aabb2[6] = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
   mjtNum aabb[6]  = {mjMAXVAL, mjMAXVAL, mjMAXVAL, -mjMAXVAL, -mjMAXVAL, -mjMAXVAL};
@@ -1153,15 +1177,15 @@ int mjc_SDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum
 
     // gradient descent - we use a special function of the two SDF as objective
     sdf.type = mjSDFTYPE_COLLISION;
-    dist = stepGradient(x, m, &sdf, d, m->opt.sdf_iterations);
+    dist2 = stepGradient(x, m, &sdf, d, m->opt.sdf_iterations);
 
     // inexact SDFs can yield spurious collisions, filter them by projecting on the midsurface
     sdf.type = mjSDFTYPE_INTERSECTION;
-    dist = stepGradient(x, m, &sdf, d, 1);
+    dist2 = stepGradient(x, m, &sdf, d, 1);
 
     // contact point and normal - we use the midsurface where SDF1=SDF2 as zero level set
     sdf.type = mjSDFTYPE_MIDSURFACE;
-    cnt = addContact(contacts, con, x, pos2, quat2, dist, cnt, m, &sdf, d, 0);
+    cnt = addPreContact(contacts, con + cnt, x, pos2, quat2, dist2, cnt, m, &sdf, d, 0);
 
     // SHOULD NOT OCCUR
     if (cnt > mjMAXCONPAIR) {
diff --git a/src/engine/engine_collision_sdf.h b/src/engine/engine_collision_sdf.h
index b50f3082..64f9b298 100644
--- a/src/engine/engine_collision_sdf.h
+++ b/src/engine/engine_collision_sdf.h
@@ -36,13 +36,13 @@ MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNu
                         const mjtNum x[3]);
 
 // collision between a height field and a signed distance field
-int mjc_HFieldSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_HFieldSDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
 
 // collision between a mesh and a signed distance field
-int mjc_MeshSDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_MeshSDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
 
 // collision between two signed distance fields
-int mjc_SDF(const mjModel* m, mjData* d, mjContact* con, int g1, int g2, mjtNum margin);
+int mjc_SDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, mjtNum margin);
 
 // collision between entire flex and SDF geom (batched processing)
 int mjc_FlexSDF(const mjModel* m, const mjData* d, mjContact* con,
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index 6a247f09..eaca8ba8 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -553,7 +553,7 @@ static mjtNum mj_geomDistanceCCD(const mjModel* m, mjData* d, int g1, int g2,
 // returns the smallest distance between two geoms
 mjtNum mj_geomDistance(const mjModel* m, mjData* d, int geom1, int geom2, mjtNum distmax,
                        mjtNum fromto[6]) {
-  mjContact con[mjMAXCONPAIR];
+  mjPreContact con[mjMAXCONPAIR];
   mjtNum dist = distmax;
   if (fromto) mju_zero(fromto, 6);
 
@@ -594,8 +594,8 @@ mjtNum mj_geomDistance(const mjModel* m, mjData* d, int geom1, int geom2, mjtNum
   // write fromto if given and a collision has been found
   if (fromto && smallest >= 0) {
     mjtNum sign = flip ? -1 : 1;
-    mju_addScl3(fromto+0, con[smallest].pos, con[smallest].frame, -0.5*sign*dist);
-    mju_addScl3(fromto+3, con[smallest].pos, con[smallest].frame, 0.5*sign*dist);
+    mju_addScl3(fromto+0, con[smallest].pos, con[smallest].normal, -0.5*sign*dist);
+    mju_addScl3(fromto+3, con[smallest].pos, con[smallest].normal, 0.5*sign*dist);
   }
 
   return dist;
diff --git a/src/engine/engine_util_spatial.c b/src/engine/engine_util_spatial.c
index bfff1a35..fb7ec3fb 100644
--- a/src/engine/engine_util_spatial.c
+++ b/src/engine/engine_util_spatial.c
@@ -504,7 +504,7 @@ void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_force,
 }
 
 
-// make 3D frame given X axis (normal) and possibly Y axis (tangent 1)
+// make 3D frame given x-axis (normal) and possibly y-axis (tangent 1)
 void mju_makeFrame(mjtNum frame[9]) {
   mjtNum tmp[3];
 
diff --git a/test/engine/engine_collision_box_test.cc b/test/engine/engine_collision_box_test.cc
index e75a6c83..0ac0ae62 100644
--- a/test/engine/engine_collision_box_test.cc
+++ b/test/engine/engine_collision_box_test.cc
@@ -15,6 +15,7 @@
 // Tests for engine/engine_collision_box.c.
 
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -45,11 +46,9 @@ TEST_F(MjCollisionBoxTest, BadContacts) {
     mj_forward(model, data);
 
     // allocate contact array and matching arrays
-    mj_markStack(data);
-    mjContact* con_raw = (mjContact*) mj_stackAllocByte(
-        data, mjMAXCONPAIR * sizeof(mjContact), alignof(mjContact));
-    int* match_raw = mj_stackAllocInt(data, mjMAXCONPAIR);
-    int* match = mj_stackAllocInt(data, data->ncon);
+    std::vector<mjPreContact> precon(mjMAXCONPAIR);
+    std::vector<int> match_raw(mjMAXCONPAIR);
+    std::vector<int> match(data->ncon);
 
     int g1 = -1;
     int g2 = -1;
@@ -73,20 +72,21 @@ TEST_F(MjCollisionBoxTest, BadContacts) {
       g2 = g2new;
 
       // call low-level box-box collider
-      int num = mjc_BoxBox(model, data, con_raw, g1, g2, con->includemargin);
+      int num = mjc_BoxBox(model, data, precon.data(), g1, g2,
+                           con->includemargin);
 
       // allocate and clear arrays marking already matched contacts
-      mju_zeroInt(match_raw, num);
-      mju_zeroInt(match, data->ncon);
+      mju_zeroInt(match_raw.data(), num);
+      mju_zeroInt(match.data(), data->ncon);
 
       // loop over raw contacts, match with contact array using pos
       int nmatched = 0;
       for (int i = 0; i < num; i++) {
         for (int j = 0; j < data->ncon; j++) {
           if (!match[j] &&
-              con_raw[i].pos[0] == data->contact[j].pos[0] &&
-              con_raw[i].pos[1] == data->contact[j].pos[1] &&
-              con_raw[i].pos[2] == data->contact[j].pos[2]) {
+              precon[i].pos[0] == data->contact[j].pos[0] &&
+              precon[i].pos[1] == data->contact[j].pos[1] &&
+              precon[i].pos[2] == data->contact[j].pos[2]) {
             match_raw[i] = match[j] = 1;
             nmatched++;
           }
@@ -97,11 +97,11 @@ TEST_F(MjCollisionBoxTest, BadContacts) {
       EXPECT_EQ(nmatched, num) << local_path;
 
       // get box info
-      const mjtNum* pos1 =  data->geom_xpos  + 3 * g1;
-      const mjtNum* mat1 =  data->geom_xmat  + 9 * g1;
+      const mjtNum* pos1 = data->geom_xpos + 3 * g1;
+      const mjtNum* mat1 = data->geom_xmat + 9 * g1;
       const mjtNum* size1 = model->geom_size + 3 * g1;
-      const mjtNum* pos2 =  data->geom_xpos  + 3 * g2;
-      const mjtNum* mat2 =  data->geom_xmat  + 9 * g2;
+      const mjtNum* pos2 = data->geom_xpos + 3 * g2;
+      const mjtNum* mat2 = data->geom_xmat + 9 * g2;
       const mjtNum* size2 = model->geom_size + 3 * g2;
       mjtNum margin = mju_max(model->geom_margin[g1], model->geom_margin[g2]);
 
@@ -117,8 +117,10 @@ TEST_F(MjCollisionBoxTest, BadContacts) {
           static mjtNum kRatio = 1.01;
 
           // is the contact outside: 1, inside: -1, within the removal width: 0
-          int out1 = mju_outsideBox(con_raw[i].pos, pos1, mat1, sz1, kRatio);
-          int out2 = mju_outsideBox(con_raw[i].pos, pos2, mat2, sz2, kRatio);
+          int out1 = mju_outsideBox(precon[i].pos, pos1, mat1, sz1,
+                                    kRatio);
+          int out2 = mju_outsideBox(precon[i].pos, pos2, mat2, sz2,
+                                    kRatio);
 
           // mark as bad if outside one box and not inside the other box
           bool outside = (out1 == 1 && out2 != -1) || (out2 == 1 && out1 != -1);
@@ -129,7 +131,6 @@ TEST_F(MjCollisionBoxTest, BadContacts) {
       }
     }
 
-    mj_freeStack(data);
     mj_deleteData(data);
     mj_deleteModel(model);
   }
@@ -146,11 +147,9 @@ TEST_F(MjCollisionBoxTest, DuplicateContacts) {
   mj_forward(model, data);
 
   // allocate contact array and matching arrays
-  mj_markStack(data);
-  mjContact* con_raw = (mjContact*) mj_stackAllocByte(
-      data, mjMAXCONPAIR * sizeof(mjContact), alignof(mjContact));
-  int* match_raw = mj_stackAllocInt(data, mjMAXCONPAIR);
-  int* match = mj_stackAllocInt(data, data->ncon);
+  std::vector<mjPreContact> precon(mjMAXCONPAIR);
+  std::vector<int> match_raw(mjMAXCONPAIR);
+  std::vector<int> match(data->ncon);
 
 
   int g1 = -1;
@@ -175,20 +174,21 @@ TEST_F(MjCollisionBoxTest, DuplicateContacts) {
     g2 = g2new;
 
     // call low-level box-box collider
-    int num = mjc_BoxBox(model, data, con_raw, g1, g2, con->includemargin);
+    int num = mjc_BoxBox(model, data, precon.data(), g1, g2,
+                         con->includemargin);
 
     // allocate and clear arrays marking already matched contacts
-    mju_zeroInt(match_raw, num);
-    mju_zeroInt(match, data->ncon);
+    mju_zeroInt(match_raw.data(), num);
+    mju_zeroInt(match.data(), data->ncon);
 
     // loop over raw contacts, match with contact array using pos
     int nmatched = 0;
     for (int i = 0; i < num; i++) {
       for (int j = 0; j < data->ncon; j++) {
         if (!match[j] &&
-            con_raw[i].pos[0] == data->contact[j].pos[0] &&
-            con_raw[i].pos[1] == data->contact[j].pos[1] &&
-            con_raw[i].pos[2] == data->contact[j].pos[2]) {
+            precon[i].pos[0] == data->contact[j].pos[0] &&
+            precon[i].pos[1] == data->contact[j].pos[1] &&
+            precon[i].pos[2] == data->contact[j].pos[2]) {
           match_raw[i] = match[j] = 1;
           nmatched++;
         }
@@ -207,9 +207,9 @@ TEST_F(MjCollisionBoxTest, DuplicateContacts) {
           if (duplicate || i == j) {
             continue;
           }
-          if (con_raw[i].pos[0] == con_raw[j].pos[0] &&
-              con_raw[i].pos[1] == con_raw[j].pos[1] &&
-              con_raw[i].pos[2] == con_raw[j].pos[2]) {
+          if (precon[i].pos[0] == precon[j].pos[0] &&
+              precon[i].pos[1] == precon[j].pos[1] &&
+              precon[i].pos[2] == precon[j].pos[2]) {
             duplicate = true;
           }
         }
@@ -220,12 +220,10 @@ TEST_F(MjCollisionBoxTest, DuplicateContacts) {
     }
   }
 
-  mj_freeStack(data);
   mj_deleteData(data);
   mj_deleteModel(model);
 }
 
-
 static const char* const kDeepFilePath =
     "engine/testdata/collision_box/boxbox_deep.xml";
 
@@ -273,7 +271,6 @@ TEST_F(MjCollisionBoxTest, BoxSphere) {
   mj_deleteModel(model);
 }
 
-
 TEST_F(MjCollisionBoxTest, BoxBoxContactDistance) {
   constexpr char xml[] = R"(
   <mujoco>
@@ -289,12 +286,12 @@ TEST_F(MjCollisionBoxTest, BoxBoxContactDistance) {
 
   mjData* data = mj_makeData(model);
   mj_kinematics(model, data);
-  mjContact contact[9];
+  mjPreContact precon[mjMAXCONPAIR];
 
   for (mjfCollision collision : {mjc_BoxBox, mjc_Convex}) {
-    int n = collision(model, data, contact, 0, 1, 0.0);
+    int n = collision(model, data, precon, 0, 1, 0.0);
     for (int i = 0; i < n; i++) {
-      EXPECT_NEAR(contact[i].dist, -0.5, MjTol(1e-8, 1e-6));
+      EXPECT_NEAR(precon[i].dist, -0.5, MjTol(1e-8, 1e-6));
     }
   }
 
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 164edc61..e66b47a3 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -1520,6 +1520,14 @@ public unsafe struct mjModel_ {
   public UInt64 signature;
 }
 
+[StructLayout(LayoutKind.Sequential)]
+public unsafe struct mjPreContact_ {
+  public double dist;
+  public fixed double pos[3];
+  public fixed double normal[3];
+  public fixed double tangent[3];
+}
+
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct mjContact_ {
   public double dist;
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index c82c2de3..8262902d 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -766,6 +766,36 @@ struct MjOption {
   bool owned_ = false;
 };
 
+struct MjPreContact {
+  ~MjPreContact();
+  MjPreContact();
+  explicit MjPreContact(mjPreContact *ptr);
+  MjPreContact(const MjPreContact &);
+  MjPreContact &operator=(const MjPreContact &);
+  std::unique_ptr<MjPreContact> copy();
+  mjPreContact* get() const;
+  void set(mjPreContact* ptr);
+  mjtNum dist() const {
+    return ptr_->dist;
+  }
+  void set_dist(mjtNum value) {
+    ptr_->dist = value;
+  }
+  emscripten::val pos() const {
+    return emscripten::val(emscripten::typed_memory_view(3, ptr_->pos));
+  }
+  emscripten::val normal() const {
+    return emscripten::val(emscripten::typed_memory_view(3, ptr_->normal));
+  }
+  emscripten::val tangent() const {
+    return emscripten::val(emscripten::typed_memory_view(3, ptr_->tangent));
+  }
+
+ private:
+  mjPreContact* ptr_;
+  bool owned_ = false;
+};
+
 struct MjSolverStat {
   ~MjSolverStat();
   MjSolverStat();
@@ -7139,7 +7169,7 @@ MjContact::~MjContact() {
     delete ptr_;
   }
 }
-MjContact::MjContact() : ptr_(new mjContact) {
+MjContact::MjContact() : ptr_(new mjContact()) {
   owned_ = true;
 }
 MjContact::MjContact(const MjContact &other) : MjContact() {
@@ -7168,7 +7198,7 @@ MjLROpt::~MjLROpt() {
     delete ptr_;
   }
 }
-MjLROpt::MjLROpt() : ptr_(new mjLROpt) {
+MjLROpt::MjLROpt() : ptr_(new mjLROpt()) {
   owned_ = true;
   mj_defaultLROpt(ptr_);
 }
@@ -7198,7 +7228,7 @@ MjOption::~MjOption() {
     delete ptr_;
   }
 }
-MjOption::MjOption() : ptr_(new mjOption) {
+MjOption::MjOption() : ptr_(new mjOption()) {
   owned_ = true;
   mj_defaultOption(ptr_);
 }
@@ -7222,13 +7252,42 @@ void MjOption::set(mjOption* ptr) {
   ptr_ = ptr;
 }
 
+MjPreContact::MjPreContact(mjPreContact *ptr) : ptr_(ptr) {}
+MjPreContact::~MjPreContact() {
+  if (owned_ && ptr_) {
+    delete ptr_;
+  }
+}
+MjPreContact::MjPreContact() : ptr_(new mjPreContact()) {
+  owned_ = true;
+}
+MjPreContact::MjPreContact(const MjPreContact &other) : MjPreContact() {
+  *ptr_ = *other.get();
+}
+MjPreContact& MjPreContact::operator=(const MjPreContact &other) {
+  if (this == &other) {
+    return *this;
+  }
+  *ptr_ = *other.get();
+  return *this;
+}
+std::unique_ptr<MjPreContact> MjPreContact::copy() {
+  return std::make_unique<MjPreContact>(*this);
+}
+mjPreContact* MjPreContact::get() const {
+  return ptr_;
+}
+void MjPreContact::set(mjPreContact* ptr) {
+  ptr_ = ptr;
+}
+
 MjSolverStat::MjSolverStat(mjSolverStat *ptr) : ptr_(ptr) {}
 MjSolverStat::~MjSolverStat() {
   if (owned_ && ptr_) {
     delete ptr_;
   }
 }
-MjSolverStat::MjSolverStat() : ptr_(new mjSolverStat) {
+MjSolverStat::MjSolverStat() : ptr_(new mjSolverStat()) {
   owned_ = true;
 }
 MjSolverStat::MjSolverStat(const MjSolverStat &other) : MjSolverStat() {
@@ -7257,7 +7316,7 @@ MjStatistic::~MjStatistic() {
     delete ptr_;
   }
 }
-MjStatistic::MjStatistic() : ptr_(new mjStatistic) {
+MjStatistic::MjStatistic() : ptr_(new mjStatistic()) {
   owned_ = true;
 }
 MjStatistic::MjStatistic(const MjStatistic &other) : MjStatistic() {
@@ -7286,7 +7345,7 @@ MjTimerStat::~MjTimerStat() {
     delete ptr_;
   }
 }
-MjTimerStat::MjTimerStat() : ptr_(new mjTimerStat) {
+MjTimerStat::MjTimerStat() : ptr_(new mjTimerStat()) {
   owned_ = true;
 }
 MjTimerStat::MjTimerStat(const MjTimerStat &other) : MjTimerStat() {
@@ -7315,7 +7374,7 @@ MjVisualGlobal::~MjVisualGlobal() {
     delete ptr_;
   }
 }
-MjVisualGlobal::MjVisualGlobal() : ptr_(new mjVisualGlobal) {
+MjVisualGlobal::MjVisualGlobal() : ptr_(new mjVisualGlobal()) {
   owned_ = true;
 }
 MjVisualGlobal::MjVisualGlobal(const MjVisualGlobal &other) : MjVisualGlobal() {
@@ -7344,7 +7403,7 @@ MjVisualHeadlight::~MjVisualHeadlight() {
     delete ptr_;
   }
 }
-MjVisualHeadlight::MjVisualHeadlight() : ptr_(new mjVisualHeadlight) {
+MjVisualHeadlight::MjVisualHeadlight() : ptr_(new mjVisualHeadlight()) {
   owned_ = true;
 }
 MjVisualHeadlight::MjVisualHeadlight(const MjVisualHeadlight &other) : MjVisualHeadlight() {
@@ -7373,7 +7432,7 @@ MjVisualMap::~MjVisualMap() {
     delete ptr_;
   }
 }
-MjVisualMap::MjVisualMap() : ptr_(new mjVisualMap) {
+MjVisualMap::MjVisualMap() : ptr_(new mjVisualMap()) {
   owned_ = true;
 }
 MjVisualMap::MjVisualMap(const MjVisualMap &other) : MjVisualMap() {
@@ -7402,7 +7461,7 @@ MjVisualQuality::~MjVisualQuality() {
     delete ptr_;
   }
 }
-MjVisualQuality::MjVisualQuality() : ptr_(new mjVisualQuality) {
+MjVisualQuality::MjVisualQuality() : ptr_(new mjVisualQuality()) {
   owned_ = true;
 }
 MjVisualQuality::MjVisualQuality(const MjVisualQuality &other) : MjVisualQuality() {
@@ -7431,7 +7490,7 @@ MjVisualRgba::~MjVisualRgba() {
     delete ptr_;
   }
 }
-MjVisualRgba::MjVisualRgba() : ptr_(new mjVisualRgba) {
+MjVisualRgba::MjVisualRgba() : ptr_(new mjVisualRgba()) {
   owned_ = true;
 }
 MjVisualRgba::MjVisualRgba(const MjVisualRgba &other) : MjVisualRgba() {
@@ -7460,7 +7519,7 @@ MjVisualScale::~MjVisualScale() {
     delete ptr_;
   }
 }
-MjVisualScale::MjVisualScale() : ptr_(new mjVisualScale) {
+MjVisualScale::MjVisualScale() : ptr_(new mjVisualScale()) {
   owned_ = true;
 }
 MjVisualScale::MjVisualScale(const MjVisualScale &other) : MjVisualScale() {
@@ -7489,7 +7548,7 @@ MjWarningStat::~MjWarningStat() {
     delete ptr_;
   }
 }
-MjWarningStat::MjWarningStat() : ptr_(new mjWarningStat) {
+MjWarningStat::MjWarningStat() : ptr_(new mjWarningStat()) {
   owned_ = true;
 }
 MjWarningStat::MjWarningStat(const MjWarningStat &other) : MjWarningStat() {
@@ -7534,7 +7593,7 @@ MjvCamera::~MjvCamera() {
     delete ptr_;
   }
 }
-MjvCamera::MjvCamera() : ptr_(new mjvCamera) {
+MjvCamera::MjvCamera() : ptr_(new mjvCamera()) {
   owned_ = true;
   mjv_defaultCamera(ptr_);
 }
@@ -7564,7 +7623,7 @@ MjvFigure::~MjvFigure() {
     delete ptr_;
   }
 }
-MjvFigure::MjvFigure() : ptr_(new mjvFigure) {
+MjvFigure::MjvFigure() : ptr_(new mjvFigure()) {
   owned_ = true;
   mjv_defaultFigure(ptr_);
 }
@@ -7594,7 +7653,7 @@ MjvGLCamera::~MjvGLCamera() {
     delete ptr_;
   }
 }
-MjvGLCamera::MjvGLCamera() : ptr_(new mjvGLCamera) {
+MjvGLCamera::MjvGLCamera() : ptr_(new mjvGLCamera()) {
   owned_ = true;
 }
 MjvGLCamera::MjvGLCamera(const MjvGLCamera &other) : MjvGLCamera() {
@@ -7623,7 +7682,7 @@ MjvGeom::~MjvGeom() {
     delete ptr_;
   }
 }
-MjvGeom::MjvGeom() : ptr_(new mjvGeom) {
+MjvGeom::MjvGeom() : ptr_(new mjvGeom()) {
   owned_ = true;
   mjv_initGeom(ptr_, mjGEOM_NONE, nullptr, nullptr, nullptr, nullptr);
 }
@@ -7653,7 +7712,7 @@ MjvLight::~MjvLight() {
     delete ptr_;
   }
 }
-MjvLight::MjvLight() : ptr_(new mjvLight) {
+MjvLight::MjvLight() : ptr_(new mjvLight()) {
   owned_ = true;
 }
 MjvLight::MjvLight(const MjvLight &other) : MjvLight() {
@@ -7682,7 +7741,7 @@ MjvOption::~MjvOption() {
     delete ptr_;
   }
 }
-MjvOption::MjvOption() : ptr_(new mjvOption) {
+MjvOption::MjvOption() : ptr_(new mjvOption()) {
   owned_ = true;
   mjv_defaultOption(ptr_);
 }
@@ -7712,7 +7771,7 @@ MjvPerturb::~MjvPerturb() {
     delete ptr_;
   }
 }
-MjvPerturb::MjvPerturb() : ptr_(new mjvPerturb) {
+MjvPerturb::MjvPerturb() : ptr_(new mjvPerturb()) {
   owned_ = true;
   mjv_defaultPerturb(ptr_);
 }
@@ -7750,7 +7809,7 @@ MjVisual::~MjVisual() {
     delete ptr_;
   }
 }
-MjVisual::MjVisual() : ptr_(new mjVisual), global(&ptr_->global), quality(&ptr_->quality), headlight(&ptr_->headlight), map(&ptr_->map), scale(&ptr_->scale), rgba(&ptr_->rgba) {
+MjVisual::MjVisual() : ptr_(new mjVisual()), global(&ptr_->global), quality(&ptr_->quality), headlight(&ptr_->headlight), map(&ptr_->map), scale(&ptr_->scale), rgba(&ptr_->rgba) {
   owned_ = true;
   mj_defaultVisual(ptr_);
 }
@@ -12375,6 +12434,13 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("tolerance", &MjOption::tolerance, &MjOption::set_tolerance, reference())
     .property("viscosity", &MjOption::viscosity, &MjOption::set_viscosity, reference())
     .property("wind", &MjOption::wind);
+  emscripten::class_<MjPreContact>("MjPreContact")
+    .constructor<>()
+    .function("copy", &MjPreContact::copy, take_ownership())
+    .property("dist", &MjPreContact::dist, &MjPreContact::set_dist, reference())
+    .property("normal", &MjPreContact::normal)
+    .property("pos", &MjPreContact::pos)
+    .property("tangent", &MjPreContact::tangent);
   emscripten::class_<MjSolverStat>("MjSolverStat")
     .constructor<>()
     .function("copy", &MjSolverStat::copy, take_ownership())
diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py
index c5a84f8b..2338dff0 100644
--- a/wasm/codegen/generators/constants.py
+++ b/wasm/codegen/generators/constants.py
@@ -292,6 +292,7 @@ STRUCTS_TO_BIND: list[str] = list(
 NO_DEFAULT_CONSTRUCTORS: tuple[str, ...] = (
     # go/keep-sorted start
     "mjContact",
+    "mjPreContact",
     "mjSolverStat",
     "mjStatistic",
     "mjTimerStat",
diff --git a/wasm/codegen/generators/structs.py b/wasm/codegen/generators/structs.py
index 0bb03066..a2731326 100644
--- a/wasm/codegen/generators/structs.py
+++ b/wasm/codegen/generators/structs.py
@@ -493,7 +493,7 @@ def build_struct_source(
 
   if not is_mjs:
     # default constructor
-    with builder.function(f"{w}::{w}() : ptr_(new {s}){fields_init}"):
+    with builder.function(f"{w}::{w}() : ptr_(new {s}()){fields_init}"):
       builder.line("owned_ = true;")
       default_func = _default_function_statement(s)
       if default_func:
diff --git a/wasm/codegen/tests/generators_test.py b/wasm/codegen/tests/generators_test.py
index 1d21bc89..3da3eb02 100644
--- a/wasm/codegen/tests/generators_test.py
+++ b/wasm/codegen/tests/generators_test.py
@@ -266,7 +266,7 @@ MjvLight::~MjvLight() {
     delete ptr_;
   }
 }
-MjvLight::MjvLight() : ptr_(new mjvLight) {
+MjvLight::MjvLight() : ptr_(new mjvLight()) {
   owned_ = true;
 }
 MjvLight::MjvLight(const MjvLight &other) : MjvLight() {

From 7c47a89cd0334d711da0ba8bf0c9d30140ad4784 Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Wed, 20 May 2026 16:25:53 -0700
Subject: [PATCH 09/61] Add asset file download indicator to the loading
 screen.

This provides visual feedback to the user while loading large models.

PiperOrigin-RevId: 918708877
Change-Id: I77ea475cce6d9718e6cf8b33cb72b3f524a6ee63
---
 src/experimental/studio/emscripten.cc | 18 +++++++++++++++++-
 src/experimental/studio/live.css      | 11 +++++++++++
 src/experimental/studio/live.html     |  4 ++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/experimental/studio/emscripten.cc b/src/experimental/studio/emscripten.cc
index 0e427626..644c9ea6 100644
--- a/src/experimental/studio/emscripten.cc
+++ b/src/experimental/studio/emscripten.cc
@@ -71,8 +71,17 @@ class AssetRegistry {
 // The caller is responsible for freeing the buffer. Returns 0 on failure.
 EM_ASYNC_JS(int, FetchUrl,
             (const char* url, char** out_data, std::int32_t* out_size), {
+              const urlStr = UTF8ToString(url);
+              const filename = urlStr.split('/').pop() || urlStr;
+              const dlEl = document.getElementById('loadingDownload');
+              const dlFileEl = document.getElementById('loadingDownloadFile');
+              if (dlEl) {
+                dlEl.style.display = 'block';
+              }
+              if (dlFileEl) {
+                dlFileEl.textContent = filename + '\u2026';
+              }
               try {
-                const urlStr = UTF8ToString(url);
                 const response = await fetch(urlStr);
                 if (!response.ok) {
                   console.error('Fetch failed: ' + response.status + ' ' +
@@ -89,6 +98,13 @@ EM_ASYNC_JS(int, FetchUrl,
               } catch (e) {
                 console.error('Fetch error:', e);
                 return 0;
+              } finally {
+                if (dlEl) {
+                  dlEl.style.display = 'none';
+                }
+                if (dlFileEl) {
+                  dlFileEl.textContent = "";
+                }
               }
             });
 
diff --git a/src/experimental/studio/live.css b/src/experimental/studio/live.css
index 74d7cfe9..6c668b27 100644
--- a/src/experimental/studio/live.css
+++ b/src/experimental/studio/live.css
@@ -23,6 +23,17 @@
   font-size: 14px;
   opacity: 0.7;
 }
+.loading-download {
+  display: none;
+  margin-top: 24px;
+  opacity: 0.6;
+}
+.loading-download-label {
+  font-size: 13px;
+}
+.loading-download-file {
+  font-size: 13px;
+}
 #dropOverlay {
   position: fixed;
   top: 0;
diff --git a/src/experimental/studio/live.html b/src/experimental/studio/live.html
index 9e101e02..81714b38 100644
--- a/src/experimental/studio/live.html
+++ b/src/experimental/studio/live.html
@@ -10,6 +10,10 @@
       <div class="loading-container">
         <div class="loading-title">Loading model&hellip;</div>
         <div class="loading-subtitle">This may take a moment for large models.</div>
+        <div id="loadingDownload" class="loading-download">
+          <div class="loading-download-label">Downloading asset file</div>
+          <div id="loadingDownloadFile" class="loading-download-file"></div>
+        </div>
       </div>
     </div>
     <div id="dropOverlay">

From ad96acfb16ec397e3719ddedd02b82a90bd1bf81 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Thu, 21 May 2026 05:15:50 -0700
Subject: [PATCH 10/61] Add internal mj_narrowphase for batch processing of
 potential geom pair collisions.

PiperOrigin-RevId: 918988042
Change-Id: I871dcdf42f8105e2c8a57e275a3775ea46c69e65
---
 src/engine/engine_collision_convex.c |  31 +-
 src/engine/engine_collision_convex.h |   3 +
 src/engine/engine_collision_driver.c | 437 ++++++++++++++++++---------
 src/engine/engine_collision_driver.h |  11 -
 4 files changed, 315 insertions(+), 167 deletions(-)

diff --git a/src/engine/engine_collision_convex.c b/src/engine/engine_collision_convex.c
index 2c52a789..30dd4b08 100644
--- a/src/engine/engine_collision_convex.c
+++ b/src/engine/engine_collision_convex.c
@@ -24,7 +24,7 @@
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_gjk.h"
-#include "engine/engine_collision_primitive.h"
+#include "engine/engine_macro.h"
 #include "engine/engine_memory.h"
 #include "engine/engine_inline.h"
 #include "engine/engine_util_blas.h"
@@ -34,6 +34,15 @@
 
 #define mjMINVAL2 (mjMINVAL * mjMINVAL)
 
+// CCD internal buffer used for batched processing; if NULL, stack memory allocated on each
+// mjc_penetration call
+static mjTHREADLOCAL void* ccd_buffer = NULL;
+
+// set CCD internal buffer
+void mjc_setCCDBuffer(void* buffer) {
+  ccd_buffer = buffer;
+}
+
 // ccd prism first dir
 static void prism_firstdir(const void* o1, const void* o2, ccd_vec3_t *vec) {
   ccdVec3Set(vec, 0, 0, 1);
@@ -85,18 +94,23 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
   mjCCDConfig config;
   mjCCDStatus status;
   mjtNum dist;
+  int nwitness = 0;
+  void* buffer = ccd_buffer;
 
   // set config
-  mj_markStack(d);
   config.max_iterations = m->opt.ccd_iterations;
   config.tolerance = m->opt.ccd_tolerance;
   config.max_contacts = ncon;
   config.dist_cutoff = 0;  // no geom distances needed
-  config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
+  if (buffer) {
+    config.buffer = buffer;
+  } else {
+    mj_markStack(d);
+    config.buffer = mj_stackAllocByte(d, mjc_ccdSize(config.max_iterations), sizeof(mjtNum));
+  }
 
   if ((dist = mjc_ccd(&config, &status, obj1, obj2)) < 0) {
-    mj_freeStack(d);
-    int nwitness = status.nx;
+    nwitness = status.nx;
     for (int i = 0; i < nwitness; i++) {
       con[i].dist = margin + dist;
       con[i].pos[0] = 0.5*(status.x1[3*i + 0] + status.x2[3*i + 0]);
@@ -106,10 +120,11 @@ static int mjc_penetration(const mjModel* m, mjData* d, mjCCDObj* obj1, mjCCDObj
       mju_normalize3(con[i].normal);
       mji_zero3(con[i].tangent);
     }
-    return nwitness;
   }
-  mj_freeStack(d);
-  return 0;
+  if (!buffer) {
+    mj_freeStack(d);
+  }
+  return nwitness;
 }
 
 
diff --git a/src/engine/engine_collision_convex.h b/src/engine/engine_collision_convex.h
index 88fda2ea..eae1e72d 100644
--- a/src/engine/engine_collision_convex.h
+++ b/src/engine/engine_collision_convex.h
@@ -124,6 +124,9 @@ int mjc_HFieldElem(const mjModel* m, mjData* d, mjContact* con, int g, int f, in
 // fix contact frame normal
 void mjc_fixNormal(const mjModel* m, const mjData* d, mjPreContact* con, int g1, int g2);
 
+// set CCD internal buffer
+void mjc_setCCDBuffer(void* buffer);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index a598602f..677a457d 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -23,6 +23,7 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_callback.h"
 #include "engine/engine_collision_convex.h"
+#include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_collision_sdf.h"
 #include "engine/engine_core_constraint.h"
@@ -155,6 +156,24 @@ int mj_maxContact(const mjModel* m, int g1, int g2, int has_margin) {
 }
 
 
+// return the margin for a given geom pair
+static inline mjtNum getMargin(const mjModel* m, int g1, int g2, int ipair) {
+  if (ipair >= 0) {
+    return mj_assignMargin(m, m->pair_margin[ipair]);
+  }
+  return mj_assignMargin(m, m->geom_margin[g1] + m->geom_margin[g2]);
+}
+
+
+// return the gap for a given geom pair
+static inline mjtNum getGap(const mjModel* m, int g1, int g2, int ipair) {
+  if (ipair >= 0) {
+    return m->pair_gap[ipair];
+  }
+  return m->geom_gap[g1] + m->geom_gap[g2];
+}
+
+
 // move arena pointer back to the end of the contact array
 static inline void resetArena(mjData* d) {
   d->parena = d->ncon * sizeof(mjContact);
@@ -167,6 +186,16 @@ static inline void resetArena(mjData* d) {
 }
 
 
+// realign the arena pointer to the specified alignment
+static inline size_t alignArena(mjData* d, size_t alignment) {
+  size_t misalignment = d->parena % alignment;
+  if (misalignment) {
+    d->parena += alignment - misalignment;
+  }
+  return d->parena;
+}
+
+
 // plane to geom_center squared distance, g1 is a plane
 static mjtNum planeGeomDist(const mjModel* m, mjData* d, int g1, int g2) {
   mjtNum* mat1 = d->geom_xmat + 9*g1;
@@ -322,6 +351,16 @@ int mj_isElemActive(const mjModel* m, int f, int e) {
 
 //----------------------------- collision detection entry point ------------------------------------
 
+// binary search between two bodyflex trees
+void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
+                    int merged, int startadr, int pairadr);
+
+// compute contacts for a batch of collision pairs contained in a buffer of
+// stride 3 ints (g1, g2, ipair)
+// if buffer is NULL, results are read from arena starting at parena
+void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena);
+
+
 // compare contact pairs by their geom/elem/vert IDs
 static inline int contactcompare(const mjContact* c1, const mjContact* c2, void* context) {
   const mjModel* m = (const mjModel*) context;
@@ -333,7 +372,7 @@ static inline int contactcompare(const mjContact* c1, const mjContact* c2, void*
   int con2_obj2 = c2->geom[1] >= 0 ? c2->geom[1] : (c2->elem[1] >= 0 ? c2->elem[1] : c2->vert[1]);
 
   // for geom:geom, reproduce the order of contacts without mj_collideTree
-  // normally sorted by (g1, g2), but in mj_collideGeoms, g1 and g2 are swapped based on geom_type
+  // normally sorted by (g1, g2), but g1 and g2 are swapped based on geom_type
   // here we undo this swapping for the purpose of sorting - needs to be done for each mjContact
   if (c1->geom[0] >= 0 && c1->geom[1] >= 0 &&
       c2->geom[0] >= 0 && c2->geom[1] >= 0) {
@@ -431,6 +470,72 @@ static void filterFlexContacts(mjData* d, int ncon_before) {
 }
 
 
+// push a candidate collision pair onto the arena
+static void pushPairArena(const mjModel* m, mjData* d, int g1, int g2, int ipair) {
+  // allocate geom pair on the arena
+  int* pair = (int*) mj_arenaAllocByte(d, 3 * sizeof(int), _Alignof(int));
+  if (!pair) {
+    mjERROR("arena too small to allocate geom pair");
+  }
+
+  if (m->geom_type[g1] > m->geom_type[g2]) {
+    pair[0] = g2;
+    pair[1] = g1;
+  } else {
+    pair[0] = g1;
+    pair[1] = g2;
+  }
+  pair[2] = ipair;
+}
+
+
+// filter candidate collision pair; return 0 if pair should be discarded
+static int filterCollisionPair(const mjModel* m, mjData* d, int g1, int g2, int ipair,
+                               int merged, int startadr, int pairadr) {
+  // merged, find matching pair
+  if (merged) {
+    for (int k=startadr; k < pairadr; k++) {
+      if ((m->pair_geom1[k] == g1 && m->pair_geom2[k] == g2) ||
+          (m->pair_geom1[k] == g2 && m->pair_geom2[k] == g1)) {
+        return 0;
+      }
+    }
+  }
+
+  if (ipair >= 0) {
+    if (mjENABLED(mjENBL_SLEEP)) {
+      int b1 = m->geom_bodyid[g1];
+      int b2 = m->geom_bodyid[g2];
+      if (d->body_awake[b1] != mjS_AWAKE && d->body_awake[b2] != mjS_AWAKE) {
+        return 0;
+      }
+    }
+  }
+
+  if (ipair < 0) {
+    if (mjcb_contactfilter) {
+      if (mjcb_contactfilter(m, d, g1, g2)) {
+        return 0;
+      }
+    } else if (filterBitmask(m->geom_contype[g1], m->geom_conaffinity[g1],
+                             m->geom_contype[g2], m->geom_conaffinity[g2])) {
+      return 0;
+    }
+  }
+
+  // bounding sphere filter
+  mjtNum margin = getMargin(m, g1, g2, ipair);
+  mjtNum gap = getGap(m, g1, g2, ipair);
+  if (mj_filterSphere(m, d, g1, g2, margin + gap)) {
+    return 0;
+  }
+
+  // highly unlikely, but check collision function is well-defined
+  int type1 = mjMIN(m->geom_type[g1], m->geom_type[g2]);
+  int type2 = mjMAX(m->geom_type[g1], m->geom_type[g2]);
+  return mjCOLLISIONFUNC[type1][type2] != NULL;
+}
+
 
 // main collision function
 void mj_collision(const mjModel* m, mjData* d) {
@@ -469,7 +574,11 @@ void mj_collision(const mjModel* m, mjData* d) {
   TM_RESTART;
 
   // process bodyflex pairs returned by broadphase, merge with predefined geom pairs
-  int pairadr = 0;
+  int pairadr = 0, ngeompair = 0, g1, g2;
+
+  // align the arena for candidate collision pairs for narrowphase
+  size_t parena = alignArena(d, _Alignof(int));
+
   for (int i=0; i < nbfpair; i++) {
     // reconstruct bodyflex pair ids
     int bf1 = (broadphasepair[i]>>16) & 0xFFFF;
@@ -490,9 +599,11 @@ void mj_collision(const mjModel* m, mjData* d) {
     // test all predefined pairs for which pair_signature <= signature
     for (; pairadr < npair && m->pair_signature[pairadr] <= signature; pairadr++) {
       merged = (m->pair_signature[pairadr] == signature);
-      int g1 = m->pair_geom1[pairadr];
-      int g2 = m->pair_geom2[pairadr];
-      mj_collideGeoms(m, d, pairadr, g1, g2);
+      g1 = m->pair_geom1[pairadr], g2 = m->pair_geom2[pairadr];
+      if (filterCollisionPair(m, d, g1, g2, pairadr, 0, 0, 0)) {
+        pushPairArena(m, d, g1, g2, pairadr);
+        ngeompair++;
+      }
     }
 
     // apply bitmask filtering at the bodyflex level
@@ -524,11 +635,20 @@ void mj_collision(const mjModel* m, mjData* d) {
 
     // process bodyflex pair: two single-geom bodies
     if (isbody1 && isbody2 && m->body_geomnum[bf1] == 1 && m->body_geomnum[bf2] == 1) {
-      mj_collideGeomPair(m, d, geomadr1, geomadr2, merged, startadr, pairadr);
+      if (filterCollisionPair(m, d, geomadr1, geomadr2, -1, merged, startadr, pairadr)) {
+        pushPairArena(m, d, geomadr1, geomadr2, -1);
+        ngeompair++;
+      }
     }
 
     // process bodyflex pair: midphase
     else if (!mjDISABLED(mjDSBL_MIDPHASE) && bvh1 >= 0 && bvh2 >= 0) {
+      // flush geom pairs before calling mj_collideTree as post sorting needs to happen
+      if (ngeompair > 0) {
+        mj_narrowphase(m, d, NULL, ngeompair, parena);
+        ngeompair = 0;
+      }
+
       int ncon_before = d->ncon;
       mj_collideTree(m, d, bf1, bf2, merged, startadr, pairadr);
       int ncon_after = d->ncon;
@@ -547,6 +667,9 @@ void mj_collision(const mjModel* m, mjData* d) {
         contactSort(d->contact + ncon_before, buf, n, (void*)m);
         mj_freeStack(d);
       }
+
+      // realign the arena after adding contacts
+      parena = alignArena(d, _Alignof(int));
     }
 
     // process bodyflex pair: all-to-all
@@ -556,15 +679,23 @@ void mj_collision(const mjModel* m, mjData* d) {
 
       // body : body
       if (isbody1 && isbody2) {
-        for (int g1=geomadr1; g1 < geomadr_end1; g1++) {
-          for (int g2=geomadr2; g2 < geomadr_end2; g2++) {
-            mj_collideGeomPair(m, d, g1, g2, merged, startadr, pairadr);
+        for (g1=geomadr1; g1 < geomadr_end1; g1++) {
+          for (g2=geomadr2; g2 < geomadr_end2; g2++) {
+            if (filterCollisionPair(m, d, g1, g2, -1, merged, startadr, pairadr)) {
+              pushPairArena(m, d, g1, g2, -1);
+              ngeompair++;
+            }
           }
         }
       }
 
       // body : flex
       else if (isbody1) {
+        if (ngeompair > 0) {
+          mj_narrowphase(m, d, NULL, ngeompair, parena);
+          ngeompair = 0;
+        }
+
         int f = bf2 - nbody;
 
         // process body geoms
@@ -599,10 +730,18 @@ void mj_collision(const mjModel* m, mjData* d) {
           }
           filterFlexContacts(d, ncon_before);
         }
+        // realign the arena after adding contacts
+        parena = alignArena(d, _Alignof(int));
       }
 
       // flex : flex
       else {
+        // flush accumulated geompairs before flex:flex processing
+        if (ngeompair > 0) {
+          mj_narrowphase(m, d, NULL, ngeompair, parena);
+          ngeompair = 0;
+        }
+
         int f1 = bf1 - nbody;
         int f2 = bf2 - nbody;
 
@@ -614,6 +753,9 @@ void mj_collision(const mjModel* m, mjData* d) {
           }
         }
         filterFlexContacts(d, ncon_before);
+
+        // realign the arena after adding contacts
+        parena = alignArena(d, _Alignof(int));
       }
     }
   }
@@ -621,9 +763,17 @@ void mj_collision(const mjModel* m, mjData* d) {
 
   // finish merging predefined geom pairs
   for (; pairadr < npair; pairadr++) {
-    int g1 = m->pair_geom1[pairadr];
-    int g2 = m->pair_geom2[pairadr];
-    mj_collideGeoms(m, d, pairadr, g1, g2);
+    g1 = m->pair_geom1[pairadr], g2 = m->pair_geom2[pairadr];
+    if (filterCollisionPair(m, d, g1, g2, pairadr, 0, 0, 0)) {
+      pushPairArena(m, d, g1, g2, pairadr);
+      ngeompair++;
+    }
+  }
+
+  // flush remaining collision pairs
+  if (ngeompair > 0) {
+    mj_narrowphase(m, d, NULL, ngeompair, parena);
+    ngeompair = 0;
   }
 
   // flex self-collisions
@@ -689,25 +839,6 @@ typedef struct  {
   int node2;
 } mjCollisionTree;
 
-
-// checks if the proposed collision pair is already present in pair_geom and calls narrow phase
-void mj_collideGeomPair(const mjModel* m, mjData* d, int g1, int g2, int merged,
-                        int startadr, int pairadr) {
-  // merged, find matching pair
-  if (merged) {
-    for (int k=startadr; k < pairadr; k++) {
-      if ((m->pair_geom1[k] == g1 && m->pair_geom2[k] == g2) ||
-          (m->pair_geom1[k] == g2 && m->pair_geom2[k] == g1)) {
-        return;
-      }
-    }
-  }
-
-  // not merged, always test
-  mj_collideGeoms(m, d, -1, g1, g2);
-}
-
-
 // oriented bounding boxes collision (see Gottschalk et al.)
 int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
                   const mjtNum xpos1[3], const mjtNum xmat1[9],
@@ -806,7 +937,6 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 }
 
 
-// binary search between two bodyflex trees
 void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                     int merged, int startadr, int pairadr) {
   int nbody = m->nbody, nbvhstatic = m->nbvhstatic;
@@ -894,7 +1024,16 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                             d->geom_xpos + 3*nodeid1, d->geom_xmat + 9*nodeid1,
                             d->geom_xpos + 3*nodeid2, d->geom_xmat + 9*nodeid2,
                             margin + gap, NULL, NULL, &initialize)) {
-            mj_collideGeomPair(m, d, nodeid1, nodeid2, merged, startadr, pairadr);
+
+            if (filterCollisionPair(m, d, nodeid1, nodeid2, -1, merged, startadr, pairadr)) {
+              int n1 = nodeid1, n2 = nodeid2;
+              if (m->geom_type[n1] > m->geom_type[n2]) {
+                n1 = nodeid2;
+                n2 = nodeid1;
+              }
+              int pair[3] = {n1, n2, -1};
+              mj_narrowphase(m, d, pair, 1, 0);
+            }
             if (mark_active) {
               d->bvh_active[node1 + bvhadr1] = 1;
               d->bvh_active[node2 + bvhadr2] = 1;
@@ -1675,139 +1814,141 @@ static void mj_makeCapsule(const mjModel* m, mjData* d, int f, const int vid[2],
 }
 
 
-// test two geoms for collision, apply filters, add to contact list
-void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2) {
-  int num, type1, type2, condim;
-  mjtNum margin, gap, friction[5], solref[mjNREF], solimp[mjNIMP];
-  mjtNum solreffriction[mjNREF] = {0};
+// compute contacts for a batch of collision pairs contained in a buffer of
+// stride 3 ints (g1, g2, ipair)
+// if buffer is NULL, results are read from arena starting at parena
+void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena) {
+  int ccd_size = mjc_ccdSize(m->opt.ccd_iterations);
+  mjtNum margin, gap;
 
-  // sleep filtering for explicit pairs
-  if (ipair >= 0) {
-    if (mjENABLED(mjENBL_SLEEP)) {
-      int b1 = m->geom_bodyid[g1];
-      int b2 = m->geom_bodyid[g2];
-      if (d->body_awake[b1] != mjS_AWAKE && d->body_awake[b2] != mjS_AWAKE) {
-        return;
-      }
-    }
-  }
-
-  // order geoms by type
-  if (m->geom_type[g1] > m->geom_type[g2]) {
-    int i = g1;
-    g1 = g2;
-    g2 = i;
-  }
-
-  // copy types and bodies
-  type1 = m->geom_type[g1];
-  type2 = m->geom_type[g2];
-
-  mjfCollision collisionFunc = mjCOLLISIONFUNC[type1][type2];
-
-  // return if no collision function
-  if (!collisionFunc) {
-    return;
-  }
-
-  // apply filters if not predefined pair
-  if (ipair < 0) {
-    // user filter if defined
-    if (mjcb_contactfilter) {
-      if (mjcb_contactfilter(m, d, g1, g2)) {
-        return;
-      }
-    }
-
-    // otherwise built-in filter
-    else if (filterBitmask(m->geom_contype[g1], m->geom_conaffinity[g1],
-                           m->geom_contype[g2], m->geom_conaffinity[g2])) {
-      return;
-    }
-  }
-
-  // set margin: dynamic or pair
-  if (ipair < 0) {
-    margin = mj_assignMargin(m, m->geom_margin[g1] + m->geom_margin[g2]);
+  // set buffer and arena pointer
+  if (!buffer) {
+    buffer = (const int*) ((char*) d->arena + parena);
   } else {
-    margin = mj_assignMargin(m, m->pair_margin[ipair]);
+    parena = d->parena;
   }
 
-  // set gap: dynamic or pair
-  if (ipair < 0) {
-    gap = m->geom_gap[g1] + m->geom_gap[g2];
-  } else {
-    gap = m->pair_gap[ipair];
+  mj_markStack(d);
+
+  // buffer store how many contacts are generated for each pair
+  int* nconbuffer = mj_stackAllocInt(d, npair);
+
+  // buffer for pair data (g1, g2, ipair, index into conbuffer)
+  int* pairbuffer = mj_stackAllocInt(d, 4 * npair);
+  int maxcon = 0;
+  for (int i = 0; i < npair; i++) {
+    int g1 = buffer[3*i + 0];
+    int g2 = buffer[3*i + 1];
+    int ipair = buffer[3*i + 2];
+
+    pairbuffer[4*i + 0] = g1;
+    pairbuffer[4*i + 1] = g2;
+    pairbuffer[4*i + 2] = ipair;
+    pairbuffer[4*i + 3] = maxcon;
+    margin = getMargin(m, g1, g2, ipair);
+    gap = getGap(m, g1, g2, ipair);
+    maxcon += mj_maxContact(m, g1, g2, margin + gap > 0);
   }
 
-  // bounding sphere filter
-  if (mj_filterSphere(m, d, g1, g2, margin + gap)) {
+  // buffer for precontact data
+  mjPreContact* conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
+
+  // buffer data has been copied to metadata on the stack;
+  // reclaim arena space so contacts can overwrite the buffer region
+  d->parena = parena;
+
+  // set buffer for nativeccd
+  mj_markStack(d);
+  mjc_setCCDBuffer(mj_stackAllocByte(d, ccd_size, sizeof(mjtNum)));
+
+  for (int i = 0; i < npair; i++) {
+    int g1 = pairbuffer[4*i + 0];
+    int g2 = pairbuffer[4*i + 1];
+    int ipair = pairbuffer[4*i + 2];
+    int idx = pairbuffer[4*i + 3];
+    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
+    margin = getMargin(m, g1, g2, ipair);
+    gap = getGap(m, g1, g2, ipair);
+    nconbuffer[i] = collision_func(m, d, conbuffer + idx, g1, g2, margin + gap);
+
+    // SHOULD NOT OCCUR
+    int expected_max = (i + 1 < npair ? pairbuffer[4*(i+1) + 3] : maxcon) - idx;
+    if (nconbuffer[i] > expected_max) {
+      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
+              "expected at most %d from mj_maxContact", nconbuffer[i], g1, g2, expected_max);
+    }
+  }
+
+  // set nativeccd buffer back to NULL
+  mjc_setCCDBuffer(NULL);
+  mj_freeStack(d);
+
+  int ncon = 0;
+  for (int i = 0; i < npair; i++) {
+    ncon += nconbuffer[i];
+  }
+
+  if (ncon == 0) {
+    mj_freeStack(d);
     return;
   }
 
-
-
-  // call collision detector to generate contacts
-  mjPreContact precon[mjMAXCONPAIR];
-  if (!(num = collisionFunc(m, d, precon, g1, g2, margin + gap))) {
-    return;
-  }
-
-  // check number of contacts, SHOULD NOT OCCUR
-  if (num > mjMAXCONPAIR) {
-    mjERROR("too many contacts returned by collision function");
-  }
-
-  mjContact* con = (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * num, _Alignof(mjContact));
+  // try allocate contact buffer in arena
+  mjContact* con =
+    (mjContact*) mj_arenaAllocByte(d, sizeof(mjContact) * ncon, _Alignof(mjContact));
   if (!con) {
     mj_warning(d, mjWARN_CONTACTFULL, d->ncon);
+    mj_freeStack(d);
     return;
   }
+  d->ncon += ncon;
 
-  // set condim, solref, solimp, friction: dynamic
-  if (ipair < 0) {
-    mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
-  }
+  // fill in contact data
+  int conpos = 0;
+  for (int i = 0; i < npair; i++) {
+    if (!(ncon = nconbuffer[i]))
+      continue;
 
-  // set condim, solref, solimp, friction: pair
-  else {
-    condim = m->pair_dim[ipair];
-    mju_copy(solref, m->pair_solref+mjNREF*ipair, mjNREF);
-    mju_copy(solimp, m->pair_solimp+mjNIMP*ipair, mjNIMP);
-    mju_copy(friction, m->pair_friction+5*ipair, 5);
+    int condim;
+    mjtNum friction[5], solref[mjNREF], solimp[mjNIMP];
+    mjtNum solreffriction[mjNREF] = {0};
+    int g1 = pairbuffer[4*i + 0];
+    int g2 = pairbuffer[4*i + 1];
+    int ipair = pairbuffer[4*i + 2];
 
-    // reference, friction directions
-    if (m->pair_solreffriction[mjNREF*ipair] || m->pair_solreffriction[mjNREF*ipair + 1]) {
-      mju_copy(solreffriction, m->pair_solreffriction+mjNREF*ipair, mjNREF);
+    if (ipair >= 0) {
+      condim = m->pair_dim[ipair];
+      mju_copy(solref, m->pair_solref+mjNREF*ipair, mjNREF);
+      mju_copy(solimp, m->pair_solimp+mjNIMP*ipair, mjNIMP);
+      mju_copy(friction, m->pair_friction+5*ipair, 5);
+      if (m->pair_solreffriction[mjNREF*ipair] || m->pair_solreffriction[mjNREF*ipair + 1]) {
+        mju_copy(solreffriction, m->pair_solreffriction+mjNREF*ipair, mjNREF);
+      }
+    } else {
+      mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
     }
+
+    mjPreContact* bc = conbuffer + pairbuffer[4*i + 3];
+    margin = getMargin(m, g1, g2, ipair);
+    for (int j=0; j < ncon; j++) {
+      mjContact* c = con + conpos + j;
+      c->dist = bc[j].dist;
+      mji_copy3(c->pos, bc[j].pos);
+      mji_copy3(c->frame, bc[j].normal);
+      mji_copy3(c->frame + 3, bc[j].tangent);
+      c->geom[0] = g1;
+      c->geom[1] = g2;
+      c->flex[0] = -1;
+      c->flex[1] = -1;
+      c->elem[0] = -1;
+      c->elem[1] = -1;
+      c->vert[0] = -1;
+      c->vert[1] = -1;
+      mj_setContact(m, c, condim, margin, solref, solreffriction, solimp, friction);
+    }
+    conpos += ncon;
   }
-
-  // add contacts returned by collision detector
-  for (int i=0; i < num; i++) {
-    // set contact parameters
-    con[i].dist = precon[i].dist;
-    mji_copy3(con[i].pos, precon[i].pos);
-    mji_copy3(con[i].frame + 0, precon[i].normal);
-    mji_copy3(con[i].frame + 3, precon[i].tangent);
-
-    con[i].geom[0] = g1;
-    con[i].geom[1] = g2;
-    con[i].flex[0] = -1;
-    con[i].flex[1] = -1;
-    con[i].elem[0] = -1;
-    con[i].elem[1] = -1;
-    con[i].vert[0] = -1;
-    con[i].vert[1] = -1;
-
-    // set remaining contact parameters
-    mj_setContact(m, con + i, condim, margin, solref, solreffriction, solimp, friction);
-  }
-
-  // add to ncon
-  d->ncon += num;
-
-  // move arena pointer back to the end of the contact array
-  resetArena(d);
+  mj_freeStack(d);
 }
 
 
diff --git a/src/engine/engine_collision_driver.h b/src/engine/engine_collision_driver.h
index e231eaaf..0e2b9edd 100644
--- a/src/engine/engine_collision_driver.h
+++ b/src/engine/engine_collision_driver.h
@@ -44,20 +44,9 @@ MJAPI int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 // is element active (for collisions)
 MJAPI int mj_isElemActive(const mjModel* m, int f, int e);
 
-// checks if pair is already present in pair_geom and calls narrow phase
-void mj_collideGeomPair(const mjModel* m, mjData* d, int g1, int g2, int merged,
-                        int startadr, int pairadr);
-
-// binary search between two bodyflex trees
-void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
-                    int merged, int startadr, int pairadr);
-
 // broad phase collision detection; return list of bodyflex pairs
 int mj_broadphase(const mjModel* m, mjData* d, int* bfpair, int maxpair);
 
-// test two geoms for collision, apply filters, add to contact list
-void mj_collideGeoms(const mjModel* m, mjData* d, int ipair, int g1, int g2);
-
 // test a plane geom and a flex for collision, add to contact list
 void mj_collidePlaneFlex(const mjModel* m, mjData* d, int g, int f);
 

From 2810edd27ab70d86d297b91430a1d0903745abe3 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 21 May 2026 06:42:56 -0700
Subject: [PATCH 11/61] Rename `mjtnum.h` to `mjtype.h`, move enum types to
 `mjtype.h`

PiperOrigin-RevId: 919020039
Change-Id: I441295c0baa0b7456f78239fbd5478d32ffc07c1
---
 CMakeLists.txt                           |   2 +-
 doc/APIreference/APIglobals.rst          |   2 +-
 doc/APIreference/APItypes.rst            |   4 +-
 doc/changelog.rst                        |   6 +-
 doc/includes/references.h                | 896 +++++++++++------------
 doc/programming/index.rst                |   5 +-
 include/mujoco/mjdata.h                  |  99 +--
 include/mujoco/mjmodel.h                 | 431 +----------
 include/mujoco/mjplugin.h                |   2 +-
 include/mujoco/mjspec.h                  |   2 +-
 include/mujoco/mjtnum.h                  |  46 --
 include/mujoco/mjtype.h                  | 574 +++++++++++++++
 include/mujoco/mjvisualize.h             |   2 +-
 include/mujoco/mujoco.h                  |   2 +-
 plugin/actuator/pid.h                    |   2 +-
 plugin/elasticity/cable.cc               |   2 +-
 plugin/elasticity/cable.h                |   2 +-
 plugin/sdf/bolt.cc                       |   2 +-
 plugin/sdf/bolt.h                        |   2 +-
 plugin/sdf/bowl.cc                       |   2 +-
 plugin/sdf/bowl.h                        |   2 +-
 plugin/sdf/gear.cc                       |   2 +-
 plugin/sdf/gear.h                        |   2 +-
 plugin/sdf/nut.cc                        |   2 +-
 plugin/sdf/nut.h                         |   2 +-
 plugin/sdf/torus.cc                      |   2 +-
 plugin/sdf/torus.h                       |   2 +-
 plugin/sensor/touch_grid.cc              |   2 +-
 plugin/sensor/touch_grid.h               |   2 +-
 python/mujoco/constants.cc               |   2 +-
 python/mujoco/introspect/enums.py        |  20 +-
 python/mujoco/serialization.h            |   2 +-
 src/engine/engine_collision_convex.h     |   2 +-
 src/engine/engine_collision_driver.h     |   2 +-
 src/engine/engine_collision_gjk.c        |   2 +-
 src/engine/engine_collision_gjk.h        |   2 +-
 src/engine/engine_collision_primitive.h  |   2 +-
 src/engine/engine_collision_sdf.c        |   2 +-
 src/engine/engine_collision_sdf.h        |   2 +-
 src/engine/engine_core_constraint.h      |   2 +-
 src/engine/engine_core_smooth.h          |   2 +-
 src/engine/engine_core_util.h            |   2 +-
 src/engine/engine_forward.h              |   2 +-
 src/engine/engine_init.c                 |   2 +-
 src/engine/engine_inline.h               |   2 +-
 src/engine/engine_memory.h               |   2 +-
 src/engine/engine_print.c                |   2 +-
 src/engine/engine_support.h              |   2 +-
 src/engine/engine_util_blas.c            |   2 +-
 src/engine/engine_util_blas.h            |   2 +-
 src/engine/engine_util_blas_avx.h        |   2 +-
 src/engine/engine_util_misc.h            |   2 +-
 src/engine/engine_util_solve.h           |   2 +-
 src/engine/engine_util_sparse.c          |   2 +-
 src/engine/engine_util_sparse.h          |   2 +-
 src/engine/engine_util_sparse_avx.h      |   2 +-
 src/engine/engine_util_spatial.h         |   2 +-
 src/user/user_api.h                      |   2 +-
 src/user/user_composite.cc               |   2 +-
 src/user/user_flexcomp.cc                |   2 +-
 src/user/user_mesh.cc                    |   2 +-
 src/user/user_model.cc                   |   2 +-
 src/user/user_model.h                    |   2 +-
 src/user/user_objects.h                  |   2 +-
 src/xml/xml_native_reader.cc             |   2 +-
 test/engine/engine_collision_gjk_test.cc |   2 +-
 test/engine/engine_collision_sdf_test.cc |   2 +-
 test/engine/engine_forward_test.cc       |   2 +-
 test/engine/engine_ray_test.cc           |   2 +-
 test/engine/engine_sensor_test.cc        |   2 +-
 test/engine/engine_util_blas_test.cc     |   2 +-
 test/engine/engine_util_spatial_test.cc  |   2 +-
 test/header_test.cc                      |   2 +-
 test/user/user_mesh_test.cc              |   2 +-
 test/user/user_objects_test.cc           |   2 +-
 test/xml/xml_native_writer_test.cc       |   2 +-
 unity/Runtime/Bindings/MjBindings.cs     |  12 +-
 77 files changed, 1117 insertions(+), 1110 deletions(-)
 delete mode 100644 include/mujoco/mjtnum.h
 create mode 100644 include/mujoco/mjtype.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2ed0692e..86fc48a1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -117,7 +117,7 @@ set(MUJOCO_HEADERS
     include/mujoco/mjsan.h
     include/mujoco/mjspec.h
     include/mujoco/mjthread.h
-    include/mujoco/mjtnum.h
+    include/mujoco/mjtype.h
     include/mujoco/mjui.h
     include/mujoco/mjvisualize.h
     include/mujoco/mjxmacro.h
diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst
index ba2d63ae..d823553a 100644
--- a/doc/APIreference/APIglobals.rst
+++ b/doc/APIreference/APIglobals.rst
@@ -415,7 +415,7 @@ indicated otherwise.
      - 1E-15
      - The minimal value allowed in any denominator, and in general any mathematical operation where 0 is not allowed.
        In almost all cases, MuJoCo silently clamps smaller values to mjMINVAL.
-       Defined in `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`_.
+       Defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`_.
    * - ``mjPI``
      - :math:`\pi`
      - The value of :math:`\pi`. This is used in various trigonometric functions, and also for conversion from degrees
diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst
index dd34f411..f7779b3e 100644
--- a/doc/APIreference/APItypes.rst
+++ b/doc/APIreference/APItypes.rst
@@ -46,7 +46,7 @@ MuJoCo defines a large number of types:
 Primitive types
 ---------------
 
-The two types below are defined in `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__.
+The three types below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
 
 
 .. _mjtNum:
@@ -63,7 +63,7 @@ double-precision version will always be available. Thus it is safe to write user
 However, our preference is to write code that works with either single or double precision. To this end we provide math
 utility functions that are always defined with the correct floating-point type.
 
-Note that changing ``mjUSESINGLE`` in ``mjtnum.h`` will not change how the library was compiled, and instead will
+Note that changing ``mjUSESINGLE`` in ``mjtype.h`` will not change how the library was compiled, and instead will
 result in numerous link errors. In general, the header files distributed with precompiled MuJoCo should never be
 changed by the user.
 
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 29b713cc..5cd5f9c1 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -73,6 +73,10 @@ General
    - The :ref:`mjfCollision` functions now populate the :ref:`mjPreContact` struct instead of the :ref:`mjContact`
      struct. The :ref:`mjPreContact` only contains the necessary fields needed for the narrowphase collision detection.
 
+   - The header file ``mjtnum.h`` was renamed to
+     `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>` and now includes all
+     enum type definitions.
+
    - MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
      ``naconmax``.
 
@@ -1194,7 +1198,7 @@ General
 10. Added :ref:`mj_setKeyframe` for saving the current state into a model keyframe.
 11. Added support for ``ball`` joints in the URDF parser ("spherical" in URDF).
 12. Replaced ``mjUSEDOUBLE`` which was previously hard-coded in
-    `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
+    `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/3577e2cf8bf841475b489aefff52276a39f24d51/include/mjtnum.h>`__
     with the build-time flag ``mjUSESINGLE``. If this symbol is not defined, MuJoCo will use double-precision floating
     point, as usual. If ``mjUSESINGLE`` is defined, MuJoCo will use single-precision floating point. See :ref:`mjtNum`.
 
diff --git a/doc/includes/references.h b/doc/includes/references.h
index 7d929ee3..7f6371a2 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -16,91 +16,6 @@
 // Error: C reference not found
 // NOLINTBEGIN
 
-typedef enum mjtState_ {            // state elements
-  mjSTATE_TIME           = 1<<0,    // time
-  mjSTATE_QPOS           = 1<<1,    // position
-  mjSTATE_QVEL           = 1<<2,    // velocity
-  mjSTATE_ACT            = 1<<3,    // actuator activation
-  mjSTATE_HISTORY        = 1<<4,    // history buffers (control, sensor)
-  mjSTATE_WARMSTART      = 1<<5,    // acceleration used for warmstart
-  mjSTATE_CTRL           = 1<<6,    // control
-  mjSTATE_QFRC_APPLIED   = 1<<7,    // applied generalized force
-  mjSTATE_XFRC_APPLIED   = 1<<8,    // applied Cartesian force/torque
-  mjSTATE_EQ_ACTIVE      = 1<<9,    // enable/disable constraints
-  mjSTATE_MOCAP_POS      = 1<<10,   // positions of mocap bodies
-  mjSTATE_MOCAP_QUAT     = 1<<11,   // orientations of mocap bodies
-  mjSTATE_USERDATA       = 1<<12,   // user data
-  mjSTATE_PLUGIN         = 1<<13,   // plugin state
-
-  mjNSTATE               = 14,      // number of state elements
-
-  // convenience values for commonly used state specifications
-  mjSTATE_PHYSICS        = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
-  mjSTATE_FULLPHYSICS    = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
-  mjSTATE_USER           = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
-                          mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
-                          mjSTATE_USERDATA,
-  mjSTATE_INTEGRATION    = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
-} mjtState;
-typedef enum mjtConstraint_ {       // type of constraint
-  mjCNSTR_EQUALITY       = 0,       // equality constraint
-  mjCNSTR_FRICTION_DOF,             // dof friction
-  mjCNSTR_FRICTION_TENDON,          // tendon friction
-  mjCNSTR_LIMIT_JOINT,              // joint limit
-  mjCNSTR_LIMIT_TENDON,             // tendon limit
-  mjCNSTR_CONTACT_FRICTIONLESS,     // frictionless contact
-  mjCNSTR_CONTACT_PYRAMIDAL,        // frictional contact, pyramidal friction cone
-  mjCNSTR_CONTACT_ELLIPTIC          // frictional contact, elliptic friction cone
-} mjtConstraint;
-typedef enum mjtConstraintState_ {  // constraint state
-  mjCNSTRSTATE_SATISFIED = 0,       // constraint satisfied, zero cost (limit, contact)
-  mjCNSTRSTATE_QUADRATIC,           // quadratic cost (equality, friction, limit, contact)
-  mjCNSTRSTATE_LINEARNEG,           // linear cost, negative side (friction)
-  mjCNSTRSTATE_LINEARPOS,           // linear cost, positive side (friction)
-  mjCNSTRSTATE_CONE                 // squared distance to cone cost (elliptic contact)
-} mjtConstraintState;
-typedef enum mjtWarning_ {          // warning types
-  mjWARN_INERTIA         = 0,       // (near) singular inertia matrix
-  mjWARN_CONTACTFULL,               // too many contacts in contact list
-  mjWARN_CNSTRFULL,                 // too many constraints
-  mjWARN_BADQPOS,                   // bad number in qpos
-  mjWARN_BADQVEL,                   // bad number in qvel
-  mjWARN_BADQACC,                   // bad number in qacc
-  mjWARN_BADCTRL,                   // bad number in ctrl
-
-  mjNWARNING                        // number of warnings
-} mjtWarning;
-typedef enum mjtTimer_ {            // internal timers
-  // main api
-  mjTIMER_STEP           = 0,       // step
-  mjTIMER_FORWARD,                  // forward
-  mjTIMER_INVERSE,                  // inverse
-
-  // breakdown of step/forward
-  mjTIMER_POSITION,                 // fwdPosition
-  mjTIMER_VELOCITY,                 // fwdVelocity
-  mjTIMER_ACTUATION,                // fwdActuation
-  mjTIMER_CONSTRAINT,               // fwdConstraint
-  mjTIMER_ADVANCE,                  // mj_Euler, mj_implicit
-
-  // breakdown of fwdPosition
-  mjTIMER_POS_KINEMATICS,           // kinematics, com, tendon, transmission
-  mjTIMER_POS_INERTIA,              // inertia computations
-  mjTIMER_POS_COLLISION,            // collision detection
-  mjTIMER_POS_MAKE,                 // make constraints
-  mjTIMER_POS_PROJECT,              // project constraints
-
-  // breakdown of mj_collision
-  mjTIMER_COL_BROAD,                // broadphase
-  mjTIMER_COL_NARROW,               // narrowphase
-
-  mjNTIMER                          // number of timers
-} mjtTimer;
-typedef enum mjtSleepState_ {       // sleep state of an object
-  mjS_STATIC = -1,                  // object is static
-  mjS_ASLEEP = 0,                   // object is asleep
-  mjS_AWAKE  = 1                    // object is awake
-} mjtSleepState;
 struct mjPreContact_ {             // contact parameters set by narrowphase collision functions
   mjtNum dist;
   mjtNum pos[3];
@@ -486,369 +401,6 @@ struct mjData_ {
   uint64_t  signature;       // also held by the mjSpec that compiled the model
 };
 typedef struct mjData_ mjData;
-typedef enum mjtDisableBit_ {     // disable default feature bitflags
-  mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
-  mjDSBL_EQUALITY     = 1<<1,     // equality constraints
-  mjDSBL_FRICTIONLOSS = 1<<2,     // joint and tendon frictionloss constraints
-  mjDSBL_LIMIT        = 1<<3,     // joint and tendon limit constraints
-  mjDSBL_CONTACT      = 1<<4,     // contact constraints
-  mjDSBL_SPRING       = 1<<5,     // passive spring forces
-  mjDSBL_DAMPER       = 1<<6,     // passive damping forces
-  mjDSBL_GRAVITY      = 1<<7,     // gravitational forces
-  mjDSBL_CLAMPCTRL    = 1<<8,     // clamp control to specified range
-  mjDSBL_WARMSTART    = 1<<9,     // warmstart constraint solver
-  mjDSBL_FILTERPARENT = 1<<10,    // remove collisions with parent body
-  mjDSBL_ACTUATION    = 1<<11,    // apply actuation forces
-  mjDSBL_REFSAFE      = 1<<12,    // integrator safety: make ref[0]>=2*timestep
-  mjDSBL_SENSOR       = 1<<13,    // sensors
-  mjDSBL_MIDPHASE     = 1<<14,    // mid-phase collision filtering
-  mjDSBL_EULERDAMP    = 1<<15,    // implicit integration of joint damping in Euler integrator
-  mjDSBL_AUTORESET    = 1<<16,    // automatic reset when numerical issues are detected
-  mjDSBL_NATIVECCD    = 1<<17,    // native convex collision detection
-  mjDSBL_ISLAND       = 1<<18,    // constraint island discovery
-  mjDSBL_MULTICCD     = 1<<19,    // multiple CCD contact points
-
-  mjNDISABLE          = 20        // number of disable flags
-} mjtDisableBit;
-typedef enum mjtEnableBit_ {      // enable optional feature bitflags
-  mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
-  mjENBL_ENERGY       = 1<<1,     // energy computation
-  mjENBL_FWDINV       = 1<<2,     // record solver statistics
-  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
-  mjENBL_SLEEP        = 1<<4,     // sleeping
-  mjENBL_DIAGEXACT    = 1<<5,     // exact diagonal of constraint inertia
-
-  mjNENABLE           = 6         // number of enable flags
-} mjtEnableBit;
-typedef enum mjtJoint_ {          // type of degree of freedom
-  mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)
-  mjJNT_BALL,                     // orientation (quat) relative to parent        (4)
-  mjJNT_SLIDE,                    // sliding distance along body-fixed axis       (1)
-  mjJNT_HINGE                     // rotation angle (rad) around body-fixed axis  (1)
-} mjtJoint;
-typedef enum mjtGeom_ {           // type of geometric shape
-  // regular geom types
-  mjGEOM_PLANE        = 0,        // plane
-  mjGEOM_HFIELD,                  // height field
-  mjGEOM_SPHERE,                  // sphere
-  mjGEOM_CAPSULE,                 // capsule
-  mjGEOM_ELLIPSOID,               // ellipsoid
-  mjGEOM_CYLINDER,                // cylinder
-  mjGEOM_BOX,                     // box
-  mjGEOM_MESH,                    // mesh
-  mjGEOM_SDF,                     // signed distance field
-
-  mjNGEOMTYPES,                   // number of regular geom types
-
-  // rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
-  mjGEOM_ARROW        = 100,      // arrow
-  mjGEOM_ARROW1,                  // arrow without wedges
-  mjGEOM_ARROW2,                  // arrow in both directions
-  mjGEOM_LINE,                    // line
-  mjGEOM_LINEBOX,                 // box with line edges
-  mjGEOM_FLEX,                    // flex
-  mjGEOM_SKIN,                    // skin
-  mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle
-
-  mjGEOM_NONE         = 1001      // missing geom type
-} mjtGeom;
-typedef enum mjtProjection_ {     // type of camera projection
-  mjPROJ_PERSPECTIVE  = 0,        // perspective
-  mjPROJ_ORTHOGRAPHIC             // orthographic
-} mjtProjection;
-typedef enum mjtCamLight_ {       // tracking mode for camera and light
-  mjCAMLIGHT_FIXED    = 0,        // pos and rot fixed in body
-  mjCAMLIGHT_TRACK,               // pos tracks body, rot fixed in global
-  mjCAMLIGHT_TRACKCOM,            // pos tracks subtree com, rot fixed in body
-  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
-  mjTEXTURE_SKYBOX                // cube texture used as skybox
-} mjtTexture;
-typedef enum mjtTextureRole_ {    // role of texture map in rendering
-  mjTEXROLE_USER      = 0,        // unspecified
-  mjTEXROLE_RGB,                  // base color (albedo)
-  mjTEXROLE_OCCLUSION,            // ambient occlusion
-  mjTEXROLE_ROUGHNESS,            // roughness
-  mjTEXROLE_METALLIC,             // metallic
-  mjTEXROLE_NORMAL,               // normal (bump) map
-  mjTEXROLE_OPACITY,              // opacity
-  mjTEXROLE_EMISSIVE,             // light emission
-  mjTEXROLE_RGBA,                 // base color, opacity
-  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
-  mjINT_IMPLICIT,                 // implicit in velocity
-  mjINT_IMPLICITFAST              // implicit in velocity, no rne derivative
-} mjtIntegrator;
-typedef enum mjtCone_ {           // type of friction cone
-  mjCONE_PYRAMIDAL     = 0,       // pyramidal
-  mjCONE_ELLIPTIC                 // elliptic
-} mjtCone;
-typedef enum mjtJacobian_ {       // type of constraint Jacobian
-  mjJAC_DENSE          = 0,       // dense
-  mjJAC_SPARSE,                   // sparse
-  mjJAC_AUTO                      // dense if nv<60, sparse otherwise
-} mjtJacobian;
-typedef enum mjtSolver_ {         // constraint solver algorithm
-  mjSOL_PGS            = 0,       // PGS    (dual)
-  mjSOL_CG,                       // CG     (primal)
-  mjSOL_NEWTON                    // Newton (primal)
-} mjtSolver;
-typedef enum mjtEq_ {             // type of equality constraint
-  mjEQ_CONNECT        = 0,        // connect two bodies at a point (ball joint)
-  mjEQ_WELD,                      // fix relative position and orientation of two bodies
-  mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
-  mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
-  mjEQ_FLEX,                      // fix all edge lengths of a flex
-  mjEQ_FLEXVERT,                  // fix all vertex lengths of a flex
-  mjEQ_FLEXSTRAIN,                // constrain strain of a trilinear/quadratic flex (B-bar)
-  mjEQ_DISTANCE                   // unsupported, will cause an error if used
-} mjtEq;
-typedef enum mjtWrap_ {           // type of tendon wrap object
-  mjWRAP_NONE         = 0,        // null object
-  mjWRAP_JOINT,                   // constant moment arm
-  mjWRAP_PULLEY,                  // pulley used to split tendon
-  mjWRAP_SITE,                    // pass through site
-  mjWRAP_SPHERE,                  // wrap around sphere
-  mjWRAP_CYLINDER                 // wrap around (infinite) cylinder
-} mjtWrap;
-typedef enum mjtTrn_ {            // type of actuator transmission
-  mjTRN_JOINT         = 0,        // force on joint
-  mjTRN_JOINTINPARENT,            // force on joint, expressed in parent frame
-  mjTRN_SLIDERCRANK,              // force via slider-crank linkage
-  mjTRN_TENDON,                   // force on tendon
-  mjTRN_SITE,                     // force on site
-  mjTRN_BODY,                     // adhesion force on a body's geoms
-
-  mjTRN_UNDEFINED     = 1000      // undefined transmission type
-} mjtTrn;
-typedef enum mjtDyn_ {            // type of actuator dynamics
-  mjDYN_NONE          = 0,        // no internal dynamics; ctrl specifies force
-  mjDYN_INTEGRATOR,               // integrator: da/dt = u
-  mjDYN_FILTER,                   // linear filter: da/dt = (u-a) / tau
-  mjDYN_FILTEREXACT,              // linear filter: da/dt = (u-a) / tau, with exact integration
-  mjDYN_MUSCLE,                   // piecewise linear filter with two time constants
-  mjDYN_DCMOTOR,                  // DC motor electrical dynamics
-  mjDYN_USER                      // user-defined dynamics type
-} mjtDyn;
-typedef enum mjtGain_ {           // type of actuator gain
-  mjGAIN_FIXED        = 0,        // fixed gain
-  mjGAIN_AFFINE,                  // const + kp*length + kv*velocity
-  mjGAIN_MUSCLE,                  // muscle FLV curve computed by mju_muscleGain()
-  mjGAIN_DCMOTOR,                 // DC motor gain: K or K/R
-  mjGAIN_USER                     // user-defined gain type
-} mjtGain;
-typedef enum mjtBias_ {           // type of actuator bias
-  mjBIAS_NONE         = 0,        // no bias
-  mjBIAS_AFFINE,                  // const + kp*length + kv*velocity
-  mjBIAS_MUSCLE,                  // muscle passive force computed by mju_muscleBias()
-  mjBIAS_DCMOTOR,                 // DC motor bias: back-EMF, cogging, LuGre friction
-  mjBIAS_USER                     // user-defined bias type
-} mjtBias;
-typedef enum mjtObj_ {            // type of MujoCo object
-  mjOBJ_UNKNOWN       = 0,        // unknown object type
-  mjOBJ_BODY,                     // body
-  mjOBJ_XBODY,                    // body, used to access regular frame instead of i-frame
-  mjOBJ_JOINT,                    // joint
-  mjOBJ_DOF,                      // dof
-  mjOBJ_GEOM,                     // geom
-  mjOBJ_SITE,                     // site
-  mjOBJ_CAMERA,                   // camera
-  mjOBJ_LIGHT,                    // light
-  mjOBJ_FLEX,                     // flex
-  mjOBJ_MESH,                     // mesh
-  mjOBJ_SKIN,                     // skin
-  mjOBJ_HFIELD,                   // heightfield
-  mjOBJ_TEXTURE,                  // texture
-  mjOBJ_MATERIAL,                 // material for rendering
-  mjOBJ_PAIR,                     // geom pair to include
-  mjOBJ_EXCLUDE,                  // body pair to exclude
-  mjOBJ_EQUALITY,                 // equality constraint
-  mjOBJ_TENDON,                   // tendon
-  mjOBJ_ACTUATOR,                 // actuator
-  mjOBJ_SENSOR,                   // sensor
-  mjOBJ_NUMERIC,                  // numeric
-  mjOBJ_TEXT,                     // text
-  mjOBJ_TUPLE,                    // tuple
-  mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN,                   // plugin instance
-
-  mjNOBJECT,                      // number of object types
-
-  // meta elements, do not appear in mjModel
-  mjOBJ_FRAME         = 100,      // frame
-  mjOBJ_DEFAULT,                  // default
-  mjOBJ_MODEL                     // entire model
-} mjtObj;
-typedef enum mjtSensor_ {         // type of sensor
-  // common robotic sensors, attached to a site
-  mjSENS_TOUCH        = 0,        // scalar contact normal forces summed over sensor zone
-  mjSENS_ACCELEROMETER,           // 3D linear acceleration, in local frame
-  mjSENS_VELOCIMETER,             // 3D linear velocity, in local frame
-  mjSENS_GYRO,                    // 3D angular velocity, in local frame
-  mjSENS_FORCE,                   // 3D force between site's body and its parent body
-  mjSENS_TORQUE,                  // 3D torque between site's body and its parent body
-  mjSENS_MAGNETOMETER,            // 3D magnetometer
-  mjSENS_RANGEFINDER,             // scalar distance to nearest geom along z-axis
-  mjSENS_CAMPROJECTION,           // pixel coordinates of a site in the camera image
-
-  // sensors related to scalar joints, tendons, actuators
-  mjSENS_JOINTPOS,                // scalar joint position (hinge and slide only)
-  mjSENS_JOINTVEL,                // scalar joint velocity (hinge and slide only)
-  mjSENS_TENDONPOS,               // scalar tendon position
-  mjSENS_TENDONVEL,               // scalar tendon velocity
-  mjSENS_ACTUATORPOS,             // scalar actuator position
-  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
-  mjSENS_BALLANGVEL,              // 3D ball joint angular velocity
-
-  // joint and tendon limit sensors, in constraint space
-  mjSENS_JOINTLIMITPOS,           // joint limit distance-margin
-  mjSENS_JOINTLIMITVEL,           // joint limit velocity
-  mjSENS_JOINTLIMITFRC,           // joint limit force
-  mjSENS_TENDONLIMITPOS,          // tendon limit distance-margin
-  mjSENS_TENDONLIMITVEL,          // tendon limit velocity
-  mjSENS_TENDONLIMITFRC,          // tendon limit force
-
-  // sensors attached to an object with spatial frame: (x)body, geom, site, camera
-  mjSENS_FRAMEPOS,                // 3D position
-  mjSENS_FRAMEQUAT,               // 4D unit quaternion orientation
-  mjSENS_FRAMEXAXIS,              // 3D unit vector: x-axis of object's frame
-  mjSENS_FRAMEYAXIS,              // 3D unit vector: y-axis of object's frame
-  mjSENS_FRAMEZAXIS,              // 3D unit vector: z-axis of object's frame
-  mjSENS_FRAMELINVEL,             // 3D linear velocity
-  mjSENS_FRAMEANGVEL,             // 3D angular velocity
-  mjSENS_FRAMELINACC,             // 3D linear acceleration
-  mjSENS_FRAMEANGACC,             // 3D angular acceleration
-
-  // sensors related to kinematic subtrees; attached to a body (which is the subtree root)
-  mjSENS_SUBTREECOM,              // 3D center of mass of subtree
-  mjSENS_SUBTREELINVEL,           // 3D linear velocity of subtree
-  mjSENS_SUBTREEANGMOM,           // 3D angular momentum of subtree
-
-  // sensors of geometric relationships
-  mjSENS_INSIDESITE,              // 1 if object is inside a site, 0 otherwise
-  mjSENS_GEOMDIST,                // signed distance between two geoms
-  mjSENS_GEOMNORMAL,              // normal direction between two geoms
-  mjSENS_GEOMFROMTO,              // segment between two geoms
-
-  // sensors for reporting contacts which occurred during the simulation
-  mjSENS_CONTACT,                 // contacts which occurred during the simulation
-
-  // global sensors
-  mjSENS_E_POTENTIAL,             // potential energy
-  mjSENS_E_KINETIC,               // kinetic energy
-  mjSENS_CLOCK,                   // simulation time
-
-  // sensors related to SDFs
-  mjSENS_TACTILE,                 // tactile sensor
-
-  // plugin-controlled sensors
-  mjSENS_PLUGIN,                  // plugin-controlled
-
-  // user-defined sensor
-  mjSENS_USER                     // sensor data provided by mjcb_sensor callback
-} mjtSensor;
-typedef enum mjtStage_ {          // computation stage
-  mjSTAGE_NONE        = 0,        // no computations
-  mjSTAGE_POS,                    // position-dependent computations
-  mjSTAGE_VEL,                    // velocity-dependent computations
-  mjSTAGE_ACC                     // acceleration/force-dependent computations
-} mjtStage;
-typedef enum mjtDataType_ {       // data type for sensors
-  mjDATATYPE_REAL     = 0,        // real values, no constraints
-  mjDATATYPE_POSITIVE,            // positive values; 0 or negative: inactive
-  mjDATATYPE_AXIS,                // 3D unit vector
-  mjDATATYPE_QUATERNION           // unit quaternion
-} mjtDataType;
-typedef enum mjtConDataField_ {   // data fields returned by contact sensors
-  mjCONDATA_FOUND     = 0,        // whether a contact was found
-  mjCONDATA_FORCE,                // contact force
-  mjCONDATA_TORQUE,               // contact torque
-  mjCONDATA_DIST,                 // contact penetration distance
-  mjCONDATA_POS,                  // contact position
-  mjCONDATA_NORMAL,               // contact frame normal
-  mjCONDATA_TANGENT,              // contact frame first tangent
-
-  mjNCONDATA                      // number of contact sensor data fields
-} mjtConDataField;
-typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
-  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
-  mjRAYDATA_DIR,                  // normalized ray direction
-  mjRAYDATA_ORIGIN,               // ray origin
-  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
-  mjRAYDATA_NORMAL,               // surface normal at intersection point
-  mjRAYDATA_DEPTH,                // depth along z-axis
-
-  mjNRAYDATA                      // number of rangefinder sensor data fields
-} mjtRayDataField;
-typedef enum mjtCamOutBit_ {      // camera output type bitflags
-  mjCAMOUT_RGB        = 1<<0,     // RGB image
-  mjCAMOUT_DEPTH      = 1<<1,     // depth image (distance from camera plane)
-  mjCAMOUT_DIST       = 1<<2,     // distance image (distance from camera origin)
-  mjCAMOUT_NORMAL     = 1<<3,     // normal image
-  mjCAMOUT_SEG        = 1<<4,     // segmentation image
-
-  mjNCAMOUT           = 5         // number of camera output types
-} mjtCamOutBit;
-typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
-  mjSAMEFRAME_NONE    = 0,        // no alignment
-  mjSAMEFRAME_BODY,               // frame is same as body frame
-  mjSAMEFRAME_INERTIA,            // frame is same as inertial frame
-  mjSAMEFRAME_BODYROT,            // frame orientation is same as body orientation
-  mjSAMEFRAME_INERTIAROT          // frame orientation is same as inertia orientation
-} mjtSameFrame;
-typedef enum mjtSleepPolicy_ {    // per-tree sleep policy
-  mjSLEEP_AUTO        = 0,        // compiler chooses sleep policy
-  mjSLEEP_AUTO_NEVER,             // compiler sleep policy: never
-  mjSLEEP_AUTO_ALLOWED,           // compiler sleep policy: allowed
-  mjSLEEP_NEVER,                  // user sleep policy: never
-  mjSLEEP_ALLOWED,                // user sleep policy: allowed
-  mjSLEEP_INIT,                   // user sleep policy: initialized asleep
-} mjtSleepPolicy;
-typedef enum mjtLRMode_ {         // mode for actuator length range computation
-  mjLRMODE_NONE       = 0,        // do not process any actuators
-  mjLRMODE_MUSCLE,                // process muscle actuators
-  mjLRMODE_MUSCLEUSER,            // process muscle and user actuators
-  mjLRMODE_ALL                    // process all actuators
-} mjtLRMode;
-typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
-  mjFLEXSELF_NONE     = 0,        // no self-collisions
-  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
-  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
-  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)
@@ -2609,6 +2161,454 @@ struct mjTask_ {        // a task that can be executed by a thread pool.
   volatile int status;  // status of the task
 };
 typedef struct mjTask_ mjTask;
+typedef enum mjtDisableBit_ {     // disable default feature bitflags
+  mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
+  mjDSBL_EQUALITY     = 1<<1,     // equality constraints
+  mjDSBL_FRICTIONLOSS = 1<<2,     // joint and tendon frictionloss constraints
+  mjDSBL_LIMIT        = 1<<3,     // joint and tendon limit constraints
+  mjDSBL_CONTACT      = 1<<4,     // contact constraints
+  mjDSBL_SPRING       = 1<<5,     // passive spring forces
+  mjDSBL_DAMPER       = 1<<6,     // passive damping forces
+  mjDSBL_GRAVITY      = 1<<7,     // gravitational forces
+  mjDSBL_CLAMPCTRL    = 1<<8,     // clamp control to specified range
+  mjDSBL_WARMSTART    = 1<<9,     // warmstart constraint solver
+  mjDSBL_FILTERPARENT = 1<<10,    // remove collisions with parent body
+  mjDSBL_ACTUATION    = 1<<11,    // apply actuation forces
+  mjDSBL_REFSAFE      = 1<<12,    // integrator safety: make ref[0]>=2*timestep
+  mjDSBL_SENSOR       = 1<<13,    // sensors
+  mjDSBL_MIDPHASE     = 1<<14,    // mid-phase collision filtering
+  mjDSBL_EULERDAMP    = 1<<15,    // implicit integration of joint damping in Euler integrator
+  mjDSBL_AUTORESET    = 1<<16,    // automatic reset when numerical issues are detected
+  mjDSBL_NATIVECCD    = 1<<17,    // native convex collision detection
+  mjDSBL_ISLAND       = 1<<18,    // constraint island discovery
+  mjDSBL_MULTICCD     = 1<<19,    // multiple CCD contact points
+
+  mjNDISABLE          = 20        // number of disable flags
+} mjtDisableBit;
+typedef enum mjtEnableBit_ {      // enable optional feature bitflags
+  mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
+  mjENBL_ENERGY       = 1<<1,     // energy computation
+  mjENBL_FWDINV       = 1<<2,     // record solver statistics
+  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
+  mjENBL_SLEEP        = 1<<4,     // sleeping
+  mjENBL_DIAGEXACT    = 1<<5,     // exact diagonal of constraint inertia
+
+  mjNENABLE           = 6         // number of enable flags
+} mjtEnableBit;
+typedef enum mjtJoint_ {          // type of degree of freedom
+  mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)
+  mjJNT_BALL,                     // orientation (quat) relative to parent        (4)
+  mjJNT_SLIDE,                    // sliding distance along body-fixed axis       (1)
+  mjJNT_HINGE                     // rotation angle (rad) around body-fixed axis  (1)
+} mjtJoint;
+typedef enum mjtGeom_ {           // type of geometric shape
+  // regular geom types
+  mjGEOM_PLANE        = 0,        // plane
+  mjGEOM_HFIELD,                  // height field
+  mjGEOM_SPHERE,                  // sphere
+  mjGEOM_CAPSULE,                 // capsule
+  mjGEOM_ELLIPSOID,               // ellipsoid
+  mjGEOM_CYLINDER,                // cylinder
+  mjGEOM_BOX,                     // box
+  mjGEOM_MESH,                    // mesh
+  mjGEOM_SDF,                     // signed distance field
+
+  mjNGEOMTYPES,                   // number of regular geom types
+
+  // rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
+  mjGEOM_ARROW        = 100,      // arrow
+  mjGEOM_ARROW1,                  // arrow without wedges
+  mjGEOM_ARROW2,                  // arrow in both directions
+  mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
+  mjGEOM_FLEX,                    // flex
+  mjGEOM_SKIN,                    // skin
+  mjGEOM_LABEL,                   // text label
+  mjGEOM_TRIANGLE,                // triangle
+
+  mjGEOM_NONE         = 1001      // missing geom type
+} mjtGeom;
+typedef enum mjtProjection_ {     // type of camera projection
+  mjPROJ_PERSPECTIVE  = 0,        // perspective
+  mjPROJ_ORTHOGRAPHIC             // orthographic
+} mjtProjection;
+typedef enum mjtCamLight_ {       // tracking mode for camera and light
+  mjCAMLIGHT_FIXED    = 0,        // pos and rot fixed in body
+  mjCAMLIGHT_TRACK,               // pos tracks body, rot fixed in global
+  mjCAMLIGHT_TRACKCOM,            // pos tracks subtree com, rot fixed in body
+  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
+  mjTEXTURE_SKYBOX                // cube texture used as skybox
+} mjtTexture;
+typedef enum mjtTextureRole_ {    // role of texture map in rendering
+  mjTEXROLE_USER      = 0,        // unspecified
+  mjTEXROLE_RGB,                  // base color (albedo)
+  mjTEXROLE_OCCLUSION,            // ambient occlusion
+  mjTEXROLE_ROUGHNESS,            // roughness
+  mjTEXROLE_METALLIC,             // metallic
+  mjTEXROLE_NORMAL,               // normal (bump) map
+  mjTEXROLE_OPACITY,              // opacity
+  mjTEXROLE_EMISSIVE,             // light emission
+  mjTEXROLE_RGBA,                 // base color, opacity
+  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
+  mjINT_IMPLICIT,                 // implicit in velocity
+  mjINT_IMPLICITFAST              // implicit in velocity, no rne derivative
+} mjtIntegrator;
+typedef enum mjtCone_ {           // type of friction cone
+  mjCONE_PYRAMIDAL     = 0,       // pyramidal
+  mjCONE_ELLIPTIC                 // elliptic
+} mjtCone;
+typedef enum mjtJacobian_ {       // type of constraint Jacobian
+  mjJAC_DENSE          = 0,       // dense
+  mjJAC_SPARSE,                   // sparse
+  mjJAC_AUTO                      // dense if nv<60, sparse otherwise
+} mjtJacobian;
+typedef enum mjtSolver_ {         // constraint solver algorithm
+  mjSOL_PGS            = 0,       // PGS    (dual)
+  mjSOL_CG,                       // CG     (primal)
+  mjSOL_NEWTON                    // Newton (primal)
+} mjtSolver;
+typedef enum mjtEq_ {             // type of equality constraint
+  mjEQ_CONNECT        = 0,        // connect two bodies at a point (ball joint)
+  mjEQ_WELD,                      // fix relative position and orientation of two bodies
+  mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
+  mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
+  mjEQ_FLEX,                      // fix all edge lengths of a flex
+  mjEQ_FLEXVERT,                  // fix all vertex lengths of a flex
+  mjEQ_FLEXSTRAIN,                // constrain strain of a trilinear/quadratic flex (B-bar)
+  mjEQ_DISTANCE                   // unsupported, will cause an error if used
+} mjtEq;
+typedef enum mjtWrap_ {           // type of tendon wrap object
+  mjWRAP_NONE         = 0,        // null object
+  mjWRAP_JOINT,                   // constant moment arm
+  mjWRAP_PULLEY,                  // pulley used to split tendon
+  mjWRAP_SITE,                    // pass through site
+  mjWRAP_SPHERE,                  // wrap around sphere
+  mjWRAP_CYLINDER                 // wrap around (infinite) cylinder
+} mjtWrap;
+typedef enum mjtTrn_ {            // type of actuator transmission
+  mjTRN_JOINT         = 0,        // force on joint
+  mjTRN_JOINTINPARENT,            // force on joint, expressed in parent frame
+  mjTRN_SLIDERCRANK,              // force via slider-crank linkage
+  mjTRN_TENDON,                   // force on tendon
+  mjTRN_SITE,                     // force on site
+  mjTRN_BODY,                     // adhesion force on a body's geoms
+
+  mjTRN_UNDEFINED     = 1000      // undefined transmission type
+} mjtTrn;
+typedef enum mjtDyn_ {            // type of actuator dynamics
+  mjDYN_NONE          = 0,        // no internal dynamics; ctrl specifies force
+  mjDYN_INTEGRATOR,               // integrator: da/dt = u
+  mjDYN_FILTER,                   // linear filter: da/dt = (u-a) / tau
+  mjDYN_FILTEREXACT,              // linear filter: da/dt = (u-a) / tau, with exact integration
+  mjDYN_MUSCLE,                   // piecewise linear filter with two time constants
+  mjDYN_DCMOTOR,                  // DC motor electrical dynamics
+  mjDYN_USER                      // user-defined dynamics type
+} mjtDyn;
+typedef enum mjtGain_ {           // type of actuator gain
+  mjGAIN_FIXED        = 0,        // fixed gain
+  mjGAIN_AFFINE,                  // const + kp*length + kv*velocity
+  mjGAIN_MUSCLE,                  // muscle FLV curve computed by mju_muscleGain()
+  mjGAIN_DCMOTOR,                 // DC motor gain: K or K/R
+  mjGAIN_USER                     // user-defined gain type
+} mjtGain;
+typedef enum mjtBias_ {           // type of actuator bias
+  mjBIAS_NONE         = 0,        // no bias
+  mjBIAS_AFFINE,                  // const + kp*length + kv*velocity
+  mjBIAS_MUSCLE,                  // muscle passive force computed by mju_muscleBias()
+  mjBIAS_DCMOTOR,                 // DC motor bias: back-EMF, cogging, LuGre friction
+  mjBIAS_USER                     // user-defined bias type
+} mjtBias;
+typedef enum mjtObj_ {            // type of MujoCo object
+  mjOBJ_UNKNOWN       = 0,        // unknown object type
+  mjOBJ_BODY,                     // body
+  mjOBJ_XBODY,                    // body, used to access regular frame instead of i-frame
+  mjOBJ_JOINT,                    // joint
+  mjOBJ_DOF,                      // dof
+  mjOBJ_GEOM,                     // geom
+  mjOBJ_SITE,                     // site
+  mjOBJ_CAMERA,                   // camera
+  mjOBJ_LIGHT,                    // light
+  mjOBJ_FLEX,                     // flex
+  mjOBJ_MESH,                     // mesh
+  mjOBJ_SKIN,                     // skin
+  mjOBJ_HFIELD,                   // heightfield
+  mjOBJ_TEXTURE,                  // texture
+  mjOBJ_MATERIAL,                 // material for rendering
+  mjOBJ_PAIR,                     // geom pair to include
+  mjOBJ_EXCLUDE,                  // body pair to exclude
+  mjOBJ_EQUALITY,                 // equality constraint
+  mjOBJ_TENDON,                   // tendon
+  mjOBJ_ACTUATOR,                 // actuator
+  mjOBJ_SENSOR,                   // sensor
+  mjOBJ_NUMERIC,                  // numeric
+  mjOBJ_TEXT,                     // text
+  mjOBJ_TUPLE,                    // tuple
+  mjOBJ_KEY,                      // keyframe
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT,                      // number of object types
+
+  // meta elements, do not appear in mjModel
+  mjOBJ_FRAME         = 100,      // frame
+  mjOBJ_DEFAULT,                  // default
+  mjOBJ_MODEL                     // entire model
+} mjtObj;
+typedef enum mjtSensor_ {         // type of sensor
+  // common robotic sensors, attached to a site
+  mjSENS_TOUCH        = 0,        // scalar contact normal forces summed over sensor zone
+  mjSENS_ACCELEROMETER,           // 3D linear acceleration, in local frame
+  mjSENS_VELOCIMETER,             // 3D linear velocity, in local frame
+  mjSENS_GYRO,                    // 3D angular velocity, in local frame
+  mjSENS_FORCE,                   // 3D force between site's body and its parent body
+  mjSENS_TORQUE,                  // 3D torque between site's body and its parent body
+  mjSENS_MAGNETOMETER,            // 3D magnetometer
+  mjSENS_RANGEFINDER,             // scalar distance to nearest geom along z-axis
+  mjSENS_CAMPROJECTION,           // pixel coordinates of a site in the camera image
+
+  // sensors related to scalar joints, tendons, actuators
+  mjSENS_JOINTPOS,                // scalar joint position (hinge and slide only)
+  mjSENS_JOINTVEL,                // scalar joint velocity (hinge and slide only)
+  mjSENS_TENDONPOS,               // scalar tendon position
+  mjSENS_TENDONVEL,               // scalar tendon velocity
+  mjSENS_ACTUATORPOS,             // scalar actuator position
+  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
+  mjSENS_BALLANGVEL,              // 3D ball joint angular velocity
+
+  // joint and tendon limit sensors, in constraint space
+  mjSENS_JOINTLIMITPOS,           // joint limit distance-margin
+  mjSENS_JOINTLIMITVEL,           // joint limit velocity
+  mjSENS_JOINTLIMITFRC,           // joint limit force
+  mjSENS_TENDONLIMITPOS,          // tendon limit distance-margin
+  mjSENS_TENDONLIMITVEL,          // tendon limit velocity
+  mjSENS_TENDONLIMITFRC,          // tendon limit force
+
+  // sensors attached to an object with spatial frame: (x)body, geom, site, camera
+  mjSENS_FRAMEPOS,                // 3D position
+  mjSENS_FRAMEQUAT,               // 4D unit quaternion orientation
+  mjSENS_FRAMEXAXIS,              // 3D unit vector: x-axis of object's frame
+  mjSENS_FRAMEYAXIS,              // 3D unit vector: y-axis of object's frame
+  mjSENS_FRAMEZAXIS,              // 3D unit vector: z-axis of object's frame
+  mjSENS_FRAMELINVEL,             // 3D linear velocity
+  mjSENS_FRAMEANGVEL,             // 3D angular velocity
+  mjSENS_FRAMELINACC,             // 3D linear acceleration
+  mjSENS_FRAMEANGACC,             // 3D angular acceleration
+
+  // sensors related to kinematic subtrees; attached to a body (which is the subtree root)
+  mjSENS_SUBTREECOM,              // 3D center of mass of subtree
+  mjSENS_SUBTREELINVEL,           // 3D linear velocity of subtree
+  mjSENS_SUBTREEANGMOM,           // 3D angular momentum of subtree
+
+  // sensors of geometric relationships
+  mjSENS_INSIDESITE,              // 1 if object is inside a site, 0 otherwise
+  mjSENS_GEOMDIST,                // signed distance between two geoms
+  mjSENS_GEOMNORMAL,              // normal direction between two geoms
+  mjSENS_GEOMFROMTO,              // segment between two geoms
+
+  // sensors for reporting contacts which occurred during the simulation
+  mjSENS_CONTACT,                 // contacts which occurred during the simulation
+
+  // global sensors
+  mjSENS_E_POTENTIAL,             // potential energy
+  mjSENS_E_KINETIC,               // kinetic energy
+  mjSENS_CLOCK,                   // simulation time
+
+  // sensors related to SDFs
+  mjSENS_TACTILE,                 // tactile sensor
+
+  // plugin-controlled sensors
+  mjSENS_PLUGIN,                  // plugin-controlled
+
+  // user-defined sensor
+  mjSENS_USER                     // sensor data provided by mjcb_sensor callback
+} mjtSensor;
+typedef enum mjtStage_ {          // computation stage
+  mjSTAGE_NONE        = 0,        // no computations
+  mjSTAGE_POS,                    // position-dependent computations
+  mjSTAGE_VEL,                    // velocity-dependent computations
+  mjSTAGE_ACC                     // acceleration/force-dependent computations
+} mjtStage;
+typedef enum mjtDataType_ {       // data type for sensors
+  mjDATATYPE_REAL     = 0,        // real values, no constraints
+  mjDATATYPE_POSITIVE,            // positive values; 0 or negative: inactive
+  mjDATATYPE_AXIS,                // 3D unit vector
+  mjDATATYPE_QUATERNION           // unit quaternion
+} mjtDataType;
+typedef enum mjtConDataField_ {   // data fields returned by contact sensors
+  mjCONDATA_FOUND     = 0,        // whether a contact was found
+  mjCONDATA_FORCE,                // contact force
+  mjCONDATA_TORQUE,               // contact torque
+  mjCONDATA_DIST,                 // contact penetration distance
+  mjCONDATA_POS,                  // contact position
+  mjCONDATA_NORMAL,               // contact frame normal
+  mjCONDATA_TANGENT,              // contact frame first tangent
+
+  mjNCONDATA                      // number of contact sensor data fields
+} mjtConDataField;
+typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
+  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
+  mjRAYDATA_DIR,                  // normalized ray direction
+  mjRAYDATA_ORIGIN,               // ray origin
+  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
+  mjRAYDATA_NORMAL,               // surface normal at intersection point
+  mjRAYDATA_DEPTH,                // depth along z-axis
+
+  mjNRAYDATA                      // number of rangefinder sensor data fields
+} mjtRayDataField;
+typedef enum mjtCamOutBit_ {      // camera output type bitflags
+  mjCAMOUT_RGB        = 1<<0,     // RGB image
+  mjCAMOUT_DEPTH      = 1<<1,     // depth image (distance from camera plane)
+  mjCAMOUT_DIST       = 1<<2,     // distance image (distance from camera origin)
+  mjCAMOUT_NORMAL     = 1<<3,     // normal image
+  mjCAMOUT_SEG        = 1<<4,     // segmentation image
+
+  mjNCAMOUT           = 5         // number of camera output types
+} mjtCamOutBit;
+typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
+  mjSAMEFRAME_NONE    = 0,        // no alignment
+  mjSAMEFRAME_BODY,               // frame is same as body frame
+  mjSAMEFRAME_INERTIA,            // frame is same as inertial frame
+  mjSAMEFRAME_BODYROT,            // frame orientation is same as body orientation
+  mjSAMEFRAME_INERTIAROT          // frame orientation is same as inertia orientation
+} mjtSameFrame;
+typedef enum mjtSleepPolicy_ {    // per-tree sleep policy
+  mjSLEEP_AUTO        = 0,        // compiler chooses sleep policy
+  mjSLEEP_AUTO_NEVER,             // compiler sleep policy: never
+  mjSLEEP_AUTO_ALLOWED,           // compiler sleep policy: allowed
+  mjSLEEP_NEVER,                  // user sleep policy: never
+  mjSLEEP_ALLOWED,                // user sleep policy: allowed
+  mjSLEEP_INIT,                   // user sleep policy: initialized asleep
+} mjtSleepPolicy;
+typedef enum mjtLRMode_ {         // mode for actuator length range computation
+  mjLRMODE_NONE       = 0,        // do not process any actuators
+  mjLRMODE_MUSCLE,                // process muscle actuators
+  mjLRMODE_MUSCLEUSER,            // process muscle and user actuators
+  mjLRMODE_ALL                    // process all actuators
+} mjtLRMode;
+typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
+  mjFLEXSELF_NONE     = 0,        // no self-collisions
+  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
+  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
+  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;
+typedef enum mjtState_ {            // state elements
+  mjSTATE_TIME           = 1<<0,    // time
+  mjSTATE_QPOS           = 1<<1,    // position
+  mjSTATE_QVEL           = 1<<2,    // velocity
+  mjSTATE_ACT            = 1<<3,    // actuator activation
+  mjSTATE_HISTORY        = 1<<4,    // history buffers (control, sensor)
+  mjSTATE_WARMSTART      = 1<<5,    // acceleration used for warmstart
+  mjSTATE_CTRL           = 1<<6,    // control
+  mjSTATE_QFRC_APPLIED   = 1<<7,    // applied generalized force
+  mjSTATE_XFRC_APPLIED   = 1<<8,    // applied Cartesian force/torque
+  mjSTATE_EQ_ACTIVE      = 1<<9,    // enable/disable constraints
+  mjSTATE_MOCAP_POS      = 1<<10,   // positions of mocap bodies
+  mjSTATE_MOCAP_QUAT     = 1<<11,   // orientations of mocap bodies
+  mjSTATE_USERDATA       = 1<<12,   // user data
+  mjSTATE_PLUGIN         = 1<<13,   // plugin state
+
+  mjNSTATE               = 14,      // number of state elements
+
+  // convenience values for commonly used state specifications
+  mjSTATE_PHYSICS        = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
+  mjSTATE_FULLPHYSICS    = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
+  mjSTATE_USER           = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
+                           mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
+                           mjSTATE_USERDATA,
+  mjSTATE_INTEGRATION    = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
+} mjtState;
+typedef enum mjtConstraint_ {       // type of constraint
+  mjCNSTR_EQUALITY       = 0,       // equality constraint
+  mjCNSTR_FRICTION_DOF,             // dof friction
+  mjCNSTR_FRICTION_TENDON,          // tendon friction
+  mjCNSTR_LIMIT_JOINT,              // joint limit
+  mjCNSTR_LIMIT_TENDON,             // tendon limit
+  mjCNSTR_CONTACT_FRICTIONLESS,     // frictionless contact
+  mjCNSTR_CONTACT_PYRAMIDAL,        // frictional contact, pyramidal friction cone
+  mjCNSTR_CONTACT_ELLIPTIC          // frictional contact, elliptic friction cone
+} mjtConstraint;
+typedef enum mjtConstraintState_ {  // constraint state
+  mjCNSTRSTATE_SATISFIED = 0,       // constraint satisfied, zero cost (limit, contact)
+  mjCNSTRSTATE_QUADRATIC,           // quadratic cost (equality, friction, limit, contact)
+  mjCNSTRSTATE_LINEARNEG,           // linear cost, negative side (friction)
+  mjCNSTRSTATE_LINEARPOS,           // linear cost, positive side (friction)
+  mjCNSTRSTATE_CONE                 // squared distance to cone cost (elliptic contact)
+} mjtConstraintState;
+typedef enum mjtWarning_ {          // warning types
+  mjWARN_INERTIA         = 0,       // (near) singular inertia matrix
+  mjWARN_CONTACTFULL,               // too many contacts in contact list
+  mjWARN_CNSTRFULL,                 // too many constraints
+  mjWARN_BADQPOS,                   // bad number in qpos
+  mjWARN_BADQVEL,                   // bad number in qvel
+  mjWARN_BADQACC,                   // bad number in qacc
+  mjWARN_BADCTRL,                   // bad number in ctrl
+
+  mjNWARNING                        // number of warnings
+} mjtWarning;
+typedef enum mjtTimer_ {            // internal timers
+  // main api
+  mjTIMER_STEP           = 0,       // step
+  mjTIMER_FORWARD,                  // forward
+  mjTIMER_INVERSE,                  // inverse
+
+  // breakdown of step/forward
+  mjTIMER_POSITION,                 // fwdPosition
+  mjTIMER_VELOCITY,                 // fwdVelocity
+  mjTIMER_ACTUATION,                // fwdActuation
+  mjTIMER_CONSTRAINT,               // fwdConstraint
+  mjTIMER_ADVANCE,                  // mj_Euler, mj_implicit
+
+  // breakdown of fwdPosition
+  mjTIMER_POS_KINEMATICS,           // kinematics, com, tendon, transmission
+  mjTIMER_POS_INERTIA,              // inertia computations
+  mjTIMER_POS_COLLISION,            // collision detection
+  mjTIMER_POS_MAKE,                 // make constraints
+  mjTIMER_POS_PROJECT,              // project constraints
+
+  // breakdown of mj_collision
+  mjTIMER_COL_BROAD,                // broadphase
+  mjTIMER_COL_NARROW,               // narrowphase
+
+  mjNTIMER                          // number of timers
+} mjtTimer;
+typedef enum mjtSleepState_ {       // sleep state of an object
+  mjS_STATIC = -1,                  // object is static
+  mjS_ASLEEP = 0,                   // object is asleep
+  mjS_AWAKE  = 1                    // object is awake
+} mjtSleepState;
 typedef enum mjtButton_ {         // mouse button
   mjBUTTON_NONE = 0,              // no button
   mjBUTTON_LEFT,                  // left button
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index 7b34cdcc..035fdaca 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -165,8 +165,9 @@ links below, to make this documentation self-contained.
    Defines the primitive types and structures needed by the OpenGL renderer.
 `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`__
    Defines the primitive types and structures needed by the UI framework.
-`mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
-   Defines MuJoCo's ``mjtNum`` floating-point type to be either ``double`` or ``float``. See :ref:`mjtNum`.
+`mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__
+   Defines primitive types and enums, including the ``mjtNum`` floating-point type to be either ``double`` or ``float``
+   (see :ref:`mjtNum`).
 `mjspec.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspec.h>`__
    Defines enums and structs used for :doc:`procedural model editing <modeledit>`.
 `mjplugin.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjplugin.h>`__
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
index 8a22911c..e7e56e9e 100644
--- a/include/mujoco/mjdata.h
+++ b/include/mujoco/mjdata.h
@@ -18,107 +18,10 @@
 #include <stddef.h>
 #include <stdint.h>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjthread.h>
 
-//---------------------------------- primitive types (mjt) -----------------------------------------
-
-typedef enum mjtState_ {            // state elements
-  mjSTATE_TIME           = 1<<0,    // time
-  mjSTATE_QPOS           = 1<<1,    // position
-  mjSTATE_QVEL           = 1<<2,    // velocity
-  mjSTATE_ACT            = 1<<3,    // actuator activation
-  mjSTATE_HISTORY        = 1<<4,    // history buffers (control, sensor)
-  mjSTATE_WARMSTART      = 1<<5,    // acceleration used for warmstart
-  mjSTATE_CTRL           = 1<<6,    // control
-  mjSTATE_QFRC_APPLIED   = 1<<7,    // applied generalized force
-  mjSTATE_XFRC_APPLIED   = 1<<8,    // applied Cartesian force/torque
-  mjSTATE_EQ_ACTIVE      = 1<<9,    // enable/disable constraints
-  mjSTATE_MOCAP_POS      = 1<<10,   // positions of mocap bodies
-  mjSTATE_MOCAP_QUAT     = 1<<11,   // orientations of mocap bodies
-  mjSTATE_USERDATA       = 1<<12,   // user data
-  mjSTATE_PLUGIN         = 1<<13,   // plugin state
-
-  mjNSTATE               = 14,      // number of state elements
-
-  // convenience values for commonly used state specifications
-  mjSTATE_PHYSICS        = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
-  mjSTATE_FULLPHYSICS    = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
-  mjSTATE_USER           = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
-                          mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
-                          mjSTATE_USERDATA,
-  mjSTATE_INTEGRATION    = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
-} mjtState;
-
-
-typedef enum mjtConstraint_ {       // type of constraint
-  mjCNSTR_EQUALITY       = 0,       // equality constraint
-  mjCNSTR_FRICTION_DOF,             // dof friction
-  mjCNSTR_FRICTION_TENDON,          // tendon friction
-  mjCNSTR_LIMIT_JOINT,              // joint limit
-  mjCNSTR_LIMIT_TENDON,             // tendon limit
-  mjCNSTR_CONTACT_FRICTIONLESS,     // frictionless contact
-  mjCNSTR_CONTACT_PYRAMIDAL,        // frictional contact, pyramidal friction cone
-  mjCNSTR_CONTACT_ELLIPTIC          // frictional contact, elliptic friction cone
-} mjtConstraint;
-
-
-typedef enum mjtConstraintState_ {  // constraint state
-  mjCNSTRSTATE_SATISFIED = 0,       // constraint satisfied, zero cost (limit, contact)
-  mjCNSTRSTATE_QUADRATIC,           // quadratic cost (equality, friction, limit, contact)
-  mjCNSTRSTATE_LINEARNEG,           // linear cost, negative side (friction)
-  mjCNSTRSTATE_LINEARPOS,           // linear cost, positive side (friction)
-  mjCNSTRSTATE_CONE                 // squared distance to cone cost (elliptic contact)
-} mjtConstraintState;
-
-
-typedef enum mjtWarning_ {          // warning types
-  mjWARN_INERTIA         = 0,       // (near) singular inertia matrix
-  mjWARN_CONTACTFULL,               // too many contacts in contact list
-  mjWARN_CNSTRFULL,                 // too many constraints
-  mjWARN_BADQPOS,                   // bad number in qpos
-  mjWARN_BADQVEL,                   // bad number in qvel
-  mjWARN_BADQACC,                   // bad number in qacc
-  mjWARN_BADCTRL,                   // bad number in ctrl
-
-  mjNWARNING                        // number of warnings
-} mjtWarning;
-
-
-typedef enum mjtTimer_ {            // internal timers
-  // main api
-  mjTIMER_STEP           = 0,       // step
-  mjTIMER_FORWARD,                  // forward
-  mjTIMER_INVERSE,                  // inverse
-
-  // breakdown of step/forward
-  mjTIMER_POSITION,                 // fwdPosition
-  mjTIMER_VELOCITY,                 // fwdVelocity
-  mjTIMER_ACTUATION,                // fwdActuation
-  mjTIMER_CONSTRAINT,               // fwdConstraint
-  mjTIMER_ADVANCE,                  // mj_Euler, mj_implicit
-
-  // breakdown of fwdPosition
-  mjTIMER_POS_KINEMATICS,           // kinematics, com, tendon, transmission
-  mjTIMER_POS_INERTIA,              // inertia computations
-  mjTIMER_POS_COLLISION,            // collision detection
-  mjTIMER_POS_MAKE,                 // make constraints
-  mjTIMER_POS_PROJECT,              // project constraints
-
-  // breakdown of mj_collision
-  mjTIMER_COL_BROAD,                // broadphase
-  mjTIMER_COL_NARROW,               // narrowphase
-
-  mjNTIMER                          // number of timers
-} mjtTimer;
-
-
-typedef enum mjtSleepState_ {       // sleep state of an object
-  mjS_STATIC = -1,                  // object is static
-  mjS_ASLEEP = 0,                   // object is asleep
-  mjS_AWAKE  = 1                    // object is awake
-} mjtSleepState;
 
 
 //------------------------------------- Contact ----------------------------------------------------
diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h
index 6e4ddb5d..2207581a 100644
--- a/include/mujoco/mjmodel.h
+++ b/include/mujoco/mjmodel.h
@@ -18,8 +18,7 @@
 #include <stddef.h>
 #include <stdint.h>
 
-
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 // global constants
 #define mjPI            3.14159265358979323846
@@ -48,434 +47,6 @@
 #define mjNISLAND       20        // number of mjData.solver arrays
 
 
-//---------------------------------- enum types (mjt) ----------------------------------------------
-
-typedef enum mjtDisableBit_ {     // disable default feature bitflags
-  mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
-  mjDSBL_EQUALITY     = 1<<1,     // equality constraints
-  mjDSBL_FRICTIONLOSS = 1<<2,     // joint and tendon frictionloss constraints
-  mjDSBL_LIMIT        = 1<<3,     // joint and tendon limit constraints
-  mjDSBL_CONTACT      = 1<<4,     // contact constraints
-  mjDSBL_SPRING       = 1<<5,     // passive spring forces
-  mjDSBL_DAMPER       = 1<<6,     // passive damping forces
-  mjDSBL_GRAVITY      = 1<<7,     // gravitational forces
-  mjDSBL_CLAMPCTRL    = 1<<8,     // clamp control to specified range
-  mjDSBL_WARMSTART    = 1<<9,     // warmstart constraint solver
-  mjDSBL_FILTERPARENT = 1<<10,    // remove collisions with parent body
-  mjDSBL_ACTUATION    = 1<<11,    // apply actuation forces
-  mjDSBL_REFSAFE      = 1<<12,    // integrator safety: make ref[0]>=2*timestep
-  mjDSBL_SENSOR       = 1<<13,    // sensors
-  mjDSBL_MIDPHASE     = 1<<14,    // mid-phase collision filtering
-  mjDSBL_EULERDAMP    = 1<<15,    // implicit integration of joint damping in Euler integrator
-  mjDSBL_AUTORESET    = 1<<16,    // automatic reset when numerical issues are detected
-  mjDSBL_NATIVECCD    = 1<<17,    // native convex collision detection
-  mjDSBL_ISLAND       = 1<<18,    // constraint island discovery
-  mjDSBL_MULTICCD     = 1<<19,    // multiple CCD contact points
-
-  mjNDISABLE          = 20        // number of disable flags
-} mjtDisableBit;
-
-
-typedef enum mjtEnableBit_ {      // enable optional feature bitflags
-  mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
-  mjENBL_ENERGY       = 1<<1,     // energy computation
-  mjENBL_FWDINV       = 1<<2,     // record solver statistics
-  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
-  mjENBL_SLEEP        = 1<<4,     // sleeping
-  mjENBL_DIAGEXACT    = 1<<5,     // exact diagonal of constraint inertia
-
-  mjNENABLE           = 6         // number of enable flags
-} mjtEnableBit;
-
-
-typedef enum mjtJoint_ {          // type of degree of freedom
-  mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)
-  mjJNT_BALL,                     // orientation (quat) relative to parent        (4)
-  mjJNT_SLIDE,                    // sliding distance along body-fixed axis       (1)
-  mjJNT_HINGE                     // rotation angle (rad) around body-fixed axis  (1)
-} mjtJoint;
-
-
-typedef enum mjtGeom_ {           // type of geometric shape
-  // regular geom types
-  mjGEOM_PLANE        = 0,        // plane
-  mjGEOM_HFIELD,                  // height field
-  mjGEOM_SPHERE,                  // sphere
-  mjGEOM_CAPSULE,                 // capsule
-  mjGEOM_ELLIPSOID,               // ellipsoid
-  mjGEOM_CYLINDER,                // cylinder
-  mjGEOM_BOX,                     // box
-  mjGEOM_MESH,                    // mesh
-  mjGEOM_SDF,                     // signed distance field
-
-  mjNGEOMTYPES,                   // number of regular geom types
-
-  // rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
-  mjGEOM_ARROW        = 100,      // arrow
-  mjGEOM_ARROW1,                  // arrow without wedges
-  mjGEOM_ARROW2,                  // arrow in both directions
-  mjGEOM_LINE,                    // line
-  mjGEOM_LINEBOX,                 // box with line edges
-  mjGEOM_FLEX,                    // flex
-  mjGEOM_SKIN,                    // skin
-  mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle
-
-  mjGEOM_NONE         = 1001      // missing geom type
-} mjtGeom;
-
-
-typedef enum mjtProjection_ {     // type of camera projection
-  mjPROJ_PERSPECTIVE  = 0,        // perspective
-  mjPROJ_ORTHOGRAPHIC             // orthographic
-} mjtProjection;
-
-
-typedef enum mjtCamLight_ {       // tracking mode for camera and light
-  mjCAMLIGHT_FIXED    = 0,        // pos and rot fixed in body
-  mjCAMLIGHT_TRACK,               // pos tracks body, rot fixed in global
-  mjCAMLIGHT_TRACKCOM,            // pos tracks subtree com, rot fixed in body
-  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
-  mjTEXTURE_SKYBOX                // cube texture used as skybox
-} mjtTexture;
-
-
-typedef enum mjtTextureRole_ {    // role of texture map in rendering
-  mjTEXROLE_USER      = 0,        // unspecified
-  mjTEXROLE_RGB,                  // base color (albedo)
-  mjTEXROLE_OCCLUSION,            // ambient occlusion
-  mjTEXROLE_ROUGHNESS,            // roughness
-  mjTEXROLE_METALLIC,             // metallic
-  mjTEXROLE_NORMAL,               // normal (bump) map
-  mjTEXROLE_OPACITY,              // opacity
-  mjTEXROLE_EMISSIVE,             // light emission
-  mjTEXROLE_RGBA,                 // base color, opacity
-  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
-  mjINT_IMPLICIT,                 // implicit in velocity
-  mjINT_IMPLICITFAST              // implicit in velocity, no rne derivative
-} mjtIntegrator;
-
-
-typedef enum mjtCone_ {           // type of friction cone
-  mjCONE_PYRAMIDAL     = 0,       // pyramidal
-  mjCONE_ELLIPTIC                 // elliptic
-} mjtCone;
-
-
-typedef enum mjtJacobian_ {       // type of constraint Jacobian
-  mjJAC_DENSE          = 0,       // dense
-  mjJAC_SPARSE,                   // sparse
-  mjJAC_AUTO                      // dense if nv<60, sparse otherwise
-} mjtJacobian;
-
-
-typedef enum mjtSolver_ {         // constraint solver algorithm
-  mjSOL_PGS            = 0,       // PGS    (dual)
-  mjSOL_CG,                       // CG     (primal)
-  mjSOL_NEWTON                    // Newton (primal)
-} mjtSolver;
-
-
-typedef enum mjtEq_ {             // type of equality constraint
-  mjEQ_CONNECT        = 0,        // connect two bodies at a point (ball joint)
-  mjEQ_WELD,                      // fix relative position and orientation of two bodies
-  mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
-  mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
-  mjEQ_FLEX,                      // fix all edge lengths of a flex
-  mjEQ_FLEXVERT,                  // fix all vertex lengths of a flex
-  mjEQ_FLEXSTRAIN,                // constrain strain of a trilinear/quadratic flex (B-bar)
-  mjEQ_DISTANCE                   // unsupported, will cause an error if used
-} mjtEq;
-
-
-typedef enum mjtWrap_ {           // type of tendon wrap object
-  mjWRAP_NONE         = 0,        // null object
-  mjWRAP_JOINT,                   // constant moment arm
-  mjWRAP_PULLEY,                  // pulley used to split tendon
-  mjWRAP_SITE,                    // pass through site
-  mjWRAP_SPHERE,                  // wrap around sphere
-  mjWRAP_CYLINDER                 // wrap around (infinite) cylinder
-} mjtWrap;
-
-
-typedef enum mjtTrn_ {            // type of actuator transmission
-  mjTRN_JOINT         = 0,        // force on joint
-  mjTRN_JOINTINPARENT,            // force on joint, expressed in parent frame
-  mjTRN_SLIDERCRANK,              // force via slider-crank linkage
-  mjTRN_TENDON,                   // force on tendon
-  mjTRN_SITE,                     // force on site
-  mjTRN_BODY,                     // adhesion force on a body's geoms
-
-  mjTRN_UNDEFINED     = 1000      // undefined transmission type
-} mjtTrn;
-
-
-typedef enum mjtDyn_ {            // type of actuator dynamics
-  mjDYN_NONE          = 0,        // no internal dynamics; ctrl specifies force
-  mjDYN_INTEGRATOR,               // integrator: da/dt = u
-  mjDYN_FILTER,                   // linear filter: da/dt = (u-a) / tau
-  mjDYN_FILTEREXACT,              // linear filter: da/dt = (u-a) / tau, with exact integration
-  mjDYN_MUSCLE,                   // piecewise linear filter with two time constants
-  mjDYN_DCMOTOR,                  // DC motor electrical dynamics
-  mjDYN_USER                      // user-defined dynamics type
-} mjtDyn;
-
-
-typedef enum mjtGain_ {           // type of actuator gain
-  mjGAIN_FIXED        = 0,        // fixed gain
-  mjGAIN_AFFINE,                  // const + kp*length + kv*velocity
-  mjGAIN_MUSCLE,                  // muscle FLV curve computed by mju_muscleGain()
-  mjGAIN_DCMOTOR,                 // DC motor gain: K or K/R
-  mjGAIN_USER                     // user-defined gain type
-} mjtGain;
-
-
-typedef enum mjtBias_ {           // type of actuator bias
-  mjBIAS_NONE         = 0,        // no bias
-  mjBIAS_AFFINE,                  // const + kp*length + kv*velocity
-  mjBIAS_MUSCLE,                  // muscle passive force computed by mju_muscleBias()
-  mjBIAS_DCMOTOR,                 // DC motor bias: back-EMF, cogging, LuGre friction
-  mjBIAS_USER                     // user-defined bias type
-} mjtBias;
-
-
-typedef enum mjtObj_ {            // type of MujoCo object
-  mjOBJ_UNKNOWN       = 0,        // unknown object type
-  mjOBJ_BODY,                     // body
-  mjOBJ_XBODY,                    // body, used to access regular frame instead of i-frame
-  mjOBJ_JOINT,                    // joint
-  mjOBJ_DOF,                      // dof
-  mjOBJ_GEOM,                     // geom
-  mjOBJ_SITE,                     // site
-  mjOBJ_CAMERA,                   // camera
-  mjOBJ_LIGHT,                    // light
-  mjOBJ_FLEX,                     // flex
-  mjOBJ_MESH,                     // mesh
-  mjOBJ_SKIN,                     // skin
-  mjOBJ_HFIELD,                   // heightfield
-  mjOBJ_TEXTURE,                  // texture
-  mjOBJ_MATERIAL,                 // material for rendering
-  mjOBJ_PAIR,                     // geom pair to include
-  mjOBJ_EXCLUDE,                  // body pair to exclude
-  mjOBJ_EQUALITY,                 // equality constraint
-  mjOBJ_TENDON,                   // tendon
-  mjOBJ_ACTUATOR,                 // actuator
-  mjOBJ_SENSOR,                   // sensor
-  mjOBJ_NUMERIC,                  // numeric
-  mjOBJ_TEXT,                     // text
-  mjOBJ_TUPLE,                    // tuple
-  mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN,                   // plugin instance
-
-  mjNOBJECT,                      // number of object types
-
-  // meta elements, do not appear in mjModel
-  mjOBJ_FRAME         = 100,      // frame
-  mjOBJ_DEFAULT,                  // default
-  mjOBJ_MODEL                     // entire model
-} mjtObj;
-
-
-typedef enum mjtSensor_ {         // type of sensor
-  // common robotic sensors, attached to a site
-  mjSENS_TOUCH        = 0,        // scalar contact normal forces summed over sensor zone
-  mjSENS_ACCELEROMETER,           // 3D linear acceleration, in local frame
-  mjSENS_VELOCIMETER,             // 3D linear velocity, in local frame
-  mjSENS_GYRO,                    // 3D angular velocity, in local frame
-  mjSENS_FORCE,                   // 3D force between site's body and its parent body
-  mjSENS_TORQUE,                  // 3D torque between site's body and its parent body
-  mjSENS_MAGNETOMETER,            // 3D magnetometer
-  mjSENS_RANGEFINDER,             // scalar distance to nearest geom along z-axis
-  mjSENS_CAMPROJECTION,           // pixel coordinates of a site in the camera image
-
-  // sensors related to scalar joints, tendons, actuators
-  mjSENS_JOINTPOS,                // scalar joint position (hinge and slide only)
-  mjSENS_JOINTVEL,                // scalar joint velocity (hinge and slide only)
-  mjSENS_TENDONPOS,               // scalar tendon position
-  mjSENS_TENDONVEL,               // scalar tendon velocity
-  mjSENS_ACTUATORPOS,             // scalar actuator position
-  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
-  mjSENS_BALLANGVEL,              // 3D ball joint angular velocity
-
-  // joint and tendon limit sensors, in constraint space
-  mjSENS_JOINTLIMITPOS,           // joint limit distance-margin
-  mjSENS_JOINTLIMITVEL,           // joint limit velocity
-  mjSENS_JOINTLIMITFRC,           // joint limit force
-  mjSENS_TENDONLIMITPOS,          // tendon limit distance-margin
-  mjSENS_TENDONLIMITVEL,          // tendon limit velocity
-  mjSENS_TENDONLIMITFRC,          // tendon limit force
-
-  // sensors attached to an object with spatial frame: (x)body, geom, site, camera
-  mjSENS_FRAMEPOS,                // 3D position
-  mjSENS_FRAMEQUAT,               // 4D unit quaternion orientation
-  mjSENS_FRAMEXAXIS,              // 3D unit vector: x-axis of object's frame
-  mjSENS_FRAMEYAXIS,              // 3D unit vector: y-axis of object's frame
-  mjSENS_FRAMEZAXIS,              // 3D unit vector: z-axis of object's frame
-  mjSENS_FRAMELINVEL,             // 3D linear velocity
-  mjSENS_FRAMEANGVEL,             // 3D angular velocity
-  mjSENS_FRAMELINACC,             // 3D linear acceleration
-  mjSENS_FRAMEANGACC,             // 3D angular acceleration
-
-  // sensors related to kinematic subtrees; attached to a body (which is the subtree root)
-  mjSENS_SUBTREECOM,              // 3D center of mass of subtree
-  mjSENS_SUBTREELINVEL,           // 3D linear velocity of subtree
-  mjSENS_SUBTREEANGMOM,           // 3D angular momentum of subtree
-
-  // sensors of geometric relationships
-  mjSENS_INSIDESITE,              // 1 if object is inside a site, 0 otherwise
-  mjSENS_GEOMDIST,                // signed distance between two geoms
-  mjSENS_GEOMNORMAL,              // normal direction between two geoms
-  mjSENS_GEOMFROMTO,              // segment between two geoms
-
-  // sensors for reporting contacts which occurred during the simulation
-  mjSENS_CONTACT,                 // contacts which occurred during the simulation
-
-  // global sensors
-  mjSENS_E_POTENTIAL,             // potential energy
-  mjSENS_E_KINETIC,               // kinetic energy
-  mjSENS_CLOCK,                   // simulation time
-
-  // sensors related to SDFs
-  mjSENS_TACTILE,                 // tactile sensor
-
-  // plugin-controlled sensors
-  mjSENS_PLUGIN,                  // plugin-controlled
-
-  // user-defined sensor
-  mjSENS_USER                     // sensor data provided by mjcb_sensor callback
-} mjtSensor;
-
-
-typedef enum mjtStage_ {          // computation stage
-  mjSTAGE_NONE        = 0,        // no computations
-  mjSTAGE_POS,                    // position-dependent computations
-  mjSTAGE_VEL,                    // velocity-dependent computations
-  mjSTAGE_ACC                     // acceleration/force-dependent computations
-} mjtStage;
-
-
-typedef enum mjtDataType_ {       // data type for sensors
-  mjDATATYPE_REAL     = 0,        // real values, no constraints
-  mjDATATYPE_POSITIVE,            // positive values; 0 or negative: inactive
-  mjDATATYPE_AXIS,                // 3D unit vector
-  mjDATATYPE_QUATERNION           // unit quaternion
-} mjtDataType;
-
-
-typedef enum mjtConDataField_ {   // data fields returned by contact sensors
-  mjCONDATA_FOUND     = 0,        // whether a contact was found
-  mjCONDATA_FORCE,                // contact force
-  mjCONDATA_TORQUE,               // contact torque
-  mjCONDATA_DIST,                 // contact penetration distance
-  mjCONDATA_POS,                  // contact position
-  mjCONDATA_NORMAL,               // contact frame normal
-  mjCONDATA_TANGENT,              // contact frame first tangent
-
-  mjNCONDATA                      // number of contact sensor data fields
-} mjtConDataField;
-
-
-typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
-  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
-  mjRAYDATA_DIR,                  // normalized ray direction
-  mjRAYDATA_ORIGIN,               // ray origin
-  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
-  mjRAYDATA_NORMAL,               // surface normal at intersection point
-  mjRAYDATA_DEPTH,                // depth along z-axis
-
-  mjNRAYDATA                      // number of rangefinder sensor data fields
-} mjtRayDataField;
-
-
-typedef enum mjtCamOutBit_ {      // camera output type bitflags
-  mjCAMOUT_RGB        = 1<<0,     // RGB image
-  mjCAMOUT_DEPTH      = 1<<1,     // depth image (distance from camera plane)
-  mjCAMOUT_DIST       = 1<<2,     // distance image (distance from camera origin)
-  mjCAMOUT_NORMAL     = 1<<3,     // normal image
-  mjCAMOUT_SEG        = 1<<4,     // segmentation image
-
-  mjNCAMOUT           = 5         // number of camera output types
-} mjtCamOutBit;
-
-
-typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
-  mjSAMEFRAME_NONE    = 0,        // no alignment
-  mjSAMEFRAME_BODY,               // frame is same as body frame
-  mjSAMEFRAME_INERTIA,            // frame is same as inertial frame
-  mjSAMEFRAME_BODYROT,            // frame orientation is same as body orientation
-  mjSAMEFRAME_INERTIAROT          // frame orientation is same as inertia orientation
-} mjtSameFrame;
-
-
-typedef enum mjtSleepPolicy_ {    // per-tree sleep policy
-  mjSLEEP_AUTO        = 0,        // compiler chooses sleep policy
-  mjSLEEP_AUTO_NEVER,             // compiler sleep policy: never
-  mjSLEEP_AUTO_ALLOWED,           // compiler sleep policy: allowed
-  mjSLEEP_NEVER,                  // user sleep policy: never
-  mjSLEEP_ALLOWED,                // user sleep policy: allowed
-  mjSLEEP_INIT,                   // user sleep policy: initialized asleep
-} mjtSleepPolicy;
-
-
-typedef enum mjtLRMode_ {         // mode for actuator length range computation
-  mjLRMODE_NONE       = 0,        // do not process any actuators
-  mjLRMODE_MUSCLE,                // process muscle actuators
-  mjLRMODE_MUSCLEUSER,            // process muscle and user actuators
-  mjLRMODE_ALL                    // process all actuators
-} mjtLRMode;
-
-
-typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
-  mjFLEXSELF_NONE     = 0,        // no self-collisions
-  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
-  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
-  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;
-
 
 //---------------------------------- mjLROpt -------------------------------------------------------
 
diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h
index 174a1f63..a0ee0076 100644
--- a/include/mujoco/mjplugin.h
+++ b/include/mujoco/mjplugin.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 
 
diff --git a/include/mujoco/mjspec.h b/include/mujoco/mjspec.h
index ab1b2447..ad33cca3 100644
--- a/include/mujoco/mjspec.h
+++ b/include/mujoco/mjspec.h
@@ -17,7 +17,7 @@
 
 #include <stddef.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 
 // this is a C-API
diff --git a/include/mujoco/mjtnum.h b/include/mujoco/mjtnum.h
deleted file mode 100644
index 845ce73b..00000000
--- a/include/mujoco/mjtnum.h
+++ /dev/null
@@ -1,46 +0,0 @@
-// Copyright 2021 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_INCLUDE_MJTNUM_H_
-#define MUJOCO_INCLUDE_MJTNUM_H_
-
-#include <stdint.h>
-
-
-//---------------------------------- floating-point definition -------------------------------------
-
-// floating point data type and minval
-#ifndef mjUSESINGLE
-  typedef double mjtNum;
-  #define mjMINVAL    1E-15       // minimum value in any denominator
-#else
-  typedef float mjtNum;
-  #define mjMINVAL    1E-15f
-#endif
-
-
-
-//---------------------------------- byte definition -----------------------------------------------
-
-typedef unsigned char mjtByte;    // used for true/false
-
-
-
-//---------------------------------- size definition -----------------------------------------------
-
-typedef int64_t mjtSize;          // used for buffer sizes
-
-
-
-#endif  // MUJOCO_INCLUDE_MJTNUM_H_
diff --git a/include/mujoco/mjtype.h b/include/mujoco/mjtype.h
new file mode 100644
index 00000000..62426a11
--- /dev/null
+++ b/include/mujoco/mjtype.h
@@ -0,0 +1,574 @@
+// Copyright 2021 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_INCLUDE_MJTYPE_H_
+#define MUJOCO_INCLUDE_MJTYPE_H_
+
+#include <stdint.h>
+
+
+//---------------------------------- floating-point definition -------------------------------------
+
+// floating point data type and minval
+#ifndef mjUSESINGLE
+  typedef double mjtNum;
+  #define mjMINVAL    1E-15       // minimum value in any denominator
+#else
+  typedef float mjtNum;
+  #define mjMINVAL    1E-15f
+#endif
+
+
+
+//---------------------------------- byte definition -----------------------------------------------
+
+typedef unsigned char mjtByte;    // used for true/false
+
+
+
+//---------------------------------- size definition -----------------------------------------------
+
+typedef int64_t mjtSize;          // used for buffer sizes
+
+
+
+//---------------------------------- enum types (mjModel) ------------------------------------------
+
+typedef enum mjtDisableBit_ {     // disable default feature bitflags
+  mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
+  mjDSBL_EQUALITY     = 1<<1,     // equality constraints
+  mjDSBL_FRICTIONLOSS = 1<<2,     // joint and tendon frictionloss constraints
+  mjDSBL_LIMIT        = 1<<3,     // joint and tendon limit constraints
+  mjDSBL_CONTACT      = 1<<4,     // contact constraints
+  mjDSBL_SPRING       = 1<<5,     // passive spring forces
+  mjDSBL_DAMPER       = 1<<6,     // passive damping forces
+  mjDSBL_GRAVITY      = 1<<7,     // gravitational forces
+  mjDSBL_CLAMPCTRL    = 1<<8,     // clamp control to specified range
+  mjDSBL_WARMSTART    = 1<<9,     // warmstart constraint solver
+  mjDSBL_FILTERPARENT = 1<<10,    // remove collisions with parent body
+  mjDSBL_ACTUATION    = 1<<11,    // apply actuation forces
+  mjDSBL_REFSAFE      = 1<<12,    // integrator safety: make ref[0]>=2*timestep
+  mjDSBL_SENSOR       = 1<<13,    // sensors
+  mjDSBL_MIDPHASE     = 1<<14,    // mid-phase collision filtering
+  mjDSBL_EULERDAMP    = 1<<15,    // implicit integration of joint damping in Euler integrator
+  mjDSBL_AUTORESET    = 1<<16,    // automatic reset when numerical issues are detected
+  mjDSBL_NATIVECCD    = 1<<17,    // native convex collision detection
+  mjDSBL_ISLAND       = 1<<18,    // constraint island discovery
+  mjDSBL_MULTICCD     = 1<<19,    // multiple CCD contact points
+
+  mjNDISABLE          = 20        // number of disable flags
+} mjtDisableBit;
+
+
+typedef enum mjtEnableBit_ {      // enable optional feature bitflags
+  mjENBL_OVERRIDE     = 1<<0,     // override contact parameters
+  mjENBL_ENERGY       = 1<<1,     // energy computation
+  mjENBL_FWDINV       = 1<<2,     // record solver statistics
+  mjENBL_INVDISCRETE  = 1<<3,     // discrete-time inverse dynamics
+  mjENBL_SLEEP        = 1<<4,     // sleeping
+  mjENBL_DIAGEXACT    = 1<<5,     // exact diagonal of constraint inertia
+
+  mjNENABLE           = 6         // number of enable flags
+} mjtEnableBit;
+
+
+typedef enum mjtJoint_ {          // type of degree of freedom
+  mjJNT_FREE          = 0,        // global position and orientation (quat)       (7)
+  mjJNT_BALL,                     // orientation (quat) relative to parent        (4)
+  mjJNT_SLIDE,                    // sliding distance along body-fixed axis       (1)
+  mjJNT_HINGE                     // rotation angle (rad) around body-fixed axis  (1)
+} mjtJoint;
+
+
+typedef enum mjtGeom_ {           // type of geometric shape
+  // regular geom types
+  mjGEOM_PLANE        = 0,        // plane
+  mjGEOM_HFIELD,                  // height field
+  mjGEOM_SPHERE,                  // sphere
+  mjGEOM_CAPSULE,                 // capsule
+  mjGEOM_ELLIPSOID,               // ellipsoid
+  mjGEOM_CYLINDER,                // cylinder
+  mjGEOM_BOX,                     // box
+  mjGEOM_MESH,                    // mesh
+  mjGEOM_SDF,                     // signed distance field
+
+  mjNGEOMTYPES,                   // number of regular geom types
+
+  // rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
+  mjGEOM_ARROW        = 100,      // arrow
+  mjGEOM_ARROW1,                  // arrow without wedges
+  mjGEOM_ARROW2,                  // arrow in both directions
+  mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
+  mjGEOM_FLEX,                    // flex
+  mjGEOM_SKIN,                    // skin
+  mjGEOM_LABEL,                   // text label
+  mjGEOM_TRIANGLE,                // triangle
+
+  mjGEOM_NONE         = 1001      // missing geom type
+} mjtGeom;
+
+
+typedef enum mjtProjection_ {     // type of camera projection
+  mjPROJ_PERSPECTIVE  = 0,        // perspective
+  mjPROJ_ORTHOGRAPHIC             // orthographic
+} mjtProjection;
+
+
+typedef enum mjtCamLight_ {       // tracking mode for camera and light
+  mjCAMLIGHT_FIXED    = 0,        // pos and rot fixed in body
+  mjCAMLIGHT_TRACK,               // pos tracks body, rot fixed in global
+  mjCAMLIGHT_TRACKCOM,            // pos tracks subtree com, rot fixed in body
+  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
+  mjTEXTURE_SKYBOX                // cube texture used as skybox
+} mjtTexture;
+
+
+typedef enum mjtTextureRole_ {    // role of texture map in rendering
+  mjTEXROLE_USER      = 0,        // unspecified
+  mjTEXROLE_RGB,                  // base color (albedo)
+  mjTEXROLE_OCCLUSION,            // ambient occlusion
+  mjTEXROLE_ROUGHNESS,            // roughness
+  mjTEXROLE_METALLIC,             // metallic
+  mjTEXROLE_NORMAL,               // normal (bump) map
+  mjTEXROLE_OPACITY,              // opacity
+  mjTEXROLE_EMISSIVE,             // light emission
+  mjTEXROLE_RGBA,                 // base color, opacity
+  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
+  mjINT_IMPLICIT,                 // implicit in velocity
+  mjINT_IMPLICITFAST              // implicit in velocity, no rne derivative
+} mjtIntegrator;
+
+
+typedef enum mjtCone_ {           // type of friction cone
+  mjCONE_PYRAMIDAL     = 0,       // pyramidal
+  mjCONE_ELLIPTIC                 // elliptic
+} mjtCone;
+
+
+typedef enum mjtJacobian_ {       // type of constraint Jacobian
+  mjJAC_DENSE          = 0,       // dense
+  mjJAC_SPARSE,                   // sparse
+  mjJAC_AUTO                      // dense if nv<60, sparse otherwise
+} mjtJacobian;
+
+
+typedef enum mjtSolver_ {         // constraint solver algorithm
+  mjSOL_PGS            = 0,       // PGS    (dual)
+  mjSOL_CG,                       // CG     (primal)
+  mjSOL_NEWTON                    // Newton (primal)
+} mjtSolver;
+
+
+typedef enum mjtEq_ {             // type of equality constraint
+  mjEQ_CONNECT        = 0,        // connect two bodies at a point (ball joint)
+  mjEQ_WELD,                      // fix relative position and orientation of two bodies
+  mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
+  mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
+  mjEQ_FLEX,                      // fix all edge lengths of a flex
+  mjEQ_FLEXVERT,                  // fix all vertex lengths of a flex
+  mjEQ_FLEXSTRAIN,                // constrain strain of a trilinear/quadratic flex (B-bar)
+  mjEQ_DISTANCE                   // unsupported, will cause an error if used
+} mjtEq;
+
+
+typedef enum mjtWrap_ {           // type of tendon wrap object
+  mjWRAP_NONE         = 0,        // null object
+  mjWRAP_JOINT,                   // constant moment arm
+  mjWRAP_PULLEY,                  // pulley used to split tendon
+  mjWRAP_SITE,                    // pass through site
+  mjWRAP_SPHERE,                  // wrap around sphere
+  mjWRAP_CYLINDER                 // wrap around (infinite) cylinder
+} mjtWrap;
+
+
+typedef enum mjtTrn_ {            // type of actuator transmission
+  mjTRN_JOINT         = 0,        // force on joint
+  mjTRN_JOINTINPARENT,            // force on joint, expressed in parent frame
+  mjTRN_SLIDERCRANK,              // force via slider-crank linkage
+  mjTRN_TENDON,                   // force on tendon
+  mjTRN_SITE,                     // force on site
+  mjTRN_BODY,                     // adhesion force on a body's geoms
+
+  mjTRN_UNDEFINED     = 1000      // undefined transmission type
+} mjtTrn;
+
+
+typedef enum mjtDyn_ {            // type of actuator dynamics
+  mjDYN_NONE          = 0,        // no internal dynamics; ctrl specifies force
+  mjDYN_INTEGRATOR,               // integrator: da/dt = u
+  mjDYN_FILTER,                   // linear filter: da/dt = (u-a) / tau
+  mjDYN_FILTEREXACT,              // linear filter: da/dt = (u-a) / tau, with exact integration
+  mjDYN_MUSCLE,                   // piecewise linear filter with two time constants
+  mjDYN_DCMOTOR,                  // DC motor electrical dynamics
+  mjDYN_USER                      // user-defined dynamics type
+} mjtDyn;
+
+
+typedef enum mjtGain_ {           // type of actuator gain
+  mjGAIN_FIXED        = 0,        // fixed gain
+  mjGAIN_AFFINE,                  // const + kp*length + kv*velocity
+  mjGAIN_MUSCLE,                  // muscle FLV curve computed by mju_muscleGain()
+  mjGAIN_DCMOTOR,                 // DC motor gain: K or K/R
+  mjGAIN_USER                     // user-defined gain type
+} mjtGain;
+
+
+typedef enum mjtBias_ {           // type of actuator bias
+  mjBIAS_NONE         = 0,        // no bias
+  mjBIAS_AFFINE,                  // const + kp*length + kv*velocity
+  mjBIAS_MUSCLE,                  // muscle passive force computed by mju_muscleBias()
+  mjBIAS_DCMOTOR,                 // DC motor bias: back-EMF, cogging, LuGre friction
+  mjBIAS_USER                     // user-defined bias type
+} mjtBias;
+
+
+typedef enum mjtObj_ {            // type of MujoCo object
+  mjOBJ_UNKNOWN       = 0,        // unknown object type
+  mjOBJ_BODY,                     // body
+  mjOBJ_XBODY,                    // body, used to access regular frame instead of i-frame
+  mjOBJ_JOINT,                    // joint
+  mjOBJ_DOF,                      // dof
+  mjOBJ_GEOM,                     // geom
+  mjOBJ_SITE,                     // site
+  mjOBJ_CAMERA,                   // camera
+  mjOBJ_LIGHT,                    // light
+  mjOBJ_FLEX,                     // flex
+  mjOBJ_MESH,                     // mesh
+  mjOBJ_SKIN,                     // skin
+  mjOBJ_HFIELD,                   // heightfield
+  mjOBJ_TEXTURE,                  // texture
+  mjOBJ_MATERIAL,                 // material for rendering
+  mjOBJ_PAIR,                     // geom pair to include
+  mjOBJ_EXCLUDE,                  // body pair to exclude
+  mjOBJ_EQUALITY,                 // equality constraint
+  mjOBJ_TENDON,                   // tendon
+  mjOBJ_ACTUATOR,                 // actuator
+  mjOBJ_SENSOR,                   // sensor
+  mjOBJ_NUMERIC,                  // numeric
+  mjOBJ_TEXT,                     // text
+  mjOBJ_TUPLE,                    // tuple
+  mjOBJ_KEY,                      // keyframe
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT,                      // number of object types
+
+  // meta elements, do not appear in mjModel
+  mjOBJ_FRAME         = 100,      // frame
+  mjOBJ_DEFAULT,                  // default
+  mjOBJ_MODEL                     // entire model
+} mjtObj;
+
+
+typedef enum mjtSensor_ {         // type of sensor
+  // common robotic sensors, attached to a site
+  mjSENS_TOUCH        = 0,        // scalar contact normal forces summed over sensor zone
+  mjSENS_ACCELEROMETER,           // 3D linear acceleration, in local frame
+  mjSENS_VELOCIMETER,             // 3D linear velocity, in local frame
+  mjSENS_GYRO,                    // 3D angular velocity, in local frame
+  mjSENS_FORCE,                   // 3D force between site's body and its parent body
+  mjSENS_TORQUE,                  // 3D torque between site's body and its parent body
+  mjSENS_MAGNETOMETER,            // 3D magnetometer
+  mjSENS_RANGEFINDER,             // scalar distance to nearest geom along z-axis
+  mjSENS_CAMPROJECTION,           // pixel coordinates of a site in the camera image
+
+  // sensors related to scalar joints, tendons, actuators
+  mjSENS_JOINTPOS,                // scalar joint position (hinge and slide only)
+  mjSENS_JOINTVEL,                // scalar joint velocity (hinge and slide only)
+  mjSENS_TENDONPOS,               // scalar tendon position
+  mjSENS_TENDONVEL,               // scalar tendon velocity
+  mjSENS_ACTUATORPOS,             // scalar actuator position
+  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
+  mjSENS_BALLANGVEL,              // 3D ball joint angular velocity
+
+  // joint and tendon limit sensors, in constraint space
+  mjSENS_JOINTLIMITPOS,           // joint limit distance-margin
+  mjSENS_JOINTLIMITVEL,           // joint limit velocity
+  mjSENS_JOINTLIMITFRC,           // joint limit force
+  mjSENS_TENDONLIMITPOS,          // tendon limit distance-margin
+  mjSENS_TENDONLIMITVEL,          // tendon limit velocity
+  mjSENS_TENDONLIMITFRC,          // tendon limit force
+
+  // sensors attached to an object with spatial frame: (x)body, geom, site, camera
+  mjSENS_FRAMEPOS,                // 3D position
+  mjSENS_FRAMEQUAT,               // 4D unit quaternion orientation
+  mjSENS_FRAMEXAXIS,              // 3D unit vector: x-axis of object's frame
+  mjSENS_FRAMEYAXIS,              // 3D unit vector: y-axis of object's frame
+  mjSENS_FRAMEZAXIS,              // 3D unit vector: z-axis of object's frame
+  mjSENS_FRAMELINVEL,             // 3D linear velocity
+  mjSENS_FRAMEANGVEL,             // 3D angular velocity
+  mjSENS_FRAMELINACC,             // 3D linear acceleration
+  mjSENS_FRAMEANGACC,             // 3D angular acceleration
+
+  // sensors related to kinematic subtrees; attached to a body (which is the subtree root)
+  mjSENS_SUBTREECOM,              // 3D center of mass of subtree
+  mjSENS_SUBTREELINVEL,           // 3D linear velocity of subtree
+  mjSENS_SUBTREEANGMOM,           // 3D angular momentum of subtree
+
+  // sensors of geometric relationships
+  mjSENS_INSIDESITE,              // 1 if object is inside a site, 0 otherwise
+  mjSENS_GEOMDIST,                // signed distance between two geoms
+  mjSENS_GEOMNORMAL,              // normal direction between two geoms
+  mjSENS_GEOMFROMTO,              // segment between two geoms
+
+  // sensors for reporting contacts which occurred during the simulation
+  mjSENS_CONTACT,                 // contacts which occurred during the simulation
+
+  // global sensors
+  mjSENS_E_POTENTIAL,             // potential energy
+  mjSENS_E_KINETIC,               // kinetic energy
+  mjSENS_CLOCK,                   // simulation time
+
+  // sensors related to SDFs
+  mjSENS_TACTILE,                 // tactile sensor
+
+  // plugin-controlled sensors
+  mjSENS_PLUGIN,                  // plugin-controlled
+
+  // user-defined sensor
+  mjSENS_USER                     // sensor data provided by mjcb_sensor callback
+} mjtSensor;
+
+
+typedef enum mjtStage_ {          // computation stage
+  mjSTAGE_NONE        = 0,        // no computations
+  mjSTAGE_POS,                    // position-dependent computations
+  mjSTAGE_VEL,                    // velocity-dependent computations
+  mjSTAGE_ACC                     // acceleration/force-dependent computations
+} mjtStage;
+
+
+typedef enum mjtDataType_ {       // data type for sensors
+  mjDATATYPE_REAL     = 0,        // real values, no constraints
+  mjDATATYPE_POSITIVE,            // positive values; 0 or negative: inactive
+  mjDATATYPE_AXIS,                // 3D unit vector
+  mjDATATYPE_QUATERNION           // unit quaternion
+} mjtDataType;
+
+
+typedef enum mjtConDataField_ {   // data fields returned by contact sensors
+  mjCONDATA_FOUND     = 0,        // whether a contact was found
+  mjCONDATA_FORCE,                // contact force
+  mjCONDATA_TORQUE,               // contact torque
+  mjCONDATA_DIST,                 // contact penetration distance
+  mjCONDATA_POS,                  // contact position
+  mjCONDATA_NORMAL,               // contact frame normal
+  mjCONDATA_TANGENT,              // contact frame first tangent
+
+  mjNCONDATA                      // number of contact sensor data fields
+} mjtConDataField;
+
+
+typedef enum mjtRayDataField_ {   // data fields returned by rangefinder sensors
+  mjRAYDATA_DIST     = 0,         // distance from ray origin to nearest surface
+  mjRAYDATA_DIR,                  // normalized ray direction
+  mjRAYDATA_ORIGIN,               // ray origin
+  mjRAYDATA_POINT,                // point at which ray intersects nearest surface
+  mjRAYDATA_NORMAL,               // surface normal at intersection point
+  mjRAYDATA_DEPTH,                // depth along z-axis
+
+  mjNRAYDATA                      // number of rangefinder sensor data fields
+} mjtRayDataField;
+
+
+typedef enum mjtCamOutBit_ {      // camera output type bitflags
+  mjCAMOUT_RGB        = 1<<0,     // RGB image
+  mjCAMOUT_DEPTH      = 1<<1,     // depth image (distance from camera plane)
+  mjCAMOUT_DIST       = 1<<2,     // distance image (distance from camera origin)
+  mjCAMOUT_NORMAL     = 1<<3,     // normal image
+  mjCAMOUT_SEG        = 1<<4,     // segmentation image
+
+  mjNCAMOUT           = 5         // number of camera output types
+} mjtCamOutBit;
+
+
+typedef enum mjtSameFrame_ {      // frame alignment of bodies with their children
+  mjSAMEFRAME_NONE    = 0,        // no alignment
+  mjSAMEFRAME_BODY,               // frame is same as body frame
+  mjSAMEFRAME_INERTIA,            // frame is same as inertial frame
+  mjSAMEFRAME_BODYROT,            // frame orientation is same as body orientation
+  mjSAMEFRAME_INERTIAROT          // frame orientation is same as inertia orientation
+} mjtSameFrame;
+
+
+typedef enum mjtSleepPolicy_ {    // per-tree sleep policy
+  mjSLEEP_AUTO        = 0,        // compiler chooses sleep policy
+  mjSLEEP_AUTO_NEVER,             // compiler sleep policy: never
+  mjSLEEP_AUTO_ALLOWED,           // compiler sleep policy: allowed
+  mjSLEEP_NEVER,                  // user sleep policy: never
+  mjSLEEP_ALLOWED,                // user sleep policy: allowed
+  mjSLEEP_INIT,                   // user sleep policy: initialized asleep
+} mjtSleepPolicy;
+
+
+typedef enum mjtLRMode_ {         // mode for actuator length range computation
+  mjLRMODE_NONE       = 0,        // do not process any actuators
+  mjLRMODE_MUSCLE,                // process muscle actuators
+  mjLRMODE_MUSCLEUSER,            // process muscle and user actuators
+  mjLRMODE_ALL                    // process all actuators
+} mjtLRMode;
+
+
+typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
+  mjFLEXSELF_NONE     = 0,        // no self-collisions
+  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
+  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
+  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;
+
+
+
+//---------------------------------- enum types (mjData) -------------------------------------------
+
+typedef enum mjtState_ {            // state elements
+  mjSTATE_TIME           = 1<<0,    // time
+  mjSTATE_QPOS           = 1<<1,    // position
+  mjSTATE_QVEL           = 1<<2,    // velocity
+  mjSTATE_ACT            = 1<<3,    // actuator activation
+  mjSTATE_HISTORY        = 1<<4,    // history buffers (control, sensor)
+  mjSTATE_WARMSTART      = 1<<5,    // acceleration used for warmstart
+  mjSTATE_CTRL           = 1<<6,    // control
+  mjSTATE_QFRC_APPLIED   = 1<<7,    // applied generalized force
+  mjSTATE_XFRC_APPLIED   = 1<<8,    // applied Cartesian force/torque
+  mjSTATE_EQ_ACTIVE      = 1<<9,    // enable/disable constraints
+  mjSTATE_MOCAP_POS      = 1<<10,   // positions of mocap bodies
+  mjSTATE_MOCAP_QUAT     = 1<<11,   // orientations of mocap bodies
+  mjSTATE_USERDATA       = 1<<12,   // user data
+  mjSTATE_PLUGIN         = 1<<13,   // plugin state
+
+  mjNSTATE               = 14,      // number of state elements
+
+  // convenience values for commonly used state specifications
+  mjSTATE_PHYSICS        = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
+  mjSTATE_FULLPHYSICS    = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
+  mjSTATE_USER           = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
+                           mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
+                           mjSTATE_USERDATA,
+  mjSTATE_INTEGRATION    = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
+} mjtState;
+
+
+typedef enum mjtConstraint_ {       // type of constraint
+  mjCNSTR_EQUALITY       = 0,       // equality constraint
+  mjCNSTR_FRICTION_DOF,             // dof friction
+  mjCNSTR_FRICTION_TENDON,          // tendon friction
+  mjCNSTR_LIMIT_JOINT,              // joint limit
+  mjCNSTR_LIMIT_TENDON,             // tendon limit
+  mjCNSTR_CONTACT_FRICTIONLESS,     // frictionless contact
+  mjCNSTR_CONTACT_PYRAMIDAL,        // frictional contact, pyramidal friction cone
+  mjCNSTR_CONTACT_ELLIPTIC          // frictional contact, elliptic friction cone
+} mjtConstraint;
+
+
+typedef enum mjtConstraintState_ {  // constraint state
+  mjCNSTRSTATE_SATISFIED = 0,       // constraint satisfied, zero cost (limit, contact)
+  mjCNSTRSTATE_QUADRATIC,           // quadratic cost (equality, friction, limit, contact)
+  mjCNSTRSTATE_LINEARNEG,           // linear cost, negative side (friction)
+  mjCNSTRSTATE_LINEARPOS,           // linear cost, positive side (friction)
+  mjCNSTRSTATE_CONE                 // squared distance to cone cost (elliptic contact)
+} mjtConstraintState;
+
+
+typedef enum mjtWarning_ {          // warning types
+  mjWARN_INERTIA         = 0,       // (near) singular inertia matrix
+  mjWARN_CONTACTFULL,               // too many contacts in contact list
+  mjWARN_CNSTRFULL,                 // too many constraints
+  mjWARN_BADQPOS,                   // bad number in qpos
+  mjWARN_BADQVEL,                   // bad number in qvel
+  mjWARN_BADQACC,                   // bad number in qacc
+  mjWARN_BADCTRL,                   // bad number in ctrl
+
+  mjNWARNING                        // number of warnings
+} mjtWarning;
+
+
+typedef enum mjtTimer_ {            // internal timers
+  // main api
+  mjTIMER_STEP           = 0,       // step
+  mjTIMER_FORWARD,                  // forward
+  mjTIMER_INVERSE,                  // inverse
+
+  // breakdown of step/forward
+  mjTIMER_POSITION,                 // fwdPosition
+  mjTIMER_VELOCITY,                 // fwdVelocity
+  mjTIMER_ACTUATION,                // fwdActuation
+  mjTIMER_CONSTRAINT,               // fwdConstraint
+  mjTIMER_ADVANCE,                  // mj_Euler, mj_implicit
+
+  // breakdown of fwdPosition
+  mjTIMER_POS_KINEMATICS,           // kinematics, com, tendon, transmission
+  mjTIMER_POS_INERTIA,              // inertia computations
+  mjTIMER_POS_COLLISION,            // collision detection
+  mjTIMER_POS_MAKE,                 // make constraints
+  mjTIMER_POS_PROJECT,              // project constraints
+
+  // breakdown of mj_collision
+  mjTIMER_COL_BROAD,                // broadphase
+  mjTIMER_COL_NARROW,               // narrowphase
+
+  mjNTIMER                          // number of timers
+} mjtTimer;
+
+
+typedef enum mjtSleepState_ {       // sleep state of an object
+  mjS_STATIC = -1,                  // object is static
+  mjS_ASLEEP = 0,                   // object is asleep
+  mjS_AWAKE  = 1                    // object is awake
+} mjtSleepState;
+
+#endif  // MUJOCO_INCLUDE_MJTYPE_H_
diff --git a/include/mujoco/mjvisualize.h b/include/mujoco/mjvisualize.h
index 68e22d21..2379f82b 100644
--- a/include/mujoco/mjvisualize.h
+++ b/include/mujoco/mjvisualize.h
@@ -17,7 +17,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 
 #define mjNGROUP        6         // number of geom, site, joint, skin groups with visflags
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index ab41686d..213afeb3 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -32,7 +32,7 @@
 #include <mujoco/mjsan.h>
 #include <mujoco/mjspec.h>
 #include <mujoco/mjthread.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
 
diff --git a/plugin/actuator/pid.h b/plugin/actuator/pid.h
index 984bf47c..0b1e831d 100644
--- a/plugin/actuator/pid.h
+++ b/plugin/actuator/pid.h
@@ -21,7 +21,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 namespace mujoco::plugin::actuator {
 
diff --git a/plugin/elasticity/cable.cc b/plugin/elasticity/cable.cc
index d9746439..5fe190d3 100644
--- a/plugin/elasticity/cable.cc
+++ b/plugin/elasticity/cable.cc
@@ -18,7 +18,7 @@
 #include <optional>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "cable.h"
 
diff --git a/plugin/elasticity/cable.h b/plugin/elasticity/cable.h
index bf2bf75d..e73d6c97 100644
--- a/plugin/elasticity/cable.h
+++ b/plugin/elasticity/cable.h
@@ -20,7 +20,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 
 
diff --git a/plugin/sdf/bolt.cc b/plugin/sdf/bolt.cc
index 4f052c06..c38cffe1 100644
--- a/plugin/sdf/bolt.cc
+++ b/plugin/sdf/bolt.cc
@@ -19,7 +19,7 @@
 #include <utility>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
 #include "bolt.h"
diff --git a/plugin/sdf/bolt.h b/plugin/sdf/bolt.h
index af04ba9a..71517b94 100644
--- a/plugin/sdf/bolt.h
+++ b/plugin/sdf/bolt.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
 
diff --git a/plugin/sdf/bowl.cc b/plugin/sdf/bowl.cc
index 2ebbc5ac..73fb8205 100644
--- a/plugin/sdf/bowl.cc
+++ b/plugin/sdf/bowl.cc
@@ -17,7 +17,7 @@
 #include <utility>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
 #include "bowl.h"
diff --git a/plugin/sdf/bowl.h b/plugin/sdf/bowl.h
index 471f685c..98eacece 100644
--- a/plugin/sdf/bowl.h
+++ b/plugin/sdf/bowl.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
 
diff --git a/plugin/sdf/gear.cc b/plugin/sdf/gear.cc
index 54288c21..ce3755fb 100644
--- a/plugin/sdf/gear.cc
+++ b/plugin/sdf/gear.cc
@@ -18,7 +18,7 @@
 #include <utility>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
diff --git a/plugin/sdf/gear.h b/plugin/sdf/gear.h
index f1f63a75..be0e6f89 100644
--- a/plugin/sdf/gear.h
+++ b/plugin/sdf/gear.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
 
diff --git a/plugin/sdf/nut.cc b/plugin/sdf/nut.cc
index 117df406..38888fbf 100644
--- a/plugin/sdf/nut.cc
+++ b/plugin/sdf/nut.cc
@@ -19,7 +19,7 @@
 #include <utility>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
 #include "nut.h"
diff --git a/plugin/sdf/nut.h b/plugin/sdf/nut.h
index 9d7377f1..933fc931 100644
--- a/plugin/sdf/nut.h
+++ b/plugin/sdf/nut.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include "sdf.h"
 
diff --git a/plugin/sdf/torus.cc b/plugin/sdf/torus.cc
index d48feb51..23f13ddd 100644
--- a/plugin/sdf/torus.cc
+++ b/plugin/sdf/torus.cc
@@ -17,7 +17,7 @@
 #include <utility>
 
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "sdf.h"
 #include "torus.h"
diff --git a/plugin/sdf/torus.h b/plugin/sdf/torus.h
index 8966d984..866160a0 100644
--- a/plugin/sdf/torus.h
+++ b/plugin/sdf/torus.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "sdf.h"
 
 namespace mujoco::plugin::sdf {
diff --git a/plugin/sensor/touch_grid.cc b/plugin/sensor/touch_grid.cc
index 98ba90f5..044bccdd 100644
--- a/plugin/sensor/touch_grid.cc
+++ b/plugin/sensor/touch_grid.cc
@@ -25,7 +25,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 
diff --git a/plugin/sensor/touch_grid.h b/plugin/sensor/touch_grid.h
index 34ad43ee..c0029065 100644
--- a/plugin/sensor/touch_grid.h
+++ b/plugin/sensor/touch_grid.h
@@ -20,7 +20,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 
 namespace mujoco::plugin::sensor {
diff --git a/python/mujoco/constants.cc b/python/mujoco/constants.cc
index 49c3ee58..593070d2 100644
--- a/python/mujoco/constants.cc
+++ b/python/mujoco/constants.cc
@@ -83,7 +83,7 @@ PYBIND11_MODULE(_constants, pymodule) {
   // from mujoco.h
   X(mjVERSION_HEADER);
 
-  // from mjtnum.h
+  // from mjtype.h
   X(mjMINVAL);
 
   #undef X
diff --git a/python/mujoco/introspect/enums.py b/python/mujoco/introspect/enums.py
index 7de73447..2247a805 100644
--- a/python/mujoco/introspect/enums.py
+++ b/python/mujoco/introspect/enums.py
@@ -508,16 +508,6 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjSDFTYPE_COLLISION', 3),
          ]),
      )),
-    ('mjtTaskStatus',
-     EnumDecl(
-         name='mjtTaskStatus',
-         declname='enum mjtTaskStatus_',
-         values=dict([
-             ('mjTASK_NEW', 0),
-             ('mjTASK_QUEUED', 1),
-             ('mjTASK_COMPLETED', 2),
-         ]),
-     )),
     ('mjtState',
      EnumDecl(
          name='mjtState',
@@ -619,6 +609,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjS_AWAKE', 1),
          ]),
      )),
+    ('mjtTaskStatus',
+     EnumDecl(
+         name='mjtTaskStatus',
+         declname='enum mjtTaskStatus_',
+         values=dict([
+             ('mjTASK_NEW', 0),
+             ('mjTASK_QUEUED', 1),
+             ('mjTASK_COMPLETED', 2),
+         ]),
+     )),
     ('mjtGeomInertia',
      EnumDecl(
          name='mjtGeomInertia',
diff --git a/python/mujoco/serialization.h b/python/mujoco/serialization.h
index d38abefd..cb234ac8 100644
--- a/python/mujoco/serialization.h
+++ b/python/mujoco/serialization.h
@@ -18,7 +18,7 @@
 #include <cstddef>
 #include <iostream>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 namespace mujoco::python::_impl {
 
diff --git a/src/engine/engine_collision_convex.h b/src/engine/engine_collision_convex.h
index eae1e72d..969151cc 100644
--- a/src/engine/engine_collision_convex.h
+++ b/src/engine/engine_collision_convex.h
@@ -26,7 +26,7 @@
 #include <mujoco/mjexport.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 // minimum number of vertices to use hill-climbing in mesh support
 #define mjMESH_HILLCLIMB_MIN 10
diff --git a/src/engine/engine_collision_driver.h b/src/engine/engine_collision_driver.h
index 0e2b9edd..4734975e 100644
--- a/src/engine/engine_collision_driver.h
+++ b/src/engine/engine_collision_driver.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c
index 4b84695c..42227694 100644
--- a/src/engine/engine_collision_gjk.c
+++ b/src/engine/engine_collision_gjk.c
@@ -19,7 +19,7 @@
 #include <stdlib.h>
 #include <string.h>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjmodel.h>
 #include "engine/engine_collision_convex.h"
 #include "engine/engine_util_blas.h"
diff --git a/src/engine/engine_collision_gjk.h b/src/engine/engine_collision_gjk.h
index 35cbf475..1464d80e 100644
--- a/src/engine/engine_collision_gjk.h
+++ b/src/engine/engine_collision_gjk.h
@@ -20,7 +20,7 @@
 
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #include "engine/engine_collision_convex.h"
 
diff --git a/src/engine/engine_collision_primitive.h b/src/engine/engine_collision_primitive.h
index b598c0e4..e2df2966 100644
--- a/src/engine/engine_collision_primitive.h
+++ b/src/engine/engine_collision_primitive.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c
index c5d093c0..490523f7 100644
--- a/src/engine/engine_collision_sdf.c
+++ b/src/engine/engine_collision_sdf.c
@@ -20,7 +20,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
diff --git a/src/engine/engine_collision_sdf.h b/src/engine/engine_collision_sdf.h
index 64f9b298..79e3d722 100644
--- a/src/engine/engine_collision_sdf.h
+++ b/src/engine/engine_collision_sdf.h
@@ -19,7 +19,7 @@
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_core_constraint.h b/src/engine/engine_core_constraint.h
index 15c03b6b..9d3f9eb0 100644
--- a/src/engine/engine_core_constraint.h
+++ b/src/engine/engine_core_constraint.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_core_smooth.h b/src/engine/engine_core_smooth.h
index 258b625f..0ea923d5 100644
--- a/src/engine/engine_core_smooth.h
+++ b/src/engine/engine_core_smooth.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_core_util.h b/src/engine/engine_core_util.h
index 36ec5fd0..9f958937 100644
--- a/src/engine/engine_core_util.h
+++ b/src/engine/engine_core_util.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_forward.h b/src/engine/engine_forward.h
index 0a3fd5e0..f0ba9e40 100644
--- a/src/engine/engine_forward.h
+++ b/src/engine/engine_forward.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_init.c b/src/engine/engine_init.c
index 15fe707d..2d6855f6 100644
--- a/src/engine/engine_init.c
+++ b/src/engine/engine_init.c
@@ -17,7 +17,7 @@
 #include <string.h>
 
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "engine/engine_util_blas.h"
 
 #ifdef _MSC_VER
diff --git a/src/engine/engine_inline.h b/src/engine/engine_inline.h
index dbe4b8f3..9532271b 100644
--- a/src/engine/engine_inline.h
+++ b/src/engine/engine_inline.h
@@ -15,7 +15,7 @@
 #ifndef MUJOCO_SRC_ENGINE_ENGINE_INLINE_H_
 #define MUJOCO_SRC_ENGINE_ENGINE_INLINE_H_
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 
 #ifdef __cplusplus
diff --git a/src/engine/engine_memory.h b/src/engine/engine_memory.h
index 11459a79..bbfd836b 100644
--- a/src/engine/engine_memory.h
+++ b/src/engine/engine_memory.h
@@ -17,8 +17,8 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
-#include <mujoco/mjtnum.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
+#include <mujoco/mjtype.h>
 #include <mujoco/mjxmacro.h>
 
 #ifdef __cplusplus
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 921c028a..54caa9b8 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -24,7 +24,7 @@
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjxmacro.h>
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_util.h"
diff --git a/src/engine/engine_support.h b/src/engine/engine_support.h
index 7e1fd911..cf112cb0 100644
--- a/src/engine/engine_support.h
+++ b/src/engine/engine_support.h
@@ -18,7 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_util_blas.c b/src/engine/engine_util_blas.c
index 8d4e7c3e..e829fea0 100644
--- a/src/engine/engine_util_blas.c
+++ b/src/engine/engine_util_blas.c
@@ -16,7 +16,7 @@
 
 #include <string.h>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #include "engine/engine_util_blas_avx.h"  // IWYU pragma: keep
 
diff --git a/src/engine/engine_util_blas.h b/src/engine/engine_util_blas.h
index 4741385f..c5f83ac3 100644
--- a/src/engine/engine_util_blas.h
+++ b/src/engine/engine_util_blas.h
@@ -18,7 +18,7 @@
 #include <math.h>
 
 #include <mujoco/mjexport.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_util_blas_avx.h b/src/engine/engine_util_blas_avx.h
index ff2af6c7..f86e34cc 100644
--- a/src/engine/engine_util_blas_avx.h
+++ b/src/engine/engine_util_blas_avx.h
@@ -22,7 +22,7 @@
 
 #include <immintrin.h>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 // res = vec*scl
 static inline
diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h
index ebf2ced7..f97bf4bb 100644
--- a/src/engine/engine_util_misc.h
+++ b/src/engine/engine_util_misc.h
@@ -17,7 +17,7 @@
 
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h
index 8cdb8ea9..1f0ef8c4 100644
--- a/src/engine/engine_util_solve.h
+++ b/src/engine/engine_util_solve.h
@@ -17,7 +17,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/engine/engine_util_sparse.c b/src/engine/engine_util_sparse.c
index 0c79a27f..bd3cbc3b 100644
--- a/src/engine/engine_util_sparse.c
+++ b/src/engine/engine_util_sparse.c
@@ -17,7 +17,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "engine/engine_memory.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_misc.h"
diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h
index 038ac5ce..ba5fc5ab 100644
--- a/src/engine/engine_util_sparse.h
+++ b/src/engine/engine_util_sparse.h
@@ -19,7 +19,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "engine/engine_util_sparse_avx.h"
 
 #ifdef __cplusplus
diff --git a/src/engine/engine_util_sparse_avx.h b/src/engine/engine_util_sparse_avx.h
index 32be6020..e9edfeff 100644
--- a/src/engine/engine_util_sparse_avx.h
+++ b/src/engine/engine_util_sparse_avx.h
@@ -23,7 +23,7 @@
 #include <immintrin.h>
 #include <string.h>
 
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 
 //------------------------------ sparse operations using avx ---------------------------------------
diff --git a/src/engine/engine_util_spatial.h b/src/engine/engine_util_spatial.h
index 6a53ddfd..8b73de09 100644
--- a/src/engine/engine_util_spatial.h
+++ b/src/engine/engine_util_spatial.h
@@ -16,7 +16,7 @@
 #define MUJOCO_SRC_ENGINE_ENGINE_UTIL_SPATIAL_H_
 
 #include <mujoco/mjexport.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 #ifdef __cplusplus
 extern "C" {
diff --git a/src/user/user_api.h b/src/user/user_api.h
index f1f216a7..cf1e26a6 100644
--- a/src/user/user_api.h
+++ b/src/user/user_api.h
@@ -21,7 +21,7 @@
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 
 
 // this is a C-API
diff --git a/src/user/user_composite.cc b/src/user/user_composite.cc
index b62d7c24..d1bebf58 100644
--- a/src/user/user_composite.cc
+++ b/src/user/user_composite.cc
@@ -21,7 +21,7 @@
 
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "cc/array_safety.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
diff --git a/src/user/user_flexcomp.cc b/src/user/user_flexcomp.cc
index 7030768c..f85aa413 100644
--- a/src/user/user_flexcomp.cc
+++ b/src/user/user_flexcomp.cc
@@ -26,7 +26,7 @@
 
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjplugin.h>
 #include "cc/array_safety.h"
 #include "engine/engine_crossplatform.h"
diff --git a/src/user/user_mesh.cc b/src/user/user_mesh.cc
index fa08d766..17c4c9b6 100644
--- a/src/user/user_mesh.cc
+++ b/src/user/user_mesh.cc
@@ -57,7 +57,7 @@
 #include <mujoco/mjmacro.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "engine/engine_crossplatform.h"  // IWYU pragma: keep
 #include "engine/engine_plugin.h"
 #include "engine/engine_util_errmem.h"
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index c9e38131..02b2b2b9 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -41,7 +41,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "cc/array_safety.h"
 #include "engine/engine_core_util.h"
diff --git a/src/user/user_model.h b/src/user/user_model.h
index 6d554649..3af7167e 100644
--- a/src/user/user_model.h
+++ b/src/user/user_model.h
@@ -30,7 +30,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "user/user_objects.h"
 
 typedef std::map<std::string, int, std::less<> > mjKeyMap;
diff --git a/src/user/user_objects.h b/src/user/user_objects.h
index cbe85e6b..ed93cf48 100644
--- a/src/user/user_objects.h
+++ b/src/user/user_objects.h
@@ -32,7 +32,7 @@
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "user/user_cache.h"
 #include "user/user_util.h"
 
diff --git a/src/xml/xml_native_reader.cc b/src/xml/xml_native_reader.cc
index 9b7ee999..7bb36644 100644
--- a/src/xml/xml_native_reader.cc
+++ b/src/xml/xml_native_reader.cc
@@ -31,7 +31,7 @@
 #include <mujoco/mujoco.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjplugin.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include "engine/engine_plugin.h"
 #include "engine/engine_support.h"
diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc
index 82ac3095..21fec2cf 100644
--- a/test/engine/engine_collision_gjk_test.cc
+++ b/test/engine/engine_collision_gjk_test.cc
@@ -24,7 +24,7 @@
 
 #include "src/engine/engine_collision_convex.h"
 #include <mujoco/mujoco.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "test/fixture.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
diff --git a/test/engine/engine_collision_sdf_test.cc b/test/engine/engine_collision_sdf_test.cc
index bcc0d1c3..1d4cf1ce 100644
--- a/test/engine/engine_collision_sdf_test.cc
+++ b/test/engine/engine_collision_sdf_test.cc
@@ -19,7 +19,7 @@
 #include <gtest/gtest.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_collision_sdf.h"
 #include "test/fixture.h"
diff --git a/test/engine/engine_forward_test.cc b/test/engine/engine_forward_test.cc
index bfcba7b6..401358c1 100644
--- a/test/engine/engine_forward_test.cc
+++ b/test/engine/engine_forward_test.cc
@@ -28,7 +28,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include <mujoco/mjxmacro.h>
 #include "src/cc/array_safety.h"
diff --git a/test/engine/engine_ray_test.cc b/test/engine/engine_ray_test.cc
index 7c78c909..0078c7db 100644
--- a/test/engine/engine_ray_test.cc
+++ b/test/engine/engine_ray_test.cc
@@ -21,7 +21,7 @@
 #include <gtest/gtest.h>
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_ray.h"
 #include "test/fixture.h"
diff --git a/test/engine/engine_sensor_test.cc b/test/engine/engine_sensor_test.cc
index 1474c856..a12fdec0 100644
--- a/test/engine/engine_sensor_test.cc
+++ b/test/engine/engine_sensor_test.cc
@@ -22,7 +22,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_support.h"
 #include "src/engine/engine_util_blas.h"
diff --git a/test/engine/engine_util_blas_test.cc b/test/engine/engine_util_blas_test.cc
index 2ef25e73..595112ee 100644
--- a/test/engine/engine_util_blas_test.cc
+++ b/test/engine/engine_util_blas_test.cc
@@ -20,7 +20,7 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include "test/fixture.h"
 
 namespace mujoco {
diff --git a/test/engine/engine_util_spatial_test.cc b/test/engine/engine_util_spatial_test.cc
index c0b7f4c6..5eda98d1 100644
--- a/test/engine/engine_util_spatial_test.cc
+++ b/test/engine/engine_util_spatial_test.cc
@@ -21,7 +21,7 @@
 #include <gtest/gtest.h>
 #include <gtest/gtest-spi.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_util_blas.h"
 #include "src/engine/engine_util_spatial.h"
diff --git a/test/header_test.cc b/test/header_test.cc
index 98cae8ca..98014a5f 100644
--- a/test/header_test.cc
+++ b/test/header_test.cc
@@ -24,7 +24,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
 #include <mujoco/mjrender.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mjxmacro.h>
diff --git a/test/user/user_mesh_test.cc b/test/user/user_mesh_test.cc
index 185b2678..9ea3a39a 100644
--- a/test/user/user_mesh_test.cc
+++ b/test/user/user_mesh_test.cc
@@ -28,7 +28,7 @@
 #include <absl/strings/str_format.h>
 #include <absl/strings/str_replace.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/cc/array_safety.h"
 #include "test/fixture.h"
diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc
index d71ba4cb..8d9f5839 100644
--- a/test/user/user_objects_test.cc
+++ b/test/user/user_objects_test.cc
@@ -24,7 +24,7 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
 
diff --git a/test/xml/xml_native_writer_test.cc b/test/xml/xml_native_writer_test.cc
index 9a70372f..afa8215b 100644
--- a/test/xml/xml_native_writer_test.cc
+++ b/test/xml/xml_native_writer_test.cc
@@ -31,7 +31,7 @@
 #include <gtest/gtest.h>
 #include <absl/strings/match.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjtnum.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mujoco.h>
 #include "src/xml/xml_numeric_format.h"
 #include "test/fixture.h"
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index e66b47a3..999c50ef 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -64,7 +64,7 @@ public const bool ADDRESS_SANITIZER = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSPEC_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true;
 public const int mjMAXTHREAD = 128;
-public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTNUM_H_ = true;
+public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTYPE_H_ = true;
 public const double mjMINVAL = 1e-15;
 public const bool THIRD_PARTY_MUJOCO_MJUI_H_ = true;
 public const int mjMAXUISECT = 10;
@@ -443,11 +443,6 @@ public enum mjtSDFType : int{
   mjSDFTYPE_MIDSURFACE = 2,
   mjSDFTYPE_COLLISION = 3,
 }
-public enum mjtTaskStatus : int{
-  mjTASK_NEW = 0,
-  mjTASK_QUEUED = 1,
-  mjTASK_COMPLETED = 2,
-}
 public enum mjtState : int{
   mjSTATE_TIME = 1,
   mjSTATE_QPOS = 2,
@@ -519,6 +514,11 @@ public enum mjtSleepState : int{
   mjS_ASLEEP = 0,
   mjS_AWAKE = 1,
 }
+public enum mjtTaskStatus : int{
+  mjTASK_NEW = 0,
+  mjTASK_QUEUED = 1,
+  mjTASK_COMPLETED = 2,
+}
 public enum mjtGeomInertia : int{
   mjINERTIA_VOLUME = 0,
   mjINERTIA_SHELL = 1,

From 88a293b5e406fdba7e33d78495d508bbd0b609c0 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 21 May 2026 08:43:24 -0700
Subject: [PATCH 12/61] Update docs to correctly reference mjtype.h

PiperOrigin-RevId: 919071817
Change-Id: I829de7a22a49e6df604b0acf601a13f67ef5fbab
---
 doc/APIreference/APItypes.rst | 4 ++--
 doc/overview.rst              | 4 ++--
 doc/programming/index.rst     | 5 ++---
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst
index f7779b3e..879d4bda 100644
--- a/doc/APIreference/APItypes.rst
+++ b/doc/APIreference/APItypes.rst
@@ -115,7 +115,7 @@ All enum types use the ``mjt`` prefix.
 Model
 ^^^^^
 
-The enums below are defined in `mjmodel.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__.
+The enums below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
 
 
 .. _mjtDisableBit:
@@ -456,7 +456,7 @@ Formulas used to combine SDFs when calling mjc_distance and mjc_gradient.
 Data
 ^^^^
 
-The enums below are defined in `mjdata.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__.
+The enums below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
 
 
 
diff --git a/doc/overview.rst b/doc/overview.rst
index 95171823..f67f5f84 100644
--- a/doc/overview.rst
+++ b/doc/overview.rst
@@ -875,8 +875,8 @@ The size of each array (``njnt`` in this case) is also given in ``mjModel``. The
 first, followed by the limits of the second joint etc. This ordering reflects the fact that all matrices in MuJoCo have
 row-major format.
 
-The available element types are defined in `mjmodel.h
-<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h#L237>`_, in the enum type :ref:`mjtObj`.
+The available element types are defined in `mjtype.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__, in the enum type :ref:`mjtObj`.
 These enums are mostly used internally. One exception are the functions :ref:`mj_name2id` and :ref:`mj_id2name` in the
 MuJoCo API, which map element names to integer ids and vice versa. These functions take an element type as input.
 
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index 035fdaca..ba6e5ea6 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -154,11 +154,10 @@ links below, to make this documentation self-contained.
    global variables, and includes all other header files except mjxmacro.h.
 `mjmodel.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__
    Defines the C structure :ref:`mjModel` which is the runtime representation of the
-   model being simulated. It also defines a number of primitive types and other structures needed to define mjModel.
+   model being simulated.
 `mjdata.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__
    Defines the C structure :ref:`mjData` which is the workspace where all computations
-   read their inputs and write their outputs. It also defines primitive types and other structures needed to define
-   mjData.
+   read their inputs and write their outputs.
 `mjvisualize.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjvisualize.h>`__
    Defines the primitive types and structures needed by the abstract visualizer.
 `mjrender.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjrender.h>`__

From 585be6cee5723712158dfd388b8361ea118757c1 Mon Sep 17 00:00:00 2001
From: Matija Kecman <matijak@google.com>
Date: Thu, 21 May 2026 09:41:29 -0700
Subject: [PATCH 13/61] Fix wireframe rendering

PiperOrigin-RevId: 919100432
Change-Id: I2931d793ab9639bdaaad11f6d77a77aa14e43b8a
---
 src/experimental/platform/hal/renderer.cc | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/experimental/platform/hal/renderer.cc b/src/experimental/platform/hal/renderer.cc
index bcb44c19..591965dd 100644
--- a/src/experimental/platform/hal/renderer.cc
+++ b/src/experimental/platform/hal/renderer.cc
@@ -232,6 +232,8 @@ void Renderer::DoRender(int width, int height) {
         draw_mode = mjDRAW_MODE_SEGMENTATION;
       } else if (scene_.flags[mjRND_DEPTH]) {
         draw_mode = mjDRAW_MODE_DEPTH;
+      } else if (scene_.flags[mjRND_WIREFRAME]) {
+        draw_mode = mjDRAW_MODE_WIREFRAME;
       }
 
       mjrRenderRequest reqs[2];
@@ -285,6 +287,8 @@ void Renderer::DoReadPixels(int width, int height, unsigned char* rgb) {
       draw_mode = mjDRAW_MODE_SEGMENTATION;
     } else if (scene_.flags[mjRND_DEPTH]) {
       draw_mode = mjDRAW_MODE_DEPTH;
+    } else if (scene_.flags[mjRND_WIREFRAME]) {
+      draw_mode = mjDRAW_MODE_WIREFRAME;
     }
 
     mjrRenderTargetConfig config;

From 15d27b363f031fc7c8ce1b42eef737e51073ec8d Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 21 May 2026 09:54:33 -0700
Subject: [PATCH 14/61] Compile-time size assertions for MuJoCo basic types

PiperOrigin-RevId: 919107499
Change-Id: I07fcba4f33c77a2a1e68528dba6d399af94d0220
---
 CMakeLists.txt                       |   1 +
 doc/changelog.rst                    |   6 ++
 doc/programming/index.rst            |   2 +
 include/mujoco/mjassert.h            | 133 +++++++++++++++++++++++++++
 include/mujoco/mujoco.h              |   1 +
 test/header_test.cc                  |  55 +----------
 unity/Runtime/Bindings/MjBindings.cs |   1 +
 7 files changed, 145 insertions(+), 54 deletions(-)
 create mode 100644 include/mujoco/mjassert.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 86fc48a1..e01c6621 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -108,6 +108,7 @@ include(MujocoMacOS)
 include(MujocoDependencies)
 
 set(MUJOCO_HEADERS
+    include/mujoco/mjassert.h
     include/mujoco/mjdata.h
     include/mujoco/mjexport.h
     include/mujoco/mjmacro.h
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 5cd5f9c1..b339da1f 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -80,6 +80,12 @@ General
    - MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
      ``naconmax``.
 
+   - Maybe-breaking: Added `mjassert.h
+     <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__, a new header containing
+     compile-time assertions that verify the sizes of MuJoCo's public types for ABI stability. This is a first step
+     towards replacing ``int`` with strongly-typed enums in the public API. If these assertions fail on your compiler or
+     platform, please report the issue on GitHub.
+
 Version 3.8.1 (May 11, 2026)
 ----------------------------
 
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index ba6e5ea6..852685cb 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -184,6 +184,8 @@ links below, to make this documentation self-contained.
    code.
 `mjsan.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjsan.h>`__
    Definitions required when building with sanitizer instrumentation.
+`mjassert.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__
+   Compile-time size assertions verifying MuJoCo ABI stability across C and C++ compilers.
 
 .. _inVersion:
 
diff --git a/include/mujoco/mjassert.h b/include/mujoco/mjassert.h
new file mode 100644
index 00000000..5c69f940
--- /dev/null
+++ b/include/mujoco/mjassert.h
@@ -0,0 +1,133 @@
+// Copyright 2026 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_MJASSERT_H_
+#define MUJOCO_MJASSERT_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjplugin.h>
+#include <mujoco/mjrender.h>
+#include <mujoco/mjspec.h>
+#include <mujoco/mjthread.h>
+#include <mujoco/mjtype.h>
+#include <mujoco/mjui.h>
+#include <mujoco/mjvisualize.h>
+
+#if defined(__cplusplus)
+  #define MJ_ASSERT_SIZE(type, size) \
+    static_assert(sizeof(type) == (size), #type " must be " #size " bytes for MuJoCo ABI stability")
+#elif defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L
+  #define MJ_ASSERT_SIZE(type, size) \
+    _Static_assert(sizeof(type) == (size), #type " must be " #size " bytes for MuJoCo ABI stability")
+#else
+  #define MJ_ASSERT_SIZE(type, size) \
+    typedef char mj_assert_##type[sizeof(type) == (size) ? 1 : -1]
+#endif
+
+// mjtnum.h
+MJ_ASSERT_SIZE(mjtByte, 1);
+MJ_ASSERT_SIZE(mjtSize, 8);
+#if !defined(mjUSESINGLE)
+MJ_ASSERT_SIZE(mjtNum, 8);
+#else
+MJ_ASSERT_SIZE(mjtNum, 4);
+#endif
+
+// mjmodel.h
+MJ_ASSERT_SIZE(mjtDisableBit, 4);
+MJ_ASSERT_SIZE(mjtEnableBit, 4);
+MJ_ASSERT_SIZE(mjtJoint, 4);
+MJ_ASSERT_SIZE(mjtGeom, 4);
+MJ_ASSERT_SIZE(mjtProjection, 4);
+MJ_ASSERT_SIZE(mjtCamLight, 4);
+MJ_ASSERT_SIZE(mjtLightType, 4);
+MJ_ASSERT_SIZE(mjtTexture, 4);
+MJ_ASSERT_SIZE(mjtTextureRole, 4);
+MJ_ASSERT_SIZE(mjtColorSpace, 4);
+MJ_ASSERT_SIZE(mjtIntegrator, 4);
+MJ_ASSERT_SIZE(mjtCone, 4);
+MJ_ASSERT_SIZE(mjtJacobian, 4);
+MJ_ASSERT_SIZE(mjtSolver, 4);
+MJ_ASSERT_SIZE(mjtEq, 4);
+MJ_ASSERT_SIZE(mjtWrap, 4);
+MJ_ASSERT_SIZE(mjtTrn, 4);
+MJ_ASSERT_SIZE(mjtDyn, 4);
+MJ_ASSERT_SIZE(mjtGain, 4);
+MJ_ASSERT_SIZE(mjtBias, 4);
+MJ_ASSERT_SIZE(mjtObj, 4);
+MJ_ASSERT_SIZE(mjtSensor, 4);
+MJ_ASSERT_SIZE(mjtStage, 4);
+MJ_ASSERT_SIZE(mjtDataType, 4);
+MJ_ASSERT_SIZE(mjtConDataField, 4);
+MJ_ASSERT_SIZE(mjtRayDataField, 4);
+MJ_ASSERT_SIZE(mjtCamOutBit, 4);
+MJ_ASSERT_SIZE(mjtSameFrame, 4);
+MJ_ASSERT_SIZE(mjtSleepPolicy, 4);
+MJ_ASSERT_SIZE(mjtLRMode, 4);
+MJ_ASSERT_SIZE(mjtFlexSelf, 4);
+MJ_ASSERT_SIZE(mjtSDFType, 4);
+
+// mjdata.h
+MJ_ASSERT_SIZE(mjtState, 4);
+MJ_ASSERT_SIZE(mjtConstraint, 4);
+MJ_ASSERT_SIZE(mjtConstraintState, 4);
+MJ_ASSERT_SIZE(mjtWarning, 4);
+MJ_ASSERT_SIZE(mjtTimer, 4);
+MJ_ASSERT_SIZE(mjtSleepState, 4);
+
+// mjspec.h
+MJ_ASSERT_SIZE(mjtGeomInertia, 4);
+MJ_ASSERT_SIZE(mjtMeshInertia, 4);
+MJ_ASSERT_SIZE(mjtMeshBuiltin, 4);
+MJ_ASSERT_SIZE(mjtBuiltin, 4);
+MJ_ASSERT_SIZE(mjtMark, 4);
+MJ_ASSERT_SIZE(mjtLimited, 4);
+MJ_ASSERT_SIZE(mjtAlignFree, 4);
+MJ_ASSERT_SIZE(mjtInertiaFromGeom, 4);
+MJ_ASSERT_SIZE(mjtOrientation, 4);
+
+// mjvisualize.h
+MJ_ASSERT_SIZE(mjtCatBit, 4);
+MJ_ASSERT_SIZE(mjtMouse, 4);
+MJ_ASSERT_SIZE(mjtPertBit, 4);
+MJ_ASSERT_SIZE(mjtCamera, 4);
+MJ_ASSERT_SIZE(mjtLabel, 4);
+MJ_ASSERT_SIZE(mjtFrame, 4);
+MJ_ASSERT_SIZE(mjtVisFlag, 4);
+MJ_ASSERT_SIZE(mjtRndFlag, 4);
+MJ_ASSERT_SIZE(mjtStereo, 4);
+
+// mjrender.h
+MJ_ASSERT_SIZE(mjtGridPos, 4);
+MJ_ASSERT_SIZE(mjtFramebuffer, 4);
+MJ_ASSERT_SIZE(mjtDepthMap, 4);
+MJ_ASSERT_SIZE(mjtFontScale, 4);
+MJ_ASSERT_SIZE(mjtFont, 4);
+
+// mjui.h
+MJ_ASSERT_SIZE(mjtButton, 4);
+MJ_ASSERT_SIZE(mjtEvent, 4);
+MJ_ASSERT_SIZE(mjtItem, 4);
+MJ_ASSERT_SIZE(mjtSection, 4);
+
+// mjthread.h
+MJ_ASSERT_SIZE(mjtTaskStatus, 4);
+
+// mjplugin.h
+MJ_ASSERT_SIZE(mjtPluginCapabilityBit, 4);
+
+#undef MJ_ASSERT_SIZE
+
+#endif  // MUJOCO_MJASSERT_H_
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 213afeb3..ac76dc0f 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -35,6 +35,7 @@
 #include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
+#include <mujoco/mjassert.h>
 
 // this is a C-API
 #ifdef __cplusplus
diff --git a/test/header_test.cc b/test/header_test.cc
index 98014a5f..3de549d2 100644
--- a/test/header_test.cc
+++ b/test/header_test.cc
@@ -21,12 +21,7 @@
 #include <vector>
 
 #include <gtest/gtest.h>
-#include <mujoco/mjdata.h>
-#include <mujoco/mjmodel.h>
-#include <mujoco/mjrender.h>
-#include <mujoco/mjtype.h>
-#include <mujoco/mjui.h>
-#include <mujoco/mjvisualize.h>
+#include <mujoco/mujoco.h>
 #include <mujoco/mjxmacro.h>
 #include "test/fixture.h"
 
@@ -58,59 +53,11 @@ void CheckAddressOrdering(
 
 using HeaderTest = MujocoTest;
 
-TEST_F(HeaderTest, IntsHave4Bytes) {
-  EXPECT_EQ(4, sizeof(int));
-}
-
 TEST_F(HeaderTest, IntsHaveAtLeast31Bits) {
   int shift_left_30 = 1 << 30;
   EXPECT_GT(shift_left_30, 0);
 }
 
-TEST_F(HeaderTest, EnumsAreInts) {
-  EXPECT_EQ(sizeof(mjtDisableBit),      sizeof(int));
-  EXPECT_EQ(sizeof(mjtEnableBit),       sizeof(int));
-  EXPECT_EQ(sizeof(mjtJoint),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtGeom),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtCamLight),        sizeof(int));
-  EXPECT_EQ(sizeof(mjtTexture),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtIntegrator),      sizeof(int));
-  EXPECT_EQ(sizeof(mjtCone),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtJacobian),        sizeof(int));
-  EXPECT_EQ(sizeof(mjtSolver),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtEq),              sizeof(int));
-  EXPECT_EQ(sizeof(mjtWrap),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtTrn),             sizeof(int));
-  EXPECT_EQ(sizeof(mjtDyn),             sizeof(int));
-  EXPECT_EQ(sizeof(mjtGain),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtBias),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtObj),             sizeof(int));
-  EXPECT_EQ(sizeof(mjtConstraint),      sizeof(int));
-  EXPECT_EQ(sizeof(mjtConstraintState), sizeof(int));
-  EXPECT_EQ(sizeof(mjtSensor),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtStage),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtDataType),        sizeof(int));
-  EXPECT_EQ(sizeof(mjtLRMode),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtWarning),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtTimer),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtGridPos),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtFramebuffer),     sizeof(int));
-  EXPECT_EQ(sizeof(mjtFontScale),       sizeof(int));
-  EXPECT_EQ(sizeof(mjtFont),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtButton),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtEvent),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtItem),            sizeof(int));
-  EXPECT_EQ(sizeof(mjtCatBit),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtMouse),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtPertBit),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtCamera),          sizeof(int));
-  EXPECT_EQ(sizeof(mjtLabel),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtFrame),           sizeof(int));
-  EXPECT_EQ(sizeof(mjtVisFlag),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtRndFlag),         sizeof(int));
-  EXPECT_EQ(sizeof(mjtStereo),          sizeof(int));
-}
-
 TEST_F(HeaderTest, MjOptionFields) {
   mjOption o;
   std::vector<std::pair<const void*, const char*>> fields;
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 999c50ef..a4d3aa22 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -24,6 +24,7 @@ public static class MujocoLib {
 
 // ----------------------------------Constants----------------------------------
 
+public const bool THIRD_PARTY_MUJOCO_MJASSERT_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MJDATA_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MJEXPORT_H_ = true;
 public const bool MUJOCO_HELPER_DLL_LOCAL = true;

From 5b73846a1ed6e4a11146ed5f0f404fd8e9f32a0e Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Thu, 21 May 2026 10:03:08 -0700
Subject: [PATCH 15/61] Upgrade MuJoCo Live loading screen

PiperOrigin-RevId: 919112273
Change-Id: I0ffdf908e32d4197e007ebcc2d9f5136b59cf6b3
---
 src/experimental/studio/live.css  | 79 +++++++++++++++++++++++++++++--
 src/experimental/studio/live.html | 28 +++++++++++
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/src/experimental/studio/live.css b/src/experimental/studio/live.css
index 6c668b27..24cac67a 100644
--- a/src/experimental/studio/live.css
+++ b/src/experimental/studio/live.css
@@ -4,11 +4,40 @@
   left: 0;
   width: 100%;
   height: 100%;
-  background: rgba(0,0,0,0.6);
+  background: radial-gradient(circle at 50% 50%, #0a0f1c 0%, #020407 100%);
   z-index: 10000;
   display: none;
   align-items: center;
   justify-content: center;
+  overflow: hidden;
+  transition: opacity 0.3s ease;
+}
+#loadingOverlay::before {
+  /* 3D Ground Grid */
+  content: '';
+  position: absolute;
+  top: 0;
+  left: 0;
+  width: 100%;
+  height: 100%;
+  background-image:
+    linear-gradient(rgba(123, 147, 232, 0.035) 1px, transparent 1px),
+    linear-gradient(90deg, rgba(123, 147, 232, 0.035) 1px, transparent 1px);
+  background-size: 50px 50px;
+  background-position: center center;
+  transform: perspective(500px) rotateX(60deg) scale(2.5);
+  transform-origin: center center;
+  opacity: 0.8;
+  pointer-events: none;
+  animation: grid-drift 60s linear infinite;
+}
+@keyframes grid-drift {
+  0% {
+    background-position: center 0;
+  }
+  100% {
+    background-position: center 1000px;
+  }
 }
 .loading-container {
   color: #fff;
@@ -16,8 +45,52 @@
   font-family: sans-serif;
   text-align: center;
 }
+.loading-visual {
+  position: relative;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  margin-bottom: 20px;
+}
+.wireframe-cage {
+  position: absolute;
+  width: 320px;
+  height: 320px;
+  transform-origin: center center;
+  animation: spin-cage 25s linear infinite;
+  opacity: 0.65;
+  z-index: 0;
+  pointer-events: none;
+}
+@keyframes spin-cage {
+  0% {
+    transform: rotate(0deg) scale(0.9);
+  }
+  50% {
+    transform: rotate(180deg) scale(1.04);
+  }
+  100% {
+    transform: rotate(360deg) scale(0.9);
+  }
+}
+.logo-container-svg {
+  position: relative;
+  width: 300px;
+  height: 100px;
+  display: flex;
+  align-items: center;
+  justify-content: center;
+  z-index: 1;
+}
+.brand-logo-svg {
+  width: 100%;
+  height: auto;
+}
 .loading-title {
+  font-size: 18px;
+  margin-top: 120px; /* Push down below the bottom vertices of the rotating cube. */
   margin-bottom: 16px;
+  color: rgba(255, 255, 255, 0.75);
 }
 .loading-subtitle {
   font-size: 14px;
@@ -29,10 +102,10 @@
   opacity: 0.6;
 }
 .loading-download-label {
-  font-size: 13px;
+  font-size: 14px;
 }
 .loading-download-file {
-  font-size: 13px;
+  font-size: 14px;
 }
 #dropOverlay {
   position: fixed;
diff --git a/src/experimental/studio/live.html b/src/experimental/studio/live.html
index 81714b38..b67756f3 100644
--- a/src/experimental/studio/live.html
+++ b/src/experimental/studio/live.html
@@ -8,8 +8,36 @@
   <body style="margin: 0; overflow: hidden">
     <div id="loadingOverlay">
       <div class="loading-container">
+        <div class="loading-visual">
+          <!-- Rotating neon light-blue physics wireframe cage. -->
+          <svg class="wireframe-cage" viewBox="0 0 200 200">
+            <g class="cage-lines">
+              <polygon points="130,10 190,130 96,190 10,96" fill="none" stroke="#00f5d4" stroke-width="0.8" opacity="0.75" />
+              <line x1="10" y1="96" x2="45" y2="106" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="45" y1="106" x2="130" y2="10" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="45" y1="106" x2="96" y2="190" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="45" y1="106" x2="190" y2="130" stroke="#00f5d4" stroke-width="0.5" opacity="0.45" />
+              <line x1="150" y1="100" x2="130" y2="10" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="150" y1="100" x2="96" y2="190" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="150" y1="100" x2="190" y2="130" stroke="#00f5d4" stroke-width="0.6" opacity="0.55" />
+              <line x1="150" y1="100" x2="45" y2="106" stroke="#00f5d4" stroke-width="0.4" opacity="0.35" />
+            </g>
+          </svg>
+
+          <!-- Centered MuJoCo SVG Logo -->
+          <div class="logo-container-svg">
+            <svg viewBox="90 30 415 90" class="brand-logo-svg">
+              <g fill="rgba(255, 255, 255, 0.98)" filter="drop-shadow(0 0 6px rgba(123, 147, 232, 0.25)) drop-shadow(0 0 1px rgba(255, 255, 255, 0.25))">
+                <path d="M136.1,118.4l-19.7-47.3l-7.1,41.6H96.5l15.4-80.9l0,0L136.4,90l25.5-58.1l0,0l13.9,80.9H163l-6.3-41.6 L136.1,118.4L136.1,118.4L136.1,118.4z M211.9,113.2c-4.5,0-8.6-0.9-12.1-2.6c-3.5-1.7-6.3-4.3-8.3-7.6c-2-3.4-3-7.5-3-12.3V55.2h13.6v34.1 c0,2.8,0.4,5,1.1,6.7c0.7,1.7,1.8,3,3.3,3.9s3.2,1.3,5.3,1.3s3.9-0.4,5.3-1.3s2.5-2.2,3.3-3.9c0.7-1.7,1.1-4,1.1-6.8V55.2h13.6 v35.5c0,4.8-1,9-3,12.3c-2,3.4-4.7,5.9-8.3,7.6C220.4,112.4,216.4,113.2,211.9,113.2L211.9,113.2z M260.1,113.6c-3,0-5.9-0.5-8.7-1.4c-2.8-0.9-5.3-2.4-7.6-4.2l6.1-10.8c1.6,1.3,3.1,2.3,4.5,2.9s2.9,1,4.2,1 c1.9,0,3.5-0.4,4.6-1.2s1.9-2,2.4-3.7c0.5-1.7,0.8-4,0.8-6.9V36.1h13.3V90c0,5.3-0.7,9.7-2.2,13.2c-1.5,3.5-3.6,6.1-6.5,7.8 C268.2,112.8,264.5,113.6,260.1,113.6L260.1,113.6z M325.8,113.2c-5.4,0-10.5-1.3-15.1-3.7c-4.6-2.5-8.4-6-11.1-10.5c-2.8-4.5-4.2-9.6-4.2-15.3s1.4-10.4,4.2-15 c2.8-4.5,6.5-8.1,11.2-10.6c4.7-2.6,9.8-3.9,15.3-3.9c5.5,0,10.6,1.3,15.2,3.9c4.6,2.6,8.4,6.1,11.1,10.6 c2.8,4.5,4.2,9.6,4.2,15.1s-1.4,10.6-4.2,15.1c-2.8,4.5-6.5,8-11.2,10.5C336.5,111.9,331.3,113.2,325.8,113.2L325.8,113.2z M325.8,100.9c3.2,0,6.1-0.7,8.6-2.2c2.6-1.5,4.6-3.5,6.1-6.1c1.5-2.6,2.2-5.5,2.2-8.7c0-3.2-0.7-6.1-2.2-8.7 c-1.5-2.7-3.5-4.7-6-6.2s-5.4-2.3-8.6-2.3s-5.9,0.8-8.5,2.3c-2.6,1.5-4.6,3.6-6.1,6.2c-1.5,2.6-2.2,5.5-2.2,8.7s0.8,6.2,2.3,8.9 c1.5,2.6,3.5,4.7,6.1,6.1C320.1,100.1,322.8,100.9,325.8,100.9L325.8,100.9z M408.1,113.6c-7.3,0-13.9-1.7-19.9-5.2c-6-3.5-10.7-8.2-14.2-14.2c-3.5-6-5.2-12.5-5.2-19.7 s1.7-13.8,5.2-19.9c3.5-6.1,8.3-10.9,14.4-14.5c6.1-3.6,13-5.4,20.5-5.4c3.4,0,6.7,0.5,9.9,1.4c3.3,0.9,6.5,2.2,9.6,4l0.1,15.4 c-3.2-2.8-6.5-4.8-9.8-6.2c-3.3-1.4-6.7-2.1-10.2-2.1c-4.8,0-9.2,1.2-13.2,3.6c-4,2.4-7.2,5.7-9.6,9.9c-2.4,4.2-3.5,8.8-3.5,13.8 s1.2,9.4,3.5,13.5c2.4,4.1,5.5,7.3,9.6,9.7c4,2.4,8.5,3.6,13.4,3.6c2.6,0,5-0.3,7.1-0.8c2.1-0.5,4.1-1.4,6.2-2.7 c2-1.2,4.2-2.9,6.7-5l-0.1,15.5c-3.1,1.7-6.4,3.1-9.9,4C415.1,113.2,411.6,113.6,408.1,113.6L408.1,113.6z M473.7,113.2c-5.4,0-10.5-1.3-15.1-3.7c-4.6-2.5-8.4-6-11.1-10.5c-2.8-4.5-4.2-9.6-4.2-15.3s1.4-10.4,4.2-15 c2.8-4.5,6.5-8.1,11.2-10.6s9.8-3.9,15.3-3.9c5.5,0,10.6,1.3,15.2,3.9c4.6,2.6,8.4,6.1,11.1,10.6c2.8,4.5,4.2,9.6,4.2,15.1 s-1.4,10.6-4.2,15.1c-2.8,4.5-6.5,8-11.2,10.5C484.4,111.9,479.3,113.2,473.7,113.2L473.7,113.2z M473.7,100.9 c3.2,0,6.1-0.7,8.6-2.2c2.6-1.5,4.6-3.5,6.1-6.1c1.5-2.6,2.2-5.5,2.2-8.7c0-3.2-0.7-6.1-2.2-8.7c-1.5-2.7-3.5-4.7-6-6.2 s-5.4-2.3-8.6-2.3c-3.2,0-5.9,0.8-8.5,2.3c-2.6,1.5-4.6,3.6-6.1,6.2c-1.5,2.6-2.2,5.5-2.2,8.7s0.8,6.2,2.3,8.9s3.5,4.7,6.1,6.1 C468,100.1,470.7,100.9,473.7,100.9L473.7,100.9z"/>
+              </g>
+            </svg>
+          </div>
+        </div>
+
         <div class="loading-title">Loading model&hellip;</div>
         <div class="loading-subtitle">This may take a moment for large models.</div>
+
+        <!-- Asset Download Indicator -->
         <div id="loadingDownload" class="loading-download">
           <div class="loading-download-label">Downloading asset file</div>
           <div id="loadingDownloadFile" class="loading-download-file"></div>

From 88b8b584295428f2103e461947d6b81c489ba8f7 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Thu, 21 May 2026 23:35:18 -0700
Subject: [PATCH 16/61] Add support for transparent PBR materials.

PiperOrigin-RevId: 919477028
Change-Id: I75cc311fc73dd5fe26b72ba8d49fd421fda455e3
---
 src/experimental/filament/CMakeLists.txt      |  2 +
 .../assets/pbr_packed_transparent.mat         | 53 ++++++++++++++++++
 .../filament/assets/pbr_transparent.mat       | 56 +++++++++++++++++++
 .../filament/compat/scene_geom_util.cc        |  1 +
 .../filament/filament/material.cc             | 31 ++++++++--
 .../filament/filament/object_manager.cc       |  3 +
 .../filament/filament/object_manager.h        |  2 +
 .../filament/render_context_filament.h        |  3 +
 src/experimental/studio/live.js               |  2 +
 9 files changed, 148 insertions(+), 5 deletions(-)
 create mode 100644 src/experimental/filament/assets/pbr_packed_transparent.mat
 create mode 100644 src/experimental/filament/assets/pbr_transparent.mat

diff --git a/src/experimental/filament/CMakeLists.txt b/src/experimental/filament/CMakeLists.txt
index 7f5aac1d..0ba891b6 100644
--- a/src/experimental/filament/CMakeLists.txt
+++ b/src/experimental/filament/CMakeLists.txt
@@ -100,7 +100,9 @@ file(COPY
 set(MATC_EXECUTABLE matc)
 set(MATERIAL_FILES
     pbr.mat
+    pbr_transparent.mat
     pbr_packed.mat
+    pbr_packed_transparent.mat
     phong_2d_fade.mat
     phong_2d.mat
     phong_2d_reflect.mat
diff --git a/src/experimental/filament/assets/pbr_packed_transparent.mat b/src/experimental/filament/assets/pbr_packed_transparent.mat
new file mode 100644
index 00000000..88b0a003
--- /dev/null
+++ b/src/experimental/filament/assets/pbr_packed_transparent.mat
@@ -0,0 +1,53 @@
+// 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.
+
+material {
+    name : pbr_packed_transparent,
+    shadingModel : lit,
+    flipUV: false,
+    blending: transparent,
+    parameters : [
+        { type : sampler2d, name : BaseColor },
+        { type : sampler2d, name : Opacity },
+        { type : sampler2d, name : Normal },
+        { type : sampler2d, name : ORM },
+        { type : sampler2d, name : Emissive },
+        { type : float4, name : BaseColorFactor },
+        { type : float, name : MetallicFactor },
+        { type : float, name : RoughnessFactor }
+    ],
+    requires : [
+        uv0
+    ]
+}
+
+fragment {
+  void material(inout MaterialInputs material) {
+    vec2 uv = getUV0();
+
+    material.normal = texture(materialParams_Normal, uv).xyz * 2.0 - 1.0;
+    prepareMaterial(material);
+
+    material.baseColor        = materialParams.BaseColorFactor;
+    material.baseColor        *= texture(materialParams_BaseColor, uv);
+    material.baseColor.a      *= texture(materialParams_Opacity, uv).r;
+    material.ambientOcclusion = texture(materialParams_ORM, uv).r;
+    material.roughness        = materialParams.RoughnessFactor;
+    material.roughness        *= texture(materialParams_ORM, uv).g;
+    material.metallic         = materialParams.MetallicFactor;
+    material.metallic         *= texture(materialParams_ORM, uv).b;
+    material.emissive         = texture(materialParams_Emissive, uv);
+    material.emissive.a       = 1.0 - material.emissive.a;
+  }
+}
diff --git a/src/experimental/filament/assets/pbr_transparent.mat b/src/experimental/filament/assets/pbr_transparent.mat
new file mode 100644
index 00000000..eb69d440
--- /dev/null
+++ b/src/experimental/filament/assets/pbr_transparent.mat
@@ -0,0 +1,56 @@
+// 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.
+
+material {
+    name : pbr_transparent,
+    shadingModel : lit,
+    flipUV: false,
+    blending: transparent,
+    parameters : [
+        { type : sampler2d, name : BaseColor },
+        { type : sampler2d, name : Opacity },
+        { type : sampler2d, name : Normal },
+        { type : sampler2d, name : Occlusion },
+        { type : sampler2d, name : Roughness },
+        { type : sampler2d, name : Metallic },
+        { type : sampler2d, name : Emissive },
+        { type : float4, name : BaseColorFactor },
+        { type : float, name : MetallicFactor },
+        { type : float, name : RoughnessFactor }
+
+    ],
+    requires : [
+        uv0
+    ]
+}
+
+fragment {
+  void material(inout MaterialInputs material) {
+    vec2 uv = getUV0();
+
+    material.normal = texture(materialParams_Normal, uv).xyz * 2.0 - 1.0;
+    prepareMaterial(material);
+
+    material.baseColor        = materialParams.BaseColorFactor;
+    material.baseColor        *= texture(materialParams_BaseColor, uv);
+    material.baseColor.a      *= texture(materialParams_Opacity, uv).r;
+    material.ambientOcclusion = texture(materialParams_Occlusion, uv).r;
+    material.roughness        = materialParams.RoughnessFactor;
+    material.roughness        *= texture(materialParams_Roughness, uv).r;
+    material.metallic         = materialParams.MetallicFactor;
+    material.metallic         *= texture(materialParams_Metallic, uv).r;
+    material.emissive         = texture(materialParams_Emissive, uv);
+    material.emissive.a       = 1.0 - material.emissive.a;
+  }
+}
diff --git a/src/experimental/filament/compat/scene_geom_util.cc b/src/experimental/filament/compat/scene_geom_util.cc
index 0d542ac2..a2c11a65 100644
--- a/src/experimental/filament/compat/scene_geom_util.cc
+++ b/src/experimental/filament/compat/scene_geom_util.cc
@@ -198,6 +198,7 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
     };
     material.color_texture = get_texture(mjTEXROLE_RGB);
     material.normal_texture = get_texture(mjTEXROLE_NORMAL);
+    material.opacity_texture = get_texture(mjTEXROLE_OPACITY);
     material.emissive_texture = get_texture(mjTEXROLE_EMISSIVE);
     material.orm_texture = get_texture(mjTEXROLE_ORM);
     material.metallic_texture = get_texture(mjTEXROLE_METALLIC);
diff --git a/src/experimental/filament/filament/material.cc b/src/experimental/filament/filament/material.cc
index b96f6f17..21a024db 100644
--- a/src/experimental/filament/filament/material.cc
+++ b/src/experimental/filament/filament/material.cc
@@ -38,15 +38,35 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
       return ObjectManager::kUnlitDecor;
     }
   } else if (material.orm_texture) {
-    return ObjectManager::kPbrPacked;
+    if (material.opacity_texture) {
+      return ObjectManager::kPbrPackedTransparent;
+    } else {
+      return ObjectManager::kPbrPacked;
+    }
   } else if (material.metallic_texture) {
-    return ObjectManager::kPbr;
+    if (material.opacity_texture) {
+      return ObjectManager::kPbrPackedTransparent;
+    } else {
+      return ObjectManager::kPbr;
+    }
   } else if (material.roughness_texture) {
-    return ObjectManager::kPbr;
+    if (material.color[3] < 1.0f) {
+      return ObjectManager::kPbrTransparent;
+    } else {
+      return ObjectManager::kPbr;
+    }
   } else if (material.metallic >= 0) {
-    return ObjectManager::kPbr;
+    if (material.color[3] < 1.0f) {
+      return ObjectManager::kPbrTransparent;
+    } else {
+      return ObjectManager::kPbr;
+    }
   } else if (material.roughness >= 0) {
-    return ObjectManager::kPbr;
+    if (material.color[3] < 1.0f) {
+      return ObjectManager::kPbrTransparent;
+    } else {
+      return ObjectManager::kPbr;
+    }
   }
 
   // Check to see if we're dealing with a mesh with texture coordinates.
@@ -163,6 +183,7 @@ void UpdateMaterialInstance(filament::MaterialInstance* instance,
   };
 
   TrySetTexture("BaseColor", material.color_texture, mjTEXROLE_RGB);
+  TrySetTexture("Opacity", material.opacity_texture, mjTEXROLE_OPACITY);
   TrySetTexture("Normal", material.normal_texture, mjTEXROLE_NORMAL);
   TrySetTexture("Metallic", material.metallic_texture, mjTEXROLE_METALLIC);
   TrySetTexture("Roughness", material.roughness_texture, mjTEXROLE_ROUGHNESS);
diff --git a/src/experimental/filament/filament/object_manager.cc b/src/experimental/filament/filament/object_manager.cc
index 1816b627..47630bb6 100644
--- a/src/experimental/filament/filament/object_manager.cc
+++ b/src/experimental/filament/filament/object_manager.cc
@@ -55,6 +55,8 @@ ObjectManager::ObjectManager(filament::Engine* engine)
     : engine_(engine) {
   materials_[kPbr] = LoadMaterial(engine, "pbr.filamat");
   materials_[kPbrPacked] = LoadMaterial(engine, "pbr_packed.filamat");
+  materials_[kPbrTransparent] = LoadMaterial(engine, "pbr_transparent.filamat");
+  materials_[kPbrPackedTransparent] = LoadMaterial(engine, "pbr_packed_transparent.filamat");
   materials_[kPhong2d] = LoadMaterial(engine, "phong_2d.filamat");
   materials_[kPhong2dFade] = LoadMaterial(engine, "phong_2d_fade.filamat");
   materials_[kPhong2dReflect] = LoadMaterial(engine, "phong_2d_reflect.filamat");
@@ -97,6 +99,7 @@ ObjectManager::ObjectManager(filament::Engine* engine)
 
   fallback_textures_[mjTEXROLE_USER] = fallback_black_;
   fallback_textures_[mjTEXROLE_RGB] = fallback_white_;
+  fallback_textures_[mjTEXROLE_OPACITY] = fallback_white_;
   fallback_textures_[mjTEXROLE_OCCLUSION] = fallback_white_;
   fallback_textures_[mjTEXROLE_ROUGHNESS] = fallback_white_;
   fallback_textures_[mjTEXROLE_METALLIC] = fallback_black_;
diff --git a/src/experimental/filament/filament/object_manager.h b/src/experimental/filament/filament/object_manager.h
index 509b3505..05ce49d5 100644
--- a/src/experimental/filament/filament/object_manager.h
+++ b/src/experimental/filament/filament/object_manager.h
@@ -42,7 +42,9 @@ class ObjectManager {
   // ObjectManager.
   enum MaterialType {
     kPbr,
+    kPbrTransparent,
     kPbrPacked,
+    kPbrPackedTransparent,
     kPhong2d,
     kPhong2dFade,
     kPhong2dReflect,
diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h
index 023cc62c..fcc587c9 100644
--- a/src/experimental/filament/render_context_filament.h
+++ b/src/experimental/filament/render_context_filament.h
@@ -627,6 +627,9 @@ struct mjrMaterial {
   // The texture containing the base color of the object.
   const mjrTexture* color_texture;
 
+  // The texture containing the opacity of the object.
+  const mjrTexture* opacity_texture;
+
   // The normal map of the object.
   const mjrTexture* normal_texture;
 
diff --git a/src/experimental/studio/live.js b/src/experimental/studio/live.js
index c479f40f..8de5230f 100644
--- a/src/experimental/studio/live.js
+++ b/src/experimental/studio/live.js
@@ -57,7 +57,9 @@ var Module = {
       "assets/ibl.ktx",
       "assets/OpenSans-Regular.ttf",
       "assets/pbr.filamat",
+      "assets/pbr_transparent.filamat",
       "assets/pbr_packed.filamat",
+      "assets/pbr_packed_transparent.filamat",
       "assets/phong_2d_fade.filamat",
       "assets/phong_2d.filamat",
       "assets/phong_2d_reflect.filamat",

From 3499adacae8a7089156f8f7dfb86c5e792cd9f83 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Fri, 22 May 2026 00:07:40 -0700
Subject: [PATCH 17/61] Store skin and flex meshes separately.

PiperOrigin-RevId: 919490079
Change-Id: I45473ee59e22465732fb4b470a94c2944fed9c37
---
 .../filament/compat/model_objects.cc          | 22 +++++++++++++++----
 .../filament/compat/model_objects.h           |  6 +++--
 .../filament/compat/scene_geom_util.cc        |  4 ++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/experimental/filament/compat/model_objects.cc b/src/experimental/filament/compat/model_objects.cc
index 989cc781..20d51231 100644
--- a/src/experimental/filament/compat/model_objects.cc
+++ b/src/experimental/filament/compat/model_objects.cc
@@ -599,7 +599,13 @@ void ModelObjects::CreateSkinFlexMesh(const mjvScene* scene, const mjvGeom& geom
   mjrMeshData data;
   mjr_defaultMeshData(&data);
   UpdateSkinFlexMeshData(&data, model_, scene, geom);
-  dynamic_meshes_.insert_or_assign(geom.objid, CreateMesh(ctx_, data));
+  if (geom.type == mjGEOM_FLEX) {
+    flexes_.insert_or_assign(geom.objid, CreateMesh(ctx_, data));
+  } else if (geom.type == mjGEOM_SKIN) {
+    skins_.insert_or_assign(geom.objid, CreateMesh(ctx_, data));
+  } else {
+    mju_error("Unsupported dynamic mesh type: %d", geom.type);
+  }
 }
 
 const mjrMesh* ModelObjects::GetMesh(int data_id) const {
@@ -624,11 +630,19 @@ const mjrMesh* ModelObjects::GetHeightField(int hfield_id) const {
   return nullptr;
 }
 
-const mjrMesh* ModelObjects::GetFlexSkinMesh(int geom_id) const {
-  if (auto it = dynamic_meshes_.find(geom_id); it != dynamic_meshes_.end()) {
+const mjrMesh* ModelObjects::GetFlexMesh(int geom_id) const {
+  if (auto it = flexes_.find(geom_id); it != flexes_.end()) {
     return it->second.get();
   }
-  mju_error("Unknown dynamic mesh %d", geom_id);
+  mju_error("Unknown flex mesh %d", geom_id);
+  return nullptr;
+}
+
+const mjrMesh* ModelObjects::GetSkinMesh(int geom_id) const {
+  if (auto it = skins_.find(geom_id); it != skins_.end()) {
+    return it->second.get();
+  }
+  mju_error("Unknown skin mesh %d", geom_id);
   return nullptr;
 }
 
diff --git a/src/experimental/filament/compat/model_objects.h b/src/experimental/filament/compat/model_objects.h
index 0cee9ad8..b1aa7a49 100644
--- a/src/experimental/filament/compat/model_objects.h
+++ b/src/experimental/filament/compat/model_objects.h
@@ -40,7 +40,8 @@ class ModelObjects {
   // Returns the cached instance of a filament object created from the mjModel.
   const mjrMesh* GetMesh(int data_id) const;
   const mjrMesh* GetHeightField(int hfield_id) const;
-  const mjrMesh* GetFlexSkinMesh(int geom_id) const;
+  const mjrMesh* GetSkinMesh(int geom_id) const;
+  const mjrMesh* GetFlexMesh(int geom_id) const;
   const mjrTexture* GetTexture(int tex_id) const;
   const mjrTexture* GetSkyboxTexture() const;
 
@@ -59,7 +60,8 @@ class ModelObjects {
   std::unordered_map<int, UniquePtr<mjrMesh>> meshes_;
   std::unordered_map<int, UniquePtr<mjrMesh>> convex_hulls_;
   std::unordered_map<int, UniquePtr<mjrMesh>> height_fields_;
-  std::unordered_map<int, UniquePtr<mjrMesh>> dynamic_meshes_;
+  std::unordered_map<int, UniquePtr<mjrMesh>> skins_;
+  std::unordered_map<int, UniquePtr<mjrMesh>> flexes_;
   std::unordered_map<int, UniquePtr<mjrTexture>> textures_;
   float specular_multiplier_ = 0.2f;
   float shininess_multiplier_ = 0.1f;
diff --git a/src/experimental/filament/compat/scene_geom_util.cc b/src/experimental/filament/compat/scene_geom_util.cc
index a2c11a65..83622ff1 100644
--- a/src/experimental/filament/compat/scene_geom_util.cc
+++ b/src/experimental/filament/compat/scene_geom_util.cc
@@ -130,7 +130,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       mjrf_setRenderableSize(renderable, geom.size);
       break;
     case mjGEOM_FLEX:
-      mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
+      mjrf_setRenderableMesh(renderable, model_objs->GetFlexMesh(geom.objid), 0, 0);
       // Flexes are defined in global space.
       std::memset(position, 0, sizeof(position));
       std::memset(rotation, 0, sizeof(rotation));
@@ -139,7 +139,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       rotation[8] = 1.f;
       break;
     case mjGEOM_SKIN:
-      mjrf_setRenderableMesh(renderable, model_objs->GetFlexSkinMesh(geom.objid), 0, 0);
+      mjrf_setRenderableMesh(renderable, model_objs->GetSkinMesh(geom.objid), 0, 0);
       // Skins are defined in global space.
       std::memset(position, 0, sizeof(position));
       std::memset(rotation, 0, sizeof(rotation));

From 8ae51e104dee2eaea13fe1395296393f6e85456d Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Fri, 22 May 2026 00:52:07 -0700
Subject: [PATCH 18/61] Move DrawTextAt from imgui_bridge to imgui_widgets.

PiperOrigin-RevId: 919508805
Change-Id: Ie311b88164d8c2ff7a8452ec2f1b8be7e7d9770e
---
 src/experimental/platform/hal/renderer.cc     |  1 +
 src/experimental/platform/ux/imgui_bridge.cc  | 32 -------------------
 src/experimental/platform/ux/imgui_bridge.h   |  4 ---
 src/experimental/platform/ux/imgui_widgets.cc | 32 +++++++++++++++++++
 src/experimental/platform/ux/imgui_widgets.h  |  4 +++
 5 files changed, 37 insertions(+), 36 deletions(-)

diff --git a/src/experimental/platform/hal/renderer.cc b/src/experimental/platform/hal/renderer.cc
index 591965dd..0acfc70c 100644
--- a/src/experimental/platform/hal/renderer.cc
+++ b/src/experimental/platform/hal/renderer.cc
@@ -32,6 +32,7 @@
 #include "experimental/filament/render_context_filament_cpp.h"
 #include "experimental/platform/hal/graphics_mode.h"
 #include "experimental/platform/ux/imgui_bridge.h"
+#include "experimental/platform/ux/imgui_widgets.h"
 #include "experimental/platform/ux/plugin.h"
 
 namespace mujoco::platform {
diff --git a/src/experimental/platform/ux/imgui_bridge.cc b/src/experimental/platform/ux/imgui_bridge.cc
index 91147fa4..ef042507 100644
--- a/src/experimental/platform/ux/imgui_bridge.cc
+++ b/src/experimental/platform/ux/imgui_bridge.cc
@@ -318,36 +318,4 @@ mjrCamera ImguiBridge::GetCamera(int width, int height) const {
   return camera;
 }
 
-static ImVec2 ClipSpaceToWindowCoordinates(float x, float y) {
-  const ImVec2& display_size = ImGui::GetIO().DisplaySize;
-  const float pos_x = display_size.x * ((x + 1) * 0.5f);
-  const float pos_y = display_size.y * (1.0f - ((y + 1) * 0.5f));
-  return ImVec2(pos_x, pos_y);
-}
-
-void DrawTextAt(const char* text, float x, float y, float z) {
-  if (x < -1 || y < -1 || x > 1 || y > 1 || z < -1 || z > 1) {
-    return;
-  }
-
-  const ImVec2 center_pos = ClipSpaceToWindowCoordinates(x, y);
-  const ImVec2 size = ImGui::CalcTextSize(text);
-
-  const ImVec2 pos = ImVec2(center_pos.x - size.x / 2, center_pos.y);
-  const ImVec2 shadow_pos = ImVec2(pos.x + 2, pos.y + 2);
-  const int flags = ImGuiWindowFlags_NoBringToFrontOnFocus |
-                    ImGuiWindowFlags_NoFocusOnAppearing |
-                    ImGuiWindowFlags_NoBackground |
-                    ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoInputs |
-                    ImGuiWindowFlags_NoNav;
-
-  ImGui::Begin("labels", nullptr, flags);
-  ImGui::BeginChild("labels", ImGui::GetIO().DisplaySize, 0, flags);
-  ImDrawList* draw_list = ImGui::GetWindowDrawList();
-  draw_list->AddText(shadow_pos, IM_COL32_BLACK, text);
-  draw_list->AddText(pos, IM_COL32_WHITE, text);
-  ImGui::EndChild();
-  ImGui::End();
-}
-
 }  // namespace mujoco
diff --git a/src/experimental/platform/ux/imgui_bridge.h b/src/experimental/platform/ux/imgui_bridge.h
index 9baefd04..f2fa8916 100644
--- a/src/experimental/platform/ux/imgui_bridge.h
+++ b/src/experimental/platform/ux/imgui_bridge.h
@@ -65,10 +65,6 @@ class ImguiBridge {
   uintptr_t next_tex_id_ = 1;
 };
 
-// Draws text at the given screen coordinates in clip space (i.e. [-1,-1,-1] to
-// [1,1,1]).
-void DrawTextAt(const char* text, float x, float y, float z);
-
 }  // namespace mujoco
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_IMGUI_BRIDGE_H_
diff --git a/src/experimental/platform/ux/imgui_widgets.cc b/src/experimental/platform/ux/imgui_widgets.cc
index cf33bc8d..2ee56a3d 100644
--- a/src/experimental/platform/ux/imgui_widgets.cc
+++ b/src/experimental/platform/ux/imgui_widgets.cc
@@ -446,4 +446,36 @@ ImPlotPairLayout ImPlot_ComputePairLayout() {
   };
 }
 
+static ImVec2 ClipSpaceToWindowCoordinates(float x, float y) {
+  const ImVec2& display_size = ImGui::GetIO().DisplaySize;
+  const float pos_x = display_size.x * ((x + 1) * 0.5f);
+  const float pos_y = display_size.y * (1.0f - ((y + 1) * 0.5f));
+  return ImVec2(pos_x, pos_y);
+}
+
+void DrawTextAt(const char* text, float x, float y, float z) {
+  if (x < -1 || y < -1 || x > 1 || y > 1 || z < -1 || z > 1) {
+    return;
+  }
+
+  const ImVec2 center_pos = ClipSpaceToWindowCoordinates(x, y);
+  const ImVec2 size = ImGui::CalcTextSize(text);
+
+  const ImVec2 pos = ImVec2(center_pos.x - size.x / 2, center_pos.y);
+  const ImVec2 shadow_pos = ImVec2(pos.x + 2, pos.y + 2);
+  const int flags = ImGuiWindowFlags_NoBringToFrontOnFocus |
+                    ImGuiWindowFlags_NoFocusOnAppearing |
+                    ImGuiWindowFlags_NoBackground |
+                    ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoInputs |
+                    ImGuiWindowFlags_NoNav;
+
+  ImGui::Begin("labels", nullptr, flags);
+  ImGui::BeginChild("labels", ImGui::GetIO().DisplaySize, 0, flags);
+  ImDrawList* draw_list = ImGui::GetWindowDrawList();
+  draw_list->AddText(shadow_pos, IM_COL32_BLACK, text);
+  draw_list->AddText(pos, IM_COL32_WHITE, text);
+  ImGui::EndChild();
+  ImGui::End();
+}
+
 }  // namespace mujoco::platform
diff --git a/src/experimental/platform/ux/imgui_widgets.h b/src/experimental/platform/ux/imgui_widgets.h
index 3ee85efa..211b213d 100644
--- a/src/experimental/platform/ux/imgui_widgets.h
+++ b/src/experimental/platform/ux/imgui_widgets.h
@@ -656,6 +656,10 @@ struct ImPlotPairLayout {
 // side-by-side; otherwise they are stacked vertically.
 ImPlotPairLayout ImPlot_ComputePairLayout();
 
+// Draws text at the given screen coordinates in clip space (i.e. [-1,-1,-1] to
+// [1,1,1]).
+void DrawTextAt(const char* text, float x, float y, float z);
+
 }  // namespace mujoco::platform
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_PLATFORM_UX_IMGUI_WIDGETS_H_

From 8469198ba5c0dac1ffbbab48eed03097b830d265 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Fri, 22 May 2026 09:28:54 -0700
Subject: [PATCH 19/61] Remove multiccd enable flag from mjSpec Colab.

PiperOrigin-RevId: 919716011
Change-Id: I746271990db1e213fee33b38e639db78bbf97a1a
---
 python/mjspec.ipynb | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/mjspec.ipynb b/python/mjspec.ipynb
index 2b98508c..49a27ec1 100644
--- a/python/mjspec.ipynb
+++ b/python/mjspec.ipynb
@@ -1653,7 +1653,6 @@
     "spec = mj.MjSpec.from_string(arena_xml)\n",
     "\n",
     "spec.option.enableflags |= mj.mjtEnableBit.mjENBL_OVERRIDE\n",
-    "spec.option.enableflags |= mj.mjtEnableBit.mjENBL_MULTICCD\n",
     "spec.option.timestep = 0.0001\n",
     "spec.compiler.degree = False\n",
     "\n",

From ab0e0091d57c17303c531ff382e82c024df8dddd Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Fri, 22 May 2026 09:34:16 -0700
Subject: [PATCH 20/61] Standardize "DO NOT EDIT" comments in generated files.

PiperOrigin-RevId: 919718422
Change-Id: I2cba07d75d37d45f3f30b23d049c944b080248ad
---
 doc/XMLschema.rst                                          | 3 +++
 include/mujoco/experimental/usd/mjcPhysics/actuator.h      | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/api.h           | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h  | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h   | 2 ++
 .../experimental/usd/mjcPhysics/equalityConnectAPI.h       | 2 ++
 .../mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h  | 2 ++
 .../mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h   | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h  | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/jointAPI.h      | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/keyframe.h      | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/materialAPI.h   | 2 ++
 .../mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h  | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h      | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/siteAPI.h       | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/tendon.h        | 2 ++
 include/mujoco/experimental/usd/mjcPhysics/tokens.h        | 2 ++
 python/mujoco/codegen/generate_cs_bindings.py              | 4 +++-
 src/experimental/usd/mjcPhysics/actuator.cpp               | 2 ++
 src/experimental/usd/mjcPhysics/collisionAPI.cpp           | 2 ++
 src/experimental/usd/mjcPhysics/equalityAPI.cpp            | 2 ++
 src/experimental/usd/mjcPhysics/equalityConnectAPI.cpp     | 2 ++
 src/experimental/usd/mjcPhysics/equalityJointAPI.cpp       | 2 ++
 src/experimental/usd/mjcPhysics/equalityWeldAPI.cpp        | 2 ++
 src/experimental/usd/mjcPhysics/imageableAPI.cpp           | 2 ++
 src/experimental/usd/mjcPhysics/jointAPI.cpp               | 2 ++
 src/experimental/usd/mjcPhysics/keyframe.cpp               | 2 ++
 src/experimental/usd/mjcPhysics/materialAPI.cpp            | 2 ++
 src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp       | 2 ++
 src/experimental/usd/mjcPhysics/sceneAPI.cpp               | 2 ++
 src/experimental/usd/mjcPhysics/siteAPI.cpp                | 2 ++
 src/experimental/usd/mjcPhysics/tendon.cpp                 | 2 ++
 src/experimental/usd/mjcPhysics/tokens.cpp                 | 2 ++
 unity/Runtime/Bindings/MjBindings.cs                       | 4 +++-
 wasm/codegen/generated/bindings.cc                         | 2 ++
 wasm/codegen/generators/binding_builder.py                 | 7 ++++++-
 36 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst
index 26eb579d..6507fbdd 100755
--- a/doc/XMLschema.rst
+++ b/doc/XMLschema.rst
@@ -1,3 +1,6 @@
+..
+  DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 .. raw:: html
 
    <div class="schema-controls" style="margin-bottom: 1em;">
diff --git a/include/mujoco/experimental/usd/mjcPhysics/actuator.h b/include/mujoco/experimental/usd/mjcPhysics/actuator.h
index cb0069b9..fc3e3fdc 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/actuator.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/actuator.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_ACTUATOR_H
 #define MJCPHYSICS_GENERATED_ACTUATOR_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/api.h b/include/mujoco/experimental/usd/mjcPhysics/api.h
index 529c31f5..79cbbded 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/api.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/api.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_API_H
 #define MJCPHYSICS_API_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h b/include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h
index 3aee44ed..ec41ead7 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_COLLISIONAPI_H
 #define MJCPHYSICS_GENERATED_COLLISIONAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h b/include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h
index 5f8aa32b..581bd91f 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_EQUALITYAPI_H
 #define MJCPHYSICS_GENERATED_EQUALITYAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h b/include/mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h
index 104d11e1..2143f1f2 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_EQUALITYCONNECTAPI_H
 #define MJCPHYSICS_GENERATED_EQUALITYCONNECTAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h b/include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h
index adfca61b..ca7aee84 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
 #define MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h b/include/mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h
index 98a7ae7a..3839dbb6 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_EQUALITYWELDAPI_H
 #define MJCPHYSICS_GENERATED_EQUALITYWELDAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h b/include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h
index 8a77eada..ec578d9c 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_IMAGEABLEAPI_H
 #define MJCPHYSICS_GENERATED_IMAGEABLEAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/jointAPI.h b/include/mujoco/experimental/usd/mjcPhysics/jointAPI.h
index 6d53472d..c600d9cd 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/jointAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/jointAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_JOINTAPI_H
 #define MJCPHYSICS_GENERATED_JOINTAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/keyframe.h b/include/mujoco/experimental/usd/mjcPhysics/keyframe.h
index 4fa0091c..572f26a5 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/keyframe.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/keyframe.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_KEYFRAME_H
 #define MJCPHYSICS_GENERATED_KEYFRAME_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/materialAPI.h b/include/mujoco/experimental/usd/mjcPhysics/materialAPI.h
index 8d30abfa..9b671935 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/materialAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/materialAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_MATERIALAPI_H
 #define MJCPHYSICS_GENERATED_MATERIALAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h b/include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h
index 431e9c38..91923540 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H
 #define MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h b/include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h
index 7d9b4b37..2e0358df 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_SCENEAPI_H
 #define MJCPHYSICS_GENERATED_SCENEAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/siteAPI.h b/include/mujoco/experimental/usd/mjcPhysics/siteAPI.h
index 8e92e9c2..0455be70 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/siteAPI.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/siteAPI.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_SITEAPI_H
 #define MJCPHYSICS_GENERATED_SITEAPI_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/tendon.h b/include/mujoco/experimental/usd/mjcPhysics/tendon.h
index 9e72ed50..8e4fab72 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/tendon.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/tendon.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_GENERATED_TENDON_H
 #define MJCPHYSICS_GENERATED_TENDON_H
 
diff --git a/include/mujoco/experimental/usd/mjcPhysics/tokens.h b/include/mujoco/experimental/usd/mjcPhysics/tokens.h
index 53796b06..e00e2daf 100644
--- a/include/mujoco/experimental/usd/mjcPhysics/tokens.h
+++ b/include/mujoco/experimental/usd/mjcPhysics/tokens.h
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #ifndef MJCPHYSICS_TOKENS_H
 #define MJCPHYSICS_TOKENS_H
 
diff --git a/python/mujoco/codegen/generate_cs_bindings.py b/python/mujoco/codegen/generate_cs_bindings.py
index c9f783c6..ed014515 100644
--- a/python/mujoco/codegen/generate_cs_bindings.py
+++ b/python/mujoco/codegen/generate_cs_bindings.py
@@ -438,7 +438,9 @@ _COPYRIGHT_HEADER = """// Copyright 2022 DeepMind Technologies Limited
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// THIS FILE IS AUTO-GENERATED
+
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 """
 
 
diff --git a/src/experimental/usd/mjcPhysics/actuator.cpp b/src/experimental/usd/mjcPhysics/actuator.cpp
index 0db68f9a..7d6b9ce8 100644
--- a/src/experimental/usd/mjcPhysics/actuator.cpp
+++ b/src/experimental/usd/mjcPhysics/actuator.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/actuator.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/collisionAPI.cpp b/src/experimental/usd/mjcPhysics/collisionAPI.cpp
index a3e1ab32..6e10e761 100644
--- a/src/experimental/usd/mjcPhysics/collisionAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/collisionAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/collisionAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/equalityAPI.cpp b/src/experimental/usd/mjcPhysics/equalityAPI.cpp
index 8caad788..fc716561 100644
--- a/src/experimental/usd/mjcPhysics/equalityAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/equalityAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/equalityAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/equalityConnectAPI.cpp b/src/experimental/usd/mjcPhysics/equalityConnectAPI.cpp
index 17442fc3..adeb228b 100644
--- a/src/experimental/usd/mjcPhysics/equalityConnectAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/equalityConnectAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/equalityJointAPI.cpp b/src/experimental/usd/mjcPhysics/equalityJointAPI.cpp
index 92607cfd..69a89b17 100644
--- a/src/experimental/usd/mjcPhysics/equalityJointAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/equalityJointAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/equalityWeldAPI.cpp b/src/experimental/usd/mjcPhysics/equalityWeldAPI.cpp
index d8cd5afa..59c354a4 100644
--- a/src/experimental/usd/mjcPhysics/equalityWeldAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/equalityWeldAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/imageableAPI.cpp b/src/experimental/usd/mjcPhysics/imageableAPI.cpp
index c36a3c5a..6f38613c 100644
--- a/src/experimental/usd/mjcPhysics/imageableAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/imageableAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/imageableAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/jointAPI.cpp b/src/experimental/usd/mjcPhysics/jointAPI.cpp
index f918008f..381f1b9f 100644
--- a/src/experimental/usd/mjcPhysics/jointAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/jointAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/jointAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/keyframe.cpp b/src/experimental/usd/mjcPhysics/keyframe.cpp
index 203d3c0d..9541c89a 100644
--- a/src/experimental/usd/mjcPhysics/keyframe.cpp
+++ b/src/experimental/usd/mjcPhysics/keyframe.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/keyframe.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/materialAPI.cpp b/src/experimental/usd/mjcPhysics/materialAPI.cpp
index 705af0d8..d0d6adf2 100644
--- a/src/experimental/usd/mjcPhysics/materialAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/materialAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/materialAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp b/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp
index 05b53877..1e002b17 100644
--- a/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/meshCollisionAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/sceneAPI.cpp b/src/experimental/usd/mjcPhysics/sceneAPI.cpp
index 6396c858..95aad5a4 100644
--- a/src/experimental/usd/mjcPhysics/sceneAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/sceneAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/sceneAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/siteAPI.cpp b/src/experimental/usd/mjcPhysics/siteAPI.cpp
index 0ec70277..d7b359a8 100644
--- a/src/experimental/usd/mjcPhysics/siteAPI.cpp
+++ b/src/experimental/usd/mjcPhysics/siteAPI.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/siteAPI.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/tendon.cpp b/src/experimental/usd/mjcPhysics/tendon.cpp
index 6f0b2e21..5446fe49 100644
--- a/src/experimental/usd/mjcPhysics/tendon.cpp
+++ b/src/experimental/usd/mjcPhysics/tendon.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/tendon.h>
 
 #include <pxr/usd/sdf/assetPath.h>
diff --git a/src/experimental/usd/mjcPhysics/tokens.cpp b/src/experimental/usd/mjcPhysics/tokens.cpp
index c278e715..cb69db80 100644
--- a/src/experimental/usd/mjcPhysics/tokens.cpp
+++ b/src/experimental/usd/mjcPhysics/tokens.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 #include <mujoco/experimental/usd/mjcPhysics/tokens.h>
 
 PXR_NAMESPACE_OPEN_SCOPE
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index a4d3aa22..667d7d33 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -12,7 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// THIS FILE IS AUTO-GENERATED
+
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 using System;
 using System.Runtime.InteropServices;
 using System.Text;
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index 8262902d..adfc48bb 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
+
 // NOLINTBEGIN(whitespace/line_length)
 // NOLINTBEGIN(whitespace/semicolon)
 
diff --git a/wasm/codegen/generators/binding_builder.py b/wasm/codegen/generators/binding_builder.py
index b0f229e4..fab3035b 100644
--- a/wasm/codegen/generators/binding_builder.py
+++ b/wasm/codegen/generators/binding_builder.py
@@ -62,7 +62,12 @@ class BindingBuilder:
       self.content_cc = common.replace_lines_containing_marker(
           self.content_cc, marker, content
       )
-    return "".join(self.content_cc)
+    text = "".join(self.content_cc)
+    return text.replace(
+        "// limitations under the License.\n\n",
+        "// limitations under the License.\n\n// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.\n\n",
+        1,
+    )
 
   def build(self, generated_path_cc: str):
     common.write_to_file(generated_path_cc, self.to_string())

From 5c21c8d28bf6db5acb9eb98b42e79d19ff3e029e Mon Sep 17 00:00:00 2001
From: Saran Tunyasuvunakool <stunya@google.com>
Date: Fri, 22 May 2026 10:25:23 -0700
Subject: [PATCH 21/61] Minor restyling of the loading screen.

PiperOrigin-RevId: 919744412
Change-Id: I3e4aa0b59d100ac85e8f915d7d733b2b3ea1ecea
---
 src/experimental/studio/live.css  | 29 +++++++++++++++--------------
 src/experimental/studio/live.html |  8 +++++---
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/src/experimental/studio/live.css b/src/experimental/studio/live.css
index 24cac67a..d72872db 100644
--- a/src/experimental/studio/live.css
+++ b/src/experimental/studio/live.css
@@ -4,7 +4,7 @@
   left: 0;
   width: 100%;
   height: 100%;
-  background: radial-gradient(circle at 50% 50%, #0a0f1c 0%, #020407 100%);
+  background: radial-gradient(circle at 50% 50%, #092057 0%, #090f23 100%);
   z-index: 10000;
   display: none;
   align-items: center;
@@ -21,8 +21,8 @@
   width: 100%;
   height: 100%;
   background-image:
-    linear-gradient(rgba(123, 147, 232, 0.035) 1px, transparent 1px),
-    linear-gradient(90deg, rgba(123, 147, 232, 0.035) 1px, transparent 1px);
+    linear-gradient(rgba(123, 147, 232, 0.2) 0.5px, transparent 0.5px),
+    linear-gradient(90deg, rgba(123, 147, 232, 0.2) 0.5px, transparent 0.5px);
   background-size: 50px 50px;
   background-position: center center;
   transform: perspective(500px) rotateX(60deg) scale(2.5);
@@ -42,7 +42,7 @@
 .loading-container {
   color: #fff;
   font-size: 24px;
-  font-family: sans-serif;
+  font-family: "Atkinson Hyperlegible Next", sans-serif;
   text-align: center;
 }
 .loading-visual {
@@ -87,25 +87,26 @@
   height: auto;
 }
 .loading-title {
-  font-size: 18px;
+  font-size: 26px;
+  font-family: "Atkinson Hyperlegible Next", sans-serif;
+  font-weight: 700;
   margin-top: 120px; /* Push down below the bottom vertices of the rotating cube. */
-  margin-bottom: 16px;
   color: rgba(255, 255, 255, 0.75);
 }
-.loading-subtitle {
-  font-size: 14px;
-  opacity: 0.7;
-}
 .loading-download {
   display: none;
-  margin-top: 24px;
+  margin-top: 16px;
   opacity: 0.6;
 }
 .loading-download-label {
-  font-size: 14px;
+  font-size: 16px;
+  font-family: "Atkinson Hyperlegible Next", sans-serif;
+  font-weight: 400;
 }
 .loading-download-file {
-  font-size: 14px;
+  font-size: 16px;
+  font-family: "Atkinson Hyperlegible Next", sans-serif;
+  font-weight: 400;
 }
 #dropOverlay {
   position: fixed;
@@ -123,7 +124,7 @@
 .drop-message {
   color: #fff;
   font-size: 24px;
-  font-family: sans-serif;
+  font-family: "Atkinson Hyperlegible Next", sans-serif;
   padding: 32px 48px;
   border: 3px dashed #fff;
   border-radius: 16px;
diff --git a/src/experimental/studio/live.html b/src/experimental/studio/live.html
index b67756f3..08d33f37 100644
--- a/src/experimental/studio/live.html
+++ b/src/experimental/studio/live.html
@@ -2,6 +2,9 @@
 <html lang="en">
   <head>
     <meta charset="utf-8" />
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Atkinson+Hyperlegible+Next:wght@200..800&display=swap" rel="stylesheet">
     <link rel="stylesheet" href="live.css">
     <title>MuJoCo Live</title>
   </head>
@@ -26,7 +29,7 @@
 
           <!-- Centered MuJoCo SVG Logo -->
           <div class="logo-container-svg">
-            <svg viewBox="90 30 415 90" class="brand-logo-svg">
+            <svg viewBox="90 30 415 90" class="brand-logo-svg" title="MuJoCo">
               <g fill="rgba(255, 255, 255, 0.98)" filter="drop-shadow(0 0 6px rgba(123, 147, 232, 0.25)) drop-shadow(0 0 1px rgba(255, 255, 255, 0.25))">
                 <path d="M136.1,118.4l-19.7-47.3l-7.1,41.6H96.5l15.4-80.9l0,0L136.4,90l25.5-58.1l0,0l13.9,80.9H163l-6.3-41.6 L136.1,118.4L136.1,118.4L136.1,118.4z M211.9,113.2c-4.5,0-8.6-0.9-12.1-2.6c-3.5-1.7-6.3-4.3-8.3-7.6c-2-3.4-3-7.5-3-12.3V55.2h13.6v34.1 c0,2.8,0.4,5,1.1,6.7c0.7,1.7,1.8,3,3.3,3.9s3.2,1.3,5.3,1.3s3.9-0.4,5.3-1.3s2.5-2.2,3.3-3.9c0.7-1.7,1.1-4,1.1-6.8V55.2h13.6 v35.5c0,4.8-1,9-3,12.3c-2,3.4-4.7,5.9-8.3,7.6C220.4,112.4,216.4,113.2,211.9,113.2L211.9,113.2z M260.1,113.6c-3,0-5.9-0.5-8.7-1.4c-2.8-0.9-5.3-2.4-7.6-4.2l6.1-10.8c1.6,1.3,3.1,2.3,4.5,2.9s2.9,1,4.2,1 c1.9,0,3.5-0.4,4.6-1.2s1.9-2,2.4-3.7c0.5-1.7,0.8-4,0.8-6.9V36.1h13.3V90c0,5.3-0.7,9.7-2.2,13.2c-1.5,3.5-3.6,6.1-6.5,7.8 C268.2,112.8,264.5,113.6,260.1,113.6L260.1,113.6z M325.8,113.2c-5.4,0-10.5-1.3-15.1-3.7c-4.6-2.5-8.4-6-11.1-10.5c-2.8-4.5-4.2-9.6-4.2-15.3s1.4-10.4,4.2-15 c2.8-4.5,6.5-8.1,11.2-10.6c4.7-2.6,9.8-3.9,15.3-3.9c5.5,0,10.6,1.3,15.2,3.9c4.6,2.6,8.4,6.1,11.1,10.6 c2.8,4.5,4.2,9.6,4.2,15.1s-1.4,10.6-4.2,15.1c-2.8,4.5-6.5,8-11.2,10.5C336.5,111.9,331.3,113.2,325.8,113.2L325.8,113.2z M325.8,100.9c3.2,0,6.1-0.7,8.6-2.2c2.6-1.5,4.6-3.5,6.1-6.1c1.5-2.6,2.2-5.5,2.2-8.7c0-3.2-0.7-6.1-2.2-8.7 c-1.5-2.7-3.5-4.7-6-6.2s-5.4-2.3-8.6-2.3s-5.9,0.8-8.5,2.3c-2.6,1.5-4.6,3.6-6.1,6.2c-1.5,2.6-2.2,5.5-2.2,8.7s0.8,6.2,2.3,8.9 c1.5,2.6,3.5,4.7,6.1,6.1C320.1,100.1,322.8,100.9,325.8,100.9L325.8,100.9z M408.1,113.6c-7.3,0-13.9-1.7-19.9-5.2c-6-3.5-10.7-8.2-14.2-14.2c-3.5-6-5.2-12.5-5.2-19.7 s1.7-13.8,5.2-19.9c3.5-6.1,8.3-10.9,14.4-14.5c6.1-3.6,13-5.4,20.5-5.4c3.4,0,6.7,0.5,9.9,1.4c3.3,0.9,6.5,2.2,9.6,4l0.1,15.4 c-3.2-2.8-6.5-4.8-9.8-6.2c-3.3-1.4-6.7-2.1-10.2-2.1c-4.8,0-9.2,1.2-13.2,3.6c-4,2.4-7.2,5.7-9.6,9.9c-2.4,4.2-3.5,8.8-3.5,13.8 s1.2,9.4,3.5,13.5c2.4,4.1,5.5,7.3,9.6,9.7c4,2.4,8.5,3.6,13.4,3.6c2.6,0,5-0.3,7.1-0.8c2.1-0.5,4.1-1.4,6.2-2.7 c2-1.2,4.2-2.9,6.7-5l-0.1,15.5c-3.1,1.7-6.4,3.1-9.9,4C415.1,113.2,411.6,113.6,408.1,113.6L408.1,113.6z M473.7,113.2c-5.4,0-10.5-1.3-15.1-3.7c-4.6-2.5-8.4-6-11.1-10.5c-2.8-4.5-4.2-9.6-4.2-15.3s1.4-10.4,4.2-15 c2.8-4.5,6.5-8.1,11.2-10.6s9.8-3.9,15.3-3.9c5.5,0,10.6,1.3,15.2,3.9c4.6,2.6,8.4,6.1,11.1,10.6c2.8,4.5,4.2,9.6,4.2,15.1 s-1.4,10.6-4.2,15.1c-2.8,4.5-6.5,8-11.2,10.5C484.4,111.9,479.3,113.2,473.7,113.2L473.7,113.2z M473.7,100.9 c3.2,0,6.1-0.7,8.6-2.2c2.6-1.5,4.6-3.5,6.1-6.1c1.5-2.6,2.2-5.5,2.2-8.7c0-3.2-0.7-6.1-2.2-8.7c-1.5-2.7-3.5-4.7-6-6.2 s-5.4-2.3-8.6-2.3c-3.2,0-5.9,0.8-8.5,2.3c-2.6,1.5-4.6,3.6-6.1,6.2c-1.5,2.6-2.2,5.5-2.2,8.7s0.8,6.2,2.3,8.9s3.5,4.7,6.1,6.1 C468,100.1,470.7,100.9,473.7,100.9L473.7,100.9z"/>
               </g>
@@ -34,8 +37,7 @@
           </div>
         </div>
 
-        <div class="loading-title">Loading model&hellip;</div>
-        <div class="loading-subtitle">This may take a moment for large models.</div>
+        <div class="loading-title">Loading model</div>
 
         <!-- Asset Download Indicator -->
         <div id="loadingDownload" class="loading-download">

From 466d068089ed49252f3e39a803fedf15556e6e3e Mon Sep 17 00:00:00 2001
From: Baruch Tabanpour <btaba@google.com>
Date: Fri, 22 May 2026 11:14:11 -0700
Subject: [PATCH 22/61] MuJoCo Studio: Support Ctrl+Option+Left Click Drag for
 perturbation

Support Ctrl+Option+Left Click Drag (maps to Ctrl+Alt+Left Click Drag in the application) to apply point force perturbations (translation) on the robot in the WebGL live viewer. This matches the Mac experience in the native MuJoCo simulate app and provides a fallback for Mac users where Right Click is difficult to trigger.

PiperOrigin-RevId: 919767954
Change-Id: I4a1d1c05784a8257d13f64f738af5b4841a615b3
---
 src/experimental/studio/app.cc | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc
index cf7d28e9..b221fa61 100644
--- a/src/experimental/studio/app.cc
+++ b/src/experimental/studio/app.cc
@@ -461,14 +461,20 @@ void App::HandleMouseEvents() {
     if (perturb_.select > 0) {
       mjtMouse action = mjMOUSE_NONE;
       if (ImGui::IsMouseDown(ImGuiMouseButton_Left)) {
-        action = io.KeyShift ? mjMOUSE_ROTATE_H : mjMOUSE_ROTATE_V;
+        if (io.KeyAlt) {
+          action = io.KeyShift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V;
+        } else {
+          action = io.KeyShift ? mjMOUSE_ROTATE_H : mjMOUSE_ROTATE_V;
+        }
       } else if (ImGui::IsMouseDown(ImGuiMouseButton_Right)) {
         action = io.KeyShift ? mjMOUSE_MOVE_H : mjMOUSE_MOVE_V;
       } else if (ImGui::IsMouseDown(ImGuiMouseButton_Middle)) {
         action = mjMOUSE_ZOOM;
       }
       const mjtPertBit active =
-          action == mjMOUSE_MOVE_V ? mjPERT_TRANSLATE : mjPERT_ROTATE;
+          (action == mjMOUSE_MOVE_V || action == mjMOUSE_MOVE_H)
+              ? mjPERT_TRANSLATE
+              : mjPERT_ROTATE;
       if (active != perturb_.active) {
         platform::InitPerturb(model(), data(), &camera_, &perturb_, active);
       }

From 7667e7110e8cc56135080624d0cac021edcb1ae7 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 23 May 2026 09:13:21 -0700
Subject: [PATCH 23/61] Introduce mjtBool for boolean types.

PiperOrigin-RevId: 920201340
Change-Id: I3b19c11a00357e06a1df08b8819832955c37f9ff
---
 doc/APIreference/APItypes.rst                  | 18 +++++++++++++++++-
 doc/programming/index.rst                      |  3 +--
 include/mujoco/mjassert.h                      | 11 ++++++-----
 include/mujoco/mjtype.h                        |  9 +++++++--
 .../mujoco/codegen/generate_spec_bindings.py   |  2 +-
 wasm/codegen/generators/constants.py           |  1 +
 wasm/codegen/generators/functions.py           |  2 +-
 7 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst
index 879d4bda..ba1a6135 100644
--- a/doc/APIreference/APItypes.rst
+++ b/doc/APIreference/APItypes.rst
@@ -84,13 +84,29 @@ changed by the user.
 mjtByte
 ^^^^^^^
 
-Byte type used to represent boolean variables.
+Byte type used to represent small integers and binary data.
 
 .. code-block:: C
 
    typedef unsigned char mjtByte;
 
 
+.. _mjtBool:
+
+mjtBool
+^^^^^^^
+
+Boolean type used to represent true/false values.
+
+.. code-block:: C
+
+   #ifndef __cplusplus
+     typedef _Bool mjtBool;
+   #else
+     typedef bool mjtBool;
+   #endif
+
+
 .. _mjtSize:
 
 mjtSize
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index 852685cb..99cc50dc 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -230,8 +230,7 @@ to which the symbol belongs. First we list the prefixes corresponding to type de
    Core simulation data structure (C struct), for example :ref:`mjModel`. If all characters
    after the prefix are capital, for example :ref:`mjMIN`, this is a macro or a symbol (#define).
 ``mjt``
-   Primitive type, for example :ref:`mjtGeom`. Except for mjtByte and mjtNum, all other
-   definitions in this family are enums.
+   Primitive type, for example :ref:`mjtNum` and :ref:`mjtGeom`. Most types in this family are enums.
 ``mjf``
    Callback function type, for example :ref:`mjfGeneric`.
 ``mjv``
diff --git a/include/mujoco/mjassert.h b/include/mujoco/mjassert.h
index 5c69f940..169cbb30 100644
--- a/include/mujoco/mjassert.h
+++ b/include/mujoco/mjassert.h
@@ -36,16 +36,17 @@
     typedef char mj_assert_##type[sizeof(type) == (size) ? 1 : -1]
 #endif
 
-// mjtnum.h
-MJ_ASSERT_SIZE(mjtByte, 1);
-MJ_ASSERT_SIZE(mjtSize, 8);
+// primitive types
 #if !defined(mjUSESINGLE)
 MJ_ASSERT_SIZE(mjtNum, 8);
 #else
 MJ_ASSERT_SIZE(mjtNum, 4);
 #endif
+MJ_ASSERT_SIZE(mjtByte, 1);
+MJ_ASSERT_SIZE(mjtBool, 1);
+MJ_ASSERT_SIZE(mjtSize, 8);
 
-// mjmodel.h
+// mjModel enums
 MJ_ASSERT_SIZE(mjtDisableBit, 4);
 MJ_ASSERT_SIZE(mjtEnableBit, 4);
 MJ_ASSERT_SIZE(mjtJoint, 4);
@@ -79,7 +80,7 @@ MJ_ASSERT_SIZE(mjtLRMode, 4);
 MJ_ASSERT_SIZE(mjtFlexSelf, 4);
 MJ_ASSERT_SIZE(mjtSDFType, 4);
 
-// mjdata.h
+// mjData enums
 MJ_ASSERT_SIZE(mjtState, 4);
 MJ_ASSERT_SIZE(mjtConstraint, 4);
 MJ_ASSERT_SIZE(mjtConstraintState, 4);
diff --git a/include/mujoco/mjtype.h b/include/mujoco/mjtype.h
index 62426a11..b3296824 100644
--- a/include/mujoco/mjtype.h
+++ b/include/mujoco/mjtype.h
@@ -15,6 +15,7 @@
 #ifndef MUJOCO_INCLUDE_MJTYPE_H_
 #define MUJOCO_INCLUDE_MJTYPE_H_
 
+#include <stdbool.h>
 #include <stdint.h>
 
 
@@ -33,9 +34,13 @@
 
 //---------------------------------- byte definition -----------------------------------------------
 
-typedef unsigned char mjtByte;    // used for true/false
-
+typedef unsigned char mjtByte;    // used for small integers and binary data
 
+#ifndef __cplusplus
+  typedef _Bool mjtBool;          // used for boolean values
+#else
+  typedef bool mjtBool;           // used for boolean values
+#endif
 
 //---------------------------------- size definition -----------------------------------------------
 
diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py
index 724325c6..5dee5684 100644
--- a/python/mujoco/codegen/generate_spec_bindings.py
+++ b/python/mujoco/codegen/generate_spec_bindings.py
@@ -22,7 +22,7 @@ from introspect import ast_nodes
 from introspect import structs
 
 
-SCALAR_TYPES = {'int', 'double', 'float', 'mjtByte', 'mjtNum'}
+SCALAR_TYPES = {'int', 'double', 'float', 'mjtByte', 'mjtBool', 'mjtNum'}
 
 # pylint: disable=bad-whitespace
 # key, parent, default, listname, objtype
diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py
index 2338dff0..fedb130f 100644
--- a/wasm/codegen/generators/constants.py
+++ b/wasm/codegen/generators/constants.py
@@ -24,6 +24,7 @@ PRIMITIVE_TYPES: Set[str] = {
     "float",
     "int",
     "mjtByte",
+    "mjtBool",
     "mjtMeshBuiltin",
     "mjtNum",
     "mjtObj",  # Adding this to the primitives because it is used as int,
diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py
index bc0f73d1..bb296d12 100644
--- a/wasm/codegen/generators/functions.py
+++ b/wasm/codegen/generators/functions.py
@@ -148,7 +148,7 @@ def get_param_string(p: ast_nodes.FunctionParameterDecl) -> str:
         return f"const String& {p.name}"
     elif (
         p.type.inner_type.name
-        in ["int", "float", "double", "mjtNum", "mjtByte"]
+        in ["int", "float", "double", "mjtNum", "mjtByte", "mjtBool"]
         and p.type.inner_type.is_const
     ):
       return f"const NumberArray& {p.name}"

From ddff7e3e2e3ae5d2f271677c2fab22179c67c896 Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 26 May 2026 02:36:12 -0700
Subject: [PATCH 24/61] MuJoCo Warp documentation: update feature parity
 section

PiperOrigin-RevId: 921334977
Change-Id: I1c20e30cde81ae18830cae894ac5aee99ac0dcb8
---
 doc/mjwarp/index.rst | 33 +++------------------------------
 1 file changed, 3 insertions(+), 30 deletions(-)

diff --git a/doc/mjwarp/index.rst b/doc/mjwarp/index.rst
index 2d4c0f83..c2cb0804 100644
--- a/doc/mjwarp/index.rst
+++ b/doc/mjwarp/index.rst
@@ -228,36 +228,9 @@ Feature Parity
 ==============
 
 MJWarp supports most of the main simulation features of MuJoCo, with a few exceptions. MJWarp will raise an exception if
-asked to copy to device an :ref:`mjModel` with field values referencing unsupported features.
-
-The following features are **not supported** in MJWarp:
-
-.. list-table::
-   :width: 90%
-   :align: left
-   :widths: 2 5
-   :header-rows: 1
-
-   * - Category
-     - Feature
-   * - :ref:`Integrator <mjtIntegrator>`
-     - ``IMPLICIT``, ``IMPLICITFAST`` not supported with fluid drag
-   * - :ref:`Solver <mjtSolver>`
-     - ``PGS``, ``noslip``, :ref:`islands <soIsland>`
-   * - Fluid Model
-     - :ref:`flEllipsoid`
-   * - :ref:`Sensors <mjtSensor>`
-     - ``GEOMDIST``, ``GEOMNORMAL``, ``GEOMFROMTO``
-   * - Flex
-     - ``VERTCOLLIDE=false``, ``INTERNAL=true``
-   * - Jacobian format
-     - ``SPARSE``
-   * - Option
-     - :ref:`contact override <COverride>`
-   * - Plugins
-     - ``All`` except ``SDF``
-   * - :ref:`User parameters <CUser>`
-     - ``All``
+asked to copy to device an :ref:`mjModel` with field values referencing unsupported features. For the most up-to-date
+feature availability, please see
+`MuJoCo API Compatibility <https://github.com/google-deepmind/mujoco_warp#mujoco-api-compatibility>`__.
 
 .. _mjwPerf:
 

From 393c1e421755e90c5dd2243e5ba81f6324f6d7aa Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Tue, 26 May 2026 05:35:15 -0700
Subject: [PATCH 25/61] Migrate `mjtByte` to `mjtBool` for boolean-valued
 fields.

PiperOrigin-RevId: 921408093
Change-Id: Icb00457836359779f0f72fc02d6a7c0dc9f6bd23
---
 doc/changelog.rst                             |  2 +
 doc/includes/references.h                     | 66 +++++++++----------
 include/mujoco/mjdata.h                       | 12 ++--
 include/mujoco/mjmodel.h                      | 36 +++++-----
 include/mujoco/mjxmacro.h                     | 48 +++++++-------
 include/mujoco/mujoco.h                       | 18 ++---
 python/mujoco/codegen/generate_cs_bindings.py |  1 +
 python/mujoco/functions.cc                    | 16 ++---
 python/mujoco/introspect/functions.py         | 22 +++----
 python/mujoco/introspect/structs.py           | 48 +++++++-------
 src/engine/engine_collision_driver.c          | 34 +++++-----
 src/engine/engine_collision_driver.h          |  2 +-
 src/engine/engine_collision_sdf.c             | 18 ++---
 src/engine/engine_derivative_fd.c             |  6 +-
 src/engine/engine_derivative_fd.h             |  4 +-
 src/engine/engine_forward.c                   |  2 +-
 src/engine/engine_print.c                     |  2 +
 src/engine/engine_ray.c                       | 16 ++---
 src/engine/engine_ray.h                       | 10 +--
 src/engine/engine_support.c                   |  6 +-
 src/engine/engine_util_misc.c                 |  6 +-
 src/engine/engine_util_solve.c                |  4 +-
 src/engine/engine_util_solve.h                |  4 +-
 src/experimental/platform/ux/imgui_widgets.cc |  8 +++
 src/experimental/platform/ux/imgui_widgets.h  |  3 +-
 src/user/user_api.cc                          |  6 +-
 src/user/user_api.h                           |  2 +-
 src/user/user_model.cc                        | 18 ++---
 test/fixture.cc                               |  4 ++
 test/user/user_objects_test.cc                |  2 +-
 test/xml/xml_native_reader_test.cc            |  2 +-
 wasm/codegen/generated/bindings.cc            | 30 ++++-----
 wasm/tests/bindings_test.ts                   |  8 +--
 33 files changed, 242 insertions(+), 224 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index b339da1f..c09fe3dd 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -21,6 +21,8 @@ General
   :ref:`mj_compile`, per-category timings (total, assets, mesh loading, convex hull, normals, inertia, BVH, octree,
   textures) are available via ``mjs_getTimer(spec)``. The :ref:`compile<saCompile>` sample prints a detailed timing
   breakdown when run without an output file.
+- Added :ref:`mjtBool` to represent boolean variables, replacing :ref:`mjtByte` across all boolean fields in
+  :ref:`mjModel`, :ref:`mjData`, and public C API function signatures.
 
 .. admonition:: Breaking API changes
    :class: attention
diff --git a/doc/includes/references.h b/doc/includes/references.h
index 7f6371a2..b158c3e1 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -123,10 +123,10 @@ struct mjData_ {
   int     nv_awake;          // number of awake dofs
 
   // flags marking lazily evaluated stages
-  mjtByte flg_energypos;     // has mj_energyPos been called
-  mjtByte flg_energyvel;     // has mj_energyVel been called
-  mjtByte flg_subtreevel;    // has mj_subtreeVel been called
-  mjtByte flg_rnepost;       // has mj_rnePostConstraint been called
+  mjtBool flg_energypos;     // has mj_energyPos been called
+  mjtBool flg_energyvel;     // has mj_energyVel been called
+  mjtBool flg_subtreevel;    // has mj_subtreeVel been called
+  mjtBool flg_rnepost;       // has mj_rnePostConstraint been called
 
   // global properties
   mjtNum  time;              // simulation time
@@ -152,7 +152,7 @@ struct mjData_ {
   mjtNum* ctrl;              // control                                          (nu x 1)
   mjtNum* qfrc_applied;      // applied generalized force                        (nv x 1)
   mjtNum* xfrc_applied;      // applied Cartesian force/torque                   (nbody x 6)
-  mjtByte* eq_active;        // enable/disable constraints                       (neq x 1)
+  mjtBool* eq_active;        // enable/disable constraints                       (neq x 1)
 
   // mocap data
   mjtNum* mocap_pos;         // positions of mocap bodies                        (nmocap x 3)
@@ -233,7 +233,7 @@ struct mjData_ {
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collision/mj_collideTree
-  mjtByte* bvh_active;       // was bounding volume checked for collision        (nbvh x 1)
+  mjtBool* bvh_active;       // was bounding volume checked for collision        (nbvh x 1)
 
   // computed by mj_updateSleep
   int*    tree_awake;        // is tree awake; 0: asleep; 1: awake               (ntree x 1)
@@ -741,9 +741,9 @@ struct mjModel_ {
   int*      jnt_bodyid;           // id of joint's body                       (njnt x 1)
   int*      jnt_actuatorid;       // actuator contributing damping / armature (njnt x 1)
   int*      jnt_group;            // group for visibility                     (njnt x 1)
-  mjtByte*  jnt_limited;          // does joint have limits                   (njnt x 1)
-  mjtByte*  jnt_actfrclimited;    // does joint have actuator force limits    (njnt x 1)
-  mjtByte*  jnt_actgravcomp;      // is gravcomp force applied via actuators  (njnt x 1)
+  mjtBool*  jnt_limited;          // does joint have limits                   (njnt x 1)
+  mjtBool*  jnt_actfrclimited;    // does joint have actuator force limits    (njnt x 1)
+  mjtBool*  jnt_actgravcomp;      // is gravcomp force applied via actuators  (njnt x 1)
   mjtNum*   jnt_solref;           // constraint solver reference: limit       (njnt x mjNREF)
   mjtNum*   jnt_solimp;           // constraint solver impedance: limit       (njnt x mjNIMP)
   mjtNum*   jnt_pos;              // local anchor position                    (njnt x 3)
@@ -842,11 +842,11 @@ struct mjModel_ {
   int*      light_targetbodyid;   // id of targeted body; -1: none            (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)
+  mjtBool*  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)
+  mjtBool*  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)
   mjtNum*   light_poscom0;        // global position rel. to sub-com in qpos0 (nlight x 3)
@@ -870,7 +870,7 @@ struct mjModel_ {
   mjtNum*   flex_friction;        // friction for (slide, spin, roll)         (nflex x 3)
   mjtNum*   flex_margin;          // geometric inflation for contact          (nflex x 1)
   mjtNum*   flex_gap;             // additional contact detection buffer      (nflex x 1)
-  mjtByte*  flex_internal;        // internal flex collision enabled          (nflex x 1)
+  mjtBool*  flex_internal;        // internal flex collision enabled          (nflex x 1)
   int*      flex_selfcollide;     // self collision mode (mjtFlexSelf)        (nflex x 1)
   int*      flex_activelayers;    // number of active element layers, 3D only (nflex x 1)
   int*      flex_passive;         // passive collisions enabled               (nflex x 1)
@@ -926,10 +926,10 @@ struct mjModel_ {
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   int*      flex_edgeequality;    // 0:none, 1:edges, 2:vertices, 3:strain    (nflex x 1)
-  mjtByte*  flex_rigid;           // are all vertices in the same body        (nflex x 1)
-  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
-  mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
-  mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
+  mjtBool*  flex_rigid;           // are all vertices in the same body        (nflex x 1)
+  mjtBool*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
+  mjtBool*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
+  mjtBool*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
   int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)
   int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
   int*      flexedge_J_rownnz;    // number of non-zeros in Jacobian row      (nflexedge x 1)
@@ -1020,7 +1020,7 @@ struct mjModel_ {
 
   // materials
   int*      mat_texid;            // indices of textures; -1: none            (nmat x mjNTEXROLE)
-  mjtByte*  mat_texuniform;       // make texture cube uniform                (nmat x 1)
+  mjtBool*  mat_texuniform;       // make texture cube uniform                (nmat x 1)
   float*    mat_texrepeat;        // texture repetition for 2d mapping        (nmat x 2)
   float*    mat_emission;         // emission (x rgb)                         (nmat x 1)
   float*    mat_specular;         // specular (x white)                       (nmat x 1)
@@ -1050,7 +1050,7 @@ struct mjModel_ {
   int*      eq_obj1id;            // id of object 1                           (neq x 1)
   int*      eq_obj2id;            // id of object 2                           (neq x 1)
   int*      eq_objtype;           // type of both objects (mjtObj)            (neq x 1)
-  mjtByte*  eq_active0;           // initial enable/disable constraint state  (neq x 1)
+  mjtBool*  eq_active0;           // initial enable/disable constraint state  (neq x 1)
   mjtNum*   eq_solref;            // constraint solver reference              (neq x mjNREF)
   mjtNum*   eq_solimp;            // constraint solver impedance              (neq x mjNIMP)
   mjtNum*   eq_data;              // numeric data for constraint              (neq x mjNEQDATA)
@@ -1066,8 +1066,8 @@ struct mjModel_ {
   int*      ten_J_rownnz;         // number of non-zeros in Jacobian row      (ntendon x 1)
   int*      ten_J_rowadr;         // row start address in colind array        (ntendon x 1)
   int*      ten_J_colind;         // column indices in sparse Jacobian        (nJten 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)
+  mjtBool*  tendon_limited;       // does tendon have length limits           (ntendon x 1)
+  mjtBool*  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)
@@ -1108,13 +1108,13 @@ struct mjModel_ {
   int*      actuator_history;     // history buffer: [nsample, interp]        (nu x 2)
   int*      actuator_historyadr;  // address in history buffer; -1: none      (nu x 1)
   mjtNum*   actuator_delay;       // delay time in seconds; 0: no delay       (nu x 1)
-  mjtByte*  actuator_ctrllimited; // is control limited                       (nu x 1)
-  mjtByte*  actuator_forcelimited;// is force limited                         (nu x 1)
-  mjtByte*  actuator_actlimited;  // is activation limited                    (nu x 1)
+  mjtBool*  actuator_ctrllimited; // is control limited                       (nu x 1)
+  mjtBool*  actuator_forcelimited;// is force limited                         (nu x 1)
+  mjtBool*  actuator_actlimited;  // is activation limited                    (nu x 1)
   mjtNum*   actuator_dynprm;      // dynamics parameters                      (nu x mjNDYN)
   mjtNum*   actuator_gainprm;     // gain parameters                          (nu x mjNGAIN)
   mjtNum*   actuator_biasprm;     // bias parameters                          (nu x mjNBIAS)
-  mjtByte*  actuator_actearly;    // step activation before force             (nu x 1)
+  mjtBool*  actuator_actearly;    // step activation before force             (nu x 1)
   mjtNum*   actuator_ctrlrange;   // range of controls                        (nu x 2)
   mjtNum*   actuator_forcerange;  // range of forces                          (nu x 2)
   mjtNum*   actuator_actrange;    // range of activations                     (nu x 2)
@@ -3373,10 +3373,10 @@ void mj_loadAllPluginLibraries(const char* directory, mjfPluginLibraryLoadCallba
 int mj_version(void);
 const char* mj_versionString(void);
 mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-              const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+              const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
               int geomid[1], mjtNum normal[3]);
 void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
-                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                 const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                  int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff);
 mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
                     const mjtNum pnt[3], const mjtNum vec[3], mjtNum normal[3]);
@@ -3386,8 +3386,8 @@ mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3
                    const mjtNum pnt[3], const mjtNum vec[3], int geomtype,
                    mjtNum normal[3]);
 mjtNum mj_rayFlex(const mjModel* m, const mjData* d, int flex_layer,
-                  mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
-                  mjtByte flg_skin, int flexid, const mjtNum pnt[3],
+                  mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
+                  mjtBool flg_skin, int flexid, const mjtNum pnt[3],
                   const mjtNum vec[3], int vertid[1], mjtNum normal[3]);
 mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
                    const mjtNum pnt[3], const mjtNum vec[3], int vertid[1]);
@@ -3582,10 +3582,10 @@ mjtNum mju_cholFactorBand(mjtNum* mat, int ntotal, int nband, int ndense,
 void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
                        int ntotal, int nband, int ndense);
 void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
-                    mjtByte flg_sym);
+                    mjtBool flg_sym);
 void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
 void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
-                       int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
+                       int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym);
 int mju_bandDiag(int i, int ntotal, int nband, int ndense);
 int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
 int mju_boxQP(mjtNum* res, mjtNum* R, int* index, const mjtNum* H, const mjtNum* g, int n,
@@ -3625,9 +3625,9 @@ 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,
+void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_centered,
                       mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
-void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
+void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_actuation,
                    mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
                    mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                    mjtNum *DmDq);
@@ -3741,7 +3741,7 @@ int mjs_setName(mjsElement* element, const char* name);
 void mjs_setBuffer(mjByteVec* dest, const void* array, int size);
 void mjs_setString(mjString* dest, const char* text);
 void mjs_setStringVec(mjStringVec* dest, const char* text);
-mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
+mjtBool mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
 void mjs_appendString(mjStringVec* dest, const char* text);
 void mjs_setInt(mjIntVec* dest, const int* array, int size);
 void mjs_appendIntVec(mjIntVecVec* dest, const int* array, int size);
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
index e7e56e9e..486c382a 100644
--- a/include/mujoco/mjdata.h
+++ b/include/mujoco/mjdata.h
@@ -146,10 +146,10 @@ struct mjData_ {
   int     nv_awake;          // number of awake dofs
 
   // flags marking lazily evaluated stages
-  mjtByte flg_energypos;     // has mj_energyPos been called
-  mjtByte flg_energyvel;     // has mj_energyVel been called
-  mjtByte flg_subtreevel;    // has mj_subtreeVel been called
-  mjtByte flg_rnepost;       // has mj_rnePostConstraint been called
+  mjtBool flg_energypos;     // has mj_energyPos been called
+  mjtBool flg_energyvel;     // has mj_energyVel been called
+  mjtBool flg_subtreevel;    // has mj_subtreeVel been called
+  mjtBool flg_rnepost;       // has mj_rnePostConstraint been called
 
   // global properties
   mjtNum  time;              // simulation time
@@ -175,7 +175,7 @@ struct mjData_ {
   mjtNum* ctrl;              // control                                          (nu x 1)
   mjtNum* qfrc_applied;      // applied generalized force                        (nv x 1)
   mjtNum* xfrc_applied;      // applied Cartesian force/torque                   (nbody x 6)
-  mjtByte* eq_active;        // enable/disable constraints                       (neq x 1)
+  mjtBool* eq_active;        // enable/disable constraints                       (neq x 1)
 
   // mocap data
   mjtNum* mocap_pos;         // positions of mocap bodies                        (nmocap x 3)
@@ -256,7 +256,7 @@ struct mjData_ {
   mjtNum* qLDiagInv;         // 1/diag(D)                                        (nv x 1)
 
   // computed by mj_collision/mj_collideTree
-  mjtByte* bvh_active;       // was bounding volume checked for collision        (nbvh x 1)
+  mjtBool* bvh_active;       // was bounding volume checked for collision        (nbvh x 1)
 
   // computed by mj_updateSleep
   int*    tree_awake;        // is tree awake; 0: asleep; 1: awake               (ntree x 1)
diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h
index 2207581a..2b1e013b 100644
--- a/include/mujoco/mjmodel.h
+++ b/include/mujoco/mjmodel.h
@@ -411,9 +411,9 @@ struct mjModel_ {
   int*      jnt_bodyid;           // id of joint's body                       (njnt x 1)
   int*      jnt_actuatorid;       // actuator contributing damping / armature (njnt x 1)
   int*      jnt_group;            // group for visibility                     (njnt x 1)
-  mjtByte*  jnt_limited;          // does joint have limits                   (njnt x 1)
-  mjtByte*  jnt_actfrclimited;    // does joint have actuator force limits    (njnt x 1)
-  mjtByte*  jnt_actgravcomp;      // is gravcomp force applied via actuators  (njnt x 1)
+  mjtBool*  jnt_limited;          // does joint have limits                   (njnt x 1)
+  mjtBool*  jnt_actfrclimited;    // does joint have actuator force limits    (njnt x 1)
+  mjtBool*  jnt_actgravcomp;      // is gravcomp force applied via actuators  (njnt x 1)
   mjtNum*   jnt_solref;           // constraint solver reference: limit       (njnt x mjNREF)
   mjtNum*   jnt_solimp;           // constraint solver impedance: limit       (njnt x mjNIMP)
   mjtNum*   jnt_pos;              // local anchor position                    (njnt x 3)
@@ -512,11 +512,11 @@ struct mjModel_ {
   int*      light_targetbodyid;   // id of targeted body; -1: none            (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)
+  mjtBool*  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)
+  mjtBool*  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)
   mjtNum*   light_poscom0;        // global position rel. to sub-com in qpos0 (nlight x 3)
@@ -540,7 +540,7 @@ struct mjModel_ {
   mjtNum*   flex_friction;        // friction for (slide, spin, roll)         (nflex x 3)
   mjtNum*   flex_margin;          // geometric inflation for contact          (nflex x 1)
   mjtNum*   flex_gap;             // additional contact detection buffer      (nflex x 1)
-  mjtByte*  flex_internal;        // internal flex collision enabled          (nflex x 1)
+  mjtBool*  flex_internal;        // internal flex collision enabled          (nflex x 1)
   int*      flex_selfcollide;     // self collision mode (mjtFlexSelf)        (nflex x 1)
   int*      flex_activelayers;    // number of active element layers, 3D only (nflex x 1)
   int*      flex_passive;         // passive collisions enabled               (nflex x 1)
@@ -596,10 +596,10 @@ struct mjModel_ {
   mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
   mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
   int*      flex_edgeequality;    // 0:none, 1:edges, 2:vertices, 3:strain    (nflex x 1)
-  mjtByte*  flex_rigid;           // are all vertices in the same body        (nflex x 1)
-  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
-  mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
-  mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
+  mjtBool*  flex_rigid;           // are all vertices in the same body        (nflex x 1)
+  mjtBool*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
+  mjtBool*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
+  mjtBool*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
   int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)
   int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
   int*      flexedge_J_rownnz;    // number of non-zeros in Jacobian row      (nflexedge x 1)
@@ -690,7 +690,7 @@ struct mjModel_ {
 
   // materials
   int*      mat_texid;            // indices of textures; -1: none            (nmat x mjNTEXROLE)
-  mjtByte*  mat_texuniform;       // make texture cube uniform                (nmat x 1)
+  mjtBool*  mat_texuniform;       // make texture cube uniform                (nmat x 1)
   float*    mat_texrepeat;        // texture repetition for 2d mapping        (nmat x 2)
   float*    mat_emission;         // emission (x rgb)                         (nmat x 1)
   float*    mat_specular;         // specular (x white)                       (nmat x 1)
@@ -720,7 +720,7 @@ struct mjModel_ {
   int*      eq_obj1id;            // id of object 1                           (neq x 1)
   int*      eq_obj2id;            // id of object 2                           (neq x 1)
   int*      eq_objtype;           // type of both objects (mjtObj)            (neq x 1)
-  mjtByte*  eq_active0;           // initial enable/disable constraint state  (neq x 1)
+  mjtBool*  eq_active0;           // initial enable/disable constraint state  (neq x 1)
   mjtNum*   eq_solref;            // constraint solver reference              (neq x mjNREF)
   mjtNum*   eq_solimp;            // constraint solver impedance              (neq x mjNIMP)
   mjtNum*   eq_data;              // numeric data for constraint              (neq x mjNEQDATA)
@@ -736,8 +736,8 @@ struct mjModel_ {
   int*      ten_J_rownnz;         // number of non-zeros in Jacobian row      (ntendon x 1)
   int*      ten_J_rowadr;         // row start address in colind array        (ntendon x 1)
   int*      ten_J_colind;         // column indices in sparse Jacobian        (nJten 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)
+  mjtBool*  tendon_limited;       // does tendon have length limits           (ntendon x 1)
+  mjtBool*  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)
@@ -778,13 +778,13 @@ struct mjModel_ {
   int*      actuator_history;     // history buffer: [nsample, interp]        (nu x 2)
   int*      actuator_historyadr;  // address in history buffer; -1: none      (nu x 1)
   mjtNum*   actuator_delay;       // delay time in seconds; 0: no delay       (nu x 1)
-  mjtByte*  actuator_ctrllimited; // is control limited                       (nu x 1)
-  mjtByte*  actuator_forcelimited;// is force limited                         (nu x 1)
-  mjtByte*  actuator_actlimited;  // is activation limited                    (nu x 1)
+  mjtBool*  actuator_ctrllimited; // is control limited                       (nu x 1)
+  mjtBool*  actuator_forcelimited;// is force limited                         (nu x 1)
+  mjtBool*  actuator_actlimited;  // is activation limited                    (nu x 1)
   mjtNum*   actuator_dynprm;      // dynamics parameters                      (nu x mjNDYN)
   mjtNum*   actuator_gainprm;     // gain parameters                          (nu x mjNGAIN)
   mjtNum*   actuator_biasprm;     // bias parameters                          (nu x mjNBIAS)
-  mjtByte*  actuator_actearly;    // step activation before force             (nu x 1)
+  mjtBool*  actuator_actearly;    // step activation before force             (nu x 1)
   mjtNum*   actuator_ctrlrange;   // range of controls                        (nu x 2)
   mjtNum*   actuator_forcerange;  // range of forces                          (nu x 2)
   mjtNum*   actuator_actrange;    // range of activations                     (nu x 2)
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
index 98b5e75b..adb7686a 100644
--- a/include/mujoco/mjxmacro.h
+++ b/include/mujoco/mjxmacro.h
@@ -318,9 +318,9 @@
     X   ( int,     jnt_bodyid,            njnt,          1                    ) \
     X   ( int,     jnt_actuatorid,        njnt,          1                    ) \
     X   ( int,     jnt_group,             njnt,          1                    ) \
-    X   ( mjtByte, jnt_limited,           njnt,          1                    ) \
-    X   ( mjtByte, jnt_actfrclimited,     njnt,          1                    ) \
-    X   ( mjtByte, jnt_actgravcomp,       njnt,          1                    ) \
+    X   ( mjtBool, jnt_limited,           njnt,          1                    ) \
+    X   ( mjtBool, jnt_actfrclimited,     njnt,          1                    ) \
+    X   ( mjtBool, jnt_actgravcomp,       njnt,          1                    ) \
     X   ( mjtNum,  jnt_solref,            njnt,          mjNREF               ) \
     X   ( mjtNum,  jnt_solimp,            njnt,          mjNIMP               ) \
     X   ( mjtNum,  jnt_pos,               njnt,          3                    ) \
@@ -419,11 +419,11 @@
     X   ( int,     light_targetbodyid,    nlight,        1                    ) \
     X   ( int,     light_type,            nlight,        1                    ) \
     X   ( int,     light_texid,           nlight,        1                    ) \
-    X   ( mjtByte, light_castshadow,      nlight,        1                    ) \
+    X   ( mjtBool, light_castshadow,      nlight,        1                    ) \
     X   ( float,   light_bulbradius,      nlight,        1                    ) \
     X   ( float,   light_intensity,       nlight,        1                    ) \
     X   ( float,   light_range,           nlight,        1                    ) \
-    X   ( mjtByte, light_active,          nlight,        1                    ) \
+    X   ( mjtBool, light_active,          nlight,        1                    ) \
     X   ( mjtNum,  light_pos,             nlight,        3                    ) \
     X   ( mjtNum,  light_dir,             nlight,        3                    ) \
     X   ( mjtNum,  light_poscom0,         nlight,        3                    ) \
@@ -447,7 +447,7 @@
     X   ( mjtNum,  flex_friction,         nflex,         3                    ) \
     X   ( mjtNum,  flex_margin,           nflex,         1                    ) \
     X   ( mjtNum,  flex_gap,              nflex,         1                    ) \
-    X   ( mjtByte, flex_internal,         nflex,         1                    ) \
+    X   ( mjtBool, flex_internal,         nflex,         1                    ) \
     X   ( int,     flex_selfcollide,      nflex,         1                    ) \
     X   ( int,     flex_activelayers,     nflex,         1                    ) \
     X   ( int,     flex_passive,          nflex,         1                    ) \
@@ -501,10 +501,10 @@
     X   ( mjtNum,  flex_edgestiffness,    nflex,         1                    ) \
     X   ( mjtNum,  flex_edgedamping,      nflex,         1                    ) \
     X   ( int,     flex_edgeequality,     nflex,         1                    ) \
-    X   ( mjtByte, flex_rigid,            nflex,         1                    ) \
-    X   ( mjtByte, flexedge_rigid,        nflexedge,     1                    ) \
-    X   ( mjtByte, flex_centered,         nflex,         1                    ) \
-    X   ( mjtByte, flex_flatskin,         nflex,         1                    ) \
+    X   ( mjtBool, flex_rigid,            nflex,         1                    ) \
+    X   ( mjtBool, flexedge_rigid,        nflexedge,     1                    ) \
+    X   ( mjtBool, flex_centered,         nflex,         1                    ) \
+    X   ( mjtBool, flex_flatskin,         nflex,         1                    ) \
     X   ( int,     flex_bvhadr,           nflex,         1                    ) \
     X   ( int,     flex_bvhnum,           nflex,         1                    ) \
     X   ( int,     flexedge_J_rownnz,     nflexedge,     1                    ) \
@@ -595,7 +595,7 @@
 
 #define MJMODEL_POINTERS_MATERIAL                                               \
     X   ( int,     mat_texid,             nmat,          mjNTEXROLE           ) \
-    X   ( mjtByte, mat_texuniform,        nmat,          1                    ) \
+    X   ( mjtBool, mat_texuniform,        nmat,          1                    ) \
     X   ( float,   mat_texrepeat,         nmat,          2                    ) \
     X   ( float,   mat_emission,          nmat,          1                    ) \
     X   ( float,   mat_specular,          nmat,          1                    ) \
@@ -625,7 +625,7 @@
     X   ( int,     eq_obj1id,             neq,           1                    ) \
     X   ( int,     eq_obj2id,             neq,           1                    ) \
     X   ( int,     eq_objtype,            neq,           1                    ) \
-    X   ( mjtByte, eq_active0,            neq,           1                    ) \
+    X   ( mjtBool, eq_active0,            neq,           1                    ) \
     X   ( mjtNum,  eq_solref,             neq,           mjNREF               ) \
     X   ( mjtNum,  eq_solimp,             neq,           mjNIMP               ) \
     X   ( mjtNum,  eq_data,               neq,           mjNEQDATA            )
@@ -641,8 +641,8 @@
     X   ( int,     ten_J_rownnz,          ntendon,       1                    ) \
     X   ( int,     ten_J_rowadr,          ntendon,       1                    ) \
     X   ( int,     ten_J_colind,          nJten,         1                    ) \
-    X   ( mjtByte, tendon_limited,        ntendon,       1                    ) \
-    X   ( mjtByte, tendon_actfrclimited,  ntendon,       1                    ) \
+    X   ( mjtBool, tendon_limited,        ntendon,       1                    ) \
+    X   ( mjtBool, tendon_actfrclimited,  ntendon,       1                    ) \
     X   ( mjtNum,  tendon_width,          ntendon,       1                    ) \
     X   ( mjtNum,  tendon_solref_lim,     ntendon,       mjNREF               ) \
     X   ( mjtNum,  tendon_solimp_lim,     ntendon,       mjNIMP               ) \
@@ -678,13 +678,13 @@
     X   ( int,     actuator_history,      nu,            2                    ) \
     X   ( int,     actuator_historyadr,   nu,            1                    ) \
     X   ( mjtNum,  actuator_delay,        nu,            1                    ) \
-    X   ( mjtByte, actuator_ctrllimited,  nu,            1                    ) \
-    X   ( mjtByte, actuator_forcelimited, nu,            1                    ) \
-    X   ( mjtByte, actuator_actlimited,   nu,            1                    ) \
+    X   ( mjtBool, actuator_ctrllimited,  nu,            1                    ) \
+    X   ( mjtBool, actuator_forcelimited, nu,            1                    ) \
+    X   ( mjtBool, actuator_actlimited,   nu,            1                    ) \
     X   ( mjtNum,  actuator_dynprm,       nu,            mjNDYN               ) \
     X   ( mjtNum,  actuator_gainprm,      nu,            mjNGAIN              ) \
     X   ( mjtNum,  actuator_biasprm,      nu,            mjNBIAS              ) \
-    X   ( mjtByte, actuator_actearly,     nu,            1                    ) \
+    X   ( mjtBool, actuator_actearly,     nu,            1                    ) \
     X   ( mjtNum,  actuator_ctrlrange,    nu,            2                    ) \
     X   ( mjtNum,  actuator_forcerange,   nu,            2                    ) \
     X   ( mjtNum,  actuator_actrange,     nu,            2                    ) \
@@ -828,7 +828,7 @@
     X   ( mjtNum,    ctrl,              nu,          1           ) \
     X   ( mjtNum,    qfrc_applied,      nv,          1           ) \
     X   ( mjtNum,    xfrc_applied,      nbody,       6           ) \
-    X   ( mjtByte,   eq_active,         neq,         1           ) \
+    X   ( mjtBool,   eq_active,         neq,         1           ) \
     X   ( mjtNum,    mocap_pos,         nmocap,      3           ) \
     X   ( mjtNum,    mocap_quat,        nmocap,      4           ) \
     X   ( mjtNum,    qacc,              nv,          1           ) \
@@ -879,7 +879,7 @@
     XNV ( mjtNum,    M,                 nC,          1           ) \
     XNV ( mjtNum,    qLD,               nC,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
-    X   ( mjtByte,   bvh_active,        nbvh,        1           ) \
+    X   ( mjtBool,   bvh_active,        nbvh,        1           ) \
     X   ( int,       tree_awake,        ntree,       1           ) \
     X   ( int,       body_awake,        nbody,       1           ) \
     X   ( int,       body_awake_ind,    nbody,       1           ) \
@@ -1032,10 +1032,10 @@
     X( int,       nbody_awake        ) \
     X( int,       nparent_awake      ) \
     X( int,       nv_awake           ) \
-    X( mjtByte,   flg_energypos      ) \
-    X( mjtByte,   flg_energyvel      ) \
-    X( mjtByte,   flg_subtreevel     ) \
-    X( mjtByte,   flg_rnepost        ) \
+    X( mjtBool,   flg_energypos      ) \
+    X( mjtBool,   flg_energyvel      ) \
+    X( mjtBool,   flg_subtreevel     ) \
+    X( mjtBool,   flg_rnepost        ) \
     X( mjtNum,    time               ) \
     X( uintptr_t, threadpool         )
 
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index ac76dc0f..95c99041 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -680,7 +680,7 @@ MJAPI const char* mj_versionString(void);
 // geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.
 // Nullable: geomgroup, geomid, normal
 MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-                    const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                    const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                     int geomid[1], mjtNum normal[3]);
 
 // Intersect multiple rays emanating from a single point, compute normals if given.
@@ -688,7 +688,7 @@ MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], cons
 // Geoms further than cutoff are ignored.
 // Nullable: geomgroup, geomid, normal
 MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
-                       const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                       const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                        int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff);
 
 // Intersect ray with hfield; return nearest distance or -1 if no intersection.
@@ -711,8 +711,8 @@ MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum
 // and also output nearest vertex id and surface normal.
 // Nullable: vertid, normal
 MJAPI mjtNum mj_rayFlex(const mjModel* m, const mjData* d, int flex_layer,
-                        mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
-                        mjtByte flg_skin, int flexid, const mjtNum pnt[3],
+                        mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
+                        mjtBool flg_skin, int flexid, const mjtNum pnt[3],
                         const mjtNum vec[3], int vertid[1], mjtNum normal[3]);
 
 // Intersect ray with skin; return nearest distance or -1 if no intersection,
@@ -1304,14 +1304,14 @@ MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 
 // Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.
 MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
-                          mjtByte flg_sym);
+                          mjtBool flg_sym);
 
 // Convert dense matrix to banded matrix.
 MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
 
 // Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.
 MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
-                             int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
+                             int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym);
 
 // Address of diagonal element i in band-dense matrix representation.
 MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
@@ -1461,7 +1461,7 @@ MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNu
 //      D: (nsensordata x 2*nv+na)
 //      C: (nsensordata x nu)
 // Nullable: A, B, C, D
-MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
+MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_centered,
                             mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
 
 // Finite differenced Jacobians of (force, sensors) = mj_inverse(state, acceleration)
@@ -1480,7 +1480,7 @@ MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg
 //     optionally computes mass matrix Jacobian DmDq
 //     flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse
 // Nullable: DfDq, DfDv, DfDa, DsDq, DsDv, DsDa, DmDq
-MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
+MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_actuation,
                          mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
                          mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                          mjtNum *DmDq);
@@ -1860,7 +1860,7 @@ MJAPI void mjs_setString(mjString* dest, const char* text);
 MJAPI void mjs_setStringVec(mjStringVec* dest, const char* text);
 
 // Set entry in string vector.
-MJAPI mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
+MJAPI mjtBool mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
 
 // Append text entry to string vector.
 MJAPI void mjs_appendString(mjStringVec* dest, const char* text);
diff --git a/python/mujoco/codegen/generate_cs_bindings.py b/python/mujoco/codegen/generate_cs_bindings.py
index ed014515..bc94dcca 100644
--- a/python/mujoco/codegen/generate_cs_bindings.py
+++ b/python/mujoco/codegen/generate_cs_bindings.py
@@ -136,6 +136,7 @@ C_TO_CS_TYPE = {
     'void': 'void',
     'mjtNum': 'double',
     'mjtByte': 'byte',
+    'mjtBool': 'byte',
     'mjtSize': 'UInt64' if ctypes.sizeof(ctypes.c_size_t) == 8 else 'UInt32',
 }
 
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index 5787d756..e509bce3 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -732,7 +732,7 @@ PYBIND11_MODULE(_functions, pymodule) {
          Eigen::Ref<const EigenVectorX> vec,
          std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
              geomgroup,
-         mjtByte flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
+         mjtBool flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
          Eigen::Ref<EigenVectorX> dist,
          std::optional<Eigen::Ref<EigenVectorX>> normal,
          int nray, mjtNum cutoff) {
@@ -758,7 +758,7 @@ PYBIND11_MODULE(_functions, pymodule) {
              const mjtNum(*vec)[3],
              std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
                  geomgroup,
-             mjtByte flg_static, int bodyexclude,
+             mjtBool flg_static, int bodyexclude,
              std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> geomid,
              std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
             return mj_ray(m, d, &(*pnt)[0], &(*vec)[0],
@@ -817,8 +817,8 @@ PYBIND11_MODULE(_functions, pymodule) {
       "mj_rayFlex",
       util::UnwrapArgs(
           [](const raw::MjModel* m, const raw::MjData* d, int flex_layer,
-             mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
-             mjtByte flg_skin, int flexid, const mjtNum(*pnt)[3],
+             mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
+             mjtBool flg_skin, int flexid, const mjtNum(*pnt)[3],
              const mjtNum(*vec)[3],
              std::optional<Eigen::Ref<Eigen::Vector<int, 1>>> vertid,
              std::optional<Eigen::Ref<Eigen::Vector<mjtNum, 3>>> normal) {
@@ -1370,7 +1370,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mju_band2Dense>(
       pymodule,
       [](Eigen::Ref<EigenArrayXX> res, Eigen::Ref<const EigenVectorX> mat,
-         int ntotal, int nband, int ndense, mjtByte flg_sym) {
+         int ntotal, int nband, int ndense, mjtBool flg_sym) {
         int nMat = (ntotal - ndense) * nband + ndense * ntotal;
         if (mat.size() != nMat) {
           throw py::type_error(
@@ -1406,7 +1406,7 @@ PYBIND11_MODULE(_functions, pymodule) {
       pymodule,
       [](Eigen::Ref<EigenVectorX> res, Eigen::Ref<const EigenArrayXX> mat,
          Eigen::Ref<const EigenArrayXX> vec, int ntotal, int nband, int ndense,
-         int nVec, mjtByte flg_sym) {
+         int nVec, mjtBool flg_sym) {
         int nMat = (ntotal - ndense) * nband + ndense * ntotal;
         if (mat.size() != nMat) {
           throw py::type_error(
@@ -1576,7 +1576,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mjd_transitionFD>(
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d,
-         mjtNum eps, mjtByte flg_centered,
+         mjtNum eps, mjtBool flg_centered,
          std::optional<Eigen::Ref<EigenArrayXX>> A,
          std::optional<Eigen::Ref<EigenArrayXX>> B,
          std::optional<Eigen::Ref<EigenArrayXX>> C,
@@ -1607,7 +1607,7 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mjd_inverseFD>(
       pymodule,
       [](const raw::MjModel* m, raw::MjData* d,
-         mjtNum eps, mjtByte flg_actuation,
+         mjtNum eps, mjtBool flg_actuation,
          std::optional<Eigen::Ref<EigenArrayXX>> DfDq,
          std::optional<Eigen::Ref<EigenArrayXX>> DfDv,
          std::optional<Eigen::Ref<EigenArrayXX>> DfDa,
diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py
index 53061ac2..c0bce30f 100644
--- a/python/mujoco/introspect/functions.py
+++ b/python/mujoco/introspect/functions.py
@@ -4088,7 +4088,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_static',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='bodyexclude',
@@ -4152,7 +4152,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_static',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='bodyexclude',
@@ -4359,19 +4359,19 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_vert',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='flg_edge',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='flg_face',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='flg_skin',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='flexid',
@@ -8479,7 +8479,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_sym',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
          ),
          doc='Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.',  # pylint: disable=line-too-long
@@ -8557,7 +8557,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_sym',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
          ),
          doc='Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.',  # pylint: disable=line-too-long
@@ -9404,7 +9404,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_centered',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='A',
@@ -9460,7 +9460,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
              ),
              FunctionParameterDecl(
                  name='flg_actuation',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
              ),
              FunctionParameterDecl(
                  name='DfDq',
@@ -11627,7 +11627,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
     ('mjs_setInStringVec',
      FunctionDecl(
          name='mjs_setInStringVec',
-         return_type=ValueType(name='mjtByte'),
+         return_type=ValueType(name='mjtBool'),
          parameters=(
              FunctionParameterDecl(
                  name='dest',
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
index 443b9f12..119219ee 100644
--- a/python/mujoco/introspect/structs.py
+++ b/python/mujoco/introspect/structs.py
@@ -1692,7 +1692,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='jnt_limited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='does joint have limits',
                  array_extent=('njnt',),
@@ -1700,7 +1700,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='jnt_actfrclimited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='does joint have actuator force limits',
                  array_extent=('njnt',),
@@ -1708,7 +1708,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='jnt_actgravcomp',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='is gravcomp force applied via actuators',
                  array_extent=('njnt',),
@@ -2404,7 +2404,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='light_castshadow',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='does light cast shadows',
                  array_extent=('nlight',),
@@ -2436,7 +2436,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='light_active',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='is light on',
                  array_extent=('nlight',),
@@ -2612,7 +2612,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='flex_internal',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='internal flex collision enabled',
                  array_extent=('nflex',),
@@ -3044,7 +3044,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='flex_rigid',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='are all vertices in the same body',
                  array_extent=('nflex',),
@@ -3052,7 +3052,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='flexedge_rigid',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='are both edge vertices in same body',
                  array_extent=('nflexedge',),
@@ -3060,7 +3060,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='flex_centered',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='are all vertex coordinates (0,0,0)',
                  array_extent=('nflex',),
@@ -3068,7 +3068,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='flex_flatskin',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='render flex skin with flat shading',
                  array_extent=('nflex',),
@@ -3716,7 +3716,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='mat_texuniform',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='make texture cube uniform',
                  array_extent=('nmat',),
@@ -3908,7 +3908,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='eq_active0',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='initial enable/disable constraint state',
                  array_extent=('neq',),
@@ -4020,7 +4020,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='tendon_limited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='does tendon have length limits',
                  array_extent=('ntendon',),
@@ -4028,7 +4028,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='tendon_actfrclimited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='does tendon have actuator force limits',
                  array_extent=('ntendon',),
@@ -4324,7 +4324,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='actuator_ctrllimited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='is control limited',
                  array_extent=('nu',),
@@ -4332,7 +4332,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='actuator_forcelimited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='is force limited',
                  array_extent=('nu',),
@@ -4340,7 +4340,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='actuator_actlimited',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='is activation limited',
                  array_extent=('nu',),
@@ -4372,7 +4372,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='actuator_actearly',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='step activation before force',
                  array_extent=('nu',),
@@ -5570,22 +5570,22 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
              StructFieldDecl(
                  name='flg_energypos',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
                  doc='has mj_energyPos been called',
              ),
              StructFieldDecl(
                  name='flg_energyvel',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
                  doc='has mj_energyVel been called',
              ),
              StructFieldDecl(
                  name='flg_subtreevel',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
                  doc='has mj_subtreeVel been called',
              ),
              StructFieldDecl(
                  name='flg_rnepost',
-                 type=ValueType(name='mjtByte'),
+                 type=ValueType(name='mjtBool'),
                  doc='has mj_rnePostConstraint been called',
              ),
              StructFieldDecl(
@@ -5690,7 +5690,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='eq_active',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='enable/disable constraints',
                  array_extent=('neq',),
@@ -6098,7 +6098,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              StructFieldDecl(
                  name='bvh_active',
                  type=PointerType(
-                     inner_type=ValueType(name='mjtByte'),
+                     inner_type=ValueType(name='mjtBool'),
                  ),
                  doc='was bounding volume checked for collision',
                  array_extent=('nbvh',),
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index 677a457d..66d0cf0f 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -843,10 +843,10 @@ typedef struct  {
 int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
                   const mjtNum xpos1[3], const mjtNum xmat1[9],
                   const mjtNum xpos2[3], const mjtNum xmat2[9], mjtNum margin,
-                  mjtNum product[36], mjtNum offset[12], mjtByte* initialize) {
+                  mjtNum product[36], mjtNum offset[12], mjtBool* initialize) {
   // get infinite dimensions (planes only)
-  mjtByte inf1[3] = {aabb1[3] >= mjMAXVAL, aabb1[4] >= mjMAXVAL, aabb1[5] >= mjMAXVAL};
-  mjtByte inf2[3] = {aabb2[3] >= mjMAXVAL, aabb2[4] >= mjMAXVAL, aabb2[5] >= mjMAXVAL};
+  mjtBool inf1[3] = {aabb1[3] >= mjMAXVAL, aabb1[4] >= mjMAXVAL, aabb1[5] >= mjMAXVAL};
+  mjtBool inf2[3] = {aabb2[3] >= mjMAXVAL, aabb2[4] >= mjMAXVAL, aabb2[5] >= mjMAXVAL};
 
   // if a bounding box is infinite, there must be a collision
   if ((inf1[0] && inf1[1] && inf1[2]) || (inf2[0] && inf2[1] && inf2[2])) {
@@ -858,7 +858,7 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
   const mjtNum *xpos[2] = {xpos1, xpos2};
   mjtNum xcenter[2][3], normal[2][3][3];
   mjtNum proj[2], radius[2];
-  mjtByte infinite[2] = {inf1[0] || inf1[1] || inf1[2], inf2[0] || inf2[1] || inf2[2]};
+  mjtBool infinite[2] = {inf1[0] || inf1[1] || inf1[2], inf2[0] || inf2[1] || inf2[2]};
 
   // compute centers in local coordinates
   if (product == NULL) {
@@ -940,8 +940,8 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                     int merged, int startadr, int pairadr) {
   int nbody = m->nbody, nbvhstatic = m->nbvhstatic;
-  mjtByte isbody1 = (bf1 < nbody);
-  mjtByte isbody2 = (bf2 < nbody);
+  mjtBool isbody1 = (bf1 < nbody);
+  mjtBool isbody2 = (bf2 < nbody);
   int f1 = isbody1 ? -1 : bf1 - nbody;
   int f2 = isbody2 ? -1 : bf2 - nbody;
   int mark_active = m->vis.global.bvactive;
@@ -957,7 +957,7 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
   // used with rotated bounding boxes (when bodies are involved)
   mjtNum product[36];  // 2 bb x 2 bb x 3 axes (body) x 3 axes (world)
   mjtNum offset[12];   // 2 bb x 2 bb x 3 axes (world)
-  mjtByte initialize = 1;
+  mjtBool initialize = true;
 
   // bitmask filter for bodyflex pair
   if (!canCollide2(m, bf1, bf2)) {
@@ -997,8 +997,8 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
     nstack--;
     int node1 = stack[nstack].node1;
     int node2 = stack[nstack].node2;
-    mjtByte isleaf1 = (child1[2*node1] < 0) && (child1[2*node1+1] < 0);
-    mjtByte isleaf2 = (child2[2*node2] < 0) && (child2[2*node2+1] < 0);
+    mjtBool isleaf1 = (child1[2*node1] < 0) && (child1[2*node1+1] < 0);
+    mjtBool isleaf2 = (child2[2*node2] < 0) && (child2[2*node2+1] < 0);
     int nodeid1 = m->bvh_nodeid[bvhadr1 + node1];
     int nodeid2 = m->bvh_nodeid[bvhadr2 + node2];
 
@@ -1035,8 +1035,8 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
               mj_narrowphase(m, d, pair, 1, 0);
             }
             if (mark_active) {
-              d->bvh_active[node1 + bvhadr1] = 1;
-              d->bvh_active[node2 + bvhadr2] = 1;
+              d->bvh_active[node1 + bvhadr1] = true;
+              d->bvh_active[node2 + bvhadr2] = true;
             }
           }
         }
@@ -1074,8 +1074,8 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
               mj_collideGeomElem(m, d, nodeid1, f2, nodeid2);
             }
             if (mark_active) {
-              d->bvh_active[node1 + bvhadr1] = 1;
-              d->bvh_active[node2 + bvhadr2] = 1;
+              d->bvh_active[node1 + bvhadr1] = true;
+              d->bvh_active[node2 + bvhadr2] = true;
             }
           }
         }
@@ -1104,8 +1104,8 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
       if (isleaf1 && isleaf2) {
         mj_collideElems(m, d, f1, nodeid1, f2, nodeid2);
         if (mark_active) {
-          d->bvh_active[node1 + bvhadr1] = 1;
-          d->bvh_active[node2 + bvhadr2] = 1;
+          d->bvh_active[node1 + bvhadr1] = true;
+          d->bvh_active[node2 + bvhadr2] = true;
         }
         continue;
       }
@@ -1118,8 +1118,8 @@ void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
     }
 
     if (mark_active) {
-      d->bvh_active[node1 + bvhadr1] = 1;
-      d->bvh_active[node2 + bvhadr2] = 1;
+      d->bvh_active[node1 + bvhadr1] = true;
+      d->bvh_active[node2 + bvhadr2] = true;
     }
 
     // keep traversing the tree
diff --git a/src/engine/engine_collision_driver.h b/src/engine/engine_collision_driver.h
index 4734975e..e3692f8f 100644
--- a/src/engine/engine_collision_driver.h
+++ b/src/engine/engine_collision_driver.h
@@ -39,7 +39,7 @@ MJAPI void mj_collision(const mjModel* m, mjData* d);
 MJAPI int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
                         const mjtNum xpos1[3], const mjtNum xmat1[9],
                         const mjtNum xpos2[3], const mjtNum xmat2[9], mjtNum margin,
-                        mjtNum product[36], mjtNum offset[12], mjtByte* initialize);
+                        mjtNum product[36], mjtNum offset[12], mjtBool* initialize);
 
 // is element active (for collisions)
 MJAPI int mj_isElemActive(const mjModel* m, int f, int e);
diff --git a/src/engine/engine_collision_sdf.c b/src/engine/engine_collision_sdf.c
index 490523f7..54e29eb1 100644
--- a/src/engine/engine_collision_sdf.c
+++ b/src/engine/engine_collision_sdf.c
@@ -772,7 +772,7 @@ static int selectFPS(const mjtNum* candidate, const mjtNum* dist, int ncandidate
                      int* selected_indices, int max_select) {
   if (ncandidate <= 0) return 0;
 
-  mjtByte selected[mjMAXCONPAIR] = {0};
+  mjtBool selected[mjMAXCONPAIR] = {false};
   mjtNum min_dist2[mjMAXCONPAIR];
   for (int i = 0; i < ncandidate; i++) {
     min_dist2[i] = mjMAXVAL;
@@ -791,7 +791,7 @@ static int selectFPS(const mjtNum* candidate, const mjtNum* dist, int ncandidate
   // iteratively select contacts using FPS
   int nselected = 0;
   while (nselected < max_select && nselected < mjMAXCONPAIR && best >= 0) {
-    selected[best] = 1;
+    selected[best] = true;
     selected_indices[nselected] = best;
     nselected++;
 
@@ -882,7 +882,7 @@ typedef struct {
 } MeshSDFContext;
 
 // process a single mesh face inline during BVH traversal
-static void processOneFace(int faceid, mjtByte* bvh_active, int node,
+static void processOneFace(int faceid, mjtBool* bvh_active, int node,
                            MeshSDFContext* ctx) {
   mjtNum corners[9];
   const mjModel* m = ctx->m;
@@ -903,7 +903,7 @@ static void processOneFace(int faceid, mjtByte* bvh_active, int node,
   processSdfCorners(corners, m, ctx->d, ctx->sdf, ctx->nstartpts,
                     ctx->candidate, ctx->dist, ctx->ncandidate);
 
-  if (bvh_active && *(ctx->ncandidate) > 0) bvh_active[node] = 1;
+  if (bvh_active && *(ctx->ncandidate) > 0) bvh_active[node] = true;
 }
 
 // Callback type for leaf node processing during BVH traversal
@@ -920,7 +920,7 @@ typedef int (*BVHLeafCallback)(int leaf_id, int node, void* ctx);
 // callback: function called for each leaf node
 // ctx: user context passed to callback
 static void traverseBVH(const mjtNum* bvh, const int* nodeid, const int* child,
-                        mjtByte* bvh_active, const mjtNum* offset, const mjtNum* rotation,
+                        mjtBool* bvh_active, const mjtNum* offset, const mjtNum* rotation,
                         const mjModel* m, const mjData* d, const mjSDF* sdf,
                         BVHLeafCallback callback, void* ctx) {
   int stack[64];
@@ -934,7 +934,7 @@ static void traverseBVH(const mjtNum* bvh, const int* nodeid, const int* child,
     if (nodeid[node] != -1) {
       if (boxIntersect(bvh + 6*node, offset, rotation, m, sdf, d)) {
         int active = callback(nodeid[node], node, ctx);
-        if (bvh_active && active) bvh_active[node] = 1;
+        if (bvh_active && active) bvh_active[node] = true;
       }
       continue;
     }
@@ -944,7 +944,7 @@ static void traverseBVH(const mjtNum* bvh, const int* nodeid, const int* child,
       continue;
     }
 
-    if (bvh_active) bvh_active[node] = 1;
+    if (bvh_active) bvh_active[node] = true;
 
     // push children
     for (int i = 0; i < 2; i++) {
@@ -1026,7 +1026,7 @@ int mjc_MeshSDF(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
     const int* nodeid = m->bvh_nodeid + bvhadr;
     const mjtNum* bvh = m->bvh_aabb + 6*bvhadr;
     const int* child = m->bvh_child + 2*bvhadr;
-    mjtByte* bvh_active = m->vis.global.bvactive ? d->bvh_active + bvhadr : NULL;
+    mjtBool* bvh_active = m->vis.global.bvactive ? d->bvh_active + bvhadr : NULL;
 
     traverseBVH(bvh, nodeid, child, bvh_active, ctx.offset, ctx.rotation,
                 m, d, ctx.sdf, meshFaceCallback, &ctx);
@@ -1354,7 +1354,7 @@ int mjc_FlexSDF(const mjModel* m, const mjData* d, mjContact* con,
   const int nbvhstatic = m->nbvhstatic;
   const mjtNum* bvh = d->bvh_aabb_dyn + 6*(bvhadr - nbvhstatic);
   const int* child = m->bvh_child + 2*bvhadr;
-  mjtByte* bvh_active = m->vis.global.bvactive ? d->bvh_active + bvhadr : NULL;
+  mjtBool* bvh_active = m->vis.global.bvactive ? d->bvh_active + bvhadr : NULL;
 
   // set up context for flex element processing
   FlexSDFContext fctx;
diff --git a/src/engine/engine_derivative_fd.c b/src/engine/engine_derivative_fd.c
index 93fbaa76..98fcc281 100644
--- a/src/engine/engine_derivative_fd.c
+++ b/src/engine/engine_derivative_fd.c
@@ -292,7 +292,7 @@ void mjd_smooth_velFD(const mjModel* m, mjData* d, mjtNum eps) {
 //   single-letter shortcuts:
 //     inputs: q=qpos, v=qvel, a=act, u=ctrl
 //     outputs: y=next_state (concatenated next qpos, qvel, act), s=sensordata
-void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
+void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_centered,
                 mjtNum* DyDq, mjtNum* DyDv, mjtNum* DyDa, mjtNum* DyDu,
                 mjtNum* DsDq, mjtNum* DsDv, mjtNum* DsDa, mjtNum* DsDu) {
   if (m->nhistory) {
@@ -539,7 +539,7 @@ void mjd_stepFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
 //      B: (2*nv+na x nu)
 //      C: (nsensordata x 2*nv+na)
 //      D: (nsensordata x nu)
-void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
+void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_centered,
                       mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D) {
   if (m->opt.integrator == mjINT_RK4) {
     mjERROR("RK4 integrator is not supported");
@@ -605,7 +605,7 @@ void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_cente
 //   notes:
 //     optionally compute mass matrix Jacobian DmDq
 //     flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse
-void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
+void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_actuation,
                    mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
                    mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                    mjtNum *DmDq) {
diff --git a/src/engine/engine_derivative_fd.h b/src/engine/engine_derivative_fd.h
index 00ad6ccc..44375491 100644
--- a/src/engine/engine_derivative_fd.h
+++ b/src/engine/engine_derivative_fd.h
@@ -33,11 +33,11 @@ MJAPI void mjd_passive_velFD(const mjModel* m, mjData* d, mjtNum eps);
 MJAPI void mj_stepSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
 
 // finite differenced transition matrices (control theory notation)
-MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte centered,
+MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool centered,
                             mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
 
 // finite differenced Jacobian of  (force, sensors) = mj_inverse(state, acceleration)
-MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
+MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_actuation,
                          mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
                          mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
                          mjtNum *DmDq);
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index f129e8d8..921fbae4 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -297,7 +297,7 @@ static mjtNum dcmotorVoltage(mjtNum ctrl, mjtNum length, mjtNum velocity,
 
 
 // clamp vector to range
-static void clampVec(mjtNum* vec, const mjtNum* range, const mjtByte* limited, int n,
+static void clampVec(mjtNum* vec, const mjtNum* range, const mjtBool* limited, int n,
                       const int* index) {
   for (int i=0; i < n; i++) {
     int j = index ? index[i] : i;
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 54caa9b8..93df2352 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -766,6 +766,7 @@ void mj_printFormattedModel(const mjModel* m, const char* filename, const char*
                                   float:   float_format,                    \
                                   int:     INT_FORMAT,                      \
                                   mjtByte: INT_FORMAT,                      \
+                                  mjtBool: INT_FORMAT,                      \
                                   mjtSize: SIZE_FORMAT,                     \
                                   default: NULL);                           \
     if (format) {                                                           \
@@ -1270,6 +1271,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
         int : INT_FORMAT,                                                     \
         mjtSize : SIZE_FORMAT,                                                \
         mjtByte : INT_FORMAT,                                                 \
+        mjtBool : INT_FORMAT,                                                 \
         default : NULL);                                                      \
     if (format) {                                                             \
       fprintf(fp, "  ");                                                      \
diff --git a/src/engine/engine_ray.c b/src/engine/engine_ray.c
index ef4bb686..056d2f86 100644
--- a/src/engine/engine_ray.c
+++ b/src/engine/engine_ray.c
@@ -66,7 +66,7 @@ static mjtNum latitude(const mjtNum vec[3]) {
 
 // eliminate geom
 static int ray_eliminate(const mjModel* m, const mjData* d, int geomid,
-                         const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude) {
+                         const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude) {
   // body exclusion
   if (m->geom_bodyid[geomid] == bodyexclude) {
     return 1;
@@ -850,14 +850,14 @@ mjtNum mju_rayTree(const mjModel* m, const mjData* d, int id, const mjtNum pnt[3
       if (sol >= 0 && (x < 0 || sol < x)) {
         x = sol;
         if (normal) mju_copy3(normal, normal_local);
-        if (mark_active) d->bvh_active[node + bvhadr] = 1;
+        if (mark_active) d->bvh_active[node + bvhadr] = true;
       }
       continue;
     }
 
     // used for rendering
     if (mark_active) {
-      d->bvh_active[node + bvhadr] = 1;
+      d->bvh_active[node + bvhadr] = true;
     }
 
     // add children to the stack
@@ -1000,8 +1000,8 @@ mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3
 
 // intersect ray with flex, return nearest vertex id, compute normal if given
 mjtNum mj_rayFlex(const mjModel* m, const mjData* d, int flex_layer,
-                  mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
-                  mjtByte flg_skin, int flexid, const mjtNum pnt[3],
+                  mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
+                  mjtBool flg_skin, int flexid, const mjtNum pnt[3],
                   const mjtNum vec[3], int vertid[1], mjtNum normal[3]) {
   int dim = m->flex_dim[flexid];
 
@@ -1306,7 +1306,7 @@ static int point_in_box(const mjtNum aabb[6], const mjtNum xpos[3],
 //  return geomid and distance (x) to nearest surface, or -1 if no intersection
 //  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion
 mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-              const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+              const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
               int geomid[1], mjtNum normal[3]) {
   int ngeom = m->ngeom;
   mjtNum dist, newdist;
@@ -1353,7 +1353,7 @@ mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtN
 
 // Initializes spherical bounding angles (geom_ba) and flag vector for a given source
 void mju_multiRayPrepare(const mjModel* m, const mjData* d, const mjtNum pnt[3],
-                         const mjtNum ray_xmat[9], const mjtByte* geomgroup, mjtByte flg_static,
+                         const mjtNum ray_xmat[9], const mjtByte* geomgroup, mjtBool flg_static,
                          int bodyexclude, mjtNum cutoff, mjtNum* geom_ba, int* geom_eliminate) {
   if (ray_xmat) {
     mjERROR("ray_xmat is currently unused, should be NULL");
@@ -1542,7 +1542,7 @@ static mjtNum mju_singleRay(const mjModel* m, mjData* d, const mjtNum pnt[3], co
 
 // performs multiple ray intersections, compute normals if given
 void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
-                 const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                 const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                  int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff) {
   mj_markStack(d);
 
diff --git a/src/engine/engine_ray.h b/src/engine/engine_ray.h
index b4fae179..7fa5de07 100644
--- a/src/engine/engine_ray.h
+++ b/src/engine/engine_ray.h
@@ -25,14 +25,14 @@ extern "C" {
 
 MJAPI void mju_multiRayPrepare(const mjModel* m, const mjData* d,
                                const mjtNum pnt[3], const mjtNum ray_xmat[9],
-                               const mjtByte* geomgroup, mjtByte flg_static,
+                               const mjtByte* geomgroup, mjtBool flg_static,
                                int bodyexclude, mjtNum cutoff, mjtNum* geom_ba,
                                int* geom_eliminate);
 
 // intersect multiple rays emanating from a single source, compute normals if given
 //  similar semantics to mj_ray, but vec, normal and dist are arrays
 MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
-                       const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                       const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                        int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff);
 
 
@@ -40,7 +40,7 @@ MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const m
 //  return geomid and distance (x) to nearest surface, or -1 if no intersection
 //  geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion
 MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
-                    const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                    const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
                     int geomid[1], mjtNum normal[3]);
 
 // intersect ray with hfield, compute normal if given
@@ -62,8 +62,8 @@ MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum
 
 // intersect ray with flex, return nearest vertex id, compute normal if given
 MJAPI mjtNum mj_rayFlex(const mjModel* m, const mjData* d, int flex_layer,
-                        mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
-                        mjtByte flg_skin, int flexid, const mjtNum pnt[3],
+                        mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
+                        mjtBool flg_skin, int flexid, const mjtNum pnt[3],
                         const mjtNum vec[3], int vertid[1], mjtNum normal[3]);
 
 // intersect ray with skin, return nearest vertex id
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index eaca8ba8..b062ec94 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -220,7 +220,7 @@ void mj_getState(const mjModel* m, const mjData* d, mjtNum* state, int sig) {
     if (element & sig) {
       int size = mj_stateElemSize(m, element);
 
-      // special handling of eq_active (mjtByte)
+      // special handling of eq_active (mjtBool)
       if (element == mjSTATE_EQ_ACTIVE) {
         int neq = m->neq;
         for (int j=0; j < neq; j++) {
@@ -288,7 +288,7 @@ void mj_setState(const mjModel* m, mjData* d, const mjtNum* state, int sig) {
     if (element & sig) {
       int size = mj_stateElemSize(m, element);
 
-      // special handling of eq_active (mjtByte)
+      // special handling of eq_active (mjtBool)
       if (element == mjSTATE_EQ_ACTIVE) {
         int neq = m->neq;
         for (int j=0; j < neq; j++) {
@@ -324,7 +324,7 @@ void mj_copyState(const mjModel* m, const mjData* src, mjData* dst, int sig) {
     if (element & sig) {
       int size = mj_stateElemSize(m, element);
 
-      // special handling of eq_active (mjtByte)
+      // special handling of eq_active (mjtBool)
       if (element == mjSTATE_EQ_ACTIVE) {
         int neq = m->neq;
         for (int j=0; j < neq; j++) {
diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c
index 723e2625..eada6949 100644
--- a/src/engine/engine_util_misc.c
+++ b/src/engine/engine_util_misc.c
@@ -33,21 +33,21 @@
 //------------------------------ tendon wrapping ---------------------------------------------------
 
 // check for intersection of two 2D line segments
-static mjtByte is_intersect(const mjtNum* p1, const mjtNum* p2,
+static mjtBool is_intersect(const mjtNum* p1, const mjtNum* p2,
                             const mjtNum* p3, const mjtNum* p4) {
   mjtNum a, b;
 
   // compute determinant, check
   mjtNum det = (p4[1]-p3[1])*(p2[0]-p1[0]) - (p4[0]-p3[0])*(p2[1]-p1[1]);
   if (mju_abs(det) < mjMINVAL) {
-    return 0;
+    return false;
   }
 
   // compute intersection point on each line
   a = ((p4[0]-p3[0])*(p1[1]-p3[1]) - (p4[1]-p3[1])*(p1[0]-p3[0])) / det;
   b = ((p2[0]-p1[0])*(p1[1]-p3[1]) - (p2[1]-p1[1])*(p1[0]-p3[0])) / det;
 
-  return ((a >= 0 && a <= 1 && b >= 0 && b <= 1) ? 1 : 0);
+  return ((a >= 0 && a <= 1 && b >= 0 && b <= 1) ? true : false);
 }
 
 
diff --git a/src/engine/engine_util_solve.c b/src/engine/engine_util_solve.c
index da50bfc7..f275b488 100644
--- a/src/engine/engine_util_solve.c
+++ b/src/engine/engine_util_solve.c
@@ -664,7 +664,7 @@ int mju_bandDiag(int i, int ntotal, int nband, int ndense) {
 
 // convert band matrix to dense matrix
 void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
-                    mjtByte flg_sym) {
+                    mjtBool flg_sym) {
   int nsparse = ntotal-ndense;
 
   // clear all
@@ -717,7 +717,7 @@ void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int n
 
 // multiply band-diagonal matrix with vector
 void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
-                       int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym) {
+                       int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym) {
   int nsparse = ntotal-ndense;
 
   // handle multiple vectors
diff --git a/src/engine/engine_util_solve.h b/src/engine/engine_util_solve.h
index 1f0ef8c4..971dfc86 100644
--- a/src/engine/engine_util_solve.h
+++ b/src/engine/engine_util_solve.h
@@ -82,14 +82,14 @@ MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 
 // convert banded matrix to dense matrix, fill upper triangle if flg_sym>0
 MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
-                          mjtByte flg_sym);
+                          mjtBool flg_sym);
 
 // convert dense matrix to banded matrix
 MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
 
 // multiply band-diagonal matrix with vector, include upper triangle if flg_sym>0
 MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
-                             int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
+                             int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym);
 
 // address of diagonal element i in band-dense matrix representation
 MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
diff --git a/src/experimental/platform/ux/imgui_widgets.cc b/src/experimental/platform/ux/imgui_widgets.cc
index 2ee56a3d..a928a6cf 100644
--- a/src/experimental/platform/ux/imgui_widgets.cc
+++ b/src/experimental/platform/ux/imgui_widgets.cc
@@ -103,6 +103,14 @@ void ImGui_DataPtrTable::DataPtr(const char* label, const mjtByte* ptr,
   }
 }
 
+void ImGui_DataPtrTable::DataPtr(const char* label, const mjtBool* ptr,
+                                 int index, int n) {
+  for (int i = 0; i < n; ++i) {
+    MakeLabel(label, i, n);
+    ImGui::Text("%s", ptr[index + i] ? "true" : "false");
+  }
+}
+
 void ImGui_DataPtrTable::DataPtr(const char* label, const mjtSize* ptr,
                                  int index, int n) {
   Numeric(label, ptr, index, n);
diff --git a/src/experimental/platform/ux/imgui_widgets.h b/src/experimental/platform/ux/imgui_widgets.h
index 211b213d..7ea7d584 100644
--- a/src/experimental/platform/ux/imgui_widgets.h
+++ b/src/experimental/platform/ux/imgui_widgets.h
@@ -174,6 +174,7 @@ class ImGui_DataPtrTable {
   // dimensionality of n.
   void DataPtr(const char* label, const char* ptr, int index, int n);
   void DataPtr(const char* label, const mjtByte* ptr, int index, int n);
+  void DataPtr(const char* label, const mjtBool* ptr, int index, int n);
   void DataPtr(const char* label, const mjtSize* ptr, int index, int n);
   void DataPtr(const char* label, const int* ptr, int index, int n);
   void DataPtr(const char* label, const float* ptr, int index, int n);
@@ -299,7 +300,7 @@ class ImGui_SpecElementTable : public ImGui_DataPtrTable {
 
     ImGui::PushID(&val);
     ImGui::SetNextItemWidth(-1.0f);
-    if constexpr (std::is_same_v<T, mjtByte>) {
+    if constexpr (std::is_same_v<T, mjtByte> || std::is_same_v<T, mjtBool>) {
       if (read_only_) {
         ImGui::Text("%s", val ? "true" : "false");
       } else {
diff --git a/src/user/user_api.cc b/src/user/user_api.cc
index d423b93c..42e0d7a6 100644
--- a/src/user/user_api.cc
+++ b/src/user/user_api.cc
@@ -2024,13 +2024,13 @@ void mjs_setString(mjString* dest, const char* text) {
 
 
 // Set specific entry in destination string vector.
-mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text) {
+mjtBool mjs_setInStringVec(mjStringVec* dest, int i, const char* text) {
   if (dest->size() <= i) {
     mju_error("Requested index in mjs_setInStringVec is out of bounds");
-    return 0;
+    return false;
   }
   dest->at(i) = std::string(text);
-  return 1;
+  return true;
 }
 
 
diff --git a/src/user/user_api.h b/src/user/user_api.h
index cf1e26a6..1ab17a40 100644
--- a/src/user/user_api.h
+++ b/src/user/user_api.h
@@ -377,7 +377,7 @@ MJAPI void mjs_setString(mjString* dest, const char* text);
 MJAPI void mjs_setStringVec(mjStringVec* dest, const char* text);
 
 // Set entry in string vector.
-MJAPI mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
+MJAPI mjtBool mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
 
 // Append text entry to string vector.
 MJAPI void mjs_appendString(mjStringVec* dest, const char* text);
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index 02b2b2b9..b42b6282 100644
--- a/src/user/user_model.cc
+++ b/src/user/user_model.cc
@@ -2800,8 +2800,8 @@ void mjCModel::CopyTree(mjModel* m) {
       // set joint fields
       m->jnt_type[jid] = pj->type;
       m->jnt_group[jid] = pj->group;
-      m->jnt_limited[jid] = (mjtByte)pj->is_limited();
-      m->jnt_actfrclimited[jid] = (mjtByte)pj->is_actfrclimited();
+      m->jnt_limited[jid] = (mjtBool)pj->is_limited();
+      m->jnt_actfrclimited[jid] = (mjtBool)pj->is_actfrclimited();
       m->jnt_actgravcomp[jid] = pj->actgravcomp;
       m->jnt_qposadr[jid] = pj->qposadr_;
       m->jnt_dofadr[jid] = pj->dofadr_;
@@ -3016,8 +3016,8 @@ void mjCModel::CopyTree(mjModel* m) {
       m->light_targetbodyid[lid] = pl->targetbodyid;
       m->light_type[lid] = pl->type;
       m->light_texid[lid] = pl->texid;
-      m->light_castshadow[lid] = (mjtByte)pl->castshadow;
-      m->light_active[lid] = (mjtByte)pl->active;
+      m->light_castshadow[lid] = (mjtBool)pl->castshadow;
+      m->light_active[lid] = (mjtBool)pl->active;
       mjuu_copyvec(m->light_pos+3*lid, pl->pos, 3);
       mjuu_copyvec(m->light_dir+3*lid, pl->dir, 3);
       m->light_bulbradius[lid] = pl->bulbradius;
@@ -3822,8 +3822,8 @@ void mjCModel::CopyObjects(mjModel* m) {
     m->tendon_num[i] = (int)pte->path.size();
     m->tendon_matid[i] = pte->matid;
     m->tendon_group[i] = pte->group;
-    m->tendon_limited[i] = (mjtByte)pte->is_limited();
-    m->tendon_actfrclimited[i] = (mjtByte)pte->is_actfrclimited();
+    m->tendon_limited[i] = (mjtBool)pte->is_limited();
+    m->tendon_actfrclimited[i] = (mjtBool)pte->is_actfrclimited();
     m->tendon_width[i] = (mjtNum)pte->width;
     mjuu_copyvec(m->tendon_solref_lim+mjNREF*i, pte->solref_limit, mjNREF);
     mjuu_copyvec(m->tendon_solimp_lim+mjNIMP*i, pte->solimp_limit, mjNIMP);
@@ -3891,9 +3891,9 @@ void mjCModel::CopyObjects(mjModel* m) {
       m->actuator_historyadr[i] = -1;
     }
 
-    m->actuator_ctrllimited[i] = (mjtByte)pac->is_ctrllimited();
-    m->actuator_forcelimited[i] = (mjtByte)pac->is_forcelimited();
-    m->actuator_actlimited[i] = (mjtByte)pac->is_actlimited();
+    m->actuator_ctrllimited[i] = (mjtBool)pac->is_ctrllimited();
+    m->actuator_forcelimited[i] = (mjtBool)pac->is_forcelimited();
+    m->actuator_actlimited[i] = (mjtBool)pac->is_actlimited();
     m->actuator_actearly[i] = pac->actearly;
     m->actuator_cranklength[i] = (mjtNum)pac->cranklength;
     mjuu_copyvec(m->actuator_gear + 6*i, pac->gear, 6);
diff --git a/test/fixture.cc b/test/fixture.cc
index 6a6b1732..33c3d8ac 100644
--- a/test/fixture.cc
+++ b/test/fixture.cc
@@ -203,6 +203,10 @@ auto Compare(unsigned char val1, unsigned char val2) {
   return val1 != val2;
 }
 
+auto Compare(bool val1, bool val2) {
+  return val1 != val2;
+}
+
 auto Compare(mjtSize val1, mjtSize val2) {
   return val1 > val2 ? val1 - val2 : val2 - val1;
 }
diff --git a/test/user/user_objects_test.cc b/test/user/user_objects_test.cc
index 8d9f5839..51cd792e 100644
--- a/test/user/user_objects_test.cc
+++ b/test/user/user_objects_test.cc
@@ -2002,7 +2002,7 @@ TEST_F(LimitedTest, JointLimited) {
 
   // see `user/testdata/auto_limits.xml` for expected values
   for (int i = 0; i < model->njnt; i++) {
-    EXPECT_EQ(model->jnt_limited[i], (mjtByte)model->jnt_user[i])
+    EXPECT_EQ(model->jnt_limited[i], (mjtBool)model->jnt_user[i])
         << i << " " << (int)model->jnt_limited[i] << " "
         << (int)model->jnt_user[i];
   }
diff --git a/test/xml/xml_native_reader_test.cc b/test/xml/xml_native_reader_test.cc
index 5f6ecef7..574da7da 100644
--- a/test/xml/xml_native_reader_test.cc
+++ b/test/xml/xml_native_reader_test.cc
@@ -2481,7 +2481,7 @@ TEST_F(ActuatorTest, ReadsByte) {
   std::array<char, 1024> error;
   mjModel* model = LoadModelFromString(xml, error.data(), error.size());
   ASSERT_THAT(model, NotNull());
-  EXPECT_EQ(*(model->actuator_actlimited), (mjtByte)(1 & 0xFF));
+  EXPECT_EQ(*(model->actuator_actlimited), true);
   mj_deleteModel(model);
 }
 
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index adfc48bb..b7157df7 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -6575,28 +6575,28 @@ struct MjData {
   void set_nv_awake(int value) {
     ptr_->nv_awake = value;
   }
-  mjtByte flg_energypos() const {
+  mjtBool flg_energypos() const {
     return ptr_->flg_energypos;
   }
-  void set_flg_energypos(mjtByte value) {
+  void set_flg_energypos(mjtBool value) {
     ptr_->flg_energypos = value;
   }
-  mjtByte flg_energyvel() const {
+  mjtBool flg_energyvel() const {
     return ptr_->flg_energyvel;
   }
-  void set_flg_energyvel(mjtByte value) {
+  void set_flg_energyvel(mjtBool value) {
     ptr_->flg_energyvel = value;
   }
-  mjtByte flg_subtreevel() const {
+  mjtBool flg_subtreevel() const {
     return ptr_->flg_subtreevel;
   }
-  void set_flg_subtreevel(mjtByte value) {
+  void set_flg_subtreevel(mjtBool value) {
     ptr_->flg_subtreevel = value;
   }
-  mjtByte flg_rnepost() const {
+  mjtBool flg_rnepost() const {
     return ptr_->flg_rnepost;
   }
-  void set_flg_rnepost(mjtByte value) {
+  void set_flg_rnepost(mjtBool value) {
     ptr_->flg_rnepost = value;
   }
   mjtNum time() const {
@@ -8977,7 +8977,7 @@ void mj_mulM2_wrapper(const MjModel& m, const MjData& d, const val& res, const N
   mj_mulM2(m.get(), d.get(), res_.data(), vec_.data());
 }
 
-void mj_multiRay_wrapper(const MjModel& m, MjData& d, const NumberArray& pnt, const NumberArray& vec, const NumberArray& geomgroup, mjtByte flg_static, int bodyexclude, const val& geomid, const val& dist, const val& normal, int nray, mjtNum cutoff) {
+void mj_multiRay_wrapper(const MjModel& m, MjData& d, const NumberArray& pnt, const NumberArray& vec, const NumberArray& geomgroup, mjtBool flg_static, int bodyexclude, const val& geomid, const val& dist, const val& normal, int nray, mjtNum cutoff) {
   UNPACK_ARRAY(mjtNum, pnt);
   UNPACK_ARRAY(mjtNum, vec);
   UNPACK_NULLABLE_ARRAY(mjtByte, geomgroup);
@@ -9052,7 +9052,7 @@ void mj_projectConstraint_wrapper(const MjModel& m, MjData& d) {
   mj_projectConstraint(m.get(), d.get());
 }
 
-mjtNum mj_ray_wrapper(const MjModel& m, const MjData& d, const NumberArray& pnt, const NumberArray& vec, const NumberArray& geomgroup, mjtByte flg_static, int bodyexclude, const val& geomid, const val& normal) {
+mjtNum mj_ray_wrapper(const MjModel& m, const MjData& d, const NumberArray& pnt, const NumberArray& vec, const NumberArray& geomgroup, mjtBool flg_static, int bodyexclude, const val& geomid, const val& normal) {
   UNPACK_ARRAY(mjtNum, pnt);
   UNPACK_ARRAY(mjtNum, vec);
   UNPACK_NULLABLE_ARRAY(mjtByte, geomgroup);
@@ -9061,7 +9061,7 @@ mjtNum mj_ray_wrapper(const MjModel& m, const MjData& d, const NumberArray& pnt,
   return mj_ray(m.get(), d.get(), pnt_.data(), vec_.data(), geomgroup_.data(), flg_static, bodyexclude, geomid_.data(), normal_.data());
 }
 
-mjtNum mj_rayFlex_wrapper(const MjModel& m, const MjData& d, int flex_layer, mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face, mjtByte flg_skin, int flexid, const NumberArray& pnt, const NumberArray& vec, const val& vertid, const val& normal) {
+mjtNum mj_rayFlex_wrapper(const MjModel& m, const MjData& d, int flex_layer, mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face, mjtBool flg_skin, int flexid, const NumberArray& pnt, const NumberArray& vec, const val& vertid, const val& normal) {
   UNPACK_ARRAY(mjtNum, pnt);
   UNPACK_ARRAY(mjtNum, vec);
   UNPACK_NULLABLE_VALUE(int, vertid);
@@ -9199,7 +9199,7 @@ std::string mj_versionString_wrapper() {
   return std::string(mj_versionString());
 }
 
-void mjd_inverseFD_wrapper(const MjModel& m, MjData& d, mjtNum eps, mjtByte flg_actuation, const val& DfDq, const val& DfDv, const val& DfDa, const val& DsDq, const val& DsDv, const val& DsDa, const val& DmDq) {
+void mjd_inverseFD_wrapper(const MjModel& m, MjData& d, mjtNum eps, mjtBool flg_actuation, const val& DfDq, const val& DfDv, const val& DfDa, const val& DsDq, const val& DsDv, const val& DsDa, const val& DmDq) {
   UNPACK_NULLABLE_VALUE(mjtNum, DfDq);
   UNPACK_NULLABLE_VALUE(mjtNum, DfDv);
   UNPACK_NULLABLE_VALUE(mjtNum, DfDa);
@@ -9237,7 +9237,7 @@ void mjd_subQuat_wrapper(const NumberArray& qa, const NumberArray& qb, const val
   mjd_subQuat(qa_.data(), qb_.data(), Da_.data(), Db_.data());
 }
 
-void mjd_transitionFD_wrapper(const MjModel& m, MjData& d, mjtNum eps, mjtByte flg_centered, const val& A, const val& B, const val& C, const val& D) {
+void mjd_transitionFD_wrapper(const MjModel& m, MjData& d, mjtNum eps, mjtBool flg_centered, const val& A, const val& B, const val& C, const val& D) {
   UNPACK_NULLABLE_VALUE(mjtNum, A);
   UNPACK_NULLABLE_VALUE(mjtNum, B);
   UNPACK_NULLABLE_VALUE(mjtNum, C);
@@ -10161,7 +10161,7 @@ void mju_axisAngle2Quat_wrapper(const val& res, const NumberArray& axis, mjtNum
   mju_axisAngle2Quat(res_.data(), axis_.data(), angle);
 }
 
-void mju_band2Dense_wrapper(const val& res, const NumberArray& mat, int ntotal, int nband, int ndense, mjtByte flg_sym) {
+void mju_band2Dense_wrapper(const val& res, const NumberArray& mat, int ntotal, int nband, int ndense, mjtBool flg_sym) {
   UNPACK_VALUE(mjtNum, res);
   UNPACK_ARRAY(mjtNum, mat);
   CHECK_SIZE(mat, (ntotal - ndense) * nband + ndense * ntotal);
@@ -10169,7 +10169,7 @@ void mju_band2Dense_wrapper(const val& res, const NumberArray& mat, int ntotal,
   mju_band2Dense(res_.data(), mat_.data(), ntotal, nband, ndense, flg_sym);
 }
 
-void mju_bandMulMatVec_wrapper(const val& res, const NumberArray& mat, const NumberArray& vec, int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym) {
+void mju_bandMulMatVec_wrapper(const val& res, const NumberArray& mat, const NumberArray& vec, int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym) {
   UNPACK_VALUE(mjtNum, res);
   UNPACK_ARRAY(mjtNum, mat);
   UNPACK_ARRAY(mjtNum, vec);
diff --git a/wasm/tests/bindings_test.ts b/wasm/tests/bindings_test.ts
index a58fd087..94fd6d01 100644
--- a/wasm/tests/bindings_test.ts
+++ b/wasm/tests/bindings_test.ts
@@ -514,7 +514,7 @@ describe('MuJoCo WASM Bindings', () => {
 
   it('should compute finite-differenced transition matrices', () => {
     const eps = 1e-6;
-    const flg_centered = 0;
+    const flg_centered = false;
     const dim = 2 * model!.nv + model!.na;
     const A = new mujoco.DoubleBuffer(dim * dim);
     const B = new mujoco.DoubleBuffer(dim * model!.nu);
@@ -1290,7 +1290,7 @@ describe('MuJoCo WASM Bindings', () => {
   // Corresponds to bindings_test.py:test_inverse_fd_none
   it('should compute inverse dynamics derivatives with null outputs', () => {
     const eps = 1e-6;
-    const flg_centered = 0;
+    const flg_centered = false;
     expect(
         () => mujoco.mjd_inverseFD(
             model!, data!, eps, flg_centered, null, null, null, null, null,
@@ -1301,7 +1301,7 @@ describe('MuJoCo WASM Bindings', () => {
   // Corresponds to bindings_test.py:test_inverse_fd
   it('should compute inverse dynamics derivatives', () => {
     const eps = 1e-6;
-    const flg_centered = 0;
+    const flg_centered = false;
 
     const nv = model!.nv;
     const nsensordata = model!.nsensordata;
@@ -1427,7 +1427,7 @@ describe('MuJoCo WASM Bindings', () => {
         expect(band.GetView()[index]).toEqual(i + 1);
       }
       const dense2 = new mujoco.DoubleBuffer(nTotal * nTotal);
-      const flgSym = 1;
+      const flgSym = true;
       mujoco.mju_band2Dense(
           dense2, band.GetView(), nTotal, nBand, nDense, flgSym);
       expectArraysEqual(new Float64Array(dense), dense2.GetView());

From da0bf9af17f796128315e4764f40becaba33011a Mon Sep 17 00:00:00 2001
From: David Hozic <davidhozic@gmail.com>
Date: Sat, 21 Feb 2026 14:53:35 +0100
Subject: [PATCH 26/61] Fix incorrect type for Python view of MjModel/tex_adrr

---
 python/mujoco/indexer_xmacro.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h
index 10c52971..23264a99 100644
--- a/python/mujoco/indexer_xmacro.h
+++ b/python/mujoco/indexer_xmacro.h
@@ -273,7 +273,7 @@
   X( int,     tex_, height,   ntex,     1 ) \
   X( int,     tex_, width,    ntex,     1 ) \
   X( int,     tex_, nchannel, ntex,     1 ) \
-  X( int,     tex_, adr,      ntex,     1 ) \
+  X( mjtSize, tex_, adr,      ntex,     1 ) \
   X( mjtByte, tex_, data,     ntexdata, 1 )
 
 #define MJMODEL_TUPLE                         \

From 57e27f94f3cb274e882e7b169f2bc1dc40d171de Mon Sep 17 00:00:00 2001
From: David Hozic <davidhozic@gmail.com>
Date: Tue, 26 May 2026 14:50:11 +0200
Subject: [PATCH 27/61] Fix compilation error

---
 src/engine/engine_print.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 93df2352..776d5281 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -480,7 +480,7 @@ static bool validateFloatFormat(const char* float_format) {
   // flag characters. allow at most one of each flag
   const char flag_characters[] = "-+ #0";
   int flag_character_counts[sizeof(flag_characters)] = { 0 };
-  char* c;
+  const char* c;
   while (c = strchr(flag_characters, float_format[cur_idx]), c != NULL) {
     int flag_idx = (c - flag_characters)/sizeof(char);
     flag_character_counts[flag_idx]++;

From 998f9960f9ed3737f3af2fd2eeeec180dc2f63a4 Mon Sep 17 00:00:00 2001
From: David Hozic <davidhozic@gmail.com>
Date: Tue, 26 May 2026 14:51:48 +0200
Subject: [PATCH 28/61] Revert indexer_xmacro.h change

---
 python/mujoco/indexer_xmacro.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/mujoco/indexer_xmacro.h b/python/mujoco/indexer_xmacro.h
index 23264a99..10c52971 100644
--- a/python/mujoco/indexer_xmacro.h
+++ b/python/mujoco/indexer_xmacro.h
@@ -273,7 +273,7 @@
   X( int,     tex_, height,   ntex,     1 ) \
   X( int,     tex_, width,    ntex,     1 ) \
   X( int,     tex_, nchannel, ntex,     1 ) \
-  X( mjtSize, tex_, adr,      ntex,     1 ) \
+  X( int,     tex_, adr,      ntex,     1 ) \
   X( mjtByte, tex_, data,     ntexdata, 1 )
 
 #define MJMODEL_TUPLE                         \

From d63a89f896f9aa5e268358e5aa5d308f409b937a Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 26 May 2026 07:22:37 -0700
Subject: [PATCH 29/61] MJX documentation: update feature parity section

PiperOrigin-RevId: 921452813
Change-Id: Ic33c20688b9875e0d2be8563377a1e009da51fa6
---
 doc/mjx.rst | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/doc/mjx.rst b/doc/mjx.rst
index 0d7f0b2d..0920e153 100644
--- a/doc/mjx.rst
+++ b/doc/mjx.rst
@@ -473,7 +473,7 @@ The following table compares feature support between MJX-Warp and MJX-JAX compar
      - All
      - ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON``
    * - :ref:`Integrator <mjtIntegrator>`
-     - All except ``IMPLICIT``
+     - All except ``IMPLICITFAST`` midpoint integrator feature
      - ``EULER``, ``RK4``, ``IMPLICITFAST`` (``IMPLICITFAST`` not supported with :doc:`fluid drag <computation/fluid>`)
    * - :ref:`Cone <mjtCone>`
      - All
@@ -503,8 +503,8 @@ The following table compares feature support between MJX-Warp and MJX-JAX compar
      - Sparse and Dense
      - Sparse and Dense
    * - Jacobian format
-     - ``DENSE`` only
      - ``DENSE`` and ``SPARSE``
+     - ``DENSE`` only
    * - Lights
      - ✓
      - Positions and directions

From 583aa5ad4f86f8e06dbc06172d43961f1afcc24b Mon Sep 17 00:00:00 2001
From: Michael Moss <mmoss@google.com>
Date: Tue, 26 May 2026 09:29:02 -0700
Subject: [PATCH 30/61] Update dependencies ahead of the 3.9.0 release.

PiperOrigin-RevId: 921512490
Change-Id: Iac2e4e52dde5ca7aeda62f033fcc03aadcaeb904
---
 cmake/MujocoDependencies.cmake | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake
index a6a0a4ca..664fc313 100644
--- a/cmake/MujocoDependencies.cmake
+++ b/cmake/MujocoDependencies.cmake
@@ -39,7 +39,7 @@ set(MUJOCO_DEP_VERSION_qhull
     CACHE STRING "Version of `qhull` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_miniz
-    3.1.1
+    d10b03cc73475af673df40f06e5cefd1d5f940d9
     CACHE STRING "Version of `miniz` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_Eigen3
@@ -58,7 +58,7 @@ set(MUJOCO_DEP_VERSION_gtest
 )
 
 set(MUJOCO_DEP_VERSION_benchmark
-    5c55f5d4f45a1b09c5d98aa63a671993ebd42c69
+    834a61fc65e8b7885fcf177f1230ae4b897118fa
     CACHE STRING "Version of `benchmark` to be fetched."
 )
 

From 6f9725a1daa14ff10afa5447f2190ad283f5036c Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Tue, 26 May 2026 15:05:54 -0700
Subject: [PATCH 31/61] Make mju_getXMLDependencies read XMLs through the
 resource provider.

The function called tinyxml2's LoadFile directly, which only works on
the OS file system. Reading through mju_openResource lets it work
against any registered backend (VFS, HTTP, github:, ...).
---
 src/xml/xml_util.cc        | 60 +++++++++++++++++++++---------------
 test/xml/xml_utils_test.cc | 62 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 24 deletions(-)

diff --git a/src/xml/xml_util.cc b/src/xml/xml_util.cc
index b224a372..e5446da3 100644
--- a/src/xml/xml_util.cc
+++ b/src/xml/xml_util.cc
@@ -222,34 +222,46 @@ void mjCopyError(char* dst, const char* src, int maxlen) {
 }
 
 void mju_getXMLDependencies(const char* filename, mjStringVec* dependencies) {
-  // load XML file or parse string
-  tinyxml2::XMLDocument doc;
-  doc.LoadFile(filename);
-
-  // error checking
-  if (doc.Error()) {
-    mju_error("Problem reading XML file '%s': %s", filename, doc.ErrorStr());
+  // Open through the resource provider so the function works against any
+  // registered backend (OS file system, VFS, HTTP, "github:", ...) rather
+  // than only the OS file system.
+  mjResource* resource = mju_openResource("", filename, nullptr, nullptr, 0);
+  if (resource == nullptr) {
+    mju_error("Could not open '%s'", filename);
   }
 
-  // get top-level element
-  tinyxml2::XMLElement *root = doc.RootElement();
-  if (!root) {
-    mju_error("XML root element not found");
-  }
-  std::unordered_set<std::string> files = {filename};
-
-  std::optional<FilePath> model_dir = std::nullopt;
-  mjResource *resource = mju_openResource("", filename, nullptr,
-                                          nullptr, 0);
-  if (resource != nullptr) {
-    const char* dir;
-    int ndir;
-    mju_getResourceDir(resource, &dir, &ndir);
-    model_dir = FilePath(std::string(dir, ndir));
+  // Read the XML bytes from the resource.
+  const void* buffer = nullptr;
+  int size = mju_readResource(resource, &buffer);
+  if (size < 0 || !size) {
     mju_closeResource(resource);
+    mju_error("Could not read '%s'", filename);
   }
-  // Get file references from include and model tags.
-  AccumulateFiles(files, root, model_dir.value());
+
+  // Capture the model directory while the resource is still open.
+  const char* dir = nullptr;
+  int ndir = 0;
+  mju_getResourceDir(resource, &dir, &ndir);
+  FilePath model_dir(std::string(dir, ndir));
+
+  // Parse from buffer and close (the parsed DOM is independent of the
+  // resource buffer once Parse returns).
+  tinyxml2::XMLDocument doc;
+  tinyxml2::XMLError err =
+      doc.Parse(static_cast<const char*>(buffer), static_cast<size_t>(size));
+  mju_closeResource(resource);
+
+  if (err != tinyxml2::XML_SUCCESS) {
+    mju_error("Problem reading XML file '%s': %s", filename,
+              doc.ErrorStr() ? doc.ErrorStr() : "");
+  }
+  tinyxml2::XMLElement* root = doc.RootElement();
+  if (!root) {
+    mju_error("XML root element not found in '%s'", filename);
+  }
+
+  std::unordered_set<std::string> files = {filename};
+  AccumulateFiles(files, root, model_dir);
 
   *dependencies = {files.begin(), files.end()};
 }
diff --git a/test/xml/xml_utils_test.cc b/test/xml/xml_utils_test.cc
index c820753d..5d49053c 100644
--- a/test/xml/xml_utils_test.cc
+++ b/test/xml/xml_utils_test.cc
@@ -13,11 +13,14 @@
 // limitations under the License.
 
 
+#include <map>
 #include <set>
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <mujoco/mjplugin.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
 
@@ -42,5 +45,64 @@ TEST_F(MujocoTest, GetXMLDependenciesTest) {
                                                 kModelPaths.end()};
   EXPECT_EQ(dependency_set, expected_dependency_set);
 }
+
+// Custom resource provider that serves XML strings from an in-memory map,
+// used to exercise the non-filesystem code path.
+namespace memxml {
+static const std::map<std::string, std::string>* g_files = nullptr;
+
+int Open(mjResource* resource) {
+  return g_files && g_files->count(resource->name) ? 1 : 0;
+}
+int Read(mjResource* resource, const void** buffer) {
+  auto it = g_files->find(resource->name);
+  if (it == g_files->end()) return -1;
+  *buffer = it->second.data();
+  return static_cast<int>(it->second.size());
+}
+void Close(mjResource* resource) {}
+}  // namespace memxml
+
+// Verifies mju_getXMLDependencies works against a non-filesystem resource
+// provider (the case the WASM/HTTP build relies on).
+TEST_F(MujocoTest, GetXMLDependenciesViaResourceProvider) {
+  const std::map<std::string, std::string> files = {
+      {"memxml:/scene.xml",
+       "<mujoco>\n"
+       "  <include file=\"child.xml\"/>\n"
+       "  <compiler meshdir=\"meshes\"/>\n"
+       "  <asset>\n"
+       "    <mesh name=\"m\" file=\"m.obj\"/>\n"
+       "  </asset>\n"
+       "</mujoco>"},
+      {"memxml:/child.xml",
+       "<mujoco>\n"
+       "  <asset>\n"
+       "    <texture name=\"t\" type=\"2d\" file=\"t.png\"/>\n"
+       "  </asset>\n"
+       "</mujoco>"},
+  };
+  memxml::g_files = &files;
+
+  mjpResourceProvider provider = {
+      .prefix = "memxml",
+      .open = memxml::Open,
+      .read = memxml::Read,
+      .close = memxml::Close,
+  };
+  ASSERT_GT(mjp_registerResourceProvider(&provider), 0);
+
+  mjStringVec dependencies;
+  mju_getXMLDependencies("memxml:/scene.xml", &dependencies);
+  std::set<std::string> dep_set(dependencies.begin(), dependencies.end());
+
+  EXPECT_THAT(dep_set, testing::UnorderedElementsAre(
+      "memxml:/scene.xml",
+      "memxml:/child.xml",
+      "memxml:/meshes/m.obj",
+      "memxml:/t.png"));
+
+  memxml::g_files = nullptr;
+}
 }  // namespace
 }  // namespace mujoco

From 249e9c08b8c04fcc1285418ea0475167caf594f2 Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Tue, 26 May 2026 15:21:01 -0700
Subject: [PATCH 32/61] MuJoCo Live: pre-fetch model assets in parallel before
 compilation.

The prefetcher asks mju_getXMLDependencies for the transitive asset
list, fetches every URL with Promise.all, and primes the bytes into the
WASM-side FetchCache. When the compiler runs, every resource it opens
is already in memory. Cold-load on a typical Menagerie model drops from
~9 s to ~2 s.
---
 src/experimental/studio/emscripten.cc | 82 ++++++++++++++++--------
 src/experimental/studio/live.js       | 91 +++++++++++++++++++++++++++
 2 files changed, 146 insertions(+), 27 deletions(-)

diff --git a/src/experimental/studio/emscripten.cc b/src/experimental/studio/emscripten.cc
index 644c9ea6..b69fae6b 100644
--- a/src/experimental/studio/emscripten.cc
+++ b/src/experimental/studio/emscripten.cc
@@ -27,6 +27,7 @@
 #include <string_view>
 #include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/graphics_mode.h"
@@ -108,8 +109,9 @@ EM_ASYNC_JS(int, FetchUrl,
               }
             });
 
-// Cache for data fetched via HTTP/HTTPS. Stores the downloaded bytes keyed by
-// the resource name (URL) so that read() can return a pointer to the data.
+// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed)
+// bytes keyed by the resource URL so that the resource provider's read()
+// can return a pointer into stable storage.
 class FetchCache {
  public:
   static FetchCache& Instance() {
@@ -117,52 +119,74 @@ class FetchCache {
     return instance;
   }
 
-  // Fetches the URL and stores the result. Returns the size (>0) on success.
+  // Returns the size of the entry for `url`, fetching it over the network
+  // if necessary. Entries that have been pre-populated via Prime() are
+  // returned without touching the network. Returns 0 on fetch failure.
   int Fetch(const char* url) {
-    char* data = nullptr;
+    if (auto it = entries_.find(url); it != entries_.end()) {
+      return static_cast<int>(it->second.size());
+    }
+    char* buf = nullptr;
     std::int32_t size = 0;
-    if (!FetchUrl(url, &data, &size)) {
+    if (!FetchUrl(url, &buf, &size)) {
       return 0;
     }
-    entries_[url] = Entry{UniquePtrWasm<char[]>(data), size};
-    return size;
+    // FetchUrl hands us a malloc'd buffer; take ownership of the bytes via
+    // a std::string and release the original allocation.
+    std::string bytes(buf, size);
+    std::free(buf);
+    return static_cast<int>(entries_.emplace(url, std::move(bytes))
+                                .first->second.size());
   }
 
-  // Returns pointer and size for a previously fetched URL.
-  int Read(const char* url, const void** buffer) {
+  // Pre-populates the cache so a subsequent Fetch() of the same URL hits
+  // memory instead of the network. First writer wins.
+  void Prime(std::string url, std::string bytes) {
+    entries_.try_emplace(std::move(url), std::move(bytes));
+  }
+
+  // Hands back a pointer into the cached bytes for the given URL.
+  // Returns -1 if the URL has not been fetched or primed.
+  int Read(const char* url, const void** buffer) const {
     auto it = entries_.find(url);
-    if (it == entries_.end()) {
-      return -1;
-    }
-    *buffer = it->second.data.get();
-    return it->second.size;
+    if (it == entries_.end()) return -1;
+    *buffer = it->second.data();
+    return static_cast<int>(it->second.size());
   }
 
-  // Frees the data for a URL.
   void Close(const char* url) { entries_.erase(url); }
 
  private:
-  struct FreeDeleter {
-    void operator()(void* p) const { std::free(p); }
-  };
-  template <typename T>
-  using UniquePtrWasm = std::unique_ptr<T, FreeDeleter>;
-
-  struct Entry {
-    UniquePtrWasm<char[]> data;
-    int size;
-  };
-  std::unordered_map<std::string, Entry> entries_;
+  std::unordered_map<std::string, std::string> entries_;
 };
 
 // ---------------------------------------------------------------------------
 
-// Javascript-facing function to register an asset.
+// Javascript-facing function to register a build-time asset (font, IBL,
+// .filamat material) that is fetched via the static AssetRegistry rather
+// than over HTTP.
 void RegisterAsset(std::string filename, std::string contents) {
   AssetRegistry::Instance().RegisterAsset(std::move(filename),
                                           std::move(contents));
 }
 
+// Javascript-facing function to pre-populate the HTTP/HTTPS fetch cache.
+// Used by the JS-side parallel prefetcher to feed already-downloaded asset
+// bytes to MuJoCo's compiler before LoadUrl is invoked, sidestepping the
+// otherwise strictly-serial EM_ASYNC_JS fetch chain.
+void PrimeFetchCache(std::string url, std::string bytes) {
+  FetchCache::Instance().Prime(std::move(url), std::move(bytes));
+}
+
+// Javascript-facing wrapper around mju_getXMLDependencies. Returns the
+// list of resource URLs that the given root XML transitively references
+// (meshes, textures, includes, ...), per MuJoCo's compiler rules.
+std::vector<std::string> GetXMLDependencies(const std::string& root_url) {
+  mjStringVec deps;
+  mju_getXMLDependencies(root_url.c_str(), &deps);
+  return deps;
+}
+
 // Javascript-facing function to initialize the app.
 void Init(const std::string& title, bool dark_theme) {
   // Note: dimensions do not matter as window will be resized to fit canvas.
@@ -305,7 +329,11 @@ void Deinit() {
 }
 
 EMSCRIPTEN_BINDINGS(studio_bindings) {
+  // Matches the C-level `mjStringVec` typedef in mjspec.h.
+  emscripten::register_vector<std::string>("mjStringVec");
   emscripten::function("registerAsset", &RegisterAsset);
+  emscripten::function("primeFetchCache", &PrimeFetchCache);
+  emscripten::function("getXMLDependencies", &GetXMLDependencies);
   emscripten::function("init", &Init);
   emscripten::function("loadFile", &LoadFile);
   emscripten::function("loadUrl", &LoadUrl);
diff --git a/src/experimental/studio/live.js b/src/experimental/studio/live.js
index 8de5230f..13bd874f 100644
--- a/src/experimental/studio/live.js
+++ b/src/experimental/studio/live.js
@@ -21,6 +21,85 @@ function hideLoading() {
   loadingOverlay.style.display = 'none';
 }
 
+// ---------------------------------------------------------------------------
+// Parallel asset prefetcher.
+//
+// MuJoCo's XML compiler synchronously opens each referenced asset via
+// mjpResourceProvider. In WASM that becomes a chain of ASYNCIFY-suspended
+// fetch() calls, strictly serial, one RTT per asset. For Menagerie models
+// that's ~80 round trips and ~10 s of avoidable wall time on a cold cache.
+//
+// To sidestep this we let MuJoCo itself enumerate the dependencies (via
+// mju_getXMLDependencies, exposed to JS as Module.getXMLDependencies), then
+// fetch them all in parallel and feed the bytes into the WASM-side
+// FetchCache via Module.primeFetchCache. When Module.loadUrl subsequently
+// runs, every resource the compiler opens is already in memory.
+//
+// Discovery rules live in src/xml/xml_util.cc; the JS only has to know
+// the scheme rewrites used by the resource providers in emscripten.cc.
+// ---------------------------------------------------------------------------
+
+// MuJoCo Live custom URL schemes that the C++ resource providers rewrite
+// before fetching. JS must do the same so primed cache keys match what
+// the providers eventually request.
+function resolveScheme(url) {
+  if (url.startsWith('github:')) {
+    return 'https://raw.githubusercontent.com/' + url.slice('github:'.length);
+  }
+  return url;
+}
+
+async function prefetchModelAssets(rootUrl, onProgress) {
+  const primed = new Set();   // URLs we've already pushed into FetchCache
+  const stats = { files: 0, bytes: 0, errors: 0 };
+
+  // Fetch + prime, dedup against `primed`. Returns true on success.
+  async function fetchAndPrime(httpUrl) {
+    if (primed.has(httpUrl)) return true;
+    primed.add(httpUrl);
+    try {
+      const resp = await fetch(httpUrl);
+      if (!resp.ok) {
+        console.warn('[prefetch] HTTP', resp.status, httpUrl);
+        stats.errors++;
+        return false;
+      }
+      const bytes = new Uint8Array(await resp.arrayBuffer());
+      Module.primeFetchCache(httpUrl, bytes);
+      stats.files++;
+      stats.bytes += bytes.length;
+      onProgress?.(stats);
+      return true;
+    } catch (err) {
+      console.warn('[prefetch]', httpUrl, err);
+      stats.errors++;
+      return false;
+    }
+  }
+
+  // Step 1: fetch the root XML so mju_getXMLDependencies (which reads it
+  // through the resource provider) hits cache instead of the network.
+  // Also gives the loading screen an immediate progress tick.
+  const rootHttp = resolveScheme(rootUrl);
+  if (!await fetchAndPrime(rootHttp)) {
+    throw new Error(`could not fetch root model ${rootHttp}`);
+  }
+
+  // Step 2: ask MuJoCo for the transitive dependency list. The call is
+  // async because reading recursively included XMLs may suspend on
+  // ASYNCIFY fetches; those XMLs end up in FetchCache as a side effect.
+  // The result is an embind mjStringVec we copy out and release.
+  const depsVec = await Module.getXMLDependencies(rootUrl);
+  const deps = [];
+  for (let i = 0; i < depsVec.size(); ++i) deps.push(depsVec.get(i));
+  depsVec.delete();
+
+  // Step 3: fetch every dependency in parallel, deduped against URLs
+  // already primed in step 1 (and earlier deps within this batch).
+  await Promise.all(deps.map((depUrl) => fetchAndPrime(resolveScheme(depUrl))));
+  return stats;
+}
+
 var Module = {
   preRun: [],
   postRun: [],
@@ -111,13 +190,25 @@ var Module = {
             // loop until the load completes, otherwise renderFrame will
             // hit "Cannot have multiple async operations in flight".
             showLoading();
+            // Surface prefetch progress on the loading screen.
+            const dlEl = document.getElementById('loadingDownload');
+            const dlFileEl = document.getElementById('loadingDownloadFile');
+            if (dlEl) dlEl.style.display = 'block';
+            const onProgress = (s) => {
+              if (dlFileEl) {
+                dlFileEl.textContent =
+                    `${s.files} files · ${(s.bytes/1024/1024).toFixed(1)} MB`;
+              }
+            };
             requestAnimationFrame(() => {
               requestAnimationFrame(async () => {
                 try {
+                  await prefetchModelAssets(modelUrl, onProgress);
                   await Module.loadUrl(modelUrl);
                 } catch (error) {
                   console.error('Failed to load model from URL:', error);
                 } finally {
+                  if (dlEl) dlEl.style.display = 'none';
                   hideLoading();
                   requestAnimationFrame(Module.animate);
                 }

From c12dc23852996ba666b58b136ac61548202ed84a Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Tue, 26 May 2026 17:06:09 -0700
Subject: [PATCH 33/61] Add x_scale to mujoco.minimize.least_squares.

Per-parameter scaling via change of variables z = x / D. Supports
'jac' (adaptive D_i = 1/||J(:,i)|| per iteration, matches scipy's TRF),
explicit array, or a positive scalar. Default 1.0 is a no-op.
---
 python/mujoco/minimize.py      | 50 +++++++++++++++++++++++++++-------
 python/mujoco/minimize_test.py | 47 ++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 10 deletions(-)

diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py
index b06f3eb1..fb829398 100644
--- a/python/mujoco/minimize.py
+++ b/python/mujoco/minimize.py
@@ -156,6 +156,7 @@ def least_squares(
     output: Optional[TextIO] = None,
     iter_callback: Optional[Callable[[List[IterLog]], None]] = None,
     check_derivatives: bool = False,
+    x_scale: Union[float, np.ndarray, str] = 1.0,
 ) -> Tuple[np.ndarray, List[IterLog]]:
   """Nonlinear Least Squares minimization with box bounds.
 
@@ -178,6 +179,13 @@ def least_squares(
     output: Optional file or StringIO to which to print messages.
     iter_callback: Optional iteration callback, takes trace argument.
     check_derivatives: Compare user-defined Jacobian and norm against fin-diff.
+    x_scale: Per-parameter scaling. Setting ``x_scale=D`` solves the problem in
+      the change of variables ``z = x / D`` and un-scales the result. ``1.0``
+      (default) is no scaling. ``'jac'`` sets ``D_i = 1 / ||J(:, i)||`` at each
+      iteration (matches scipy's ``least_squares(method='trf', x_scale='jac')``).
+      An array of shape ``(n,)`` (or a positive scalar) is used as ``D``
+      directly. Note that ``mu``, ``mu_min`` and ``mu_max`` then act on the
+      scaled subproblem.
 
   Returns:
     x: best solution found
@@ -203,6 +211,22 @@ def least_squares(
   mu = np.float64(0.0)  # Optimistically start with no regularization.
   n_reduc = 0  # Number of sequential mu reductions.
 
+  # Resolve x_scale -> D of shape (n, 1). For 'jac', D is refreshed each iter.
+  adaptive_scale = isinstance(x_scale, str)
+  if adaptive_scale and x_scale != 'jac':
+    raise ValueError(
+        f"x_scale must be 'jac', a positive scalar, or array, got {x_scale!r}."
+    )
+  if adaptive_scale:
+    D = np.ones((n, 1))
+  else:
+    D = np.asarray(x_scale, dtype=np.float64)
+    if D.shape not in ((), (n,)):
+      raise ValueError(f'x_scale array must have shape ({n},), got {D.shape}.')
+    if not np.all(np.isfinite(D)) or np.any(D <= 0):
+      raise ValueError('x_scale must be positive and finite.')
+    D = (np.ones(n) * D).reshape(n, 1)
+
   # Initialize logging.
   trace = []
   n_res = 0
@@ -284,8 +308,13 @@ def least_squares(
     if i == 0 and check_derivatives and not isinstance(norm, Quadratic):
       check_norm(r, norm, eps, output)
 
-    # Get gradient, Gauss-Newton Hessian.
-    grad, hess = norm.grad_hess(r, jac)
+    # Refresh D for adaptive ('jac') scaling: column-norm preconditioner.
+    if adaptive_scale:
+      col_norms = np.linalg.norm(jac, axis=0)
+      D = (1.0 / np.maximum(col_norms, eps)).reshape(n, 1)
+
+    # Gradient/Hessian in scaled coords (jac * D.T scales columns).
+    grad, hess = norm.grad_hess(r, jac * D.T)
 
     # Get free (unclamped) gradient.
     if bounds is None:
@@ -304,9 +333,9 @@ def least_squares(
         print('Zero gradient norm: exact minimum found?', file=output)
       break
 
-    # Bounds relative to x
-    dlower = None if bounds is None else bounds[0] - x
-    dupper = None if bounds is None else bounds[1] - x
+    # Bounds relative to x, expressed in scaled coords (dz = dx / D).
+    dlower = None if bounds is None else (bounds[0] - x) / D
+    dupper = None if bounds is None else (bounds[1] - x) / D
 
     # Find reduction satisfying Armijo's rule.
     armijo = -1
@@ -329,8 +358,8 @@ def least_squares(
       if status != Status.MAX_ITER:
         break
 
-      # New candidate, residual.
-      xnew = x + dx
+      # New candidate (D * dx is the x-space step).
+      xnew = x + D * dx
       t_start = time.time()
       rnew = residual(xnew)
       t_res += time.time() - t_start
@@ -362,8 +391,9 @@ def least_squares(
     else:
       reduction_ratio = reduction / expected_reduction
 
-    # Iteration message.
-    dx_norm = np.linalg.norm(dx)
+    # Iteration message. Step printed in x-space (D * dx).
+    step = D * dx
+    dx_norm = np.linalg.norm(step)
     if verbose >= Verbosity.ITER.value:
       logmu = np.log10(mu) if mu > 0 else -np.inf
       message = (
@@ -377,7 +407,7 @@ def least_squares(
     log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)
     if verbose >= Verbosity.FULLITER.value:
       log = dataclasses.replace(
-          log, residual=r, jacobian=jac, grad=grad, step=dx
+          log, residual=r, jacobian=jac, grad=grad / D, step=step
       )
     trace.append(log)
     if iter_callback is not None:
diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py
index 1792643d..77694a3f 100644
--- a/python/mujoco/minimize_test.py
+++ b/python/mujoco/minimize_test.py
@@ -319,6 +319,53 @@ class MinimizeTest(absltest.TestCase):
     self.assertIn('User-provided norm gradient matches', out.getvalue())
     self.assertIn('User-provided norm Hessian matches', out.getvalue())
 
+  def test_x_scale_default_is_no_op(self) -> None:
+    """x_scale=1.0 produces an identical trace to the default unscaled run."""
+    def residual(x):
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
+
+    x0 = np.array((0.0, 0.0))
+    x_default, trace_default = minimize.least_squares(
+        x0, residual, verbose=minimize.Verbosity.SILENT)
+    x_one, trace_one = minimize.least_squares(
+        x0, residual, x_scale=1.0, verbose=minimize.Verbosity.SILENT)
+    np.testing.assert_array_equal(x_default, x_one)
+    self.assertEqual(len(trace_default), len(trace_one))
+
+  def test_x_scale_reaches_same_minimum(self) -> None:
+    """Both 'jac' and an explicit array reach the unscaled run's minimum."""
+    def residual(x):
+      return np.stack([1 - x[0, :], 10 * (x[1, :] - x[0, :] ** 2)])
+
+    x0 = np.array((0.0, 0.0))
+    x_unscaled, _ = minimize.least_squares(
+        x0, residual, verbose=minimize.Verbosity.SILENT)
+    x_jac, _ = minimize.least_squares(
+        x0, residual, x_scale='jac', verbose=minimize.Verbosity.SILENT)
+    x_array, _ = minimize.least_squares(
+        x0, residual, x_scale=np.array([10.0, 0.1]),
+        verbose=minimize.Verbosity.SILENT)
+    np.testing.assert_allclose(x_jac, x_unscaled, atol=1e-6)
+    np.testing.assert_allclose(x_array, x_unscaled, atol=1e-6)
+
+  def test_x_scale_validation(self) -> None:
+    """Invalid x_scale arguments raise ValueError."""
+    def residual(x):
+      return x
+
+    x0 = np.array((1.0, 1.0))
+    bad_values = [
+        'bogus',                          # unknown string
+        np.array([1.0, -1.0]),            # non-positive entry
+        np.array([np.inf, 1.0]),          # non-finite entry
+        np.array([1.0]),                  # wrong shape
+    ]
+    for bad in bad_values:
+      with self.assertRaises(ValueError):
+        minimize.least_squares(
+            x0, residual, x_scale=bad,
+            verbose=minimize.Verbosity.SILENT)
+
 
 if __name__ == '__main__':
   absltest.main()

From f6cd0234fd840b195088c54cc18b78de2f61fa2f Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 27 May 2026 02:00:00 -0700
Subject: [PATCH 34/61] Change tactile sensor normal component to report
 penetration depth.

The normal component of the tactile sensor now reports the maximum penetration depth at each taxel, instead of a derived normal force. The depth is negated so that positive values indicate penetration.

PiperOrigin-RevId: 921980899
Change-Id: Ide4103c0aff465e25a81cfdda650f6a0e2da9e0b
---
 doc/XMLreference.rst                            | 17 ++++++++---------
 doc/changelog.rst                               |  2 ++
 .../mjx/third_party/mujoco_warp/_src/sensor.py  |  8 ++------
 src/engine/engine_sensor.c                      | 10 ++--------
 test/engine/engine_sensor_test.cc               |  8 ++++----
 5 files changed, 18 insertions(+), 27 deletions(-)

diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst
index 85b889df..87a6d928 100644
--- a/doc/XMLreference.rst
+++ b/doc/XMLreference.rst
@@ -8514,15 +8514,14 @@ Extraction
    :width: 30%
    :target: https://github.com/google-deepmind/mujoco/blob/main/model/tactile/tactile.xml
 
-The tactile sensor returns the penetration pressure and the sliding velocities in the tangent frame at given points
-between the geom associated with the sensor and the SDF geoms in contact with it. We define the penetration pressure as
-a function of the penetration depth :math:`p(d) = \frac{d}{d_{max}-d}`, which is zero at the surface and goes to
-infinity as the maximum depth is reached. The sensor is associated with a geom and a mesh. It is activated by the
-contact between its associated geom and other geoms. The vertices of the mesh, when positioned in the geom frame, are
-the points at which sensor values are computed, so the dimension of the output is 3 times the number of vertices in the
-mesh. The mesh must have 3 normal vectors per vertex, which are used to compute the tangent frame. If the penetration
-depth is positive (no contact), then all values are 0 for the corresponding vertex. Only contacts with geoms of type SDF
-contribute to the sensor output. The sensor can be visualized by enabling the visualization of contact points.
+The tactile sensor returns the maximum penetration depth and the sliding velocities in the tangent frame at given points
+between the geom associated with the sensor and the SDF geoms in contact with it. The sensor is associated with a geom
+and a mesh. It is activated by the contact between its associated geom and other geoms. The vertices of the mesh, when
+positioned in the geom frame, are the points at which sensor values are computed, so the dimension of the output is 3
+times the number of vertices in the mesh. The mesh must have 3 normal vectors per vertex, which are used to compute the
+tangent frame. If the penetration depth is positive (no contact), then all values are 0 for the corresponding vertex.
+Only contacts with geoms of type SDF contribute to the sensor output. The sensor can be visualized by enabling the
+visualization of contact points.
 
 .. _sensor-tactile-geom:
 
diff --git a/doc/changelog.rst b/doc/changelog.rst
index c09fe3dd..dfbd187e 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -79,6 +79,8 @@ General
      `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>` and now includes all
      enum type definitions.
 
+   - The :ref:`tactile<sensor-tactile>` sensor now reports raw depth instead of an estimated pressure.
+
    - MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
      ``naconmax``.
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
index 0610612b..472fa617 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
@@ -2252,18 +2252,14 @@ def _sensor_tactile(
     )
     vel_rel = vel_sensor - vel_other
 
-    kMaxDepth = 0.05
-    pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
-    force = wp.mul(normal, pressure)
-
     forceT = wp.vec3(0.0, 0.0, 0.0)
-    forceT[0] = wp.dot(force, normal)
+    forceT[0] = -depth
     if has_frame:
       forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
       forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
 
     dim = sensor_dim[sensor_id] // 3
-    wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
 
diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c
index d6a2eb13..83db1d49 100644
--- a/src/engine/engine_sensor.c
+++ b/src/engine/engine_sensor.c
@@ -172,14 +172,8 @@ static void* tactile_taxel_batch(void* args) {
                           mesh_normal[normal_stride*j + 2]};
       mju_rotVecQuat(normal, normal, m->mesh_quat + 4 * mesh_id);
 
-      // get contact force
-      mjtNum force[3];
-      mjtNum kMaxDepth = 0.05;
-      mjtNum pressure = depth / mju_max(kMaxDepth - depth, mjMINVAL);
-      mju_scl3(force, normal, pressure);
-
-      // accumulate into forcesT (disjoint writes per taxel j)
-      t->forcesT[0*ncon + j] += mju_dot3(force, normal);
+      // take max penetration depth (SDF distance is negative; negate for positive output)
+      t->forcesT[0*ncon + j] = mju_max(t->forcesT[0*ncon + j], -depth);
 
       if (has_frame) {
         mjtNum tang1[3] = {mesh_normal[normal_stride*j + 3],
diff --git a/test/engine/engine_sensor_test.cc b/test/engine/engine_sensor_test.cc
index a12fdec0..a71c5061 100644
--- a/test/engine/engine_sensor_test.cc
+++ b/test/engine/engine_sensor_test.cc
@@ -1702,7 +1702,7 @@ TEST_F(SensorTest, TactileSkipTangents) {
   EXPECT_EQ(data->time, 0.0);
   EXPECT_GT(data->ncon, 0) << "No contacts generated";
 
-  // Tactile sensor layout: [normal_forces..., tang1_forces..., tang2_forces...]
+  // Tactile sensor layout: [depths..., tang1_vel..., tang2_vel...]
   int ntaxel = model->nsensordata / 3;
   ASSERT_EQ(model->nsensordata % 3, 0) << "Sensor dim should be divisible by 3";
 
@@ -1712,13 +1712,13 @@ TEST_F(SensorTest, TactileSkipTangents) {
         << "Tangent component at index " << i << " should be 0";
   }
 
-  // Normal force components: verify count, sign, and magnitude ~-0.8
+  // Penetration depth components: verify count, sign, and magnitude ~0.2
   int nonzero_count = 0;
   for (int i = 0; i < ntaxel; i++) {
     if (data->sensordata[i] != 0) {
       nonzero_count++;
-      EXPECT_NEAR(data->sensordata[i], -0.8, 0.1)
-          << "Normal force at taxel " << i;
+      EXPECT_NEAR(data->sensordata[i], 0.2, 0.1)
+          << "Penetration depth at taxel " << i;
     }
   }
   EXPECT_EQ(nonzero_count, 2) << "Expected 2 taxels in contact";

From 237c17e48539b6c90bf90d3161547cbdcbfaa1e0 Mon Sep 17 00:00:00 2001
From: Michael Moss <mmoss@google.com>
Date: Wed, 27 May 2026 02:46:48 -0700
Subject: [PATCH 35/61] Update changelog for the 3.9.0 release.

PiperOrigin-RevId: 922002472
Change-Id: I92489261c59cff10b29def784698868e67aa43da
---
 doc/changelog.rst | 127 +++++++++++++++++++++++-----------------------
 1 file changed, 64 insertions(+), 63 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index dfbd187e..a2243f40 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,93 +2,94 @@
 Changelog
 =========
 
-Upcoming version (not yet released)
------------------------------------
+Version 3.9.0 (May 27, 2026)
+----------------------------
 
 General
 ^^^^^^^
-- Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
-  dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
-- Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
-  constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
-  This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
-  :ref:`Exact diagonal <soExactDiag>` for details.
-- The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
-  release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
-  undesirable time dependence.
-- Flexes are now allowed to sleep, with the exception of completely passive (constraint-free) flexes.
-- Added compiler timing diagnostics via the new :ref:`mjtCTimer` enum and the :ref:`mjs_getTimer` C API. After
-  :ref:`mj_compile`, per-category timings (total, assets, mesh loading, convex hull, normals, inertia, BVH, octree,
-  textures) are available via ``mjs_getTimer(spec)``. The :ref:`compile<saCompile>` sample prints a detailed timing
-  breakdown when run without an output file.
-- Added :ref:`mjtBool` to represent boolean variables, replacing :ref:`mjtByte` across all boolean fields in
-  :ref:`mjModel`, :ref:`mjData`, and public C API function signatures.
+1. Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
+   dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
+2. Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
+   constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
+   This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
+   :ref:`Exact diagonal <soExactDiag>` for details.
+3. The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
+   release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
+   undesirable time dependence.
+4. Flexes are now allowed to sleep, with the exception of completely passive (constraint-free) flexes.
+5. Added compiler timing diagnostics via the new :ref:`mjtCTimer` enum and the :ref:`mjs_getTimer` C API. After
+   :ref:`mj_compile`, per-category timings (total, assets, mesh loading, convex hull, normals, inertia, BVH, octree,
+   textures) are available via ``mjs_getTimer(spec)``. The :ref:`compile<saCompile>` sample prints a detailed timing
+   breakdown when run without an output file.
+6. Added :ref:`mjtBool` to represent boolean variables, replacing :ref:`mjtByte` across all boolean fields in
+   :ref:`mjModel`, :ref:`mjData`, and public C API function signatures.
 
 .. admonition:: Breaking API changes
    :class: attention
 
-   - The semantics of the contact ``margin`` and ``gap`` parameters have been redesigned for conceptual clarity and
-     consistency with `Newton <https://github.com/newton-physics/newton>`__. See the new
-     :ref:`margin and gap<coMarginGap>` documentation section for details.
+   7. The semantics of the contact ``margin`` and ``gap`` parameters have been redesigned for conceptual clarity and
+      consistency with `Newton <https://github.com/newton-physics/newton>`__. See the new
+      :ref:`margin and gap<coMarginGap>` documentation section for details.
 
-     Previously, ``margin`` controlled the *detection threshold* (contacts exist when ``dist < margin``) and ``gap``
-     was subtracted from it to produce the *force threshold* (forces generated when ``dist < margin - gap``). This was
-     unintuitive: users expected ``margin`` to mean geometric inflation and ``gap`` to mean a spatial gap.
+      Previously, ``margin`` controlled the *detection threshold* (contacts exist when ``dist < margin``) and ``gap``
+      was subtracted from it to produce the *force threshold* (forces generated when ``dist < margin - gap``). This was
+      unintuitive: users expected ``margin`` to mean geometric inflation and ``gap`` to mean a spatial gap.
 
-     Under the new semantics, ``margin`` is the geometric inflation of the geom surface and ``gap`` is an additional
-     detection buffer beyond the inflated surface:
+      Under the new semantics, ``margin`` is the geometric inflation of the geom surface and ``gap`` is an additional
+      detection buffer beyond the inflated surface:
 
-     - **Detection**: contacts are created when ``dist < margin + gap``.
-     - **Force generation**: constraint forces are applied when ``dist < margin``.
-     - **Inactive contacts**: contacts with ``margin < dist ≤ margin + gap`` are included in ``mjData.contact`` but
-       generate no force (``efc_address = -1``). This is useful for :ref:`adhesion<actuator-adhesion>` actuators and
-       custom callbacks.
+      - **Detection**: contacts are created when ``dist < margin + gap``.
+      - **Force generation**: constraint forces are applied when ``dist < margin``.
+      - **Inactive contacts**: contacts with ``margin < dist ≤ margin + gap`` are included in ``mjData.contact`` but
+        generate no force (``efc_address = -1``). This is useful for :ref:`adhesion<actuator-adhesion>` actuators and
+        custom callbacks.
 
-     With the default values ``margin = 0``, ``gap = 0``, the behavior is unchanged.
+      With the default values ``margin = 0``, ``gap = 0``, the behavior is unchanged.
 
-     .. image:: images/modeling/margin_gap_light.svg
-        :width: 80%
-        :align: center
-        :class: only-light
+      .. image:: images/modeling/margin_gap_light.svg
+         :width: 80%
+         :align: center
+         :class: only-light
 
-     .. image:: images/modeling/margin_gap_dark.svg
-        :width: 80%
-        :align: center
-        :class: only-dark
+      .. image:: images/modeling/margin_gap_dark.svg
+         :width: 80%
+         :align: center
+         :class: only-dark
 
-     |
+      |
 
-     **Migration:** Models that use the default ``gap="0"`` (the vast majority) require no changes. For models with
-     ``gap > 0``, apply the following transformation to preserve identical behavior:
+      **Migration:** Models that use the default ``gap="0"`` (the vast majority) require no changes. For models with
+      ``gap > 0``, apply the following transformation to preserve identical behavior:
 
-     .. code-block::
+      .. code-block::
 
-        margin_new = margin_old - gap_old
-        gap_new    = gap_old
+         margin_new = margin_old - gap_old
+         gap_new    = gap_old
 
-     For example, a geom with the old attributes ``margin="0.1" gap="0.1"`` should be changed to
-     ``margin="0" gap="0.1"``.
+      For example, a geom with the old attributes ``margin="0.1" gap="0.1"`` should be changed to
+      ``margin="0" gap="0.1"``.
 
-     Negative ``margin`` values are now permitted (corresponding to ``gap > margin`` under the old semantics). The
-     constraint ``margin + gap >= 0`` should be maintained to ensure valid collision detection.
+      Negative ``margin`` values are now permitted (corresponding to ``gap > margin`` under the old semantics). The
+      constraint ``margin + gap >= 0`` should be maintained to ensure valid collision detection.
 
-   - The :ref:`mjfCollision` functions now populate the :ref:`mjPreContact` struct instead of the :ref:`mjContact`
-     struct. The :ref:`mjPreContact` only contains the necessary fields needed for the narrowphase collision detection.
+   8. The :ref:`mjfCollision` functions now populate the :ref:`mjPreContact` struct instead of the :ref:`mjContact`
+      struct. The :ref:`mjPreContact` only contains the necessary fields needed for the narrowphase collision detection.
 
-   - The header file ``mjtnum.h`` was renamed to
-     `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>` and now includes all
-     enum type definitions.
+   9. The header file ``mjtnum.h`` was renamed to
+      `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>` and now includes all
+      enum type definitions.
 
-   - The :ref:`tactile<sensor-tactile>` sensor now reports raw depth instead of an estimated pressure.
+   10. The :ref:`tactile<sensor-tactile>` sensor now reports raw depth instead of an estimated pressure.
 
-   - MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
-     ``naconmax``.
+   11. MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
+       ``naconmax``.
+
+   12. Maybe-breaking: Added `mjassert.h
+       <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__, a new header containing
+       compile-time assertions that verify the sizes of MuJoCo's public types for ABI stability. This is a first step
+       towards replacing ``int`` with strongly-typed enums in the public API. If these assertions fail on your compiler or
+       platform, please report the issue on GitHub.
 
-   - Maybe-breaking: Added `mjassert.h
-     <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__, a new header containing
-     compile-time assertions that verify the sizes of MuJoCo's public types for ABI stability. This is a first step
-     towards replacing ``int`` with strongly-typed enums in the public API. If these assertions fail on your compiler or
-     platform, please report the issue on GitHub.
 
 Version 3.8.1 (May 11, 2026)
 ----------------------------

From 79aeceeaa3a6ce8f10533a85a9e8e42e0536d69d Mon Sep 17 00:00:00 2001
From: Michael Moss <mmoss@google.com>
Date: Wed, 27 May 2026 03:49:54 -0700
Subject: [PATCH 36/61] Update MuJoCo version to 3.10.0 following the 3.9.0
 release

PiperOrigin-RevId: 922027849
Change-Id: I52dd7fb5eeb962c2670ab5dbbe7ce192ac9332c5
---
 CMakeLists.txt                                 | 2 +-
 dist/mujoco.rc                                 | 8 ++++----
 dist/simulate.rc                               | 8 ++++----
 doc/APIreference/APIglobals.rst                | 2 +-
 doc/unity.rst                                  | 4 ++--
 include/mujoco/mujoco.h                        | 2 +-
 mjx/pyproject.toml                             | 8 ++++----
 python/mujoco/CMakeLists.txt                   | 4 ++--
 python/mujoco/mjpython/Info.plist              | 8 ++++----
 python/pyproject.toml                          | 6 +++---
 sample/CMakeLists.txt                          | 2 +-
 simulate/CMakeLists.txt                        | 2 +-
 src/engine/engine_support.c                    | 4 ++--
 unity/Editor/Bindings/MujocoBinaryRetriever.cs | 4 ++--
 unity/Runtime/Bindings/MjBindings.cs           | 2 +-
 unity/package.json                             | 2 +-
 16 files changed, 34 insertions(+), 34 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index e01c6621..b50248f8 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -29,7 +29,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco
-  VERSION 3.9.0
+  VERSION 3.10.0
   DESCRIPTION "MuJoCo Physics Simulator"
   HOMEPAGE_URL "https://mujoco.org"
 )
diff --git a/dist/mujoco.rc b/dist/mujoco.rc
index b9f1eff7..f8ebe6a0 100644
--- a/dist/mujoco.rc
+++ b/dist/mujoco.rc
@@ -1,6 +1,6 @@
 1 VERSIONINFO
-FILEVERSION 3,9,0,0
-PRODUCTVERSION 3,9,0,0
+FILEVERSION 3,10,0,0
+PRODUCTVERSION 3,10,0,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -9,9 +9,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "3.9.0"
+      VALUE "ProductVersion", "3.10.0"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "3.9.0"
+      VALUE "FileVersion", "3.10.0"
       VALUE "InternalName", "mujoco.dll"
       VALUE "OriginalFilename", "mujoco.dll"
       VALUE "CompanyName", "Google DeepMind"
diff --git a/dist/simulate.rc b/dist/simulate.rc
index dd3f0778..94161040 100644
--- a/dist/simulate.rc
+++ b/dist/simulate.rc
@@ -1,8 +1,8 @@
 MUJOCO ICON "mujoco.ico"
 
 1 VERSIONINFO
-FILEVERSION 3,9,0,0
-PRODUCTVERSION 3,9,0,0
+FILEVERSION 3,10,0,0
+PRODUCTVERSION 3,10,0,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -11,9 +11,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "3.9.0"
+      VALUE "ProductVersion", "3.10.0"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "3.9.0"
+      VALUE "FileVersion", "3.10.0"
       VALUE "InternalName", "simulate.exe"
       VALUE "OriginalFilename", "simulate.exe"
       VALUE "CompanyName", "Google DeepMind"
diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst
index d823553a..80c758af 100644
--- a/doc/APIreference/APIglobals.rst
+++ b/doc/APIreference/APIglobals.rst
@@ -388,7 +388,7 @@ Defined in `mujoco.h <https://github.com/google-deepmind/mujoco/blob/main/includ
      - value
      - description
    * - ``mjVERSION_HEADER``
-     - 3009000
+     - 3010000
      - The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P"
        using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001.
        The API function :ref:`mj_version` returns a number with the same meaning
diff --git a/doc/unity.rst b/doc/unity.rst
index d2e02a57..76b505e6 100644
--- a/doc/unity.rst
+++ b/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.9.0.dylib`` (it can be
+``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.10.0.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.9.0/lib/libmujoco.so.3.9.0`` and rename it as ``libmujoco.so``.
+``~/.mujoco/mujoco-3.10.0/lib/libmujoco.so.3.10.0`` and rename it as ``libmujoco.so``.
 
 Windows
 _______
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 95c99041..ee3ab87d 100644
--- a/include/mujoco/mujoco.h
+++ b/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 3009000
+#define mjVERSION_HEADER 3010000
 
 // needed to define size_t, fabs and log10
 #include <stdlib.h>
diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index aa5f1573..43609582 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-mjx"
-version = "3.9.0"
+version = "3.10.0"
 authors = [
     {name = "Google DeepMind", email = "mujoco@deepmind.com"},
 ]
@@ -30,7 +30,7 @@ dependencies = [
     "etils[epath]",
     "jax",
     "jaxlib",
-    "mujoco>=3.9.0.dev0",
+    "mujoco>=3.10.0.dev0",
     "scipy",
     "trimesh",
 ]
@@ -50,9 +50,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.9.0"
+Documentation = "https://mujoco.readthedocs.io/en/3.10.0"
 Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
-Changelog = "https://mujoco.readthedocs.io/en/3.9.0/changelog.html"
+Changelog = "https://mujoco.readthedocs.io/en/3.10.0/changelog.html"
 
 [tool.isort]
 force_single_line = true
diff --git a/python/mujoco/CMakeLists.txt b/python/mujoco/CMakeLists.txt
index 8810b3b1..7c7b09ba 100644
--- a/python/mujoco/CMakeLists.txt
+++ b/python/mujoco/CMakeLists.txt
@@ -86,7 +86,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.9.0.dylib
+          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.10.0.dylib
       )
       target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
     endif()
@@ -94,7 +94,7 @@ if(NOT TARGET mujoco)
 
   if(NOT MUJOCO_FRAMEWORK)
     find_library(
-      MUJOCO_LIBRARY mujoco mujoco.3.9.0 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
+      MUJOCO_LIBRARY mujoco mujoco.3.10.0 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
     )
     find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
     message("MuJoCo is at ${MUJOCO_LIBRARY}")
diff --git a/python/mujoco/mjpython/Info.plist b/python/mujoco/mjpython/Info.plist
index 3bb2afb2..751d3e42 100644
--- a/python/mujoco/mjpython/Info.plist
+++ b/python/mujoco/mjpython/Info.plist
@@ -7,13 +7,13 @@
     <key>CFBundleIdentifier</key>
     <string>org.mujoco.mjpython</string>
     <key>CFBundleVersion</key>
-    <string>3.9.0</string>
+    <string>3.10.0</string>
     <key>CFBundleGetInfoString</key>
-    <string>3.9.0</string>
+    <string>3.10.0</string>
     <key>CFBundleLongVersionString</key>
-    <string>3.9.0</string>
+    <string>3.10.0</string>
     <key>CFBundleShortVersionString</key>
-    <string>3.9.0</string>
+    <string>3.10.0</string>
     <key>CFBundleExecutable</key>
     <string>mjpython</string>
     <key>CFBundleIconFile</key>
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 1e7e3632..e66c406f 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mujoco"
-version = "3.9.0"
+version = "3.10.0"
 authors = [
     {name = "Google DeepMind", email = "mujoco@deepmind.com"},
 ]
@@ -35,9 +35,9 @@ dynamic = ["readme", "scripts"]
 
 [project.urls]
 Homepage = "https://github.com/google-deepmind/mujoco"
-Documentation = "https://mujoco.readthedocs.io/en/3.9.0"
+Documentation = "https://mujoco.readthedocs.io/en/3.10.0"
 Repository = "https://github.com/google-deepmind/mujoco"
-Changelog = "https://mujoco.readthedocs.io/en/3.9.0/changelog.html"
+Changelog = "https://mujoco.readthedocs.io/en/3.10.0/changelog.html"
 
 [tool.setuptools]
 include-package-data = false
diff --git a/sample/CMakeLists.txt b/sample/CMakeLists.txt
index 88e9cb87..ee30c449 100644
--- a/sample/CMakeLists.txt
+++ b/sample/CMakeLists.txt
@@ -24,7 +24,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco_samples
-  VERSION 3.9.0
+  VERSION 3.10.0
   DESCRIPTION "MuJoCo samples binaries"
   HOMEPAGE_URL "https://mujoco.org"
 )
diff --git a/simulate/CMakeLists.txt b/simulate/CMakeLists.txt
index f12701bb..70f32e41 100644
--- a/simulate/CMakeLists.txt
+++ b/simulate/CMakeLists.txt
@@ -29,7 +29,7 @@ set(MUJOCO_DEP_VERSION_lodepng
 
 project(
   mujoco_simulate
-  VERSION 3.9.0
+  VERSION 3.10.0
   DESCRIPTION "MuJoCo simulate binaries"
   HOMEPAGE_URL "https://mujoco.org"
 )
diff --git a/src/engine/engine_support.c b/src/engine/engine_support.c
index b062ec94..73e62190 100644
--- a/src/engine/engine_support.c
+++ b/src/engine/engine_support.c
@@ -43,8 +43,8 @@
 
 //-------------------------- Constants -------------------------------------------------------------
 
- #define mjVERSION 3009000
-#define mjVERSIONSTRING "3.9.0"
+ #define mjVERSION 3010000
+#define mjVERSIONSTRING "3.10.0"
 
 // names of disable flags
 const char* mjDISABLESTRING[mjNDISABLE] = {
diff --git a/unity/Editor/Bindings/MujocoBinaryRetriever.cs b/unity/Editor/Bindings/MujocoBinaryRetriever.cs
index 1ab0f725..1b4da168 100644
--- a/unity/Editor/Bindings/MujocoBinaryRetriever.cs
+++ b/unity/Editor/Bindings/MujocoBinaryRetriever.cs
@@ -37,7 +37,7 @@ public class MujocoBinaryRetriever {
           if (AssetDatabase.LoadMainAssetAtPath(mujocoPath + "/mujoco.dylib") == null) {
             File.Copy(
                 "/Applications/MuJoCo.app/Contents/Frameworks" +
-                "/mujoco.framework/Versions/Current/libmujoco.3.9.0.dylib",
+                "/mujoco.framework/Versions/Current/libmujoco.3.10.0.dylib",
                 mujocoPath + "/mujoco.dylib");
             AssetDatabase.Refresh();
           }
@@ -45,7 +45,7 @@ public class MujocoBinaryRetriever {
           if (AssetDatabase.LoadMainAssetAtPath(mujocoPath + "/libmujoco.so") == null) {
             File.Copy(
                 Environment.GetFolderPath(Environment.SpecialFolder.UserProfile) +
-                "/.mujoco/mujoco-3.9.0/lib/libmujoco.so.3.9.0",
+                "/.mujoco/mujoco-3.10.0/lib/libmujoco.so.3.10.0",
                 mujocoPath + "/libmujoco.so");
             AssetDatabase.Refresh();
           }
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 667d7d33..37aa73f1 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -116,7 +116,7 @@ public const int mjMAXLINEPNT = 1001;
 public const int mjMAXPLANEGRID = 200;
 public const bool THIRD_PARTY_MUJOCO_MJXMACRO_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MUJOCO_H_ = true;
-public const int mjVERSION_HEADER = 3009000;
+public const int mjVERSION_HEADER = 3010000;
 
 
 // ------------------------------------Enums------------------------------------
diff --git a/unity/package.json b/unity/package.json
index d533ede4..9a4cb940 100644
--- a/unity/package.json
+++ b/unity/package.json
@@ -1,7 +1,7 @@
 {
   "name": "org.mujoco",
   "displayName": "MuJoCo",
-  "version": "3.9.0",
+  "version": "3.10.0",
   "description": "MuJoCo importer and runtime plug-in",
   "dependencies": {},
   "author": {

From 0607dc450cd409a7569a18fd6e757d9c584476ba Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 20 Feb 2026 19:56:31 +0530
Subject: [PATCH 37/61] security: fix integer overflow in safeAddToBufferSize
 on MSVC

The MSVC fallback path in safeAddToBufferSize() performed unchecked
arithmetic (type_size*nr*nc) on attacker-controlled values read from
.mjb binary model files. This could cause integer overflow, leading
to an undersized heap allocation followed by a heap buffer overflow
when data is copied into the buffer.

The fix adds manual overflow detection using SIZE_MAX/INTPTR_MAX
comparisons, matching the behavior of the existing __builtin_*_overflow
path used on GCC/Clang.

Also adds a regression test that crafts a binary model buffer with
overflow-inducing size fields and asserts safe rejection.
---
 src/engine/engine_io.c        | 29 +++++++++++++++++--
 test/engine/engine_io_test.cc | 53 +++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 6a1e10fe..8ec5a268 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -184,9 +184,32 @@ static mjtSize safeAddToBufferSize(intptr_t* offset, mjtSize* nbuffer,
   if (__builtin_add_overflow(*nbuffer, to_add, nbuffer)) return 0;
   if (__builtin_add_overflow(*offset, to_add, offset)) return 0;
 #else
-  // TODO: offer a safe implementation for MSVC or other compilers that don't have the builtins
-  *nbuffer += SKIP(*offset) + type_size*nr*nc;
-  *offset += SKIP(*offset) + type_size*nr*nc;
+  // safe overflow checks for MSVC and other compilers without __builtin_*_overflow
+  {
+    size_t product;
+    size_t to_add;
+    size_t skip = SKIP(*offset);
+
+    // nc * nr
+    if (nr > 0 && (size_t)nc > SIZE_MAX / (size_t)nr) return 0;
+    product = (size_t)nc * (size_t)nr;
+
+    // product * type_size
+    if (type_size > 0 && product > SIZE_MAX / type_size) return 0;
+    product *= type_size;
+
+    // product + SKIP(*offset)
+    if (product > SIZE_MAX - skip) return 0;
+    to_add = product + skip;
+
+    // *nbuffer + to_add
+    if ((size_t)*nbuffer > SIZE_MAX - to_add) return 0;
+    *nbuffer += to_add;
+
+    // *offset + to_add
+    if (*offset > 0 && to_add > (size_t)(INTPTR_MAX - *offset)) return 0;
+    *offset += to_add;
+  }
 #endif
 
   return 1;
diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index 0100ef4c..1e93810c 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -955,5 +955,58 @@ TEST_F(EngineIoTest, RedZoneAlignmentTest) {
 }
 #endif
 
+// Regression test: crafted binary model with overflow-inducing sizes must be
+// safely rejected (not cause heap overflow). This exercises the overflow checks
+// in safeAddToBufferSize on all compilers including MSVC.
+TEST_F(EngineIoTest, LoadModelBufferRejectsOverflowingSizes) {
+  // construct a minimal valid-looking .mjb header
+  int header[5];
+  header[0] = 20;              // ID
+  header[1] = sizeof(mjtNum);  // floating point size
+  // We need the correct nsize and nptr from the current build.
+  // Rather than hardcoding, we create and save a trivial model, then mutate
+  // the sizes to trigger overflow.
+
+  constexpr char xml[] = "<mujoco />";
+  std::array<char, 1024> error;
+  mjModel* model = LoadModelFromString(xml, error.data(), error.size());
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error.data();
+
+  // save model to a buffer
+  int bufsize = mj_sizeModel(model);
+  ASSERT_GT(bufsize, 0);
+  std::vector<char> buffer(bufsize);
+  mj_saveModel(model, nullptr, buffer.data(), bufsize);
+  mj_deleteModel(model);
+
+  // locate the size fields in the buffer (after 5-int header)
+  const int header_bytes = 5 * sizeof(int);
+  ASSERT_GT(bufsize, header_bytes + 77 * (int)sizeof(mjtSize));
+
+  // mutate a size field to an extremely large value that would overflow
+  // when multiplied by sizeof(type). ntexdata is a good candidate since it
+  // is a byte count field and gets multiplied by sizeof(mjtByte)==1, but other
+  // fields multiply by sizeof(int) or sizeof(mjtNum), making overflow easier.
+  mjtSize* sizes = reinterpret_cast<mjtSize*>(buffer.data() + header_bytes);
+
+  // save original value, set to overflow-inducing value
+  mjtSize original = sizes[49];  // ntexdata index (approximate)
+  sizes[49] = static_cast<mjtSize>(SIZE_MAX / 2);
+
+  // also need to update the nbuffer field (last size) to avoid the early
+  // nbuffer mismatch check — but the overflow should be caught earlier
+  // in safeAddToBufferSize/mj_makeModel before we reach that check.
+
+  // attempt to load — should return NULL, not crash
+  mjModel* bad_model = mj_loadModelBuffer(buffer.data(), bufsize);
+  EXPECT_THAT(bad_model, IsNull())
+      << "Expected mj_loadModelBuffer to reject overflow-inducing sizes";
+
+  // clean up if somehow it succeeded
+  if (bad_model) {
+    mj_deleteModel(bad_model);
+  }
+}
+
 }  // namespace
 }  // namespace mujoco

From c0b8aa9decba72e65cd97840951e23a9f5855915 Mon Sep 17 00:00:00 2001
From: Matija Kecman <matijak@google.com>
Date: Wed, 27 May 2026 06:42:17 -0700
Subject: [PATCH 38/61] Make studio python more like studio C++

  - C++ Platform UX: Added a "Copy Camera" button to copy XML camera definitions to clipboard.
  - C++ Platform UX: Extracted `SetSpeedIndex` and `GetExpectedLabelWidth` into reusable platform helpers.
  - C++ Platform UX: Fixed combo box widths and removed auto-hide tab bar flag from dockspace.
  - C++ Studio App: Deduplicated label width and speed index logic by adopting platform helpers.
  - C++ Studio App: Updated the toolbar layout to use a 2-column table.
  - Python Bindings: Added pybind definitions for `set_camera_index`, `set_speed_index`, and `camera_to_string`.
  - Python Studio App: Aligned toolbar UI with C++ using a 2-column table and fixed immediate theme switching.
  - Python Studio App: Added a Help menu displaying the MuJoCo version and Stats toggle.
  - Python Events: Added `-`/`=` shortcuts for simulation speed and `Esc`/`[`/`]` shortcuts for camera selection.

PiperOrigin-RevId: 922097319
Change-Id: Iec74afe37bf5c7a1ebf1cab4141e47fff5fe9274
---
 src/experimental/platform/ux/gui.cc           | 33 ++++++++++--
 src/experimental/platform/ux/gui.h            |  4 ++
 src/experimental/platform/ux/imgui_widgets.cc | 17 ++++++
 src/experimental/platform/ux/imgui_widgets.h  |  3 ++
 src/experimental/studio/app.cc                | 54 +++----------------
 src/experimental/studio/app.h                 |  2 -
 6 files changed, 59 insertions(+), 54 deletions(-)

diff --git a/src/experimental/platform/ux/gui.cc b/src/experimental/platform/ux/gui.cc
index 2892b185..9a155613 100644
--- a/src/experimental/platform/ux/gui.cc
+++ b/src/experimental/platform/ux/gui.cc
@@ -316,8 +316,7 @@ ImVec4 ConfigureDockingLayout() {
     ImGui::Begin("Dockspace", nullptr, kWorkspaceFlags);
     const ImGuiDockNodeFlags kDockSpaceFlags =
         ImGuiDockNodeFlags_PassthruCentralNode |
-        ImGuiDockNodeFlags_NoDockingOverCentralNode |
-        ImGuiDockNodeFlags_AutoHideTabBar;
+        ImGuiDockNodeFlags_NoDockingOverCentralNode;
     ImGui::DockSpace(root, ImVec2(0.0f, 0.0f), kDockSpaceFlags);
     ImGui::End();
   }
@@ -426,9 +425,7 @@ void StepControlGui(const mjModel* model, StepControl* step_control,
                         ImGuiComboFlags_NoArrowButton)) {
     for (int n = 0; n < kPercentRealTime.size(); n++) {
       if (ImGui::Selectable(kPercentRealTime[n], (speed_index == n))) {
-        speed_index = std::clamp<int>(n, 0, kPercentRealTime.size() - 1);
-        float speed = std::stof(kPercentRealTime[speed_index]);
-        step_control->SetSpeed(speed);
+        SetSpeedIndex(step_control, speed_index, n);
       }
     }
     ImGui::EndCombo();
@@ -442,6 +439,17 @@ void StepControlGui(const mjModel* model, StepControl* step_control,
   }
 }
 
+void SetSpeedIndex(StepControl* step_control, int& speed_index,
+                   int request_idx) {
+  if (!step_control || request_idx == speed_index || kPercentRealTime.empty()) {
+    return;
+  }
+
+  speed_index = std::clamp<int>(request_idx, 0, kPercentRealTime.size() - 1);
+  float speed = std::stof(kPercentRealTime[speed_index]);
+  step_control->SetSpeed(speed);
+}
+
 bool ThemeSelectGui(GuiTheme* theme, const ImVec2& size) {
   static constexpr const char* ICON_DARKMODE = ICON_FA_CIRCLE;
   static constexpr const char* ICON_LIGHTMODE = ICON_FA_CIRCLE_O;
@@ -471,6 +479,7 @@ bool LabelSelectionGui(mjvOption* opts) {
   const std::string label_preview =
       opts->label == 0 ? std::string(ICON_LABEL) + " Label"
                        : std::string(ICON_LABEL) + " " + kLabelNames[opts->label];
+  ImGui::SetNextItemWidth(GetExpectedLabelWidth());
   if (ImGui::BeginCombo("##Label", label_preview.c_str(),
                         ImGuiComboFlags_NoArrowButton)) {
     for (int n = 0; n < IM_ARRAYSIZE(kLabelNames); n++) {
@@ -494,6 +503,7 @@ bool FrameSelectionGui(mjvOption* opts) {
   const std::string frame_preview =
       opts->frame == 0 ? std::string(ICON_FRAME) + " Frame"
                        : std::string(ICON_FRAME) + " " + kFrameNames[opts->frame];
+  ImGui::SetNextItemWidth(GetExpectedLabelWidth());
   if (ImGui::BeginCombo("##Frame", frame_preview.c_str(),
                         ImGuiComboFlags_NoArrowButton)) {
     for (int n = 0; n < IM_ARRAYSIZE(kFrameNames); n++) {
@@ -530,6 +540,19 @@ static std::string GetCameraName(const mjModel* model, const mjvCamera& camera,
 bool CameraSelectionGui(const mjModel* model, mjData* data, mjvCamera& camera,
                         int& index) {
   static constexpr const char* ICON_CAMERA = ICON_FA_CAMERA;
+  static constexpr const char* ICON_COPY_CAMERA = ICON_FA_COPY;
+
+  // Copy camera button.
+  const float btn_size = ImGui::GetFrameHeight();
+  const ImVec2 square_size(btn_size, btn_size);
+  if (ImGui::Button(ICON_COPY_CAMERA, square_size)) {
+    std::string camera_string = CameraToString(data, &camera);
+    MaybeSaveToClipboard(camera_string);
+  }
+  ImGui::SetItemTooltip("%s", "Copy Camera");
+  ImGui::SameLine(0, 0);
+
+  ImGui::SetNextItemWidth(GetExpectedLabelWidth());
 
   auto select = [&](int type, int idx) {
     if (ImGui::Selectable(GetCameraName(model, camera, type).c_str(),
diff --git a/src/experimental/platform/ux/gui.h b/src/experimental/platform/ux/gui.h
index 3599a93d..3edd0425 100644
--- a/src/experimental/platform/ux/gui.h
+++ b/src/experimental/platform/ux/gui.h
@@ -84,6 +84,10 @@ static constexpr std::array<const char*, 31> kPercentRealTime = {
 void StepControlGui(const mjModel* model, StepControl* step_control,
                     int& speed_index);
 
+// Sets the simulation speed index and updates the StepControl object.
+void SetSpeedIndex(StepControl* step_control, int& speed_index,
+                   int request_idx);
+
 // UX for selecting the GUI theme.
 bool ThemeSelectGui(GuiTheme* theme, const ImVec2& size = ImVec2(0, 0));
 
diff --git a/src/experimental/platform/ux/imgui_widgets.cc b/src/experimental/platform/ux/imgui_widgets.cc
index a928a6cf..c21e3a99 100644
--- a/src/experimental/platform/ux/imgui_widgets.cc
+++ b/src/experimental/platform/ux/imgui_widgets.cc
@@ -386,6 +386,23 @@ void MaybeSaveToClipboard(const std::string& contents) {
   ImGui::SetClipboardText(contents.c_str());
 }
 
+float GetExpectedLabelWidth() {
+  static float expected_label_width = 0;
+  if (expected_label_width == 0) {
+    int longest = 0;
+    const char* longest_label = "";
+    for (int i = 0; i < mjNVISFLAG; ++i) {
+      int length = static_cast<int>(std::strlen(mjVISSTRING[i][0]));
+      if (length > longest) {
+        longest_label = mjVISSTRING[i][0];
+        longest = length;
+      }
+    }
+    expected_label_width = ImGui::CalcTextSize(longest_label).x + 16;
+  }
+  return expected_label_width;
+}
+
 ImPlotFlags ImPlot_SetupPlotFlags(ImVec2 plot_size) {
   ImPlotFlags flags = ImPlotFlags_None;
   if (plot_size.x > 0 && plot_size.y > 0) {
diff --git a/src/experimental/platform/ux/imgui_widgets.h b/src/experimental/platform/ux/imgui_widgets.h
index 7ea7d584..53ea5bbc 100644
--- a/src/experimental/platform/ux/imgui_widgets.h
+++ b/src/experimental/platform/ux/imgui_widgets.h
@@ -617,6 +617,9 @@ inline void EndBoxSection() { ImGui::EndTable(); }
 // Saves the given contents to the clipboard if the clipboard is available.
 void MaybeSaveToClipboard(const std::string& contents);
 
+// Returns the expected width of a UI label based on MuJoCo visualization flags.
+float GetExpectedLabelWidth();
+
 // Returns plot flags with title/legend conditionally hidden when the plot
 // area is too small.  `plot_size` is the final rendered size of the plot.
 ImPlotFlags ImPlot_SetupPlotFlags(ImVec2 plot_size);
diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc
index b221fa61..d24df9a8 100644
--- a/src/experimental/studio/app.cc
+++ b/src/experimental/studio/app.cc
@@ -826,14 +826,7 @@ void App::SaveSettings() {
 }
 
 void App::SetSpeedIndex(int idx) {
-  if (idx == tmp_.speed_index || platform::kPercentRealTime.empty()) {
-    return;
-  }
-
-  tmp_.speed_index =
-      std::clamp<int>(idx, 0, platform::kPercentRealTime.size() - 1);
-  float speed = std::stof(platform::kPercentRealTime[tmp_.speed_index]);
-  step_control_.SetSpeed(speed);
+  platform::SetSpeedIndex(&step_control_, tmp_.speed_index, idx);
 }
 
 void App::MoveCamera(platform::CameraMotion motion, mjtNum reldx,
@@ -1022,7 +1015,7 @@ void App::BuildGui() {
 }
 
 void App::ModelOptionsGui() {
-  const float min_width = GetExpectedLabelWidth();
+  const float min_width = platform::GetExpectedLabelWidth();
   const ImGuiChildFlags child_flags =
       ImGuiChildFlags_AutoResizeY | ImGuiChildFlags_AlwaysAutoResize;
   const ImGuiTreeNodeFlags node_flags =
@@ -1064,7 +1057,7 @@ void App::DataInspectorGui() {
     return;
   }
 
-  const float min_width = GetExpectedLabelWidth();
+  const float min_width = platform::GetExpectedLabelWidth();
   const ImGuiChildFlags child_flags =
       ImGuiChildFlags_AutoResizeY | ImGuiChildFlags_AlwaysAutoResize;
   const ImGuiTreeNodeFlags node_flags =
@@ -1424,15 +1417,14 @@ void App::ToolBarGui() {
               ImVec2(ImGui::GetStyle().ItemSpacing.x * 2.0f,
                      ImGui::GetStyle().ItemSpacing.y));
 
-    const float label_width = GetExpectedLabelWidth();
-    const float copy_btn_width = ImGui::CalcTextSize(ICON_COPY_CAMERA).x +
-                                 ImGui::GetStyle().FramePadding.x * 2;
+    const float label_width = platform::GetExpectedLabelWidth();
+    const float copy_btn_width = ImGui::GetFrameHeight();
     const float theme_width =
         ImGui::CalcTextSize(platform::ICON_FA_CIRCLE_O).x +
         ImGui::GetStyle().FramePadding.x * 2;
     const float sp = ImGui::GetStyle().ItemSpacing.x;
-    const float right_width = label_width + sp + label_width + sp +
-                              label_width + sp + copy_btn_width + sp +
+    const float right_width = copy_btn_width + label_width + sp +
+                              label_width + sp + label_width + sp +
                               theme_width;
     const float separator_width = ImGui::GetFrameHeight() * .6f;
 
@@ -1467,26 +1459,15 @@ void App::ToolBarGui() {
 
     ImGui::TableNextColumn();
 
-    if (ImGui::Button(ICON_COPY_CAMERA, square_size)) {
-      std::string camera_string = platform::CameraToString(data(), &camera_);
-      platform::MaybeSaveToClipboard(camera_string);
-    }
-    ImGui::SetItemTooltip("%s", "Copy Camera");
-
-    ImGui::SameLine(0, 0);
-    ImGui::SetNextItemWidth(GetExpectedLabelWidth());
     platform::CameraSelectionGui(model(), data(), camera_, ui_.camera_idx);
 
     ImGui::SameLine();
-    ImGui::SetNextItemWidth(GetExpectedLabelWidth());
     platform::LabelSelectionGui(&vis_options_);
 
     ImGui::SameLine();
-    ImGui::SetNextItemWidth(GetExpectedLabelWidth());
     platform::FrameSelectionGui(&vis_options_);
 
     ImGui::SameLine();
-    ImGui::SetNextItemWidth(GetExpectedLabelWidth());
     if (platform::ThemeSelectGui(&ui_.theme, square_size)) {
       platform::SetupTheme(ui_.theme);
       ImGui::GetIO().WantSaveIniSettings = true;
@@ -1836,27 +1817,6 @@ void App::FileDialogGui() {
   }
 }
 
-float App::GetExpectedLabelWidth() {
-  // Find the longest label which we'll use to set the minimum toggle button
-  // width. This isn't perfect because we may have labels that are longer, but
-  // it's a good enough approximation.
-  if (tmp_.expected_label_width == 0) {
-    int longest = 0;
-    const char* longest_label = "";
-    for (int i = 0; i < mjNVISFLAG; ++i) {
-      int length = static_cast<int>(strlen(mjVISSTRING[i][0]));
-      if (length > longest) {
-        longest_label = mjVISSTRING[i][0];
-        longest = length;
-      }
-    }
-    // Pad the width a bit to account for how the labels will be displayed
-    // (e.g. as button labels or besides checkboxes).
-    tmp_.expected_label_width = ImGui::CalcTextSize(longest_label).x + 16;
-  }
-  return tmp_.expected_label_width;
-}
-
 App::UiState::Dict App::UiState::ToDict() const {
   return {
       {"theme", std::to_string(static_cast<int>(theme))},
diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h
index 175b3fee..9819c059 100644
--- a/src/experimental/studio/app.h
+++ b/src/experimental/studio/app.h
@@ -224,8 +224,6 @@ class App {
   void SpecExplorerGui();
   void SpecEditorGui();
 
-  float GetExpectedLabelWidth();
-
   mjSpec* spec() { return model_holder_->spec(); }
   mjModel* model() { return model_holder_->model(); }
   mjData* data() { return model_holder_->data(); }

From 36ea0c466f70d522b4a7a1910f8f5b91fdbd69cb Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Wed, 27 May 2026 06:54:23 -0700
Subject: [PATCH 39/61] Support multiscene, multiview rendering.

Update Renderable so that it can switch its own draw states (material
instance, shadow casting, reflection textures, etc.) per render request.
This allows us to use the same Renderable with different effects from
different views within the same frame.

PiperOrigin-RevId: 922102817
Change-Id: I1cf02fddb7a00345a741db1bee88e716517765f4
---
 .../filament/compat/scene_bridge.cc           |   4 +-
 .../filament/compat/scene_geom_util.cc        |  40 +---
 .../filament/compat/scene_geom_util.h         |   2 +-
 .../filament/filament/filament_context.cc     |  10 +
 .../filament/filament/material.cc             |  56 ++++-
 src/experimental/filament/filament/material.h |   6 +
 .../filament/filament/renderable.cc           | 224 +++++++++++++-----
 .../filament/filament/renderable.h            |  47 +++-
 .../filament/filament/scene_view.cc           | 118 +++------
 .../filament/filament/scene_view.h            |  23 +-
 src/experimental/filament/filament_util.h     |   6 +-
 .../filament/render_context_filament.h        |   4 -
 12 files changed, 321 insertions(+), 219 deletions(-)

diff --git a/src/experimental/filament/compat/scene_bridge.cc b/src/experimental/filament/compat/scene_bridge.cc
index 61a24d2f..b8869bbd 100644
--- a/src/experimental/filament/compat/scene_bridge.cc
+++ b/src/experimental/filament/compat/scene_bridge.cc
@@ -129,8 +129,8 @@ void SceneBridge::Update(const mjrRect& viewport, const mjvScene* scene) {
       model_objects_->CreateSkinFlexMesh(scene, *geom);
     }
 
-    UniquePtr<mjrRenderable> renderable = CreateGeomRenderable(
-        *geom, ctx_, model_objects_.get(), headpos, scene->flags);
+    UniquePtr<mjrRenderable> renderable =
+        CreateGeomRenderable(*geom, ctx_, model_objects_.get(), scene->flags);
 
     mjrf_addRenderableToScene(scene_.get(), renderable.get());
     renderables_.push_back(std::move(renderable));
diff --git a/src/experimental/filament/compat/scene_geom_util.cc b/src/experimental/filament/compat/scene_geom_util.cc
index 83622ff1..17ad8e6d 100644
--- a/src/experimental/filament/compat/scene_geom_util.cc
+++ b/src/experimental/filament/compat/scene_geom_util.cc
@@ -40,12 +40,6 @@ static float GetPlaneTileSize(const mjModel* model, int matid,
   }
 }
 
-static bool IsBehind(const float* headpos, const float* pos, const float* mat) {
-  return ((headpos[0] - pos[0]) * mat[2] + (headpos[1] - pos[1]) * mat[5] +
-              (headpos[2] - pos[2]) * mat[8] <
-          0.0f);
-}
-
 static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
                               ModelObjects* model_objs) {
   const mjModel* model = model_objs->GetModel();
@@ -69,19 +63,8 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
       break;
     case mjGEOM_PLANE: {
       mjrf_setRenderableGeomMesh(renderable, geom_type, nstack, nslice, nquad);
-
-      float size[3];
-      std::memcpy(size, &geom.size, 3 * sizeof(float));
-      const bool is_infinite = !(size[0] > 0 && size[1] > 0);
-      if (is_infinite) {
-        // Infinite planes are scaled to match the tile size used by
-        // re-centering in engine_vis_visualize.c.
-        const float plane_scale = static_cast<float>(mjMAXPLANEGRID) / 2.0f;
-        size[0] = plane_scale;
-        size[1] = plane_scale;
-      }
       // Planes only define an xy size, so set the z-dimension to 1.0f.
-      size[2] = 1.0f;
+      const float size[3] = {geom.size[0], geom.size[1], 1.0f};
       mjrf_setRenderableSize(renderable, size);
       break;
     }
@@ -160,7 +143,7 @@ static void PrepareGeomMeshes(mjrRenderable* renderable, const mjvGeom& geom,
 }
 
 static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
-                               ModelObjects* model_objs, const float headpos[3],
+                               ModelObjects* model_objs,
                                const mjtByte render_flags[mjNRNDFLAG]) {
   const mjModel* model = model_objs->GetModel();
 
@@ -175,21 +158,8 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
   material.color[1] = geom.rgba[1];
   material.color[2] = geom.rgba[2];
   material.color[3] = geom.rgba[3];
-  if (geom.type == mjGEOM_PLANE) {
-    if (IsBehind(headpos, geom.pos, geom.mat)) {
-      material.color[3] *= 0.3;
-      mjrf_setRenderableReceiveShadows(renderable, false);
-      material.reflective = false;
-    } else {
-      mjrf_setRenderableReceiveShadows(renderable, true);
-      material.reflective = geom.reflectance > 0 && material.color[3] == 1.0f;
-    }
-  }
+
   mjrf_setRenderableLayerMask(renderable, geom.category);
-  if (geom.category == mjCAT_DECOR) {
-    mjrf_setRenderableCastShadows(renderable, false);
-    mjrf_setRenderableReceiveShadows(renderable, false);
-  }
 
   if (geom.matid >= 0 && geom.matid < model->nmat) {
     auto get_texture = [&](int role) -> const mjrTexture* {
@@ -310,12 +280,12 @@ static void UpdateGeomMaterial(mjrRenderable* renderable, const mjvGeom& geom,
 
 UniquePtr<mjrRenderable> CreateGeomRenderable(
     const mjvGeom& geom, mjrfContext* ctx, ModelObjects* model_objs,
-    const float headpos[3], const mjtByte render_flags[mjNRNDFLAG]) {
+    const mjtByte render_flags[mjNRNDFLAG]) {
   mjrRenderableParams params;
   mjr_defaultRenderableParams(&params);
   auto renderable = CreateRenderable(ctx, params);
   PrepareGeomMeshes(renderable.get(), geom, model_objs);
-  UpdateGeomMaterial(renderable.get(), geom, model_objs, headpos, render_flags);
+  UpdateGeomMaterial(renderable.get(), geom, model_objs, render_flags);
   return renderable;
 }
 }  // namespace mujoco
diff --git a/src/experimental/filament/compat/scene_geom_util.h b/src/experimental/filament/compat/scene_geom_util.h
index 89e516a2..53911c79 100644
--- a/src/experimental/filament/compat/scene_geom_util.h
+++ b/src/experimental/filament/compat/scene_geom_util.h
@@ -26,7 +26,7 @@ namespace mujoco {
 // Creates a Renderable from the given mjvGeom.
 UniquePtr<mjrRenderable> CreateGeomRenderable(
     const mjvGeom& geom, mjrfContext* ctx, ModelObjects* model_objs,
-    const float headpos[3], const mjtByte render_flags[mjNRNDFLAG]);
+    const mjtByte render_flags[mjNRNDFLAG]);
 
 }  // namespace mujoco
 
diff --git a/src/experimental/filament/filament/filament_context.cc b/src/experimental/filament/filament/filament_context.cc
index e381d402..e5682486 100644
--- a/src/experimental/filament/filament/filament_context.cc
+++ b/src/experimental/filament/filament/filament_context.cc
@@ -18,7 +18,9 @@
 #include <cstdint>
 #include <memory>
 #include <span>
+#include <unordered_map>
 #include <utility>
+#include <vector>
 
 #include <filament/Engine.h>
 #include <filament/IndexBuffer.h>
@@ -94,6 +96,14 @@ mjrFrameHandle FilamentContext::Render(
 
   ValidateSwapChains(requests);
 
+  std::unordered_map<mjrScene*, std::vector<const mjrRenderRequest*>> scene_to_requests;
+  for (const mjrRenderRequest& request : requests) {
+    scene_to_requests[request.scene].push_back(&request);
+  }
+  for (auto& [scene, requests] : scene_to_requests) {
+    SceneView::downcast(scene)->PrepareToRender(requests);
+  }
+
   bool render_began = false;
   mjrRenderTarget* current_target = nullptr;
   for (const mjrRenderRequest& request : requests) {
diff --git a/src/experimental/filament/filament/material.cc b/src/experimental/filament/filament/material.cc
index 21a024db..8c8b5460 100644
--- a/src/experimental/filament/filament/material.cc
+++ b/src/experimental/filament/filament/material.cc
@@ -13,6 +13,9 @@
 // limitations under the License.
 
 #include "experimental/filament/filament/material.h"
+#include <cstddef>
+#include <cstdint>
+#include <functional>
 
 #include <filament/Color.h>
 #include <filament/Material.h>
@@ -29,6 +32,51 @@
 
 namespace mujoco {
 
+template <class T>
+static void Combine(uint64_t& seed, const T* v) {
+  seed ^= std::hash<T>()(*v) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+}
+
+template <typename T>
+static uint64_t hash(const T& obj) {
+  static_assert(std::is_trivially_copyable_v<T>,
+                "Only trivially copyable types are hashable.");
+  int num_bytes = sizeof(T);
+  const std::byte* ptr = reinterpret_cast<const std::byte*>(&obj);
+  uint64_t seed = 0;
+  while (num_bytes >= sizeof(uint64_t)) {
+    Combine(seed, reinterpret_cast<const uint64_t*>(ptr));
+    ptr += sizeof(uint64_t);
+    num_bytes -= sizeof(uint64_t);
+  }
+  while (num_bytes >= sizeof(uint32_t)) {
+    Combine(seed, reinterpret_cast<const uint32_t*>(ptr));
+    ptr += sizeof(uint32_t);
+    num_bytes -= sizeof(uint32_t);
+  }
+  while (num_bytes >= sizeof(uint16_t)) {
+    Combine(seed, reinterpret_cast<const uint16_t*>(ptr));
+    ptr += sizeof(uint16_t);
+    num_bytes -= sizeof(uint16_t);
+  }
+  while (num_bytes >= sizeof(uint8_t)) {
+    Combine(seed, reinterpret_cast<const uint8_t*>(ptr));
+    ptr += sizeof(uint8_t);
+    num_bytes -= sizeof(uint8_t);
+  }
+  return seed;
+}
+
+uint64_t BuildMaterialKey(ObjectManager::MaterialType material_type,
+                         const mjrMaterial& material) {
+  // Normally, hashing the struct by memory would be a problem because of
+  // padding and other uninitialized data. However, we do a memset(0) on the
+  // entire structure in mjr_defaultMaterial so this should be safe.
+  uint64_t key = hash(material);
+  Combine(key, &material_type);
+  return key;
+}
+
 ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
                                             const Mesh* mesh) {
   if (material.decor_ux) {
@@ -80,7 +128,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   if (color_texture == nullptr) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhongColorFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhongColorReflect;
     } else {
       return ObjectManager::kPhongColor;
@@ -88,7 +136,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else if (color_texture->GetSamplerType() == mjTEXTURE_CUBE) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhongCubeFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhongCubeReflect;
     } else {
       return ObjectManager::kPhongCube;
@@ -96,7 +144,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else if (has_texcoords) {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhong2dUvFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhong2dUvReflect;
     } else {
       return ObjectManager::kPhong2dUv;
@@ -104,7 +152,7 @@ ObjectManager::MaterialType GetMaterialType(const mjrMaterial& material,
   } else {
     if (material.color[3] < 1.0f) {
       return ObjectManager::kPhong2dFade;
-    } else if (material.reflective) {
+    } else if (material.reflectance > 0) {
       return ObjectManager::kPhong2dReflect;
     } else {
       return ObjectManager::kPhong2d;
diff --git a/src/experimental/filament/filament/material.h b/src/experimental/filament/filament/material.h
index a131633d..9bb0db64 100644
--- a/src/experimental/filament/filament/material.h
+++ b/src/experimental/filament/filament/material.h
@@ -15,6 +15,7 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATERIAL_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATERIAL_H_
 
+#include <cstdint>
 #include <filament/Engine.h>
 #include <filament/MaterialInstance.h>
 #include "experimental/filament/filament/mesh.h"
@@ -23,6 +24,11 @@
 
 namespace mujoco {
 
+// Generates a hash for the given material (type) and material parameters.
+using MaterialKey = uint64_t;
+MaterialKey BuildMaterialKey(ObjectManager::MaterialType type,
+                             const mjrMaterial& material);
+
 // Returns a MaterialType that best matches the given material data and mesh.
 ObjectManager::MaterialType GetMaterialType(
     const mjrMaterial& material, const Mesh* mesh);
diff --git a/src/experimental/filament/filament/renderable.cc b/src/experimental/filament/filament/renderable.cc
index f181b82c..04665f22 100644
--- a/src/experimental/filament/filament/renderable.cc
+++ b/src/experimental/filament/filament/renderable.cc
@@ -15,13 +15,18 @@
 #include "experimental/filament/filament/renderable.h"
 
 #include <cstdint>
+#include <functional>
+#include <memory>
 #include <numbers>
+#include <span>
 
 #include <filament/Engine.h>
 #include <filament/Material.h>
 #include <filament/RenderableManager.h>
 #include <filament/Scene.h>
 #include <filament/TransformManager.h>
+#include <math/TMatHelpers.h>
+#include <math/TVecHelpers.h>
 #include <math/mat3.h>
 #include <math/mat4.h>
 #include <math/vec2.h>
@@ -29,10 +34,12 @@
 #include <math/vec4.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/filament_util.h"
 #include "experimental/filament/filament/builtins.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
@@ -57,7 +64,6 @@ Renderable::Renderable(filament::Engine* engine,
 Renderable::~Renderable() noexcept {
   filament::Engine* engine = GetEngine();
   utils::EntityManager& em = utils::EntityManager::get();
-
   for (Part& part : parts_) {
     if (assigned_scene_) {
       assigned_scene_->remove(part.entity);
@@ -65,11 +71,8 @@ Renderable::~Renderable() noexcept {
     engine->destroy(part.entity);
     em.destroy(part.entity);
   }
-  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
-    if (instances_[i] != nullptr) {
-      engine->destroy(instances_[i]);
-      instances_[i] = nullptr;
-    }
+  for (auto& instance : instances_) {
+    engine->destroy(instance.second);
   }
 }
 
@@ -134,9 +137,6 @@ void Renderable::InitPartEntity(Part& part) {
   } else {
     builder.culling(false);
   }
-  if (instances_[static_cast<int>(draw_mode_)] != nullptr) {
-    builder.material(0, instances_[static_cast<int>(draw_mode_)]);
-  }
   builder.castShadows(params_.cast_shadows);
   builder.receiveShadows(params_.receive_shadows);
   builder.layerMask(0xff, params_.layer_mask);
@@ -168,7 +168,14 @@ void Renderable::UpdateTransform() {
   }
 
   filament::TransformManager& tm = GetEngine()->getTransformManager();
-  if (get_transform_fn_) {
+  if (geom_type_ == mjGEOM_PLANE && (trs_.size[0] <= 0 || trs_.size[1] <= 0)) {
+    infinite_plane_ = true;
+    const mat4f transform =
+        filament::math::mat4f(trs_.rotation, trs_.translation);
+    for (Part& part : parts_) {
+      tm.setTransform(tm.getInstance(part.entity), transform);
+    }
+  } else if (get_transform_fn_) {
     for (int i = 0; i < parts_.size(); ++i) {
       const mat4f transform = get_transform_fn_(i, trs_);
       tm.setTransform(tm.getInstance(parts_[i].entity), transform);
@@ -220,68 +227,171 @@ void Renderable::RemoveFromScene(filament::Scene* scene) {
 
 void Renderable::UpdateMaterial(const mjrMaterial& material) {
   material_ = material;
-
-  const Mesh* mesh = !parts_.empty() ? parts_[0].mesh : nullptr;
-
-  const ObjectManager::MaterialType type = GetMaterialType(material, mesh);
-  AssignMaterial(mjDRAW_MODE_COLOR, type);
-  if (!material_.decor_ux) {
-    AssignMaterial(mjDRAW_MODE_WIREFRAME, type);
-    AssignMaterial(mjDRAW_MODE_DEPTH, ObjectManager::kUnlitDepth);
-    AssignMaterial(mjDRAW_MODE_SEGMENTATION, ObjectManager::kUnlitSegmentation);
-  }
-
-  for (int i = 0; i < mjNUM_DRAW_MODES; ++i) {
-    if (instances_[i]) {
-      UpdateMaterialInstance(instances_[i], material_, object_mgr_);
-    }
-  }
-}
-
-void Renderable::AssignMaterial(mjrDrawMode mode,
-                                ObjectManager::MaterialType material_type) {
-  const int index = static_cast<int>(mode);
-
-  filament::Material* material = object_mgr_->GetMaterial(material_type);
-  if (instances_[index]) {
-    if (instances_[index]->getMaterial() == material) {
-      // The correct material is already assigned, do nothing.
-      return;
-    } else {
-      GetEngine()->destroy(instances_[index]);
-      instances_[index] = nullptr;
-    }
-  }
-  if (material) {
-    instances_[index] = material->createInstance();
-    if (geom_type_ == mjGEOM_PLANE || geom_type_ == mjGEOM_TRIANGLE) {
-      instances_[index]->setCullingMode(
-          filament::MaterialInstance::CullingMode::NONE);
-    }
-  }
 }
 
 const mjrMaterial& Renderable::GetMaterial() const {
   return material_;
 }
 
-void Renderable::SetDrawMode(mjrDrawMode mode) {
-  // Only SceneObjects support non-color draw modes.
-  if (material_.decor_ux) {
-    mode = mjDRAW_MODE_COLOR;
+void Renderable::Prepare(std::span<const mjrRenderRequest*> requests) {
+  // We assume BindMaterialInstance will be called with the same requests in
+  // the same order. As such, we'll just store the draw state in a deque rather
+  // than trying to perform any kind of matching with the requests.
+  if (!draw_queue_.empty()) {
+    mju_error(
+        "Expected all previous draw calls to be bound before preparing new "
+        "draw calls.");
+  }
+  curr_state_ = DrawState();
+
+  int num_reflections = 0;
+  for (const mjrRenderRequest* request : requests) {
+    DrawState draw_state;
+    mjrMaterial material = material_;
+
+    draw_state.wireframe = (request->draw_mode == mjDRAW_MODE_WIREFRAME);
+    if (material.decor_ux) {
+      draw_state.cast_shadows = false;
+      draw_state.receive_shadows = false;
+    }
+
+    if (geom_type_ == mjGEOM_PLANE) {
+      const float3 camera_pos = ReadFloat3(request->camera.pos);
+      const float3 position = transform_[3].xyz;
+      const float3 forward = transform_[2].xyz;
+      const bool is_behind_camera = dot(camera_pos - position, forward) < 0;
+      if (is_behind_camera) {
+        material.color[3] *= 0.3;
+        draw_state.receive_shadows = false;
+      }
+    }
+
+    const bool reflective = request->draw_mode == mjDRAW_MODE_COLOR &&
+                            request->enable_reflections &&
+                            material.reflectance > 0.0;
+    if (reflective) {
+      // Allocate a render target as needed and assign it to the material's
+      // reflection texture.
+      if (reflect_targets_.size() == num_reflections) {
+        mjrRenderTargetConfig config;
+        mjr_defaultRenderTargetConfig(&config);
+        config.color_format = mjPIXEL_FORMAT_RGBA8;
+        config.depth_format = mjPIXEL_FORMAT_DEPTH32F;
+        reflect_targets_.push_back(
+            std::make_unique<RenderTarget>(GetEngine(), config));
+      }
+      RenderTarget* target = reflect_targets_[num_reflections].get();
+      target->Prepare(request->viewport.width, request->viewport.height);
+      material.reflection_texture = target->GetColorTexture();
+      draw_state.reflection_idx = num_reflections;
+      ++num_reflections;
+    }
+
+    draw_state.material_key =
+        PrepareMaterialInstance(material, request->draw_mode);
+    draw_queue_.push_back(draw_state);
+  }
+}
+
+MaterialKey Renderable::PrepareMaterialInstance(const mjrMaterial& material,
+                                                mjrDrawMode draw_mode) {
+  ObjectManager::MaterialType type;
+  if (draw_mode == mjDRAW_MODE_COLOR || draw_mode == mjDRAW_MODE_WIREFRAME) {
+    const Mesh* mesh = !parts_.empty() ? parts_[0].mesh : nullptr;
+    type = GetMaterialType(material, mesh);
+  } else if (draw_mode == mjDRAW_MODE_DEPTH) {
+    type = ObjectManager::kUnlitDepth;
+  } else if (draw_mode == mjDRAW_MODE_SEGMENTATION) {
+    type = ObjectManager::kUnlitSegmentation;
+  } else {
+    mju_error("Invalid draw mode: %d", draw_mode);
+    return 0;
   }
 
-  SetWireframe(mode == mjDRAW_MODE_WIREFRAME);
+  const MaterialKey key = BuildMaterialKey(type, material);
 
-  filament::MaterialInstance* instance = instances_[static_cast<int>(mode)];
-  if (instance) {
+  auto it = instances_.find(key);
+  if (it == instances_.end()) {
+    filament::MaterialInstance* instance =
+        object_mgr_->GetMaterial(type)->createInstance();
+    UpdateMaterialInstance(instance, material, object_mgr_);
+    if (geom_type_ == mjGEOM_PLANE || geom_type_ == mjGEOM_TRIANGLE) {
+      instance->setCullingMode(filament::MaterialInstance::CullingMode::NONE);
+    }
+    instances_[key] = instance;
+  }
+  return key;
+}
+
+RenderTarget* Renderable::GetReflectionTarget() const {
+  if (curr_state_.reflection_idx < 0) {
+    return nullptr;
+  } else {
+    return reflect_targets_[curr_state_.reflection_idx].get();
+  }
+}
+
+void Renderable::BindMaterialInstance(const mjrRenderRequest& request) {
+  if (draw_queue_.empty()) {
+    mju_error("No material instances to bind.");
+  }
+
+  if (geom_type_ == mjGEOM_PLANE && infinite_plane_) {
+    // Emulate an infinite plane by recentering a large quad in world space
+    // relative to the camera. We use the shared mjMAXPLANEGRID value as the
+    // size of the quad to ensure the texture scaling matches.
+    static constexpr float kInfiniteScale = 0.5f * mjMAXPLANEGRID;
+    const mat4f scaling =
+        mat4f::scaling(float3{kInfiniteScale, kInfiniteScale, 1.0f});
+
+    const float3 camera_pos = ReadFloat3(request.camera.pos);
+    const float3 plane_origin = transform_[3].xyz;
+    const mat3f plane_rotation = transform_.upperLeft();
+
+    const float3 vec = camera_pos - plane_origin;
+    const float3 plane_x = normalize(plane_rotation[0]);
+    const float3 plane_y = normalize(plane_rotation[1]);
+
+    // Project camera position onto the plane's local XY axes.
+    float dx = dot(vec, plane_x);
+    float dy = dot(vec, plane_y);
+
+    // Quantize based on uv_scale.
+    const float tile_size[] = {kInfiniteScale / material_.uv_scale[0],
+                               kInfiniteScale / material_.uv_scale[1]};
+    dx = tile_size[0] * mju_round(dx / tile_size[0]);
+    dy = tile_size[1] * mju_round(dy / tile_size[1]);
+
+    // Calculate the new center quad as a displacement from the plane origin.
+    const float3 displacement = dx * plane_x + dy * plane_y;
+    const float3 center = plane_origin + displacement;
+
+    const mat4f transform = mat4f(plane_rotation, center) * scaling;
+    filament::TransformManager& tm = GetEngine()->getTransformManager();
+    for (Part& part : parts_) {
+      tm.setTransform(tm.getInstance(part.entity), transform);
+    }
+  }
+
+  const DrawState& state = draw_queue_.front();
+  SetCastShadows(state.cast_shadows);
+  SetReceiveShadows(state.receive_shadows);
+  SetWireframe(state.wireframe);
+  SetMaterialInstance(state.material_key);
+  curr_state_ = state;
+  draw_queue_.pop_front();
+}
+
+void Renderable::SetMaterialInstance(MaterialKey key) {
+  if (key != curr_state_.material_key) {
+    filament::MaterialInstance* instance = instances_[key];
     filament::RenderableManager& rm = GetEngine()->getRenderableManager();
     for (Part& part : parts_) {
       filament::RenderableManager::Instance ri = rm.getInstance(part.entity);
       rm.setMaterialInstanceAt(ri, 0, instance);
     }
+    curr_state_.material_key = key;
   }
-  draw_mode_ = mode;
 }
 
 std::uint8_t Renderable::SetLayerMask(std::uint8_t mask) {
diff --git a/src/experimental/filament/filament/renderable.h b/src/experimental/filament/filament/renderable.h
index ab0cd3b2..d1366cf8 100644
--- a/src/experimental/filament/filament/renderable.h
+++ b/src/experimental/filament/filament/renderable.h
@@ -16,7 +16,11 @@
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_RENDERABLE_H_
 
 #include <cstdint>
+#include <deque>
 #include <functional>
+#include <memory>
+#include <span>
+#include <unordered_map>
 #include <vector>
 
 #include <filament/Engine.h>
@@ -26,8 +30,10 @@
 #include <math/vec3.h>
 #include <utils/Entity.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
+#include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/render_context_filament.h"
 
 namespace mujoco {
@@ -94,15 +100,23 @@ class Renderable : public mjrRenderable {
   // Removes this renderable from the filament Scene.
   void RemoveFromScene(filament::Scene* scene);
 
-  // Determines how this renderable will be drawn. See mjrDrawMode for details.
-  void SetDrawMode(mjrDrawMode mode);
-
   // Updates the parameters and textures of the material for this renderable.
   void UpdateMaterial(const mjrMaterial& material);
 
   // Returns this renderable's current material.
   const mjrMaterial& GetMaterial() const;
 
+  // Prepares the material instances for the given render requests.
+  void Prepare(std::span<const mjrRenderRequest*> requests);
+
+  // Binds the material instance for the given render request.
+  void BindMaterialInstance(const mjrRenderRequest& request);
+
+  // Returns the render target used for reflections, if any. The main render
+  // function will use this target to render a reflection for this renderable
+  // based on the current camera position.
+  RenderTarget* GetReflectionTarget() const;
+
   static Renderable* downcast(mjrRenderable* renderable) {
     return static_cast<Renderable*>(renderable);
   }
@@ -132,6 +146,21 @@ class Renderable : public mjrRenderable {
     }
   };
 
+  // For each render request in a batch, we may need to render the object
+  // differently (e.g. a different reflection texture for each camera). However,
+  // we cannot change the material instance during the actual frame rendering
+  // (i.e. between beginFrame/endFrame). Instead, we can switch out the material
+  // instance entirely for each request. We keep track of which instance, as
+  // well as other render state, to use with each render request.
+  struct DrawState {
+    MaterialKey material_key;
+    // Index to the reflection target to use, if any.
+    int reflection_idx = -1;
+    bool cast_shadows = true;
+    bool receive_shadows = true;
+    bool wireframe = false;
+  };
+
   // When composing a multi-part renderable, each Entity will have its own
   // transform offset based on the transform of the Renderable itself.
   using GetTransformFn = std::function<filament::math::mat4f(int, const Trs&)>;
@@ -140,22 +169,26 @@ class Renderable : public mjrRenderable {
   void InitPartEntity(Part& part);
   void UpdateTransform();
 
-  void AssignMaterial(mjrDrawMode mode,
-                      ObjectManager::MaterialType material_type);
+  MaterialKey PrepareMaterialInstance(const mjrMaterial& material,
+                                      mjrDrawMode draw_mode);
+  void SetMaterialInstance(MaterialKey key);
 
   filament::Engine* GetEngine();
 
   ObjectManager* object_mgr_;
   mjrRenderableParams params_;
-  filament::MaterialInstance* instances_[mjNUM_DRAW_MODES] = {nullptr};
   mjtGeom geom_type_ = mjGEOM_NONE;
   mjrMaterial material_;
-  mjrDrawMode draw_mode_ = mjDRAW_MODE_COLOR;
+  std::unordered_map<MaterialKey, filament::MaterialInstance*> instances_;
+  std::vector<std::unique_ptr<RenderTarget>> reflect_targets_;
+  std::deque<DrawState> draw_queue_;
+  DrawState curr_state_;
   filament::Scene* assigned_scene_ = nullptr;
   std::vector<Part> parts_;
   filament::math::mat4f transform_;
   GetTransformFn get_transform_fn_;
   Trs trs_;
+  bool infinite_plane_ = false;
   bool wireframe_ = false;
 };
 
diff --git a/src/experimental/filament/filament/scene_view.cc b/src/experimental/filament/filament/scene_view.cc
index c794a298..f76f7e55 100644
--- a/src/experimental/filament/filament/scene_view.cc
+++ b/src/experimental/filament/filament/scene_view.cc
@@ -14,11 +14,9 @@
 
 #include "experimental/filament/filament/scene_view.h"
 
-#include <algorithm>
-#include <array>
-#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <span>
 #include <string_view>
 
 #include <filament/ColorGrading.h>
@@ -145,6 +143,7 @@ SceneView::SceneView(filament::Engine* engine, const mjrSceneParams& params)
   reflect_view_->setPostProcessingEnabled(false);
   reflect_view_->setFrontFaceWindingInverted(true);
   reflect_view_->setVisibleLayers(0xff, params.reflection_layer_mask);
+  reflect_view_->setMultiSampleAntiAliasingOptions({.enabled = false});
 
   // Rotate the fog to align with mujoco's +Z up space.
   auto fog = main_view_->getFogEntity();
@@ -167,7 +166,6 @@ SceneView::~SceneView() {
   }
   lights_.clear();
   renderables_.clear();
-  reflect_targets_.clear();
   engine_->destroyCameraComponent(reflect_camera_->getEntity());
   engine_->destroy(reflect_view_);
   engine_->destroyCameraComponent(camera_->getEntity());
@@ -194,18 +192,11 @@ void SceneView::RemoveFromScene(Light* light) {
 void SceneView::AddToScene(Renderable* renderable) {
   if (renderables_.insert(renderable).second) {
     renderable->AddToScene(scene_);
-    if (renderable->GetMaterial().reflective) {
-      AddReflectiveRenderable(renderable);
-    }
   }
 }
 
 void SceneView::RemoveFromScene(Renderable* renderable) {
   if (renderables_.erase(renderable)) {
-    auto it = std::find(reflectives_.begin(), reflectives_.end(), renderable);
-    if (it != reflectives_.end()) {
-      reflectives_.erase(it);
-    }
     renderable->RemoveFromScene(scene_);
   }
 }
@@ -224,34 +215,30 @@ void SceneView::SetSkybox(const Texture* skybox_texture) {
   }
 }
 
+void SceneView::PrepareToRender(std::span<const mjrRenderRequest*> requests) {
+  for (const mjrRenderRequest* request : requests) {
+    if (request->scene != this) {
+      mju_error("Invalid scene for SceneView::PrepareToRender.");
+    }
+  }
+  for (Renderable* renderable : renderables_) {
+    renderable->Prepare(requests);
+  }
+}
+
 void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& request) {
   if (request.scene != this) {
     mju_error("Invalid scene for SceneView::Render.");
   }
 
-  if (request.enable_reflections) {
-    EnableReflections();
-  } else {
-    DisableReflections();
-  }
-
   filament::Viewport viewport(request.viewport.left, request.viewport.bottom,
                               request.viewport.width, request.viewport.height);
-  main_view_->setViewport(viewport);
-  depth_segment_view_->setViewport(viewport);
-  reflect_view_->setViewport(viewport);
-
-  SetupCamera(request.camera, viewport, camera_);
-
-  for (auto& iter : renderables_) {
-    iter->SetDrawMode(request.draw_mode);
-  }
-
   filament::View* view = main_view_;
   if (request.draw_mode == mjDRAW_MODE_DEPTH ||
       request.draw_mode == mjDRAW_MODE_SEGMENTATION) {
     view = depth_segment_view_;
   }
+  view->setViewport(viewport);
   view->setShadowingEnabled(request.enable_shadows);
   view->setPostProcessingEnabled(request.enable_post_processing);
 
@@ -264,25 +251,30 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
     view->setMultiSampleAntiAliasingOptions({.enabled = false});
   }
 
-  // Render reflection passes.
-  if (request.draw_mode == mjDRAW_MODE_COLOR && reflections_enabled_) {
-    for (size_t i = 0; i < reflectives_.size(); ++i) {
-      Renderable* renderable = reflectives_[i];
+  SetupCamera(request.camera, viewport, camera_);
+
+  for (auto& iter : renderables_) {
+    iter->BindMaterialInstance(request);
+
+    if (RenderTarget* target = iter->GetReflectionTarget()) {
+      viewport.left = 0;
+      viewport.bottom = 0;
+      reflect_view_->setViewport(viewport);
 
       // We assume the 0th entity is the reflective entity.
-      mat4 transform(renderable->GetTransform());
+      mat4 transform(iter->GetTransform());
       SetupReflectionCamera(transform, camera_, reflect_camera_);
 
       // Hide reflective surface from its own reflection pass.
-      std::uint8_t previous_layer_mask = renderable->SetLayerMask(0x00);
+      std::uint8_t previous_layer_mask = iter->SetLayerMask(0x00);
 
       // Render the reflection to its render target.
-      reflect_view_->setRenderTarget(
-          reflect_targets_[i]->GetFilamentRenderTarget());
+      reflect_view_->setRenderTarget(target->GetFilamentRenderTarget());
       renderer->render(reflect_view_);
+      reflect_view_->setRenderTarget(nullptr);
 
       // Unhide the reflective surface.
-      renderable->SetLayerMask(previous_layer_mask);
+      iter->SetLayerMask(previous_layer_mask);
     }
   }
 
@@ -296,33 +288,6 @@ void SceneView::Render(filament::Renderer* renderer, const mjrRenderRequest& req
   }
 }
 
-void SceneView::AddReflectiveRenderable(Renderable* renderable) {
-  const int index = reflectives_.size();
-  reflectives_.push_back(renderable);
-
-  // Ensure we have the same number of render targets as we do reflective
-  // renderables.
-  while (reflect_targets_.size() < reflectives_.size()) {
-    mjrRenderTargetConfig config;
-    mjr_defaultRenderTargetConfig(&config);
-
-    config.color_format = mjPIXEL_FORMAT_RGBA8;
-    config.depth_format = mjPIXEL_FORMAT_DEPTH32F;
-    reflect_targets_.push_back(std::make_unique<RenderTarget>(engine_, config));
-  }
-
-  // Prepare a render target for the reflective renderable.
-  auto viewport = reflect_view_->getViewport();
-  auto& target = reflect_targets_[index];
-  target->Prepare(viewport.width, viewport.height);
-
-  if (reflections_enabled_) {
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = target->GetColorTexture();
-    renderable->UpdateMaterial(material);
-  }
-}
-
 void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
   auto tone_mapper = CreateToneMapper(opts.tone_mapper);
   auto color_grading = ToBuilder(color_grading_options_)
@@ -336,33 +301,6 @@ void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {
   color_grading_options_ = opts;
 }
 
-void SceneView::EnableReflections() {
-  if (reflections_enabled_) {
-    return;
-  }
-
-  reflections_enabled_ = true;
-  for (int i = 0; i < reflectives_.size(); ++i) {
-    Renderable* renderable = reflectives_[i];
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = reflect_targets_[i]->GetColorTexture();
-    renderable->UpdateMaterial(material);
-  }
-}
-
-void SceneView::DisableReflections() {
-  if (!reflections_enabled_) {
-    return;
-  }
-
-  reflections_enabled_ = false;
-  for (Renderable* renderable : reflectives_) {
-    mjrMaterial material = renderable->GetMaterial();
-    material.reflection_texture = nullptr;
-    renderable->UpdateMaterial(material);
-  }
-}
-
 filament::View* SceneView::GetDefaultRenderView() {
   return main_view_;
 }
diff --git a/src/experimental/filament/filament/scene_view.h b/src/experimental/filament/filament/scene_view.h
index 9237ccdb..60c585e2 100644
--- a/src/experimental/filament/filament/scene_view.h
+++ b/src/experimental/filament/filament/scene_view.h
@@ -15,10 +15,8 @@
 #ifndef MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 #define MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_SCENE_VIEW_H_
 
-#include <array>
-#include <memory>
+#include <span>
 #include <unordered_set>
-#include <vector>
 
 #include <filament/Camera.h>
 #include <filament/ColorGrading.h>
@@ -29,7 +27,6 @@
 #include "experimental/filament/filament/color_grading_options.h"
 #include "experimental/filament/filament/light.h"
 #include "experimental/filament/filament/renderable.h"
-#include "experimental/filament/filament/render_target.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/render_context_filament.h"
 
@@ -55,6 +52,12 @@ class SceneView : public mjrScene {
   void RemoveFromScene(Renderable* renderable);
   void SetSkybox(const Texture* skybox_texture);
 
+  // Performs necessary preparations in order to render the given requests.
+  // Assumes that the Render() function will be called the same number of times
+  // and in the same order with the given requests.
+  void PrepareToRender(std::span<const mjrRenderRequest*> requests);
+
+  // Fulfills the given render request using the renderer.
   void Render(filament::Renderer* renderer, const mjrRenderRequest& request);
 
   // Returns the filament Engine managing the scene.
@@ -80,13 +83,6 @@ class SceneView : public mjrScene {
   }
 
  private:
-  void EnableReflections();
-  void DisableReflections();
-
-  // Marks a renderable as reflective. Reflective renderables have to be
-  // rendered in their own passes to create the reflective texture.
-  void AddReflectiveRenderable(Renderable* renderable);
-
   filament::Engine* engine_ = nullptr;
   filament::Scene* scene_ = nullptr;
   filament::Camera* camera_ = nullptr;
@@ -103,11 +99,6 @@ class SceneView : public mjrScene {
   // Custom view and camera for reflective surfaces.
   filament::View* reflect_view_ = nullptr;
   filament::Camera* reflect_camera_ = nullptr;
-
-  // The list of reflective renderables and their corresponding render targets.
-  bool reflections_enabled_ = true;
-  std::vector<Renderable*> reflectives_;
-  std::vector<std::unique_ptr<RenderTarget>> reflect_targets_;
 };
 }  // namespace mujoco
 
diff --git a/src/experimental/filament/filament_util.h b/src/experimental/filament/filament_util.h
index c8398522..7cc9ee33 100644
--- a/src/experimental/filament/filament_util.h
+++ b/src/experimental/filament/filament_util.h
@@ -52,9 +52,9 @@ template <typename T>
 inline filament::math::mat3f ReadMat3(const T* arr, int index = 0) {
   // clang-format off
   const T* ptr = arr + (9 * index);
-  return filament::math::mat3f(ptr[0], ptr[3], ptr[6],
-                               ptr[1], ptr[4], ptr[7],
-                               ptr[2], ptr[5], ptr[8]);
+  return filament::math::mat3f(ptr[0], ptr[1], ptr[2],
+                               ptr[3], ptr[4], ptr[5],
+                               ptr[6], ptr[7], ptr[8]);
   // clang-format on
 }
 
diff --git a/src/experimental/filament/render_context_filament.h b/src/experimental/filament/render_context_filament.h
index fcc587c9..98204f75 100644
--- a/src/experimental/filament/render_context_filament.h
+++ b/src/experimental/filament/render_context_filament.h
@@ -139,8 +139,6 @@ typedef enum mjrDrawMode_ {
   mjDRAW_MODE_SEGMENTATION,
 } mjrDrawMode;
 
-enum { mjNUM_DRAW_MODES = 4 };  // Number of modes in `mjrDrawMode`.
-
 // Parameters describing the camera to use for rendering an image.
 typedef mjvGLCamera mjrCamera;
 
@@ -614,8 +612,6 @@ struct mjrMaterial {
   // The emissive (glow) factor of the object.
   float emissive;
 
-  // Whether or not the object is a reflective surface. Only applies to planes.
-  mjtByte reflective;
   // The blend factor to use for reflective surfaces. A value of 1.0 means that
   // the surface is fully reflective (i.e. a mirror).
   float reflectance;

From af4be63cd8ac26eaf49ee5269097c8e6d050e8f0 Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Wed, 27 May 2026 09:10:43 -0700
Subject: [PATCH 40/61] Add xmacros for MjSpec.

PiperOrigin-RevId: 922169069
Change-Id: I01426596e49196996dfa324d5ab3849a90de2e3a
---
 CMakeLists.txt                       |   1 +
 doc/APIreference/APIglobals.rst      |   5 +-
 doc/programming/index.rst            |   6 +-
 include/mujoco/mjspecmacro.h         | 603 +++++++++++++++++++++++++++
 test/CMakeLists.txt                  |   2 +
 test/mjspecmacro_test.cc             | 145 +++++++
 unity/Runtime/Bindings/MjBindings.cs |   1 +
 7 files changed, 760 insertions(+), 3 deletions(-)
 create mode 100644 include/mujoco/mjspecmacro.h
 create mode 100644 test/mjspecmacro_test.cc

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b50248f8..0423d3f4 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -117,6 +117,7 @@ set(MUJOCO_HEADERS
     include/mujoco/mjrender.h
     include/mujoco/mjsan.h
     include/mujoco/mjspec.h
+    include/mujoco/mjspecmacro.h
     include/mujoco/mjthread.h
     include/mujoco/mjtype.h
     include/mujoco/mjui.h
diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst
index 80c758af..9890735b 100644
--- a/doc/APIreference/APIglobals.rst
+++ b/doc/APIreference/APIglobals.rst
@@ -706,7 +706,8 @@ X Macros
 ^^^^^^^^
 
 The X Macros are not needed in most user projects. They are used internally to allocate the model, and are also
-available for users who know how to use this programming technique. See the header file `mjxmacro.h
-<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`_ for the actual definitions. They are
+available for users who know how to use this programming technique. See the header files `mjxmacro.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`_ and `mjspecmacro.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspecmacro.h>`_ for the actual definitions. They are
 particularly useful in writing MuJoCo wrappers for scripting languages, where dynamic structures matching the MuJoCo
 data structures need to be constructed programmatically.
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index 99cc50dc..24c2402f 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -151,7 +151,7 @@ links below, to make this documentation self-contained.
 
 `mujoco.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mujoco.h>`__
    This is the main header file and must be included in all programs using MuJoCo. It defines all API functions and
-   global variables, and includes all other header files except mjxmacro.h.
+   global variables, and includes all other header files except mjxmacro.h and mjspecmacro.h.
 `mjmodel.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__
    Defines the C structure :ref:`mjModel` which is the runtime representation of the
    model being simulated.
@@ -179,6 +179,10 @@ links below, to make this documentation self-contained.
    This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
    automate the mapping of mjModel and mjData into scripting languages, as well as other operations that require
    accessing all fields of mjModel and mjData.
+`mjspecmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspecmacro.h>`__
+   This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
+   automate the mapping of mjSpec and its element structs into scripting languages, as well as other operations that require
+   accessing all fields of mjSpec during procedural model editing.
 `mjexport.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjexport.h>`__
    Macros used for exporting public symbols from the MuJoCo library. This header should not be used directly by client
    code.
diff --git a/include/mujoco/mjspecmacro.h b/include/mujoco/mjspecmacro.h
new file mode 100644
index 00000000..81e861e1
--- /dev/null
+++ b/include/mujoco/mjspecmacro.h
@@ -0,0 +1,603 @@
+// Copyright 2026 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_MJSPECMACRO_H_
+#define MUJOCO_MJSPECMACRO_H_
+
+
+//-------------------------------- mjsElement ------------------------------------------------------
+
+#define MJSELEMENT_FIELDS       \
+    X( mjtObj,   elemtype,  1 ) \
+    X( uint64_t, signature, 1 )
+
+
+//-------------------------------- mjsCompiler -----------------------------------------------------
+
+#define MJSCOMPILER_FIELDS                  \
+    X   ( mjtByte,   autolimits,        1 ) \
+    X   ( double,    boundmass,         1 ) \
+    X   ( double,    boundinertia,      1 ) \
+    X   ( double,    settotalmass,      1 ) \
+    X   ( mjtByte,   balanceinertia,    1 ) \
+    X   ( mjtByte,   fitaabb,           1 ) \
+    X   ( mjtByte,   degree,            1 ) \
+    XVEC( char,      eulerseq,          3 ) \
+    X   ( mjtByte,   discardvisual,     1 ) \
+    X   ( mjtByte,   usethread,         1 ) \
+    X   ( mjtByte,   fusestatic,        1 ) \
+    X   ( int,       inertiafromgeom,   1 ) \
+    XVEC( int,       inertiagrouprange, 2 ) \
+    X   ( mjtByte,   saveinertial,      1 ) \
+    X   ( int,       alignfree,         1 ) \
+    X   ( mjLROpt,   LRopt,             1 ) \
+    X   ( mjString*, meshdir,           1 ) \
+    X   ( mjString*, texturedir,        1 )
+
+
+//-------------------------------- mjSpec ----------------------------------------------------------
+
+#define MJSPEC_FIELDS                          \
+    X( mjsElement*, element,               1 ) \
+    X( mjString*,   modelname,             1 ) \
+    X( mjsCompiler, compiler,              1 ) \
+    X( mjtByte,     strippath,             1 ) \
+    X( mjOption,    option,                1 ) \
+    X( mjVisual,    visual,                1 ) \
+    X( mjStatistic, stat,                  1 ) \
+    X( mjtSize,     memory,                1 ) \
+    X( int,         nemax,                 1 ) \
+    X( int,         nuserdata,             1 ) \
+    X( int,         nuser_body,            1 ) \
+    X( int,         nuser_jnt,             1 ) \
+    X( int,         nuser_geom,            1 ) \
+    X( int,         nuser_site,            1 ) \
+    X( int,         nuser_cam,             1 ) \
+    X( int,         nuser_tendon,          1 ) \
+    X( int,         nuser_actuator,        1 ) \
+    X( int,         nuser_sensor,          1 ) \
+    X( int,         nkey,                  1 ) \
+    X( int,         njmax,                 1 ) \
+    X( int,         nconmax,               1 ) \
+    X( mjtSize,     nstack,                1 ) \
+    X( mjString*,   comment,               1 ) \
+    X( mjString*,   modelfiledir,          1 ) \
+    X( mjtByte,     hasImplicitPluginElem, 1 )
+
+
+//-------------------------------- mjsOrientation --------------------------------------------------
+
+#define MJSORIENTATION_FIELDS            \
+    X   ( mjtOrientation, type,      1 ) \
+    XVEC( double,         axisangle, 4 ) \
+    XVEC( double,         xyaxes,    6 ) \
+    XVEC( double,         zaxis,     3 ) \
+    XVEC( double,         euler,     3 )
+
+
+//-------------------------------- mjsPlugin -------------------------------------------------------
+
+#define MJSPLUGIN_FIELDS             \
+    X( mjsElement*, element,     1 ) \
+    X( mjString*,   name,        1 ) \
+    X( mjString*,   plugin_name, 1 ) \
+    X( mjtByte,     active,      1 ) \
+    X( mjString*,   info,        1 )
+
+
+//-------------------------------- mjsBody ---------------------------------------------------------
+
+#define MJSBODY_FIELDS                           \
+    X   ( mjsElement*,     element,          1 ) \
+    X   ( mjString*,       childclass,       1 ) \
+    XVEC( double,          pos,              3 ) \
+    XVEC( double,          quat,             4 ) \
+    X   ( mjsOrientation,  alt,              1 ) \
+    X   ( double,          mass,             1 ) \
+    XVEC( double,          ipos,             3 ) \
+    XVEC( double,          iquat,            4 ) \
+    XVEC( double,          inertia,          3 ) \
+    X   ( mjsOrientation,  ialt,             1 ) \
+    XVEC( double,          fullinertia,      6 ) \
+    X   ( mjtByte,         mocap,            1 ) \
+    X   ( double,          gravcomp,         1 ) \
+    X   ( mjtSleepPolicy,  sleep,            1 ) \
+    X   ( mjDoubleVec*,    userdata,         1 ) \
+    X   ( mjtByte,         explicitinertial, 1 ) \
+    X   ( mjsPlugin,       plugin,           1 ) \
+    X   ( mjString*,       info,             1 )
+
+
+//-------------------------------- mjsFrame --------------------------------------------------------
+
+#define MJSFRAME_FIELDS                   \
+    X   ( mjsElement*,    element,    1 ) \
+    X   ( mjString*,      childclass, 1 ) \
+    XVEC( double,         pos,        3 ) \
+    XVEC( double,         quat,       4 ) \
+    X   ( mjsOrientation, alt,        1 ) \
+    X   ( mjString*,      info,       1 )
+
+
+//-------------------------------- mjsJoint --------------------------------------------------------
+
+#define MJSJOINT_FIELDS                               \
+    X   ( mjsElement*,   element,         1         ) \
+    X   ( mjtJoint,      type,            1         ) \
+    XVEC( double,        pos,             3         ) \
+    XVEC( double,        axis,            3         ) \
+    X   ( double,        ref,             1         ) \
+    X   ( int,           align,           1         ) \
+    XVEC( double,        stiffness,       mjNPOLY+1 ) \
+    X   ( double,        springref,       1         ) \
+    XVEC( double,        springdamper,    2         ) \
+    X   ( int,           limited,         1         ) \
+    XVEC( double,        range,           2         ) \
+    X   ( double,        margin,          1         ) \
+    XVEC( mjtNum,        solref_limit,    mjNREF    ) \
+    XVEC( mjtNum,        solimp_limit,    mjNIMP    ) \
+    X   ( int,           actfrclimited,   1         ) \
+    XVEC( double,        actfrcrange,     2         ) \
+    X   ( double,        armature,        1         ) \
+    XVEC( double,        damping,         mjNPOLY+1 ) \
+    X   ( double,        frictionloss,    1         ) \
+    XVEC( mjtNum,        solref_friction, mjNREF    ) \
+    XVEC( mjtNum,        solimp_friction, mjNIMP    ) \
+    X   ( int,           group,           1         ) \
+    X   ( mjtByte,       actgravcomp,     1         ) \
+    X   ( mjDoubleVec*,  userdata,        1         ) \
+    X   ( mjString*,     info,            1         )
+
+
+//-------------------------------- mjsGeom ---------------------------------------------------------
+
+#define MJSGEOM_FIELDS                              \
+    X   ( mjsElement*,    element,         1      ) \
+    X   ( mjtGeom,        type,            1      ) \
+    XVEC( double,         pos,             3      ) \
+    XVEC( double,         quat,            4      ) \
+    X   ( mjsOrientation, alt,             1      ) \
+    XVEC( double,         fromto,          6      ) \
+    XVEC( double,         size,            3      ) \
+    X   ( int,            contype,         1      ) \
+    X   ( int,            conaffinity,     1      ) \
+    X   ( int,            condim,          1      ) \
+    X   ( int,            priority,        1      ) \
+    XVEC( double,         friction,        3      ) \
+    X   ( double,         solmix,          1      ) \
+    XVEC( mjtNum,         solref,          mjNREF ) \
+    XVEC( mjtNum,         solimp,          mjNIMP ) \
+    X   ( double,         margin,          1      ) \
+    X   ( double,         gap,             1      ) \
+    X   ( double,         mass,            1      ) \
+    X   ( double,         density,         1      ) \
+    X   ( mjtGeomInertia, typeinertia,     1      ) \
+    X   ( mjtNum,         fluid_ellipsoid, 1      ) \
+    XVEC( mjtNum,         fluid_coefs,     5      ) \
+    X   ( mjString*,      material,        1      ) \
+    XVEC( float,          rgba,            4      ) \
+    X   ( int,            group,           1      ) \
+    X   ( mjString*,      hfieldname,      1      ) \
+    X   ( mjString*,      meshname,        1      ) \
+    X   ( double,         fitscale,        1      ) \
+    X   ( mjDoubleVec*,   userdata,        1      ) \
+    X   ( mjsPlugin,      plugin,          1      ) \
+    X   ( mjString*,      info,            1      )
+
+
+//-------------------------------- mjsSite ---------------------------------------------------------
+
+#define MJSSITE_FIELDS                  \
+    X   ( mjsElement*,    element,  1 ) \
+    XVEC( double,         pos,      3 ) \
+    XVEC( double,         quat,     4 ) \
+    X   ( mjsOrientation, alt,      1 ) \
+    XVEC( double,         fromto,   6 ) \
+    XVEC( double,         size,     3 ) \
+    X   ( mjtGeom,        type,     1 ) \
+    X   ( mjString*,      material, 1 ) \
+    X   ( int,            group,    1 ) \
+    XVEC( float,          rgba,     4 ) \
+    X   ( mjDoubleVec*,   userdata, 1 ) \
+    X   ( mjString*,      info,     1 )
+
+
+//-------------------------------- mjsCamera -------------------------------------------------------
+
+#define MJSCAMERA_FIELDS                        \
+    X   ( mjsElement*,    element,          1 ) \
+    XVEC( double,         pos,              3 ) \
+    XVEC( double,         quat,             4 ) \
+    X   ( mjsOrientation, alt,              1 ) \
+    X   ( mjtCamLight,    mode,             1 ) \
+    X   ( mjString*,      targetbody,       1 ) \
+    X   ( mjtProjection,  proj,             1 ) \
+    XVEC( int,            resolution,       2 ) \
+    X   ( int,            output,           1 ) \
+    X   ( double,         fovy,             1 ) \
+    X   ( double,         ipd,              1 ) \
+    XVEC( float,          intrinsic,        4 ) \
+    XVEC( float,          sensor_size,      2 ) \
+    XVEC( float,          focal_length,     2 ) \
+    XVEC( float,          focal_pixel,      2 ) \
+    XVEC( float,          principal_length, 2 ) \
+    XVEC( float,          principal_pixel,  2 ) \
+    X   ( mjDoubleVec*,   userdata,         1 ) \
+    X   ( mjString*,      info,             1 )
+
+
+//-------------------------------- mjsLight --------------------------------------------------------
+
+#define MJSLIGHT_FIELDS                  \
+    X   ( mjsElement*,  element,     1 ) \
+    XVEC( double,       pos,         3 ) \
+    XVEC( double,       dir,         3 ) \
+    X   ( mjtCamLight,  mode,        1 ) \
+    X   ( mjString*,    targetbody,  1 ) \
+    X   ( mjtByte,      active,      1 ) \
+    X   ( mjtLightType, type,        1 ) \
+    X   ( mjString*,    texture,     1 ) \
+    X   ( mjtByte,      castshadow,  1 ) \
+    X   ( float,        bulbradius,  1 ) \
+    X   ( float,        intensity,   1 ) \
+    X   ( float,        range,       1 ) \
+    XVEC( float,        attenuation, 3 ) \
+    X   ( float,        cutoff,      1 ) \
+    X   ( float,        exponent,    1 ) \
+    XVEC( float,        ambient,     3 ) \
+    XVEC( float,        diffuse,     3 ) \
+    XVEC( float,        specular,    3 ) \
+    X   ( mjString*,    info,        1 )
+
+
+//-------------------------------- mjsFlex ---------------------------------------------------------
+
+#define MJSFLEX_FIELDS                           \
+    X   ( mjsElement*,   element,       1      ) \
+    X   ( int,           contype,       1      ) \
+    X   ( int,           conaffinity,   1      ) \
+    X   ( int,           condim,        1      ) \
+    X   ( int,           priority,      1      ) \
+    XVEC( double,        friction,      3      ) \
+    X   ( double,        solmix,        1      ) \
+    XVEC( mjtNum,        solref,        mjNREF ) \
+    XVEC( mjtNum,        solimp,        mjNIMP ) \
+    X   ( double,        margin,        1      ) \
+    X   ( double,        gap,           1      ) \
+    X   ( int,           dim,           1      ) \
+    X   ( double,        radius,        1      ) \
+    XVEC( double,        size,          3      ) \
+    X   ( mjtByte,       internal,      1      ) \
+    X   ( mjtByte,       flatskin,      1      ) \
+    X   ( int,           selfcollide,   1      ) \
+    X   ( int,           passive,       1      ) \
+    X   ( int,           activelayers,  1      ) \
+    X   ( int,           group,         1      ) \
+    X   ( double,        edgestiffness, 1      ) \
+    X   ( double,        edgedamping,   1      ) \
+    XVEC( float,         rgba,          4      ) \
+    X   ( mjString*,     material,      1      ) \
+    X   ( double,        young,         1      ) \
+    X   ( double,        poisson,       1      ) \
+    X   ( double,        damping,       1      ) \
+    X   ( double,        thickness,     1      ) \
+    X   ( int,           elastic2d,     1      ) \
+    XVEC( int,           cellcount,     3      ) \
+    X   ( int,           order,         1      ) \
+    X   ( mjStringVec*,  nodebody,      1      ) \
+    X   ( mjStringVec*,  vertbody,      1      ) \
+    X   ( mjDoubleVec*,  node,          1      ) \
+    X   ( mjDoubleVec*,  vert,          1      ) \
+    X   ( mjIntVec*,     elem,          1      ) \
+    X   ( mjFloatVec*,   texcoord,      1      ) \
+    X   ( mjIntVec*,     elemtexcoord,  1      ) \
+    X   ( mjString*,     info,          1      )
+
+
+//-------------------------------- mjsMesh ---------------------------------------------------------
+
+#define MJSMESH_FIELDS                           \
+    X   ( mjsElement*,     element,          1 ) \
+    X   ( mjString*,       content_type,     1 ) \
+    X   ( mjString*,       file,             1 ) \
+    XVEC( double,          refpos,           3 ) \
+    XVEC( double,          refquat,          4 ) \
+    XVEC( double,          scale,            3 ) \
+    X   ( mjtMeshInertia,  inertia,          1 ) \
+    X   ( mjtByte,         smoothnormal,     1 ) \
+    X   ( mjtByte,         needsdf,          1 ) \
+    X   ( int,             maxhullvert,      1 ) \
+    X   ( mjFloatVec*,     uservert,         1 ) \
+    X   ( mjFloatVec*,     usernormal,       1 ) \
+    X   ( mjFloatVec*,     usertexcoord,     1 ) \
+    X   ( mjIntVec*,       userface,         1 ) \
+    X   ( mjIntVec*,       userfacenormal,   1 ) \
+    X   ( mjIntVec*,       userfacetexcoord, 1 ) \
+    X   ( mjsPlugin,       plugin,           1 ) \
+    X   ( mjString*,       material,         1 ) \
+    X   ( int,             octree_maxdepth,  1 ) \
+    X   ( mjString*,       info,             1 )
+
+
+//-------------------------------- mjsHField -------------------------------------------------------
+
+#define MJSHFIELD_FIELDS                  \
+    X   ( mjsElement*,  element,      1 ) \
+    X   ( mjString*,    content_type, 1 ) \
+    X   ( mjString*,    file,         1 ) \
+    XVEC( double,       size,         4 ) \
+    X   ( int,          nrow,         1 ) \
+    X   ( int,          ncol,         1 ) \
+    X   ( mjFloatVec*,  userdata,     1 ) \
+    X   ( mjString*,    info,         1 )
+
+
+//-------------------------------- mjsSkin ---------------------------------------------------------
+
+#define MJSSKIN_FIELDS                    \
+    X   ( mjsElement*,    element,    1 ) \
+    X   ( mjString*,      file,       1 ) \
+    X   ( mjString*,      material,   1 ) \
+    XVEC( float,          rgba,       4 ) \
+    X   ( float,          inflate,    1 ) \
+    X   ( int,            group,      1 ) \
+    X   ( mjFloatVec*,    vert,       1 ) \
+    X   ( mjFloatVec*,    texcoord,   1 ) \
+    X   ( mjIntVec*,      face,       1 ) \
+    X   ( mjStringVec*,   bodyname,   1 ) \
+    X   ( mjFloatVec*,    bindpos,    1 ) \
+    X   ( mjFloatVec*,    bindquat,   1 ) \
+    X   ( mjIntVecVec*,   vertid,     1 ) \
+    X   ( mjFloatVecVec*, vertweight, 1 ) \
+    X   ( mjString*,      info,       1 )
+
+
+//-------------------------------- mjsTexture ------------------------------------------------------
+
+#define MJSTEXTURE_FIELDS                    \
+    X   ( mjsElement*,    element,      1  ) \
+    X   ( mjtTexture,     type,         1  ) \
+    X   ( mjtColorSpace,  colorspace,   1  ) \
+    X   ( int,            builtin,      1  ) \
+    X   ( int,            mark,         1  ) \
+    XVEC( double,         rgb1,         3  ) \
+    XVEC( double,         rgb2,         3  ) \
+    XVEC( double,         markrgb,      3  ) \
+    X   ( double,         random,       1  ) \
+    X   ( int,            height,       1  ) \
+    X   ( int,            width,        1  ) \
+    X   ( int,            nchannel,     1  ) \
+    X   ( mjString*,      content_type, 1  ) \
+    X   ( mjString*,      file,         1  ) \
+    XVEC( int,            gridsize,     2  ) \
+    XVEC( char,           gridlayout,   12 ) \
+    X   ( mjStringVec*,   cubefiles,    1  ) \
+    X   ( mjByteVec*,     data,         1  ) \
+    X   ( mjtByte,        hflip,        1  ) \
+    X   ( mjtByte,        vflip,        1  ) \
+    X   ( mjString*,      info,         1  )
+
+
+//-------------------------------- mjsMaterial -----------------------------------------------------
+
+#define MJSMATERIAL_FIELDS               \
+    X   ( mjsElement*,  element,     1 ) \
+    X   ( mjStringVec*, textures,    1 ) \
+    X   ( mjtByte,      texuniform,  1 ) \
+    XVEC( float,        texrepeat,   2 ) \
+    X   ( float,        emission,    1 ) \
+    X   ( float,        specular,    1 ) \
+    X   ( float,        shininess,   1 ) \
+    X   ( float,        reflectance, 1 ) \
+    X   ( float,        metallic,    1 ) \
+    X   ( float,        roughness,   1 ) \
+    XVEC( float,        rgba,        4 ) \
+    X   ( mjString*,    info,        1 )
+
+
+//-------------------------------- mjsPair ---------------------------------------------------------
+
+#define MJSPAIR_FIELDS                          \
+    X   ( mjsElement*, element,        1      ) \
+    X   ( mjString*,   geomname1,      1      ) \
+    X   ( mjString*,   geomname2,      1      ) \
+    X   ( int,         condim,         1      ) \
+    XVEC( mjtNum,      solref,         mjNREF ) \
+    XVEC( mjtNum,      solreffriction, mjNREF ) \
+    XVEC( mjtNum,      solimp,         mjNIMP ) \
+    X   ( double,      margin,         1      ) \
+    X   ( double,      gap,            1      ) \
+    XVEC( double,      friction,       5      ) \
+    X   ( mjString*,   info,           1      )
+
+
+//-------------------------------- mjsExclude ------------------------------------------------------
+
+#define MJSEXCLUDE_FIELDS          \
+    X( mjsElement*, element,   1 ) \
+    X( mjString*,   bodyname1, 1 ) \
+    X( mjString*,   bodyname2, 1 ) \
+    X( mjString*,   info,      1 )
+
+
+//-------------------------------- mjsEquality -----------------------------------------------------
+
+#define MJSEQUALITY_FIELDS                  \
+    X   ( mjsElement*, element, 1         ) \
+    X   ( mjtEq,       type,    1         ) \
+    XVEC( double,      data,    mjNEQDATA ) \
+    X   ( mjtByte,     active,  1         ) \
+    X   ( mjString*,   name1,   1         ) \
+    X   ( mjString*,   name2,   1         ) \
+    X   ( mjtObj,      objtype, 1         ) \
+    XVEC( mjtNum,      solref,  mjNREF    ) \
+    XVEC( mjtNum,      solimp,  mjNIMP    ) \
+    X   ( mjString*,   info,    1         )
+
+
+//-------------------------------- mjsTendon -------------------------------------------------------
+
+#define MJSTENDON_FIELDS                              \
+    X   ( mjsElement*,   element,         1         ) \
+    XVEC( double,        stiffness,       mjNPOLY+1 ) \
+    XVEC( double,        springlength,    2         ) \
+    XVEC( double,        damping,         mjNPOLY+1 ) \
+    X   ( double,        frictionloss,    1         ) \
+    XVEC( mjtNum,        solref_friction, mjNREF    ) \
+    XVEC( mjtNum,        solimp_friction, mjNIMP    ) \
+    X   ( double,        armature,        1         ) \
+    X   ( int,           limited,         1         ) \
+    X   ( int,           actfrclimited,   1         ) \
+    XVEC( double,        range,           2         ) \
+    XVEC( double,        actfrcrange,     2         ) \
+    X   ( double,        margin,          1         ) \
+    XVEC( mjtNum,        solref_limit,    mjNREF    ) \
+    XVEC( mjtNum,        solimp_limit,    mjNIMP    ) \
+    X   ( mjString*,     material,        1         ) \
+    X   ( double,        width,           1         ) \
+    XVEC( float,         rgba,            4         ) \
+    X   ( int,           group,           1         ) \
+    X   ( mjDoubleVec*,  userdata,        1         ) \
+    X   ( mjString*,     info,            1         )
+
+
+//-------------------------------- mjsWrap ---------------------------------------------------------
+
+#define MJSWRAP_FIELDS           \
+    X( mjsElement*, element, 1 ) \
+    X( mjtWrap,     type,    1 ) \
+    X( mjString*,   info,    1 )
+
+
+//-------------------------------- mjsActuator -----------------------------------------------------
+
+#define MJSACTUATOR_FIELDS                          \
+    X   ( mjsElement*,   element,       1         ) \
+    X   ( mjtGain,       gaintype,      1         ) \
+    XVEC( double,        gainprm,       mjNGAIN   ) \
+    X   ( mjtBias,       biastype,      1         ) \
+    XVEC( double,        biasprm,       mjNGAIN   ) \
+    X   ( mjtDyn,        dyntype,       1         ) \
+    XVEC( double,        dynprm,        mjNDYN    ) \
+    X   ( int,           actdim,        1         ) \
+    X   ( mjtByte,       actearly,      1         ) \
+    X   ( mjtTrn,        trntype,       1         ) \
+    XVEC( double,        gear,          6         ) \
+    X   ( mjString*,     target,        1         ) \
+    X   ( mjString*,     refsite,       1         ) \
+    X   ( mjString*,     slidersite,    1         ) \
+    X   ( double,        cranklength,   1         ) \
+    XVEC( double,        lengthrange,   2         ) \
+    X   ( double,        inheritrange,  1         ) \
+    XVEC( double,        damping,       mjNPOLY+1 ) \
+    X   ( double,        armature,      1         ) \
+    X   ( int,           ctrllimited,   1         ) \
+    XVEC( double,        ctrlrange,     2         ) \
+    X   ( int,           forcelimited,  1         ) \
+    XVEC( double,        forcerange,    2         ) \
+    X   ( int,           actlimited,    1         ) \
+    XVEC( double,        actrange,      2         ) \
+    X   ( int,           group,         1         ) \
+    X   ( int,           nsample,       1         ) \
+    X   ( int,           interp,        1         ) \
+    X   ( double,        delay,         1         ) \
+    X   ( mjDoubleVec*,  userdata,      1         ) \
+    X   ( mjsPlugin,     plugin,        1         ) \
+    X   ( mjString*,     info,          1         )
+
+
+//-------------------------------- mjsSensor -------------------------------------------------------
+
+#define MJSSENSOR_FIELDS                      \
+    X   ( mjsElement*,   element,   1       ) \
+    X   ( mjtSensor,     type,      1       ) \
+    X   ( mjtObj,        objtype,   1       ) \
+    X   ( mjString*,     objname,   1       ) \
+    X   ( mjtObj,        reftype,   1       ) \
+    X   ( mjString*,     refname,   1       ) \
+    XVEC( int,           intprm,    mjNSENS ) \
+    X   ( mjtDataType,   datatype,  1       ) \
+    X   ( mjtStage,      needstage, 1       ) \
+    X   ( int,           dim,       1       ) \
+    X   ( double,        cutoff,    1       ) \
+    X   ( double,        noise,     1       ) \
+    X   ( int,           nsample,   1       ) \
+    X   ( int,           interp,    1       ) \
+    X   ( double,        delay,     1       ) \
+    XVEC( double,        interval,  2       ) \
+    X   ( mjDoubleVec*,  userdata,  1       ) \
+    X   ( mjsPlugin,     plugin,    1       ) \
+    X   ( mjString*,     info,      1       )
+
+
+//-------------------------------- mjsNumeric ------------------------------------------------------
+
+#define MJSNUMERIC_FIELDS         \
+    X( mjsElement*,  element, 1 ) \
+    X( mjDoubleVec*, data,    1 ) \
+    X( int,          size,    1 ) \
+    X( mjString*,    info,    1 )
+
+
+//-------------------------------- mjsText ---------------------------------------------------------
+
+#define MJSTEXT_FIELDS           \
+    X( mjsElement*, element, 1 ) \
+    X( mjString*,   data,    1 ) \
+    X( mjString*,   info,    1 )
+
+
+//-------------------------------- mjsTuple --------------------------------------------------------
+
+#define MJSTUPLE_FIELDS           \
+    X( mjsElement*,  element, 1 ) \
+    X( mjIntVec*,    objtype, 1 ) \
+    X( mjStringVec*, objname, 1 ) \
+    X( mjDoubleVec*, objprm,  1 ) \
+    X( mjString*,    info,    1 )
+
+
+//-------------------------------- mjsKey ----------------------------------------------------------
+
+#define MJSKEY_FIELDS             \
+    X( mjsElement*,  element, 1 ) \
+    X( double,       time,    1 ) \
+    X( mjDoubleVec*, qpos,    1 ) \
+    X( mjDoubleVec*, qvel,    1 ) \
+    X( mjDoubleVec*, act,     1 ) \
+    X( mjDoubleVec*, mpos,    1 ) \
+    X( mjDoubleVec*, mquat,   1 ) \
+    X( mjDoubleVec*, ctrl,    1 ) \
+    X( mjString*,    info,    1 )
+
+
+//-------------------------------- mjsDefault ------------------------------------------------------
+
+#define MJSDEFAULT_FIELDS          \
+    X( mjsElement*,  element,  1 ) \
+    X( mjsJoint*,    joint,    1 ) \
+    X( mjsGeom*,     geom,     1 ) \
+    X( mjsSite*,     site,     1 ) \
+    X( mjsCamera*,   camera,   1 ) \
+    X( mjsLight*,    light,    1 ) \
+    X( mjsFlex*,     flex,     1 ) \
+    X( mjsMesh*,     mesh,     1 ) \
+    X( mjsMaterial*, material, 1 ) \
+    X( mjsPair*,     pair,     1 ) \
+    X( mjsEquality*, equality, 1 ) \
+    X( mjsTendon*,   tendon,   1 ) \
+    X( mjsActuator*, actuator, 1 )
+
+
+#endif  // MUJOCO_MJSPECMACRO_H_
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index a59b55e0..1fda2951 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -108,6 +108,8 @@ mujoco_test(fixture_test)
 
 mujoco_test(header_test)
 
+mujoco_test(mjspecmacro_test)
+
 mujoco_test(pipeline_test)
 
 add_subdirectory(benchmark)
diff --git a/test/mjspecmacro_test.cc b/test/mjspecmacro_test.cc
new file mode 100644
index 00000000..d94b42ba
--- /dev/null
+++ b/test/mjspecmacro_test.cc
@@ -0,0 +1,145 @@
+// Copyright 2026 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 X macros in mjspecmacro.h.
+
+#include <cstddef>
+#include <type_traits>
+
+#include <gtest/gtest.h>
+#include <mujoco/mjspec.h>
+#include <mujoco/mjspecmacro.h>
+
+namespace mujoco {
+namespace {
+
+namespace expected {
+#define X(type, name, dim) type name;
+#define XVEC(type, name, dim) type name[dim];
+
+struct mjsElement { MJSELEMENT_FIELDS };
+struct mjsCompiler { MJSCOMPILER_FIELDS };
+struct mjSpec { MJSPEC_FIELDS };
+struct mjsOrientation { MJSORIENTATION_FIELDS };
+struct mjsPlugin { MJSPLUGIN_FIELDS };
+struct mjsBody { MJSBODY_FIELDS };
+struct mjsFrame { MJSFRAME_FIELDS };
+struct mjsJoint { MJSJOINT_FIELDS };
+struct mjsGeom { MJSGEOM_FIELDS };
+struct mjsSite { MJSSITE_FIELDS };
+struct mjsCamera { MJSCAMERA_FIELDS };
+struct mjsLight { MJSLIGHT_FIELDS };
+struct mjsFlex { MJSFLEX_FIELDS };
+struct mjsMesh { MJSMESH_FIELDS };
+struct mjsHField { MJSHFIELD_FIELDS };
+struct mjsSkin { MJSSKIN_FIELDS };
+struct mjsTexture { MJSTEXTURE_FIELDS };
+struct mjsMaterial { MJSMATERIAL_FIELDS };
+struct mjsPair { MJSPAIR_FIELDS };
+struct mjsExclude { MJSEXCLUDE_FIELDS };
+struct mjsEquality { MJSEQUALITY_FIELDS };
+struct mjsTendon { MJSTENDON_FIELDS };
+struct mjsWrap { MJSWRAP_FIELDS };
+struct mjsActuator { MJSACTUATOR_FIELDS };
+struct mjsSensor { MJSSENSOR_FIELDS };
+struct mjsNumeric { MJSNUMERIC_FIELDS };
+struct mjsText { MJSTEXT_FIELDS };
+struct mjsTuple { MJSTUPLE_FIELDS };
+struct mjsKey { MJSKEY_FIELDS };
+struct mjsDefault { MJSDEFAULT_FIELDS };
+
+#undef XVEC
+#undef X
+}  // namespace expected
+
+template <typename T, size_t N>
+struct ArrayOrScalar {
+  using type = T[N];
+};
+template <typename T>
+struct ArrayOrScalar<T, 1> {
+  using type = T;
+};
+template <typename T, size_t N>
+using ArrayOrScalarT = typename ArrayOrScalar<T, N>::type;
+
+#define XIMPL(type, name, dim)                                                 \
+  static_assert(                                                               \
+      std::is_same_v<decltype(RealStruct::name), ArrayOrScalarT<type, dim>>,   \
+      "Type mismatch for " #name);                                             \
+  static_assert(offsetof(RealStruct, name) == offsetof(ExpectedStruct, name),  \
+                "Offset mismatch for " #name);
+
+#define X(type, name, dim)                                                     \
+  static_assert(dim == 1, "use XVEC for non-scalar fields");                   \
+  XIMPL(type, name, dim)
+
+#define XVEC(type, name, dim)                                                  \
+  static_assert(dim > 1, "use X for scalar fields");                           \
+  XIMPL(type, name, dim)
+
+#define CHECK_STRUCT(Name, Macro)                                              \
+  struct Check##Name {                                                         \
+    using RealStruct = Name;                                                   \
+    using ExpectedStruct = expected::Name;                                     \
+    static_assert(sizeof(RealStruct) == sizeof(ExpectedStruct),                \
+                  "sizeof mismatch for " #Name);                               \
+    static_assert(alignof(RealStruct) == alignof(ExpectedStruct),              \
+                  "alignof mismatch for " #Name);                              \
+    Macro                                                                      \
+  };
+
+CHECK_STRUCT(mjsElement, MJSELEMENT_FIELDS)
+CHECK_STRUCT(mjsCompiler, MJSCOMPILER_FIELDS)
+CHECK_STRUCT(mjSpec, MJSPEC_FIELDS)
+CHECK_STRUCT(mjsOrientation, MJSORIENTATION_FIELDS)
+CHECK_STRUCT(mjsPlugin, MJSPLUGIN_FIELDS)
+CHECK_STRUCT(mjsBody, MJSBODY_FIELDS)
+CHECK_STRUCT(mjsFrame, MJSFRAME_FIELDS)
+CHECK_STRUCT(mjsJoint, MJSJOINT_FIELDS)
+CHECK_STRUCT(mjsGeom, MJSGEOM_FIELDS)
+CHECK_STRUCT(mjsSite, MJSSITE_FIELDS)
+CHECK_STRUCT(mjsCamera, MJSCAMERA_FIELDS)
+CHECK_STRUCT(mjsLight, MJSLIGHT_FIELDS)
+CHECK_STRUCT(mjsFlex, MJSFLEX_FIELDS)
+CHECK_STRUCT(mjsMesh, MJSMESH_FIELDS)
+CHECK_STRUCT(mjsHField, MJSHFIELD_FIELDS)
+CHECK_STRUCT(mjsSkin, MJSSKIN_FIELDS)
+CHECK_STRUCT(mjsTexture, MJSTEXTURE_FIELDS)
+CHECK_STRUCT(mjsMaterial, MJSMATERIAL_FIELDS)
+CHECK_STRUCT(mjsPair, MJSPAIR_FIELDS)
+CHECK_STRUCT(mjsExclude, MJSEXCLUDE_FIELDS)
+CHECK_STRUCT(mjsEquality, MJSEQUALITY_FIELDS)
+CHECK_STRUCT(mjsTendon, MJSTENDON_FIELDS)
+CHECK_STRUCT(mjsWrap, MJSWRAP_FIELDS)
+CHECK_STRUCT(mjsActuator, MJSACTUATOR_FIELDS)
+CHECK_STRUCT(mjsSensor, MJSSENSOR_FIELDS)
+CHECK_STRUCT(mjsNumeric, MJSNUMERIC_FIELDS)
+CHECK_STRUCT(mjsText, MJSTEXT_FIELDS)
+CHECK_STRUCT(mjsTuple, MJSTUPLE_FIELDS)
+CHECK_STRUCT(mjsKey, MJSKEY_FIELDS)
+CHECK_STRUCT(mjsDefault, MJSDEFAULT_FIELDS)
+
+#undef CHECK_STRUCT
+#undef XVEC
+#undef X
+#undef XIMPL
+
+TEST(MjspecmacroTest, CompileTimeChecks) {
+  // Verifies that the test binary compiled successfully and all static_asserts passed.
+  EXPECT_TRUE(true);
+}
+
+}  // namespace
+}  // namespace mujoco
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 37aa73f1..0a6ba9b4 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -65,6 +65,7 @@ public const int mjMAXMATERIAL = 1000;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSAN_H_ = true;
 public const bool ADDRESS_SANITIZER = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSPEC_H_ = true;
+public const bool THIRD_PARTY_MUJOCO_MJSPECMACRO_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true;
 public const int mjMAXTHREAD = 128;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTYPE_H_ = true;

From f6c85287bf625b7991fdd9a3a7ade530e4b47074 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Wed, 27 May 2026 10:04:37 -0700
Subject: [PATCH 41/61] Optimize mj_flex performance.

Two performance optimizations for interpolated flex objects:

1. Hoist loop-invariant stride calculations in `mju_cellLookup` out of nested loops.
   This reduces multiplications from 16 to 6 (linear) and 54 to 12 (quadratic)
   per vertex.

2. Extract and consolidate optimized 3D interpolation logic into a new
   reusable utility `mju_evalBasisArray` in `engine_util_misc.c`. This function
   uses nested loops and precomputed 1D shape functions to avoid expensive
   dynamic `phi` calls and branching, and leverages stack-buffered outputs to
   eliminate compiler pointer-aliasing barriers. We propagate this optimization
   to both kinematics (`mju_interpolate3D`) and constraint setup
   (`engine_core_constraint.c`).

Together these changes yield a ~50% overall speedup in `mj_fwdKinematics` for interpolated flexes with ~10k vertices in the collision meshes and ~10 nodes in the deformation grid.

PiperOrigin-RevId: 922198093
Change-Id: I9c804e65cd532e2f9ac6e9a71d01c30cd64186fa
---
 src/engine/engine_core_constraint.c | 26 ++++++++---
 src/engine/engine_util_misc.c       | 70 ++++++++++++++++++++++++++---
 src/engine/engine_util_misc.h       |  3 ++
 3 files changed, 88 insertions(+), 11 deletions(-)

diff --git a/src/engine/engine_core_constraint.c b/src/engine/engine_core_constraint.c
index 2b7a1dbf..e7a521a0 100644
--- a/src/engine/engine_core_constraint.c
+++ b/src/engine/engine_core_constraint.c
@@ -290,13 +290,27 @@ static int mj_vertBodyWeight(const mjModel* m, const mjData* d, int f, int* v,
   int nstart = m->flex_nodeadr[f];
   int nb = 0;
 
-  for (int j = 0; j < npc; j++) {
-    mjtNum w = mju_evalBasis(local, j, order);
-    if (w < 1e-5) {
-      continue;
+  if (npc > 27) {
+    for (int j = 0; j < npc; j++) {
+      mjtNum w = mju_evalBasis(local, j, order);
+      if (w < 1e-5) {
+        continue;
+      }
+      if (bweight) bweight[nb] = sign * w;
+      body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]];
+    }
+  } else {
+    mjtNum basis[27];
+    mju_evalBasisArray(basis, local, order);
+
+    for (int j = 0; j < npc; j++) {
+      mjtNum w = basis[j];
+      if (w < 1e-5) {
+        continue;
+      }
+      if (bweight) bweight[nb] = sign * w;
+      body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]];
     }
-    if (bweight) bweight[nb] = sign * w;
-    body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]];
   }
 
   return nb;
diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c
index eada6949..65b0d6c2 100644
--- a/src/engine/engine_util_misc.c
+++ b/src/engine/engine_util_misc.c
@@ -574,6 +574,49 @@ mjtNum mju_evalBasis(const mjtNum x[3], int i, int order) {
   }
 }
 
+// evaluate the basis functions at x for all nodes in the cell
+void mju_evalBasisArray(mjtNum* basis, const mjtNum x[3], int order) {
+  if (order == 1) {
+    mjtNum p[3][2] = {
+      {1 - x[0], x[0]},
+      {1 - x[1], x[1]},
+      {1 - x[2], x[2]}
+    };
+    int j = 0;
+    for (int i0=0; i0<2; i0++) {
+      mjtNum w0 = p[0][i0];
+      for (int i1=0; i1<2; i1++) {
+        mjtNum w01 = w0 * p[1][i1];
+        for (int i2=0; i2<2; i2++) {
+          basis[j++] = w01 * p[2][i2];
+        }
+      }
+    }
+  } else if (order == 2) {
+    mjtNum p[3][3];
+    for (int d=0; d<3; d++) {
+      for (int i=0; i<3; i++) {
+        p[d][i] = phi(x[d], i, 2);
+      }
+    }
+    int j = 0;
+    for (int i0=0; i0<3; i0++) {
+      mjtNum w0 = p[0][i0];
+      for (int i1=0; i1<3; i1++) {
+        mjtNum w01 = w0 * p[1][i1];
+        for (int i2=0; i2<3; i2++) {
+          basis[j++] = w01 * p[2][i2];
+        }
+      }
+    }
+  } else {
+    int npoint = (order + 1) * (order + 1) * (order + 1);
+    for (int j=0; j < npoint; j++) {
+      basis[j] = mju_evalBasis(x, j, order);
+    }
+  }
+}
+
 // map global parametric coord to cell-local coord and build node indices
 //   coord: [0,1]^3 parametric coordinates
 //   cellnum: cell counts (cx, cy, cz)
@@ -600,16 +643,21 @@ int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNu
 
   // build node indices for this cell
   if (nodeindices) {
+    int gi_base = ci * order;
+    int gj_base = cj * order;
+    int gk_base = ck * order;
     int ny_g = cy * order + 1;
     int nz_g = cz * order + 1;
     int ni = 0;
     for (int li = 0; li <= order; li++) {
+      int gi = gi_base + li;
+      int gi_stride = gi * ny_g * nz_g;
       for (int lj = 0; lj <= order; lj++) {
+        int gj = gj_base + lj;
+        int gj_stride = gi_stride + gj * nz_g;
         for (int lk = 0; lk <= order; lk++) {
-          int gi = ci*order + li;
-          int gj = cj*order + lj;
-          int gk = ck*order + lk;
-          nodeindices[ni++] = gi*ny_g*nz_g + gj*nz_g + gk;
+          int gk = gk_base + lk;
+          nodeindices[ni++] = gj_stride + gk;
         }
       }
     }
@@ -624,9 +672,21 @@ int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNu
 void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order,
                        const int* nodeindices) {
   int npoint = (order + 1) * (order + 1) * (order + 1);
+
+  if (npoint > 27) {
+    for (int j=0; j < npoint; j++) {
+      int idx = nodeindices ? nodeindices[j] : j;
+      mju_addToScl3(res, coeff+3*idx, mju_evalBasis(x, j, order));
+    }
+    return;
+  }
+
+  mjtNum basis[27];
+  mju_evalBasisArray(basis, x, order);
+
   for (int j=0; j < npoint; j++) {
     int idx = nodeindices ? nodeindices[j] : j;
-    mju_addToScl3(res, coeff+3*idx, mju_evalBasis(x, j, order));
+    mju_addToScl3(res, coeff+3*idx, basis[j]);
   }
 }
 
diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h
index f97bf4bb..42aa3704 100644
--- a/src/engine/engine_util_misc.h
+++ b/src/engine/engine_util_misc.h
@@ -89,6 +89,9 @@ MJAPI void mju_defGradient(mjtNum res[9], const mjtNum p[3], const mjtNum* dof,
 // evaluate the basis function at x for the i-th node
 MJAPI mjtNum mju_evalBasis(const mjtNum x[3], int i, int order);
 
+// evaluate the basis functions at x for all nodes in the cell
+MJAPI void mju_evalBasisArray(mjtNum* basis, const mjtNum x[3], int order);
+
 // map global parametric coord to cell-local coord and build node indices
 MJAPI int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3],
                          int* nodeindices);

From 4035c66428817c9b1e857cc77482127c5ad6a98d Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Wed, 27 May 2026 11:31:53 -0700
Subject: [PATCH 42/61] Fix multiccd edge bug that occurs in certain
 situations.

PiperOrigin-RevId: 922251109
Change-Id: I74c08ae49d5dfd2a41b2dcd28eb9f2b5ba113216
---
 src/engine/engine_collision_gjk.c        | 10 ++++++-
 test/engine/engine_collision_gjk_test.cc | 38 ++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/engine/engine_collision_gjk.c b/src/engine/engine_collision_gjk.c
index 42227694..8ad8d782 100644
--- a/src/engine/engine_collision_gjk.c
+++ b/src/engine/engine_collision_gjk.c
@@ -2169,8 +2169,16 @@ static void multicontact(Polytope* pt, Face* face, mjCCDStatus* status,
 
   // face1 is an edge; clip face1 against face2
   if (edgecon1) {
-    scl3(approx_dir, n2 + 3*j, norm3(dir));
+    scl3(approx_dir, n2 + 3*j, -norm3(dir));
     polygonClip(status, face2, nface2, face1, nface1, n2 + 3*j, approx_dir);
+    // x1 and x2 must be flipped as we flipped the faces in polygonClip
+    int nx = status->nx;
+    for (int k = 0; k < nx; k++) {
+      mjtNum tmp[3];
+      copy3(tmp, status->x1 + 3*k);
+      copy3(status->x1 + 3*k, status->x2 + 3*k);
+      copy3(status->x2 + 3*k, tmp);
+    }
     return;
   }
 
diff --git a/test/engine/engine_collision_gjk_test.cc b/test/engine/engine_collision_gjk_test.cc
index 21fec2cf..f46d0c0f 100644
--- a/test/engine/engine_collision_gjk_test.cc
+++ b/test/engine/engine_collision_gjk_test.cc
@@ -2028,5 +2028,43 @@ TEST_F(MjGjkTest, CylinderBoxMargin) {
   mj_deleteModel(model);
 }
 
+TEST_F(MjGjkTest, BoxEdgeFlipped) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <geom name="geom1" pos="1.10164554 -0.11389316 0.74"
+            quat="-0.348312918 0 0 0.937378318" type="box" size="0.65 0.48 0.04"/>
+      <geom name="geom2" type="box" size="0.1 1.2 1.4" pos="1.4 0 1.425"/>
+    </worldbody>
+  </mujoco>)";
+
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << "Failed to load model: " << error;
+
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  int g1 = mj_name2id(model, mjOBJ_GEOM, "geom1");
+  int g2 = mj_name2id(model, mjOBJ_GEOM, "geom2");
+
+  mjCCDStatus status;
+  std::vector<mjtNum> dir, pos;
+  mjtNum dist;
+  int ncons = Penetration(status, dist, dir, pos, model, data, g1, g2, 0, 1000);
+
+  EXPECT_EQ(ncons, 2);
+
+  EXPECT_NEAR(status.x1[0], 1.907368, kTolerance);
+  EXPECT_NEAR(status.x1[1], -0.052973, kTolerance);
+  EXPECT_NEAR(status.x1[2], 0.700000, kTolerance);
+  EXPECT_NEAR(status.x2[0], 1.30000, kTolerance);
+  EXPECT_NEAR(status.x2[1], -0.052973, kTolerance);
+  EXPECT_NEAR(status.x2[2], 0.700000, kTolerance);
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco

From aa4a43da1732d73c6fb89ae60893fac5ece72fef Mon Sep 17 00:00:00 2001
From: Kevin Zakka <kevinarmandzakka@gmail.com>
Date: Wed, 27 May 2026 12:55:14 -0700
Subject: [PATCH 43/61] Address reviewer feedback: default x_scale=None, drop
 TRF qualifier.

---
 python/mujoco/minimize.py | 23 ++++++++++++-----------
 1 file changed, 12 insertions(+), 11 deletions(-)

diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py
index fb829398..520fef11 100644
--- a/python/mujoco/minimize.py
+++ b/python/mujoco/minimize.py
@@ -156,7 +156,7 @@ def least_squares(
     output: Optional[TextIO] = None,
     iter_callback: Optional[Callable[[List[IterLog]], None]] = None,
     check_derivatives: bool = False,
-    x_scale: Union[float, np.ndarray, str] = 1.0,
+    x_scale: Optional[Union[float, np.ndarray, str]] = None,
 ) -> Tuple[np.ndarray, List[IterLog]]:
   """Nonlinear Least Squares minimization with box bounds.
 
@@ -179,13 +179,12 @@ def least_squares(
     output: Optional file or StringIO to which to print messages.
     iter_callback: Optional iteration callback, takes trace argument.
     check_derivatives: Compare user-defined Jacobian and norm against fin-diff.
-    x_scale: Per-parameter scaling. Setting ``x_scale=D`` solves the problem in
-      the change of variables ``z = x / D`` and un-scales the result. ``1.0``
-      (default) is no scaling. ``'jac'`` sets ``D_i = 1 / ||J(:, i)||`` at each
-      iteration (matches scipy's ``least_squares(method='trf', x_scale='jac')``).
-      An array of shape ``(n,)`` (or a positive scalar) is used as ``D``
-      directly. Note that ``mu``, ``mu_min`` and ``mu_max`` then act on the
-      scaled subproblem.
+    x_scale: Per-parameter scaling (matches scipy's ``x_scale``). Setting
+      ``x_scale=D`` solves the problem in the change of variables ``z = x / D``
+      and un-scales the result. ``None`` (default) or ``1.0`` is no scaling.
+      ``'jac'`` sets ``D_i = 1 / ||J(:, i)||`` at each iteration. An array of
+      shape ``(n,)`` (or a positive scalar) is used as ``D`` directly. Note
+      that ``mu``, ``mu_min`` and ``mu_max`` then act on the scaled subproblem.
 
   Returns:
     x: best solution found
@@ -211,13 +210,15 @@ def least_squares(
   mu = np.float64(0.0)  # Optimistically start with no regularization.
   n_reduc = 0  # Number of sequential mu reductions.
 
-  # Resolve x_scale -> D of shape (n, 1). For 'jac', D is refreshed each iter.
+  # Resolve x_scale -> D of shape (n, 1). For 'jac', D is refreshed each iter;
+  # None means no scaling (D = 1).
   adaptive_scale = isinstance(x_scale, str)
   if adaptive_scale and x_scale != 'jac':
     raise ValueError(
-        f"x_scale must be 'jac', a positive scalar, or array, got {x_scale!r}."
+        f"x_scale must be None, 'jac', a positive scalar, or array, got "
+        f'{x_scale!r}.'
     )
-  if adaptive_scale:
+  if x_scale is None or adaptive_scale:
     D = np.ones((n, 1))
   else:
     D = np.asarray(x_scale, dtype=np.float64)

From 106f1c5dbbbf1d4fc39dc39d6c590d7a919b9d15 Mon Sep 17 00:00:00 2001
From: Tarik Kelestemur <tkelestemur@gmail.com>
Date: Wed, 27 May 2026 17:50:31 -0400
Subject: [PATCH 44/61] fix mjx codegen and rerun

---
 mjx/mujoco/mjx/codegen/generate_warp_shim.py  |  7 ++++---
 mjx/mujoco/mjx/codegen/generate_warp_types.py |  9 ++++++---
 mjx/mujoco/mjx/warp/bvh.py                    | 11 ++++++----
 mjx/mujoco/mjx/warp/collision_driver.py       | 11 ++++++----
 mjx/mujoco/mjx/warp/forward.py                | 10 ++++++----
 mjx/mujoco/mjx/warp/render.py                 | 11 ++++++----
 mjx/mujoco/mjx/warp/smooth.py                 | 11 ++++++----
 mjx/mujoco/mjx/warp/types.py                  | 20 ++++++++++++++++++-
 8 files changed, 63 insertions(+), 27 deletions(-)

diff --git a/mjx/mujoco/mjx/codegen/generate_warp_shim.py b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
index 4c163153..2596c725 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_shim.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_shim.py
@@ -215,7 +215,7 @@ def _warp_function(
       info = mjwarp_field_info[f]
       expected_type = _clean_type(info.expected_type)
       fn_args_model.append((f'{f}: {expected_type},', info.param_order))
-      fn_assignments.append(f'  _m.{f.replace('__', '.')} = {f}')
+      fn_assignments.append(f"  _m.{f.replace('__', '.')} = {f}")
     fn_args_model = sorted(fn_args_model, key=lambda x: x[1])
     fn_args_model = ['# Model'] + [f[0] for f in fn_args_model]
 
@@ -233,7 +233,7 @@ def _warp_function(
       param_order = info.param_order
       expected_type = _clean_type(info.expected_type)
       fn_args_data.append((f'{f}: {expected_type},', (is_out, param_order)))
-      fn_assignments.append(f'  _d.{f.replace('__', '.')} = {f}')
+      fn_assignments.append(f"  _d.{f.replace('__', '.')} = {f}")
     fn_args_data = sorted(fn_args_data, key=lambda x: x[1])
     fn_args_data = ['# Data'] + [f[0] for f in fn_args_data]
 
@@ -418,6 +418,7 @@ def create_jax_warp_shim(
   fn_args_raw_str = '\n'.join(['    ' + arg for arg in fn_args_raw])
   warp_fn_args = [arg.split(':')[0] for arg in fn_args_raw if '#' not in arg]  # pytype: disable=attribute-error
 
+  fn_assignments_str = '\n'.join(fn_assignments)
   src += f"""
 @ffi.format_args_for_warp
 def _{fn_name}_shim(
@@ -428,7 +429,7 @@ def _{fn_name}_shim(
   _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
-{'\n'.join(fn_assignments)}
+{fn_assignments_str}
   {fn_call}
   """
   src += '\n\n'
diff --git a/mjx/mujoco/mjx/codegen/generate_warp_types.py b/mjx/mujoco/mjx/codegen/generate_warp_types.py
index 374e7464..466ee281 100644
--- a/mjx/mujoco/mjx/codegen/generate_warp_types.py
+++ b/mjx/mujoco/mjx/codegen/generate_warp_types.py
@@ -67,7 +67,8 @@ def _to_py_string(value, indent=0):
   next_indent_str = '  ' * (indent + 1)
   if isinstance(value, tuple):
     items = [_to_py_string(item, indent) for item in value]
-    return f'({', '.join(items)})'
+    joined_items = ', '.join(items)
+    return f'({joined_items})'
 
   if isinstance(value, type):
     if value.__module__ == 'builtins':
@@ -79,12 +80,14 @@ def _to_py_string(value, indent=0):
         f'\n{next_indent_str}{repr(k)}: {_to_py_string(v, indent + 1)}'
         for k, v in sorted(value.items(), key=lambda x: x[0])
     ]
-    return f'{{{','.join(items)}\n{indent_str}}}'
+    joined_items = ','.join(items)
+    return f'{{{joined_items}\n{indent_str}}}'
 
   if isinstance(value, set):
     items = sorted([_to_py_string(item, indent) for item in value])
     items = [f'\n{next_indent_str}{item}' for item in items]
-    return f'{{{",".join(items)}\n{indent_str}}}'
+    joined_items = ','.join(items)
+    return f'{{{joined_items}\n{indent_str}}}'
 
   return repr(value)
 
diff --git a/mjx/mujoco/mjx/warp/bvh.py b/mjx/mujoco/mjx/warp/bvh.py
index 7f648694..1301aa88 100644
--- a/mjx/mujoco/mjx/warp/bvh.py
+++ b/mjx/mujoco/mjx/warp/bvh.py
@@ -14,17 +14,19 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
+import warp as wp
+
 from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 from mujoco.mjx.warp import ffi
 from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.render_context import RenderContextPytree
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
-import warp as wp
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _refit_bvh_shim(
     # Model
diff --git a/mjx/mujoco/mjx/warp/collision_driver.py b/mjx/mujoco/mjx/warp/collision_driver.py
index 47b94ed2..e3d6b595 100644
--- a/mjx/mujoco/mjx/warp/collision_driver.py
+++ b/mjx/mujoco/mjx/warp/collision_driver.py
@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+from mujoco.mjx.warp import ffi
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _collision_shim(
     # Model
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index 0adf9236..52a84999 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+from mujoco.mjx.warp import ffi
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
diff --git a/mjx/mujoco/mjx/warp/render.py b/mjx/mujoco/mjx/warp/render.py
index cf6754a1..3fb5a2ea 100644
--- a/mjx/mujoco/mjx/warp/render.py
+++ b/mjx/mujoco/mjx/warp/render.py
@@ -14,17 +14,19 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
+import warp as wp
+
 from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 from mujoco.mjx.warp import ffi
 from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
 from mujoco.mjx.warp.render_context import RenderContextPytree
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
-import warp as wp
-
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _render_shim(
     # Model
diff --git a/mjx/mujoco/mjx/warp/smooth.py b/mjx/mujoco/mjx/warp/smooth.py
index ce765d8b..862b6d07 100644
--- a/mjx/mujoco/mjx/warp/smooth.py
+++ b/mjx/mujoco/mjx/warp/smooth.py
@@ -14,15 +14,17 @@
 # ==============================================================================
 
 """DO NOT EDIT. This file is auto-generated."""
+
 import dataclasses
 import functools
+
 import jax
-from mujoco.mjx._src import types
-from mujoco.mjx.warp import ffi
-import mujoco.mjx.third_party.mujoco_warp as mjwarp
-from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 import warp as wp
 
+from mujoco.mjx._src import types
+import mujoco.mjx.third_party.mujoco_warp as mjwarp
+from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+from mujoco.mjx.warp import ffi
 
 _m = mjwarp.Model(
     **{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
     **{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
 )
 
+
 @ffi.format_args_for_warp
 def _kinematics_shim(
     # Model
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index f64b6fba..df4098d2 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -15,14 +15,18 @@
 """MJX Warp types.
 DO NOT EDIT. This file is auto-generated.
 """
+
 import dataclasses
 import typing
 from typing import Tuple
+
 import jax
 from jax import tree_util
 from jax.interpreters import batching
-from mujoco.mjx._src import dataclasses as mjx_dataclasses
 import numpy as np
+
+from mujoco.mjx._src import dataclasses as mjx_dataclasses
+
 if typing.TYPE_CHECKING:
   GraphMode = int
 
@@ -33,6 +37,7 @@ if typing.TYPE_CHECKING:
 else:
   try:
     from warp._src.jax_experimental.ffi import GraphMode
+
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
 
     Callback = mjwp_types.Callback
@@ -41,6 +46,7 @@ else:
     Callback = None
 PyTreeNode = mjx_dataclasses.PyTreeNode
 
+
 @dataclasses.dataclass(frozen=True)
 @tree_util.register_pytree_node_class
 class TileSet:
@@ -52,6 +58,7 @@ class TileSet:
     adr: address of each tile in the set
     size: size of all the tiles in this set
   """
+
   adr: np.ndarray
   size: int
 
@@ -72,6 +79,7 @@ class BlockDim:
 
   TODO(team): experimental and may be removed
   """
+
   actuator_velocity: int
   cholesky_factorize: int
   cholesky_factorize_solve: int
@@ -102,10 +110,13 @@ class BlockDim:
 
 class StatisticWarp(PyTreeNode):
   """Derived fields from Statistic."""
+
   meaninertia: jax.Array
 
+
 class OptionWarp(PyTreeNode):
   """Derived fields from Option."""
+
   broadphase: int
   broadphase_filter: int
   ccd_iterations: int
@@ -120,8 +131,10 @@ class OptionWarp(PyTreeNode):
   sdf_initpoints: int
   sdf_iterations: int
 
+
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
+
   M_colind: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
@@ -283,8 +296,10 @@ class ModelWarp(PyTreeNode):
   wrap_site_adr: np.ndarray
   wrap_site_pair_adr: np.ndarray
 
+
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
+
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -363,6 +378,8 @@ class DataWarp(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   shape = property(lambda self: self.cacc.shape)
+
+
 DATA_NON_VMAP = {
     'contact__dim',
     'contact__dist',
@@ -390,6 +407,7 @@ DATA_NON_VMAP = {
     'nworld',
 }
 
+
 def _to_elt(cont, _, d, axis):
   return DataWarp(**{
       f.name: (

From b9fb817af870e75fb5e6bba9ba4bec767c2c0431 Mon Sep 17 00:00:00 2001
From: Sam Haves <shaves@google.com>
Date: Thu, 28 May 2026 02:56:01 -0700
Subject: [PATCH 45/61] Lexically normalize resource path in mjz_decoder
 resource provider.

PiperOrigin-RevId: 922654438
Change-Id: Ieb46230a2aba46a5e711bc5b24aec982f71e56ca
---
 doc/changelog.rst          |  8 ++++++++
 src/xml/mjz/mjz_decoder.cc | 13 ++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index a2243f40..b07ba2b0 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+Bug fixes
+^^^^^^^^^
+- Fixed a bug in the ``mjz`` :ref:`decoder <mjpDecoder>` where unnormalized paths would fail to be read.
+
+
 Version 3.9.0 (May 27, 2026)
 ----------------------------
 
diff --git a/src/xml/mjz/mjz_decoder.cc b/src/xml/mjz/mjz_decoder.cc
index f2864518..f2bd64bb 100644
--- a/src/xml/mjz/mjz_decoder.cc
+++ b/src/xml/mjz/mjz_decoder.cc
@@ -90,7 +90,7 @@ class ZipArchiveProvider : public mjpResourceProvider {
     // with the same name as the archive itself. Failing that, look for an XML
     // file within a subdirectory with the same name as the archive.
     const std::filesystem::path path(name_);
-    root_model_ = (path / path.stem()).string() + ".xml";
+    root_model_ = (path / path.stem()).generic_string() + ".xml";
     if (!Contains(root_model_)) {
       root_model_ = (path / path.stem() / path.stem()).string() + ".xml";
       if (!Contains(root_model_)) {
@@ -110,12 +110,14 @@ class ZipArchiveProvider : public mjpResourceProvider {
     };
     open = [](mjResource* resource) {
       ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
-      const bool found = self->Contains(resource->name);
+      auto path = std::filesystem::path(resource->name).lexically_normal();
+      const bool found = self->Contains(path.string());
       return found ? 1 : 0;
     };
     read = [](mjResource* resource, const void** buffer) {
       ZipArchiveProvider* self = (ZipArchiveProvider*)resource->provider;
-      std::span<char> bytes = self->Read(resource->name);
+      auto path = std::filesystem::path(resource->name).lexically_normal();
+      std::span<char> bytes = self->Read(path.string());
       *buffer = bytes.data();
       return static_cast<int>(bytes.size());
     };
@@ -138,6 +140,11 @@ class ZipArchiveProvider : public mjpResourceProvider {
 
   // Returns true if the archive contains a file with the given name/path.
   bool Contains(std::string_view name) const {
+    // Lexically normalized path might strip the archive path.
+    // a/b.mjz/../../../c.xml -> ../c.xml
+    if (!name.starts_with(name_)) {
+      return false;
+    }
     const std::string_view filename = name.substr(name_.size() + 1);
     return files_.find(filename.data()) != files_.end();
   }

From 4c381635e118db803038cec3b48846b36165d638 Mon Sep 17 00:00:00 2001
From: Alessio Quaglino <quaglino@google.com>
Date: Thu, 28 May 2026 05:58:47 -0700
Subject: [PATCH 46/61] Add mjs_makeFlex to the MuJoCo C and Python APIs.

- Add mjs_makeFlex C API that wraps mjCFlexcomp::Make(), creating
  a flex with auto-generated bodies, joints, and equality constraints
  from parameters (type, dim, dof, count, cellcount, spacing, scale,
  radius, mass, inertiabox, equality, rigid, flatskin, elastic2d,
  pos, quat, origin, file, vfs).
- Add Python binding body.make_flex() in specs.cc with keyword args.
- Refactor flex_spec_provision.py to use make_flex() instead of XML
  string templating + separate spec + attach. Eliminates textwrap
  dependency.
- Merge user_flexcomp into user_model BUILD target to resolve
  circular dependency.

PiperOrigin-RevId: 922725304
Change-Id: Ice13787c45f3173407fec659f9590a6453bc65dd
---
 doc/APIreference/functions.rst        |  11 ++
 doc/changelog.rst                     |   7 +-
 doc/includes/references.h             |   6 ++
 include/mujoco/mujoco.h               |   9 ++
 python/mujoco/introspect/functions.py | 138 ++++++++++++++++++++++++++
 python/mujoco/specs.cc                | 124 +++++++++++++++++++++++
 python/mujoco/specs_test.py           |  70 +++++++++++++
 src/user/user_api.cc                  | 116 ++++++++++++++++++++++
 src/user/user_api.h                   |   8 ++
 wasm/codegen/generated/bindings.cc    |  17 ++++
 wasm/codegen/generators/functions.py  |   2 +-
 11 files changed, 506 insertions(+), 2 deletions(-)

diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index 298e5134..4daaaf44 100644
--- a/doc/APIreference/functions.rst
+++ b/doc/APIreference/functions.rst
@@ -4496,6 +4496,17 @@ Add sensor.
 
 Add flex.
 
+.. _mjs_makeFlex:
+
+`mjs_makeFlex <#mjs_makeFlex>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mjs_makeFlex
+
+Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.
+
+*Nullable:* ``type``, ``dof``, ``count``, ``cellcount``, ``spacing``, ``scale``, ``pos``, ``quat``, ``origin``, ``file``, ``vfs``
+
 .. _mjs_addPair:
 
 `mjs_addPair <#mjs_addPair>`__
diff --git a/doc/changelog.rst b/doc/changelog.rst
index b07ba2b0..699d0dd0 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -5,11 +5,16 @@ Changelog
 Upcoming version (not yet released)
 -----------------------------------
 
+General
+^^^^^^^
+- Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp<body-flexcomp>` element for
+  programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as
+  ``body.make_flex()`` in Python.
+
 Bug fixes
 ^^^^^^^^^
 - Fixed a bug in the ``mjz`` :ref:`decoder <mjpDecoder>` where unnormalized paths would fail to be read.
 
-
 Version 3.9.0 (May 27, 2026)
 ----------------------------
 
diff --git a/doc/includes/references.h b/doc/includes/references.h
index b158c3e1..60b8d699 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -3678,6 +3678,12 @@ int mjs_delete(mjSpec* spec, mjsElement* element);
 mjsActuator* mjs_addActuator(mjSpec* s, const mjsDefault* def);
 mjsSensor* mjs_addSensor(mjSpec* s);
 mjsFlex* mjs_addFlex(mjSpec* s);
+mjsFlex* mjs_makeFlex(mjsBody* body, const char* name, const char* type, int dim,
+                      const char* dof, const int count[3], const int cellcount[3],
+                      const double spacing[3], const double scale[3], double radius,
+                      double mass, double inertiabox, int equality, int rigid, int flatskin,
+                      int elastic2d, const double pos[3], const double quat[4],
+                      const double origin[3], const char* file, const mjVFS* vfs);
 mjsPair* mjs_addPair(mjSpec* s, const mjsDefault* def);
 mjsExclude* mjs_addExclude(mjSpec* s);
 mjsEquality* mjs_addEquality(mjSpec* s, const mjsDefault* def);
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index ee3ab87d..4a1384a7 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -1667,6 +1667,15 @@ MJAPI mjsSensor* mjs_addSensor(mjSpec* s);
 // Add flex.
 MJAPI mjsFlex* mjs_addFlex(mjSpec* s);
 
+// Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.
+// Nullable: type, dof, count, cellcount, spacing, scale, pos, quat, origin, file, vfs
+MJAPI mjsFlex* mjs_makeFlex(mjsBody* body, const char* name, const char* type, int dim,
+                            const char* dof, const int count[3], const int cellcount[3],
+                            const double spacing[3], const double scale[3], double radius,
+                            double mass, double inertiabox, int equality, int rigid, int flatskin,
+                            int elastic2d, const double pos[3], const double quat[4],
+                            const double origin[3], const char* file, const mjVFS* vfs);
+
 // Add contact pair.
 // Nullable: def
 MJAPI mjsPair* mjs_addPair(mjSpec* s, const mjsDefault* def);
diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py
index c0bce30f..d2323d9e 100644
--- a/python/mujoco/introspect/functions.py
+++ b/python/mujoco/introspect/functions.py
@@ -10366,6 +10366,144 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Add flex.',
      )),
+    ('mjs_makeFlex',
+     FunctionDecl(
+         name='mjs_makeFlex',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjsFlex'),
+         ),
+         parameters=(
+             FunctionParameterDecl(
+                 name='body',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjsBody'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='name',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='type',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='dim',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='dof',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='count',
+                 type=ArrayType(
+                     inner_type=ValueType(name='int', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='cellcount',
+                 type=ArrayType(
+                     inner_type=ValueType(name='int', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='spacing',
+                 type=ArrayType(
+                     inner_type=ValueType(name='double', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='scale',
+                 type=ArrayType(
+                     inner_type=ValueType(name='double', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='radius',
+                 type=ValueType(name='double'),
+             ),
+             FunctionParameterDecl(
+                 name='mass',
+                 type=ValueType(name='double'),
+             ),
+             FunctionParameterDecl(
+                 name='inertiabox',
+                 type=ValueType(name='double'),
+             ),
+             FunctionParameterDecl(
+                 name='equality',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='rigid',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='flatskin',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='elastic2d',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='pos',
+                 type=ArrayType(
+                     inner_type=ValueType(name='double', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='quat',
+                 type=ArrayType(
+                     inner_type=ValueType(name='double', is_const=True),
+                     extents=(4,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='origin',
+                 type=ArrayType(
+                     inner_type=ValueType(name='double', is_const=True),
+                     extents=(3,),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='file',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+                 nullable=True,
+             ),
+             FunctionParameterDecl(
+                 name='vfs',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjVFS', is_const=True),
+                 ),
+                 nullable=True,
+             ),
+         ),
+         doc='Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.',  # pylint: disable=line-too-long
+     )),
     ('mjs_addPair',
      FunctionDecl(
          name='mjs_addPair',
diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index abe101d9..738ab73b 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -842,6 +842,130 @@ PYBIND11_MODULE(_specs, m) {
         return FindAllImpl(self, objtype, true);
       },
       py::return_value_policy::reference_internal);
+  mjsBody.def(
+      "make_flex",
+      [](raw::MjsBody& self,
+         const std::string& name,
+         std::optional<std::string> type,
+         int dim,
+         std::optional<std::string> dof,
+         std::optional<std::vector<int>> count,
+         std::optional<std::vector<int>> cellcount,
+         std::optional<std::vector<double>> spacing,
+         std::optional<std::vector<double>> scale,
+         double radius,
+         double mass,
+         double inertiabox,
+         int equality,
+         int rigid,
+         int flatskin,
+         int elastic2d,
+         std::optional<std::vector<double>> pos,
+         std::optional<std::vector<double>> quat,
+         std::optional<std::vector<double>> origin,
+         std::optional<std::string> file,
+         MjVfs* vfs) -> raw::MjsFlex* {
+        const char* type_str = type.has_value() ? type->c_str() : nullptr;
+        const char* dof_str = dof.has_value() ? dof->c_str() : nullptr;
+        const char* file_str = file.has_value() ? file->c_str() : nullptr;
+        const mjVFS* vfs_ptr = vfs ? vfs->get() : nullptr;
+
+        int count_arr[3] = {10, 10, 10};
+        if (count.has_value()) {
+          if (count->size() != 3) {
+            throw pybind11::value_error("count must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) count_arr[i] = (*count)[i];
+        }
+        const int* count_ptr = count.has_value() ? count_arr : nullptr;
+
+        int cellcount_arr[3] = {-1, -1, -1};
+        if (cellcount.has_value()) {
+          if (cellcount->size() != 3) {
+            throw pybind11::value_error("cellcount must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) cellcount_arr[i] = (*cellcount)[i];
+        }
+        const int* cellcount_ptr =
+            cellcount.has_value() ? cellcount_arr : nullptr;
+
+        double spacing_arr[3] = {0.02, 0.02, 0.02};
+        if (spacing.has_value()) {
+          if (spacing->size() != 3) {
+            throw pybind11::value_error("spacing must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) spacing_arr[i] = (*spacing)[i];
+        }
+        const double* spacing_ptr = spacing.has_value() ? spacing_arr : nullptr;
+
+        double scale_arr[3] = {1, 1, 1};
+        if (scale.has_value()) {
+          if (scale->size() != 3) {
+            throw pybind11::value_error("scale must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) scale_arr[i] = (*scale)[i];
+        }
+        const double* scale_ptr = scale.has_value() ? scale_arr : nullptr;
+
+        double pos_arr[3] = {0, 0, 0};
+        if (pos.has_value()) {
+          if (pos->size() != 3) {
+            throw pybind11::value_error("pos must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) pos_arr[i] = (*pos)[i];
+        }
+        const double* pos_ptr = pos.has_value() ? pos_arr : nullptr;
+
+        double quat_arr[4] = {1, 0, 0, 0};
+        if (quat.has_value()) {
+          if (quat->size() != 4) {
+            throw pybind11::value_error("quat must have 4 elements");
+          }
+          for (int i = 0; i < 4; i++) quat_arr[i] = (*quat)[i];
+        }
+        const double* quat_ptr = quat.has_value() ? quat_arr : nullptr;
+
+        double origin_arr[3] = {0, 0, 0};
+        if (origin.has_value()) {
+          if (origin->size() != 3) {
+            throw pybind11::value_error("origin must have 3 elements");
+          }
+          for (int i = 0; i < 3; i++) origin_arr[i] = (*origin)[i];
+        }
+        const double* origin_ptr = origin.has_value() ? origin_arr : nullptr;
+
+        auto out = mjs_makeFlex(
+            &self, name.c_str(), type_str, dim, dof_str,
+            count_ptr, cellcount_ptr, spacing_ptr, scale_ptr,
+            radius, mass, inertiabox, equality, rigid, flatskin, elastic2d,
+            pos_ptr, quat_ptr, origin_ptr, file_str, vfs_ptr);
+        if (!out) {
+          raw::MjSpec* spec = mjs_getSpec(self.element);
+          throw pybind11::value_error(mjs_getError(spec));
+        }
+        return out;
+      },
+      py::arg("name"),
+      py::arg("type") = py::none(),
+      py::arg("dim") = 3,
+      py::arg("dof") = py::none(),
+      py::arg("count") = py::none(),
+      py::arg("cellcount") = py::none(),
+      py::arg("spacing") = py::none(),
+      py::arg("scale") = py::none(),
+      py::arg("radius") = 0.0,
+      py::arg("mass") = 1.0,
+      py::arg("inertiabox") = 0.005,
+      py::arg("equality") = 0,
+      py::arg("rigid") = 0,
+      py::arg("flatskin") = 0,
+      py::arg("elastic2d") = 0,
+      py::arg("pos") = py::none(),
+      py::arg("quat") = py::none(),
+      py::arg("origin") = py::none(),
+      py::arg("file") = py::none(),
+      py::arg("vfs") = py::none(),
+      py::return_value_policy::reference_internal);
   mjsBody.def(
       "find_child",
       [](raw::MjsBody& self, std::string& name) -> raw::MjsBody* {
diff --git a/python/mujoco/specs_test.py b/python/mujoco/specs_test.py
index 7e1baa04..8bae2be1 100644
--- a/python/mujoco/specs_test.py
+++ b/python/mujoco/specs_test.py
@@ -2032,5 +2032,75 @@ class SpecsTest(absltest.TestCase):
     with self.assertRaises(mujoco.FatalError):
       spec.encode(filename, model)
 
+  def test_make_flex_grid(self):
+    # Create a spec with a flexcomp grid.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body(name='flex_body')
+    flex = body.make_flex(
+        name='test_flex',
+        type='grid',
+        dim=3,
+        count=[4, 4, 4],
+        spacing=[0.05, 0.05, 0.05],
+        mass=0.5,
+        equality=1,
+    )
+    self.assertIsNotNone(flex)
+    model = spec.compile()
+    self.assertIsNotNone(model)
+    self.assertGreater(model.nflex, 0)
+
+    # Verify elastic2d is forwarded to Make() and produces shell-mode
+    # strain constraints (equality=3 + elastic2d=2 triggers shell path).
+    spec2 = mujoco.MjSpec()
+    body2 = spec2.worldbody.add_body(name='shell_body')
+    flex2 = body2.make_flex(
+        name='shell_flex',
+        type='grid',
+        dim=3,
+        count=[3, 3, 3],
+        spacing=[0.1, 0.1, 0.1],
+        dof='trilinear',
+        cellcount=[2, 2, 1],
+        mass=0.5,
+        equality=3,  # strain
+        elastic2d=2,  # bend
+    )
+    flex2.young = 1e3
+    flex2.thickness = 0.01
+    flex2.selfcollide = mujoco.mjtFlexSelf.mjFLEXSELF_NONE
+    self.assertIsNotNone(flex2)
+    model2 = spec2.compile()
+    self.assertIsNotNone(model2)
+    # Shell mode creates face-based constraints; verify they exist.
+    self.assertGreater(model2.neq, 0)
+
+  def test_make_flex_defaults(self):
+    # Create a spec with minimal flexcomp args (defaults).
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body(name='flex_body')
+    flex = body.make_flex(name='default_flex', equality=1)
+    self.assertIsNotNone(flex)
+    model = spec.compile()
+    self.assertIsNotNone(model)
+
+  def test_make_flex_with_pos_quat(self):
+    # Create a spec with flexcomp that has a pose.
+    spec = mujoco.MjSpec()
+    body = spec.worldbody.add_body(name='flex_body')
+    flex = body.make_flex(
+        name='posed_flex',
+        type='grid',
+        dim=2,
+        count=[3, 3, 1],
+        spacing=[0.1, 0.1, 0.1],
+        pos=[1.0, 2.0, 3.0],
+        quat=[1.0, 0.0, 0.0, 0.0],
+        equality=1,
+    )
+    self.assertIsNotNone(flex)
+    model = spec.compile()
+    self.assertIsNotNone(model)
+
 if __name__ == '__main__':
   absltest.main()
diff --git a/src/user/user_api.cc b/src/user/user_api.cc
index 42e0d7a6..0f657ff9 100644
--- a/src/user/user_api.cc
+++ b/src/user/user_api.cc
@@ -33,6 +33,7 @@
 #include <mujoco/mujoco.h>
 #include "engine/engine_support.h"
 #include "user/user_cache.h"
+#include "user/user_flexcomp.h"
 #include "user/user_model.h"
 #include "user/user_objects.h"
 #include "user/user_resource.h"
@@ -599,6 +600,121 @@ mjsFlex* mjs_addFlex(mjSpec* s) {
 
 
 
+// helper: convert type string to mjtFcompType
+static mjtFcompType FlexcompTypeFromStr(const char* type) {
+  if (!type || !strcmp(type, "grid"))    return mjFCOMPTYPE_GRID;
+  if (!strcmp(type, "box"))              return mjFCOMPTYPE_BOX;
+  if (!strcmp(type, "cylinder"))         return mjFCOMPTYPE_CYLINDER;
+  if (!strcmp(type, "ellipsoid"))        return mjFCOMPTYPE_ELLIPSOID;
+  if (!strcmp(type, "square"))           return mjFCOMPTYPE_SQUARE;
+  if (!strcmp(type, "disc"))             return mjFCOMPTYPE_DISC;
+  if (!strcmp(type, "circle"))           return mjFCOMPTYPE_CIRCLE;
+  if (!strcmp(type, "mesh"))             return mjFCOMPTYPE_MESH;
+  if (!strcmp(type, "gmsh"))             return mjFCOMPTYPE_GMSH;
+  if (!strcmp(type, "direct"))           return mjFCOMPTYPE_DIRECT;
+  return mjFCOMPTYPE_GRID;  // default
+}
+
+// helper: convert dof string to mjtDof
+static mjtDof FlexcompDofFromStr(const char* dof) {
+  if (!dof || !strcmp(dof, "full"))      return mjFCOMPDOF_FULL;
+  if (!strcmp(dof, "radial"))            return mjFCOMPDOF_RADIAL;
+  if (!strcmp(dof, "trilinear"))         return mjFCOMPDOF_TRILINEAR;
+  if (!strcmp(dof, "quadratic"))         return mjFCOMPDOF_QUADRATIC;
+  if (!strcmp(dof, "2d"))                return mjFCOMPDOF_2D;
+  return mjFCOMPDOF_FULL;  // default
+}
+
+
+// add flexcomp: create flex with auto-generated bodies/joints
+mjsFlex* mjs_makeFlex(mjsBody* body, const char* name, const char* type, int dim,
+                            const char* dof, const int count[3], const int cellcount[3],
+                            const double spacing[3], const double scale[3], double radius,
+                            double mass, double inertiabox, int equality, int rigid, int flatskin,
+                            int elastic2d, const double pos[3], const double quat[4],
+                            const double origin[3], const char* file, const mjVFS* vfs) {
+  if (!body || !name) {
+    mju_error("mjs_makeFlex: body and name must not be null");
+    return nullptr;
+  }
+
+  mjCModel* model = static_cast<mjCBody*>(body->element)->model;
+
+  // create temporary flexcomp with defaults
+  mjCFlexcomp fcomp;
+  fcomp.name = name;
+  fcomp.type = FlexcompTypeFromStr(type);
+  fcomp.doftype = FlexcompDofFromStr(dof);
+
+  // topology
+  if (count) {
+    fcomp.count[0] = count[0];
+    fcomp.count[1] = count[1];
+    fcomp.count[2] = count[2];
+  }
+  if (cellcount) {
+    fcomp.cellcount[0] = cellcount[0];
+    fcomp.cellcount[1] = cellcount[1];
+    fcomp.cellcount[2] = cellcount[2];
+  }
+  if (spacing) {
+    fcomp.spacing[0] = spacing[0];
+    fcomp.spacing[1] = spacing[1];
+    fcomp.spacing[2] = spacing[2];
+  }
+  if (scale) {
+    fcomp.scale[0] = scale[0];
+    fcomp.scale[1] = scale[1];
+    fcomp.scale[2] = scale[2];
+  }
+  if (origin) {
+    fcomp.origin[0] = origin[0];
+    fcomp.origin[1] = origin[1];
+    fcomp.origin[2] = origin[2];
+  }
+
+  // physics
+  fcomp.def.spec.flex->dim = dim;
+  fcomp.def.spec.flex->radius = radius;
+  if (mass > 0) fcomp.mass = mass;
+  if (inertiabox > 0) fcomp.inertiabox = inertiabox;
+  fcomp.equality = equality;
+  fcomp.rigid = rigid;
+  fcomp.def.spec.flex->flatskin = flatskin;
+  fcomp.def.spec.flex->elastic2d = elastic2d;
+
+  // pose
+  if (pos) {
+    fcomp.pos[0] = pos[0];
+    fcomp.pos[1] = pos[1];
+    fcomp.pos[2] = pos[2];
+  }
+  if (quat) {
+    fcomp.quat[0] = quat[0];
+    fcomp.quat[1] = quat[1];
+    fcomp.quat[2] = quat[2];
+    fcomp.quat[3] = quat[3];
+  }
+
+  // file
+  if (file) {
+    fcomp.file = file;
+  }
+
+  // call Make
+  char error[500] = "";
+  if (!fcomp.Make(body, error, sizeof(error), vfs)) {
+    model->SetError(mjCError(nullptr, "%s", error));
+    return nullptr;
+  }
+
+  // return the flex that was created (last flex in model)
+  mjCFlex* flex = model->Flexes().back();
+  return &flex->spec;
+}
+
+
+
 // add frame to body
 mjsFrame* mjs_addFrame(mjsBody* bodyspec, mjsFrame* parentframe) {
   mjCFrame* parentframeC = 0;
diff --git a/src/user/user_api.h b/src/user/user_api.h
index 1ab17a40..53b250e7 100644
--- a/src/user/user_api.h
+++ b/src/user/user_api.h
@@ -118,6 +118,14 @@ MJAPI mjsSensor* mjs_addSensor(mjSpec* s);
 // Add flex.
 MJAPI mjsFlex* mjs_addFlex(mjSpec* s);
 
+// Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.
+MJAPI mjsFlex* mjs_makeFlex(mjsBody* body, const char* name, const char* type, int dim,
+                            const char* dof, const int count[3], const int cellcount[3],
+                            const double spacing[3], const double scale[3], double radius,
+                            double mass, double inertiabox, int equality, int rigid, int flatskin,
+                            int elastic2d, const double pos[3], const double quat[4],
+                            const double origin[3], const char* file, const mjVFS* vfs);
+
 // Add contact pair.
 MJAPI mjsPair* mjs_addPair(mjSpec* s, const mjsDefault* def);
 
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index b7157df7..80d0675a 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -9952,6 +9952,22 @@ int mjs_isWarning_wrapper(MjSpec& s) {
   return mjs_isWarning(s.get());
 }
 
+std::optional<MjsFlex> mjs_makeFlex_wrapper(MjsBody& body, const String& name, const StringOrNull& type, int dim, const StringOrNull& dof, const NumberArray& count, const NumberArray& cellcount, const NumberArray& spacing, const NumberArray& scale, double radius, double mass, double inertiabox, int equality, int rigid, int flatskin, int elastic2d, const NumberArray& pos, const NumberArray& quat, const NumberArray& origin, const StringOrNull& file, const MjVFS& vfs) {
+  CHECK_VAL(name);
+  UNPACK_NULLABLE_ARRAY(int, count);
+  UNPACK_NULLABLE_ARRAY(int, cellcount);
+  UNPACK_NULLABLE_ARRAY(double, spacing);
+  UNPACK_NULLABLE_ARRAY(double, scale);
+  UNPACK_NULLABLE_ARRAY(double, pos);
+  UNPACK_NULLABLE_ARRAY(double, quat);
+  UNPACK_NULLABLE_ARRAY(double, origin);
+  mjsFlex* result = mjs_makeFlex(body.get(), name.as<const std::string>().data(), type.as<const std::string>().data(), dim, dof.as<const std::string>().data(), count_.data(), cellcount_.data(), spacing_.data(), scale_.data(), radius, mass, inertiabox, equality, rigid, flatskin, elastic2d, pos_.data(), quat_.data(), origin_.data(), file.as<const std::string>().data(), vfs.get());
+  if (result == nullptr) {
+    return std::nullopt;
+  }
+  return MjsFlex(result);
+}
+
 int mjs_makeMesh_wrapper(MjsMesh& mesh, mjtMeshBuiltin builtin, const val& params, int nparams) {
   UNPACK_VALUE(double, params);
   return mjs_makeMesh(mesh.get(), builtin, params_.data(), nparams);
@@ -13474,6 +13490,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mjs_getWrapSideSite", &mjs_getWrapSideSite_wrapper);
   function("mjs_getWrapTarget", &mjs_getWrapTarget_wrapper);
   function("mjs_isWarning", &mjs_isWarning_wrapper);
+  function("mjs_makeFlex", &mjs_makeFlex_wrapper);
   function("mjs_makeMesh", &mjs_makeMesh_wrapper);
   function("mjs_nextChild", &mjs_nextChild_wrapper);
   function("mjs_nextElement", &mjs_nextElement_wrapper);
diff --git a/wasm/codegen/generators/functions.py b/wasm/codegen/generators/functions.py
index bb296d12..f7939a83 100644
--- a/wasm/codegen/generators/functions.py
+++ b/wasm/codegen/generators/functions.py
@@ -143,7 +143,7 @@ def get_param_string(p: ast_nodes.FunctionParameterDecl) -> str:
     # Pointer to primitive value parameters or arrays
     if p.type.inner_type.name == "char":
       if p.nullable:
-        return f"const NullableString& {p.name}"
+        return f"const StringOrNull& {p.name}"
       else:
         return f"const String& {p.name}"
     elif (

From cae5a608d5d2153fd16979b518f4319427e4308c Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Thu, 28 May 2026 07:00:29 -0700
Subject: [PATCH 47/61] Import NVIDIA/warp 1.13.0 from GitHub.

PiperOrigin-RevId: 922749092
Change-Id: Ide815e4224e64af69e7c3a58102508bc81d22770
---
 mjx/cuda_requirements.txt                     | 10 +++----
 mjx/mujoco/mjx/_src/io.py                     |  4 +--
 .../warp/_src/jax_experimental/custom_call.py | 12 +++++----
 .../warp/_src/jax_experimental/ffi.py         | 27 ++++++++++---------
 mjx/pyproject.toml                            |  2 +-
 5 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt
index 1db7abfb..a726944a 100644
--- a/mjx/cuda_requirements.txt
+++ b/mjx/cuda_requirements.txt
@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
 jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
     --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
     --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
-warp-lang==1.12.1 \
-    --hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
-    --hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
-    --hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
-    --hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
+warp-lang==1.13.0 \
+    --hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
+    --hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
+    --hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
+    --hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index c1e2049e..dc10af8c 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     wp_field.shape = wp_field.shape[1:]
   # warp scalars
   wp_dtype = type(wp_field)
-  if wp_dtype in wp.types.warp_type_to_np_dtype:
-    return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
+  if wp_dtype in wp._src.types.warp_type_to_np_dtype:
+    return wp.dtype_to_numpy(wp_dtype)(wp_field)
 
   # warp arrays
   if isinstance(wp_field, wp.array):
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
index dba3f715..6f63ba66 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
 
     # Launch the kernel.
     wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
     )
 
 
 def _create_jax_warp_primitive():
     import jax  # noqa: PLC0415
-    from jax._src.interpreters import batching  # noqa: PLC0415
+    from jax.interpreters import batching  # noqa: PLC0415
     from jax.interpreters import mlir  # noqa: PLC0415
     from jax.interpreters.mlir import ir  # noqa: PLC0415
     from jaxlib.hlo_helpers import custom_call  # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
         # Figure out the number of outputs.
         wp_kernel = _registered_kernels[params["kernel"]]
         output_count = len(wp_kernel.adj.args) - len(args)
-        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
         size = shape[dim]
-        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
         # Create the batched primitive.
         return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
 
@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
     def base_type_to_jax_ir(warp_dtype):
         warp_to_jax_dict = {
             wp.float16: ir.F16Type.get(),
+            wp.bfloat16: ir.BF16Type.get(),
             wp.float32: ir.F32Type.get(),
             wp.float64: ir.F64Type.get(),
             wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
     def base_type_is_compatible(warp_type, jax_ir_type):
         jax_ir_to_warp = {
             "f16": wp.float16,
+            "bf16": wp.bfloat16,
             "f32": wp.float32,
             "f64": wp.float64,
             "i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
         # Figure out the types and shapes of the input arrays.
         arg_strings = []
         operand_layouts = []
-        for actual, warg in zip(args, wp_kernel.adj.args):
+        for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
             wtype = warg.type
             rtt = ir.RankedTensorType(actual.type)
 
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
index d931dded..ac7b0f26 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
@@ -8,8 +8,8 @@ import ctypes
 import inspect
 import threading
 import traceback
+from collections.abc import Callable
 from enum import IntEnum
-from typing import Callable
 
 import jax
 
@@ -25,6 +25,7 @@ from warp._src.types import (
     type_size_in_bytes,
     type_to_warp,
 )
+from warp._src.utils import warn
 
 from .xla_ffi import *
 
@@ -249,19 +250,18 @@ class FfiKernel:
                 # check dtype
                 if input_value.dtype != input_arg.jax_scalar_type:
                     raise TypeError(
-                            f"Invalid data type for array argument '{input_arg.name}',"
-                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
+                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                     )
                 # check ndim
                 if input_value.ndim != input_arg.jax_ndim:
                     raise TypeError(
-                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                     )
                 # check inner dims
                 for d in range(input_arg.dtype_ndim):
                     if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                         raise TypeError(
-                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                         )
             else:
                 # make sure scalar is not a traced variable, should be static
@@ -455,6 +455,7 @@ class FfiKernel:
                         hooks.forward_smem_bytes,
                         kernel_params,
                         stream,
+                        None,  # apic_info
                     )
                 else:
                     wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
         except Exception as e:
             print(traceback.format_exc())
             return create_ffi_error(
-                call_frame.contents.api,
-                XLA_FFI_Error_Code.UNKNOWN,
-                f"FFI callback error: {type(e).__name__}: {e}",
+                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
             )
 
 
@@ -617,7 +616,9 @@ class FfiCallable:
         self.input_output_aliases = input_output_aliases
 
         # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
+                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
+        )
 
         self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
         ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1071,7 +1072,7 @@ class FfiCallable:
                         call_desc.capture = capture
 
                     else:
-                        # not capturing or on CPU
+                        # not capturing
                         self.func(*arg_list)
 
         except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
                 try:
                     gi.zero_()
                 except Exception as e:
-                    wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
+                    warn(f"Failed to zero gradient array: {e}", stacklevel=2)
                     raise e
 
         # NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
         non_static_inputs, output_vals_tuple = residuals
 
         input_vals = list(non_static_inputs)
-        for i, v in zip(static_args, nondiff_vals):
+        for i, v in zip(static_args, nondiff_vals, strict=True):
             input_vals.insert(i, v)
 
         # Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(
 
         out_dims_map = {}
         param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
-        for name, val in zip(differentiable_input_names, non_static_inputs):
+        for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
             ann = param_ann.get(name)
             if ann is None:
                 continue
diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index 43609582..fbf96bcf 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.12.1",
+    "warp-lang==1.13.0",
 ]
 dev = [
     "isort",

From 50ec8e84def856676c42275e872f6d8ee81b1f5b Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Thu, 28 May 2026 21:46:15 +0530
Subject: [PATCH 48/61] test: intercept mju_warning in overflow regression test

---
 test/engine/engine_io_test.cc | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index 1e93810c..db4a9d7c 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -987,20 +987,29 @@ TEST_F(EngineIoTest, LoadModelBufferRejectsOverflowingSizes) {
   // when multiplied by sizeof(type). ntexdata is a good candidate since it
   // is a byte count field and gets multiplied by sizeof(mjtByte)==1, but other
   // fields multiply by sizeof(int) or sizeof(mjtNum), making overflow easier.
-  mjtSize* sizes = reinterpret_cast<mjtSize*>(buffer.data() + header_bytes);
+  // Use memcpy to avoid undefined behavior from unaligned access.
+  const int size_offset = header_bytes + 49 * sizeof(mjtSize);
 
-  // save original value, set to overflow-inducing value
-  mjtSize original = sizes[49];  // ntexdata index (approximate)
-  sizes[49] = static_cast<mjtSize>(SIZE_MAX / 2);
+  // set to overflow-inducing value
+  mjtSize overflow_val = static_cast<mjtSize>(SIZE_MAX / 2);
+  std::memcpy(buffer.data() + size_offset, &overflow_val, sizeof(mjtSize));
 
   // also need to update the nbuffer field (last size) to avoid the early
   // nbuffer mismatch check — but the overflow should be caught earlier
   // in safeAddToBufferSize/mj_makeModel before we reach that check.
 
+  // intercept mju_warning because the test framework translates it to ADD_FAILURE
+  static bool warning_triggered = false;
+  warning_triggered = false;
+  mju_user_warning = [](const char* msg) {
+    warning_triggered = true;
+  };
+
   // attempt to load — should return NULL, not crash
   mjModel* bad_model = mj_loadModelBuffer(buffer.data(), bufsize);
   EXPECT_THAT(bad_model, IsNull())
       << "Expected mj_loadModelBuffer to reject overflow-inducing sizes";
+  EXPECT_TRUE(warning_triggered) << "Expected a warning about invalid sizes";
 
   // clean up if somehow it succeeded
   if (bad_model) {

From b935d4153c755bcd0b6f3d0e0f59f1df134a3284 Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Thu, 28 May 2026 10:09:07 -0700
Subject: [PATCH 49/61] Add new mju_threadpool API function, and delete old
 threading API.

PiperOrigin-RevId: 922838541
Change-Id: Id9f7e0fb298ffde61fcc49a802dc78971858ce51
---
 CMakeLists.txt                            |   7 +-
 doc/APIreference/APIglobals.rst           |   4 -
 doc/APIreference/APItypes.rst             |  30 ---
 doc/APIreference/functions.rst            |  55 +---
 doc/changelog.rst                         |  12 +-
 doc/includes/references.h                 |  30 +--
 doc/programming/index.rst                 |   7 +-
 doc/programming/samples.rst               |   2 +-
 doc/programming/simulation.rst            |   7 +-
 include/mujoco/mjassert.h                 |   4 -
 include/mujoco/mjdata.h                   |  10 +-
 include/mujoco/mjthread.h                 |  42 ---
 include/mujoco/mjxmacro.h                 |   5 +-
 include/mujoco/mujoco.h                   |  20 +-
 python/mujoco/functions.cc                |   3 +
 python/mujoco/introspect/enums.py         |  10 -
 python/mujoco/introspect/functions.py     |  88 +-----
 python/mujoco/introspect/structs.py       |  59 +---
 sample/testspeed.cc                       |   7 +-
 simulate/simulate.cc                      |  34 ++-
 simulate/simulate.h                       |   8 +-
 src/engine/CMakeLists.txt                 |   2 +
 src/engine/engine_collision_driver.c      | 107 +++++---
 src/engine/engine_forward.c               | 124 +--------
 src/engine/engine_io.c                    |   7 +-
 src/engine/engine_memory.c                | 136 +++++-----
 src/engine/engine_memory.h                |   3 -
 src/engine/engine_sensor.c                |  45 ++--
 src/engine/engine_thread.cc               | 190 +++++++++++++
 src/engine/engine_thread.h                |  41 +++
 src/thread/CMakeLists.txt                 |  23 --
 src/thread/thread_pool.cc                 | 312 ----------------------
 src/thread/thread_pool.h                  |  85 ------
 src/thread/thread_queue.h                 | 152 -----------
 src/thread/thread_task.cc                 |  33 ---
 src/thread/thread_task.h                  |  49 ----
 test/benchmark/ccd_benchmark_test.cc      |  83 +++---
 test/benchmark/thread_performance_test.cc |  26 +-
 test/engine/engine_io_test.cc             |  62 +----
 test/engine/engine_thread_test.cc         |  11 +-
 test/thread/CMakeLists.txt                |  17 --
 test/thread/thread_pool_test.cc           | 142 ----------
 test/thread/thread_queue_test.cc          |  46 ----
 unity/Runtime/Bindings/MjBindings.cs      | 141 +---------
 wasm/codegen/generated/bindings.cc        |  32 ++-
 wasm/codegen/generators/constants.py      |  14 -
 wasm/tests/enums_test.ts                  |   4 -
 47 files changed, 576 insertions(+), 1755 deletions(-)
 delete mode 100644 include/mujoco/mjthread.h
 create mode 100644 src/engine/engine_thread.cc
 create mode 100644 src/engine/engine_thread.h
 delete mode 100644 src/thread/CMakeLists.txt
 delete mode 100644 src/thread/thread_pool.cc
 delete mode 100644 src/thread/thread_pool.h
 delete mode 100644 src/thread/thread_queue.h
 delete mode 100644 src/thread/thread_task.cc
 delete mode 100644 src/thread/thread_task.h
 delete mode 100644 test/thread/CMakeLists.txt
 delete mode 100644 test/thread/thread_pool_test.cc
 delete mode 100644 test/thread/thread_queue_test.cc

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0423d3f4..9a9a34b4 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -118,7 +118,6 @@ set(MUJOCO_HEADERS
     include/mujoco/mjsan.h
     include/mujoco/mjspec.h
     include/mujoco/mjspecmacro.h
-    include/mujoco/mjthread.h
     include/mujoco/mjtype.h
     include/mujoco/mjui.h
     include/mujoco/mjvisualize.h
@@ -135,6 +134,7 @@ endif()
 
 # Emscripten does not support SHARED libs
 if(NOT EMSCRIPTEN)
+  find_package(Threads REQUIRED)
   add_library(mujoco SHARED ${MUJOCO_RESOURCE_FILES})
 else()
   add_library(mujoco STATIC ${MUJOCO_RESOURCE_FILES})
@@ -160,7 +160,6 @@ add_subdirectory(src/engine)
 add_subdirectory(src/user)
 add_subdirectory(src/xml)
 add_subdirectory(src/xml/mjz)
-add_subdirectory(src/thread)
 if(NOT EMSCRIPTEN AND NOT MUJOCO_USE_FILAMENT_MJR_COMPAT)
   add_subdirectory(src/render/classic)
   add_subdirectory(src/ui)
@@ -204,6 +203,10 @@ target_link_libraries(
           tinyxml2
 )
 
+if(NOT EMSCRIPTEN)
+  target_link_libraries(mujoco PRIVATE Threads::Threads)
+endif()
+
 set_target_properties(
   mujoco PROPERTIES VERSION "${mujoco_VERSION}" PUBLIC_HEADER "${MUJOCO_HEADERS}"
 )
diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst
index 9890735b..2f6a4f69 100644
--- a/doc/APIreference/APIglobals.rst
+++ b/doc/APIreference/APIglobals.rst
@@ -452,10 +452,6 @@ indicated otherwise.
    * - ``mjMINAWAKE``
      - 10
      - The minimum number of timesteps that must pass after a tree is awoken, before it is allowed to go back to sleep.
-   * - ``mjMAXTHREAD``
-     - 128
-     - Maximum number of OS threads that can be used in a thread pool.
-       Defined in `mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`_.
 
 
 .. _glNumericSizes:
diff --git a/doc/APIreference/APItypes.rst b/doc/APIreference/APItypes.rst
index ba1a6135..fbc2e6d9 100644
--- a/doc/APIreference/APItypes.rst
+++ b/doc/APIreference/APItypes.rst
@@ -1013,36 +1013,6 @@ Asset cache used by the compiler to avoid repeated slow recompilation. See :ref:
 
 .. mujoco-include:: mjCache
 
-.. _mjtTaskStatus:
-
-mjtTaskStatus
-~~~~~~~~~~~~~
-
-Status values for :ref:`mjTask`.
-
-.. mujoco-include:: mjtTaskStatus
-
-.. _mjTask:
-
-mjTask
-~~~~~~
-
-This is a representation of a task to be run asynchronously inside of an :ref:`mjThreadPool` . It is created in the
-:ref:`mju_threadPoolEnqueue` method of the :ref:`mjThreadPool`  and is used to join the task at completion.
-The ``status`` field uses values from :ref:`mjtTaskStatus`.
-
-.. mujoco-include:: mjTask
-
-.. _mjThreadPool:
-
-mjThreadPool
-~~~~~~~~~~~~
-
-This is the data structure of the threadpool. It can only be constructed programmatically, and does not
-have an analog in MJCF. In order to enable multi-threaded calculations, a pointer to an existing :ref:`mjThreadPool`
-should be assigned to the ``mjData.threadpool``.
-
-.. mujoco-include:: mjThreadPool
 
 .. _tyStatStructure:
 
diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index 4daaaf44..7a409cc2 100644
--- a/doc/APIreference/functions.rst
+++ b/doc/APIreference/functions.rst
@@ -3375,59 +3375,14 @@ If no match, return NULL.
 
 Threads
 ^^^^^^^
-.. _mju_threadPoolCreate:
+.. _mju_threadpool:
 
-`mju_threadPoolCreate <#mju_threadPoolCreate>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+`mju_threadpool <#mju_threadpool>`__
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. mujoco-include:: mju_threadPoolCreate
+.. mujoco-include:: mju_threadpool
 
-Create a thread pool with the specified number of threads running.
-
-.. _mju_bindThreadPool:
-
-`mju_bindThreadPool <#mju_bindThreadPool>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mju_bindThreadPool
-
-Adds a thread pool to mjData and configures it for multi-threaded use.
-
-.. _mju_threadPoolEnqueue:
-
-`mju_threadPoolEnqueue <#mju_threadPoolEnqueue>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mju_threadPoolEnqueue
-
-Enqueue a task in a thread pool.
-
-.. _mju_threadPoolDestroy:
-
-`mju_threadPoolDestroy <#mju_threadPoolDestroy>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mju_threadPoolDestroy
-
-Destroy a thread pool.
-
-.. _mju_defaultTask:
-
-`mju_defaultTask <#mju_defaultTask>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mju_defaultTask
-
-Initialize an mjTask.
-
-.. _mju_taskJoin:
-
-`mju_taskJoin <#mju_taskJoin>`__
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. mujoco-include:: mju_taskJoin
-
-Wait for a task to complete.
+Create a thread pool with nthread worker threads.
 
 .. _Standardmath:
 
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 699d0dd0..c5b5377d 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -10,6 +10,16 @@ General
 - Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp<body-flexcomp>` element for
   programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as
   ``body.make_flex()`` in Python.
+- Added :ref:`mju_threadpool`, a new function for creating a thread pool on an ``mjData`` instance. When a thread pool
+  is initialized, parts of the simulation pipeline, such as collision detection and constraint solving across islands,
+  are parallelized. The thread pool is automatically destroyed when the ``mjData`` is freed.
+
+  .. admonition:: Breaking API changes
+     :class: attention
+
+     - The header file ``mjthread.h`` was removed along with the engine threading API.
+
+     **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
 
 Bug fixes
 ^^^^^^^^^
@@ -1708,7 +1718,7 @@ New features
    If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
    see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
 
-5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
+5. Added ``mjThreadPool`` and ``mjTask`` which allow for multi-threaded operations within the MuJoCo engine
    pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
 
    - Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
diff --git a/doc/includes/references.h b/doc/includes/references.h
index 60b8d699..1ce969d7 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -89,9 +89,12 @@ struct mjData_ {
   // arena pointer
   size_t  parena;            // first available byte in arena
 
+  // threading
+  uintptr_t threadpool;      // thread pool pointer
+  mjtBool threadlock;        // disable stack freeing during threaded execution
+
   // memory utilization statistics
   mjtSize maxuse_stack;                       // maximum stack allocation in bytes (mutable)
-  mjtSize maxuse_threadstack[mjMAXTHREAD];    // maximum stack allocation per thread in bytes
   mjtSize maxuse_arena;                       // maximum arena allocation in bytes
   int     maxuse_con;                         // maximum number of contacts
   int     maxuse_efc;                         // maximum number of scalar constraints
@@ -394,9 +397,6 @@ struct mjData_ {
   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
 };
@@ -2146,21 +2146,6 @@ typedef struct mjsDefault_ {       // default specification
   mjsTendon* tendon;               // tendon defaults
   mjsActuator* actuator;           // actuator defaults
 } mjsDefault;
-typedef enum mjtTaskStatus_ {  // status values for mjTask
-  mjTASK_NEW = 0,              // newly created
-  mjTASK_QUEUED,               // enqueued in a thread pool
-  mjTASK_COMPLETED             // completed execution
-} mjtTaskStatus;
-struct mjThreadPool_ {
-  int nworker;  // number of workers in the pool
-};
-typedef struct mjThreadPool_ mjThreadPool;
-struct mjTask_ {        // a task that can be executed by a thread pool.
-  mjfTask func;         // pointer to the function that implements the task
-  void* args;           // arguments to func
-  volatile int status;  // status of the task
-};
-typedef struct mjTask_ mjTask;
 typedef enum mjtDisableBit_ {     // disable default feature bitflags
   mjDSBL_CONSTRAINT   = 1<<0,     // entire constraint solver
   mjDSBL_EQUALITY     = 1<<1,     // equality constraints
@@ -3658,12 +3643,7 @@ void mju_getResourceDir(mjResource* resource, const char** dir, int* ndir);
 int mju_isModifiedResource(const mjResource* resource, const char* timestamp);
 mjSpec* mju_decodeResource(mjResource* resource, const char* content_type,
                            const mjVFS* vfs);
-mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
-void mju_bindThreadPool(mjData* d, void* thread_pool);
-void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
-void mju_threadPoolDestroy(mjThreadPool* thread_pool);
-void mju_defaultTask(mjTask* task);
-void mju_taskJoin(mjTask* task);
+void mju_threadpool(mjData* d, int nthread);
 mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
                        const char* prefix, const char* suffix);
 mjsBody* mjs_addBody(mjsBody* body, const mjsDefault* def);
diff --git a/doc/programming/index.rst b/doc/programming/index.rst
index 24c2402f..25aa9ee9 100644
--- a/doc/programming/index.rst
+++ b/doc/programming/index.rst
@@ -31,9 +31,8 @@ OpenGL renderer
    state-of-the-art rendering engines (and can be replaced with such an engine if desired) but nevertheless it provides
    efficient and informative 3D rendering.
 Thread
-   The threading framework is written in C++ and exposed in C. It provides a :ref:`mjThreadPool<mjThreadPool>` interface
-   to process tasks asynchronously. To enable use in MuJoCo, create a thread pool and assign it to the
-   ``mjData.threadpool`` field.
+   The threading framework is written in C++ and exposed in C. It provides a thread pool interface
+   to process tasks asynchronously. To enable use in MuJoCo, call ``mju_threadpool``.
 UI framework
    The UI framework is written in C. UI elements are rendered in OpenGL. It has its own event
    mechanism and abstract hooks for keyboard and mouse input. The code samples use it with GLFW, but it can also be used
@@ -171,8 +170,6 @@ links below, to make this documentation self-contained.
    Defines enums and structs used for :doc:`procedural model editing <modeledit>`.
 `mjplugin.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjplugin.h>`__
    Defines data structures required by :ref:`engine plugins<exPlugin>`.
-`mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`__
-   Defines data structures and functions required by :ref:`thread<Thread>`.
 `mjmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmacro.h>`__
    Defines C macros that are useful in user code.
 `mjxmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`__
diff --git a/doc/programming/samples.rst b/doc/programming/samples.rst
index 8a964c86..ee6c7feb 100644
--- a/doc/programming/samples.rst
+++ b/doc/programming/samples.rst
@@ -60,7 +60,7 @@ Where the command line arguments are
   keyframe named "test" is present in the model, it is used as the initial state.
 - The ``ctrlnoise`` argument prevents models from settling into a static state where, due to warmstarts, one can
   measure artificially faster simulation.
-- When ``npoolthread > 1`` is specified, an engine-internal :ref:`mjThreadPool` is created with the specified number of
+- When ``npoolthread > 1`` is specified, an engine-internal thread pool is created with the specified number of
   threads, to speed up simulation of large scenes. Note that while it is possible to use both ``nthread`` and
   ``npoolthread``, the scenarios for which one would want these different types of multithreading are usually mutually
   exclusive.
diff --git a/doc/programming/simulation.rst b/doc/programming/simulation.rst
index 702a9768..2f636221 100644
--- a/doc/programming/simulation.rst
+++ b/doc/programming/simulation.rst
@@ -562,10 +562,9 @@ external force computed by inverse dynamics.
 Multi-threading
 ~~~~~~~~~~~~~~~
 
-MuJoCo has experimental support for within-step multi-threading. When a :ref:`mjThreadPool` is assigned to
-``mjData.threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
-:ref:`islands<siSleep>` — can be distributed across worker threads. Note that within-step threading currently has
-significant memory overhead and is still a work in progress.
+MuJoCo has support for within-step multi-threading. When a thread pool is initialized via
+``mju_threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
+:ref:`islands<siSleep>` — can be distributed across worker threads.
 
 The more common and well-supported use of multi-threading is to speed up sampling operations that are
 common in more advanced applications. Simulation is inherently serial over time (the output of one mj_step is the
diff --git a/include/mujoco/mjassert.h b/include/mujoco/mjassert.h
index 169cbb30..36651ac5 100644
--- a/include/mujoco/mjassert.h
+++ b/include/mujoco/mjassert.h
@@ -20,7 +20,6 @@
 #include <mujoco/mjplugin.h>
 #include <mujoco/mjrender.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjthread.h>
 #include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
@@ -123,9 +122,6 @@ MJ_ASSERT_SIZE(mjtEvent, 4);
 MJ_ASSERT_SIZE(mjtItem, 4);
 MJ_ASSERT_SIZE(mjtSection, 4);
 
-// mjthread.h
-MJ_ASSERT_SIZE(mjtTaskStatus, 4);
-
 // mjplugin.h
 MJ_ASSERT_SIZE(mjtPluginCapabilityBit, 4);
 
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
index 486c382a..2ee67f95 100644
--- a/include/mujoco/mjdata.h
+++ b/include/mujoco/mjdata.h
@@ -20,7 +20,7 @@
 
 #include <mujoco/mjtype.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjthread.h>
+
 
 
 
@@ -112,9 +112,12 @@ struct mjData_ {
   // arena pointer
   size_t  parena;            // first available byte in arena
 
+  // threading
+  uintptr_t threadpool;      // thread pool pointer
+  mjtBool threadlock;        // disable stack freeing during threaded execution
+
   // memory utilization statistics
   mjtSize maxuse_stack;                       // maximum stack allocation in bytes (mutable)
-  mjtSize maxuse_threadstack[mjMAXTHREAD];    // maximum stack allocation per thread in bytes
   mjtSize maxuse_arena;                       // maximum arena allocation in bytes
   int     maxuse_con;                         // maximum number of contacts
   int     maxuse_efc;                         // maximum number of scalar constraints
@@ -417,9 +420,6 @@ struct mjData_ {
   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
 };
diff --git a/include/mujoco/mjthread.h b/include/mujoco/mjthread.h
deleted file mode 100644
index 153ff0cc..00000000
--- a/include/mujoco/mjthread.h
+++ /dev/null
@@ -1,42 +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_INCLUDE_MJTHREAD_H_
-#define MUJOCO_INCLUDE_MJTHREAD_H_
-
-#define mjMAXTHREAD 128        // maximum number of threads in a thread pool
-
-typedef enum mjtTaskStatus_ {  // status values for mjTask
-  mjTASK_NEW = 0,              // newly created
-  mjTASK_QUEUED,               // enqueued in a thread pool
-  mjTASK_COMPLETED             // completed execution
-} mjtTaskStatus;
-
-// function pointer type for mjTask
-typedef void* (*mjfTask)(void*);
-
-// An opaque type representing a thread pool.
-struct mjThreadPool_ {
-  int nworker;  // number of workers in the pool
-};
-typedef struct mjThreadPool_ mjThreadPool;
-
-struct mjTask_ {        // a task that can be executed by a thread pool.
-  mjfTask func;         // pointer to the function that implements the task
-  void* args;           // arguments to func
-  volatile int status;  // status of the task
-};
-typedef struct mjTask_ mjTask;
-
-#endif  // MUJOCO_INCLUDE_MJTHREAD_H_
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
index adb7686a..04539f5e 100644
--- a/include/mujoco/mjxmacro.h
+++ b/include/mujoco/mjxmacro.h
@@ -1014,6 +1014,7 @@
     X( size_t,    pstack             ) \
     X( size_t,    pbase              ) \
     X( size_t,    parena             ) \
+    X( uintptr_t, threadpool         ) \
     X( size_t,    maxuse_stack       ) \
     X( size_t,    maxuse_arena       ) \
     X( int,       maxuse_con         ) \
@@ -1036,13 +1037,11 @@
     X( mjtBool,   flg_energyvel      ) \
     X( mjtBool,   flg_subtreevel     ) \
     X( mjtBool,   flg_rnepost        ) \
-    X( mjtNum,    time               ) \
-    X( uintptr_t, threadpool         )
+    X( mjtNum,    time               )
 
 
 // vector fields of mjData
 #define MJDATA_VECTOR                                                \
-    X( size_t,         maxuse_threadstack, mjMAXTHREAD,  1         ) \
     X( mjSolverStat,   solver,             mjNISLAND,    mjNSOLVER ) \
     X( int,            solver_niter,       mjNISLAND,    1         ) \
     X( int,            solver_nnz,         mjNISLAND,    1         ) \
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 4a1384a7..1af40764 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -31,7 +31,6 @@
 #include <mujoco/mjrender.h>
 #include <mujoco/mjsan.h>
 #include <mujoco/mjspec.h>
-#include <mujoco/mjthread.h>
 #include <mujoco/mjtype.h>
 #include <mujoco/mjui.h>
 #include <mujoco/mjvisualize.h>
@@ -1593,23 +1592,8 @@ MJAPI mjSpec* mju_decodeResource(mjResource* resource, const char* content_type,
 
 //---------------------------------- Threads -------------------------------------------------------
 
-// Create a thread pool with the specified number of threads running.
-MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
-
-// Adds a thread pool to mjData and configures it for multi-threaded use.
-MJAPI void mju_bindThreadPool(mjData* d, void* thread_pool);
-
-// Enqueue a task in a thread pool.
-MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
-
-// Destroy a thread pool.
-MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
-
-// Initialize an mjTask.
-MJAPI void mju_defaultTask(mjTask* task);
-
-// Wait for a task to complete.
-MJAPI void mju_taskJoin(mjTask* task);
+// Create a thread pool with nthread worker threads.
+MJAPI void mju_threadpool(mjData* d, int nthread);
 
 
 //---------------------------------- Attachment ----------------------------------------------------
diff --git a/python/mujoco/functions.cc b/python/mujoco/functions.cc
index e509bce3..05dfa362 100644
--- a/python/mujoco/functions.cc
+++ b/python/mujoco/functions.cc
@@ -725,6 +725,9 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_version>(pymodule);
   Def<traits::mj_versionString>(pymodule);
 
+  // Thread pool
+  Def<traits::mju_threadpool>(pymodule);
+
   // Ray collision
   Def<traits::mj_multiRay>(
       pymodule,
diff --git a/python/mujoco/introspect/enums.py b/python/mujoco/introspect/enums.py
index 2247a805..3611253e 100644
--- a/python/mujoco/introspect/enums.py
+++ b/python/mujoco/introspect/enums.py
@@ -609,16 +609,6 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjS_AWAKE', 1),
          ]),
      )),
-    ('mjtTaskStatus',
-     EnumDecl(
-         name='mjtTaskStatus',
-         declname='enum mjtTaskStatus_',
-         values=dict([
-             ('mjTASK_NEW', 0),
-             ('mjTASK_QUEUED', 1),
-             ('mjTASK_COMPLETED', 2),
-         ]),
-     )),
     ('mjtGeomInertia',
      EnumDecl(
          name='mjtGeomInertia',
diff --git a/python/mujoco/introspect/functions.py b/python/mujoco/introspect/functions.py
index d2323d9e..5deefe3c 100644
--- a/python/mujoco/introspect/functions.py
+++ b/python/mujoco/introspect/functions.py
@@ -9985,23 +9985,9 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Find the decoder for a resource and return the decoded spec. The caller takes ownership of the spec and is responsible for cleaning it up.',  # pylint: disable=line-too-long
      )),
-    ('mju_threadPoolCreate',
+    ('mju_threadpool',
      FunctionDecl(
-         name='mju_threadPoolCreate',
-         return_type=PointerType(
-             inner_type=ValueType(name='mjThreadPool'),
-         ),
-         parameters=(
-             FunctionParameterDecl(
-                 name='number_of_threads',
-                 type=ValueType(name='size_t'),
-             ),
-         ),
-         doc='Create a thread pool with the specified number of threads running.',  # pylint: disable=line-too-long
-     )),
-    ('mju_bindThreadPool',
-     FunctionDecl(
-         name='mju_bindThreadPool',
+         name='mju_threadpool',
          return_type=ValueType(name='void'),
          parameters=(
              FunctionParameterDecl(
@@ -10011,75 +9997,11 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
              FunctionParameterDecl(
-                 name='thread_pool',
-                 type=PointerType(
-                     inner_type=ValueType(name='void'),
-                 ),
+                 name='nthread',
+                 type=ValueType(name='int'),
              ),
          ),
-         doc='Adds a thread pool to mjData and configures it for multi-threaded use.',  # pylint: disable=line-too-long
-     )),
-    ('mju_threadPoolEnqueue',
-     FunctionDecl(
-         name='mju_threadPoolEnqueue',
-         return_type=ValueType(name='void'),
-         parameters=(
-             FunctionParameterDecl(
-                 name='thread_pool',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjThreadPool'),
-                 ),
-             ),
-             FunctionParameterDecl(
-                 name='task',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjTask'),
-                 ),
-             ),
-         ),
-         doc='Enqueue a task in a thread pool.',
-     )),
-    ('mju_threadPoolDestroy',
-     FunctionDecl(
-         name='mju_threadPoolDestroy',
-         return_type=ValueType(name='void'),
-         parameters=(
-             FunctionParameterDecl(
-                 name='thread_pool',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjThreadPool'),
-                 ),
-             ),
-         ),
-         doc='Destroy a thread pool.',
-     )),
-    ('mju_defaultTask',
-     FunctionDecl(
-         name='mju_defaultTask',
-         return_type=ValueType(name='void'),
-         parameters=(
-             FunctionParameterDecl(
-                 name='task',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjTask'),
-                 ),
-             ),
-         ),
-         doc='Initialize an mjTask.',
-     )),
-    ('mju_taskJoin',
-     FunctionDecl(
-         name='mju_taskJoin',
-         return_type=ValueType(name='void'),
-         parameters=(
-             FunctionParameterDecl(
-                 name='task',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjTask'),
-                 ),
-             ),
-         ),
-         doc='Wait for a task to complete.',
+         doc='Create a thread pool with nthread worker threads.',
      )),
     ('mjs_attach',
      FunctionDecl(
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
index 119219ee..8fa54694 100644
--- a/python/mujoco/introspect/structs.py
+++ b/python/mujoco/introspect/structs.py
@@ -5104,42 +5104,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
-    ('mjThreadPool',
-     StructDecl(
-         name='mjThreadPool',
-         declname='struct mjThreadPool_',
-         fields=(
-             StructFieldDecl(
-                 name='nworker',
-                 type=ValueType(name='int'),
-                 doc='number of workers in the pool',
-             ),
-         ),
-     )),
-    ('mjTask',
-     StructDecl(
-         name='mjTask',
-         declname='struct mjTask_',
-         fields=(
-             StructFieldDecl(
-                 name='func',
-                 type=ValueType(name='mjfTask'),
-                 doc='pointer to the function that implements the task',
-             ),
-             StructFieldDecl(
-                 name='args',
-                 type=PointerType(
-                     inner_type=ValueType(name='void'),
-                 ),
-                 doc='arguments to func',
-             ),
-             StructFieldDecl(
-                 name='status',
-                 type=ValueType(name='int', is_volatile=True),
-                 doc='status of the task',
-             ),
-         ),
-     )),
     ('mjPreContact',
      StructDecl(
          name='mjPreContact',
@@ -5422,19 +5386,21 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='size_t'),
                  doc='first available byte in arena',
              ),
+             StructFieldDecl(
+                 name='threadpool',
+                 type=ValueType(name='uintptr_t'),
+                 doc='thread pool pointer',
+             ),
+             StructFieldDecl(
+                 name='threadlock',
+                 type=ValueType(name='mjtBool'),
+                 doc='disable stack freeing during threaded execution',
+             ),
              StructFieldDecl(
                  name='maxuse_stack',
                  type=ValueType(name='mjtSize'),
                  doc='maximum stack allocation in bytes (mutable)',
              ),
-             StructFieldDecl(
-                 name='maxuse_threadstack',
-                 type=ArrayType(
-                     inner_type=ValueType(name='mjtSize'),
-                     extents=(128,),
-                 ),
-                 doc='maximum stack allocation per thread in bytes',
-             ),
              StructFieldDecl(
                  name='maxuse_arena',
                  type=ValueType(name='mjtSize'),
@@ -6903,11 +6869,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='constraint force',
                  array_extent=('nidof',),
              ),
-             StructFieldDecl(
-                 name='threadpool',
-                 type=ValueType(name='uintptr_t'),
-                 doc='thread pool pointer',
-             ),
              StructFieldDecl(
                  name='signature',
                  type=ValueType(name='uint64_t'),
diff --git a/sample/testspeed.cc b/sample/testspeed.cc
index f14e6572..0cedb767 100644
--- a/sample/testspeed.cc
+++ b/sample/testspeed.cc
@@ -218,8 +218,7 @@ int main(int argc, char** argv) {
 
     // make and bind threadpool
     if (npoolthread > 1) {
-      mjThreadPool* threadpool = mju_threadPoolCreate(npoolthread);
-      mju_bindThreadPool(d[id], threadpool);
+      mju_threadpool(d[id], npoolthread);
     }
   }
 
@@ -348,11 +347,7 @@ int main(int argc, char** argv) {
 
   // free per-thread data
   for (int id=0; id < nthread; id++) {
-    mjThreadPool* threadpool = (mjThreadPool*) d[id]->threadpool;
     mj_deleteData(d[id]);
-    if (threadpool) {
-      mju_threadPoolDestroy(threadpool);
-    }
   }
 
   // finalize
diff --git a/simulate/simulate.cc b/simulate/simulate.cc
index 4249d7a9..61940816 100644
--- a/simulate/simulate.cc
+++ b/simulate/simulate.cc
@@ -1642,33 +1642,36 @@ void UiEvent(mjuiState* state) {
     // simulation section
     else if (it && it->sectionid==SECT_SIMULATION) {
       switch (it->itemid) {
-      case 1:             // Reset
+      case 1:             // Threadpool
+        sim->pending_.update_threadpool = true;
+        break;
+      case 2:             // Reset
         sim->pending_.reset = true;
         break;
 
-      case 2:             // Reload
+      case 3:             // Reload
         sim->uiloadrequest.fetch_add(1);
         break;
 
-      case 3:             // Align
+      case 4:             // Align
         sim->pending_.align = true;
         break;
 
-      case 4:             // Copy key
+      case 5:             // Copy key
         sim->pending_.copy_key = true;
         sim->pending_.copy_key_full_precision = sim->platform_ui->IsShiftKeyPressed();
         break;
 
-      case 5:             // Adjust key
-      case 6:             // Load key
+      case 6:             // Adjust key
+      case 7:             // Load key
         sim->pending_.load_key = true;
         break;
 
-      case 7:             // Save key
+      case 8:             // Save key
         sim->pending_.save_key = true;
         break;
 
-      case 11:            // History scrubber
+      case 12:            // History scrubber
         sim->run = 0;
         sim->pending_.load_from_history = true;
         mjui0_update_section(sim, SECT_SIMULATION);
@@ -2174,6 +2177,11 @@ void Simulate::Sync(bool state_only) {
     pending_.print_data = std::nullopt;
   }
 
+  if (pending_.update_threadpool) {
+    mju_threadpool(d_, nthread);
+    pending_.update_threadpool = false;
+  }
+
   if (pending_.reset) {
     mj_resetData(m_, d_);
     mj_forward(m_, d_);
@@ -2561,13 +2569,13 @@ void Simulate::LoadOnRenderThread() {
   }
 
   // set keyframe range and divisions
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.range[0] = 0;
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.range[1] = mjMAX(0, this->m_->nkey - 1);
-  this->ui0.sect[SECT_SIMULATION].item[5].slider.divisions = mjMAX(1, this->m_->nkey - 1);
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.range[0] = 0;
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.range[1] = mjMAX(0, this->m_->nkey - 1);
+  this->ui0.sect[SECT_SIMULATION].item[6].slider.divisions = mjMAX(1, this->m_->nkey - 1);
 
   // set scrubber range and divisions
-  this->ui0.sect[SECT_SIMULATION].item[11].slider.range[0] = 1 - nhistory_;
-  this->ui0.sect[SECT_SIMULATION].item[11].slider.divisions = nhistory_;
+  this->ui0.sect[SECT_SIMULATION].item[12].slider.range[0] = 1 - nhistory_;
+  this->ui0.sect[SECT_SIMULATION].item[12].slider.divisions = nhistory_;
 
   // detect image sensors for visualization
   DetectImageSensors(this, this->m_);
diff --git a/simulate/simulate.h b/simulate/simulate.h
index 27727ee7..367d95a3 100644
--- a/simulate/simulate.h
+++ b/simulate/simulate.h
@@ -153,6 +153,7 @@ class Simulate {
     std::optional<std::string> save_mjb;
     std::optional<std::string> print_model;
     std::optional<std::string> print_data;
+    bool update_threadpool;
     bool reset;
     bool align;
     bool copy_key;
@@ -205,6 +206,10 @@ class Simulate {
   // simulation
   int run = 1;
 
+
+  // number of workers in threadpool
+  int nthread = 0;
+
   // atomics for cross-thread messages
   std::atomic_int exitrequest = 0;
   std::atomic_int droploadrequest = 0;
@@ -320,9 +325,10 @@ class Simulate {
 
 
   // simulation section of UI
-  const mjuiDef def_simulation[14] = {
+  const mjuiDef def_simulation[15] = {
     {mjITEM_SECTION,   "Simulation",    mjPRESERVE, nullptr,     "AS"},
     {mjITEM_RADIO,     "",              5, &this->run,           "Pause\nRun"},
+    {mjITEM_EDITINT,   "Num threads",   5, &this->nthread,       "1"},
     {mjITEM_BUTTON,    "Reset",         2, nullptr,              " #259"},
     {mjITEM_BUTTON,    "Reload",        5, nullptr,              "CL"},
     {mjITEM_BUTTON,    "Align",         2, nullptr,              "CA"},
diff --git a/src/engine/CMakeLists.txt b/src/engine/CMakeLists.txt
index e61c1dc4..baa23887 100644
--- a/src/engine/CMakeLists.txt
+++ b/src/engine/CMakeLists.txt
@@ -75,6 +75,8 @@ set(MUJOCO_ENGINE_SRCS
     engine_sort.h
     engine_support.c
     engine_support.h
+    engine_thread.cc
+    engine_thread.h
     engine_util_blas.c
     engine_util_blas.h
     engine_util_blas_avx.h
diff --git a/src/engine/engine_collision_driver.c b/src/engine/engine_collision_driver.c
index 66d0cf0f..dd70d99e 100644
--- a/src/engine/engine_collision_driver.c
+++ b/src/engine/engine_collision_driver.c
@@ -26,6 +26,7 @@
 #include "engine/engine_collision_gjk.h"
 #include "engine/engine_collision_primitive.h"
 #include "engine/engine_collision_sdf.h"
+#include "engine/engine_thread.h"
 #include "engine/engine_core_constraint.h"
 #include "engine/engine_core_util.h"
 #include "engine/engine_inline.h"
@@ -178,10 +179,8 @@ static inline mjtNum getGap(const mjModel* m, int g1, int g2, int ipair) {
 static inline void resetArena(mjData* d) {
   d->parena = d->ncon * sizeof(mjContact);
 #ifdef ADDRESS_SANITIZER
-  if (!d->threadpool) {
     ASAN_POISON_MEMORY_REGION(
       (char*)d->arena + d->parena, d->narena - d->pstack - d->parena);
-  }
 #endif
 }
 
@@ -937,6 +936,7 @@ int mj_collideOBB(const mjtNum aabb1[6], const mjtNum aabb2[6],
 }
 
 
+// binary search between two bodyflex trees
 void mj_collideTree(const mjModel* m, mjData* d, int bf1, int bf2,
                     int merged, int startadr, int pairadr) {
   int nbody = m->nbody, nbvhstatic = m->nbvhstatic;
@@ -1814,13 +1814,67 @@ static void mj_makeCapsule(const mjModel* m, mjData* d, int f, const int vid[2],
 }
 
 
+// struct for collision task
+typedef struct {
+  mjPreContact* conbuffer;  // pre-contact buffer returned by collision functions
+  int* nconbuffer;          // contact count for each collision pair
+  char* epabuffer;          // buffer for nativeccd
+  int ccd_size;             // size of nativeccd buffer
+  const int* pairbuffer;    // collision pairs (g1, g2, ipair, index into conbuffer)
+  int npair;                // number of collision pairs
+  int chunksize;            // number of pairs to process per task
+  int maxcon;               // maximum number of contacts (size of conbuffer)
+} mjContactArg;
+
+
+static void collisionTask(const mjModel* m, mjData* d, void* arg, int thread_id, int idx) {
+  mjContactArg* conargs = (mjContactArg*)arg;
+  mjPreContact* conbuffer = conargs->conbuffer;
+  char* epabuffer = conargs->epabuffer;
+  int chunksize = conargs->chunksize;
+  int globalidx = chunksize * idx;
+  const int* pair = conargs->pairbuffer + 4 * globalidx;
+  int* ncon = conargs->nconbuffer + chunksize * idx;
+
+  int npair = conargs->npair;
+  int n = mjMIN(chunksize, npair - globalidx);
+
+  mjc_setCCDBuffer(epabuffer + thread_id * conargs->ccd_size);
+  for (int i = 0; i < n; i++) {
+    int g1 = pair[4*i + 0];
+    int g2 = pair[4*i + 1];
+    int ipair = pair[4*i + 2];
+    int conpos = pair[4*i + 3];
+
+    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
+    mjtNum margin = getMargin(m, g1, g2, ipair);
+    mjtNum gap = getGap(m, g1, g2, ipair);
+    ncon[i] = collision_func(m, d, conbuffer + conpos, g1, g2, margin + gap);
+
+    // SHOULD NOT OCCUR
+    int expected_max = (globalidx + i + 1 < npair ? pair[4*(i+1) + 3] : conargs->maxcon) - conpos;
+    if (ncon[i] > expected_max) {
+      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
+              "expected at most %d from mj_maxContact", ncon[i], g1, g2, expected_max);
+    }
+  }
+  mjc_setCCDBuffer(NULL);
+}
+
+
 // compute contacts for a batch of collision pairs contained in a buffer of
 // stride 3 ints (g1, g2, ipair)
 // if buffer is NULL, results are read from arena starting at parena
 void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, size_t parena) {
+  int nthread = mju_numThread(d);
   int ccd_size = mjc_ccdSize(m->opt.ccd_iterations);
   mjtNum margin, gap;
 
+  // try to balance load of 5 chunks per thread (chunksize should be divisible by 16)
+  int chunksize = npair / mjMAX(1, 5 * nthread);
+  chunksize = mjMAX(16, (chunksize + 15) & ~15);  // round up to next 16
+  int nchunk = (npair + chunksize - 1) / chunksize;
+
   // set buffer and arena pointer
   if (!buffer) {
     buffer = (const int*) ((char*) d->arena + parena);
@@ -1830,9 +1884,6 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
 
   mj_markStack(d);
 
-  // buffer store how many contacts are generated for each pair
-  int* nconbuffer = mj_stackAllocInt(d, npair);
-
   // buffer for pair data (g1, g2, ipair, index into conbuffer)
   int* pairbuffer = mj_stackAllocInt(d, 4 * npair);
   int maxcon = 0;
@@ -1850,42 +1901,30 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
     maxcon += mj_maxContact(m, g1, g2, margin + gap > 0);
   }
 
-  // buffer for precontact data
-  mjPreContact* conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
-
   // buffer data has been copied to metadata on the stack;
   // reclaim arena space so contacts can overwrite the buffer region
   d->parena = parena;
 
-  // set buffer for nativeccd
-  mj_markStack(d);
-  mjc_setCCDBuffer(mj_stackAllocByte(d, ccd_size, sizeof(mjtNum)));
+  mjContactArg arg;
+  arg.ccd_size = ccd_size;
+  arg.pairbuffer = pairbuffer;
+  arg.nconbuffer = mjSTACKALLOC(d, npair, int);
+  arg.conbuffer = mjSTACKALLOC(d, maxcon, mjPreContact);
+  arg.npair = npair;
+  arg.chunksize = chunksize;
+  arg.maxcon = maxcon;
 
-  for (int i = 0; i < npair; i++) {
-    int g1 = pairbuffer[4*i + 0];
-    int g2 = pairbuffer[4*i + 1];
-    int ipair = pairbuffer[4*i + 2];
-    int idx = pairbuffer[4*i + 3];
-    mjfCollision collision_func = mjCOLLISIONFUNC[m->geom_type[g1]][m->geom_type[g2]];
-    margin = getMargin(m, g1, g2, ipair);
-    gap = getGap(m, g1, g2, ipair);
-    nconbuffer[i] = collision_func(m, d, conbuffer + idx, g1, g2, margin + gap);
-
-    // SHOULD NOT OCCUR
-    int expected_max = (i + 1 < npair ? pairbuffer[4*(i+1) + 3] : maxcon) - idx;
-    if (nconbuffer[i] > expected_max) {
-      mjERROR("collision function returned %d contacts for geom pair (%d, %d), "
-              "expected at most %d from mj_maxContact", nconbuffer[i], g1, g2, expected_max);
-    }
+  // dispatch narrowphase to threads with local stack allocation for EPA
+  {
+    mj_markStack(d);
+    arg.epabuffer = mj_stackAllocByte(d, ccd_size * nthread, sizeof(mjtNum));
+    mju_dispatch(m, d, collisionTask, &arg, nchunk);
+    mj_freeStack(d);
   }
 
-  // set nativeccd buffer back to NULL
-  mjc_setCCDBuffer(NULL);
-  mj_freeStack(d);
-
   int ncon = 0;
   for (int i = 0; i < npair; i++) {
-    ncon += nconbuffer[i];
+    ncon += arg.nconbuffer[i];
   }
 
   if (ncon == 0) {
@@ -1906,7 +1945,7 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
   // fill in contact data
   int conpos = 0;
   for (int i = 0; i < npair; i++) {
-    if (!(ncon = nconbuffer[i]))
+    if (!(ncon = arg.nconbuffer[i]))
       continue;
 
     int condim;
@@ -1928,7 +1967,7 @@ void mj_narrowphase(const mjModel* m, mjData* d, const int* buffer, int npair, s
       mj_contactParam(m, &condim, solref, solimp, friction, g1, g2, -1, -1);
     }
 
-    mjPreContact* bc = conbuffer + pairbuffer[4*i + 3];
+    mjPreContact* bc = arg.conbuffer + pairbuffer[4*i+3];
     margin = getMargin(m, g1, g2, ipair);
     for (int j=0; j < ncon; j++) {
       mjContact* c = con + conpos + j;
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index 921fbae4..59dcd459 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -44,8 +44,7 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_solve.h"
 #include "engine/engine_util_sparse.h"
-#include "thread/thread_pool.h"
-#include "thread/thread_task.h"
+#include "engine/engine_thread.h"
 
 
 
@@ -116,28 +115,6 @@ void mj_checkAcc(const mjModel* m, mjData* d) {
 
 //-------------------------- solver components -----------------------------------------------------
 
-// args for internal functions in mj_fwdPosition
-struct mjFwdPositionArgs_ {
-  const mjModel* m;
-  mjData* d;
-};
-typedef struct mjFwdPositionArgs_ mjFwdPositionArgs;
-
-// wrapper for mj_crb and mj_factorM
-void* mj_inertialThreaded(void* args) {
-  mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args;
-  mj_makeM(forward_args->m, forward_args->d);
-  mj_factorM(forward_args->m, forward_args->d);
-  return NULL;
-}
-
-// wrapper for mj_collision
-void* mj_collisionThreaded(void* args) {
-  mjFwdPositionArgs* forward_args = (mjFwdPositionArgs*) args;
-  mj_collision(forward_args->m, forward_args->d);
-  return NULL;
-}
-
 // kinematics-related computations
 void mj_fwdKinematics(const mjModel* m, mjData* d) {
   mj_kinematics(m, d);
@@ -162,36 +139,12 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
 
   TM_END(mjTIMER_POS_KINEMATICS);
 
-  // no threadpool: inertia and collision on main thread
-  if (!d->threadpool) {
-    // inertia, timed internally (POS_INERTIA)
-    mj_makeM(m, d);
-    mj_factorM(m, d);
+  // inertia, timed internally (POS_INERTIA)
+  mj_makeM(m, d);
+  mj_factorM(m, d);
 
-    // collision, timed internally (POS_COLLISION)
-    mj_collision(m, d);
-  }
-
-  // have threadpool: inertia and collision on separate threads
-  else {
-    mjTask tasks[2];
-    mjFwdPositionArgs forward_args;
-    forward_args.m = m;
-    forward_args.d = d;
-
-    mju_defaultTask(&tasks[0]);
-    tasks[0].func = mj_inertialThreaded;
-    tasks[0].args = &forward_args;
-    mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[0]);
-
-    mju_defaultTask(&tasks[1]);
-    tasks[1].func = mj_collisionThreaded;
-    tasks[1].args = &forward_args;
-    mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[1]);
-
-    mju_taskJoin(&tasks[0]);
-    mju_taskJoin(&tasks[1]);
-  }
+  // collision, timed internally (POS_COLLISION)
+  mj_collision(m, d);
 
   if (mj_wakeCollision(m, d)) {
     mj_updateSleep(m, d);
@@ -909,51 +862,13 @@ static void warmstart(const mjModel* m, mjData* d) {
 }
 
 
-// struct encapsulating arguments to thread task
-struct mjSolIslandArgs_ {
-  const mjModel* m;
-  mjData* d;
-  int island;
-};
-typedef struct mjSolIslandArgs_ mjSolIslandArgs;
-
-// extract arguments, pass to CG solver
-static void* CG_wrapper(void* args) {
-  mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
-  mj_solCG_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
-  return NULL;
-}
-
-// extract arguments, pass to Newton solver
-static void* Newton_wrapper(void* args) {
-  mjSolIslandArgs* solargs = (mjSolIslandArgs*) args;
-  mj_solNewton_island(solargs->m, solargs->d, solargs->island, solargs->m->opt.iterations);
-  return NULL;
-}
-
-// CG solver, multi-threaded over islands
-static void solve_threaded(const mjModel* m, mjData* d, int flg_Newton) {
-  mj_markStack(d);
-  // allocate array of arguments to be passed to threads
-  mjSolIslandArgs* sol_island_args = mjSTACKALLOC(d, d->nisland, mjSolIslandArgs);
-  mjTask* tasks = mjSTACKALLOC(d, d->nisland, mjTask);
-
-  for (int island = 0; island < d->nisland; ++island) {
-    sol_island_args[island].m = m;
-    sol_island_args[island].d = d;
-    sol_island_args[island].island = island;
-
-    mju_defaultTask(&tasks[island]);
-    tasks[island].func = flg_Newton ? Newton_wrapper : CG_wrapper;
-    tasks[island].args = &sol_island_args[island];
-    mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[island]);
+// mju_dispatch callback: solve one island
+static void solveIslandTask(const mjModel* m, mjData* d, void* arg, int thread_id, int island) {
+  if (m->opt.solver == mjSOL_NEWTON) {
+    mj_solNewton_island(m, d, island, m->opt.iterations);
+  } else {
+    mj_solCG_island(m, d, island, m->opt.iterations);
   }
-
-  for (int island = 0; island < d->nisland; ++island) {
-    mju_taskJoin(&tasks[island]);
-  }
-
-  mj_freeStack(d);
 }
 
 
@@ -1009,20 +924,7 @@ void mj_fwdConstraint(const mjModel* m, mjData* d) {
       mju_gather(d->iefc_force,      d->efc_force,       d->map_iefc2efc, nefc);
       mju_gather(d->iefc_aref,       d->efc_aref,        d->map_iefc2efc, nefc);
 
-      // solve per island, with or without threads
-      if (!d->threadpool) {
-        // no threadpool, loop over islands
-        for (int island=0; island < nisland; island++) {
-          if (m->opt.solver == mjSOL_NEWTON) {
-            mj_solNewton_island(m, d, island, m->opt.iterations);
-          } else {
-            mj_solCG_island(m, d, island, m->opt.iterations);
-          }
-        }
-      } else {
-        // have threadpool, solve using threads
-        solve_threaded(m, d, m->opt.solver == mjSOL_NEWTON);
-      }
+      mju_dispatch(m, d, solveIslandTask, NULL, nisland);
 
       // copy back solver outputs (scatter dofs since ni <= nv)
       mju_scatter(d->qacc,            d->iacc,            d->map_idof2dof, nidof);
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 6a1e10fe..ff23d33a 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -33,6 +33,7 @@
 #include "engine/engine_memory.h"
 #include "engine/engine_plugin.h"
 #include "engine/engine_sleep.h"
+#include "engine/engine_thread.h"
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
@@ -1081,6 +1082,7 @@ void mj_makeRawData(mjData** dest, const mjModel* m) {
 
   // clear threadpool
   d->threadpool = 0;
+  d->threadlock = 0;
 
   // clear nplugin (overwritten by _initPlugin)
   d->nplugin = 0;
@@ -1140,6 +1142,7 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
   *dest = *src;
   dest->buffer = save_buffer;
   dest->arena = save_arena;
+  dest->threadpool = 0;
   mj_setPtrData(m, dest);
 
   // save plugin_data, since the X macro copying block below will override it
@@ -1239,8 +1242,6 @@ mjData* mj_copyDataVisual(mjData* dest, const mjModel* m, const mjData* src, int
     }
   }
 
-  dest->threadpool = src->threadpool;
-
   return dest;
 }
 
@@ -1300,7 +1301,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
 
   // clear memory utilization stats
   d->maxuse_stack = 0;
-  memset(d->maxuse_threadstack, 0, mjMAXTHREAD*sizeof(mjtSize));
   d->maxuse_arena = 0;
   d->maxuse_con = 0;
   d->maxuse_efc = 0;
@@ -1572,6 +1572,7 @@ void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key) {
 // de-allocate mjData
 void mj_deleteData(mjData* d) {
   if (d) {
+    mju_threadpool(d, 0);
     freeDataBuffers(d);
     mju_free(d);
   }
diff --git a/src/engine/engine_memory.c b/src/engine/engine_memory.c
index 741575b7..9ecdea1e 100644
--- a/src/engine/engine_memory.c
+++ b/src/engine/engine_memory.c
@@ -16,6 +16,7 @@
 
 #include <inttypes.h>  // IWYU pragma: keep
 #include <limits.h>
+#include <stdatomic.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
@@ -26,7 +27,7 @@
 #include <mujoco/mjsan.h>  // IWYU pragma: keep
 #include "engine/engine_crossplatform.h"
 #include "engine/engine_util_errmem.h"
-#include "thread/thread_pool.h"
+
 
 #ifdef ADDRESS_SANITIZER
   #include <sanitizer/asan_interface.h>
@@ -57,25 +58,20 @@ static inline size_t fastmod(size_t a, size_t b) {
   return a % b;
 }
 
+typedef struct {
+  uintptr_t bottom;      // first memory address available to the stack
+  uintptr_t top;         // current memory address used by the stack
+  uintptr_t limit;       // top limit of the stack (stack grows down)
+  uintptr_t stack_base;  // current stack base for mark and free stack
+} mjStackInfo;
+
+
 typedef struct {
   size_t pbase;   // value of d->pbase immediately before mj_markStack
   size_t pstack;  // value of d->pstack immediately before mj_markStack
   void* pc;       // program counter of the call site of mj_markStack (only set when under asan)
 } mjStackFrame;
 
-static void maybe_lock_alloc_mutex(mjData* d) {
-  if (d->threadpool != 0) {
-    mju_threadPoolLockAllocMutex((mjThreadPool*)d->threadpool);
-  }
-}
-
-static void maybe_unlock_alloc_mutex(mjData* d) {
-  if (d->threadpool != 0) {
-    mju_threadPoolUnlockAllocMutex((mjThreadPool*)d->threadpool);
-  }
-}
-
-
 static inline mjStackInfo get_stack_info_from_data(const mjData* d) {
   mjStackInfo stack_info;
   stack_info.bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
@@ -110,14 +106,12 @@ static size_t stack_usage_redzone(const mjStackInfo* stack_info) {
 
 // allocate memory from the mjData arena
 void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
-  maybe_lock_alloc_mutex(d);
   size_t misalignment = fastmod(d->parena, alignment);
   size_t padding = misalignment ? alignment - misalignment : 0;
 
   // check size
   size_t bytes_available = d->narena - d->pstack;
   if (mjUNLIKELY(d->parena + padding + bytes > bytes_available)) {
-    maybe_unlock_alloc_mutex(d);
     return NULL;
   }
 
@@ -125,16 +119,8 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
 
   // under ASAN, get stack usage from red zone
 #ifdef ADDRESS_SANITIZER
-  mjStackInfo stack_info;
-  mjStackInfo* stack_info_ptr;
-  if (!d->threadpool) {
-    stack_info = get_stack_info_from_data(d);
-    stack_info_ptr = &stack_info;
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    stack_info_ptr = mju_getStackInfoForThread(d, thread_id);
-  }
-  stack_usage = stack_usage_redzone(stack_info_ptr);
+  mjStackInfo stack_info = get_stack_info_from_data(d);
+  stack_usage = stack_usage_redzone(&stack_info);
 #endif
 
   // allocate, update max, return pointer to buffer
@@ -150,7 +136,6 @@ void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment) {
   __msan_allocated_memory(result, bytes);
 #endif
 
-  maybe_unlock_alloc_mutex(d);
   return result;
 }
 
@@ -212,13 +197,8 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
 
   // update max usage statistics
   stack_info->top = new_top_ptr;
-  if (!d->threadpool) {
-    d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
-    d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    d->maxuse_threadstack[thread_id] = mjMAX(d->maxuse_threadstack[thread_id], usage);
-  }
+  d->maxuse_stack = mjMAX(d->maxuse_stack, usage);
+  d->maxuse_arena = mjMAX(d->maxuse_arena, usage + d->parena);
 
   return (void*)start_ptr;
 }
@@ -228,18 +208,46 @@ static inline void* stackallocinternal(mjData* d, mjStackInfo* stack_info, size_
 // declared inline so that modular arithmetic with specific alignments can be optimized out
 static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
                                const char* caller, int line) {
-  // single threaded allocation
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
-    d->pstack = stack_info.bottom - stack_info.top;
-    return result;
+  // size zero: no-op
+  if (!size) {
+    return NULL;
   }
 
-  // multi threaded allocation
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  return stackallocinternal(d, stack_info, size, alignment, caller, line);
+  // call in mju_dispatch: atomically reserve space on the stack
+  if (d->threadlock) {
+    size_t alloc_size = size + alignment - 1 + 2 * mjREDZONE;
+    size_t old_pstack = atomic_fetch_add_explicit(
+        (_Atomic size_t*)&d->pstack, alloc_size, memory_order_relaxed);
+
+    // check for stack overflow
+    size_t stack_available_bytes = (size_t)d->narena - d->parena;
+    if (mjUNLIKELY(old_pstack + alloc_size > stack_available_bytes)) {
+      char info[1024];
+      if (caller) {
+        snprintf(info, sizeof(info), " at %s, line %d", caller, line);
+      } else {
+        info[0] = '\0';
+      }
+      mju_error(
+          "mj_stackAlloc: out of memory, stack overflow%s (threadlock)\n"
+          "  max = %" PRIuPTR ", available = %" PRIuPTR ", requested = %" PRIuPTR
+          "\n nefc = %d, ncon = %d",
+          info, (uintptr_t)stack_available_bytes,
+          (uintptr_t)(stack_available_bytes - old_pstack),
+          (uintptr_t)alloc_size, d->nefc, d->ncon);
+    }
+
+    uintptr_t bottom = (uintptr_t)d->arena + (uintptr_t)d->narena;
+    uintptr_t start_ptr = bottom - old_pstack - size - mjREDZONE;
+    start_ptr -= fastmod(start_ptr, alignment);
+    ASAN_UNPOISON_MEMORY_REGION((void*)start_ptr, size);
+    return (void*)start_ptr;
+  }
+
+  mjStackInfo stack_info = get_stack_info_from_data(d);
+  void* result = stackallocinternal(d, &stack_info, size, alignment, caller, line);
+  d->pstack = stack_info.bottom - stack_info.top;
+  return result;
 }
 
 
@@ -268,17 +276,15 @@ void mj_markStack(mjData* d)
 void mj__markStack(mjData* d)
 #endif
 {
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    markstackinternal(d, &stack_info);
-    d->pstack = stack_info.bottom - stack_info.top;
-    d->pbase = stack_info.stack_base;
+  // no-op if called from mju_dispatch
+  if (d->threadlock) {
     return;
   }
 
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  markstackinternal(d, stack_info);
+  mjStackInfo stack_info = get_stack_info_from_data(d);
+  markstackinternal(d, &stack_info);
+  d->pstack = stack_info.bottom - stack_info.top;
+  d->pbase = stack_info.stack_base;
 }
 
 
@@ -319,30 +325,14 @@ void mj_freeStack(mjData* d)
 void mj__freeStack(mjData* d)
 #endif
 {
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    freestackinternal(&stack_info);
-    d->pstack = stack_info.bottom - stack_info.top;
-    d->pbase = stack_info.stack_base;
+  if (d->threadlock) {
     return;
   }
 
-  size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-  mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-  freestackinternal(stack_info);
-}
-
-
-// returns the number of bytes available on the stack
-size_t mj_stackBytesAvailable(mjData* d) {
-  if (!d->threadpool) {
-    mjStackInfo stack_info = get_stack_info_from_data(d);
-    return stack_info.top - stack_info.limit;
-  } else {
-    size_t thread_id = mju_threadPoolCurrentWorkerId((mjThreadPool*)d->threadpool);
-    mjStackInfo* stack_info = mju_getStackInfoForThread(d, thread_id);
-    return stack_info->top - stack_info->limit;
-  }
+  mjStackInfo stack_info = get_stack_info_from_data(d);
+  freestackinternal(&stack_info);
+  d->pstack = stack_info.bottom - stack_info.top;
+  d->pbase = stack_info.stack_base;
 }
 
 
diff --git a/src/engine/engine_memory.h b/src/engine/engine_memory.h
index bbfd836b..f7969eb4 100644
--- a/src/engine/engine_memory.h
+++ b/src/engine/engine_memory.h
@@ -49,9 +49,6 @@ void mj__freeStack(mjData* d) __attribute__((noinline));
 
 #endif  // ADDRESS_SANITIZER
 
-// returns the number of bytes available on the stack
-MJAPI size_t mj_stackBytesAvailable(mjData* d);
-
 // allocate bytes on the stack
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
 
diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c
index 83db1d49..c8309e77 100644
--- a/src/engine/engine_sensor.c
+++ b/src/engine/engine_sensor.c
@@ -35,8 +35,7 @@
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_spatial.h"
-#include "thread/thread_pool.h"
-#include "thread/thread_task.h"
+#include "engine/engine_thread.h"
 
 
 
@@ -64,8 +63,6 @@ mjPARTIAL_SORT(ContactSelect, ContactInfo, ContactInfoCompare);
 
 // arguments for parallel tactile sensor computation
 typedef struct mjTactileTaskArgs_ {
-  const mjModel* m;
-  mjData* d;
   int sensor_id;
   int mesh_id;
   int geom_id;
@@ -80,10 +77,8 @@ typedef struct mjTactileTaskArgs_ {
 
 
 // worker function for parallel tactile computation over taxel batches
-static void* tactile_taxel_batch(void* args) {
+static void* tactile_taxel_batch(const mjModel* m, mjData* d, void* args) {
   mjTactileTaskArgs* t = (mjTactileTaskArgs*)args;
-  const mjModel* m = t->m;
-  mjData* d = t->d;
   int mesh_id = t->mesh_id;
   int geom_id = t->geom_id;
   int parent_weld = t->parent_weld;
@@ -193,6 +188,12 @@ static void* tactile_taxel_batch(void* args) {
 }
 
 
+static void tactileTask(const mjModel* m, mjData* d, void* arg, int thread_id, int task_id) {
+  mjTactileTaskArgs* args_array = (mjTactileTaskArgs*)arg;
+  tactile_taxel_batch(m, d, &args_array[task_id]);
+}
+
+
 // apply cutoff to sensor i, clamping values in data buffer
 static void apply_cutoff(const mjModel* m, int i, mjtNum* data) {
   mjtNum cutoff = m->sensor_cutoff[i];
@@ -1261,18 +1262,15 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
       // threshold for parallelization (taxel count below which sequential is faster)
       const int kTactileParallelThreshold = 1000;
 
-      // parallel path: use threadpool to process taxel batches
-      if (d->threadpool && ncon >= kTactileParallelThreshold) {
-        int nthreads = mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool);
-        int batch_size = (ncon + nthreads - 1) / nthreads;
-        int ntasks = (ncon + batch_size - 1) / batch_size;
+      // parallel path: use mj_batch to process taxel batches
+      int nthread = mju_numThread(d);
+      if (nthread > 0 && ncon >= kTactileParallelThreshold) {
+        int batch_size = (ncon + nthread - 1) / nthread;
+        int ntask = (ncon + batch_size - 1) / batch_size;
 
-        mjTask* tasks = mjSTACKALLOC(d, ntasks, mjTask);
-        mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntasks, mjTactileTaskArgs);
+        mjTactileTaskArgs* task_args = mjSTACKALLOC(d, ntask, mjTactileTaskArgs);
 
-        for (int t = 0; t < ntasks; t++) {
-          task_args[t].m = m;
-          task_args[t].d = d;
+        for (int t = 0; t < ntask; t++) {
           task_args[t].sensor_id = i;
           task_args[t].mesh_id = mesh_id;
           task_args[t].geom_id = geom_id;
@@ -1283,22 +1281,13 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
           task_args[t].start_taxel = t * batch_size;
           task_args[t].end_taxel = mju_min((t+1) * batch_size, ncon);
           task_args[t].forcesT = forcesT;
-
-          mju_defaultTask(&tasks[t]);
-          tasks[t].func = tactile_taxel_batch;
-          tasks[t].args = &task_args[t];
-          mju_threadPoolEnqueue((mjThreadPool*)d->threadpool, &tasks[t]);
         }
 
-        for (int t = 0; t < ntasks; t++) {
-          mju_taskJoin(&tasks[t]);
-        }
+        mju_dispatch(m, d, tactileTask, task_args, ntask);
       }
       // sequential path: call tactile_taxel_batch with full range
       else {
         mjTactileTaskArgs args;
-        args.m = m;
-        args.d = d;
         args.sensor_id = i;
         args.mesh_id = mesh_id;
         args.geom_id = geom_id;
@@ -1309,7 +1298,7 @@ static void mj_computeSensorAcc(const mjModel* m, mjData* d, int i, mjtNum* sens
         args.start_taxel = 0;
         args.end_taxel = ncon;
         args.forcesT = forcesT;
-        tactile_taxel_batch(&args);
+        tactile_taxel_batch(m, d, &args);
       }
 
       // compute sensor output
diff --git a/src/engine/engine_thread.cc b/src/engine/engine_thread.cc
new file mode 100644
index 00000000..3f521d21
--- /dev/null
+++ b/src/engine/engine_thread.cc
@@ -0,0 +1,190 @@
+// Copyright 2026 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 "engine/engine_thread.h"
+
+#include <atomic>
+#include <cstdint>
+#include <thread>
+#include <vector>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmacro.h>
+#include <mujoco/mjmodel.h>
+#include "engine/engine_memory.h"
+
+// context for thread pool stored on mjData
+class ThreadPoolContext {
+ public:
+  explicit ThreadPoolContext(int nthread) : threads_(nthread) {
+    for (int i = 0; i < nthread; i++) {
+      threads_[i] = std::thread(&ThreadPoolContext::Worker, this, i + 1);
+    }
+  }
+
+  // non-copyable, non-movable
+  ThreadPoolContext(const ThreadPoolContext&) = delete;
+  ThreadPoolContext& operator=(const ThreadPoolContext&) = delete;
+
+  ~ThreadPoolContext() {
+    signal_.store(0, std::memory_order_release);
+    signal_.notify_all();
+    for (auto& thread : threads_) {
+      if (thread.joinable()) {
+        thread.join();
+      }
+    }
+  }
+
+  // dispatch tasks to the thread pool and work on them on the main thread
+  void Dispatch(const mjModel* model, mjData* data, mjTaskFunc func, void* arg,
+                int ntask) {
+    func_ = func;
+    model_ = model;
+    data_ = data;
+    arg_ = arg;
+    ntask_ = ntask;
+    next_.store(0, std::memory_order_relaxed);
+    ndone_.store(0, std::memory_order_relaxed);
+    signal_.store(-signal_.load(std::memory_order_relaxed),
+                 std::memory_order_release);
+    signal_.notify_all();
+
+    // process tasks on main thread
+    while (true) {
+      int taskId = next_.fetch_add(1, std::memory_order_relaxed);
+      if (taskId >= ntask_) {
+        break;
+      }
+      func_(model_, data_, arg_, 0, taskId);
+    }
+
+    // busy wait for rest of workers to finish
+    int nthread = threads_.size();
+    while (ndone_.load(std::memory_order_acquire) < nthread) {
+    }
+  }
+
+  int ThreadCount() const { return threads_.size(); }
+
+ private:
+  // worker loop for each worker thread
+  void Worker(int threadId) {
+    int status = 1;
+
+    // main loop waiting for next batch of tasks
+    while (true) {
+      // wait until signal atomic is notified and sign flips
+      signal_.wait(status, std::memory_order_acquire);
+
+      // if signal was set to zero, halt
+      status = signal_.load(std::memory_order_acquire);
+      if (status == 0) {
+        return;
+      }
+
+      // subloop to process tasks for the current batch
+      while (true) {
+        int taskId = next_.fetch_add(1, std::memory_order_relaxed);
+        if (taskId >= ntask_) {
+          break;
+        }
+        func_(model_, data_, arg_, threadId, taskId);
+      }
+
+      // let main thread know this worker is done
+      ndone_.fetch_add(1, std::memory_order_release);
+    }
+  }
+
+  // arguments for the current batch set by Dispatch
+  const mjModel* model_;
+  mjData* data_;
+  mjTaskFunc func_;
+  void* arg_;
+  int ntask_;  // total number of tasks for workers to do
+
+  // atomic for each worker to grab the next task
+  std::atomic<int> next_{0};
+
+  // atomic counter for number of workers who completed their tasks
+  alignas(64) std::atomic<int> ndone_{0};
+
+  // alternating signal from -1, 1 to start / halt the worker threads,
+  // set to 0 to force all workers to exit
+  std::atomic<int> signal_{1};
+
+  std::vector<std::thread> threads_;
+};
+
+
+
+// create a thread pool with nthread threads
+void mju_threadpool(mjData* d, int nthread) {
+  if (d->threadpool) {
+    ThreadPoolContext* ctx =
+        reinterpret_cast<ThreadPoolContext*>(d->threadpool);
+    // same size, nothing to do
+    if (nthread == ctx->ThreadCount()) {
+        return;
+    }
+    delete ctx;
+    d->threadpool = 0;  // null out in case nthread == 0
+  }
+
+  if (nthread >= 1) {
+    d->threadpool = reinterpret_cast<uintptr_t>(new ThreadPoolContext(nthread));
+  }
+}
+
+
+// dispatch ntask tasks to the thread pool; passes arg into func along with
+// thread_id and task_id
+void mju_dispatch(const mjModel* m, mjData* d, mjTaskFunc func, void* arg,
+                  int ntask) {
+  // no thread pool or trivial number of tasks: run on main thread
+  if (!d->threadpool || ntask < 2) {
+    for (int i = 0; i < ntask; i++) {
+      func(m, d, arg, 0, i);
+    }
+    return;
+  }
+
+  ThreadPoolContext& ctx = *reinterpret_cast<ThreadPoolContext*>(d->threadpool);
+
+  // lock mjData and mark stack frame, memory will be freed after thread completion
+  if (!d->threadlock) {
+    mj_markStack(d);
+    d->threadlock = true;
+  }
+
+  ctx.Dispatch(m, d, func, arg, ntask);
+
+  if (d->threadlock) {
+    // update max usage statistics
+    d->maxuse_stack = mjMAX(d->maxuse_stack, d->pstack);
+    d->maxuse_arena = mjMAX(d->maxuse_arena, d->pstack + d->parena);
+
+    // unlock mjData and free stack used during worker execution
+    d->threadlock = false;
+    mj_freeStack(d);
+  }
+}
+
+
+// return total number of threads in the pool (including main thread)
+int mju_numThread(const mjData* d) {
+  ThreadPoolContext* ctx = reinterpret_cast<ThreadPoolContext*>(d->threadpool);
+  return ctx ? ctx->ThreadCount() + 1 : 1;
+}
diff --git a/src/engine/engine_thread.h b/src/engine/engine_thread.h
new file mode 100644
index 00000000..52f8ba64
--- /dev/null
+++ b/src/engine/engine_thread.h
@@ -0,0 +1,41 @@
+// Copyright 2026 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_SRC_ENGINE_ENGINE_THREAD_H_
+#define MUJOCO_SRC_ENGINE_ENGINE_THREAD_H_
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjexport.h>
+#include <mujoco/mjmodel.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// dispatch function for mju_dispatch
+typedef void (*mjTaskFunc)(const mjModel* m, mjData* d, void* arg, int thread_id, int task_id);
+
+// create a thread pool with nthread worker threads.
+MJAPI void mju_threadpool(mjData* d, int nthread);
+
+// return total number of threads in the pool (including main thread)
+MJAPI int mju_numThread(const mjData* d);
+
+// dispatch ntask tasks to the thread pool; passes arg into func along with thread_id and task_id
+MJAPI void mju_dispatch(const mjModel* m, mjData* d, mjTaskFunc func, void* arg, int ntask);
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // MUJOCO_SRC_ENGINE_ENGINE_THREAD_H_
diff --git a/src/thread/CMakeLists.txt b/src/thread/CMakeLists.txt
deleted file mode 100644
index e320c20b..00000000
--- a/src/thread/CMakeLists.txt
+++ /dev/null
@@ -1,23 +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
-#
-#     https://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.
-
-set(MUJOCO_THREAD_SRCS
-    thread_pool.cc
-    thread_pool.h
-    thread_queue.h
-    thread_task.cc
-    thread_task.h
-)
-
-target_sources(mujoco PRIVATE ${MUJOCO_THREAD_SRCS})
diff --git a/src/thread/thread_pool.cc b/src/thread/thread_pool.cc
deleted file mode 100644
index af117f6a..00000000
--- a/src/thread/thread_pool.cc
+++ /dev/null
@@ -1,312 +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 "thread/thread_pool.h"
-
-#include <stdint.h>
-
-#include <algorithm>
-#include <atomic>
-#include <cstddef>
-#include <memory>
-#include <mutex>
-#include <thread>
-#include <utility>
-#include <vector>
-
-#include <mujoco/mjsan.h>  // IWYU pragma: keep
-#include <mujoco/mjthread.h>
-#include <mujoco/mujoco.h>
-#include "engine/engine_crossplatform.h"
-#include "engine/engine_util_errmem.h"
-#include "thread/thread_queue.h"
-#include "thread/thread_task.h"
-
-namespace mujoco {
-namespace {
-constexpr size_t kThreadPoolQueueSize = 640;
-
-// Each thread being run will be assigned a worker_id.
-//   0: main thread
-//   1->n: workers
-thread_local size_t worker_id = 0;
-
-struct WorkerThread {
-  // Shutdown function passed to running threads to ensure clean shutdown.
-  static void* ShutdownFunction(void* args) {
-    return nullptr;
-  }
-
-  // Thread for the worker.
-  std::unique_ptr<std::thread> thread_;
-
-  // An mjTask for shutting down this worker.
-  mjTask shutdown_task_ {
-    &ShutdownFunction,
-    nullptr,
-    mjTASK_NEW
-  };
-};
-}  // namespace
-
-// Concrete C++ class definition for mjThreadPool.
-// (The public mjThreadPool C struct is an opaque one.)
-class ThreadPoolImpl : public mjThreadPool {
- public:
-  ThreadPoolImpl(int num_worker) : mjThreadPool{num_worker} {
-    // initialize worker threads
-    for (int i = 0; i < std::min(num_worker, mjMAXTHREAD); ++i) {
-      WorkerThread worker{
-        std::make_unique<std::thread>(ThreadPoolWorker, this, i)};
-      workers_.push_back(std::move(worker));
-    }
-  }
-
-  size_t NumberOfThreads() {
-    return workers_.size();
-  }
-
-  // start a task in the threadpool
-  void Enqueue(mjTask* task) {
-    if (mjUNLIKELY(GetAtomicTaskStatus(task).exchange(mjTASK_QUEUED) !=
-                   mjTASK_NEW)) {
-      mjERROR("task->status is not mjTASK_NEW");
-    }
-    lockless_queue_.push(task);
-  }
-
-  // shutdown the threadpool
-  void Shutdown() {
-    if (shutdown_) {
-      return;
-    }
-
-    shutdown_ = true;
-    std::vector<mjTask> shutdown_tasks(workers_.size());
-    for (auto& worker : workers_) {
-      Enqueue(&worker.shutdown_task_);
-    }
-
-    for (auto& worker : workers_) {
-      worker.thread_->join();
-    }
-  }
-
-  // registers a worker ID for a given thread
-  void RegisterWorker(const size_t input_worker_id) {
-    worker_id = input_worker_id;
-  }
-
-  // gets the worker id of the current thread
-  size_t GetWorkerId() {
-    return worker_id;
-  }
-
-  void LockAlloc() {
-    alloc_mutex_.lock();
-  }
-
-  void UnlockAlloc() {
-    alloc_mutex_.unlock();
-  }
-
-  bool IsThreadPoolBound() {
-    return thread_pool_bound_;
-  }
-
-  void BindThreadPool() {
-    thread_pool_bound_ = true;
-  }
-
-  ~ThreadPoolImpl() { Shutdown(); }
-
- private:
-  // method executed by running threads
-  static void ThreadPoolWorker(
-      ThreadPoolImpl* thread_pool, const size_t thread_index) {
-    worker_id = thread_index + 1;
-    while (!thread_pool->shutdown_) {
-      auto task = static_cast<mjTask*>(thread_pool->lockless_queue_.pop());
-      task->args = task->func(task->args);
-      GetAtomicTaskStatus(task).store(mjTASK_COMPLETED);
-    }
-  }
-
-  // indicates whether the thread pool is being shut down
-  std::atomic<bool> shutdown_ = false;
-
-  // OS threads that are running in this pool
-  std::vector<WorkerThread> workers_;
-
-  // queue of tasks to execute
-  mujoco::LocklessQueue<void*, kThreadPoolQueueSize> lockless_queue_;
-
-  // Mutex to protect arena allocations.
-  std::mutex alloc_mutex_;
-
-  // Whether or not a ThreadPool was bound using mju_bindThreadPool.
-  bool thread_pool_bound_ = false;
-};
-
-// create a thread pool
-mjThreadPool* mju_threadPoolCreate(size_t number_of_threads) {
-  return reinterpret_cast<mjThreadPool*>(new ThreadPoolImpl(number_of_threads));
-}
-
-// gets the number of shards the stack is currently broken into
-static size_t GetNumberOfShards(mjData* d) {
-  if (!d->threadpool) {
-    return 1;
-  }
-  return mju_threadPoolNumberOfThreads((mjThreadPool*)d->threadpool) + 1;
-}
-
-// returns the stack information for the specified thread's shard
-mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id) {
-  auto thread_pool = (ThreadPoolImpl*)d->threadpool;
-  if (!thread_pool || !thread_pool->IsThreadPoolBound()) {
-    mju_error("Thread Pool not bound, use mju_bindThreadPool to add an mjThreadPool to mjData");
-  }
-
-  // number of threads running in the threadpool plus the main thread
-  size_t number_of_shards = GetNumberOfShards(d);
-
-  // size of entire arena/stack in bytes
-  size_t total_arena_size_bytes = d->narena;
-
-  // set the shard cursor to the end of the arena
-  uintptr_t end_of_arena_ptr = (uintptr_t)d->arena + total_arena_size_bytes;
-
-  // each thread including the main one will get an equal shard of the stack
-  size_t bytes_per_shard = total_arena_size_bytes / (2 * (number_of_shards));
-
-  // ensure the shard is larger than the cache line
-  size_t misalignment = bytes_per_shard % mju_getDestructiveInterferenceSize();
-
-  if (misalignment != 0) {
-    bytes_per_shard += mju_getDestructiveInterferenceSize() - misalignment;
-  }
-
-  if (bytes_per_shard * number_of_shards > total_arena_size_bytes) {
-    mju_error("Arena is not large enough for %zu shards", number_of_shards);
-  }
-
-  uintptr_t result = (end_of_arena_ptr - (thread_id + 1) * bytes_per_shard);
-
-  // align the end of the shard to be mjStackInfo.
-  misalignment = result % alignof(mjStackInfo);
-  result -= misalignment;
-#ifdef ADDRESS_SANITIZER
-  // Ensure StackInfo is always accessible
-  ASAN_UNPOISON_MEMORY_REGION((void*)result, sizeof(mjStackInfo));
-#endif
-  return (mjStackInfo*) result;
-}
-
-// shards the stack for each thread
-static void  ConfigureMultiThreadedStack(mjData* d) {
-  if (!d->threadpool) {
-    mju_error("No thread pool specified for multithreaded operation");
-  }
-
-  size_t number_of_shards = GetNumberOfShards(d);
-
-  // current top of the stack
-  uintptr_t current_limit = (uintptr_t)d->arena + d->narena - d->pstack;
-
-  // set the shard cursor to the end of the arena
-  uintptr_t begin_shard_cursor_ptr = (uintptr_t)d->arena + d->narena;
-
-  for (size_t shard_index = 0; shard_index < number_of_shards; ++shard_index) {
-    mjStackInfo* end_shard_cursor_ptr = mju_getStackInfoForThread(d, shard_index);
-#ifdef ADDRESS_SANITIZER
-  // unpoison stack info
-  ASAN_UNPOISON_MEMORY_REGION((void*)end_shard_cursor_ptr, sizeof(mjStackInfo));
-#endif
-    // handle the main thread's stack which may already have data in it
-    if (shard_index == 0) {
-      // abort if the current stack is already larger than the portion of the stack
-      // that would be reserved for the main thread
-      if ((uintptr_t)end_shard_cursor_ptr > current_limit) {
-        mju_error("mj_bindThreadPool: sharding stack - existing stack larger than shard size: current_size = %zu, "
-                  "max_size = %zu", current_limit, (uintptr_t) end_shard_cursor_ptr);
-      }
-      end_shard_cursor_ptr->top = current_limit;
-      end_shard_cursor_ptr->stack_base = d->pbase;
-    } else {
-      // all other stacks are empty because threads have not been used yet
-      end_shard_cursor_ptr->top = begin_shard_cursor_ptr;
-      end_shard_cursor_ptr->stack_base = 0;
-    }
-
-    end_shard_cursor_ptr->bottom = begin_shard_cursor_ptr;
-    end_shard_cursor_ptr->limit = (uintptr_t)end_shard_cursor_ptr + sizeof(mjStackInfo);
-    begin_shard_cursor_ptr = (uintptr_t)end_shard_cursor_ptr - 1;
-  }
-}
-
-// adds a thread pool to mjData and configures it for multi-threaded use.
-void mju_bindThreadPool(mjData* d, void* thread_pool) {
-  if (d->threadpool) {
-    mju_error("Thread Pool already bound to mjData");
-  }
-
-  d->threadpool = (uintptr_t) thread_pool;
-  ((ThreadPoolImpl*)thread_pool)->BindThreadPool();
-  ConfigureMultiThreadedStack(d);
-}
-
-// gets the number of running threads in the thread pool.
-size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  return thread_pool_impl->NumberOfThreads();
-}
-
-size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  return thread_pool_impl->GetWorkerId();
-}
-
-// start a task in the threadpool
-void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  thread_pool_impl->Enqueue(task);
-}
-
-// shutdown the threadpool and free the memory
-void mju_threadPoolDestroy(mjThreadPool* thread_pool) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  thread_pool_impl->Shutdown();
-  delete thread_pool_impl;
-}
-
-// locks the allocation mutex to protect Stack and Arena allocations
-void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  thread_pool_impl->LockAlloc();
-}
-
-// unlocks the allocation mutex to protect Stack and Arena allocations
-void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool) {
-  auto thread_pool_impl = static_cast<ThreadPoolImpl*>(thread_pool);
-  thread_pool_impl->UnlockAlloc();
-}
-
-// Get the destructive interference size for the architecture.
-size_t mju_getDestructiveInterferenceSize(void) {
-  // return std::hardware_destructive_interference_size;
-  return 128;
-}
-
-}  // namespace mujoco
diff --git a/src/thread/thread_pool.h b/src/thread/thread_pool.h
deleted file mode 100644
index 54d852d9..00000000
--- a/src/thread/thread_pool.h
+++ /dev/null
@@ -1,85 +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_SRC_THREAD_THREAD_POOL_H_
-#define MUJOCO_SRC_THREAD_THREAD_POOL_H_
-
-#include <stddef.h>
-
-#include <mujoco/mjexport.h>
-#include <mujoco/mjthread.h>
-#include <mujoco/mujoco.h>
-
-#ifdef __cplusplus
-namespace mujoco {
-extern "C" {
-#endif
-
-// MultiThreaded Stack will be an approximately 50/50 split of the entire buffer, with a little
-// wiggle for alignment and caching concerns. The basic layout is to reuse the existing single
-// threaded markers, and then create shards for each thread to use as its stack.
-// Not to scale.
-// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
-// |Used Arena|Free Arena |Shard1     |Shard1     |Shard1     |Shard0    |Shard0     |Shard0     |
-// |%%%%%%%%%%|           |StackInfo  |Free Stack |Used Stack |StackInfo |Free Stack |Used Stack |
-// |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
-// |%%%%%%%%%%|           |           |           |%%%%%%%%%%%|          |           |%%%%%%%%%%%|
-// |----------|-----------|-----------|-----------|-----------|----------|-----------|-----------|
-// d->arena   d->parena   d->pstack   shard1->stack_info      shard1->bottom_of_stack            shard1->bottom_of_stack
-//                        shard1->stack_info                  shard0->stack_info     shard0->current_stack
-//                                    shard1->top_of_stack               shard1->top_of_stack
-//                                                shard1->current_stack
-typedef struct {
-  uintptr_t bottom;          // First memory address available to the stack
-  uintptr_t top;             // Current memory address used by the stack
-  uintptr_t limit;             // Top limit of the stack (note this is smaller than bottom, stack grows down)
-  uintptr_t stack_base;               // Current stack base for mark and free stack
-} mjStackInfo;
-
-// Create a thread pool with the specified number of threads running.
-MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
-
-// Returns the stack information for the specified thread's shard.
-mjStackInfo* mju_getStackInfoForThread(mjData* d, size_t thread_id);
-
-// Adds a thread pool to mjData and configures it for multi-threaded use.
-MJAPI void mju_bindThreadPool(mjData* d, void* thread_pool);
-
-// Gets the number of running threads in the thread pool.
-MJAPI size_t mju_threadPoolNumberOfThreads(mjThreadPool* thread_pool);
-
-// Gets the ID of the current thread being executed
-MJAPI size_t mju_threadPoolCurrentWorkerId(mjThreadPool* thread_pool);
-
-// Enqueue a task in a thread pool.
-MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
-
-// Locks the allocation mutex to protect Arena allocations.
-MJAPI void mju_threadPoolLockAllocMutex(mjThreadPool* thread_pool);
-
-// Unlocks the allocation mutex to protect Arena allocations.
-MJAPI void mju_threadPoolUnlockAllocMutex(mjThreadPool* thread_pool);
-
-// Destroy a thread pool.
-MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
-
-// Get the destructive interference size for the architecture.
-MJAPI size_t mju_getDestructiveInterferenceSize(void);
-
-#ifdef __cplusplus
-}  // extern "C"
-}  // namespace mujoco
-#endif  // __cplusplus
-
-#endif  // MUJOCO_SRC_THREAD_THREAD_POOL_H_
diff --git a/src/thread/thread_queue.h b/src/thread/thread_queue.h
deleted file mode 100644
index 3178f8ec..00000000
--- a/src/thread/thread_queue.h
+++ /dev/null
@@ -1,152 +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_SRC_THREAD_LOCKLESS_QUEUE_H_
-#define MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
-
-#include <atomic>
-#include <climits>
-#include <cstddef>
-#include <thread>
-
-namespace mujoco {
-
-// A Lockless Queue allows for sending information quickly between different
-// threads. This is a Multi-Producer Multi-Consumer Lockless Queue allowing for
-// multiple threads to be adding items to the queue while multiple threads are
-// consuming items from the queue. Internally it uses a Ring Buffer for storage
-// so it will not grow as items are added. Push will block if the Queue is full
-// and Pop will block if it is empty.
-//
-// For a basic overview of this category of structures:
-// https://www.linuxjournal.com/content/lock-free-multi-producer-multi-consumer-queue-ring-buffer
-template  <typename T, size_t buffer_capacity>
-class LocklessQueue {
- public:
-  bool full() const {
-    return full_internal(
-        convert_to_index(read_cursor_), convert_to_index(write_cursor_));
-  }
-
-  bool empty() const {
-    return maximum_read_cursor_ == read_cursor_;
-  }
-
-  // Push an element into the queue.
-  void push(const T& input) {
-    // Reserve a slot in the queue
-    size_t current_write_cursor;
-    size_t dummy_current_write_cursor;
-    size_t next_write_cursor;
-    size_t current_write_index;
-    size_t current_read_index;
-    do {
-      // Check if the queue is full.
-      do {
-        current_write_cursor = write_cursor_.load();
-        current_write_index = convert_to_index(current_write_cursor);
-        next_write_cursor = get_next_cursor(current_write_cursor);
-
-        current_read_index = convert_to_index(read_cursor_.load());
-      } while (full_internal(current_read_index, current_write_index));
-
-      // Once it's not full, attempt to grab a slot to write.
-      dummy_current_write_cursor = current_write_cursor;
-    } while (!write_cursor_.compare_exchange_weak(
-                 dummy_current_write_cursor, next_write_cursor));
-
-    // Write the entry.
-    buffer_[current_write_index].store(input);
-
-    // Increment maximum read cursor. Note here it has to wait if the compare
-    // and exchange fails as another thread might not have completed its write.
-    do {
-      dummy_current_write_cursor = current_write_cursor;
-    } while (!maximum_read_cursor_.compare_exchange_weak(
-                 dummy_current_write_cursor, next_write_cursor));
-  }
-
-  // Pop an element from the queue.
-  T pop() {
-    size_t current_read_cursor;
-    size_t dummy_current_read_cursor;
-    size_t current_read_index;
-    size_t next_read_cursor;
-    size_t current_maximum_read_cursor;
-    size_t current_maximum_read_index;
-    bool empty = false;
-    T result;
-    do {
-      // Wait until the queue has an element
-      do {
-        if (empty) {
-          std::this_thread::yield();
-        }
-        current_read_cursor = read_cursor_.load();
-        current_maximum_read_cursor = maximum_read_cursor_.load();
-
-        current_read_index = convert_to_index(current_read_cursor);
-        current_maximum_read_index = convert_to_index(
-            current_maximum_read_cursor);
-
-        empty = empty_internal(
-            current_read_index, current_maximum_read_index);
-      } while (empty);
-
-      next_read_cursor = get_next_cursor(current_read_cursor);
-
-      // Attempt to grab the element, if unsuccessful then wait for the next
-      // element to arrive.
-      result = buffer_[current_read_index].load();
-      dummy_current_read_cursor = current_read_cursor;
-    } while (!read_cursor_.compare_exchange_weak(
-                 dummy_current_read_cursor, next_read_cursor));
-
-    return result;
-  }
-
- private:
-  size_t convert_to_index(size_t input) const {
-    return input % internal_buffer_capacity_;
-  }
-
-  size_t get_next_cursor(size_t input) const {
-    return (input + 1) % cursor_max_;
-  }
-
-  size_t get_next_index(size_t input) const {
-    return convert_to_index(get_next_cursor(input));
-  }
-
-  bool full_internal(size_t read_index, size_t write_index) const {
-    return get_next_index(write_index) == read_index;
-  }
-
-  bool empty_internal(size_t read_index, size_t write_index) const {
-    return read_index == write_index;
-  }
-
-  const size_t internal_buffer_capacity_ = buffer_capacity + 1;
-  const size_t cursor_max_ = UINT_MAX - (UINT_MAX % internal_buffer_capacity_);
-
-  std::atomic<size_t> read_cursor_ = 0;
-  std::atomic<size_t> write_cursor_ = 0;
-  std::atomic<size_t> maximum_read_cursor_ = 0;
-
-  std::atomic<T> buffer_[(buffer_capacity + 1)];
-};
-
-}  // namespace mujoco
-
-#endif  // MUJOCO_SRC_THREAD_LOCKLESS_QUEUE_H_
diff --git a/src/thread/thread_task.cc b/src/thread/thread_task.cc
deleted file mode 100644
index 77ce2bfd..00000000
--- a/src/thread/thread_task.cc
+++ /dev/null
@@ -1,33 +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 "thread/thread_task.h"
-
-#include <thread>
-
-#include <mujoco/mjthread.h>
-
-namespace mujoco {
-void mju_defaultTask(mjTask* task) {
-  task->func = nullptr;
-  task->args = nullptr;
-  task->status = mjTASK_NEW;
-}
-
-void mju_taskJoin(mjTask* task) {
-  while (GetAtomicTaskStatus(task) != mjTASK_COMPLETED) {
-    std::this_thread::yield();
-  }
-}
-}  // namespace mujoco
diff --git a/src/thread/thread_task.h b/src/thread/thread_task.h
deleted file mode 100644
index 0ed6f468..00000000
--- a/src/thread/thread_task.h
+++ /dev/null
@@ -1,49 +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_SRC_THREAD_THREAD_TASK_H_
-#define MUJOCO_SRC_THREAD_THREAD_TASK_H_
-
-#include <mujoco/mjexport.h>
-#include <mujoco/mjthread.h>
-
-#ifdef __cplusplus
-#include <atomic>
-#include <new>
-#include <type_traits>
-namespace mujoco {
-extern "C" {
-#endif
-
-// Initialize an mjTask.
-MJAPI void mju_defaultTask(mjTask* task);
-
-// Wait for a task to complete.
-MJAPI void mju_taskJoin(mjTask* task);
-
-#ifdef __cplusplus
-}  // extern "C"
-
-using TaskStatus = std::remove_volatile_t<decltype(mjTask::status)>;
-inline std::atomic<TaskStatus>& GetAtomicTaskStatus(mjTask* task) {
-  static_assert(sizeof(std::atomic<TaskStatus>) == sizeof(TaskStatus));
-  static_assert(alignof(std::atomic<TaskStatus>) == alignof(TaskStatus));
-  static_assert(std::atomic<TaskStatus>::is_always_lock_free);
-  return *std::launder(reinterpret_cast<std::atomic<TaskStatus>*>(
-      const_cast<TaskStatus*>(&task->status)));
-}
-}  // namespace mujoco
-#endif  // __cplusplus
-
-#endif  // MUJOCO_SRC_THREAD_THREAD_TASK_H_
diff --git a/test/benchmark/ccd_benchmark_test.cc b/test/benchmark/ccd_benchmark_test.cc
index a7768dce..9f1ad4bf 100644
--- a/test/benchmark/ccd_benchmark_test.cc
+++ b/test/benchmark/ccd_benchmark_test.cc
@@ -15,8 +15,8 @@
 // A benchmark for parsing and compiling models from XML.
 
 #include <cstddef>
-#include <vector>
 #include <string>
+#include <vector>
 
 #include <benchmark/benchmark.h>
 #include <absl/base/attributes.h>
@@ -36,6 +36,9 @@ static const int kNumWarmupSteps = 1000;
 // number of steps to benchmark (before resetting state)
 static const int kBatchSize = 50;
 
+// number of threads to test
+static const int kNumThreads = 9;
+
 static const char kBoxMeshPath[] =
     "../test/engine/testdata/collision_convex/perf/boxmesh.xml";
 static const char kBoxBoxPath[] =
@@ -47,11 +50,10 @@ static const char kMixedPath[] =
 
 class TestHarness {
  public:
-  TestHarness(const char* xml_path, std::string label, int disable_flags = 0) {
+  TestHarness(const char* xml_path, std::string label) {
     // Fail test if there are any mujoco errors
     MujocoErrorTestGuard guard;
     model_ = LoadModelFromPath(xml_path);
-    model_->opt.disableflags |= disable_flags;
     data_ = mj_makeData(model_);
     for (int i=0; i < kNumWarmupSteps; i++) {
       mj_step(model_, data_);
@@ -62,13 +64,18 @@ class TestHarness {
     int size = mj_stateSize(model_, spec_);
     initial_state_.resize(size);
     mj_getState(model_, data_, initial_state_.data(), spec_);
-    label_ = label;
+    name_ = label;
   }
 
   void Reset() {
     mj_setState(model_, data_, initial_state_.data(), spec_);
   }
 
+  void SetThreads(int nthread) {
+    mju_threadpool(data_, nthread);
+    nthread_ = nthread;
+  }
+
   void RunBenchmark(benchmark::State& state) {
     std::size_t ncon = 0;
     while (state.KeepRunningBatch(kBatchSize)) {
@@ -81,7 +88,9 @@ class TestHarness {
       }
     }
 
-    state.SetLabel(label_);
+    std::string label = name_ + " " + std::to_string(nthread_ + 1) +
+                        " thread(s)";
+    state.SetLabel(label);
     state.SetItemsProcessed(ncon);  // report number of contacts per second
   }
 
@@ -91,7 +100,8 @@ class TestHarness {
   }
 
  private:
-  std::string label_;
+  std::string name_;
+  int nthread_ = 0;
   int spec_;
   mjModel* model_;
   mjData* data_;
@@ -103,62 +113,47 @@ class TestHarness {
 // separately in CPU profiles (and don't get replaced with raw calls to
 // run_parse_benchmark).
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
-    BM_BoxMesh_NativeCCD(benchmark::State& state) {
-  static TestHarness harness(kBoxMeshPath, "boxmesh.xml (nativeccd)");
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxMesh(benchmark::State& state) {
+  int nthread = state.range(0);
+  static TestHarness harness(kBoxMeshPath, "boxmesh.xml");
+  harness.SetThreads(nthread);
   harness.RunBenchmark(state);
 }
-BENCHMARK(BM_BoxMesh_NativeCCD);
-
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
-    BM_BoxMesh_LibCCD(benchmark::State& state) {
-  static TestHarness harness(kBoxMeshPath, "boxmesh.xml (libccd)",
-                             mjDSBL_NATIVECCD);
-  harness.RunBenchmark(state);
-}
-BENCHMARK(BM_BoxMesh_LibCCD);
+BENCHMARK(BM_BoxMesh)->Arg(0)->Arg(kNumThreads);
 
 void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox(benchmark::State& state) {
-  static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)", mjDSBL_NATIVECCD);
+  int nthread = state.range(0);
+  static TestHarness harness(kBoxBoxPath, "box.xml (BoxBox)");
+  harness.SetThreads(nthread);
   harness.RunBenchmark(state);
 }
-BENCHMARK(BM_BoxBox);
+BENCHMARK(BM_BoxBox)->Arg(0)->Arg(kNumThreads);
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBox_NativeCCD(benchmark::State& state) {
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_BoxBoxConvex(benchmark::State& state) {
+  int nthread = state.range(0);
   mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_Convex;
   static TestHarness harness(kBoxBoxPath, "box.xml (NativeCCD)");
+  harness.SetThreads(nthread);
   harness.RunBenchmark(state);
   mjCOLLISIONFUNC[mjGEOM_BOX][mjGEOM_BOX] = mjc_BoxBox;
 }
-BENCHMARK(BM_BoxBox_NativeCCD);
+BENCHMARK(BM_BoxBoxConvex)->Arg(0)->Arg(kNumThreads);
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
-    BM_Ellipsoid_NativeCCD(benchmark::State& state) {
-  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (nativeccd)");
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Ellipsoid(benchmark::State& state) {
+  int nthread = state.range(0);
+  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml");
+  harness.SetThreads(nthread);
   harness.RunBenchmark(state);
 }
-BENCHMARK(BM_Ellipsoid_NativeCCD);
+BENCHMARK(BM_Ellipsoid)->Arg(0)->Arg(kNumThreads);
 
-void ABSL_ATTRIBUTE_NO_TAIL_CALL
-    BM_Ellipsoid_LibCCD(benchmark::State& state) {
-  static TestHarness harness(kEllipsoidPath, "ellipsoid.xml (libccd)",
-                             mjDSBL_NATIVECCD);
+void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed(benchmark::State& state) {
+  int nthread = state.range(0);
+  static TestHarness harness(kMixedPath, "mixed.xml");
+  harness.SetThreads(nthread);
   harness.RunBenchmark(state);
 }
-BENCHMARK(BM_Ellipsoid_LibCCD);
-
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_NativeCCD(benchmark::State& state) {
-  static TestHarness harness(kMixedPath, "mixed.xml (nativeccd)");
-  harness.RunBenchmark(state);
-}
-BENCHMARK(BM_Mixed_NativeCCD);
-
-void ABSL_ATTRIBUTE_NO_TAIL_CALL BM_Mixed_LibCCD(benchmark::State& state) {
-  static TestHarness harness(kMixedPath, "mixed.xml (libccd)",
-                             mjDSBL_NATIVECCD);
-  harness.RunBenchmark(state);
-}
-BENCHMARK(BM_Mixed_LibCCD);
+BENCHMARK(BM_Mixed)->Arg(0)->Arg(kNumThreads);
 
 }  // namespace
 }  // namespace mujoco
diff --git a/test/benchmark/thread_performance_test.cc b/test/benchmark/thread_performance_test.cc
index aa4c04a8..0bb3df47 100644
--- a/test/benchmark/thread_performance_test.cc
+++ b/test/benchmark/thread_performance_test.cc
@@ -17,7 +17,6 @@
 
 #include <benchmark/benchmark.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjthread.h>
 #include <mujoco/mujoco.h>
 #include "test/fixture.h"
 
@@ -30,6 +29,9 @@ static const int kNumWarmupSteps = 500;
 // number of steps to benchmark (before resetting state)
 static const int kBatchSize = 50;
 
+// number of threads to test
+static const int kNumThreads = 6;
+
 void BM_StepHumanoid200(benchmark::State& state) {
   int nthread = state.range(0);
   std::string label = std::to_string(nthread) + " thread(s)";
@@ -43,10 +45,8 @@ void BM_StepHumanoid200(benchmark::State& state) {
   model->opt.disableflags &= ~mjDSBL_ISLAND;  // enable islands
 
   mjData* data = mj_makeData(model);
-  mjThreadPool* threadpool = nullptr;
-  if (nthread > 1) {
-    threadpool = mju_threadPoolCreate(nthread);
-    mju_bindThreadPool(data, threadpool);
+  if (nthread) {
+    mju_threadpool(data, nthread);
   }
 
   // warm-up rollout to get a steady state
@@ -73,9 +73,6 @@ void BM_StepHumanoid200(benchmark::State& state) {
   state.SetLabel(label);
   state.SetItemsProcessed(state.iterations());
   mj_deleteData(data);
-  if (threadpool) {
-    mju_threadPoolDestroy(threadpool);
-  }
 }
 
 void BM_Step22Humanoids(benchmark::State& state) {
@@ -90,10 +87,8 @@ void BM_Step22Humanoids(benchmark::State& state) {
   model->opt.disableflags &= ~mjDSBL_ISLAND;  // enable islands
 
   mjData* data = mj_makeData(model);
-  mjThreadPool* threadpool = nullptr;
-  if (nthread > 1) {
-    threadpool = mju_threadPoolCreate(nthread);
-    mju_bindThreadPool(data, threadpool);
+  if (nthread) {
+    mju_threadpool(data, nthread);
   }
 
   // warm-up rollout to get a steady state
@@ -123,12 +118,9 @@ void BM_Step22Humanoids(benchmark::State& state) {
   state.SetLabel(label);
   state.SetItemsProcessed(state.iterations());
   mj_deleteData(data);
-  if (threadpool) {
-    mju_threadPoolDestroy(threadpool);
-  }
 }
 
-BENCHMARK(BM_StepHumanoid200)->Arg(1)->Arg(6);
-BENCHMARK(BM_Step22Humanoids)->Arg(1)->Arg(6);
+BENCHMARK(BM_StepHumanoid200)->Arg(0)->Arg(kNumThreads);
+BENCHMARK(BM_Step22Humanoids)->Arg(0)->Arg(kNumThreads);
 }  // namespace
 }  // namespace mujoco
diff --git a/test/engine/engine_io_test.cc b/test/engine/engine_io_test.cc
index 0100ef4c..c83ff769 100644
--- a/test/engine/engine_io_test.cc
+++ b/test/engine/engine_io_test.cc
@@ -32,7 +32,7 @@
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
 #include "src/engine/engine_util_errmem.h"
-#include "src/thread/thread_pool.h"
+#include "src/engine/engine_thread.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -804,32 +804,18 @@ TEST_F(EngineIoTest, VeryLargeMemory) {
   }
 }
 
-struct TestFunctionArgs_ {
-  mjData* d;
-  int input;
-  int stack_output;
-  int arena_output;
-  size_t output_thread_worker;
+struct TestFunctionArgs {
+  int stack_output[1000];
+  int arena_output[1000];
 };
-typedef TestFunctionArgs_ TestFunctionArgs;
 
-void* TestFunction(void* args) {
+void TestFunction(const mjModel* m, mjData* d, void* args, int i, int j) {
   TestFunctionArgs* test_args = static_cast<TestFunctionArgs*>(args);
-  test_args->output_thread_worker =
-      mju_threadPoolCurrentWorkerId((mjThreadPool*)test_args->d->threadpool);
-  mj_markStack(test_args->d);
-  int* test_ints = mj_stackAllocInt(test_args->d, 10);
-  test_ints[0] = test_args->input;
-  test_args->stack_output = test_ints[0];
-
-  int* test_arena_ints =
-      (int*)mj_arenaAllocByte(test_args->d, sizeof(int) * 10, alignof(int));
-  test_arena_ints[0] = test_args->input;
-  test_args->arena_output = test_arena_ints[0];
-
-
-  mj_freeStack(test_args->d);
-  return nullptr;
+  mj_markStack(d);
+  int* test_ints = mj_stackAllocInt(d, 10);
+  test_ints[0] = j;
+  test_args->stack_output[j] = test_ints[0];
+  mj_freeStack(d);
 }
 
 TEST_F(EngineIoTest, TestStackShardingForThreads) {
@@ -846,38 +832,18 @@ TEST_F(EngineIoTest, TestStackShardingForThreads) {
 
   mjData* data = mj_makeData(model);
   ASSERT_THAT(data, NotNull());
-  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
-  mju_bindThreadPool(data, thread_pool);
+  mju_threadpool(data, 10);
 
   constexpr int kTasks = 1000;
-  TestFunctionArgs test_function_args[kTasks];
-  mjTask tasks[kTasks];
-  for (int i = 0; i < kTasks; ++i) {
-    test_function_args[i].d = data;
-    test_function_args[i].input = i;
-    mju_defaultTask(&tasks[i]);
-    tasks[i].func = TestFunction;
-    tasks[i].args = &test_function_args[i];
-    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
-  }
-
-  mj_markStack(data);
-  int* test_ints = mj_stackAllocInt(data, 10);
-  test_ints[0] = 1;
-  mj_freeStack(data);
+  TestFunctionArgs test_function_args;
+  mju_dispatch(model, data, TestFunction, &test_function_args, kTasks);
 
   for (int i = 0; i < kTasks; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].stack_output);
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].arena_output);
+    EXPECT_EQ(i, test_function_args.stack_output[i]);
   }
 
   mj_deleteData(data);
   mj_deleteModel(model);
-  mju_threadPoolDestroy(thread_pool);
 }
 
 #ifdef ADDRESS_SANITIZER
diff --git a/test/engine/engine_thread_test.cc b/test/engine/engine_thread_test.cc
index fd4ceae3..384f8ede 100644
--- a/test/engine/engine_thread_test.cc
+++ b/test/engine/engine_thread_test.cc
@@ -22,10 +22,9 @@
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include <mujoco/mjmodel.h>
-#include <mujoco/mjthread.h>
 #include <mujoco/mjxmacro.h>
 #include <mujoco/mujoco.h>
-#include "src/thread/thread_pool.h"
+#include "src/engine/engine_thread.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -61,8 +60,7 @@ TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
   mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
 
   // bind a threadpool to the data_threaded
-  mjThreadPool* threadpool = mju_threadPoolCreate(10);
-  mju_bindThreadPool(data_threaded, threadpool);
+  mju_threadpool(data_threaded, 10);
 
   for (int i = 0; i < 10; ++i) {
     mj_step(model, data);
@@ -83,7 +81,6 @@ TEST_F(ThreadTest, SingleAndMultiThreadedMatch) {
   mj_deleteModel(model);
   mj_deleteData(data_threaded);
   mj_deleteModel(model_threaded);
-  mju_threadPoolDestroy(threadpool);
 }
 
 TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
@@ -116,8 +113,7 @@ TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
   mj_setState(model_threaded, data_threaded, initial_state.data(), spec);
 
   // bind a threadpool to the data_threaded
-  mjThreadPool* threadpool = mju_threadPoolCreate(10);
-  mju_bindThreadPool(data_threaded, threadpool);
+  mju_threadpool(data_threaded, 10);
 
   for (int i = 0; i < 10; ++i) {
     mj_step(model, data);
@@ -138,7 +134,6 @@ TEST_F(ThreadTest, IslandSingleAndMultiThreadedMatch) {
   mj_deleteModel(model);
   mj_deleteData(data_threaded);
   mj_deleteModel(model_threaded);
-  mju_threadPoolDestroy(threadpool);
 }
 
 }  // namespace
diff --git a/test/thread/CMakeLists.txt b/test/thread/CMakeLists.txt
deleted file mode 100644
index ada7b7ef..00000000
--- a/test/thread/CMakeLists.txt
+++ /dev/null
@@ -1,17 +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
-#
-#     https://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.
-
-mujoco_test(thread_pool_test)
-
-mujoco_test(thread_queue_test)
diff --git a/test/thread/thread_pool_test.cc b/test/thread/thread_pool_test.cc
deleted file mode 100644
index 2532ec65..00000000
--- a/test/thread/thread_pool_test.cc
+++ /dev/null
@@ -1,142 +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 <atomic>
-#include <condition_variable>
-#include <memory>
-#include <mutex>
-#include <thread>
-
-#include <gtest/gtest.h>
-#include <mujoco/mujoco.h>
-
-namespace mujoco {
-namespace {
-
-struct TestFunctionArgs_ {
-  int input;
-  // make this atomic to avoid red-herring tsan failures.
-  std::atomic<int> output;
-};
-typedef struct TestFunctionArgs_ TestFunctionArgs;
-
-void* test_function(void* args) {
-  TestFunctionArgs* test_function_args = static_cast<TestFunctionArgs*>(args);
-  if (!test_function_args) {
-    return nullptr;
-  }
-  test_function_args->output = test_function_args->input;
-  return nullptr;
-}
-
-TEST(TestMjThreadPool, TestMjThreadPool10Threads) {
-  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
-
-  constexpr int kTasks = 1000;
-  TestFunctionArgs test_function_args[kTasks];
-  mjTask tasks[kTasks];
-  for (int i = 0; i < kTasks; ++i) {
-    test_function_args[i].input = i;
-    mju_defaultTask(&tasks[i]);
-    tasks[i].func = test_function;
-    tasks[i].args = &test_function_args[i];
-    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
-  }
-  mju_threadPoolDestroy(thread_pool);
-}
-
-TEST(TestMjThreadPool, TestMjThreadPool100Threads) {
-  mjThreadPool* thread_pool = mju_threadPoolCreate(100);
-
-  constexpr int kTasks = 1000;
-  TestFunctionArgs test_function_args[kTasks];
-  mjTask tasks[kTasks];
-  for (int i = 0; i < kTasks; ++i) {
-    test_function_args[i].input = i;
-    mju_defaultTask(&tasks[i]);
-    tasks[i].func = test_function;
-    tasks[i].args = &test_function_args[i];
-    mju_threadPoolEnqueue(thread_pool, &tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
-  }
-
-  mju_threadPoolDestroy(thread_pool);
-}
-
-TEST(TestMjThreadPool, TestMjThreadPoolManyWriters) {
-  mjThreadPool* thread_pool = mju_threadPoolCreate(10);
-
-  constexpr int kTasks = 20;
-  TestFunctionArgs test_function_args[kTasks];
-  mjTask tasks[kTasks];
-  std::unique_ptr<std::thread> enqueue_threads[kTasks];
-
-  // add tasks to the thread pool from many threads
-  std::condition_variable start_cv;
-  std::mutex start_mutex;
-  bool start = false;
-  for (int i = 0; i < kTasks; ++i) {
-    test_function_args[i].input = i;
-    mju_defaultTask(&tasks[i]);
-    tasks[i].func = &test_function;
-    tasks[i].args = &test_function_args[i];
-
-    enqueue_threads[i] = std::make_unique<std::thread>([&, i] {
-      // synchronize all threads adding to the thread_pool at the same time
-      {
-        std::unique_lock<std::mutex> lock(start_mutex);
-        start_cv.wait(lock, [&] { return start; });
-      }
-      // enqueue outside the lock, to get some concurrency
-      mju_threadPoolEnqueue(thread_pool, &tasks[i]);
-    });
-  }
-  {
-    std::unique_lock<std::mutex> lock(start_mutex);
-    start = true;
-  }
-  start_cv.notify_all();
-
-  for (int i = 0; i < kTasks; ++i) {
-    enqueue_threads[i]->join();
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    mju_taskJoin(&tasks[i]);
-  }
-
-  for (int i = 0; i < kTasks; ++i) {
-    EXPECT_EQ(test_function_args[i].input, test_function_args[i].output);
-  }
-
-  mju_threadPoolDestroy(thread_pool);
-}
-
-}  // namespace
-}  // namespace mujoco
diff --git a/test/thread/thread_queue_test.cc b/test/thread/thread_queue_test.cc
deleted file mode 100644
index 7b68b6f9..00000000
--- a/test/thread/thread_queue_test.cc
+++ /dev/null
@@ -1,46 +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 "src/thread/thread_queue.h"
-
-#include <cstddef>
-
-#include <gtest/gtest.h>
-
-namespace mujoco {
-namespace {
-
-constexpr size_t kBufferCapacity = 640;
-
-TEST(TestMujocoLocklessQueue, TestMujocoLocklessQueue) {
-  LocklessQueue<void*, 640> queue;
-  EXPECT_TRUE(queue.empty());
-  int test_integers[kBufferCapacity];
-  for (int h = 0; h < 10; ++h) {
-    for (int i = 0; i < kBufferCapacity; ++i) {
-      test_integers[i] = i;
-      queue.push(&test_integers[i]);
-    }
-    EXPECT_TRUE(queue.full());
-
-    for (int i = 0; i < kBufferCapacity; ++i) {
-      void* output_ptr = queue.pop();
-      ASSERT_EQ(output_ptr, &test_integers[i]);
-    }
-    EXPECT_TRUE(queue.empty());
-  }
-}
-
-}  // namespace
-}  // namespace mujoco
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index 0a6ba9b4..e25c0986 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -66,8 +66,6 @@ public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSAN_H_ = true;
 public const bool ADDRESS_SANITIZER = true;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJSPEC_H_ = true;
 public const bool THIRD_PARTY_MUJOCO_MJSPECMACRO_H_ = true;
-public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTHREAD_H_ = true;
-public const int mjMAXTHREAD = 128;
 public const bool THIRD_PARTY_MUJOCO_INCLUDE_MJTYPE_H_ = true;
 public const double mjMINVAL = 1e-15;
 public const bool THIRD_PARTY_MUJOCO_MJUI_H_ = true;
@@ -518,11 +516,6 @@ public enum mjtSleepState : int{
   mjS_ASLEEP = 0,
   mjS_AWAKE = 1,
 }
-public enum mjtTaskStatus : int{
-  mjTASK_NEW = 0,
-  mjTASK_QUEUED = 1,
-  mjTASK_COMPLETED = 2,
-}
 public enum mjtGeomInertia : int{
   mjINERTIA_VOLUME = 0,
   mjINERTIA_SHELL = 1,
@@ -1586,135 +1579,9 @@ public unsafe struct mjData_ {
   public UIntPtr pstack;
   public UIntPtr pbase;
   public UIntPtr parena;
+  public UIntPtr threadpool;
+  public byte threadlock;
   public UInt64 maxuse_stack;
-  public UInt64 maxuse_threadstack0;
-  public UInt64 maxuse_threadstack1;
-  public UInt64 maxuse_threadstack2;
-  public UInt64 maxuse_threadstack3;
-  public UInt64 maxuse_threadstack4;
-  public UInt64 maxuse_threadstack5;
-  public UInt64 maxuse_threadstack6;
-  public UInt64 maxuse_threadstack7;
-  public UInt64 maxuse_threadstack8;
-  public UInt64 maxuse_threadstack9;
-  public UInt64 maxuse_threadstack10;
-  public UInt64 maxuse_threadstack11;
-  public UInt64 maxuse_threadstack12;
-  public UInt64 maxuse_threadstack13;
-  public UInt64 maxuse_threadstack14;
-  public UInt64 maxuse_threadstack15;
-  public UInt64 maxuse_threadstack16;
-  public UInt64 maxuse_threadstack17;
-  public UInt64 maxuse_threadstack18;
-  public UInt64 maxuse_threadstack19;
-  public UInt64 maxuse_threadstack20;
-  public UInt64 maxuse_threadstack21;
-  public UInt64 maxuse_threadstack22;
-  public UInt64 maxuse_threadstack23;
-  public UInt64 maxuse_threadstack24;
-  public UInt64 maxuse_threadstack25;
-  public UInt64 maxuse_threadstack26;
-  public UInt64 maxuse_threadstack27;
-  public UInt64 maxuse_threadstack28;
-  public UInt64 maxuse_threadstack29;
-  public UInt64 maxuse_threadstack30;
-  public UInt64 maxuse_threadstack31;
-  public UInt64 maxuse_threadstack32;
-  public UInt64 maxuse_threadstack33;
-  public UInt64 maxuse_threadstack34;
-  public UInt64 maxuse_threadstack35;
-  public UInt64 maxuse_threadstack36;
-  public UInt64 maxuse_threadstack37;
-  public UInt64 maxuse_threadstack38;
-  public UInt64 maxuse_threadstack39;
-  public UInt64 maxuse_threadstack40;
-  public UInt64 maxuse_threadstack41;
-  public UInt64 maxuse_threadstack42;
-  public UInt64 maxuse_threadstack43;
-  public UInt64 maxuse_threadstack44;
-  public UInt64 maxuse_threadstack45;
-  public UInt64 maxuse_threadstack46;
-  public UInt64 maxuse_threadstack47;
-  public UInt64 maxuse_threadstack48;
-  public UInt64 maxuse_threadstack49;
-  public UInt64 maxuse_threadstack50;
-  public UInt64 maxuse_threadstack51;
-  public UInt64 maxuse_threadstack52;
-  public UInt64 maxuse_threadstack53;
-  public UInt64 maxuse_threadstack54;
-  public UInt64 maxuse_threadstack55;
-  public UInt64 maxuse_threadstack56;
-  public UInt64 maxuse_threadstack57;
-  public UInt64 maxuse_threadstack58;
-  public UInt64 maxuse_threadstack59;
-  public UInt64 maxuse_threadstack60;
-  public UInt64 maxuse_threadstack61;
-  public UInt64 maxuse_threadstack62;
-  public UInt64 maxuse_threadstack63;
-  public UInt64 maxuse_threadstack64;
-  public UInt64 maxuse_threadstack65;
-  public UInt64 maxuse_threadstack66;
-  public UInt64 maxuse_threadstack67;
-  public UInt64 maxuse_threadstack68;
-  public UInt64 maxuse_threadstack69;
-  public UInt64 maxuse_threadstack70;
-  public UInt64 maxuse_threadstack71;
-  public UInt64 maxuse_threadstack72;
-  public UInt64 maxuse_threadstack73;
-  public UInt64 maxuse_threadstack74;
-  public UInt64 maxuse_threadstack75;
-  public UInt64 maxuse_threadstack76;
-  public UInt64 maxuse_threadstack77;
-  public UInt64 maxuse_threadstack78;
-  public UInt64 maxuse_threadstack79;
-  public UInt64 maxuse_threadstack80;
-  public UInt64 maxuse_threadstack81;
-  public UInt64 maxuse_threadstack82;
-  public UInt64 maxuse_threadstack83;
-  public UInt64 maxuse_threadstack84;
-  public UInt64 maxuse_threadstack85;
-  public UInt64 maxuse_threadstack86;
-  public UInt64 maxuse_threadstack87;
-  public UInt64 maxuse_threadstack88;
-  public UInt64 maxuse_threadstack89;
-  public UInt64 maxuse_threadstack90;
-  public UInt64 maxuse_threadstack91;
-  public UInt64 maxuse_threadstack92;
-  public UInt64 maxuse_threadstack93;
-  public UInt64 maxuse_threadstack94;
-  public UInt64 maxuse_threadstack95;
-  public UInt64 maxuse_threadstack96;
-  public UInt64 maxuse_threadstack97;
-  public UInt64 maxuse_threadstack98;
-  public UInt64 maxuse_threadstack99;
-  public UInt64 maxuse_threadstack100;
-  public UInt64 maxuse_threadstack101;
-  public UInt64 maxuse_threadstack102;
-  public UInt64 maxuse_threadstack103;
-  public UInt64 maxuse_threadstack104;
-  public UInt64 maxuse_threadstack105;
-  public UInt64 maxuse_threadstack106;
-  public UInt64 maxuse_threadstack107;
-  public UInt64 maxuse_threadstack108;
-  public UInt64 maxuse_threadstack109;
-  public UInt64 maxuse_threadstack110;
-  public UInt64 maxuse_threadstack111;
-  public UInt64 maxuse_threadstack112;
-  public UInt64 maxuse_threadstack113;
-  public UInt64 maxuse_threadstack114;
-  public UInt64 maxuse_threadstack115;
-  public UInt64 maxuse_threadstack116;
-  public UInt64 maxuse_threadstack117;
-  public UInt64 maxuse_threadstack118;
-  public UInt64 maxuse_threadstack119;
-  public UInt64 maxuse_threadstack120;
-  public UInt64 maxuse_threadstack121;
-  public UInt64 maxuse_threadstack122;
-  public UInt64 maxuse_threadstack123;
-  public UInt64 maxuse_threadstack124;
-  public UInt64 maxuse_threadstack125;
-  public UInt64 maxuse_threadstack126;
-  public UInt64 maxuse_threadstack127;
   public UInt64 maxuse_arena;
   public int maxuse_con;
   public int maxuse_efc;
@@ -5926,7 +5793,6 @@ public unsafe struct mjData_ {
   public int* efc_state;
   public double* efc_force;
   public double* ifrc_constraint;
-  public UIntPtr threadpool;
   public UInt64 signature;
 }
 
@@ -7725,8 +7591,5 @@ public static unsafe extern void mjd_subQuat(double* qa, double* qb, double* Da,
 
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mjd_quatIntegrate(double* vel, double scale, double* Dquat, double* Dvel, double* Dscale);
-
-[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
-public static unsafe extern void mju_bindThreadPool(mjData_* d, void* thread_pool);
 }
 }
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index 80d0675a..735c4b1e 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -6455,15 +6455,24 @@ struct MjData {
   void set_parena(size_t value) {
     ptr_->parena = value;
   }
+  uintptr_t threadpool() const {
+    return ptr_->threadpool;
+  }
+  void set_threadpool(uintptr_t value) {
+    ptr_->threadpool = value;
+  }
+  mjtBool threadlock() const {
+    return ptr_->threadlock;
+  }
+  void set_threadlock(mjtBool value) {
+    ptr_->threadlock = value;
+  }
   int maxuse_stack() const {
     return static_cast<int>(ptr_->maxuse_stack);
   }
   void set_maxuse_stack(int value) {
     ptr_->maxuse_stack = static_cast<mjtSize>(value);
   }
-  emscripten::val maxuse_threadstack() const {
-    return emscripten::val(emscripten::typed_memory_view(128, ptr_->maxuse_threadstack));
-  }
   int maxuse_arena() const {
     return static_cast<int>(ptr_->maxuse_arena);
   }
@@ -7095,12 +7104,6 @@ struct MjData {
   emscripten::val ifrc_constraint() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->ifrc_constraint));
   }
-  uintptr_t threadpool() const {
-    return ptr_->threadpool;
-  }
-  void set_threadpool(uintptr_t value) {
-    ptr_->threadpool = value;
-  }
   uint64_t signature() const {
     return ptr_->signature;
   }
@@ -10766,6 +10769,10 @@ void mju_symmetrize_wrapper(const val& res, const NumberArray& mat, int n) {
   mju_symmetrize(res_.data(), mat_.data(), n);
 }
 
+void mju_threadpool_wrapper(MjData& d, int nthread) {
+  mju_threadpool(d.get(), nthread);
+}
+
 void mju_transformSpatial_wrapper(const val& res, const NumberArray& vec, int flg_force, const NumberArray& newpos, const NumberArray& oldpos, const NumberArray& rotnew2old) {
   UNPACK_VALUE(mjtNum, res);
   UNPACK_ARRAY(mjtNum, vec);
@@ -11467,10 +11474,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .value("mjSTEREO_NONE", mjSTEREO_NONE)
     .value("mjSTEREO_QUADBUFFERED", mjSTEREO_QUADBUFFERED)
     .value("mjSTEREO_SIDEBYSIDE", mjSTEREO_SIDEBYSIDE);
-  enum_<mjtTaskStatus>("mjtTaskStatus")
-    .value("mjTASK_NEW", mjTASK_NEW)
-    .value("mjTASK_QUEUED", mjTASK_QUEUED)
-    .value("mjTASK_COMPLETED", mjTASK_COMPLETED);
   enum_<mjtTexture>("mjtTexture")
     .value("mjTEXTURE_2D", mjTEXTURE_2D)
     .value("mjTEXTURE_CUBE", mjTEXTURE_CUBE)
@@ -11734,7 +11737,6 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("maxuse_con", &MjData::maxuse_con, &MjData::set_maxuse_con, reference())
     .property("maxuse_efc", &MjData::maxuse_efc, &MjData::set_maxuse_efc, reference())
     .property("maxuse_stack", &MjData::maxuse_stack, &MjData::set_maxuse_stack, reference())
-    .property("maxuse_threadstack", &MjData::maxuse_threadstack)
     .property("mocap_pos", &MjData::mocap_pos)
     .property("mocap_quat", &MjData::mocap_quat)
     .property("moment_colind", &MjData::moment_colind)
@@ -11804,6 +11806,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("ten_wrapadr", &MjData::ten_wrapadr)
     .property("ten_wrapnum", &MjData::ten_wrapnum)
     .property("tendon_efcadr", &MjData::tendon_efcadr)
+    .property("threadlock", &MjData::threadlock, &MjData::set_threadlock, reference())
     .property("threadpool", &MjData::threadpool, &MjData::set_threadpool, reference())
     .property("time", &MjData::time, &MjData::set_time, reference())
     .property("timer", &MjData::timer, reference())
@@ -13610,6 +13613,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
   function("mju_sum", &mju_sum_wrapper);
   function("mju_sym2dense", &mju_sym2dense_wrapper);
   function("mju_symmetrize", &mju_symmetrize_wrapper);
+  function("mju_threadpool", &mju_threadpool_wrapper);
   function("mju_transformSpatial", &mju_transformSpatial_wrapper);
   function("mju_transpose", &mju_transpose_wrapper);
   function("mju_trnVecPose", &mju_trnVecPose_wrapper);
diff --git a/wasm/codegen/generators/constants.py b/wasm/codegen/generators/constants.py
index fedb130f..086746b3 100644
--- a/wasm/codegen/generators/constants.py
+++ b/wasm/codegen/generators/constants.py
@@ -98,17 +98,6 @@ _SKIPPED_WRITABLE_ERROR: tuple[str, ...] = (
     "mj_printSchema",
 )
 
-# Omitted thread management functions
-_SKIPPED_THREAD_FUNCTIONS: tuple[str, ...] = (
-    # go/keep-sorted start
-    "mju_bindThreadPool",
-    "mju_defaultTask",
-    "mju_taskJoin",
-    "mju_threadPoolCreate",
-    "mju_threadPoolDestroy",
-    "mju_threadPoolEnqueue",
-    # go/keep-sorted end
-)
 
 # Omitted asset cache functions
 _SKIPPED_ASSET_CACHE_FUNCTIONS: tuple[str, ...] = (
@@ -207,7 +196,6 @@ _SKIPPED_UTILITY_FUNCTIONS: tuple[str, ...] = (
 # List of functions that should be skipped during the code generation process.
 SKIPPED_FUNCTIONS: tuple[str, ...] = (
     _SKIPPED_CLASS_METHODS
-    + _SKIPPED_THREAD_FUNCTIONS
     + _SKIPPED_MEMORY_FUNCTIONS
     + _SKIPPED_PLUGIN_FUNCTIONS
     + _SKIPPED_GETTERS_AND_SETTERS
@@ -223,8 +211,6 @@ SKIPPED_STRUCTS: tuple[str, ...] = (
     # go/keep-sorted start
     "mjCache",
     "mjSDF",
-    "mjTask",
-    "mjThreadPool",
     "mjUI",
     "mjVFS",
     "mjrContext",
diff --git a/wasm/tests/enums_test.ts b/wasm/tests/enums_test.ts
index d80fb48e..6f16e9f4 100644
--- a/wasm/tests/enums_test.ts
+++ b/wasm/tests/enums_test.ts
@@ -152,10 +152,6 @@ describe('Enums', () => {
     expect(mujoco.mjtSDFType).toBeDefined();
   });
 
-  it('mjtTaskStatus should exist', () => {
-    expect(mujoco.mjtTaskStatus).toBeDefined();
-  });
-
   it('mjtState should exist', () => {
     expect(mujoco.mjtState).toBeDefined();
   });

From 9436972a220c74f1c9dbe542f8070f74638ba7ef Mon Sep 17 00:00:00 2001
From: Google DeepMind <noreply@deepmind.com>
Date: Thu, 28 May 2026 11:50:32 -0700
Subject: [PATCH 50/61] Import NVIDIA/warp 1.13.0 from GitHub.

PiperOrigin-RevId: 922893824
Change-Id: Ib91ffb03257c614c48feda096822f723323be54d
---
 mjx/cuda_requirements.txt                     | 10 +++----
 mjx/mujoco/mjx/_src/io.py                     |  4 +--
 .../warp/_src/jax_experimental/custom_call.py | 12 ++++-----
 .../warp/_src/jax_experimental/ffi.py         | 27 +++++++++----------
 mjx/pyproject.toml                            |  2 +-
 5 files changed, 26 insertions(+), 29 deletions(-)

diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt
index a726944a..1db7abfb 100644
--- a/mjx/cuda_requirements.txt
+++ b/mjx/cuda_requirements.txt
@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
 jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
     --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
     --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
-warp-lang==1.13.0 \
-    --hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
-    --hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
-    --hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
-    --hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8
+warp-lang==1.12.1 \
+    --hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
+    --hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
+    --hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
+    --hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index dc10af8c..c1e2049e 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     wp_field.shape = wp_field.shape[1:]
   # warp scalars
   wp_dtype = type(wp_field)
-  if wp_dtype in wp._src.types.warp_type_to_np_dtype:
-    return wp.dtype_to_numpy(wp_dtype)(wp_field)
+  if wp_dtype in wp.types.warp_type_to_np_dtype:
+    return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
 
   # warp arrays
   if isinstance(wp_field, wp.array):
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
index 6f63ba66..dba3f715 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
 
     # Launch the kernel.
     wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
     )
 
 
 def _create_jax_warp_primitive():
     import jax  # noqa: PLC0415
-    from jax.interpreters import batching  # noqa: PLC0415
+    from jax._src.interpreters import batching  # noqa: PLC0415
     from jax.interpreters import mlir  # noqa: PLC0415
     from jax.interpreters.mlir import ir  # noqa: PLC0415
     from jaxlib.hlo_helpers import custom_call  # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
         # Figure out the number of outputs.
         wp_kernel = _registered_kernels[params["kernel"]]
         output_count = len(wp_kernel.adj.args) - len(args)
-        shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
         size = shape[dim]
-        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
         # Create the batched primitive.
         return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
 
@@ -211,7 +211,6 @@ def _create_jax_warp_primitive():
     def base_type_to_jax_ir(warp_dtype):
         warp_to_jax_dict = {
             wp.float16: ir.F16Type.get(),
-            wp.bfloat16: ir.BF16Type.get(),
             wp.float32: ir.F32Type.get(),
             wp.float64: ir.F64Type.get(),
             wp.int8: ir.IntegerType.get_signless(8),
@@ -233,7 +232,6 @@ def _create_jax_warp_primitive():
     def base_type_is_compatible(warp_type, jax_ir_type):
         jax_ir_to_warp = {
             "f16": wp.float16,
-            "bf16": wp.bfloat16,
             "f32": wp.float32,
             "f64": wp.float64,
             "i8": wp.int8,
@@ -326,7 +324,7 @@ def _create_jax_warp_primitive():
         # Figure out the types and shapes of the input arrays.
         arg_strings = []
         operand_layouts = []
-        for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
+        for actual, warg in zip(args, wp_kernel.adj.args):
             wtype = warg.type
             rtt = ir.RankedTensorType(actual.type)
 
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
index ac7b0f26..d931dded 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
@@ -8,8 +8,8 @@ import ctypes
 import inspect
 import threading
 import traceback
-from collections.abc import Callable
 from enum import IntEnum
+from typing import Callable
 
 import jax
 
@@ -25,7 +25,6 @@ from warp._src.types import (
     type_size_in_bytes,
     type_to_warp,
 )
-from warp._src.utils import warn
 
 from .xla_ffi import *
 
@@ -250,18 +249,19 @@ class FfiKernel:
                 # check dtype
                 if input_value.dtype != input_arg.jax_scalar_type:
                     raise TypeError(
-                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
+                            f"Invalid data type for array argument '{input_arg.name}',"
+                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                     )
                 # check ndim
                 if input_value.ndim != input_arg.jax_ndim:
                     raise TypeError(
-                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                     )
                 # check inner dims
                 for d in range(input_arg.dtype_ndim):
                     if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                         raise TypeError(
-                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                         )
             else:
                 # make sure scalar is not a traced variable, should be static
@@ -455,7 +455,6 @@ class FfiKernel:
                         hooks.forward_smem_bytes,
                         kernel_params,
                         stream,
-                        None,  # apic_info
                     )
                 else:
                     wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -468,7 +467,9 @@ class FfiKernel:
         except Exception as e:
             print(traceback.format_exc())
             return create_ffi_error(
-                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
+                call_frame.contents.api,
+                XLA_FFI_Error_Code.UNKNOWN,
+                f"FFI callback error: {type(e).__name__}: {e}",
             )
 
 
@@ -616,9 +617,7 @@ class FfiCallable:
         self.input_output_aliases = input_output_aliases
 
         # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
-                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
-        )
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
 
         self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
         ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1072,7 +1071,7 @@ class FfiCallable:
                         call_desc.capture = capture
 
                     else:
-                        # not capturing
+                        # not capturing or on CPU
                         self.func(*arg_list)
 
         except Exception as e:
@@ -1360,7 +1359,7 @@ def jax_kernel(
                 try:
                     gi.zero_()
                 except Exception as e:
-                    warn(f"Failed to zero gradient array: {e}", stacklevel=2)
+                    wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
                     raise e
 
         # NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1446,7 +1445,7 @@ def jax_kernel(
         non_static_inputs, output_vals_tuple = residuals
 
         input_vals = list(non_static_inputs)
-        for i, v in zip(static_args, nondiff_vals, strict=True):
+        for i, v in zip(static_args, nondiff_vals):
             input_vals.insert(i, v)
 
         # Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1462,7 +1461,7 @@ def jax_kernel(
 
         out_dims_map = {}
         param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
-        for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
+        for name, val in zip(differentiable_input_names, non_static_inputs):
             ann = param_ann.get(name)
             if ann is None:
                 continue
diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index fbf96bcf..43609582 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.13.0",
+    "warp-lang==1.12.1",
 ]
 dev = [
     "isort",

From b612d352d4ee4d9178dec6225a52134a4cec5c4e Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Thu, 28 May 2026 12:56:19 -0700
Subject: [PATCH 51/61] Fix GitHub Actions for threading.

PiperOrigin-RevId: 922926969
Change-Id: I968e6b130b6125dbc0ec65ba913d04292270b563
---
 src/engine/engine_crossplatform.h | 10 ++++++++++
 src/engine/engine_memory.c        |  4 +---
 src/engine/engine_thread.cc       | 25 ++++++++++++++++++++++++-
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/engine/engine_crossplatform.h b/src/engine/engine_crossplatform.h
index ee9ad61c..7fa8478b 100644
--- a/src/engine/engine_crossplatform.h
+++ b/src/engine/engine_crossplatform.h
@@ -69,6 +69,16 @@
   #endif
 #endif
 
+// Atomics helper for size_t.
+#if defined(_MSC_VER) && !defined(__clang__)
+  #include <intrin.h>
+  #define mj_atomic_add_size_t(ptr, val) \
+      (size_t)_InterlockedExchangeAdd64((__int64 volatile*)(ptr), (__int64)(val))
+#else
+  #define mj_atomic_add_size_t(ptr, val) \
+      __atomic_fetch_add(ptr, val, __ATOMIC_RELAXED)
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
diff --git a/src/engine/engine_memory.c b/src/engine/engine_memory.c
index 9ecdea1e..d0c044e0 100644
--- a/src/engine/engine_memory.c
+++ b/src/engine/engine_memory.c
@@ -16,7 +16,6 @@
 
 #include <inttypes.h>  // IWYU pragma: keep
 #include <limits.h>
-#include <stdatomic.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <stdio.h>
@@ -216,8 +215,7 @@ static inline void* stackalloc(mjData* d, size_t size, size_t alignment,
   // call in mju_dispatch: atomically reserve space on the stack
   if (d->threadlock) {
     size_t alloc_size = size + alignment - 1 + 2 * mjREDZONE;
-    size_t old_pstack = atomic_fetch_add_explicit(
-        (_Atomic size_t*)&d->pstack, alloc_size, memory_order_relaxed);
+    size_t old_pstack = mj_atomic_add_size_t(&d->pstack, alloc_size);
 
     // check for stack overflow
     size_t stack_available_bytes = (size_t)d->narena - d->parena;
diff --git a/src/engine/engine_thread.cc b/src/engine/engine_thread.cc
index 3f521d21..7eeae5cb 100644
--- a/src/engine/engine_thread.cc
+++ b/src/engine/engine_thread.cc
@@ -24,6 +24,10 @@
 #include <mujoco/mjmodel.h>
 #include "engine/engine_memory.h"
 
+#if defined(__cpp_lib_atomic_wait)
+
+// ----------------------------- multithreaded implementation --------------------------------------
+
 // context for thread pool stored on mjData
 class ThreadPoolContext {
  public:
@@ -129,7 +133,6 @@ class ThreadPoolContext {
 };
 
 
-
 // create a thread pool with nthread threads
 void mju_threadpool(mjData* d, int nthread) {
   if (d->threadpool) {
@@ -188,3 +191,23 @@ int mju_numThread(const mjData* d) {
   ThreadPoolContext* ctx = reinterpret_cast<ThreadPoolContext*>(d->threadpool);
   return ctx ? ctx->ThreadCount() + 1 : 1;
 }
+
+#else
+
+// -------------------------------------- fallback  ------------------------------------------------
+
+void mju_threadpool(mjData* d, int nthread) {
+  // No-op
+}
+
+void mju_dispatch(const mjModel* m, mjData* d, mjTaskFunc func, void* arg, int ntask) {
+  for (int i = 0; i < ntask; i++) {
+    func(m, d, arg, 0, i);
+  }
+}
+
+int mju_numThread(const mjData* d) {
+  return 1;
+}
+
+#endif  // __cpp_lib_atomic_wait

From df6942eb3955a7c28b62ebd050dd2d3105e8e2d7 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Thu, 28 May 2026 13:56:09 -0700
Subject: [PATCH 52/61] Support float32 precision for `mjtNum` in Python
 bindings.

PiperOrigin-RevId: 922959317
Change-Id: I384b3bdeabdc4d3f66c71cae41dabfaf51387be7
---
 python/mujoco/__init__.py                     |   3 +
 python/mujoco/bindings_test.py                | 202 +++++++++---------
 .../mujoco/codegen/generate_spec_bindings.py  |   6 +-
 python/mujoco/constants.cc                    |   2 +
 python/mujoco/serialization.h                 |   4 +-
 python/mujoco/specs.cc                        |   6 +
 6 files changed, 119 insertions(+), 104 deletions(-)

diff --git a/python/mujoco/__init__.py b/python/mujoco/__init__.py
index 799a13a5..fdd86372 100644
--- a/python/mujoco/__init__.py
+++ b/python/mujoco/__init__.py
@@ -59,6 +59,9 @@ from mujoco._functions import *
 from mujoco._specs import *
 from mujoco._structs import *
 
+import numpy as np
+MJTNUM_DTYPE = np.float32 if MJTNUM_BYTES == 4 else np.float64  # pylint: disable=undefined-variable
+
 try:
   from mujoco._render import *  # pytype: disable=import-error
 except ImportError:
diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index a79b2d1a..cd1ab24f 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -25,6 +25,8 @@ from absl.testing import parameterized
 import mujoco
 import numpy as np
 
+DTYPE = mujoco.MJTNUM_DTYPE
+
 TEST_XML = r"""
 <mujoco model="test">
   <compiler coordinate="local" angle="radian" eulerseq="xyz"/>
@@ -155,21 +157,21 @@ class MuJoCoBindingsTest(parameterized.TestCase):
       )
 
   def test_can_read_array(self):
-    np.testing.assert_array_equal(
+    np.testing.assert_allclose(
         self.model.body_pos,
         [[0, 0, 0], [0, 0, 0.1], [0, 0, 0], [0, 0, 0], [42.0, 0, 42.0]],
     )
 
   def test_can_set_array(self):
     self.data.qpos = 0.12345
-    np.testing.assert_array_equal(
+    np.testing.assert_allclose(
         self.data.qpos, [0.12345] * len(self.data.qpos)
     )
 
   def test_array_is_a_view(self):
     qpos_ref = self.data.qpos
     self.data.qpos = 0.789
-    np.testing.assert_array_equal(qpos_ref, [0.789] * len(self.data.qpos))
+    np.testing.assert_allclose(qpos_ref, [0.789] * len(self.data.qpos))
 
   # This test is disabled on PyPy as it uses sys.getrefcount
   # However PyPy is not officially supported by MuJoCo
@@ -340,24 +342,22 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     )
 
   def test_named_indexing_repr_in_data(self):
-    expected_repr = """<_MjDataGeomViews
-  id: 1
-  name: 'mybox'
-  xmat: array([0., 0., 0., 0., 0., 0., 0., 0., 0.])
-  xpos: array([0., 0., 0.])
->"""
-    self.assertEqual(expected_repr, repr(self.data.geom('mybox')))
+    view_repr = repr(self.data.geom('mybox'))
+    self.assertStartsWith(view_repr, '<_MjDataGeomViews')
+    self.assertIn("name: 'mybox'", view_repr)
+    self.assertIn('xpos:', view_repr)
+    self.assertEndsWith(view_repr, '>')
 
   def test_named_indexing_body_repr_in_data(self):
     view_repr = repr(self.data.body('mybox'))
     self.assertStartsWith(view_repr, '<_MjDataBodyViews')
-    self.assertIn('xpos: array([0., 0., 0.])', view_repr)
+    self.assertIn('xpos:', view_repr)
     self.assertEndsWith(view_repr, '>')
 
   def test_named_indexing_repr_in_model(self):
     view_repr = repr(self.model.geom('mybox'))
     self.assertStartsWith(view_repr, '<_MjModelGeomViews')
-    self.assertIn('size: array([0.1, 0.1, 0.1])', view_repr)
+    self.assertIn('size:', view_repr)
     self.assertEndsWith(view_repr, '>')
 
   def test_addresses_differ_between_structs(self):
@@ -385,25 +385,25 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     self.assertIn('MjVisual', vis_repr)
 
   def test_mjmodel_can_read_and_write_opt(self):
-    self.assertEqual(self.model.opt.timestep, 0.002)
-    np.testing.assert_array_equal(self.model.opt.gravity, [0, 0, -9.81])
+    np.testing.assert_allclose(self.model.opt.timestep, 0.002)
+    np.testing.assert_allclose(self.model.opt.gravity, [0, 0, -9.81])
 
     opt = self.model.opt
     self.model.opt.timestep = 0.001
-    self.assertEqual(opt.timestep, 0.001)
+    np.testing.assert_allclose(opt.timestep, 0.001)
 
     gravity = opt.gravity
     self.model.opt.gravity[1] = 0.1
-    np.testing.assert_array_equal(gravity, [0, 0.1, -9.81])
+    np.testing.assert_allclose(gravity, [0, 0.1, -9.81])
     self.model.opt.gravity = 0.2
-    np.testing.assert_array_equal(gravity, [0.2, 0.2, 0.2])
+    np.testing.assert_allclose(gravity, [0.2, 0.2, 0.2])
 
   def test_mjmodel_can_read_and_write_stat(self):
     self.assertNotEqual(self.model.stat.meanmass, 0)
 
     stat = self.model.stat
     self.model.stat.meanmass = 1.2
-    self.assertEqual(stat.meanmass, 1.2)
+    np.testing.assert_allclose(stat.meanmass, 1.2)
 
   def test_mjmodel_can_read_and_write_vis(self):
     self.assertEqual(self.model.vis.quality.shadowsize, 51)
@@ -435,8 +435,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
   def test_mjoption_can_make_default(self):
     opt = mujoco.MjOption()
-    self.assertEqual(opt.timestep, 0.002)
-    np.testing.assert_array_equal(opt.gravity, [0, 0, -9.81])
+    np.testing.assert_allclose(opt.timestep, 0.002)
+    np.testing.assert_allclose(opt.gravity, [0, 0, -9.81])
 
   def test_mjoption_can_copy(self):
     opt1 = mujoco.MjOption()
@@ -444,14 +444,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     opt1.gravity = 2
 
     opt2 = copy.copy(opt1)
-    self.assertEqual(opt2.timestep, 0.001)
-    np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
+    np.testing.assert_allclose(opt2.timestep, 0.001)
+    np.testing.assert_allclose(opt2.gravity, [2, 2, 2])
 
     # Make sure opt2 is actually a copy.
     opt1.timestep = 0.005
     opt1.gravity = 5
-    self.assertEqual(opt2.timestep, 0.001)
-    np.testing.assert_array_equal(opt2.gravity, [2, 2, 2])
+    np.testing.assert_allclose(opt2.timestep, 0.001)
+    np.testing.assert_allclose(opt2.gravity, [2, 2, 2])
 
   def test_mjmodel_can_copy(self):
     model_copy = copy.copy(self.model)
@@ -469,8 +469,8 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     # Make sure it's a copy.
     self.model.geom_size[1] = 0.5
-    np.testing.assert_array_equal(self.model.geom_size[1], [0.5, 0.5, 0.5])
-    np.testing.assert_array_equal(model_copy.geom_size[1], [0.1, 0.1, 0.1])
+    np.testing.assert_allclose(self.model.geom_size[1], [0.5, 0.5, 0.5])
+    np.testing.assert_allclose(model_copy.geom_size[1], [0.1, 0.1, 0.1])
 
   def test_mjdata_can_copy(self):
     self.data.qpos = [0, 0, 0.1*np.sqrt(2) - 0.001,
@@ -598,21 +598,21 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     expected_pos = []
     for contact in self.data.contact:
-      expected_pos.append(np.random.uniform(size=3))
+      expected_pos.append(np.array(np.random.uniform(size=3), dtype=DTYPE))
       contact.pos = expected_pos[-1]
     self.assertLen(expected_pos, expected_ncon)
     np.testing.assert_array_equal(self.data.contact.pos, expected_pos)
 
     expected_friction = []
     for contact in self.data.contact:
-      expected_friction.append(np.random.uniform(size=5))
+      expected_friction.append(np.array(np.random.uniform(size=5), dtype=DTYPE))
       contact.friction = expected_friction[-1]
     self.assertLen(expected_friction, expected_ncon)
     np.testing.assert_array_equal(self.data.contact.friction, expected_friction)
 
     expected_H = []  # pylint: disable=invalid-name
     for contact in self.data.contact:
-      expected_H.append(np.random.uniform(size=36))
+      expected_H.append(np.array(np.random.uniform(size=36), dtype=DTYPE))
       contact.H = expected_H[-1]
     self.assertLen(expected_H, expected_ncon)
     np.testing.assert_array_equal(self.data.contact.H, expected_H)
@@ -749,7 +749,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     expected = np.array([1, 1, 0]) / np.sqrt(2)
 
     # Check that the output argument works, and that the binding returns None.
-    res = np.empty(3, np.float64)
+    res = np.empty(3, DTYPE)
     self.assertIsNone(mujoco.mju_rotVecQuat(res, vec, quat))
     np.testing.assert_allclose(res, expected)
 
@@ -759,7 +759,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
 
     # Check that the res argument must have the right size.
     with self.assertRaises(TypeError):
-      mujoco.mju_rotVecQuat(np.zeros(4, np.float64), vec, quat)
+      mujoco.mju_rotVecQuat(np.zeros(4, DTYPE), vec, quat)
 
     # Check that the vec argument must have the right size.
     with self.assertRaises(TypeError):
@@ -792,24 +792,24 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     sig = mujoco.mjtState.mjSTATE_INTEGRATION
     size = mujoco.mj_stateSize(self.model, sig)
 
-    state_bad_size = np.empty(size + 1, np.float64)
+    state_bad_size = np.empty(size + 1, DTYPE)
     expected_message = 'state size should equal mj_stateSize(m, sig)'
     with self.assertRaisesWithLiteralMatch(TypeError, expected_message):
       mujoco.mj_getState(self.model, self.data, state_bad_size, sig)
 
     # Get initial state.
-    state0 = np.empty(size, np.float64)
+    state0 = np.empty(size, DTYPE)
     mujoco.mj_getState(self.model, self.data, state0, sig)
 
     # Step, get next state.
     mujoco.mj_step(self.model, self.data)
-    state1a = np.empty(size, np.float64)
+    state1a = np.empty(size, DTYPE)
     mujoco.mj_getState(self.model, self.data, state1a, sig)
 
     # Reset to initial state, step again, get state again.
     mujoco.mj_setState(self.model, self.data, state0, sig)
     mujoco.mj_step(self.model, self.data)
-    state1b = np.empty(size, np.float64)
+    state1b = np.empty(size, DTYPE)
     mujoco.mj_getState(self.model, self.data, state1b, sig)
 
     # Expect next states to be equal.
@@ -818,7 +818,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     # Test mj_copyState
     data2 = mujoco.MjData(self.model)
     mujoco.mj_copyState(self.model, self.data, data2, sig)
-    state1c = np.empty(size, np.float64)
+    state1c = np.empty(size, DTYPE)
     mujoco.mj_getState(self.model, data2, state1c, sig)
     np.testing.assert_array_equal(state1a, state1c)
 
@@ -853,14 +853,14 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     np.testing.assert_array_equal(act, self.data.act)
 
   def test_mj_angmomMat(self):  # pylint: disable=invalid-name
-    self.data.qvel = np.ones(self.model.nv, np.float64)
+    self.data.qvel = np.ones(self.model.nv, DTYPE)
     mujoco.mj_forward(self.model, self.data)
     mujoco.mj_subtreeVel(self.model, self.data)
 
-    mat = np.empty((3, 10), np.float64)
+    mat = np.empty((3, 10), DTYPE)
     mujoco.mj_angmomMat(self.model, self.data, mat, 0)
-    np.testing.assert_almost_equal(
-        mat @ self.data.qvel, self.data.subtree_angmom[0, :]
+    np.testing.assert_allclose(
+        mat @ self.data.qvel, self.data.subtree_angmom[0, :], atol=1e-5
     )
 
   def test_mj_jacSite(self):  # pylint: disable=invalid-name
@@ -868,7 +868,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     site_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, 'mysite')
 
     # Call mj_jacSite with only jacp.
-    jacp = np.empty((3, 10), np.float64)
+    jacp = np.empty((3, 10), DTYPE)
     mujoco.mj_jacSite(self.model, self.data, jacp, None, site_id)
 
     expected_jacp = np.array([
@@ -879,7 +879,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
     np.testing.assert_array_equal(jacp, expected_jacp)
 
     # Call mj_jacSite with only jacr.
-    jacr = np.empty((3, 10), np.float64)
+    jacr = np.empty((3, 10), DTYPE)
     mujoco.mj_jacSite(self.model, self.data, None, jacr, site_id)
 
     expected_jacr = np.array([
@@ -941,8 +941,8 @@ Euler integrator, semi-implicit in velocity.
     )
 
   def test_float_constant(self):
-    self.assertEqual(mujoco.mjMAXVAL, 1e10)
-    self.assertEqual(mujoco.mjMINVAL, 1e-15)
+    np.testing.assert_allclose(mujoco.mjMAXVAL, 1e10, rtol=1e-6)
+    np.testing.assert_allclose(mujoco.mjMINVAL, 1e-15, rtol=1e-6)
 
   def test_string_constants(self):
     self.assertLen(mujoco.mjDISABLESTRING, mujoco.mjtDisableBit.mjNDISABLE)
@@ -1306,7 +1306,7 @@ Euler integrator, semi-implicit in velocity.
         geomid,
     )
     # Check that named arguments work
-    normal = np.zeros(3, np.float64)
+    normal = np.zeros(3, DTYPE)
     mujoco.mj_ray(
         m=self.model,
         d=self.data,
@@ -1321,14 +1321,14 @@ Euler integrator, semi-implicit in velocity.
 
   def test_mju_ray_geom(self):
     # Test mju_rayGeom with a plane at origin
-    pos = np.zeros(3)
-    mat = np.eye(3).flatten()
-    size = np.array([10.0, 10.0, 1.0])
-    pnt = np.array([5.0, 5.0, 5.0])
+    pos = np.zeros(3, DTYPE)
+    mat = np.eye(3, dtype=DTYPE).flatten()
+    size = np.array([10.0, 10.0, 1.0], dtype=DTYPE)
+    pnt = np.array([5.0, 5.0, 5.0], dtype=DTYPE)
     # Normalize direction for Euclidean distance
-    vec = np.array([-1.0, -1.0, -1.0])
+    vec = np.array([-1.0, -1.0, -1.0], dtype=DTYPE)
     vec = vec / np.linalg.norm(vec)
-    normal = np.zeros(3)
+    normal = np.zeros(3, DTYPE)
 
     # Call with normal argument
     dist = mujoco.mju_rayGeom(
@@ -1347,12 +1347,12 @@ Euler integrator, semi-implicit in velocity.
   def test_mj_multi_ray(self):
     nray = 3
     geom1 = np.zeros(1, np.int32)
-    pnt = np.array([-0.3, 0, 0.1])
-    vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], np.float64)
+    pnt = np.array([-0.3, 0, 0.1], DTYPE)
+    vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], DTYPE)
     dist_ex = np.array([0.2, -1, 0.1])
     geom_ex = np.array([1, -1, 0])
     geomid = np.zeros(nray, np.int32)
-    dist = np.zeros(nray, np.float64)
+    dist = np.zeros(nray, DTYPE)
 
     mujoco.mj_forward(self.model, self.data)
     mujoco.mj_multiRay(
@@ -1400,23 +1400,23 @@ Euler integrator, semi-implicit in velocity.
 
   def test_geom_distance(self):
     mujoco.mj_forward(self.model, self.data)
-    fromto = np.empty(6, np.float64)
+    fromto = np.empty(6, DTYPE)
     dist = mujoco.mj_geomDistance(self.model, self.data, 0, 2, 200, fromto)
-    self.assertEqual(dist, 41.9)
-    np.testing.assert_array_equal(
-        fromto, np.array((42.0, 0.0, 0.0, 42.0, 0.0, 41.9))
+    np.testing.assert_allclose(dist, 41.9, atol=1e-5)
+    np.testing.assert_allclose(
+        fromto, np.array((42.0, 0.0, 0.0, 42.0, 0.0, 41.9)), atol=1e-5
     )
 
   def test_inverse_fd(self):
     eps = 1e-6
     flg_centered = 0
-    df_dq = np.zeros((self.model.nv, self.model.nv))
-    df_dv = np.zeros((self.model.nv, self.model.nv))
-    df_da = np.zeros((self.model.nv, self.model.nv))
-    ds_dq = np.zeros((self.model.nv, self.model.nsensordata))
-    ds_dv = np.zeros((self.model.nv, self.model.nsensordata))
-    ds_da = np.zeros((self.model.nv, self.model.nsensordata))
-    dm_dq = np.zeros((self.model.nv, self.model.nM))
+    df_dq = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE)
+    df_dv = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE)
+    df_da = np.zeros((self.model.nv, self.model.nv), dtype=DTYPE)
+    ds_dq = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE)
+    ds_dv = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE)
+    ds_da = np.zeros((self.model.nv, self.model.nsensordata), dtype=DTYPE)
+    dm_dq = np.zeros((self.model.nv, self.model.nM), dtype=DTYPE)
     mujoco.mjd_inverseFD(
         self.model,
         self.data,
@@ -1440,11 +1440,11 @@ Euler integrator, semi-implicit in velocity.
   def test_mjd_sub_quat(self):
     quat1 = np.array((0.2, 0.3, 0.3, 0.4))
     quat2 = np.array((0.1, 0.2, 0.4, 0.5))
-    d1 = np.empty(9, np.float64)
-    d2 = np.empty(9, np.float64)
+    d1 = np.empty(9, DTYPE)
+    d2 = np.empty(9, DTYPE)
     mujoco.mjd_subQuat(quat1, quat2, d1, d2)
-    d3 = np.empty((3, 3), np.float64)
-    d4 = np.empty((3, 3), np.float64)
+    d3 = np.empty((3, 3), DTYPE)
+    d4 = np.empty((3, 3), DTYPE)
     mujoco.mjd_subQuat(quat1, quat2, None, d3)
     mujoco.mjd_subQuat(quat1, quat2, d4, None)
     np.testing.assert_array_equal(d2, d3.flatten())
@@ -1453,9 +1453,9 @@ Euler integrator, semi-implicit in velocity.
   def test_mjd_quat_integrate(self):
     scale = 0.1
     vel = np.array((0.2, 0.3, 0.3))
-    d_quat = np.empty(9, np.float64)
-    d_vel = np.empty(9, np.float64)
-    d_h = np.empty(3, np.float64)
+    d_quat = np.empty(9, DTYPE)
+    d_vel = np.empty(9, DTYPE)
+    d_h = np.empty(3, DTYPE)
     mujoco.mjd_quatIntegrate(vel, scale, d_quat, d_vel, d_h)
 
   def test_banded(self):
@@ -1467,23 +1467,23 @@ Euler integrator, semi-implicit in velocity.
         [0, 2.0, 0, 0.2],
         [0, 0, 3.0, 0.3],
         [0.1, 0.2, 0.3, 4.0],
-    ])
-    band = np.zeros(n_band * (n_total - n_dense) + n_dense * n_total)
+    ], dtype=DTYPE)
+    band = np.zeros(n_band * (n_total - n_dense) + n_dense * n_total, dtype=DTYPE)
     mujoco.mju_dense2Band(band, dense, n_total, n_band, n_dense)
     for i in range(4):
       index = mujoco.mju_bandDiag(i, n_total, n_band, n_dense)
       self.assertEqual(band[index], i + 1)
-    dense2 = np.zeros((n_total, n_total))
+    dense2 = np.zeros((n_total, n_total), dtype=DTYPE)
     flg_sym = 1
     mujoco.mju_band2Dense(dense2, band, n_total, n_band, n_dense, flg_sym)
     np.testing.assert_array_equal(dense, dense2)
-    vec = np.array([[2.0], [2.0], [3.0], [4.0]])
+    vec = np.array([[2.0], [2.0], [3.0], [4.0]], dtype=DTYPE)
     res = np.zeros_like(vec)
     n_vec = 1
     mujoco.mju_bandMulMatVec(
         res, band, vec, n_total, n_band, n_dense, n_vec, flg_sym
     )
-    np.testing.assert_array_equal(res, dense @ vec)
+    np.testing.assert_allclose(res, dense @ vec)
     diag_add = 0
     diag_mul = 0
     mujoco.mju_cholFactorBand(
@@ -1494,31 +1494,31 @@ Euler integrator, semi-implicit in velocity.
 
   def test_mju_box_qp(self):
     n = 5
-    res = np.zeros(n)
-    r = np.zeros((n, n + 7))
+    res = np.zeros(n, DTYPE)
+    r = np.zeros((n, n + 7), DTYPE)
     index = np.zeros(n, np.int32)
-    h = np.eye(n)
-    g = np.ones((n,))
-    lower = -np.ones((n,))
-    upper = np.ones((n,))
+    h = np.eye(n, dtype=DTYPE)
+    g = np.ones((n,), dtype=DTYPE)
+    lower = -np.ones((n,), dtype=DTYPE)
+    upper = np.ones((n,), dtype=DTYPE)
     rank = mujoco.mju_boxQP(res, r, index, h, g, lower, upper)
     self.assertGreater(rank, -1)
 
   def test_mju_fill(self):
-    res = np.empty(3, np.float64)
+    res = np.empty(3, DTYPE)
     mujoco.mju_fill(res, 1.5)
-    np.testing.assert_array_equal(res, np.full(3, 1.5))
+    np.testing.assert_allclose(res, np.full(3, 1.5))
 
   def test_mju_eye(self):
-    eye4 = np.empty((4, 4), np.float64)
+    eye4 = np.empty((4, 4), DTYPE)
     mujoco.mju_eye(eye4)
     np.testing.assert_array_equal(eye4, np.eye(4))
 
   def test_mju_symmetrize(self):
-    mat = np.linspace(0, 1, 16).reshape(4, 4)
-    res = np.empty((4, 4), np.float64)
+    mat = np.linspace(0, 1, 16, dtype=DTYPE).reshape(4, 4)
+    res = np.empty((4, 4), DTYPE)
     mujoco.mju_symmetrize(res, mat)
-    np.testing.assert_array_equal(res, 0.5 * (mat + mat.T))
+    np.testing.assert_allclose(res, 0.5 * (mat + mat.T))
 
   def test_mju_clip(self):
     self.assertEqual(mujoco.mju_clip(1.5, 1.0, 2.0), 1.5)
@@ -1532,12 +1532,12 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(mujoco.mju_mulVecMatVec(vec1, mat, vec2), 204.0)
 
   def test_mju_dense_to_sparse(self):
-    mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
+    mat = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]], dtype=DTYPE)
     expected_vals = np.array([1.0, 2.0, 3.0])
     expected_rownnz = np.array([1, 2])
     expected_rowadr = np.array([0, 1])
     expected_colind = np.array([1, 0, 2])
-    vals = np.zeros(3)
+    vals = np.zeros(3, DTYPE)
     row_nnz = np.zeros(2, np.int32)
     row_adr = np.zeros(2, np.int32)
     col_ind = np.zeros(3, np.int32)
@@ -1550,21 +1550,21 @@ Euler integrator, semi-implicit in velocity.
 
   def test_mju_sparse_to_dense(self):
     expected = np.array([[0.0, 1.0, 0.0], [2.0, 0.0, 3.0]])
-    mat = np.array((1.0, 2.0, 3.0))
+    mat = np.array((1.0, 2.0, 3.0), dtype=DTYPE)
     rownnz = np.array([1, 2])
     rowadr = np.array([0, 1])
     colind = np.array([1, 0, 2])
-    res = np.zeros((2, 3))
+    res = np.zeros((2, 3), dtype=DTYPE)
     mujoco.mju_sparse2dense(res, mat, rownnz, rowadr, colind)
     np.testing.assert_array_equal(res, expected)
 
   def test_mju_euler_to_quat(self):
-    quat = np.zeros(4)
-    euler = np.array([0, np.pi / 2, 0])
+    quat = np.zeros(4, DTYPE)
+    euler = np.array([0, np.pi / 2, 0], DTYPE)
     seq = 'xyz'
     mujoco.mju_euler2Quat(quat, euler, seq)
     expected_quat = np.array([np.sqrt(0.5), 0, np.sqrt(0.5), 0.0])
-    np.testing.assert_almost_equal(quat, expected_quat)
+    np.testing.assert_almost_equal(quat, expected_quat, decimal=5)
 
     error = 'mju_euler2Quat: seq must contain exactly 3 characters'
     with self.assertRaisesWithLiteralMatch(mujoco.FatalError, error):
@@ -1798,21 +1798,21 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(nhistory, 3)
     self.assertEqual(dim, 3)  # accelerometer has dim=3
     times = np.array([0.0, 0.01, 0.02])
-    values = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float64)
+    values = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=DTYPE)
     mujoco.mj_initSensorHistory(model, data, 0, times, values, phase=0.0)
 
     # Read back a value using zero-order hold
     # mj_readSensor auto-subtracts delay: lookup_time = read_time - delay
     # delay = 0.01, so:
     #   read_time=0.02 -> lookup at 0.01 -> value [4, 5, 6]
-    result = np.zeros(dim)
+    result = np.zeros(dim, DTYPE)
     mujoco.mj_readSensor(model, data, 0, 0.02, result, interp=0)
     # ZOH returns value at t=0.01
     np.testing.assert_array_equal(result, [4, 5, 6])
 
     # Test with times=None (uses existing timestamps)
     new_values = np.array([
-        [10, 11, 12], [13, 14, 15], [16, 17, 18]], dtype=np.float64)
+        [10, 11, 12], [13, 14, 15], [16, 17, 18]], dtype=DTYPE)
     mujoco.mj_initSensorHistory(model, data, 0, None, new_values, phase=0.0)
     # read_time=0.02 -> lookup at 0.01 -> value [13, 14, 15]
     mujoco.mj_readSensor(model, data, 0, 0.02, result, interp=0)
@@ -1861,7 +1861,7 @@ Euler integrator, semi-implicit in velocity.
 
     # After mj_makeData, reading from the delay buffer gives all zeros.
     # For a quaternion sensor, this is invalid data.
-    result = np.zeros(dim)
+    result = np.zeros(dim, DTYPE)
     mujoco.mj_readSensor(model, data, 0, delay, result, interp=0)
     np.testing.assert_array_equal(result, [0, 0, 0, 0])
 
@@ -1883,7 +1883,7 @@ Euler integrator, semi-implicit in velocity.
 
     # Now reading from the delay buffer gives the valid identity quaternion.
     mujoco.mj_readSensor(model, data, 0, delay, result, interp=0)
-    np.testing.assert_array_equal(result, [1, 0, 0, 0])
+    np.testing.assert_allclose(result, [1, 0, 0, 0])
 
   def test_poly_stiffness(self):
     xml = r"""
diff --git a/python/mujoco/codegen/generate_spec_bindings.py b/python/mujoco/codegen/generate_spec_bindings.py
index 5dee5684..2b1d4b14 100644
--- a/python/mujoco/codegen/generate_spec_bindings.py
+++ b/python/mujoco/codegen/generate_spec_bindings.py
@@ -156,8 +156,10 @@ def _array_binding_code(
   if classname == 'mjSpec':  # raw mjSpec has a wrapper
     rawclassname = classname.replace('mjS', 'MjS')
     fullvarname = 'ptr->' + varname
-  if innertype == 'double' or innertype == 'mjtNum':
-    innertype = 'MjDouble'  # custom Eigen type
+  if innertype == 'mjtNum':
+    innertype = 'MjNum'  # custom Eigen type for mjtNum fields
+  elif innertype == 'double':
+    innertype = 'MjDouble'  # custom Eigen type for double fields
   elif innertype == 'float':
     innertype = 'MjFloat'  # custom Eigen type
   elif innertype == 'int':
diff --git a/python/mujoco/constants.cc b/python/mujoco/constants.cc
index 593070d2..023b78cb 100644
--- a/python/mujoco/constants.cc
+++ b/python/mujoco/constants.cc
@@ -15,6 +15,7 @@
 #include <utility>
 #include <vector>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtype.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
 #include <pybind11/cast.h>
@@ -94,6 +95,7 @@ PYBIND11_MODULE(_constants, pymodule) {
   pymodule.attr("mjFRAMESTRING") = MakeTuple(mjFRAMESTRING);
   pymodule.attr("mjVISSTRING") = MakeTuple(mjVISSTRING);
   pymodule.attr("mjRNDSTRING") = MakeTuple(mjRNDSTRING);
+  pymodule.attr("MJTNUM_BYTES") = sizeof(mjtNum);
 }
 }  // namespace
 }  // namespace mujoco::python
diff --git a/python/mujoco/serialization.h b/python/mujoco/serialization.h
index cb234ac8..b387b455 100644
--- a/python/mujoco/serialization.h
+++ b/python/mujoco/serialization.h
@@ -17,6 +17,7 @@
 
 #include <cstddef>
 #include <iostream>
+#include <type_traits>
 
 #include <mujoco/mjtype.h>
 
@@ -30,7 +31,8 @@ namespace mujoco::python::_impl {
 
 static_assert(sizeof(char) == 1);
 static_assert(sizeof(int) == 4);
-static_assert(sizeof(mjtNum) == 8);
+static_assert((std::is_same_v<mjtNum, float> && sizeof(mjtNum) == 4) ||
+              (std::is_same_v<mjtNum, double> && sizeof(mjtNum) == 8));
 
 inline void WriteChar(std::ostream& output, char c) {
   output.write(&c, 1);
diff --git a/python/mujoco/specs.cc b/python/mujoco/specs.cc
index 738ab73b..f8abf3ad 100644
--- a/python/mujoco/specs.cc
+++ b/python/mujoco/specs.cc
@@ -59,6 +59,9 @@ using MjDouble9 = Eigen::Map<Eigen::Matrix<double, 9, 1>>;
 using MjDouble10 = Eigen::Map<Eigen::Matrix<double, 10, 1>>;
 using MjDouble11 = Eigen::Map<Eigen::Matrix<double, 11, 1>>;
 using MjDoubleVec = Eigen::Map<Eigen::VectorXd>;
+using MjNum2 = Eigen::Map<Eigen::Matrix<mjtNum, 2, 1>>;
+using MjNum3 = Eigen::Map<Eigen::Matrix<mjtNum, 3, 1>>;
+using MjNum5 = Eigen::Map<Eigen::Matrix<mjtNum, 5, 1>>;
 
 using MjIntRef2 = Eigen::Ref<const Eigen::Vector2i>;
 using MjIntRef3 = Eigen::Ref<const Eigen::Vector3i>;
@@ -74,6 +77,9 @@ using MjDoubleRef9 = Eigen::Ref<const Eigen::Matrix<double, 9, 1>>;
 using MjDoubleRef10 = Eigen::Ref<const Eigen::Matrix<double, 10, 1>>;
 using MjDoubleRef11 = Eigen::Ref<const Eigen::Matrix<double, 11, 1>>;
 using MjDoubleRefVec = Eigen::Ref<const Eigen::VectorXd>;
+using MjNumRef2 = Eigen::Ref<const Eigen::Matrix<mjtNum, 2, 1>>;
+using MjNumRef3 = Eigen::Ref<const Eigen::Matrix<mjtNum, 3, 1>>;
+using MjNumRef5 = Eigen::Ref<const Eigen::Matrix<mjtNum, 5, 1>>;
 
 template <typename LoadFunc>
 static raw::MjSpec* LoadSpecFileImpl(

From c7d340be4f71a9ba11148acc6354f91a1b0ccdff Mon Sep 17 00:00:00 2001
From: Kyle Bayes <kylebayes@google.com>
Date: Fri, 29 May 2026 03:52:24 -0700
Subject: [PATCH 53/61] Cap box<->box collisions.

PiperOrigin-RevId: 923312868
Change-Id: I801d76cf1d1d7b9ed1535d3e3c723a5e8af4162a
---
 src/engine/engine_collision_box.c | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/src/engine/engine_collision_box.c b/src/engine/engine_collision_box.c
index 6cf09568..f1a96934 100644
--- a/src/engine/engine_collision_box.c
+++ b/src/engine/engine_collision_box.c
@@ -1402,13 +1402,15 @@ int mjc_BoxBox(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2, m
   }
 
   // consolidate good
-  int i = 0;
+  int ncon = 0;
   for (int j=0; j < num; j++) {
     if (dupe[j] == 0) {
-      con[i] = tmp[j];
-      i++;
+      con[ncon++] = tmp[j];
+      if (ncon >= 8) {
+        break;
+      }
     }
   }
 
-  return i;
+  return ncon;
 }

From 8eaf174a88341f2c8abb32e7fa3e2e383c407a85 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Fri, 29 May 2026 06:30:03 -0700
Subject: [PATCH 54/61] Change "Num threads" UI element to a slider.

PiperOrigin-RevId: 923378753
Change-Id: Ie5f53faf6ad90e0e1a1148f004b1867a57e6fb5e
---
 simulate/simulate.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/simulate/simulate.h b/simulate/simulate.h
index 367d95a3..acb218b5 100644
--- a/simulate/simulate.h
+++ b/simulate/simulate.h
@@ -328,7 +328,7 @@ class Simulate {
   const mjuiDef def_simulation[15] = {
     {mjITEM_SECTION,   "Simulation",    mjPRESERVE, nullptr,     "AS"},
     {mjITEM_RADIO,     "",              5, &this->run,           "Pause\nRun"},
-    {mjITEM_EDITINT,   "Num threads",   5, &this->nthread,       "1"},
+    {mjITEM_SLIDERINT, "Num threads",   5, &this->nthread,       "0 10"},
     {mjITEM_BUTTON,    "Reset",         2, nullptr,              " #259"},
     {mjITEM_BUTTON,    "Reload",        5, nullptr,              "CL"},
     {mjITEM_BUTTON,    "Align",         2, nullptr,              "CA"},

From f63b92da5ec806e47a61e14b46fce9eb46299d27 Mon Sep 17 00:00:00 2001
From: Haroon Qureshi <haroonq@google.com>
Date: Fri, 29 May 2026 08:08:47 -0700
Subject: [PATCH 55/61] Clear draw queue instead of failing if rendering
 skipped.

PiperOrigin-RevId: 923420183
Change-Id: I34aac76810a7bf0e4a0394ebe27b45a8fc32f110
---
 src/experimental/filament/filament/renderable.cc | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/experimental/filament/filament/renderable.cc b/src/experimental/filament/filament/renderable.cc
index 04665f22..93d0e93f 100644
--- a/src/experimental/filament/filament/renderable.cc
+++ b/src/experimental/filament/filament/renderable.cc
@@ -237,11 +237,7 @@ void Renderable::Prepare(std::span<const mjrRenderRequest*> requests) {
   // We assume BindMaterialInstance will be called with the same requests in
   // the same order. As such, we'll just store the draw state in a deque rather
   // than trying to perform any kind of matching with the requests.
-  if (!draw_queue_.empty()) {
-    mju_error(
-        "Expected all previous draw calls to be bound before preparing new "
-        "draw calls.");
-  }
+  draw_queue_.clear();
   curr_state_ = DrawState();
 
   int num_reflections = 0;

From 5d782a2bb8569f2c79059c9845cae5147dd684a2 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Fri, 29 May 2026 13:51:02 -0700
Subject: [PATCH 56/61] Fix memory overallocation in sparse primal solvers.

PiperOrigin-RevId: 923586079
Change-Id: I8d484b840063c0a7a8a79c5bdc9074a57d2ee55f
---
 src/engine/engine_solver.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c
index dd18ae55..b9c112a1 100644
--- a/src/engine/engine_solver.c
+++ b/src/engine/engine_solver.c
@@ -1072,7 +1072,7 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
   // local sizes and flags
   int nv = ctx->nv;
   int nefc = ctx->nefc;
-  int nJ = ctx->is_sparse ? d->nJ : 0;
+  int nJ = ctx->is_sparse ? ctx->nJ : 0;
   int is_sparse = ctx->is_sparse;
   int is_elliptic = ctx->is_elliptic;
 

From 96bf8aea8134dc6afa7503c4cb130d23f7dc163b Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 30 May 2026 03:07:31 -0700
Subject: [PATCH 57/61] Move island-specific sparse matrices from arena to
 stack.

PiperOrigin-RevId: 923850345
Change-Id: I9683d7554b15b7cd8c45a8dce7814640aa266452
---
 doc/changelog.rst                    |   6 +-
 doc/includes/references.h            |  11 ---
 include/mujoco/mjdata.h              |  11 ---
 include/mujoco/mjxmacro.h            |  11 ---
 python/mujoco/introspect/structs.py  |  88 ------------------
 src/engine/engine_island.c           |  35 -------
 src/engine/engine_print.c            |  47 +++++-----
 src/engine/engine_solver.c           | 127 ++++++++++++++++++-------
 test/engine/engine_island_test.cc    | 133 ---------------------------
 unity/Runtime/Bindings/MjBindings.cs |  11 ---
 wasm/codegen/generated/bindings.cc   |  44 ---------
 11 files changed, 122 insertions(+), 402 deletions(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index c5b5377d..5862cf42 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -19,7 +19,11 @@ General
 
      - The header file ``mjthread.h`` was removed along with the engine threading API.
 
-     **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
+       **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
+
+     - Moved island sparse matrix construction from :ref:`mj_island` (single threaded) into :ref:`mj_fwdConstraint`
+       (multi-threaded). The island-specific matrices ``iM, iLD, iefc_J`` were removed from the arena and are now
+       allocated on the stack.
 
 Bug fixes
 ^^^^^^^^^
diff --git a/doc/includes/references.h b/doc/includes/references.h
index 1ce969d7..03f478c5 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -341,12 +341,6 @@ struct mjData_ {
   // 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)
@@ -361,11 +355,6 @@ struct mjData_ {
   // 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)
-  mjtNum* iefc_J;            // constraint Jacobian                              (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)
diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
index 2ee67f95..aab0709a 100644
--- a/include/mujoco/mjdata.h
+++ b/include/mujoco/mjdata.h
@@ -364,12 +364,6 @@ struct mjData_ {
   // 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)
@@ -384,11 +378,6 @@ struct mjData_ {
   // 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)
-  mjtNum* iefc_J;            // constraint Jacobian                              (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)
diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
index 04539f5e..ec3e0da3 100644
--- a/include/mujoco/mjxmacro.h
+++ b/include/mujoco/mjxmacro.h
@@ -969,12 +969,6 @@
     X  ( int,     map_idof2dof,      MJ_M(nv),       1 ) \
     X  ( mjtNum,  ifrc_smooth,       MJ_D(nidof),    1 ) \
     X  ( mjtNum,  iacc_smooth,       MJ_D(nidof),    1 ) \
-    XNV( int,     iM_rownnz,         MJ_D(nidof),    1 ) \
-    XNV( int,     iM_rowadr,         MJ_D(nidof),    1 ) \
-    XNV( int,     iM_colind,         MJ_M(nC),       1 ) \
-    XNV( mjtNum,  iM,                MJ_M(nC),       1 ) \
-    XNV( mjtNum,  iLD,               MJ_M(nC),       1 ) \
-    X  ( mjtNum,  iLDiagInv,         MJ_D(nidof),    1 ) \
     X  ( mjtNum,  iacc,              MJ_D(nidof),    1 ) \
     X  ( int,     efc_island,        MJ_D(nefc),     1 ) \
     X  ( int,     island_ne,         MJ_D(nisland),  1 ) \
@@ -985,11 +979,6 @@
     X  ( int,     map_iefc2efc,      MJ_D(nefc),     1 ) \
     X  ( int,     iefc_type,         MJ_D(nefc),     1 ) \
     X  ( int,     iefc_id,           MJ_D(nefc),     1 ) \
-    XNV( int,     iefc_J_rownnz,     MJ_D(nefc),     1 ) \
-    XNV( int,     iefc_J_rowadr,     MJ_D(nefc),     1 ) \
-    XNV( int,     iefc_J_rowsuper,   MJ_D(nefc),     1 ) \
-    XNV( int,     iefc_J_colind,     MJ_D(nJ),       1 ) \
-    XNV( mjtNum,  iefc_J,            MJ_D(nJ),       1 ) \
     X  ( mjtNum,  iefc_frictionloss, MJ_D(nefc),     1 ) \
     X  ( mjtNum,  iefc_D,            MJ_D(nefc),     1 ) \
     X  ( mjtNum,  iefc_R,            MJ_D(nefc),     1 ) \
diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py
index 8fa54694..78f97fd0 100644
--- a/python/mujoco/introspect/structs.py
+++ b/python/mujoco/introspect/structs.py
@@ -6541,54 +6541,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='unconstrained acceleration',
                  array_extent=('nidof',),
              ),
-             StructFieldDecl(
-                 name='iM_rownnz',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='inertia: non-zeros in each row',
-                 array_extent=('nidof',),
-             ),
-             StructFieldDecl(
-                 name='iM_rowadr',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='inertia: address of each row in iM_colind',
-                 array_extent=('nidof',),
-             ),
-             StructFieldDecl(
-                 name='iM_colind',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='inertia: column indices of non-zeros',
-                 array_extent=('nC',),
-             ),
-             StructFieldDecl(
-                 name='iM',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjtNum'),
-                 ),
-                 doc='total inertia (sparse)',
-                 array_extent=('nC',),
-             ),
-             StructFieldDecl(
-                 name='iLD',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjtNum'),
-                 ),
-                 doc="L'*D*L factorization of M (sparse)",
-                 array_extent=('nC',),
-             ),
-             StructFieldDecl(
-                 name='iLDiagInv',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjtNum'),
-                 ),
-                 doc='1/diag(D)',
-                 array_extent=('nidof',),
-             ),
              StructFieldDecl(
                  name='iacc',
                  type=PointerType(
@@ -6669,46 +6621,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='id of object of specified type',
                  array_extent=('nefc',),
              ),
-             StructFieldDecl(
-                 name='iefc_J_rownnz',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='number of non-zeros in constraint Jacobian row',
-                 array_extent=('nefc',),
-             ),
-             StructFieldDecl(
-                 name='iefc_J_rowadr',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='row start address in colind array',
-                 array_extent=('nefc',),
-             ),
-             StructFieldDecl(
-                 name='iefc_J_rowsuper',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='number of subsequent rows in supernode',
-                 array_extent=('nefc',),
-             ),
-             StructFieldDecl(
-                 name='iefc_J_colind',
-                 type=PointerType(
-                     inner_type=ValueType(name='int'),
-                 ),
-                 doc='column indices in constraint Jacobian',
-                 array_extent=('nJ',),
-             ),
-             StructFieldDecl(
-                 name='iefc_J',
-                 type=PointerType(
-                     inner_type=ValueType(name='mjtNum'),
-                 ),
-                 doc='constraint Jacobian',
-                 array_extent=('nJ',),
-             ),
              StructFieldDecl(
                  name='iefc_frictionloss',
                  type=PointerType(
diff --git a/src/engine/engine_island.c b/src/engine/engine_island.c
index a8aff1ce..88ee12e9 100644
--- a/src/engine/engine_island.c
+++ b/src/engine/engine_island.c
@@ -526,15 +526,6 @@ void mj_island(const mjModel* m, mjData* d) {
     d->island_dofadr[i] = d->map_idof2dof[d->island_idofadr[i]];
   }
 
-  // inertia: block-diagonalize both iLD <- qLD and iM <- M
-  mju_blockDiagSparse(d->iLD, d->iM_rownnz, d->iM_rowadr, d->iM_colind,
-                      d->qLD,  m->M_rownnz, m->M_rowadr, m->M_colind,
-                      nidof, nisland,
-                      d->map_idof2dof, d->map_dof2idof,
-                      d->island_idofadr, d->island_idofadr,
-                      d->iM, d->M);
-  mju_gather(d->iLDiagInv, d->qLDiagInv, d->map_idof2dof, nidof);
-
 
   // ------------------------------------- constraints ---------------------------------------------
 
@@ -578,32 +569,6 @@ void mj_island(const mjModel* m, mjData* d) {
   // SHOULD NOT OCCUR
   if (!mju_compare(island_nefc2, d->island_nefc, nisland)) mjERROR("island_nefc miscount");
 
-  // dense: block-diagonalize Jacobian
-  if (!mj_isSparse(m)) {
-    mju_blockDiag(d->iefc_J, d->efc_J,
-                  nv, nidof, nisland,
-                  d->map_iefc2efc, d->map_idof2dof,
-                  d->island_nefc, d->island_nv,
-                  d->island_iefcadr, d->island_idofadr);
-  }
-
-  // sparse
-  else {
-    // block-diagonalize Jacobian
-    mju_blockDiagSparse(d->iefc_J, d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind,
-                        d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
-                        nefc, nisland,
-                        d->map_iefc2efc, d->map_dof2idof,
-                        d->island_iefcadr, d->island_idofadr, NULL, NULL);
-
-    // recompute rowsuper per island
-    for (int island=0; island < nisland; island++) {
-      int adr = d->island_iefcadr[island];
-      mju_superSparse(d->island_nefc[island], d->iefc_J_rowsuper + adr,
-                      d->iefc_J_rownnz + adr, d->iefc_J_rowadr + adr, d->iefc_J_colind);
-    }
-  }
-
   // copy position-dependent efc vectors required by solver
   mju_gatherInt(d->iefc_type, d->efc_type, d->map_iefc2efc, nefc);
   mju_gatherInt(d->iefc_id, d->efc_id, d->map_iefc2efc, nefc);
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 776d5281..f6ca332b 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -190,11 +190,12 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
 
 
 // print block-diagonal dense matrix, embedded in a larger matrix
-static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
+static void printBlockArray(const char* str, const mjtNum* data, int nc,
                             int nisland, const int* island_nr, const int* island_nc,
                             const int* island_r, const int* island_c,
+                            const int* map_r, const int* map_c,
                             FILE* fp, const char* float_format) {
-  if (!data || !nr || !nc) {
+  if (!data || !nisland) {
     return;
   }
 
@@ -209,7 +210,6 @@ static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
     int bnc = island_nc[b];
     int r_start = island_r[b];
     int c_start = island_c[b];
-    const mjtNum* data_ptr = data + r_start * nc;
 
     // print rows for this block
     for (int r_block = 0; r_block < bnr; r_block++) {
@@ -220,10 +220,13 @@ static void printBlockArray(const char* str, const mjtNum* data, int nr, int nc,
         fprintf(fp, " ");
       }
 
+      int row = map_r[r_start + r_block];
+
       // block data
       for (int c = 0; c < bnc; c++) {
+        int col = map_c[c_start + c];
         fprintf(fp, " ");
-        fprintf(fp, float_format, *data_ptr++);
+        fprintf(fp, float_format, data[row * nc + col]);
       }
 
       // trailing dots
@@ -318,6 +321,7 @@ void mj_printBlockSparsity(const char* str, int nr, int nc, int nisland,
                            const int* island_col_offset,
                            const int* entity_island,
                            const int* map_row_to_entity,
+                           const int* map_col_to_entity,
                            const int* rownnz, const int* rowadr, const int* colind,
                            const int* rowsuper, FILE* fp) {
   // if no rows / columns, or too many columns to be visually useful, return
@@ -343,27 +347,25 @@ void mj_printBlockSparsity(const char* str, int nr, int nc, int nisland,
 
     int c_start = island_col_offset[island];
     int bnc = island_block_ncols[island];
-    int current_nnz = 0;
-    int adr = rowadr[r];
+    int adr = rowadr[entity_r];
+    int nnz = rownnz[entity_r];
     char nz_char = (island < 10) ? ('0' + island) : 'x';
 
-    for (int c = 0; c < nc; c++) {  // c is the global column index
+    for (int c = 0; c < nc; c++) {  // c is the block-space column index
       bool nonzero = false;
       if (c >= c_start && c < c_start + bnc) {
-        int c_block = c - c_start;  // c_block is the island-local column index
-
-        // search for c_block in colind for the current row r
-        while (current_nnz < rownnz[r] && colind[adr + current_nnz] < c_block) {
-          current_nnz++;
-        }
-        if (current_nnz < rownnz[r] && colind[adr + current_nnz] == c_block) {
-          nonzero = true;
+        int target_col = map_col_to_entity[c];
+        for (int i = 0; i < nnz; i++) {
+          if (colind[adr + i] == target_col) {
+            nonzero = true;
+            break;
+          }
         }
       }
       fprintf(fp, "%c", nonzero ? nz_char : ' ');
     }
     fprintf(fp, " |");
-    if (rowsuper && rowsuper[r] > 0) fprintf(fp, " %d", rowsuper[r]);
+    if (rowsuper && rowsuper[entity_r] > 0) fprintf(fp, " %d", rowsuper[entity_r]);
     fprintf(fp, "\n");
   }
   for (int c = 0; c < nc + 2; c++) fprintf(fp, "-");
@@ -1477,8 +1479,8 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
     mj_printBlockSparsity("iM: block-diagonal inertia (nnzs are island ids)",
                           d->nidof, d->nidof, d->nisland,
                           d->island_nv, d->island_idofadr,
-                          d->dof_island, d->map_idof2dof,
-                          d->iM_rownnz, d->iM_rowadr, d->iM_colind, NULL, fp);
+                          d->dof_island, d->map_idof2dof, d->map_idof2dof,
+                          m->M_rownnz, m->M_rowadr, m->M_colind, NULL, fp);
   }
 
   if (!mju_isZero(d->qHDiagInv, m->nv)) {
@@ -1555,9 +1557,10 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
   if (!mj_isSparse(m)) {
     printArray2d("EFC_J", d->nefc, m->nv, d->efc_J, fp, float_format);
     if (d->nisland) {
-      printBlockArray("IEFC_J", d->iefc_J, d->nefc, d->nidof,
+      printBlockArray("IEFC_J", d->efc_J, m->nv,
                       d->nisland, d->island_nefc, d->island_nv,
                       d->island_iefcadr, d->island_idofadr,
+                      d->map_iefc2efc, d->map_idof2dof,
                       fp, float_format);
     }
     printArray2d("EFC_AR", d->nefc, d->nefc, d->efc_AR, fp, float_format);
@@ -1577,9 +1580,9 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
       mj_printBlockSparsity("IEFC_J: block-diagonalized constraint Jacobian (nnzs are island ids)",
                             d->nefc, d->nidof, d->nisland,
                             d->island_nv, d->island_idofadr,
-                            d->efc_island, d->map_iefc2efc,
-                            d->iefc_J_rownnz, d->iefc_J_rowadr, d->iefc_J_colind,
-                            d->iefc_J_rowsuper, fp);
+                            d->efc_island, d->map_iefc2efc, d->map_idof2dof,
+                            d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
+                            d->efc_J_rowsuper, fp);
     }
 
     if (mj_isDual(m)) {
diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c
index b9c112a1..5c7fa0ea 100644
--- a/src/engine/engine_solver.c
+++ b/src/engine/engine_solver.c
@@ -877,12 +877,12 @@ typedef struct {
   mjtNum* qacc;
 
   // inertia
-  const int* M_rownnz;
-  const int* M_rowadr;
-  const int* M_colind;
-  const mjtNum* M;
-  const mjtNum* qLD;
-  const mjtNum* qLDiagInv;
+  int* M_rownnz;
+  int* M_rowadr;
+  int* M_colind;
+  mjtNum* M;
+  mjtNum* qLD;
+  mjtNum* qLDiagInv;
 
   // efc arrays
   const mjtNum* efc_D;
@@ -895,11 +895,11 @@ typedef struct {
   int* efc_state;
 
   // Jacobians
-  const int* J_rownnz;
-  const int* J_rowadr;
-  const int* J_rowsuper;
-  const int* J_colind;
-  const mjtNum* J;
+  int* J_rownnz;
+  int* J_rowadr;
+  int* J_rowsuper;
+  int* J_colind;
+  mjtNum* J;
   int* JT_rownnz;
   int* JT_rowadr;
   int* JT_rowsuper;
@@ -1031,14 +1031,6 @@ static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* c
     ctx->qacc_smooth      = d->iacc_smooth       + idofadr;
     ctx->qacc             = d->iacc              + idofadr;
 
-    // inertia
-    ctx->M_rownnz         = d->iM_rownnz         + idofadr;
-    ctx->M_rowadr         = d->iM_rowadr         + idofadr;
-    ctx->M_colind         = d->iM_colind;
-    ctx->M                = d->iM;
-    ctx->qLD              = d->iLD;
-    ctx->qLDiagInv        = d->iLDiagInv         + idofadr;
-
     // efc arrays
     int iefcadr           = d->island_iefcadr[island];
     ctx->efc_D            = d->iefc_D            + iefcadr;
@@ -1049,32 +1041,44 @@ static void PrimalPointers(const mjModel* m, const mjData* d, mjPrimalContext* c
     ctx->efc_type         = d->iefc_type         + iefcadr;
     ctx->efc_force        = d->iefc_force        + iefcadr;
     ctx->efc_state        = d->iefc_state        + iefcadr;
-
-    // Jacobians
-    if (!ctx->is_sparse) {
-      ctx->J              = d->iefc_J + d->nidof * iefcadr;
-    } else {
-      ctx->J_rownnz       = d->iefc_J_rownnz     + iefcadr;
-      ctx->J_rowadr       = d->iefc_J_rowadr     + iefcadr;
-      ctx->J_rowsuper     = d->iefc_J_rowsuper   + iefcadr;
-      ctx->J_colind       = d->iefc_J_colind;
-      ctx->J              = d->iefc_J;
-      ctx->nJ             = ctx->J_rowadr[ctx->nefc-1] + ctx->J_rownnz[ctx->nefc-1]
-                            - ctx->J_rowadr[0];
-    }
   }
 }
 
 
 // allocate fixed-size arrays in mjPrimalContext
 //  mj_{mark/free}Stack in calling function!
-static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
+static void PrimalAllocate(const mjModel* m, mjData* d, mjPrimalContext* ctx, int flg_Newton) {
   // local sizes and flags
   int nv = ctx->nv;
   int nefc = ctx->nefc;
-  int nJ = ctx->is_sparse ? ctx->nJ : 0;
   int is_sparse = ctx->is_sparse;
   int is_elliptic = ctx->is_elliptic;
+  int nJ = is_sparse ? d->nJ : 0;
+
+  // compute island matrix sizes if needed
+  int nC = 0;
+  if (ctx->island >= 0) {
+    // count nC: number of nonzeros in M block of island (always sparse)
+    int island = ctx->island;
+    int idofadr = d->island_idofadr[island];
+    for (int i = 0; i < nv; i++) {
+      int dof = d->map_idof2dof[idofadr + i];
+      nC += m->M_rownnz[dof];
+    }
+
+    // count nJ: number of nonzeros in J block of island (sparse or dense)
+    if (is_sparse) {
+      nJ = 0;
+      int iefcadr = d->island_iefcadr[island];
+      for (int i = 0; i < nefc; i++) {
+        int efc = d->map_iefc2efc[iefcadr + i];
+        nJ += d->efc_J_rownnz[efc];
+      }
+    } else {
+      nJ = nefc * nv;
+    }
+    ctx->nJ = nJ;
+  }
 
   // compute mjtNum block size
   size_t nNum = 5*nefc + 5*nv;         // common arrays
@@ -1090,6 +1094,11 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
     nNum += 3*nv;                      // CG arrays
   }
 
+  // add island matrix sizes
+  if (ctx->island >= 0) {
+    nNum += 2 * nC + nv + nJ;          // iM, iLD, iLDiagInv, iefc_J
+  }
+
   // compute int block size
   size_t nInt = nefc;                  // oldstate
   if (is_sparse) {
@@ -1097,10 +1106,58 @@ static void PrimalAllocate(mjData* d, mjPrimalContext* ctx, int flg_Newton) {
     if (flg_Newton) nInt += 8*nv;      // Newton sparse
   }
 
+  // add island matrix sizes
+  if (ctx->island >= 0) {
+    nInt += 2 * nv + nC;               // iM_{rownnz, rowadr, colind}
+    if (is_sparse) {
+      nInt += 3 * nefc + nJ;           // iefc_J_{rownnz, rowadr, rowsuper, colind}
+    }
+  }
+
   // allocate mjtNum and int blocks
   mjtNum* numblock = mjSTACKALLOC(d, nNum, mjtNum);
   int* intblock    = mjSTACKALLOC(d, nInt, int);
 
+  // populate island matrices if needed
+  if (ctx->island >= 0) {
+    int island = ctx->island;
+
+    int idofadr = d->island_idofadr[island];
+    int iefcadr = d->island_iefcadr[island];
+
+    ctx->M_rownnz = intblock; intblock += nv;
+    ctx->M_rowadr = intblock; intblock += nv;
+    ctx->M_colind = intblock; intblock += nC;
+
+    ctx->M = numblock; numblock += nC;
+    ctx->qLD = numblock; numblock += nC;
+    ctx->qLDiagInv = numblock; numblock += nv;
+
+    mju_blockSparse(ctx->qLD, ctx->M_rownnz, ctx->M_rowadr, ctx->M_colind,
+                    d->qLD, m->M_rownnz, m->M_rowadr, m->M_colind,
+                    nv, d->map_idof2dof + idofadr, d->map_dof2idof,
+                    d->island_idofadr[island], 0, ctx->M, d->M);
+    mju_gather(ctx->qLDiagInv, d->qLDiagInv, d->map_idof2dof + idofadr, nv);
+
+    ctx->J = numblock; numblock += nJ;
+    if (!is_sparse) {
+      mju_block(ctx->J, d->efc_J, m->nv, nv, nefc,
+                d->map_iefc2efc + iefcadr, d->map_idof2dof + idofadr);
+    } else {
+      ctx->J_rownnz = intblock; intblock += nefc;
+      ctx->J_rowadr = intblock; intblock += nefc;
+      ctx->J_rowsuper = intblock; intblock += nefc;
+      ctx->J_colind = intblock; intblock += nJ;
+
+      mju_blockSparse(ctx->J, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind,
+                      d->efc_J, d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind,
+                      nefc, d->map_iefc2efc + iefcadr, d->map_dof2idof,
+                      d->island_idofadr[island], 0, NULL, NULL);
+
+      mju_superSparse(nefc, ctx->J_rowsuper, ctx->J_rownnz, ctx->J_rowadr, ctx->J_colind);
+    }
+  }
+
   // carve mjtNum block
   ctx->Jaref  = numblock;  numblock += nefc;
   ctx->Jv     = numblock;  numblock += nefc;
@@ -1962,7 +2019,7 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
 
   // make context
   PrimalPointers(m, d, &ctx, island);
-  PrimalAllocate(d, &ctx, flg_Newton);
+  PrimalAllocate(m, d, &ctx, flg_Newton);
 
   // local copies
   int nv   = ctx.nv;
diff --git a/test/engine/engine_island_test.cc b/test/engine/engine_island_test.cc
index e6592b2b..227b350d 100644
--- a/test/engine/engine_island_test.cc
+++ b/test/engine/engine_island_test.cc
@@ -400,139 +400,6 @@ TEST_F(IslandTest, IslandFlex) {
   mj_deleteModel(model);
 }
 
-static const char* const k2H100Path = "engine/testdata/island/2humanoid100.xml";
-
-TEST_F(IslandTest, IslandJacobian) {
-  for (const char* local_path : {kIlslandEfcPath, k2H100Path}) {
-    const std::string xml_path = GetTestDataFilePath(local_path);
-    char error[1024];
-    mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
-    ASSERT_THAT(m, NotNull()) << error;
-    int jac0 = m->opt.jacobian;
-    mjData* d = mj_makeData(m);
-
-    for (mjtNum t_stop : {0.0, 0.2, 2.0}) {
-      while (d->time < t_stop) {
-        mj_step(m, d);
-      }
-
-      for (mjtJacobian jac : {mjJAC_DENSE, mjJAC_SPARSE}) {
-        m->opt.jacobian = jac;
-        mj_forward(m, d);
-
-        int nv = m->nv;
-        int nefc = d->nefc;
-        int nisland = d->nisland;
-        int nidof = d->nidof;
-
-        mjtNum* J = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nv);
-        mjtNum* iJ = (mjtNum*)mju_malloc(sizeof(mjtNum) * nefc * nidof);
-
-        // get local dense Jacobian
-        if (jac == mjJAC_DENSE) {
-          mju_copy(J, d->efc_J, nefc * nv);
-          mju_copy(iJ, d->iefc_J, nefc * nidof);
-        } else {
-          mju_sparse2dense(J, d->efc_J, nefc, nv, d->efc_J_rownnz,
-                           d->efc_J_rowadr, d->efc_J_colind);
-        }
-
-        // compare random access in efc_J to contiguous memory in iefc_J
-        for (int island=0; island < nisland; island++) {
-          int idof = d->island_idofadr[island];
-          int iefc = d->island_iefcadr[island];
-          int nefc_island = d->island_nefc[island];
-          int nv_island = d->island_nv[island];
-
-          // === test J
-
-          // get pointer to J_island, dense (nefc_island x nv_island) submatrix
-          mjtNum* J_island;
-          if (jac == mjJAC_DENSE) {
-            // point to starting address of island in efc_J
-            J_island = iJ + iefc * nidof;
-          } else {
-            // dense copy of island in iJ (here used as scratch)
-            mju_sparse2dense(iJ, d->iefc_J, nefc_island, nv_island,
-                             d->iefc_J_rownnz + iefc,
-                             d->iefc_J_rowadr + iefc,
-                             d->iefc_J_colind);
-            J_island = iJ;
-          }
-
-          // sequential memory in J_island equals random access memory in J
-          for (int i=0; i < nefc_island; i++) {
-            for (int j=0; j < nv_island; j++) {
-              int efc = d->map_iefc2efc[iefc + i];
-              int dof = d->map_idof2dof[idof + j];
-              EXPECT_EQ(J_island[i * nv_island + j], J[efc * nv + dof]);
-            }
-          }
-        }
-
-        mju_free(iJ);
-        mju_free(J);
-      }
-
-      // reset opt.jacobian to initial value
-      m->opt.jacobian = jac0;
-    }
-
-    mj_deleteData(d);
-    mj_deleteModel(m);
-  }
-}
-
-TEST_F(IslandTest, IslandInertia) {
-  for (const char* local_path : {kIlslandEfcPath, k2H100Path}) {
-    const std::string xml_path = GetTestDataFilePath(local_path);
-    char error[1024];
-    mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
-    ASSERT_THAT(m, NotNull()) << error;
-    int nv = m->nv;
-    mjData* d = mj_makeData(m);
-    mjtNum* M = (mjtNum*)mju_malloc(sizeof(mjtNum) * nv * nv);
-
-    for (mjtNum t_stop : {0.0, 0.2, 2.0}) {
-      while (d->time < t_stop) {
-        mj_step(m, d);
-      }
-      mj_forward(m, d);
-
-      int nisland = d->nisland;
-
-      // get dense inertia (lower only)
-      mj_fullM(m, M, d->qM);
-
-      // compare iM sub-matrix to full M
-      for (int island=0; island < nisland; island++) {
-        int nvi = d->island_nv[island];
-        mjtNum* Mi = (mjtNum*)mju_malloc(sizeof(mjtNum) * nvi * nvi);
-
-        int adr = d->island_idofadr[island];
-        mju_sparse2dense(Mi, d->iM, nvi, nvi,
-                         d->iM_rownnz + adr,
-                         d->iM_rowadr + adr,
-                         d->iM_colind);
-
-        // compare Mi to M (lower triangle only)
-        for (int i=0; i < nvi; i++) {
-          for (int j=0; j <= i; j++) {
-            int dofi = d->map_idof2dof[adr + j];
-            int dofj = d->map_idof2dof[adr + i];
-            EXPECT_EQ(Mi[i * nvi + j], M[dofi * nv + dofj]);
-          }
-        }
-        mju_free(Mi);
-      }
-    }
-
-    mju_free(M);
-    mj_deleteData(d);
-    mj_deleteModel(m);
-  }
-}
-
 TEST_F(IslandTest, IslandEfcElliptic) {
   const std::string xml_path = GetTestDataFilePath(kIlslandEfcPath);
   char error[1024];
diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
index e25c0986..e9baa2ea 100644
--- a/unity/Runtime/Bindings/MjBindings.cs
+++ b/unity/Runtime/Bindings/MjBindings.cs
@@ -5752,12 +5752,6 @@ public unsafe struct mjData_ {
   public int* map_idof2dof;
   public double* ifrc_smooth;
   public double* iacc_smooth;
-  public int* iM_rownnz;
-  public int* iM_rowadr;
-  public int* iM_colind;
-  public double* iM;
-  public double* iLD;
-  public double* iLDiagInv;
   public double* iacc;
   public int* efc_island;
   public int* island_ne;
@@ -5768,11 +5762,6 @@ public unsafe struct mjData_ {
   public int* map_iefc2efc;
   public int* iefc_type;
   public int* iefc_id;
-  public int* iefc_J_rownnz;
-  public int* iefc_J_rowadr;
-  public int* iefc_J_rowsuper;
-  public int* iefc_J_colind;
-  public double* iefc_J;
   public double* iefc_frictionloss;
   public double* iefc_D;
   public double* iefc_R;
diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc
index 735c4b1e..eae8ae44 100644
--- a/wasm/codegen/generated/bindings.cc
+++ b/wasm/codegen/generated/bindings.cc
@@ -6981,24 +6981,6 @@ struct MjData {
   emscripten::val iacc_smooth() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->iacc_smooth));
   }
-  emscripten::val iM_rownnz() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->iM_rownnz));
-  }
-  emscripten::val iM_rowadr() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->iM_rowadr));
-  }
-  emscripten::val iM_colind() const {
-    return emscripten::val(emscripten::typed_memory_view(model->nC, ptr_->iM_colind));
-  }
-  emscripten::val iM() const {
-    return emscripten::val(emscripten::typed_memory_view(model->nC, ptr_->iM));
-  }
-  emscripten::val iLD() const {
-    return emscripten::val(emscripten::typed_memory_view(model->nC, ptr_->iLD));
-  }
-  emscripten::val iLDiagInv() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->iLDiagInv));
-  }
   emscripten::val iacc() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nidof, ptr_->iacc));
   }
@@ -7029,21 +7011,6 @@ struct MjData {
   emscripten::val iefc_id() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_id));
   }
-  emscripten::val iefc_J_rownnz() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_J_rownnz));
-  }
-  emscripten::val iefc_J_rowadr() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_J_rowadr));
-  }
-  emscripten::val iefc_J_rowsuper() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_J_rowsuper));
-  }
-  emscripten::val iefc_J_colind() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nJ, ptr_->iefc_J_colind));
-  }
-  emscripten::val iefc_J() const {
-    return emscripten::val(emscripten::typed_memory_view(ptr_->nJ, ptr_->iefc_J));
-  }
   emscripten::val iefc_frictionloss() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nefc, ptr_->iefc_frictionloss));
   }
@@ -11694,20 +11661,9 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("geom_xmat", &MjData::geom_xmat)
     .property("geom_xpos", &MjData::geom_xpos)
     .property("history", &MjData::history)
-    .property("iLD", &MjData::iLD)
-    .property("iLDiagInv", &MjData::iLDiagInv)
-    .property("iM", &MjData::iM)
-    .property("iM_colind", &MjData::iM_colind)
-    .property("iM_rowadr", &MjData::iM_rowadr)
-    .property("iM_rownnz", &MjData::iM_rownnz)
     .property("iacc", &MjData::iacc)
     .property("iacc_smooth", &MjData::iacc_smooth)
     .property("iefc_D", &MjData::iefc_D)
-    .property("iefc_J", &MjData::iefc_J)
-    .property("iefc_J_colind", &MjData::iefc_J_colind)
-    .property("iefc_J_rowadr", &MjData::iefc_J_rowadr)
-    .property("iefc_J_rownnz", &MjData::iefc_J_rownnz)
-    .property("iefc_J_rowsuper", &MjData::iefc_J_rowsuper)
     .property("iefc_R", &MjData::iefc_R)
     .property("iefc_aref", &MjData::iefc_aref)
     .property("iefc_force", &MjData::iefc_force)

From d7e872b6f9f812664489002946588557311596af Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 30 May 2026 03:40:59 -0700
Subject: [PATCH 58/61] Add UI control for MuJoCo threadpool size.

PiperOrigin-RevId: 923860258
Change-Id: I11dc5d7904f1c3657d7e1dd2b0e3c7864c0c7ead
---
 src/experimental/studio/app.cc | 18 ++++++++++++++++++
 src/experimental/studio/app.h  |  2 ++
 2 files changed, 20 insertions(+)

diff --git a/src/experimental/studio/app.cc b/src/experimental/studio/app.cc
index d24df9a8..3c6a64f2 100644
--- a/src/experimental/studio/app.cc
+++ b/src/experimental/studio/app.cc
@@ -28,6 +28,7 @@
 #include <span>
 #include <string>
 #include <string_view>
+#include <thread>
 #include <utility>
 #include <vector>
 
@@ -227,6 +228,7 @@ void App::OnModelLoaded(std::string filename, ModelKind model_kind) {
       plugin->post_model_loaded(plugin, model_path_.c_str());
     }
   });
+  tmp_.update_threadpool = true;
 }
 
 void App::UpdateFilePaths(const std::string& resolved_path) {
@@ -268,6 +270,11 @@ void App::UpdatePhysics() {
     return;
   }
 
+  if (tmp_.update_threadpool) {
+    mju_threadpool(data(), ui_.nthread);
+    tmp_.update_threadpool = false;
+  }
+
   bool stepped = false;
   platform::ForEachPlugin<platform::ModelPlugin>([&](auto* plugin) {
     if (plugin->do_update) {
@@ -1457,6 +1464,15 @@ void App::ToolBarGui() {
     ImGui::SameLine(0, separator_width);
     platform::StepControlGui(model(), &step_control_, tmp_.speed_index);
 
+    ImGui::SameLine(0, separator_width);
+    ImGui::SetNextItemWidth(120);
+    ImGui::BeginDisabled(std::thread::hardware_concurrency() <= 1);
+    if (ImGui::SliderInt("##NumThreads", &ui_.nthread, 0, 8, "%d threads")) {
+      tmp_.update_threadpool = true;
+    }
+    ImGui::EndDisabled();
+    ImGui::SetItemTooltip("%s", "Number of threads in threadpool");
+
     ImGui::TableNextColumn();
 
     platform::CameraSelectionGui(model(), data(), camera_, ui_.camera_idx);
@@ -1823,6 +1839,7 @@ App::UiState::Dict App::UiState::ToDict() const {
       {"font_scale", std::to_string(font_scale)},
       {"window_width", std::to_string(window_width)},
       {"window_height", std::to_string(window_height)},
+      {"nthread", std::to_string(nthread)},
   };
 }
 
@@ -1834,5 +1851,6 @@ void App::UiState::FromDict(const Dict& dict) {
   window_width = ReadIniValue(dict, "window_width", window_width);
   window_height = ReadIniValue(dict, "window_height", window_height);
   font_scale = ReadIniValue(dict, "font_scale", font_scale);
+  nthread = ReadIniValue(dict, "nthread", nthread);
 }
 }  // namespace mujoco::studio
diff --git a/src/experimental/studio/app.h b/src/experimental/studio/app.h
index 9819c059..fa678ad2 100644
--- a/src/experimental/studio/app.h
+++ b/src/experimental/studio/app.h
@@ -114,6 +114,7 @@ class App {
     float font_scale = 1.0f;
     int window_width = 0;
     int window_height = 0;
+    int nthread = 0;
 
     using Dict = std::unordered_map<std::string, std::string>;
     Dict ToDict() const;
@@ -124,6 +125,7 @@ class App {
   struct UiTempState {
     bool should_exit = false;
     bool first_frame = true;
+    bool update_threadpool = false;
 
     // Windows.
     bool help = false;

From 3a6e0ee9b385a52110440ad51542c2888657e13d Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Sat, 30 May 2026 04:11:38 -0700
Subject: [PATCH 59/61] Reduce effect of infinite planes on `center` and
 `extent` statistics

PiperOrigin-RevId: 923868757
Change-Id: I3630d29f05a19e0767549284abbe8a8d9386894b
---
 src/engine/engine_setconst.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c
index 06f5f313..0098feaa 100644
--- a/src/engine/engine_setconst.c
+++ b/src/engine/engine_setconst.c
@@ -1080,7 +1080,7 @@ static void setStat(mjModel* m, mjData* d) {
 
       // infinite in both directions
       else {
-        rbound = 1;
+        rbound = 0.01;
       }
     } else if (m->geom_type[i] == mjGEOM_HFIELD) {
       int j = m->geom_dataid[i];

From b9c7a4b81a186b4fd8ceced592a57e44df35ad1b Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Sat, 30 May 2026 04:47:29 -0700
Subject: [PATCH 60/61] Import NVIDIA/warp 1.13.0 from GitHub.

PiperOrigin-RevId: 923878216
Change-Id: If9df5e0672e956eb62874fb6f4f74ba3a02977b3
---
 mjx/cuda_requirements.txt                     | 10 +++----
 mjx/mujoco/mjx/_src/io.py                     |  4 +--
 .../warp/_src/jax_experimental/custom_call.py | 12 +++++----
 .../warp/_src/jax_experimental/ffi.py         | 27 ++++++++++---------
 mjx/pyproject.toml                            |  2 +-
 5 files changed, 29 insertions(+), 26 deletions(-)

diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt
index 1db7abfb..a726944a 100644
--- a/mjx/cuda_requirements.txt
+++ b/mjx/cuda_requirements.txt
@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
 jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
     --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
     --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
-warp-lang==1.12.1 \
-    --hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
-    --hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
-    --hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
-    --hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
+warp-lang==1.13.0 \
+    --hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
+    --hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
+    --hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
+    --hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index c1e2049e..dc10af8c 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     wp_field.shape = wp_field.shape[1:]
   # warp scalars
   wp_dtype = type(wp_field)
-  if wp_dtype in wp.types.warp_type_to_np_dtype:
-    return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
+  if wp_dtype in wp._src.types.warp_type_to_np_dtype:
+    return wp.dtype_to_numpy(wp_dtype)(wp_field)
 
   # warp arrays
   if isinstance(wp_field, wp.array):
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
index dba3f715..6f63ba66 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
 
     # Launch the kernel.
     wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
     )
 
 
 def _create_jax_warp_primitive():
     import jax  # noqa: PLC0415
-    from jax._src.interpreters import batching  # noqa: PLC0415
+    from jax.interpreters import batching  # noqa: PLC0415
     from jax.interpreters import mlir  # noqa: PLC0415
     from jax.interpreters.mlir import ir  # noqa: PLC0415
     from jaxlib.hlo_helpers import custom_call  # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
         # Figure out the number of outputs.
         wp_kernel = _registered_kernels[params["kernel"]]
         output_count = len(wp_kernel.adj.args) - len(args)
-        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
         size = shape[dim]
-        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
         # Create the batched primitive.
         return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
 
@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
     def base_type_to_jax_ir(warp_dtype):
         warp_to_jax_dict = {
             wp.float16: ir.F16Type.get(),
+            wp.bfloat16: ir.BF16Type.get(),
             wp.float32: ir.F32Type.get(),
             wp.float64: ir.F64Type.get(),
             wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
     def base_type_is_compatible(warp_type, jax_ir_type):
         jax_ir_to_warp = {
             "f16": wp.float16,
+            "bf16": wp.bfloat16,
             "f32": wp.float32,
             "f64": wp.float64,
             "i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
         # Figure out the types and shapes of the input arrays.
         arg_strings = []
         operand_layouts = []
-        for actual, warg in zip(args, wp_kernel.adj.args):
+        for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
             wtype = warg.type
             rtt = ir.RankedTensorType(actual.type)
 
diff --git a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
index d931dded..ac7b0f26 100644
--- a/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
+++ b/mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
@@ -8,8 +8,8 @@ import ctypes
 import inspect
 import threading
 import traceback
+from collections.abc import Callable
 from enum import IntEnum
-from typing import Callable
 
 import jax
 
@@ -25,6 +25,7 @@ from warp._src.types import (
     type_size_in_bytes,
     type_to_warp,
 )
+from warp._src.utils import warn
 
 from .xla_ffi import *
 
@@ -249,19 +250,18 @@ class FfiKernel:
                 # check dtype
                 if input_value.dtype != input_arg.jax_scalar_type:
                     raise TypeError(
-                            f"Invalid data type for array argument '{input_arg.name}',"
-                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
+                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                     )
                 # check ndim
                 if input_value.ndim != input_arg.jax_ndim:
                     raise TypeError(
-                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                     )
                 # check inner dims
                 for d in range(input_arg.dtype_ndim):
                     if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                         raise TypeError(
-                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                         )
             else:
                 # make sure scalar is not a traced variable, should be static
@@ -455,6 +455,7 @@ class FfiKernel:
                         hooks.forward_smem_bytes,
                         kernel_params,
                         stream,
+                        None,  # apic_info
                     )
                 else:
                     wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
         except Exception as e:
             print(traceback.format_exc())
             return create_ffi_error(
-                call_frame.contents.api,
-                XLA_FFI_Error_Code.UNKNOWN,
-                f"FFI callback error: {type(e).__name__}: {e}",
+                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
             )
 
 
@@ -617,7 +616,9 @@ class FfiCallable:
         self.input_output_aliases = input_output_aliases
 
         # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
+                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
+        )
 
         self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
         ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1071,7 +1072,7 @@ class FfiCallable:
                         call_desc.capture = capture
 
                     else:
-                        # not capturing or on CPU
+                        # not capturing
                         self.func(*arg_list)
 
         except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
                 try:
                     gi.zero_()
                 except Exception as e:
-                    wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
+                    warn(f"Failed to zero gradient array: {e}", stacklevel=2)
                     raise e
 
         # NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
         non_static_inputs, output_vals_tuple = residuals
 
         input_vals = list(non_static_inputs)
-        for i, v in zip(static_args, nondiff_vals):
+        for i, v in zip(static_args, nondiff_vals, strict=True):
             input_vals.insert(i, v)
 
         # Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(
 
         out_dims_map = {}
         param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
-        for name, val in zip(differentiable_input_names, non_static_inputs):
+        for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
             ann = param_ann.get(name)
             if ann is None:
                 continue
diff --git a/mjx/pyproject.toml b/mjx/pyproject.toml
index 43609582..fbf96bcf 100644
--- a/mjx/pyproject.toml
+++ b/mjx/pyproject.toml
@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.12.1",
+    "warp-lang==1.13.0",
 ]
 dev = [
     "isort",

From 50e823e91c6754ddb6dcd1d397b0247b01529c4d Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Sat, 30 May 2026 07:43:56 -0700
Subject: [PATCH 61/61] Import google-deepmind/mujoco_warp from GitHub.

PiperOrigin-RevId: 923922783
Change-Id: I0b407347a70a0e2024aae7ed5bab1592dc40a074
---
 mjx/mujoco/mjx/_src/io.py                     |   13 +-
 mjx/mujoco/mjx/_src/io_test.py                |    2 +-
 .../mjx/third_party/mujoco_warp/__init__.py   |    4 +
 .../mujoco_warp/_src/block_cholesky.py        |  109 +-
 .../mjx/third_party/mujoco_warp/_src/cli.py   |   29 +-
 .../mujoco_warp/_src/collision_convex.py      |   52 +-
 .../mujoco_warp/_src/collision_core.py        |    4 +-
 .../mujoco_warp/_src/collision_driver.py      |   10 +-
 .../mujoco_warp/_src/collision_gjk.py         |   97 +-
 .../mujoco_warp/_src/collision_sdf.py         |   92 +-
 .../mujoco_warp/_src/constraint.py            |  245 +-
 .../mujoco_warp/_src/derivative.py            |  423 +-
 .../third_party/mujoco_warp/_src/forward.py   |  107 +-
 .../third_party/mujoco_warp/_src/history.py   |  925 ++++
 .../third_party/mujoco_warp/_src/inverse.py   |   10 +-
 .../mjx/third_party/mujoco_warp/_src/io.py    |  432 +-
 .../third_party/mujoco_warp/_src/island.py    |  841 ++-
 .../mjx/third_party/mujoco_warp/_src/ray.py   |   14 +-
 .../third_party/mujoco_warp/_src/render.py    |   33 +-
 .../third_party/mujoco_warp/_src/sensor.py    |   28 +-
 .../third_party/mujoco_warp/_src/smooth.py    |  241 +-
 .../third_party/mujoco_warp/_src/solver.py    | 2337 +++++++-
 .../third_party/mujoco_warp/_src/support.py   |  380 +-
 .../mjx/third_party/mujoco_warp/_src/types.py |  301 +-
 .../third_party/mujoco_warp/_src/util_pkg.py  |    6 +-
 .../third_party/mujoco_warp/pyproject.toml    |    4 +-
 mjx/mujoco/mjx/warp/__init__.py               |    6 +
 mjx/mujoco/mjx/warp/forward.py                | 4699 +++++++++--------
 mjx/mujoco/mjx/warp/forward_test.py           |    4 +-
 mjx/mujoco/mjx/warp/smooth_test.py            |    4 +-
 mjx/mujoco/mjx/warp/types.py                  |  224 +-
 31 files changed, 8980 insertions(+), 2696 deletions(-)
 create mode 100644 mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index dc10af8c..70d7f973 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -272,6 +272,10 @@ def _put_option(
 
 
   if impl == types.Impl.WARP:
+    if not mjxw.mjwp_io.ENABLE_ISLANDS:
+      fields['disableflags'] = types.DisableBit(
+          fields['disableflags'] | mjwp_types.DisableBit.ISLAND
+      )
     impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
     return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
 
@@ -1272,8 +1276,11 @@ def _get_data_into_warp(
       else:
         value = getattr(d_i, field.name)
 
-      if field.name in ('ne', 'nl', 'nf'):
-        pass
+      if field.name in ('ne', 'nl', 'nf', 'nisland', 'nidof'):
+        if isinstance(value, np.ndarray) and value.size == 0:
+          value = 0
+        else:
+          value = int(value)
       elif field.name in ('nefc', 'ncon'):
         value = {'nefc': nefc, 'ncon': ncon}[field.name]
       elif field.name.endswith('xmat') or field.name == 'ximat':
@@ -1287,9 +1294,11 @@ def _get_data_into_warp(
           'contact',
           'qM',
           'qLD',
+          'qLU',
           'qLDiagInv',
           'ten_J',
           'flexedge_J',
+          'M',
       ):
         continue
       if field.name.startswith('efc_'):
diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py
index 21e1910c..bd39f5f4 100644
--- a/mjx/mujoco/mjx/_src/io_test.py
+++ b/mjx/mujoco/mjx/_src/io_test.py
@@ -530,7 +530,7 @@ class DataIOTest(parameterized.TestCase):
     elif impl == 'warp':
       qm = np.zeros((m.nv, m.nv), dtype=np.float64)
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      np.testing.assert_allclose(dx._impl.qM, qm)
+      np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
       # TODO(taylorhowell): test efc__J
       np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
index 44224545..ec62bc20 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
@@ -47,6 +47,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
 from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as rungekutta4
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
 from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_ctrl_history as init_ctrl_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import init_sensor_history as init_sensor_history
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_ctrl as read_ctrl
+from mujoco.mjx.third_party.mujoco_warp._src.history import read_sensor as read_sensor
 from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
 from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
 from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
index 855f9686..7faaa578 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
@@ -19,48 +19,77 @@ import warp as wp
 
 
 @lru_cache(maxsize=None)
-def create_blocked_cholesky_func(block_size: int):
+def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
   @wp.func
-  def blocked_cholesky_func(
+  def blocked_cholesky_factorize_solve_func(
     # In:
     A: wp.array2d[float],
+    b: wp.array2d[float],
     matrix_size: int,
     # Out:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
+    x: wp.array2d[float],
   ):
-    """Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
+    """Block Cholesky factorization and solve while keeping the forward RHS live."""
+    rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    It returns a lower-triangular matrix L such that A = L L^T.
-    """
-    # Process the matrix in blocks along its leading dimension.
     for k in range(0, matrix_size, block_size):
       end = k + block_size
+      rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(k, 0))
 
-      # Load current diagonal block A[k:end, k:end]
-      # and update with contributions from previously computed blocks.
-      A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
+      A_kk_tile = wp.tile_load(
+        A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
+      )
 
       for j in range(0, k, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-        wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
+        U_block = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+        )
+        wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
 
-      # Compute the Cholesky factorization for the block
-      wp.tile_cholesky_inplace(A_kk_tile)
-      wp.tile_store(L, A_kk_tile, offset=(k, k))
+        y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
+
+      wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
+      wp.tile_store(U, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
+
+      wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), rhs_view)
 
-      # Process the blocks below the current block
       for i in range(end, matrix_size, block_size):
-        A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
+        A_ki_tile = wp.tile_load(
+          A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
+        )
 
         for j in range(0, k, block_size):
-          L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
-          L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
-          wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
+          U_jk_tile = wp.tile_load(
+            U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
+          )
+          U_ji_tile = wp.tile_load(
+            U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+          )
+          wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
 
-        wp.tile_lower_solve_inplace(A_kk_tile, wp.tile_transpose(A_ik_tile))
-        wp.tile_store(L, A_ik_tile, offset=(i, k))
+        wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
+        wp.tile_store(U, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
 
-  return blocked_cholesky_func
+    for i in range(matrix_size - block_size, -1, -block_size):
+      i_end = i + block_size
+      tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
+      for j in range(i_end, matrix_size, block_size):
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
+        x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
+
+    wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
+
+  return blocked_cholesky_factorize_solve_func
 
 
 @lru_cache(maxsize=None)
@@ -68,41 +97,49 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
   @wp.func
   def blocked_cholesky_solve_func(
     # In:
-    L: wp.array2d[float],
+    U: wp.array2d[float],
     b: wp.array2d[float],
     matrix_size: int,
     # Out:
     x: wp.array2d[float],
   ):
-    """Block Cholesky factorization and solve.
+    """Block Cholesky solve.
 
-    Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
+    Solves A x = b given the Cholesky factor U (A = U^T U) using blocked forward and backward
     substitution.
     """
     rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
 
-    # Forward substitution: solve L y = b
+    # Forward substitution: solve U^T y = b
     for i in range(0, matrix_size, block_size):
       rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(0, i, block_size):
-        L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
+        U_block = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
+        )
         y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
+        wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
 
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
-      wp.tile_lower_solve_inplace(L_tile, rhs_view)
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
+      wp.tile_lower_solve_inplace(wp.tile_transpose(U_tile), rhs_view)
 
-    # Backward substitution: solve L^T x = y
+    # Backward substitution: solve U x = y
     for i in range(matrix_size - block_size, -1, -block_size):
       i_end = i + block_size
       tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
       for j in range(i_end, matrix_size, block_size):
-        L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
+        U_tile = wp.tile_load(
+          U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
+        )
         x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
-        wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
-      L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
+        wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
+      U_tile = wp.tile_load(
+        U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
+      )
 
-      wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
+      wp.tile_upper_solve_inplace(U_tile, tmp_tile)
 
     wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
index 0e427408..425b54bb 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
@@ -25,6 +25,7 @@ from absl import flags
 from etils import epath
 
 import mujoco.mjx.third_party.mujoco_warp as mjw
+from mujoco.mjx.third_party.mujoco_warp._src import io
 from mujoco.mjx.third_party.mujoco_warp._src import warp_util
 from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
 from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -42,6 +43,12 @@ KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulatio
 EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
 NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
 NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
+ENABLE_ISLANDS = flags.DEFINE_bool(
+  "enable_islands",
+  False,
+  "Enable constraint islands solver",
+)
+
 
 DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
 REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -52,6 +59,12 @@ RENDER_RGB = flags.DEFINE_bool("render_rgb", True, "render RGB image")
 RENDER_DEPTH = flags.DEFINE_bool("render_depth", True, "render depth image")
 RENDER_TEXTURES = flags.DEFINE_bool("render_textures", True, "use textures")
 RENDER_SHADOWS = flags.DEFINE_bool("render_shadows", False, "use shadows")
+RENDER_BACKFACE_CULLING = flags.DEFINE_bool(
+  "render_backface_culling",
+  True,
+  "enable renderer backface culling (RenderContext.enable_backface_culling)",
+)
+RENDER_SKYBOX = flags.DEFINE_bool("render_skybox", True, "render skybox texture if available")
 
 
 def load_model(path: epath.Path) -> mujoco.MjModel:
@@ -128,6 +141,8 @@ def init_structs(
   fn: Callable[..., None], mjm: mujoco.MjModel
 ) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
   """Initialize device structs."""
+  io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
+
   mjd = mujoco.MjData(mjm)
   ctrls = None
   if REPLAY.value:
@@ -152,12 +167,14 @@ def init_structs(
 
     rc = mjw.create_render_context(
       mjm,
-      NWORLD.value,
-      (RENDER_WIDTH.value, RENDER_HEIGHT.value),
-      RENDER_RGB.value,
-      RENDER_DEPTH.value,
-      RENDER_TEXTURES.value,
-      RENDER_SHADOWS.value,
+      nworld=NWORLD.value,
+      cam_res=(RENDER_WIDTH.value, RENDER_HEIGHT.value),
+      render_rgb=RENDER_RGB.value,
+      render_depth=RENDER_DEPTH.value,
+      use_textures=RENDER_TEXTURES.value,
+      use_shadows=RENDER_SHADOWS.value,
+      enable_backface_culling=RENDER_BACKFACE_CULLING.value,
+      render_skybox=RENDER_SKYBOX.value,
     )
 
     return m, d, rc, ctrls
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
index a444e399..9bf7ce5e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
@@ -23,6 +23,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_param
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import epa_phase
+from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import gjk_phase
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
 from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
 from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
@@ -35,7 +37,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -715,6 +717,7 @@ def ccd_kernel_builder(
     opt_ccd_tolerance: wp.array[float],
     # Data in:
     naconmax_in: int,
+    naccdmax_in: int,
     # In:
     epa_vert_in: wp.array2d[wp.vec3],
     epa_vert_index_in: wp.array2d[int],
@@ -737,7 +740,7 @@ def ccd_kernel_builder(
     geom2: Geom,
     geoms: wp.vec2i,
     worldid: int,
-    ccdid: int,
+    nccd_in: wp.array[int],
     margin: float,
     gap: float,
     condim: int,
@@ -774,30 +777,47 @@ def ccd_kernel_builder(
     if is_collision_sensor:
       cutoff = 1.0e32
     else:
-      if wp.static(_NEW_GAP_SEMANTICS):
+      if wp.static(NEW_GAP_SEMANTICS):
         cutoff = gap
       else:
         cutoff = 0.0
-    dist, ncollision, w1, w2, multiccd_idx = ccd(
+    needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
       opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
       cutoff,
       gjk_iterations,
-      epa_iterations,
       geom1,
       geom2,
       geomtype1,
       geomtype2,
       x1,
       x2,
-      epa_vert_in[ccdid],
-      epa_vert_index_in[ccdid],
-      epa_face_in[ccdid],
-      epa_pr_in[ccdid],
-      epa_norm2_in[ccdid],
-      epa_horizon_in[ccdid],
     )
 
-    if wp.static(_NEW_GAP_SEMANTICS):
+    ccdid = int(-1)
+    multiccd_idx = int(-1)
+
+    if needs_epa:
+      ccdid = wp.atomic_add(nccd_in, geomgeomid, 1)
+      if ccdid >= naccdmax_in:
+        wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
+        return 0
+      dist, ncollision, w1, w2, multiccd_idx = epa_phase(
+        opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
+        epa_iterations,
+        gjk_result,
+        geom1,
+        geom2,
+        geomtype1,
+        geomtype2,
+        epa_vert_in[ccdid],
+        epa_vert_index_in[ccdid],
+        epa_face_in[ccdid],
+        epa_pr_in[ccdid],
+        epa_norm2_in[ccdid],
+        epa_horizon_in[ccdid],
+      )
+
+    if wp.static(NEW_GAP_SEMANTICS):
       if dist >= gap and pairid[1] == -1:
         return 0
     else:
@@ -990,11 +1010,6 @@ def ccd_kernel_builder(
     if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
       return
 
-    ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
-    if ccdid >= naccdmax_in:
-      wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
-      return
-
     worldid = collision_worldid_in[collisionid]
 
     _, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -1046,6 +1061,7 @@ def ccd_kernel_builder(
     eval_ccd_write_contact(
       opt_ccd_tolerance,
       naconmax_in,
+      naccdmax_in,
       epa_vert_in,
       epa_vert_index_in,
       epa_face_in,
@@ -1067,7 +1083,7 @@ def ccd_kernel_builder(
       geom2,
       geoms,
       worldid,
-      ccdid,
+      nccd_in,
       margin,
       gap,
       condim,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
index 41feb836..35bb0f27 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
@@ -22,8 +22,8 @@ import warp as wp
 
 from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
-from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
@@ -219,7 +219,7 @@ def write_contact(
     contact_frame_out[cid] = frame_in
     contact_geom_out[cid] = geoms_in
     contact_worldid_out[cid] = worldid_in
-    if wp.static(_NEW_GAP_SEMANTICS):
+    if wp.static(NEW_GAP_SEMANTICS):
       includemargin = margin_in
     else:
       includemargin = margin_in - gap_in
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
index df75adc7..7c2fe28e 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
@@ -375,6 +375,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
   return upper
 
 
+@cache_kernel
 def _sap_project(opt_broadphase: int):
   @wp.kernel(module="unique", enable_backward=False)
   def sap_project(
@@ -531,6 +532,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
   return kernel
 
 
+@cache_kernel
 def _segmented_sort(tile_size: int):
   @wp.kernel(module="unique")
   def segmented_sort(
@@ -629,7 +631,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
   # assumes each geom has 5 other geoms (batched over all worlds)
   nsweep = 5 * nworldgeom
   wp.launch(
-    kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    kernel=_sap_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=nsweep,
     inputs=[
       m.ngeom,
@@ -718,7 +722,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
   `contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
   """
   wp.launch(
-    _nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
+    _nxn_broadphase(
+      m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
+    ),
     dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
     inputs=[
       m.geom_type,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
index 21e61b7c..f299cd78 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
@@ -2135,8 +2135,12 @@ def multicontact(
 
   # face1 is an edge; clip face1 against face2
   if is_edge_contact_geom1:
-    approx_dir = wp.norm_l2(dir) * n2[j]
-    return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
+    approx_dir = -wp.norm_l2(dir) * n2[j]
+    nclipped, clipped1, clipped2 = _polygon_clip(
+      plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped
+    )
+    # the faces were flipped in calling _polygon_clip so we need to flip them back
+    return nclipped, clipped2, clipped1
 
   # face2 is an edge; clip face2 against face1
   if is_edge_contact_geom2:
@@ -2197,34 +2201,29 @@ def _inflate(
 
 
 @wp.func
-def ccd(
+def gjk_phase(
   # In:
   tolerance: float,
   cutoff: float,
   gjk_iterations: int,
-  epa_iterations: int,
   geom1: Geom,
   geom2: Geom,
   geomtype1: int,
   geomtype2: int,
   x_1: wp.vec3,
   x_2: wp.vec3,
-  vert: wp.array[wp.vec3],
-  vert_index: wp.array[int],
-  face: wp.array[int],
-  face_pr: wp.array[wp.vec3],
-  face_norm2: wp.array[float],
-  horizon: wp.array[int],
-) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
-  """General convex collision detection via GJK/EPA."""
+) -> Tuple[bool, float, int, wp.vec3, wp.vec3, GJKResult, Geom, Geom]:
+  """Run GJK phase of CCD."""
   full_margin1 = 0.0
   full_margin2 = 0.0
   size1 = 0.0
   size2 = 0.0
+  empty = GJKResult()
 
   # determine if the geoms being tested are discrete
   is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
 
+  # special handling for sphere and capsule (shrink to point and line respectively)
   if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
     size1 = geom1.size[0]
     full_margin1 = size1 + 0.5 * geom1.margin
@@ -2237,7 +2236,6 @@ def ccd(
     geom2.margin = 0.0
     geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
 
-  # special handling for sphere and capsule (shrink to point and line respectively)
   if size1 + size2 > 0.0:
     cutoff += full_margin1 + full_margin2
     result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
@@ -2245,11 +2243,11 @@ def ccd(
     # shallow penetration, inflate contact
     if result.dist > tolerance:
       if result.dist == FLOAT_MAX:
-        return result.dist, 1, result.x1, result.x2, -1
+        return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
       dist, x1, x2 = _inflate(result, geom1, geom2, geomtype1, geomtype2, full_margin1, full_margin2)
-      return dist, 1, x1, x2, -1
+      return False, dist, 1, x1, x2, empty, geom1, geom2
 
-    # deep penetration, reset initial conditions and rerun GJK + EPA
+    # deep penetration: reset initial conditions and rerun GJK + EPA
     geom1.margin = full_margin1 - size1
     geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
     geom2.margin = full_margin2 - size2
@@ -2260,8 +2258,29 @@ def ccd(
 
   # no penetration depth to recover
   if result.dist > tolerance or result.dim < 2:
-    return result.dist, 1, result.x1, result.x2, -1
+    return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
 
+  return True, result.dist, 1, result.x1, result.x2, result, geom1, geom2
+
+
+@wp.func
+def epa_phase(
+  # In:
+  tolerance: float,
+  epa_iterations: int,
+  result: GJKResult,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
   pt = Polytope()
   pt.nface = 0
   pt.nvert = 0
@@ -2330,6 +2349,7 @@ def ccd(
   if pt.status:
     return result.dist, 1, result.x1, result.x2, -1
 
+  is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
   dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
   if idx == -1:
     return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2343,3 +2363,46 @@ def ccd(
     idx = -1
 
   return dist, 1, x1, x2, idx
+
+
+@wp.func
+def ccd(
+  # In:
+  tolerance: float,
+  cutoff: float,
+  gjk_iterations: int,
+  epa_iterations: int,
+  geom1: Geom,
+  geom2: Geom,
+  geomtype1: int,
+  geomtype2: int,
+  x_1: wp.vec3,
+  x_2: wp.vec3,
+  vert: wp.array[wp.vec3],
+  vert_index: wp.array[int],
+  face: wp.array[int],
+  face_pr: wp.array[wp.vec3],
+  face_norm2: wp.array[float],
+  horizon: wp.array[int],
+) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
+  """General convex collision detection via GJK/EPA."""
+  needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
+    tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, x_1, x_2
+  )
+  if not needs_epa:
+    return dist, ncontact, x1, x2, -1
+  return epa_phase(
+    tolerance,
+    epa_iterations,
+    result,
+    geom1,
+    geom2,
+    geomtype1,
+    geomtype2,
+    vert,
+    vert_index,
+    face,
+    face_pr,
+    face_norm2,
+    horizon,
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
index 5937cf6e..a96c8985 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collisio
 from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
 from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
 from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import Data
 from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -183,6 +184,24 @@ def ellipsoid(p: wp.vec3, size: wp.vec3) -> float:
   return k0 * (k0 - 1.0) / denom
 
 
+@wp.func
+def capsule(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  return wp.length(diff) - r
+
+
+@wp.func
+def cylinder(p: wp.vec3, size: wp.vec3) -> float:
+  r = size[0]
+  h = size[1]
+  dx = wp.length(wp.vec2(p[0], p[1])) - r
+  dy = wp.abs(p[2]) - h
+  return wp.min(wp.max(dx, dy), 0.0) + wp.length(wp.vec2(wp.max(dx, 0.0), wp.max(dy, 0.0)))
+
+
 @wp.func
 def grad_sphere(p: wp.vec3) -> wp.vec3:
   c = wp.length(p)
@@ -230,6 +249,52 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
   return raw_grad / wp.length(raw_grad)
 
 
+@wp.func
+def grad_capsule(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  h = size[1]
+  pz_clamped = wp.clamp(p[2], -h, h)
+  diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
+  c = wp.length(diff)
+  if c > MJ_MINVAL:
+    return diff / c
+  else:
+    return wp.vec3(0.0)
+
+
+@wp.func
+def grad_cylinder(p: wp.vec3, size: wp.vec3) -> wp.vec3:
+  r = size[0]
+  h = size[1]
+
+  radial_dist = wp.length(wp.vec2(p[0], p[1]))
+  if radial_dist > MJ_MINVAL:
+    u = wp.vec3(p[0] / radial_dist, p[1] / radial_dist, 0.0)
+  else:
+    u = wp.vec3(0.0)
+
+  w = wp.vec3(0.0, 0.0, wp.sign(p[2]))
+
+  dx = radial_dist - r
+  dy = wp.abs(p[2]) - h
+
+  if dx > 0.0 and dy > 0.0:
+    v = wp.vec2(dx, dy)
+    len_v = wp.length(v)
+    if len_v > MJ_MINVAL:
+      return u * (dx / len_v) + w * (dy / len_v)
+    else:
+      return wp.vec3(0.0)
+  elif dx > 0.0:
+    return u
+  elif dy > 0.0:
+    return w
+  else:
+    if dx > dy:
+      return u
+    else:
+      return w
+
+
 @wp.func
 def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
   """User-defined SDF function.
@@ -287,16 +352,17 @@ def find_oct(
     # check if the node is a leaf
     # child indices are relative to root (mesh_octadr offset)
     child0 = oct_child[node][0]
+    # Evaluate this hot leaf predicate eagerly to avoid branch-heavy codegen.
     if (
-      child0 == -1
-      and oct_child[node][1] == -1
-      and oct_child[node][2] == -1
-      and oct_child[node][3] == -1
-      and oct_child[node][4] == -1
-      and oct_child[node][5] == -1
-      and oct_child[node][6] == -1
-      and oct_child[node][7] == -1
-    ):
+      int(child0 == -1)
+      & int(oct_child[node][1] == -1)
+      & int(oct_child[node][2] == -1)
+      & int(oct_child[node][3] == -1)
+      & int(oct_child[node][4] == -1)
+      & int(oct_child[node][5] == -1)
+      & int(oct_child[node][6] == -1)
+      & int(oct_child[node][7] == -1)
+    ) != 0:
       for j in range(8):
         if not grad:
           rx[j] = (
@@ -394,6 +460,10 @@ def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data:
     return p[2]
   elif type == GeomType.SPHERE:
     return sphere(p, attr_vec3)
+  elif type == GeomType.CAPSULE:
+    return capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
@@ -452,6 +522,10 @@ def sdf_grad(
     return grad
   elif type == GeomType.SPHERE:
     return grad_sphere(p)
+  elif type == GeomType.CAPSULE:
+    return grad_capsule(p, attr_vec3)
+  elif type == GeomType.CYLINDER:
+    return grad_cylinder(p, attr_vec3)
   elif type == GeomType.BOX:
     return grad_box(p, attr_vec3)
   elif type == GeomType.ELLIPSOID:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
index 59a7ad72..bcb1acf5 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
@@ -134,7 +134,10 @@ def _equality_connect(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   eq_obj1id: wp.array[int],
@@ -154,6 +157,9 @@ def _equality_connect(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -214,6 +220,7 @@ def _equality_connect(
 
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvel = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -282,6 +289,42 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -294,7 +337,9 @@ def _equality_connect(
       efc_J_out[worldid, 0, sparseid1] = j1mj2[1]
       efc_J_out[worldid, 0, sparseid2] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, da]
+      qvel = qvel_in[worldid, da]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
       nnz += 1
   else:
@@ -326,11 +371,49 @@ def _equality_connect(
       )
       j1mj2 = jacp1 - jacp2
 
+      jacp1_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, _ = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+      j1mj2_dot = jacp1_dot - jacp2_dot
+
       efc_J_out[worldid, efcid0, dofid] = j1mj2[0]
       efc_J_out[worldid, efcid1, dofid] = j1mj2[1]
       efc_J_out[worldid, efcid2, dofid] = j1mj2[2]
 
-      Jqvel += j1mj2 * qvel_in[worldid, dofid]
+      qvel = qvel_in[worldid, dofid]
+      Jqvel += j1mj2 * qvel
+      Jdotv += j1mj2_dot * qvel
 
   body_invweight0_id = worldid % body_invweight0.shape[0]
   invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -368,6 +451,8 @@ def _equality_connect(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcidi] -= Jdotv[i]
+
 
 @wp.kernel
 def _equality_joint(
@@ -812,7 +897,10 @@ def _equality_weld(
   body_dofnum: wp.array[int],
   body_dofadr: wp.array[int],
   body_invweight0: wp.array2d[wp.vec2],
+  jnt_type: wp.array[int],
+  jnt_dofadr: wp.array[int],
   dof_bodyid: wp.array[int],
+  dof_jntid: wp.array[int],
   dof_parentid: wp.array[int],
   site_bodyid: wp.array[int],
   site_quat: wp.array2d[wp.quat],
@@ -834,6 +922,9 @@ def _equality_weld(
   site_xpos_in: wp.array2d[wp.vec3],
   subtree_com_in: wp.array2d[wp.vec3],
   cdof_in: wp.array2d[wp.spatial_vector],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  subtree_linvel_in: wp.array2d[wp.vec3],
   njmax_in: int,
   njmax_nnz_in: int,
   # Data out:
@@ -903,9 +994,43 @@ def _equality_weld(
     quat = math.mul_quat(xquat_in[worldid, body1], relpose)
     quat1 = math.quat_inv(xquat_in[worldid, body2])
 
+    # quat1 = quat_inv(xquat_in[worldid, body2])
+    q2 = xquat_in[worldid, body2]
+    quat1 = wp.quat(q2[0], -q2[1], -q2[2], -q2[3])
+
+  # compute rotational Jdotv helper quaternions
+  omega1 = wp.spatial_top(cvel_in[worldid, body1])
+  omega2 = wp.spatial_top(cvel_in[worldid, body2])
+  domega = omega1 - omega2
+
+  omega1_q = wp.quat(0.0, omega1[0], omega1[1], omega1[2])
+  omega2_q = wp.quat(0.0, omega2[0], omega2[1], omega2[2])
+  domega_q = wp.quat(0.0, domega[0], domega[1], domega[2])
+
+  if is_site:
+    qdot0r = math.mul_quat(omega1_q, quat) * 0.5
+    qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
+    qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
+
+  else:
+    # qdot0 = mul_quat(xquat_in[worldid, body1], omega1_q) * 0.5
+    u7 = xquat_in[worldid, body1]
+    qdot0 = math.mul_quat(omega1_q, xquat_in[worldid, body1]) * 0.5
+    qdot0r = math.mul_quat(qdot0, relpose)
+    q1_non_site = xquat_in[worldid, body2]
+    qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
+
+    negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
+    negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
+
   # compute Jacobian difference (opposite of contact: 0 - 1)
   Jqvelp = wp.vec3f(0.0, 0.0, 0.0)
   Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_p = wp.vec3f(0.0, 0.0, 0.0)
+  Jdotv_r0 = wp.vec3f(0.0, 0.0, 0.0)
 
   if is_sparse:
     # TODO(team): pre-compute number of non-zeros
@@ -985,6 +1110,44 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        da,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        da,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       sparseid0 = rowadr + nnz
       sparseid1 = rowadr + rownnz + nnz
       sparseid2 = rowadr + 2 * rownnz + nnz
@@ -1008,6 +1171,8 @@ def _equality_weld(
 
       Jqvelp += jacdifp * qvel_in[worldid, da]
       Jqvelr += jacdifr * qvel_in[worldid, da]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, da]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, da]
 
       nnz += 1
   else:
@@ -1047,12 +1212,52 @@ def _equality_weld(
       jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
       jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
 
+      jacp1_dot, jacr1_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos1,
+        body1,
+        dofid,
+        worldid,
+      )
+      jacp2_dot, jacr2_dot = support.jac_dot_dof(
+        body_parentid,
+        body_rootid,
+        jnt_type,
+        jnt_dofadr,
+        dof_bodyid,
+        dof_jntid,
+        body_isdofancestor,
+        subtree_com_in,
+        cdof_in,
+        cvel_in,
+        cdof_dot_in,
+        pos2,
+        body2,
+        dofid,
+        worldid,
+      )
+
+      jacdifp_dot = jacp1_dot - jacp2_dot
+      jacdifr_dot = jacr1_dot - jacr2_dot
+
       efc_J_out[worldid, efcid3, dofid] = jacdifr[0]
       efc_J_out[worldid, efcid4, dofid] = jacdifr[1]
       efc_J_out[worldid, efcid5, dofid] = jacdifr[2]
 
       Jqvelp += jacdifp * qvel_in[worldid, dofid]
       Jqvelr += jacdifr * qvel_in[worldid, dofid]
+      Jdotv_p += jacdifp_dot * qvel_in[worldid, dofid]
+      Jdotv_r0 += jacdifr_dot * qvel_in[worldid, dofid]
 
   # error is difference in global position and orientation
   cpos = pos1 - pos2
@@ -1070,6 +1275,22 @@ def _equality_weld(
 
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
+  djrdv_q = wp.quat(0.0, Jdotv_r0[0], Jdotv_r0[1], Jdotv_r0[2])
+
+  # Term 1: negqdot1 * domega * q0r
+  t1a = math.mul_quat(negqdot1, domega_q)
+  t1 = math.mul_quat(t1a, quat)
+
+  # Term 2: negq1 * djrdv * q0r
+  t2a = math.mul_quat(negq1, djrdv_q)
+  t2 = math.mul_quat(t2a, quat)
+
+  # Term 3: negq1 * domega * qdot0r
+  t3a = math.mul_quat(negq1, domega_q)
+  t3 = math.mul_quat(t3a, qdot0r)
+
+  Jdotv_r = wp.vec3(t1[1] + t2[1] + t3[1], t1[2] + t2[2] + t3[2], t1[3] + t2[3] + t3[3]) * 0.5 * torquescale
+
   for i in range(3):
     _efc_row(
       opt_disableflags,
@@ -1096,6 +1317,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + i] -= Jdotv_p[i]
+
   invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1]
 
   for i in range(3):
@@ -1124,6 +1347,8 @@ def _equality_weld(
       efc_frictionloss_out,
     )
 
+    efc_aref_out[worldid, efcid + 3 + i] -= Jdotv_r[i]
+
 
 @wp.kernel
 def _friction_dof(
@@ -1680,6 +1905,7 @@ def _limit_tendon(
     )
 
 
+@cache_kernel
 def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
   IS_SPARSE = is_sparse
@@ -1803,6 +2029,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_sparse(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -1973,6 +2200,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   TILE_SIZE = tile_size
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2117,6 +2345,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
   return kernel
 
 
+@cache_kernel
 def _efc_contact_update(cone_type: types.ConeType):
   IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
 
@@ -2297,7 +2526,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.eq_obj1id,
@@ -2316,6 +2548,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],
@@ -2351,7 +2586,10 @@ def make_constraint(m: types.Model, d: types.Data):
           m.body_dofnum,
           m.body_dofadr,
           m.body_invweight0,
+          m.jnt_type,
+          m.jnt_dofadr,
           m.dof_bodyid,
+          m.dof_jntid,
           m.dof_parentid,
           m.site_bodyid,
           m.site_quat,
@@ -2372,6 +2610,9 @@ def make_constraint(m: types.Model, d: types.Data):
           d.site_xpos,
           d.subtree_com,
           d.cdof,
+          d.cvel,
+          d.cdof_dot,
+          d.subtree_linvel,
           d.njmax,
           d.njmax_nnz,
         ],
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
index 4a241f3b..28c8193f 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
@@ -15,6 +15,7 @@
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -24,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
 from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
 from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
 from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -60,13 +62,47 @@ def _qderiv_actuator_passive_vel(
   actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
   actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
 
+  bias = float(0.0)
+
   if actuator_gaintype[actid] == GainType.AFFINE:
     gain = actuator_gainprm[actuator_gainprm_id, actid][2]
+  elif actuator_gaintype[actid] == GainType.DCMOTOR:
+    gain = 0.0
+    dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
+    gainprm = actuator_gainprm[actuator_gainprm_id, actid]
+    te = dynprm[0]
+
+    # controller velocity derivative: dV/dω
+    input_mode = int(gainprm[8])
+    dVdw = 0.0
+    if input_mode == 1:
+      dVdw = -gainprm[6]  # position: -kd
+    elif input_mode == 2:
+      dVdw = -gainprm[4]  # velocity: -kp
+
+    if te > 0.0:
+      # stateful current with actearly: d(K*next_act)/dω
+      # includes both back-EMF (-K) and controller (dVdw) through act_dot
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      s = 1.0 - wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / te)
+      bias += K * (dVdw - K) * s / R
+    elif dVdw != 0.0:
+      # stateless: controller terms only (back-EMF handled in bias block)
+      R = wp.max(MJ_MINVAL, gainprm[0])
+      K = gainprm[1]
+      bias += K * dVdw / R
+
+    # LuGre: force includes -sigma1*z_dot, z_dot = a*z + v
+    # d(sigma1*z_dot)/dv = sigma1*(da/dv*z + 1), ignoring higher-order da/dv*z
+    sigma1 = dynprm[6]
+    if sigma1 > 0.0:
+      bias -= sigma1
   else:
     gain = 0.0
 
   if actuator_biastype[actid] == BiasType.AFFINE:
-    bias = actuator_biasprm[actuator_biasprm_id, actid][2]
+    bias += actuator_biasprm[actuator_biasprm_id, actid][2]
   elif actuator_biastype[actid] == BiasType.DCMOTOR:
     dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
     te = dynprm[0]
@@ -86,11 +122,7 @@ def _qderiv_actuator_passive_vel(
         Ta = dynprm[4]
         R *= 1.0 + alpha * (T + Ta - T0)
 
-      bias = -K * K / wp.max(MJ_MINVAL, R)
-    else:
-      bias = 0.0
-  else:
-    bias = 0.0
+      bias += -K * K / wp.max(MJ_MINVAL, R)
 
   if bias == 0.0 and gain == 0.0:
     vel_out[worldid, actid] = 0.0
@@ -151,15 +183,15 @@ def _qderiv_actuator_passive_actuation_dense(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
   qderiv_contrib = float(0.0)
   for actid in range(nu):
     vel = vel_in[worldid, actid]
@@ -195,8 +227,7 @@ def _qderiv_actuator_passive_actuation_dense(
 @wp.kernel
 def _qderiv_actuator_passive_actuation_sparse(
   # Model:
-  M_rownnz: wp.array[int],
-  M_rowadr: wp.array[int],
+  M_elemid: wp.array2d[int],
   # Data in:
   moment_rownnz_in: wp.array2d[int],
   moment_rowadr_in: wp.array2d[int],
@@ -204,7 +235,6 @@ def _qderiv_actuator_passive_actuation_sparse(
   actuator_moment_in: wp.array2d[float],
   # In:
   vel_in: wp.array2d[float],
-  qMj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
@@ -231,19 +261,10 @@ def _qderiv_actuator_passive_actuation_sparse(
         continue
       dofj = moment_colind_in[worldid, rowadrj]
 
-      contrib = moment_i * moment_j * vel
-
-      # Search the corresponding elemid
-      # TODO: This could be precalculated for improved performance
-      row = dofi
-      col = dofj
-      row_startk = M_rowadr[row] - 1
-      row_nnz = M_rownnz[row]
-      for k in range(row_nnz):
-        row_startk += 1
-        if qMj[row_startk] == col:
-          wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
-          break
+      elemid = M_elemid[dofi, dofj]
+      if elemid >= 0:
+        contrib = moment_i * moment_j * vel
+        wp.atomic_add(qDeriv_out[worldid, 0], elemid, contrib)
 
 
 @wp.kernel
@@ -254,23 +275,29 @@ def _qderiv_actuator_passive(
   dof_damping: wp.array2d[float],
   dof_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   qvel_in: wp.array2d[float],
-  qM_in: wp.array3d[float],
+  M_in: wp.array3d[float],
   # In:
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   qDeriv_in: wp.array3d[float],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
 
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
+
+  madr = M_elemid[dofiid, dofjid]
 
   if is_sparse:
-    qderiv = qDeriv_in[worldid, 0, elemid]
+    if madr >= 0:
+      qderiv = qDeriv_in[worldid, 0, madr]
+    else:
+      qderiv = 0.0
   else:
     qderiv = qDeriv_in[worldid, dofiid, dofjid]
 
@@ -283,12 +310,13 @@ def _qderiv_actuator_passive(
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] = M_in[worldid, 0, madr] - qderiv
   else:
-    qM = qM_in[worldid, dofiid, dofjid] - qderiv
-    qDeriv_out[worldid, dofiid, dofjid] = qM
+    M = M_in[worldid, dofiid, dofjid] - qderiv
+    qDeriv_out[worldid, dofiid, dofjid] = M
     if dofiid != dofjid:
-      qDeriv_out[worldid, dofjid, dofiid] = qM
+      qDeriv_out[worldid, dofjid, dofiid] = M
 
 
 # TODO(team): improve performance with tile operations?
@@ -303,18 +331,19 @@ def _qderiv_tendon_damping(
   tendon_damping: wp.array2d[float],
   tendon_dampingpoly: wp.array2d[wp.vec2],
   is_sparse: bool,
+  M_elemid: wp.array2d[int],
   # Data in:
   ten_J_in: wp.array2d[float],
   ten_velocity_in: wp.array2d[float],
   # In:
-  qMi: wp.array[int],
-  qMj: wp.array[int],
+  Mi: wp.array[int],
+  Mj: wp.array[int],
   # Out:
   qDeriv_out: wp.array3d[float],
 ):
   worldid, elemid = wp.tid()
-  dofiid = qMi[elemid]
-  dofjid = qMj[elemid]
+  dofiid = Mi[elemid]
+  dofjid = Mj[elemid]
 
   qderiv = float(0.0)
   tendon_damping_id = worldid % tendon_damping.shape[0]
@@ -343,14 +372,283 @@ def _qderiv_tendon_damping(
 
   qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
 
+  madr = M_elemid[dofiid, dofjid]
+
   if is_sparse:
-    qDeriv_out[worldid, 0, elemid] -= qderiv
+    if madr >= 0:
+      qDeriv_out[worldid, 0, madr] -= qderiv
   else:
     qDeriv_out[worldid, dofiid, dofjid] -= qderiv
     if dofiid != dofjid:
       qDeriv_out[worldid, dofjid, dofiid] -= qderiv
 
 
+@wp.kernel
+def deriv_rne_cvel_cdof_dot(
+  # Model:
+  body_parentid: wp.array[int],
+  body_jntnum: wp.array[int],
+  body_jntadr: wp.array[int],
+  body_dofadr: wp.array[int],
+  jnt_type: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcvel_out: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cvel)/d(qvel_k) and d(cdof_dot)/d(qvel_k).
+
+  Mirrors the accumulation order of comvel for each joint type.
+
+  Dcdof_dot for rotation DOFs of free joints (dofid+0..2) is zero because the
+  forward pass sets cdof_dot[dofid+0..2] = 0.  The Dcdof_dot array is
+  zero-initialized so no explicit write is needed.
+  """
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  jntid = body_jntadr[bodyid]
+  jntnum = body_jntnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  cdof = cdof_in[worldid]
+
+  # Initialize from parent
+  cvel_k = Dcvel_out[worldid, pid, dofid]
+
+  if jntnum == 0:
+    Dcvel_out[worldid, bodyid, dofid] = cvel_k
+    return
+
+  dof_i = dofadr
+
+  for j in range(jntid, jntid + jntnum):
+    jnttype = jnt_type[j]
+
+    if jnttype == 0:  # FREE
+      # rotation DOFs (dof_i+0..2) contribute to cvel
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      # cdof_dot for rotation DOFs is zero (set in forward kinematics),
+      # so Dcdof_dot for rotation DOFs is zero (from wp.zeros init)
+
+      # derivative of cdof_dot for translation DOFs 3,4,5
+      Dcdof_dot_out[worldid, dof_i + 3, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 3])
+      Dcdof_dot_out[worldid, dof_i + 4, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 4])
+      Dcdof_dot_out[worldid, dof_i + 5, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 5])
+
+      # translation DOFs (dof_i+3..5) contribute to cvel
+      if dofid >= dof_i + 3 and dofid < dof_i + 6:
+        cvel_k += cdof[dofid]
+
+      dof_i += 6
+
+    elif jnttype == 1:  # BALL
+      Dcdof_dot_out[worldid, dof_i + 0, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 0])
+      Dcdof_dot_out[worldid, dof_i + 1, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 1])
+      Dcdof_dot_out[worldid, dof_i + 2, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 2])
+
+      if dofid >= dof_i and dofid < dof_i + 3:
+        cvel_k += cdof[dofid]
+
+      dof_i += 3
+    else:  # HINGE or SLIDE
+      Dcdof_dot_out[worldid, dof_i, dofid] = math.motion_cross(cvel_k, cdof[dof_i])
+
+      if dofid == dof_i:
+        cvel_k += cdof[dof_i]
+
+      dof_i += 1
+
+  Dcvel_out[worldid, bodyid, dofid] = cvel_k
+
+
+@wp.kernel
+def deriv_rne_cacc_cfrcbody_forward(
+  # Model:
+  body_parentid: wp.array[int],
+  body_dofnum: wp.array[int],
+  body_dofadr: wp.array[int],
+  # Data in:
+  qvel_in: wp.array2d[float],
+  cinert_in: wp.array2d[vec10],
+  cvel_in: wp.array2d[wp.spatial_vector],
+  cdof_dot_in: wp.array2d[wp.spatial_vector],
+  # In:
+  body_tree_: wp.array[int],
+  Dcvel_in: wp.array3d[wp.spatial_vector],
+  Dcdof_dot_in: wp.array3d[wp.spatial_vector],
+  # Out:
+  Dcacc_out: wp.array3d[wp.spatial_vector],
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Forward pass: compute d(cacc)/d(qvel_k) and d(cfrc_body)/d(qvel_k)."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  dofadr = body_dofadr[bodyid]
+  dofnum = body_dofnum[bodyid]
+  pid = body_parentid[bodyid]
+
+  qvel = qvel_in[worldid]
+
+  dcacc = Dcacc_out[worldid, pid, dofid]
+
+  for j in range(dofadr, dofadr + dofnum):
+    # Term 1: d(cdof_dot * qvel)/d(qvel_k) when j == dofid
+    if j == dofid:
+      dcacc += cdof_dot_in[worldid, j]
+
+    # Term 2: cdof_dot depends on cvel which depends on qvel_k
+    dcdofdot = Dcdof_dot_in[worldid, j, dofid]
+    dcacc += dcdofdot * qvel[j]
+
+  Dcacc_out[worldid, bodyid, dofid] = dcacc
+
+  # d(cfrc_body)/d(qvel_k)
+  cinert = cinert_in[worldid, bodyid]
+  cvel = cvel_in[worldid, bodyid]
+  dcvel = Dcvel_in[worldid, bodyid, dofid]
+
+  # term1 = cinert * d(cacc)/d(qvel_k)
+  term1 = math.inert_vec(cinert, dcacc)
+
+  # term2 = d(cvel x* (cinert * cvel))/d(qvel_k)
+  cinert_cvel = math.inert_vec(cinert, cvel)
+  cinert_dcvel = math.inert_vec(cinert, dcvel)
+  term2 = math.motion_cross_force(dcvel, cinert_cvel) + math.motion_cross_force(cvel, cinert_dcvel)
+
+  Dcfrcbody_out[worldid, bodyid, dofid] = term1 + term2
+
+
+@wp.kernel
+def deriv_rne_cfrcbody_backward(
+  # Model:
+  body_parentid: wp.array[int],
+  # In:
+  body_tree_: wp.array[int],
+  # Out:
+  Dcfrcbody_out: wp.array3d[wp.spatial_vector],
+):
+  """Backward pass: accumulate d(cfrc_body) from children to parents."""
+  worldid, nodeid, dofid = wp.tid()
+  bodyid = body_tree_[nodeid]
+  pid = body_parentid[bodyid]
+
+  # body_tree never contains bodyid=0 (worldbody), so pid >= 0 is always valid.
+  # Siblings at the same level may share a parent; atomic_add handles this.
+  val = Dcfrcbody_out[worldid, bodyid, dofid]
+  wp.atomic_add(Dcfrcbody_out[worldid, pid], dofid, val)
+
+
+@wp.kernel
+def deriv_rne_body2jnt_sparse(
+  # Model:
+  dof_bodyid: wp.array[int],
+  # Data in:
+  cdof_in: wp.array2d[wp.spatial_vector],
+  # In:
+  timestep: wp.array[float],
+  Di: wp.array[int],
+  Dj: wp.array[int],
+  Dcfrcbody_in: wp.array3d[wp.spatial_vector],
+  flg_subtract: bool,
+  # Out:
+  qDeriv_out: wp.array3d[float],
+):
+  """Project body-space RNE derivatives into joint-space qDeriv (sparse)."""
+  worldid, elemid = wp.tid()
+  dt = timestep[worldid % timestep.shape[0]]
+
+  i = Di[elemid]
+  j = Dj[elemid]
+
+  body_i = dof_bodyid[i]
+  dcfrc = Dcfrcbody_in[worldid, body_i, j]
+  term = wp.dot(cdof_in[worldid, i], dcfrc)
+
+  if flg_subtract:
+    wp.atomic_sub(qDeriv_out[worldid, 0], elemid, dt * term)
+  else:
+    wp.atomic_add(qDeriv_out[worldid, 0], elemid, dt * term)
+
+
+def deriv_rne_vel(m: Model, d: Data, out: wp.array3d[float], flg_subtract: bool = False):
+  """Compute RNE velocity derivatives and add/subtract from the output.
+
+  Implements the analytical derivative of inverse-dynamics Coriolis/centrifugal
+  forces with respect to joint velocities.
+
+  Args:
+    m: The model (device).
+    d: The data (device).
+    out: D-structure output array (nworld, 1, nD) to accumulate RNE terms into.
+    flg_subtract: If True, subtract the RNE derivatives from output instead of adding them.
+  """
+  # TODO(team): consider caching these allocations
+  Dcvel = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcdof_dot = wp.zeros((d.nworld, m.nv, m.nv), dtype=wp.spatial_vector)
+  Dcacc = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+  Dcfrcbody = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
+
+  # Forward pass 1: compute Dcvel and Dcdof_dot
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cvel_cdof_dot,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_jntnum,
+        m.body_jntadr,
+        m.body_dofadr,
+        m.jnt_type,
+        d.cdof,
+        body_tree,
+      ],
+      outputs=[Dcvel, Dcdof_dot],
+    )
+
+  # Forward pass 2: compute Dcacc and Dcfrcbody
+  for body_tree in m.body_tree:
+    wp.launch(
+      deriv_rne_cacc_cfrcbody_forward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[
+        m.body_parentid,
+        m.body_dofnum,
+        m.body_dofadr,
+        d.qvel,
+        d.cinert,
+        d.cvel,
+        d.cdof_dot,
+        body_tree,
+        Dcvel,
+        Dcdof_dot,
+      ],
+      outputs=[Dcacc, Dcfrcbody],
+    )
+
+  # Backward pass: accumulate Dcfrcbody from children to parents
+  for body_tree in reversed(m.body_tree):
+    wp.launch(
+      deriv_rne_cfrcbody_backward,
+      dim=(d.nworld, body_tree.size, m.nv),
+      inputs=[m.body_parentid, body_tree],
+      outputs=[Dcfrcbody],
+    )
+
+  # Project body-space derivatives into joint-space qDeriv (always sparse D-structure)
+  wp.launch(
+    deriv_rne_body2jnt_sparse,
+    dim=(d.nworld, m.qD_fullm_i.size),
+    inputs=[m.dof_bodyid, d.cdof, m.opt.timestep, m.qD_fullm_i, m.qD_fullm_j, Dcfrcbody, flg_subtract],
+    outputs=[out],
+  )
+
+
 @event_scope
 def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   """Analytical derivative of smooth forces w.r.t. velocities.
@@ -358,12 +656,10 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
+    out: M - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
   """
-  qMi = m.qM_fullm_i
-  qMj = m.qM_fullm_j
-
-  # TODO(team): implicit requires different sparsity structure
+  Mi = m.M_fullm_i
+  Mj = m.M_fullm_j
 
   if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
     # TODO(team): only clear elements not set by _qderiv_actuator_passive
@@ -399,41 +695,49 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         wp.launch(
           _qderiv_actuator_passive_actuation_sparse,
           dim=(d.nworld, m.nu),
-          inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
+          inputs=[
+            m.M_elemid,
+            d.moment_rownnz,
+            d.moment_rowadr,
+            d.moment_colind,
+            d.actuator_moment,
+            vel,
+          ],
           outputs=[out],
         )
       else:
         wp.launch(
           _qderiv_actuator_passive_actuation_dense,
-          dim=(d.nworld, qMi.size),
-          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
+          dim=(d.nworld, Mi.size),
+          inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, Mi, Mj],
           outputs=[out],
         )
     wp.launch(
       _qderiv_actuator_passive,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.opt.timestep,
         m.opt.disableflags,
         m.dof_damping,
         m.dof_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.qvel,
-        d.qM,
-        qMi,
-        qMj,
+        d.M,
+        Mi,
+        Mj,
         out,
       ],
       outputs=[out],
     )
   else:
-    # TODO(team): directly utilize qM for these settings
-    wp.copy(out, d.qM)
+    # TODO(team): directly utilize M for these settings
+    wp.copy(out, d.M)
 
   if not (m.opt.disableflags & DisableBit.DAMPER):
     wp.launch(
       _qderiv_tendon_damping,
-      dim=(d.nworld, qMi.size),
+      dim=(d.nworld, Mi.size),
       inputs=[
         m.ntendon,
         m.opt.timestep,
@@ -443,12 +747,11 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
         m.tendon_damping,
         m.tendon_dampingpoly,
         m.is_sparse,
+        m.M_elemid,
         d.ten_J,
         d.ten_velocity,
-        qMi,
-        qMj,
+        Mi,
+        Mj,
       ],
       outputs=[out],
     )
-
-  # TODO(team): rne derivative
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
index c9b3c153..00f22815 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
@@ -20,12 +20,14 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
 from mujoco.mjx.third_party.mujoco_warp._src import constraint
 from mujoco.mjx.third_party.mujoco_warp._src import derivative
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import passive
 from mujoco.mjx.third_party.mujoco_warp._src import sensor
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import solver
+from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src import util_misc
 from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
 from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
@@ -302,6 +304,9 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
     outputs=[d.qpos],
   )
 
+  # advance history buffers before time advance
+  history.insert_ctrl_history(m, d)
+
   wp.launch(
     _next_time,
     dim=d.nworld,
@@ -346,17 +351,18 @@ def _compute_damping_deriv(
 def _euler_damp_qfrc_sparse(
   # Model:
   opt_timestep: wp.array[float],
-  dof_Madr: wp.array[int],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # In:
   damp_deriv: wp.array2d[float],
   # Out:
-  qM_integration_out: wp.array3d[float],
+  M_integration_out: wp.array3d[float],
 ):
   worldid, tid = wp.tid()
   timestep = opt_timestep[worldid % opt_timestep.shape[0]]
 
-  adr = dof_Madr[tid]
-  qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
+  adr = M_rowadr[tid] + M_rownnz[tid] - 1
+  M_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
 
 
 @cache_kernel
@@ -366,7 +372,7 @@ def _tile_euler_dense(tile: TileSet):
     # Model:
     opt_timestep: wp.array[float],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     efc_Ma_in: wp.array2d[float],
     # In:
     damp_deriv: wp.array2d[float],
@@ -379,14 +385,14 @@ def _tile_euler_dense(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr_in[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     damping_tile = wp.tile_load(damp_deriv[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     damping_scaled = damping_tile * timestep
     qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
 
     Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
-    L_tile = wp.tile_cholesky(qm_integration_tile)
-    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
+    L_tile = wp.tile_cholesky(qm_integration_tile, fill_mode="upper")
+    qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile, fill_mode="upper")
     wp.tile_store(qacc_out[worldid], qacc_tile, offset=(dofid))
 
   return euler_dense
@@ -409,22 +415,22 @@ def euler(m: Model, d: Data):
     )
 
     if m.is_sparse:
-      qM = wp.clone(d.qM)
+      M = wp.clone(d.M)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
       wp.launch(
         _euler_damp_qfrc_sparse,
         dim=(d.nworld, m.nv),
-        inputs=[m.opt.timestep, m.dof_Madr, damp_deriv],
-        outputs=[qM],
+        inputs=[m.opt.timestep, m.M_rownnz, m.M_rowadr, damp_deriv],
+        outputs=[M],
       )
-      smooth.factor_solve_i(m, d, qM, qLD, qLDiagInv, qacc, d.efc.Ma)
+      smooth.factor_solve_i(m, d, M, qLD, qLDiagInv, qacc, d.efc.Ma)
     else:
-      for tile in m.qM_tiles:
+      for tile in m.M_tiles:
         wp.launch_tiled(
           _tile_euler_dense(tile),
           dim=(d.nworld, tile.adr.size),
-          inputs=[m.opt.timestep, d.qM, d.efc.Ma, damp_deriv, tile.adr],
+          inputs=[m.opt.timestep, d.M, d.efc.Ma, damp_deriv, tile.adr],
           outputs=[qacc],
           block_dim=m.block_dim.euler_dense,
         )
@@ -573,16 +579,62 @@ def rungekutta4(m: Model, d: Data):
   _advance(m, d, qacc_rk, qvel_rk)
 
 
+@wp.kernel
+def _map_m2d(
+  # Model:
+  mapM2D: wp.array[int],
+  is_sparse: bool,
+  # In:
+  qDi: wp.array[int],
+  qDj: wp.array[int],
+  qH_M: wp.array3d[float],
+  # Data out:
+  qLU_out: wp.array3d[float],
+):
+  worldid, elemid = wp.tid()
+  if is_sparse:
+    m_idx = mapM2D[elemid]
+    if m_idx >= 0:
+      qLU_out[worldid, 0, elemid] = qH_M[worldid, 0, m_idx]
+    else:
+      qLU_out[worldid, 0, elemid] = 0.0
+  else:
+    i = qDi[elemid]
+    j = qDj[elemid]
+    qLU_out[worldid, 0, elemid] = qH_M[worldid, i, j]
+
+
 @event_scope
 def implicit(m: Model, d: Data):
   """Integrates fully implicit in velocity."""
-  if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
+  if m.opt.integrator == IntegratorType.IMPLICIT:
+    qH_M = wp.empty(d.M.shape, dtype=float)
+
+    # 1. Compute M - dt * qDeriv_smooth in M-structure
+    derivative.deriv_smooth_vel(m, d, qH_M)
+
+    # 2. Map M-structure to D-structure
+    wp.launch(
+      _map_m2d,
+      dim=(d.nworld, m.nD),
+      inputs=[m.mapM2D, m.is_sparse, m.qD_fullm_i, m.qD_fullm_j, qH_M],
+      outputs=[d.qLU],
+    )
+
+    # 3. Compute RNE derivatives, scale by timestep, and subtract in-place from qLU
+    derivative.deriv_rne_vel(m, d, d.qLU, flg_subtract=True)
+
+    # 4. Factorize and solve: qacc = qLU \ Ma
+    qacc = wp.empty((d.nworld, m.nv), dtype=float)
+    smooth.factor_solve_lu(m, d, d.qLU, qacc, d.efc.Ma)
+    _advance(m, d, qacc)
+  elif ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
       qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
-      qDeriv = wp.empty(d.qM.shape, dtype=float)
-      qLD = wp.empty(d.qM.shape, dtype=float)
+      qDeriv = wp.empty(d.M.shape, dtype=float)
+      qLD = wp.empty(d.M.shape, dtype=float)
     qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -613,8 +665,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
   if m.opt.run_collision_detection:
     collision_driver.collision(m, d)
   constraint.make_constraint(m, d)
-  # TODO(team): remove False after island features are more complete
-  if False and not (m.opt.disableflags & DisableBit.ISLAND):
+  if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
     island.island(m, d)
   smooth.transmission(m, d)
 
@@ -1098,6 +1149,13 @@ def fwd_actuation(m: Model, d: Data):
     d.actuator_force.zero_()
     return
 
+  # read delayed ctrl (or direct copy if no delay)
+  if m.nhistory > 0:
+    ctrl = wp.empty((d.nworld, m.nu), dtype=float)
+    history.read_ctrl_delayed(m, d, ctrl)
+  else:
+    ctrl = d.ctrl
+
   wp.launch(
     _actuator_force,
     dim=(d.nworld, m.nu),
@@ -1122,7 +1180,7 @@ def fwd_actuation(m: Model, d: Data):
       m.actuator_acc0,
       m.actuator_lengthrange,
       d.act,
-      d.ctrl,
+      ctrl,
       d.actuator_length,
       d.actuator_velocity,
       m.opt.disableflags & DisableBit.CLAMPCTRL,
@@ -1221,7 +1279,7 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
   xfrc_accumulate(m, d, d.qfrc_smooth)
 
   if factorize:
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
   else:
     smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
 
@@ -1266,7 +1324,7 @@ def step(m: Model, d: Data):
     euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
     rungekutta4(m, d)
-  elif m.opt.integrator == IntegratorType.IMPLICITFAST:
+  elif m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
     implicit(m, d)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
@@ -1307,8 +1365,7 @@ def step2(m: Model, d: Data):
   sensor.sensor_acc(m, d)
 
   # integrate with Euler or implicitfast
-  # TODO(team): implicit
-  if m.opt.integrator == IntegratorType.IMPLICITFAST:
+  if m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
     implicit(m, d)
   else:
     # note: RK4 defaults to Euler
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
new file mode 100644
index 00000000..42149578
--- /dev/null
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
@@ -0,0 +1,925 @@
+# Copyright 2026 The Newton Developers
+#
+# 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.
+# ==============================================================================
+
+import warp as wp
+
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import Data
+from mujoco.mjx.third_party.mujoco_warp._src.types import Model
+
+wp.set_module_options({"enable_backward": False})
+
+
+@wp.func
+def _history_physical_index(cursor: int, n: int, logical: int) -> int:
+  """Convert logical index (0=oldest, n-1=newest) to physical index."""
+  return (cursor + 1 + logical) % n
+
+
+@wp.func
+def _history_find_index(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  cursor: int,
+  t: float,
+) -> int:
+  """Find logical index i such that times[i-1] < t <= times[i].
+
+  Returns 0 if t <= times[oldest], n if t > times[newest].
+  Uses circular binary search matching MuJoCo C historyFindIndex.
+  """
+  times_offset = buf_offset + 2
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # before or at first element
+  if t <= t_oldest:
+    return 0
+
+  # after last element
+  if t > t_newest:
+    return n
+
+  # circular binary search: find smallest logical i such that times[phys(i)] >= t
+  lo = int(0)
+  hi = int(n - 1)
+  while hi - lo > 1:
+    mid = int((lo + hi) >> 1)
+    mid_phys = _history_physical_index(cursor, n, mid)
+    if buf[worldid, times_offset + mid_phys] < t:
+      lo = mid
+    else:
+      hi = mid
+
+  return hi
+
+
+@wp.func
+def _history_read_scalar(
+  # In:
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  interp: int,
+) -> float:
+  """Read a scalar value from history buffer at time t.
+
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest
+  if t <= t_oldest + 1e-6:
+    return buf[worldid, values_offset + oldest_phys]
+
+  # extrapolate after newest
+  if t >= t_newest - 1e-6:
+    return buf[worldid, values_offset + newest_phys]
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    return buf[worldid, values_offset + phys_i]
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    return buf[worldid, values_offset + phys_lo]
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+  v_lo = buf[worldid, values_offset + phys_lo]
+  v_hi = buf[worldid, values_offset + phys_hi]
+
+  # linear interpolation
+  if interp == 1:
+    return v_lo + alpha * (v_hi - v_lo)
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  # finite-differenced Catmull-Rom slopes, 0 at endpoints
+  m_lo = 0.0
+  if i > 1:
+    phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+    dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+    m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev]) / dt_lo
+
+  m_hi = 0.0
+  if i < n - 1:
+    phys_hi_next = _history_physical_index(cursor, n, i + 1)
+    dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+    m_hi = (buf[worldid, values_offset + phys_hi_next] - v_lo) / dt_hi
+
+  return h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+
+
+@wp.func
+def _history_read_vector(
+  # In:
+  adr: int,
+  buf: wp.array2d[float],
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  interp: int,
+  # Data out:
+  sensordata_out: wp.array2d[float],
+) -> int:
+  """Read a vector value from history buffer at time t into sensordata.
+
+  Returns 1 on success (value written to sensordata).
+  interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
+  """
+  cursor = int(buf[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  oldest_phys = _history_physical_index(cursor, n, 0)
+  newest_phys = _history_physical_index(cursor, n, n - 1)
+  t_oldest = buf[worldid, times_offset + oldest_phys]
+  t_newest = buf[worldid, times_offset + newest_phys]
+
+  # extrapolate before oldest: copy oldest
+  if t <= t_oldest + 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + oldest_phys * dim + d]
+    return 1
+
+  # extrapolate after newest: copy newest
+  if t >= t_newest - 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + newest_phys * dim + d]
+    return 1
+
+  # find bracketing index
+  i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
+  phys_i = _history_physical_index(cursor, n, i)
+
+  # exact match
+  if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_i * dim + d]
+    return 1
+
+  phys_lo = _history_physical_index(cursor, n, i - 1)
+  phys_hi = phys_i
+
+  # zero-order hold
+  if interp == 0:
+    for d in range(dim):
+      sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_lo * dim + d]
+    return 1
+
+  dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
+  alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
+
+  # linear interpolation
+  if interp == 1:
+    for d in range(dim):
+      v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+      v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+      sensordata_out[worldid, adr + d] = v_lo + alpha * (v_hi - v_lo)
+    return 1
+
+  # cubic spline interpolation (Catmull-Rom)
+  alpha2 = alpha * alpha
+  alpha3 = alpha2 * alpha
+  h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
+  h10 = alpha3 - 2.0 * alpha2 + alpha
+  h01 = -2.0 * alpha3 + 3.0 * alpha2
+  h11 = alpha3 - alpha2
+
+  for d in range(dim):
+    v_lo = buf[worldid, values_offset + phys_lo * dim + d]
+    v_hi = buf[worldid, values_offset + phys_hi * dim + d]
+
+    # finite-differenced Catmull-Rom slopes, 0 at endpoints
+    m_lo = 0.0
+    if i > 1:
+      phys_lo_prev = _history_physical_index(cursor, n, i - 2)
+      dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
+      m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev * dim + d]) / dt_lo
+
+    m_hi = 0.0
+    if i < n - 1:
+      phys_hi_next = _history_physical_index(cursor, n, i + 1)
+      dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
+      m_hi = (buf[worldid, values_offset + phys_hi_next * dim + d] - v_lo) / dt_hi
+
+    sensordata_out[worldid, adr + d] = h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
+  return 1
+
+
+@wp.func
+def _history_insert_scalar(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  t: float,
+  value: float,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a scalar value into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      buf_out[worldid, values_offset + phys_i] = value
+      return
+
+  # older than oldest: replace oldest
+  if i == 0:
+    oldest_phys = _history_physical_index(cursor, n, 0)
+    buf_out[worldid, times_offset + oldest_phys] = t
+    buf_out[worldid, values_offset + oldest_phys] = value
+    return
+
+  # newer than newest: advance cursor
+  if i == n:
+    cursor = (cursor + 1) % n
+    buf_out[worldid, buf_offset + 1] = float(cursor)
+    buf_out[worldid, times_offset + cursor] = t
+    buf_out[worldid, values_offset + cursor] = value
+    return
+
+  # out-of-order: shift [1, i-1] left, insert at i-1
+  for j in range(i - 1):
+    src_phys = _history_physical_index(cursor, n, j + 1)
+    dst_phys = _history_physical_index(cursor, n, j)
+    buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+    buf_out[worldid, values_offset + dst_phys] = buf_out[worldid, values_offset + src_phys]
+  insert_phys = _history_physical_index(cursor, n, i - 1)
+  buf_out[worldid, times_offset + insert_phys] = t
+  buf_out[worldid, values_offset + insert_phys] = value
+
+
+@wp.func
+def _history_insert_vector(
+  # In:
+  worldid: int,
+  buf_offset: int,
+  n: int,
+  dim: int,
+  t: float,
+  src: wp.array2d[float],
+  src_adr: int,
+  # Out:
+  buf_out: wp.array2d[float],
+):
+  """Insert a vector value from src[worldid, src_adr:src_adr+dim] into history buffer at time t."""
+  cursor = int(buf_out[worldid, buf_offset + 1])
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + n
+
+  i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
+
+  slot_phys = -1
+
+  # exact match
+  if i < n:
+    phys_i = _history_physical_index(cursor, n, i)
+    if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
+      slot_phys = phys_i
+
+  if slot_phys < 0:
+    if i == 0:
+      # older than oldest: replace oldest
+      slot_phys = _history_physical_index(cursor, n, 0)
+      buf_out[worldid, times_offset + slot_phys] = t
+    elif i == n:
+      # newer than newest: advance cursor
+      cursor = (cursor + 1) % n
+      buf_out[worldid, buf_offset + 1] = float(cursor)
+      slot_phys = cursor
+      buf_out[worldid, times_offset + slot_phys] = t
+    else:
+      # out-of-order: shift [1, i-1] left, insert at i-1
+      for j in range(i - 1):
+        src_phys = _history_physical_index(cursor, n, j + 1)
+        dst_phys = _history_physical_index(cursor, n, j)
+        buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
+        for d in range(dim):
+          buf_out[worldid, values_offset + dst_phys * dim + d] = buf_out[worldid, values_offset + src_phys * dim + d]
+      slot_phys = _history_physical_index(cursor, n, i - 1)
+      buf_out[worldid, times_offset + slot_phys] = t
+
+  # copy values
+  for d in range(dim):
+    buf_out[worldid, values_offset + slot_phys * dim + d] = src[worldid, src_adr + d]
+
+
+@wp.kernel
+def _read_ctrl_delayed_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  ctrl_out: wp.array2d[float],
+):
+  """Read delayed ctrl for each actuator."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  delay = actuator_delay[uid]
+
+  if nsample == 0 or delay == 0.0:
+    # no delay: direct copy
+    ctrl_out[worldid, uid] = ctrl_in[worldid, uid]
+  else:
+    interp = hist[1]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    ctrl_out[worldid, uid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp)
+
+
+@wp.kernel
+def _insert_ctrl_history_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # Data in:
+  time_in: wp.array[float],
+  ctrl_in: wp.array2d[float],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current ctrl into history buffers."""
+  worldid, uid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = actuator_historyadr[uid]
+  t = time_in[worldid]
+  value = ctrl_in[worldid, uid]
+  _history_insert_scalar(worldid, buf_offset, nsample, t, value, history_out)
+
+
+@wp.kernel
+def _insert_sensor_history_stage(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  """Insert current sensor values into history buffers for specific sensor IDs."""
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample == 0:
+    return
+
+  buf_offset = sensor_historyadr[sid]
+  dim = sensor_dim[sid]
+  interval_val = sensor_interval[sid]
+  period = interval_val[0]
+  t = time_in[worldid]
+
+  if period > 0.0:
+    # interval mode: check if condition is satisfied
+    time_prev = history_out[worldid, buf_offset]  # user slot stores time_prev
+    if time_prev + period <= t:
+      # advance time_prev by exact period
+      history_out[worldid, buf_offset] = time_prev + period
+      # insert sensor value
+      _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+  else:
+    _history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
+
+
+@wp.kernel
+def _apply_sensor_delay_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  sensor_interval: wp.array[wp.vec2],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  # In:
+  sensor_ids: wp.array[int],
+  # Data out:
+  sensordata_out: wp.array2d[float],
+):
+  """Apply delay/interval logic for sensors after computation.
+
+  TODO(team): Revisit always-compute decision for computationally expensive sensors
+  with interval/period (e.g., raytracers)
+  """
+  worldid, idx = wp.tid()
+  sid = sensor_ids[idx]
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  if nsample <= 0:
+    return
+
+  delay = sensor_delay[sid]
+  dim = sensor_dim[sid]
+  interp = hist[1]
+  buf_offset = sensor_historyadr[sid]
+  t = time_in[worldid]
+
+  if delay > 0.0:
+    # delay > 0: read delayed value from buffer
+    _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t - delay, interp, sensordata_out)
+  else:
+    # interval-only (delay == 0, interval > 0): check interval condition
+    interval_val = sensor_interval[sid]
+    period = interval_val[0]
+    if period > 0.0:
+      time_prev = history_in[worldid, buf_offset]  # user slot
+      if time_prev + period > t:
+        # interval condition not satisfied: read from buffer
+        _history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t, interp, sensordata_out)
+      # else: interval condition satisfied, keep computed value
+
+
+def read_ctrl_delayed(m: Model, d: Data, ctrl: wp.array2d[float]):
+  """Read delayed ctrl values for all actuators."""
+  if m.nhistory == 0:
+    wp.copy(ctrl, d.ctrl)
+    return
+
+  wp.launch(
+    _read_ctrl_delayed_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      d.time,
+      d.history,
+      d.ctrl,
+    ],
+    outputs=[ctrl],
+  )
+
+
+def insert_ctrl_history(m: Model, d: Data):
+  """Insert current ctrl values into history buffers."""
+  if m.nhistory == 0 or m.nu == 0:
+    return
+
+  wp.launch(
+    _insert_ctrl_history_kernel,
+    dim=(d.nworld, m.nu),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      d.time,
+      d.ctrl,
+    ],
+    outputs=[d.history],
+  )
+
+
+def apply_sensor_delay(m: Model, d: Data, sensorid: wp.array[int]):
+  """Apply delay/interval logic for given sensors after computation.
+
+  Matches MuJoCo C architecture where the delayed read (mj_sensorPos) occurs
+  before the fresh value insert (mj_advance). We save fresh sensordata,
+  overwrite with delayed values, then insert the saved fresh values.
+  """
+  if m.nhistory == 0 or sensorid.shape[0] == 0:
+    return
+
+  # Save fresh sensordata before delay overwrite
+  fresh_sensordata = wp.empty_like(d.sensordata)
+  wp.copy(fresh_sensordata, d.sensordata)
+
+  # Read delayed values from buffer → overwrite sensordata
+  wp.launch(
+    _apply_sensor_delay_kernel,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      d.history,
+      sensorid,
+    ],
+    outputs=[d.sensordata],
+  )
+
+  # Insert saved fresh sensor values into history buffers
+  wp.launch(
+    _insert_sensor_history_stage,
+    dim=(d.nworld, sensorid.shape[0]),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      m.sensor_interval,
+      d.time,
+      fresh_sensordata,
+      sensorid,
+    ],
+    outputs=[d.history],
+  )
+
+
+@wp.kernel
+def _read_ctrl_kernel(
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  actuator_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  ctrl_in: wp.array2d[float],
+  # In:
+  uid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  hist = actuator_history[uid]
+  nsample = hist[0]
+
+  if nsample == 0:
+    result_out[worldid] = ctrl_in[worldid, uid]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = actuator_delay[uid]
+    buf_offset = actuator_historyadr[uid]
+    t = time_in[worldid] - delay
+    result_out[worldid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp_val)
+
+
+def read_ctrl(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed ctrl for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld,).
+  """
+  wp.launch(
+    _read_ctrl_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      m.actuator_delay,
+      time,
+      d.history,
+      d.ctrl,
+      ctrlid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _read_sensor_kernel(
+  # Model:
+  sensor_dim: wp.array[int],
+  sensor_adr: wp.array[int],
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_delay: wp.array[float],
+  # Data in:
+  time_in: wp.array[float],
+  history_in: wp.array2d[float],
+  sensordata_in: wp.array2d[float],
+  # In:
+  sid: int,
+  interp: int,
+  # Out:
+  result_out: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  hist = sensor_history[sid]
+  nsample = hist[0]
+  dim = sensor_dim[sid]
+  adr = sensor_adr[sid]
+
+  if nsample == 0:
+    for i in range(dim):
+      result_out[worldid, i] = sensordata_in[worldid, adr + i]
+  else:
+    interp_val = interp
+    if interp_val < 0:
+      interp_val = hist[1]
+    delay = sensor_delay[sid]
+    buf_offset = sensor_historyadr[sid]
+    t = time_in[worldid] - delay
+    _history_read_vector(
+      0,  # write to result_out starting at index 0 (not global sensor adr)
+      history_in,
+      worldid,
+      buf_offset,
+      nsample,
+      dim,
+      t,
+      interp_val,
+      result_out,
+    )
+
+
+def read_sensor(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  time: wp.array[float],
+  interp: int,
+  result: wp.array2d[float],
+):
+  """Read delayed sensor for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    time: query time per world (nworld,).
+    interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
+    result: output buffer (nworld, dim).
+  """
+  wp.launch(
+    _read_sensor_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_dim,
+      m.sensor_adr,
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_delay,
+      time,
+      d.history,
+      d.sensordata,
+      sensorid,
+      interp,
+    ],
+    outputs=[result],
+  )
+
+
+@wp.kernel
+def _init_ctrl_history_kernel(
+  # kernel_analyzer: off
+  # Model:
+  actuator_history: wp.array[wp.vec2i],
+  actuator_historyadr: wp.array[int],
+  # In:
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+  # kernel_analyzer: on
+):
+  """Initialize history buffer for 1 actuator across all worlds."""
+  worldid = wp.tid()
+
+  nsample = actuator_history[ctrlid][0]
+  buf_offset = actuator_historyadr[ctrlid]
+
+  # preserve user slot
+  user = history_out[worldid, buf_offset]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    history_out[worldid, values_offset + i] = values[worldid, i]
+
+  # restore user slot
+  history_out[worldid, buf_offset] = user
+
+
+def init_ctrl_history(
+  m: Model,
+  d: Data,
+  ctrlid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+):
+  """Initialize history buffer for 1 actuator across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    ctrlid: actuator index.
+    times: timestamps or None (nsample,).
+    values: ctrl values (nworld, nsample).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_ctrl_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.actuator_history,
+      m.actuator_historyadr,
+      ctrlid,
+      times,
+      values,
+      has_times,
+    ],
+    outputs=[d.history],
+  )
+
+
+# kernel_analyzer: off
+@wp.kernel
+def _init_sensor_history_kernel(
+  # Model:
+  sensor_history: wp.array[wp.vec2i],
+  sensor_historyadr: wp.array[int],
+  sensor_dim_arr: wp.array[int],
+  # In:
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+  has_times: int,
+  # Data out:
+  history_out: wp.array2d[float],
+):
+  # kernel_analyzer: on
+  """Initialize history buffer for 1 sensor across all worlds."""
+  worldid = wp.tid()
+
+  nsample = sensor_history[sensorid][0]
+  dim = sensor_dim_arr[sensorid]
+  buf_offset = sensor_historyadr[sensorid]
+
+  # set user slot (phase = last computation time for interval sensors)
+  history_out[worldid, buf_offset] = phase[worldid]
+
+  # cursor = 0 (samples in order, newest at index nsample-1)
+  history_out[worldid, buf_offset + 1] = float(nsample - 1)
+
+  times_offset = buf_offset + 2
+  values_offset = buf_offset + 2 + nsample
+
+  for i in range(nsample):
+    if has_times != 0:
+      history_out[worldid, times_offset + i] = times[i]
+    else:
+      history_out[worldid, times_offset + i] = -MJ_MAXVAL
+    for j in range(dim):
+      history_out[worldid, values_offset + i * dim + j] = values[worldid, i * dim + j]
+
+
+def init_sensor_history(
+  m: Model,
+  d: Data,
+  sensorid: int,
+  times: wp.array[float],
+  values: wp.array2d[float],
+  phase: wp.array[float],
+):
+  """Initialize history buffer for 1 sensor across all worlds.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    sensorid: sensor index.
+    times: timestamps or None (nsample,).
+    values: sensor values (nworld, nsample * dim).
+    phase: user slot value per world (nworld,).
+
+  Raises:
+    ValueError: If times are not strictly increasing.
+  """
+  has_times = 0 if times is None else 1
+  if times is not None:
+    t_np = times.numpy()
+    for i in range(len(t_np) - 1):
+      if t_np[i + 1] - t_np[i] < MJ_MINVAL:
+        raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
+  if times is None:
+    times = wp.empty(0, dtype=float)
+
+  wp.launch(
+    _init_sensor_history_kernel,
+    dim=(d.nworld,),
+    inputs=[
+      m.sensor_history,
+      m.sensor_historyadr,
+      m.sensor_dim,
+      sensorid,
+      times,
+      values,
+      phase,
+      has_times,
+    ],
+    outputs=[d.history],
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
index 82b7cf4a..388ad1aa 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
@@ -95,9 +95,9 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
 
     # TODO(team): qacc = d.qacc if (m.dof_damping == 0.0).all()
 
-    # set qfrc = (d.qM + m.opt.timestep * diag(m.dof_damping)) * d.qacc
+    # set qfrc = (d.M + m.opt.timestep * diag(m.dof_damping)) * d.qacc
 
-    # d.qM @ d.qacc
+    # d.M @ d.qacc
     support.mul_m(m, d, qfrc, d.qacc)
 
     # qfrc += m.opt.timestep * damp_deriv * d.qacc
@@ -109,16 +109,16 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
     )
   elif m.opt.integrator == IntegratorType.IMPLICITFAST:
     if m.is_sparse:
-      qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
+      qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
     else:
       qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
     derivative.deriv_smooth_vel(m, d, qDeriv)
     mul_m(m, d, qfrc, d.qacc, M=qDeriv)
-    smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, qacc, qfrc)
+    smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, qacc, qfrc)
   else:
     raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
 
-  # solve for qacc: qfrc = d.qM @ d.qacc
+  # solve for qacc: qfrc = d.M @ d.qacc
   smooth.solve_m(m, d, qacc, qfrc)
 
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
index 309f9e41..ffd9f8f3 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
@@ -33,6 +33,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
 from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
 from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 
+# TODO(team): remove after improving island solver performance
+ENABLE_ISLANDS = False
+
 
 def _is_array_spec(typ) -> bool:
   """Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -64,6 +67,61 @@ def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
   return array
 
 
+def _create_constraint(
+  mjm,
+  nworld: int,
+  njmax: int,
+  njmax_nnz: int,
+  sizes: dict,
+  island_enabled: bool,
+  mjd=None,
+) -> types.Constraint:
+  """Construct a types.Constraint with standard and island local fields allocated properly."""
+  efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
+  sparse = is_sparse(mjm)
+
+  for f in dataclasses.fields(types.Constraint):
+    if f.name == "itype":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iid":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    elif f.name == "iJ_rownnz":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_rowadr":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
+    elif f.name == "iJ_colind":
+      efc_kwargs[f.name] = wp.empty((nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, 0, 0), dtype=int)
+    elif f.name == "iJ":
+      efc_kwargs[f.name] = wp.empty(
+        (nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, njmax if island_enabled else 0, mjm.nv),
+        dtype=float,
+      )
+    elif f.name == "iD":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iaref":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "ifrictionloss":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "iforce":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
+    elif f.name == "istate":
+      efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
+    else:
+      if f.name in efc_kwargs:
+        continue
+
+      if mjd is not None:
+        shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
+        val = np.zeros(shape, dtype=f.type.dtype)
+        if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
+          val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
+        efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
+      else:
+        efc_kwargs[f.name] = _create_array(None, f.type, sizes)
+
+  return types.Constraint(**efc_kwargs)
+
+
 def is_sparse(mjm: mujoco.MjModel) -> bool:
   if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
     if mjm.nv > 32:
@@ -183,6 +241,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
   opt = types.Option(**opt_kwargs)
 
+  # islands are disabled by default while performance is being improved
+  # override by setting io.ENABLE_ISLANDS = True
+  # TODO(team): remove after improving island solver performance
+  if not ENABLE_ISLANDS:
+    opt.disableflags |= types.DisableBit.ISLAND
+
   # C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
   # adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
   # bouncing around the optimal solution
@@ -271,7 +335,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE))
   m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0]
   m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0]
-  m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv)
+  m.dof_tri_row, m.dof_tri_col = np.triu_indices(mjm.nv)
 
   # precompute body_isdofancestor: which DOFs affect each body
   # TODO: Investigate alternative approach such as bitmap
@@ -585,14 +649,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
     for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
   ]
 
-  # qM_tiles records the block diagonal structure of qM
+  # M_tiles records the block diagonal structure of M
   tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
   tiles = {}
   for i in range(len(tile_corners)):
     tile_beg = tile_corners[i]
     tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
     tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
-  m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
+  m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
 
   # qLD_updates has dof tree ordering of qLD updates for sparse factor m
   qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
@@ -620,41 +684,69 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
   m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
   m.qLD_level_offsets = level_offsets
 
-  # indices for sparse qM_fullm (used in solver)
-  m.qM_fullm_i, m.qM_fullm_j = [], []
+  # Indices for sparse M_fullm (used in solver). M_fullm_i/j are built by
+  # walking dof_parentid for each dof, so for joint types whose internal block
+  # MuJoCo stores diagonal-only in the compact (M_rownnz, M_rowadr) layout
+  # (e.g. free joints), the chain-aware layout here has more entries per row
+  # than the compact layout.
+  m.M_fullm_i, m.M_fullm_j = [], []
   for i in range(mjm.nv):
     j = i
     while j > -1:
-      m.qM_fullm_i.append(i)
-      m.qM_fullm_j.append(j)
+      m.M_fullm_i.append(i)
+      m.M_fullm_j.append(j)
       j = mjm.dof_parentid[j]
+  # M_elemid maps (row, col) -> madr index in the native CSR M layout
+  M_elemid = np.full((mjm.nv, mjm.nv), -1, dtype=np.int32)
+  for i in range(mjm.nv):
+    rowadr = mjm.M_rowadr[i]
+    rownnz = mjm.M_rownnz[i]
+    for k in range(rownnz):
+      madr = rowadr + k
+      col = int(mjm.M_colind[madr])
+      M_elemid[i, col] = madr
+  m.M_elemid = M_elemid
+
+  upper_j, upper_i = np.triu_indices(mjm.nv)
+  upper_elemid = M_elemid[upper_i, upper_j]
+  valid_mask = upper_elemid != -1
+  m.M_fullm_upper_i = upper_j[valid_mask].tolist()
+  m.M_fullm_upper_j = upper_i[valid_mask].tolist()
+  m.M_fullm_upper_elemid = upper_elemid[valid_mask].tolist()
+
+  # indices for sparse qD_fullm (used in RNE derivatives)
+  # D-structure is the full square sparsity pattern (both upper and lower triangle)
+  m.qD_fullm_i, m.qD_fullm_j = [], []
+  for i in range(mjm.nv):
+    rowadr = mjm.D_rowadr[i]
+    rownnz = mjm.D_rownnz[i]
+    for k in range(rownnz):
+      m.qD_fullm_i.append(i)
+      m.qD_fullm_j.append(int(mjm.D_colind[rowadr + k]))
+  m.nD = mjm.nD
 
   # Gather-based sparse mul_m: for each row, all (col, madr) including diagonal
   row_elements = [[] for _ in range(mjm.nv)]
 
-  # Add diagonal
   for i in range(mjm.nv):
-    row_elements[i].append((i, mjm.dof_Madr[i]))
-
-  # Add off-diagonals: ancestors (lower) and descendants (upper)
-  for i in range(mjm.nv):
-    madr_ij, j = mjm.dof_Madr[i], i
-    while True:
-      madr_ij, j = madr_ij + 1, mjm.dof_parentid[j]
-      if j == -1:
-        break
-      row_elements[i].append((j, madr_ij))  # row i gathers M[i,j] * vec[j]
-      row_elements[j].append((i, madr_ij))  # row j gathers M[j,i] * vec[i]
+    rowadr = mjm.M_rowadr[i]
+    rownnz = mjm.M_rownnz[i]
+    for k in range(rownnz):
+      madr = rowadr + k
+      col = int(mjm.M_colind[madr])
+      row_elements[i].append((col, madr))  # row i gathers M[i,col] * vec[col]
+      if i != col:
+        row_elements[col].append((i, madr))  # row col gathers M[i,col] * vec[i]
 
   # Flatten into CSR-like arrays
-  m.qM_mulm_rowadr = [0]
-  m.qM_mulm_col = []
-  m.qM_mulm_madr = []
+  m.M_mulm_rowadr = [0]
+  m.M_mulm_col = []
+  m.M_mulm_madr = []
   for i in range(mjm.nv):
     for col, madr in row_elements[i]:
-      m.qM_mulm_col.append(col)
-      m.qM_mulm_madr.append(madr)
-    m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
+      m.M_mulm_col.append(col)
+      m.M_mulm_madr.append(madr)
+    m.M_mulm_rowadr.append(len(m.M_mulm_col))
 
   m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
   m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
@@ -886,6 +978,42 @@ def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int
   return default
 
 
+def _allocate_island_arrays(
+  mjm: mujoco.MjModel,
+  d: types.Data,
+  nworld: int,
+  njmax: int,
+  island_enabled: bool,
+  mjd: mujoco.MjData,
+):
+  ntree_size = mjm.ntree if island_enabled else 0
+  nv_size = mjm.nv if island_enabled else 0
+  njmax_size = njmax if island_enabled else 0
+
+  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
+  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+  d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
+
+  d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_nf = wp.empty((nworld, ntree_size), dtype=int)
+  d.island_efcadr = wp.empty((nworld, ntree_size), dtype=int)
+  d.nidof = wp.empty((nworld if island_enabled else 0,), dtype=int)
+  d.map_dof2idof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_idof2dof = wp.empty((nworld, nv_size), dtype=int)
+  d.map_efc2iefc = wp.empty((nworld, njmax_size), dtype=int)
+  d.map_iefc2efc = wp.empty((nworld, njmax_size), dtype=int)
+
+  d.dof_islandid = wp.empty((nworld, nv_size), dtype=int)
+  d.efc_islandid = wp.empty((nworld, njmax_size), dtype=int)
+  d.iqacc = wp.empty((nworld, nv_size), dtype=float)
+  d.iqacc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_smooth = wp.empty((nworld, nv_size), dtype=float)
+  d.iqfrc_constraint = wp.empty((nworld, nv_size), dtype=float)
+
+
 def make_data(
   mjm: mujoco.MjModel,
   nworld: int = 1,
@@ -967,7 +1095,8 @@ def make_data(
 
   contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
   contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
-  efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
+
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS)
 
   if is_sparse(mjm):
     efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1005,7 +1134,7 @@ def make_data(
     "njmax": njmax,
     "njmax_pad": sizes["njmax_pad"],
     "njmax_nnz": njmax_nnz,
-    "qM": None,
+    "M": None,
     "qLD": None,
     # world body
     "xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
@@ -1024,6 +1153,24 @@ def make_data(
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1033,15 +1180,13 @@ def make_data(
   d = types.Data(**d_kwargs)
 
   if is_sparse(mjm):
-    d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
+    d.M = wp.zeros((nworld, 1, mjm.nC), dtype=float)
     d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
   else:
-    d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
+    d.M = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
     d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
 
-  # island discovery arrays
-  d.nisland = wp.zeros((nworld,), dtype=int)
-  d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
 
   return d
 
@@ -1176,16 +1321,7 @@ def put_data(
   # create efc
   efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
 
-  for f in dataclasses.fields(types.Constraint):
-    if f.name in efc_kwargs:
-      continue
-    shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
-    val = np.zeros(shape, dtype=f.type.dtype)
-    if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
-      val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
-    efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
-
-  efc = types.Constraint(**efc_kwargs)
+  efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS, mjd)
 
   if is_sparse(mjm):
     J_rownnz = np.zeros(njmax, dtype=np.int32)
@@ -1236,12 +1372,30 @@ def put_data(
     "njmax_nnz": njmax_nnz,
     # fields set after initialization:
     "solver_niter": None,
-    "qM": None,
+    "M": None,
     "qLD": None,
     "nacon": None,
     # island arrays
     "nisland": None,
     "tree_island": None,
+    "dof_island": None,
+    "island_dofadr": None,
+    "island_nv": None,
+    "island_nefc": None,
+    "island_ne": None,
+    "island_nf": None,
+    "island_efcadr": None,
+    "nidof": None,
+    "map_dof2idof": None,
+    "map_idof2dof": None,
+    "map_efc2iefc": None,
+    "map_iefc2efc": None,
+    "dof_islandid": None,
+    "efc_islandid": None,
+    "iqacc": None,
+    "iqacc_smooth": None,
+    "iqfrc_smooth": None,
+    "iqfrc_constraint": None,
   }
   for f in dataclasses.fields(types.Data):
     if f.name in d_kwargs:
@@ -1256,20 +1410,25 @@ def put_data(
   d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
 
   if is_sparse(mjm):
-    d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.M), dtype=float)
+    else:
+      d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.qM[mjm.mapM2M]), dtype=float)
     d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
   else:
-    qM = np.zeros((mjm.nv, mjm.nv))
-    mujoco.mj_fullM(mjm, qM, mjd.qM)
-    qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+    M = np.zeros((mjm.nv, mjm.nv))
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      mujoco.mju_sym2dense(M, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
+      qLD = np.linalg.cholesky(M).T if (mjd.M != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
+    else:
+      mujoco.mj_fullM(mjm, M, mjd.qM)
+      qLD = np.linalg.cholesky(M).T if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
     padding = sizes["nv_pad"] - mjm.nv
-    qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
-    d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
+    M_padded = np.pad(M, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
+    d.M = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), M_padded), dtype=float)
     d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
 
-  # island arrays
-  d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
-  d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int)
+  _allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
 
   d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
 
@@ -1403,6 +1562,9 @@ def get_data_into(
   result.qfrc_constraint[:] = d.qfrc_constraint.numpy()[world_id]
   result.qfrc_inverse[:] = d.qfrc_inverse.numpy()[world_id]
 
+  if mjm.nhistory > 0:
+    result.history[:] = d.history.numpy()[world_id]
+
   # contact
   result.contact.dist[:ncon] = d.contact.dist.numpy()[ncon_filter]
   result.contact.pos[:ncon] = d.contact.pos.numpy()[ncon_filter]
@@ -1417,17 +1579,27 @@ def get_data_into(
   result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
 
   if is_sparse(mjm):
-    result.qM[:] = d.qM.numpy()[world_id, 0]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      result.M[:] = d.M.numpy()[world_id, 0]
+    else:
+      result.qM[mjm.mapM2M] = d.M.numpy()[world_id, 0]
     result.qLD[:] = d.qLD.numpy()[world_id, 0]
   else:
-    qM = d.qM.numpy()[world_id]
-    adr = 0
-    for i in range(mjm.nv):
-      j = i
-      while j >= 0:
-        result.qM[adr] = qM[i, j]
-        j = mjm.dof_parentid[j]
-        adr += 1
+    M = d.M.numpy()[world_id]
+    if check_version("mujoco>=3.8.1.dev910242375"):
+      for i in range(mjm.nv):
+        adr = mjm.M_rowadr[i]
+        for k in range(mjm.M_rownnz[i]):
+          col = mjm.M_colind[adr + k]
+          result.M[adr + k] = M[i, col]
+    else:
+      adr = 0
+      for i in range(mjm.nv):
+        j = i
+        while j >= 0:
+          result.qM[adr] = M[i, j]
+          j = mjm.dof_parentid[j]
+          adr += 1
     mujoco.mj_factorM(mjm, result)
 
   if nefc > 0:
@@ -1466,6 +1638,7 @@ def get_data_into(
   result.efc_frictionloss[:] = d.efc.frictionloss.numpy()[world_id, efc_idx]
   result.efc_state[:] = d.efc.state.numpy()[world_id, efc_idx]
   result.efc_force[:] = d.efc.force.numpy()[world_id, efc_idx]
+  result.efc_island[:] = d.efc.island.numpy()[world_id, efc_idx]
 
   # rne_postconstraint
   result.cacc[:] = d.cacc.numpy()[world_id]
@@ -1487,8 +1660,51 @@ def get_data_into(
   # islands
   nisland = d.nisland.numpy()[world_id]
   result.nisland = nisland
-  if nisland:
+  if d.tree_island.shape[1] > 0 and nisland:
     result.tree_island[:] = d.tree_island.numpy()[world_id]
+    result.dof_island[:] = d.dof_island.numpy()[world_id]
+    result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
+    result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
+    result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
+    result.island_ne[:nisland] = d.island_ne.numpy()[world_id, :nisland]
+    result.island_nf[:nisland] = d.island_nf.numpy()[world_id, :nisland]
+    result.island_iefcadr[:nisland] = d.island_efcadr.numpy()[world_id, :nisland]
+    nv = mjm.nv
+    result.map_dof2idof[:nv] = d.map_dof2idof.numpy()[world_id, :nv]
+    result.map_idof2dof[:nv] = d.map_idof2dof.numpy()[world_id, :nv]
+    result.map_efc2iefc[:nefc] = d.map_efc2iefc.numpy()[world_id, :nefc]
+    result.map_iefc2efc[:nefc] = d.map_iefc2efc.numpy()[world_id, :nefc]
+
+    result.iefc_type[:nefc] = d.efc.itype.numpy()[world_id, :nefc]
+    result.iefc_id[:nefc] = d.efc.iid.numpy()[world_id, :nefc]
+    result.iefc_D[:nefc] = d.efc.iD.numpy()[world_id, :nefc]
+    result.iefc_aref[:nefc] = d.efc.iaref.numpy()[world_id, :nefc]
+    result.iefc_frictionloss[:nefc] = d.efc.ifrictionloss.numpy()[world_id, :nefc]
+    result.iefc_state[:nefc] = d.efc.istate.numpy()[world_id, :nefc]
+    result.iefc_force[:nefc] = d.efc.iforce.numpy()[world_id, :nefc]
+
+    if is_sparse(mjm):
+      iefc_J = np.zeros((nefc, mjm.nv))
+      mujoco.mju_sparse2dense(
+        iefc_J,
+        d.efc.iJ.numpy()[world_id, 0],
+        d.efc.iJ_rownnz.numpy()[world_id, :nefc],
+        d.efc.iJ_rowadr.numpy()[world_id, :nefc],
+        d.efc.iJ_colind.numpy()[world_id, 0],
+      )
+    else:
+      iefc_J = d.efc.iJ.numpy()[world_id, :nefc, : mjm.nv]
+
+    if mujoco.mj_isSparse(mjm):
+      mujoco.mju_dense2sparse(
+        result.iefc_J,
+        iefc_J,
+        result.iefc_J_rownnz,
+        result.iefc_J_rowadr,
+        result.iefc_J_colind,
+      )
+    else:
+      result.iefc_J[: nefc * mjm.nv] = iefc_J.flatten()
 
 
 def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@@ -1511,14 +1727,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
     xfrc_applied_out[worldid, bodyid][elemid] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
-  def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
+  def reset_M(reset_in: wp.array[bool], M_out: wp.array3d[float]):
     worldid, elemid1, elemid2 = wp.tid()
 
     if wp.static(reset is not None):
       if not reset_in[worldid]:
         return
 
-    qM_out[worldid, elemid1, elemid2] = 0.0
+    M_out[worldid, elemid1, elemid2] = 0.0
 
   @wp.kernel(module="unique", enable_backward=False)
   def reset_nworld(
@@ -1670,10 +1886,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
 
   wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
   wp.launch(
-    reset_qM,
-    dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]),
+    reset_M,
+    dim=(d.nworld, d.M.shape[1], d.M.shape[2]),
     inputs=[reset_input],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
   # set mocap_pos/quat = body_pos/quat for mocap bodies
@@ -1786,11 +2002,12 @@ def _copy_tendon_length0(
 def _compute_meaninertia(
   nv: int,
   is_sparse: bool,
-  dof_Madr_in: wp.array[int],
-  qM_in: wp.array3d[float],
+  M_rownnz_in: wp.array[int],
+  M_rowadr_in: wp.array[int],
+  M_in: wp.array3d[float],
   meaninertia_out: wp.array[float],
 ):
-  """Compute mean diagonal inertia from qM at qpos0."""
+  """Compute mean diagonal inertia from M at qpos0."""
   worldid = wp.tid()
 
   if nv == 0:
@@ -1800,12 +2017,12 @@ def _compute_meaninertia(
   total = float(0.0)
   for i in range(nv):
     if is_sparse:
-      # Sparse: qM is flattened lower triangular, diagonal at dof_Madr[i]
-      madr = dof_Madr_in[i]
-      total += qM_in[worldid, 0, madr]
+      # Sparse: M is in CSR format, diagonal at M_rowadr_in[i] + M_rownnz_in[i] - 1
+      madr = M_rowadr_in[i] + M_rownnz_in[i] - 1
+      total += M_in[worldid, 0, madr]
     else:
-      # Dense: qM is 2D matrix, diagonal at [i,i]
-      total += qM_in[worldid, i, i]
+      # Dense: M is 2D matrix, diagonal at [i,i]
+      total += M_in[worldid, i, i]
 
   meaninertia_out[worldid % meaninertia_out.shape[0]] = total / float(nv)
 
@@ -2288,11 +2505,11 @@ def set_const_0(m: types.Model, d: types.Data):
   smooth.factor_m(m, d)
   smooth.transmission(m, d)
 
-  # Compute meaninertia from qM diagonal at qpos0
+  # Compute meaninertia from M diagonal at qpos0
   wp.launch(
     _compute_meaninertia,
     dim=d.nworld,
-    inputs=[m.nv, m.is_sparse, m.dof_Madr, d.qM],
+    inputs=[m.nv, m.is_sparse, m.M_rownnz, m.M_rowadr, d.M],
     outputs=[m.stat.meaninertia],
   )
 
@@ -2452,27 +2669,31 @@ def set_const(m: types.Model, d: types.Data):
 
   Model fields that can be modified safely with set_const:
 
-    Field                            | Notes
-    ---------------------------------|----------------------------------------------
-    qpos0, qpos_spring               |
-    body_mass, body_inertia,         | Mass and inertia are usually scaled together
-    body_ipos, body_iquat            | since inertia is sum(m * r^2).
-    body_pos, body_quat              | Unsafe for static bodies (invalidates BVH).
-    body_gravcomp                    | If changing from 0 to >0 bodies, required.
-    dof_armature                     |
-    eq_data                          | For connect/weld, offsets computed if not set.
-    hfield_size                      |
-    tendon_stiffness, tendon_damping | Only if changing from/to zero.
-    actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
+  ==================================  ==============================================
+  Field                               Notes
+  ==================================  ==============================================
+  qpos0, qpos_spring
+  body_mass, body_inertia,            Mass and inertia are usually scaled together
+  body_ipos, body_iquat               since inertia is sum(m * r^2).
+  body_pos, body_quat                 Unsafe for static bodies (invalidates BVH).
+  body_gravcomp                       If changing from 0 to >0 bodies, required.
+  dof_armature
+  eq_data                             For connect/weld, offsets computed if not set.
+  hfield_size
+  tendon_stiffness, tendon_damping    Only if changing from/to zero.
+  actuator_gainprm, actuator_biasprm  For position actuators with dampratio.
+  ==================================  ==============================================
 
   For selective updates, use the sub-functions directly based on what changed:
 
-    Modified Field  | Call
-    ----------------|------------------
-    body_mass       | set_const
-    body_gravcomp   | set_const_fixed
-    body_inertia    | set_const_0
-    qpos0           | set_const_0
+  ==============  ===============
+  Modified Field  Call
+  ==============  ===============
+  body_mass       set_const
+  body_gravcomp   set_const_fixed
+  body_inertia    set_const_0
+  qpos0           set_const_0
+  ==============  ===============
 
   Computes:
     - Fixed quantities (via set_const_fixed):
@@ -2621,6 +2842,9 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
       else:
         val = typ(val)
 
+      if attr == "disableflags" and isinstance(obj, types.Option) and not ENABLE_ISLANDS:
+        val = int(val) | types.DisableBit.ISLAND
+
       setattr(obj, attr, val)
 
 
@@ -2724,11 +2948,13 @@ def create_render_context(
   render_seg: list[bool] | bool | None = None,
   use_textures: bool = True,
   use_shadows: bool = False,
+  use_ambient_lighting: bool = True,
   enabled_geom_groups: list[int] = [0, 1, 2],
   cam_active: list[bool] | None = None,
   flex_render_smooth: bool = True,
   use_precomputed_rays: bool = True,
   render_skybox: bool = False,
+  enable_backface_culling: bool = True,
 ) -> types.RenderContext:
   """Creates a render context on device.
 
@@ -2743,6 +2969,8 @@ def create_render_context(
       If None, uses the MuJoCo model values.
     use_textures: Whether to use textures.
     use_shadows: Whether to use shadows.
+    use_ambient_lighting: Whether to add the renderer's hemispheric ambient
+      lighting term before applying model lights.
     enabled_geom_groups: The geom groups to render.
     cam_active: List of booleans indicating which cameras to include in rendering.
                 If None, all cameras are included.
@@ -2751,6 +2979,10 @@ def create_render_context(
                           When using domain randomization for camera intrinsics, set to False.
     render_skybox: Whether to shade missed rays with the MuJoCo skybox texture.
                    Requires the model to contain a texture with type `mjTEXTURE_SKYBOX`.
+    enable_backface_culling: Drop primitive-ray hits whose normal faces away from
+                             the ray (ray origin inside the geom). Matches MuJoCo's
+                             mesh-ray rule. Default True. Disable for a small
+                             performance gain when no camera is ever inside a geom.
 
   Returns:
     The render context containing rendering fields and output arrays on device.
@@ -2758,13 +2990,7 @@ def create_render_context(
   mjd = mujoco.MjData(mjm)
   mujoco.mj_forward(mjm, mjd)
 
-  constructor = "sah"
-  if check_version("warp>=1.13.0.dev20260325"):
-    # TODO: The cubql constructor and is_cubql_available exist only in
-    # recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
-    _cubql_avail = getattr(wp, "is_cubql_available", None)
-    if callable(_cubql_avail) and _cubql_avail():
-      constructor = "cubql"
+  constructor = "cubql"
 
   # Mesh BVHs – build for all meshes so per-world variants are available
   nmesh = mjm.nmesh
@@ -2937,6 +3163,7 @@ def create_render_context(
     cam_id_map=wp.array(active_cam_indices, dtype=int),
     use_textures=use_textures,
     use_shadows=use_shadows,
+    use_ambient_lighting=use_ambient_lighting,
     background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
     use_precomputed_rays=use_precomputed_rays,
     render_skybox=render_skybox,
@@ -2982,6 +3209,7 @@ def create_render_context(
     render_seg=wp.array(render_seg, dtype=bool),
     znear=znear,
     total_rays=int(total),
+    enable_backface_culling=enable_backface_culling,
   )
 
   bvh.build_scene_bvh(mjm, mjd, rc, nworld)
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
index f7fb3a10..0e92b034 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
@@ -18,6 +18,7 @@ import warp as wp
 from mujoco.mjx.third_party.mujoco_warp._src import types
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
+from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 
@@ -35,11 +36,15 @@ def _tree_edges(
   eq_obj1id: wp.array[int],
   eq_obj2id: wp.array[int],
   eq_objtype: wp.array[int],
+  is_sparse: bool,
   # Data in:
   nefc_in: wp.array[int],
   contact_geom_in: wp.array[wp.vec2i],
   efc_type_in: wp.array2d[int],
   efc_id_in: wp.array2d[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
   efc_J_in: wp.array3d[float],
   njmax_in: int,
   # Out:
@@ -131,26 +136,39 @@ def _tree_edges(
   first_tree = int(-1)
   has_cross_edge = int(0)
 
-  for dof in range(nv):
-    # TODO(team): sparse efc_J
-    # TODO(team): tree dof skip
-    J_val = efc_J_in[worldid, efcid, dof]
-    if J_val != 0.0:
-      tree = dof_treeid[dof]
-      if tree < 0:
+  count = nv
+  rowadr = 0
+  if is_sparse:
+    count = efc_J_rownnz_in[worldid, efcid]
+    rowadr = efc_J_rowadr_in[worldid, efcid]
+
+  for i in range(count):
+    dof = i
+    if is_sparse:
+      sparseid = rowadr + i
+      dof = efc_J_colind_in[worldid, 0, sparseid]
+    else:
+      J_val = efc_J_in[worldid, efcid, dof]
+      if J_val == 0.0:
         continue
-      if first_tree == -1:
-        first_tree = tree
-      elif tree != first_tree:
-        t1 = wp.min(first_tree, tree)
-        t2 = wp.max(first_tree, tree)
-        wp.atomic_max(tree_tree, worldid, t1, t2, 1)
-        has_cross_edge = 1
+
+    tree = dof_treeid[dof]
+    if tree < 0:
+      continue
+    if first_tree == -1:
+      first_tree = tree
+    elif tree != first_tree:
+      t1 = wp.min(first_tree, tree)
+      t2 = wp.max(first_tree, tree)
+      wp.atomic_max(tree_tree, worldid, t1, t2, 1)
+      wp.atomic_max(tree_tree, worldid, t2, t1, 1)
+      has_cross_edge = 1
 
   if first_tree >= 0 and has_cross_edge == 0:
     wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
 
 
+@event_scope
 def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
   """Compute tree-tree adjacency matrix."""
   tree_tree.zero_()
@@ -168,10 +186,14 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
       m.eq_obj1id,
       m.eq_obj2id,
       m.eq_objtype,
+      m.is_sparse,
       d.nefc,
       d.contact.geom,
       d.efc.type,
       d.efc.id,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
       d.efc.J,
       d.njmax,
     ],
@@ -242,6 +264,19 @@ def _flood_fill(
   nisland_out[worldid] = nisland
 
 
+@event_scope
+def flood_fill(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
+  d.tree_island.fill_(-1)
+  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+
+  wp.launch(
+    _flood_fill,
+    dim=d.nworld,
+    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
+    outputs=[d.nisland, d.tree_island, stack_scratch],
+  )
+
+
 @event_scope
 def island(m: types.Model, d: types.Data):
   """Discover constraint islands."""
@@ -254,12 +289,778 @@ def island(m: types.Model, d: types.Data):
   tree_edges(m, d, tree_tree)
 
   # Step 2: DFS flood fill
-  d.tree_island.fill_(-1)
-  stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
+  flood_fill(m, d, tree_tree)
+
+
+@wp.kernel
+def _island_count_dofs(
+  dof_treeid: wp.array[int],
+  tree_island_in: wp.array2d[int],
+  dof_island_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  island_id = tree_island_in[worldid, dof_treeid[dofid]]
+  dof_island_out[worldid, dofid] = island_id
+  if island_id >= 0:
+    wp.atomic_add(island_nv_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_scan_sizes(
+  nisland_in: wp.array[int],
+  island_idofadr_out: wp.array2d[int],
+  island_nv_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    nidof_out[worldid] = 0
+    return
+
+  # Scan DOFs and Constraints
+  island_idofadr_out[worldid, 0] = 0
+  island_iefcadr_out[worldid, 0] = 0
+  for i in range(1, nisland):
+    island_idofadr_out[worldid, i] = island_idofadr_out[worldid, i - 1] + island_nv_inout[worldid, i - 1]
+    island_iefcadr_out[worldid, i] = island_iefcadr_out[worldid, i - 1] + island_nefc_inout[worldid, i - 1]
+
+  nidof = island_idofadr_out[worldid, nisland - 1] + island_nv_inout[worldid, nisland - 1]
+  nidof_out[worldid] = nidof
+
+  # Reset for recount
+  for i in range(nisland):
+    island_nv_inout[worldid, i] = 0
+    island_nefc_inout[worldid, i] = 0
+
+
+@wp.kernel
+def _island_map_dofs(
+  nv: int,
+  dof_island_in: wp.array2d[int],
+  island_idofadr_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  island_nv_inout: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+  unconstrained_cnt_inout: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  island_id = dof_island_in[worldid, dofid]
+
+  if island_id >= 0:
+    local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
+    idof = island_idofadr_in[worldid, island_id] + local_idx
+    idof_islandid_out[worldid, idof] = island_id
+  else:
+    cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
+    idof = nidof + cnt
+
+  map_dof2idof_out[worldid, dofid] = idof
+  map_idof2dof_out[worldid, idof] = dofid
+
+
+@wp.kernel
+def _island_count_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_tree_in: wp.array2d[int],
+  tree_island_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  efc_island_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_tree = efc_tree_in[worldid, efcid]
+  if efc_tree < 0:
+    efc_island_out[worldid, efcid] = -1
+    return
+  island_id = tree_island_in[worldid, efc_tree]
+  efc_island_out[worldid, efcid] = island_id
+
+  if island_id >= 0:
+    wp.atomic_add(island_nefc_out, worldid, island_id, 1)
+
+    efc_type = efc_type_in[worldid, efcid]
+    if efc_type == ConstraintType.EQUALITY:
+      wp.atomic_add(island_ne_out, worldid, island_id, 1)
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      wp.atomic_add(island_nf_out, worldid, island_id, 1)
+
+
+@wp.kernel
+def _island_map_constraints(
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  efc_island_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  efc_type_in: wp.array2d[int],
+  # Counters (inout):
+  island_ne_mapped_inout: wp.array2d[int],
+  island_nf_mapped_inout: wp.array2d[int],
+  island_nother_mapped_inout: wp.array2d[int],
+  island_nefc_inout: wp.array2d[int],
+  # Out:
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  island_id = efc_island_in[worldid, efcid]
+  if island_id >= 0:
+    efc_type = efc_type_in[worldid, efcid]
+
+    # 1. Determine local index and absolute index ic based on category
+    if efc_type == ConstraintType.EQUALITY:
+      local_idx = wp.atomic_add(island_ne_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + local_idx
+    elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
+      local_idx = wp.atomic_add(island_nf_mapped_inout, worldid, island_id, 1)
+      ic = island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + local_idx
+    else:
+      local_idx = wp.atomic_add(island_nother_mapped_inout, worldid, island_id, 1)
+      ic = (
+        island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + island_nf_in[worldid, island_id] + local_idx
+      )
+
+    # 2. Increment overall island_nefc counter to reconstruct d.island_nefc
+    wp.atomic_add(island_nefc_inout, worldid, island_id, 1)
+
+    # 3. Store mappings
+    map_efc2iefc_out[worldid, efcid] = ic
+    map_iefc2efc_out[worldid, ic] = efcid
+    iefc_islandid_out[worldid, ic] = island_id
+
+
+@wp.kernel
+def _island_scan_sparse_rows(
+  nisland_in: wp.array[int],
+  efc_J_rownnz_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  island_iefcadr_in: wp.array2d[int],
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_J_rownnz_out: wp.array2d[int],
+  iefc_J_rowadr_out: wp.array2d[int],
+):
+  worldid = wp.tid()
+
+  nisland = nisland_in[worldid]
+  if nisland == 0:
+    return
+
+  total_gathered_efc = island_iefcadr_in[worldid, nisland - 1] + island_nefc_in[worldid, nisland - 1]
+
+  running_rowadr = int(0)
+  for ic in range(total_gathered_efc):
+    c = map_iefc2efc_in[worldid, ic]
+    rownnz = efc_J_rownnz_in[worldid, c]
+    iefc_J_rowadr_out[worldid, ic] = running_rowadr
+    iefc_J_rownnz_out[worldid, ic] = rownnz
+    running_rowadr += rownnz
+
+
+@wp.kernel
+def _compute_efc_tree(
+  # Model:
+  nv: int,
+  body_treeid: wp.array[int],
+  jnt_dofadr: wp.array[int],
+  dof_treeid: wp.array[int],
+  geom_bodyid: wp.array[int],
+  site_bodyid: wp.array[int],
+  eq_type: wp.array[int],
+  eq_obj1id: wp.array[int],
+  eq_obj2id: wp.array[int],
+  eq_objtype: wp.array[int],
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  contact_geom_in: wp.array[wp.vec2i],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_J_in: wp.array3d[float],
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  njmax_in: int,
+  # Out:
+  efc_tree_out: wp.array2d[int],
+):
+  """Compute the first non-negative tree for each constraint."""
+  worldid, efcid = wp.tid()
+
+  if efcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  efc_type = efc_type_in[worldid, efcid]
+  efc_id = efc_id_in[worldid, efcid]
+
+  tree = int(-1)
+  use_generic = int(0)
+
+  # equality (connect/weld)
+  if efc_type == ConstraintType.EQUALITY:
+    eq_t = eq_type[efc_id]
+
+    if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
+      b1 = eq_obj1id[efc_id]
+      b2 = eq_obj2id[efc_id]
+
+      if eq_objtype[efc_id] == ObjType.SITE:
+        b1 = site_bodyid[b1]
+        b2 = site_bodyid[b2]
+
+      t1 = body_treeid[b1]
+      t2 = body_treeid[b2]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # JOINT, TENDON, FLEX: generic scan
+      use_generic = 1
+
+  # joint friction
+  elif efc_type == ConstraintType.FRICTION_DOF:
+    tree = dof_treeid[efc_id]
+
+  # joint limit
+  elif efc_type == ConstraintType.LIMIT_JOINT:
+    tree = dof_treeid[jnt_dofadr[efc_id]]
+
+  # contact
+  elif (
+    efc_type == ConstraintType.CONTACT_FRICTIONLESS
+    or efc_type == ConstraintType.CONTACT_PYRAMIDAL
+    or efc_type == ConstraintType.CONTACT_ELLIPTIC
+  ):
+    geom_pair = contact_geom_in[efc_id]
+    g1 = geom_pair[0]
+    g2 = geom_pair[1]
+
+    if g1 >= 0 and g2 >= 0:
+      t1 = body_treeid[geom_bodyid[g1]]
+      t2 = body_treeid[geom_bodyid[g2]]
+      if t1 >= 0:
+        tree = t1
+      else:
+        tree = t2
+    else:
+      # flex contacts: generic scan
+      use_generic = 1
+
+  else:
+    # generic: scan Jacobian row
+    use_generic = 1
+
+  if use_generic:
+    count = nv
+    rowadr = 0
+    if is_sparse:
+      count = efc_J_rownnz_in[worldid, efcid]
+      rowadr = efc_J_rowadr_in[worldid, efcid]
+
+    for i in range(count):
+      dof = i
+      if is_sparse:
+        sparseid = rowadr + i
+        dof = efc_J_colind_in[worldid, 0, sparseid]
+      else:
+        J_val = efc_J_in[worldid, efcid, dof]
+        if J_val == 0.0:
+          continue
+
+      t = dof_treeid[dof]
+      if t >= 0:
+        tree = t
+        break
+
+  efc_tree_out[worldid, efcid] = tree
+
+
+@wp.kernel
+def _gather_efc_and_jacobian(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  efc_D_in: wp.array2d[float],
+  efc_type_in: wp.array2d[int],
+  efc_id_in: wp.array2d[int],
+  efc_frictionloss_in: wp.array2d[float],
+  efc_aref_in: wp.array2d[float],
+  efc_J_in: wp.array3d[float],
+  # Sparse Jacobian arrays:
+  efc_J_rownnz_in: wp.array2d[int],
+  efc_J_rowadr_in: wp.array2d[int],
+  efc_J_colind_in: wp.array3d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  # In:
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  map_dof2idof_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  # Out:
+  iefc_D_out: wp.array2d[float],
+  iefc_type_out: wp.array2d[int],
+  iefc_id_out: wp.array2d[int],
+  iefc_frictionloss_out: wp.array2d[float],
+  iefc_aref_out: wp.array2d[float],
+  iefc_J_out: wp.array3d[float],
+  iefc_J_colind_out: wp.array3d[int],
+):
+  """Gather constraint arrays and Jacobian into island-local dense order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+
+  # Scalar constraint fields
+  iefc_D_out[worldid, iefcid] = efc_D_in[worldid, c]
+  iefc_type_out[worldid, iefcid] = efc_type_in[worldid, c]
+  iefc_id_out[worldid, iefcid] = efc_id_in[worldid, c]
+  iefc_frictionloss_out[worldid, iefcid] = efc_frictionloss_in[worldid, c]
+  iefc_aref_out[worldid, iefcid] = efc_aref_in[worldid, c]
+
+  # Jacobian: convert to island-local dense format
+  nid = nidof_in[worldid]
+
+  if is_sparse:
+    rownnz = efc_J_rownnz_in[worldid, c]
+    rowadr_in = efc_J_rowadr_in[worldid, c]
+    rowadr_out = iefc_J_rowadr_in[worldid, iefcid]
+    for i in range(rownnz):
+      sparseid_in = rowadr_in + i
+      sparseid_out = rowadr_out + i
+      dof = efc_J_colind_in[worldid, 0, sparseid_in]
+      idof = map_dof2idof_in[worldid, dof]
+      # Store in sparse iefc_J_out and iefc_J_colind_out
+      iefc_J_out[worldid, 0, sparseid_out] = efc_J_in[worldid, 0, sparseid_in]
+      iefc_J_colind_out[worldid, 0, sparseid_out] = idof
+  else:
+    # Dense path: reorder rows by iefc, columns by idof
+    for idof in range(nid):
+      dof = map_idof2dof_in[worldid, idof]
+      iefc_J_out[worldid, iefcid, idof] = efc_J_in[worldid, c, dof]
+
+
+@wp.kernel
+def _gather_dof_arrays(
+  # Data in:
+  qacc_in: wp.array2d[float],
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  # In:
+  nidof_in: wp.array[int],
+  map_idof2dof_in: wp.array2d[int],
+  # Out:
+  iacc_out: wp.array2d[float],
+  iacc_smooth_out: wp.array2d[float],
+  ifrc_smooth_out: wp.array2d[float],
+):
+  """Gather DOF arrays into island-local order."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  dof = map_idof2dof_in[worldid, idofid]
+  iacc_out[worldid, idofid] = qacc_in[worldid, dof]
+  iacc_smooth_out[worldid, idofid] = qacc_smooth_in[worldid, dof]
+  ifrc_smooth_out[worldid, idofid] = qfrc_smooth_in[worldid, dof]
+
+
+@wp.kernel
+def _scatter_dof_arrays(
+  # In:
+  qacc_smooth_in: wp.array2d[float],
+  qfrc_smooth_in: wp.array2d[float],
+  dof_island_in: wp.array2d[int],
+  iacc_in: wp.array2d[float],
+  ifrc_constraint_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  map_dof2idof_in: wp.array2d[int],
+  scatter_Ma: bool,
+  # Out:
+  qacc_out: wp.array2d[float],
+  qfrc_constraint_out: wp.array2d[float],
+  Ma_out: wp.array2d[float],
+):
+  """Scatter island results to global arrays, copy qacc_smooth for non-island DOFs."""
+  worldid, dofid = wp.tid()
+
+  if dof_island_in[worldid, dofid] < 0:
+    qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
+    qfrc_constraint_out[worldid, dofid] = 0.0
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid]
+  else:
+    idof = map_dof2idof_in[worldid, dofid]
+    qacc_out[worldid, dofid] = iacc_in[worldid, idof]
+    qfrc_constraint_out[worldid, dofid] = ifrc_constraint_in[worldid, idof]
+    if scatter_Ma:
+      Ma_out[worldid, dofid] = iMa_in[worldid, idof]
+
+
+@wp.kernel
+def _scatter_efc_arrays(
+  # In:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  map_iefc2efc_in: wp.array2d[int],
+  iefc_force_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  # Out:
+  efc_force_out: wp.array2d[float],
+  efc_state_out: wp.array2d[int],
+):
+  """Scatter island-local constraint results back to global order."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  c = map_iefc2efc_in[worldid, iefcid]
+  efc_force_out[worldid, c] = iefc_force_in[worldid, iefcid]
+  efc_state_out[worldid, c] = iefc_state_in[worldid, iefcid]
+
+
+@wp.kernel
+def _init_island_arrays(
+  island_idofadr_out: wp.array2d[int],
+  island_nv_out: wp.array2d[int],
+  island_nefc_out: wp.array2d[int],
+  island_ne_out: wp.array2d[int],
+  island_nf_out: wp.array2d[int],
+  island_iefcadr_out: wp.array2d[int],
+  nidof_out: wp.array[int],
+):
+  worldid, islandid = wp.tid()
+  island_nv_out[worldid, islandid] = 0
+  island_nefc_out[worldid, islandid] = 0
+  island_ne_out[worldid, islandid] = 0
+  island_nf_out[worldid, islandid] = 0
+  island_idofadr_out[worldid, islandid] = 0
+  island_iefcadr_out[worldid, islandid] = 0
+  if islandid == 0:
+    nidof_out[worldid] = 0
+
+
+@wp.kernel
+def _init_dof_arrays(
+  dof_island_out: wp.array2d[int],
+  map_dof2idof_out: wp.array2d[int],
+  map_idof2dof_out: wp.array2d[int],
+  idof_islandid_out: wp.array2d[int],
+):
+  worldid, dofid = wp.tid()
+  dof_island_out[worldid, dofid] = -1
+  map_dof2idof_out[worldid, dofid] = 0
+  map_idof2dof_out[worldid, dofid] = 0
+  idof_islandid_out[worldid, dofid] = -1
+
+
+@wp.kernel
+def _init_efc_arrays(
+  efc_island_out: wp.array2d[int],
+  map_efc2iefc_out: wp.array2d[int],
+  map_iefc2efc_out: wp.array2d[int],
+  iefc_islandid_out: wp.array2d[int],
+  efc_tree_out: wp.array2d[int],
+):
+  worldid, efcid = wp.tid()
+  efc_island_out[worldid, efcid] = -1
+  map_efc2iefc_out[worldid, efcid] = 0
+  map_iefc2efc_out[worldid, efcid] = 0
+  iefc_islandid_out[worldid, efcid] = -1
+  efc_tree_out[worldid, efcid] = -1
+
+
+@event_scope
+def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Compute DOF/constraint island mappings after island discovery.
+
+  Populates island solver context arrays via ctx: nv, nefc, ne, nf,
+  iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
+  map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
+  and d.island_dofadr.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext.
+  """
+  # Ensure dof_islandid / efc_islandid are allocated at the right shape
+  if d.dof_islandid.shape[1] != m.nv:
+    d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
+  if d.efc_islandid.shape[1] != d.njmax:
+    d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
+
+  # Ensure island-local DOF arrays are allocated at the right shape
+  if d.iqacc.shape[1] != m.nv:
+    nw = d.nworld
+    d.iqacc = wp.empty((nw, m.nv), dtype=float)
+    d.iqacc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_smooth = wp.empty((nw, m.nv), dtype=float)
+    d.iqfrc_constraint = wp.empty((nw, m.nv), dtype=float)
+  wp.launch(
+    _init_island_arrays,
+    dim=(d.nworld, m.ntree),
+    inputs=[],
+    outputs=[
+      d.island_dofadr,
+      d.island_nv,
+      d.island_nefc,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.nidof,
+    ],
+  )
+  wp.launch(
+    _init_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[],
+    outputs=[d.dof_island, d.map_dof2idof, d.map_idof2dof, d.dof_islandid],
+  )
+  efc_tree = wp.empty((d.nworld, d.njmax), dtype=int)
+  wp.launch(
+    _init_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[],
+    outputs=[d.efc.island, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid, efc_tree],
+  )
 
   wp.launch(
-    _flood_fill,
-    dim=d.nworld,
-    inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
-    outputs=[d.nisland, d.tree_island, stack_scratch],
+    _compute_efc_tree,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.nv,
+      m.body_treeid,
+      m.jnt_dofadr,
+      m.dof_treeid,
+      m.geom_bodyid,
+      m.site_bodyid,
+      m.eq_type,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.is_sparse,
+      d.nefc,
+      d.contact.geom,
+      d.efc.type,
+      d.efc.id,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.njmax,
+    ],
+    outputs=[efc_tree],
+  )
+
+  # 1. Count DOFs per island
+  wp.launch(
+    _island_count_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.dof_treeid, d.tree_island],
+    outputs=[d.dof_island, d.island_nv],
+  )
+
+  # 2. Count Constraints per island
+  wp.launch(
+    _island_count_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[d.nefc, d.njmax, efc_tree, d.tree_island, d.efc.type],
+    outputs=[d.efc.island, d.island_nefc, d.island_ne, d.island_nf],
+  )
+
+  # 3. Scan sizes and reset counters for mapping
+  wp.launch(
+    _island_scan_sizes,
+    dim=d.nworld,
+    inputs=[d.nisland],
+    outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
+  )
+
+  # 4. Map DOFs
+  unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
+  wp.launch(
+    _island_map_dofs,
+    dim=(d.nworld, m.nv),
+    inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
+    outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
+  )
+
+  # 5. Map Constraints
+  ne_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nf_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+  nother_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
+
+  wp.launch(
+    _island_map_constraints,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.efc.island,
+      d.island_efcadr,
+      d.island_ne,
+      d.island_nf,
+      d.efc.type,
+      ne_mapped,
+      nf_mapped,
+      nother_mapped,
+    ],
+    outputs=[d.island_nefc, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid],
+  )
+
+  # 6. Scan Sparse Rows (if sparse)
+  if m.is_sparse:
+    wp.launch(
+      _island_scan_sparse_rows,
+      dim=d.nworld,
+      inputs=[d.nisland, d.efc.J_rownnz, d.island_nefc, d.island_efcadr, d.map_iefc2efc],
+      outputs=[d.efc.iJ_rownnz, d.efc.iJ_rowadr],
+    )
+
+
+@event_scope
+def gather_island_inputs(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Gather constraint and DOF arrays into island-local order.
+
+  Populates d.iefc (D, type, id, frictionloss, aref, J, J_colind) and
+  d.iqacc, d.iqacc_smooth, d.iqfrc_smooth.
+
+  Must be called after compute_island_mapping() and before per-island solving.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext whose arrays are populated.
+  """
+  # Gather constraint arrays and dense Jacobian (fused)
+  wp.launch(
+    _gather_efc_and_jacobian,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.efc.D,
+      d.efc.type,
+      d.efc.id,
+      d.efc.frictionloss,
+      d.efc.aref,
+      d.efc.J,
+      d.efc.J_rownnz,
+      d.efc.J_rowadr,
+      d.efc.J_colind,
+      d.efc.iJ_rowadr,
+      d.njmax,
+      d.map_iefc2efc,
+      d.map_idof2dof,
+      d.map_dof2idof,
+      d.nidof,
+    ],
+    outputs=[
+      d.efc.iD,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.ifrictionloss,
+      d.efc.iaref,
+      d.efc.iJ,
+      d.efc.iJ_colind,
+    ],
+  )
+
+  # Gather DOF arrays
+  wp.launch(
+    _gather_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc,
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.nidof,
+      d.map_idof2dof,
+    ],
+    outputs=[
+      d.iqacc,
+      d.iqacc_smooth,
+      d.iqfrc_smooth,
+    ],
+  )
+
+
+@event_scope
+def scatter_island_results(m: types.Model, d: types.Data, ctx: IslandSolverContext, scatter_Ma: bool):
+  """Scatter island-local solver results back to global arrays.
+
+  Reads ctx qacc, qfrc_constraint, Ma, d.efc iforce, istate and
+  writes them back to d.qacc, d.qfrc_constraint, d.efc.Ma, d.efc.force,
+  d.efc.state. Unconstrained DOFs receive qacc_smooth and zero qfrc.
+
+  Args:
+    m: Model.
+    d: Data.
+    ctx: IslandSolverContext which contains results.
+    scatter_Ma: Whether to scatter Ma for Euler/implicit integrators.
+  """
+  # Scatter DOF results (and optionally Ma for Euler/implicit integrators)
+  wp.launch(
+    _scatter_dof_arrays,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.qacc_smooth,
+      d.qfrc_smooth,
+      d.dof_island,
+      d.iqacc,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.map_dof2idof,
+      scatter_Ma,
+    ],
+    outputs=[
+      d.qacc,
+      d.qfrc_constraint,
+      d.efc.Ma,
+    ],
+  )
+
+  # Scatter constraint results (force and state from island-local arrays)
+  wp.launch(
+    _scatter_efc_arrays,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      d.map_iefc2efc,
+      d.efc.iforce,
+      d.efc.istate,
+    ],
+    outputs=[
+      d.efc.force,
+      d.efc.state,
+    ],
   )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
index a320644a..d9b20861 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
@@ -707,10 +707,14 @@ def ray_mesh_with_bvh(
   pnt: wp.vec3,
   vec: wp.vec3,
   max_t: float,
+  cull_backfaces: bool,
 ) -> Tuple[float, wp.vec3, float, float, int, int]:
   """Returns intersection information for ray mesh intersections.
 
   Requires wp.Mesh be constructed and their ids to be passed.
+
+  When ``cull_backfaces`` is True, the function rejects exit-face hits in the
+  local-space frame. This matches MuJoCo OpenGL rendering's backface culling rule.
   """
   t = float(-1.0)
   u = float(0.0)
@@ -722,7 +726,12 @@ def ray_mesh_with_bvh(
   lpnt, lvec = _ray_map(pos, mat, pnt, vec)
   hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
 
-  if hit and wp.dot(lvec, n) < 0.0:  # Backface culling in local space
+  if not hit:
+    return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
+
+  # Front-face hit, or back-face hit when cull is disabled: rotate the
+  # local-space normal into world space and return the hit.
+  if (not cull_backfaces) or wp.dot(lvec, n) < 0.0:
     normal = mat @ n
     normal = wp.normalize(normal)
     return t, normal, u, v, f, mesh_geom_id
@@ -1055,6 +1064,8 @@ def _ray_geom_mesh_bvh(
 
     if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
       bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
+      # Public ray API (mjw.ray / mjw.rays) preserves MuJoCo's mj_ray cull rule:
+      # rangefinder sensors and user-facing ray casts always cull back-faces.
       t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
         bvh_ids,
         geom_dataid[worldid % geom_dataid.shape[0], geomid],
@@ -1063,6 +1074,7 @@ def _ray_geom_mesh_bvh(
         pnt,
         vec,
         min_dist,
+        True,
       )
       if t >= 0.0 and t < min_dist:
         return t, n
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
index ba3b4cfe..7afb140b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
@@ -180,6 +180,7 @@ def cast_ray(
   flex_group_root: wp.array2d[int],
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
+  cull_backfaces: bool,
 ) -> Tuple[int, float, wp.vec3, float, float, int, int]:
   dist = float(MJ_MAXVAL)
   normal = wp.vec3(0.0, 0.0, 0.0)
@@ -232,6 +233,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -281,6 +283,7 @@ def cast_ray(
         ray_origin_world,
         ray_dir_world,
         dist,
+        cull_backfaces,
       )
     if gtype == GeomType.FLEX:
       hit_geom_id = -2
@@ -308,6 +311,11 @@ def cast_ray(
         if d >= 0.0:
           hit_mesh_id = flexid
 
+    # Backface cull: drop exit-face hits when the ray origin is inside the geom,
+    # matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < dist:
       dist = d
       normal = n
@@ -349,6 +357,7 @@ def cast_ray_first_hit(
   ray_origin_world: wp.vec3,
   ray_dir_world: wp.vec3,
   max_dist: float,
+  cull_backfaces: bool,
 ) -> bool:
   """A simpler version of casting rays that only checks for the first hit."""
   query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
@@ -387,6 +396,7 @@ def cast_ray_first_hit(
         ray_origin_world,
         ray_dir_world,
         max_dist,
+        cull_backfaces,
       )
     if gtype == GeomType.SPHERE:
       d, n = ray_sphere(
@@ -467,6 +477,11 @@ def cast_ray_first_hit(
         )
         d = 0.0 if hit else -1.0
 
+    # Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
+    # hits and skips branches with a zero-vector normal (mesh/flex anyhit).
+    if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
+      d = -1.0
+
     if d >= 0.0 and d < max_dist:
       return True
 
@@ -507,6 +522,7 @@ def compute_lighting(
   lightdir: wp.vec3,
   normal: wp.vec3,
   hitpoint: wp.vec3,
+  cull_backfaces: bool,
 ) -> float:
   light_contribution = float(0.0)
 
@@ -571,6 +587,7 @@ def compute_lighting(
       shadow_origin,
       L,
       max_t,
+      cull_backfaces,
     )
 
     if shadow_hit:
@@ -725,6 +742,7 @@ def render(m: Model, d: Data, rc: RenderContext):
       flex_group_root,
       ray_origin_world,
       ray_dir_world,
+      wp.static(rc.enable_backface_culling),
     )
 
     if render_seg[cam_idx] and geom_id != -1:
@@ -798,12 +816,14 @@ def render(m: Model, d: Data, rc: RenderContext):
             )
             base_color = wp.cw_mul(base_color, tex_color)
 
-    len_n = wp.length(normal)
-    n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
-    n = wp.normalize(n)
-    hemispheric = 0.5 * (n[2] + 1.0)
-    ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
-    result = 0.5 * wp.cw_mul(base_color, ambient_color)
+    result = wp.vec3(0.0, 0.0, 0.0)
+    if wp.static(rc.use_ambient_lighting):
+      len_n = wp.length(normal)
+      n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
+      n = wp.normalize(n)
+      hemispheric = 0.5 * (n[2] + 1.0)
+      ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
+      result = 0.5 * wp.cw_mul(base_color, ambient_color)
 
     # Apply lighting and shadows
     for l in range(wp.static(m.nlight)):
@@ -837,6 +857,7 @@ def render(m: Model, d: Data, rc: RenderContext):
         light_xdir_in[worldid, l],
         normal,
         hit_point,
+        wp.static(rc.enable_backface_culling),
       )
       result = result + base_color * light_contribution
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
index 472fa617..619e55cf 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
@@ -17,6 +17,7 @@ from typing import Any, Tuple
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import history
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import ray
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
@@ -26,6 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
 from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
+from mujoco.mjx.third_party.mujoco_warp._src.types import TACTILE_DEPTH_SEMANTICS
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
 from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -917,6 +919,10 @@ def sensor_pos(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for position sensors
+  history.apply_sensor_delay(m, d, m.sensor_pos_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitpos_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.POS)
 
@@ -1459,6 +1465,10 @@ def sensor_vel(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for velocity sensors
+  history.apply_sensor_delay(m, d, m.sensor_vel_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitvel_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.VEL)
 
@@ -2253,13 +2263,23 @@ def _sensor_tactile(
     vel_rel = vel_sensor - vel_other
 
     forceT = wp.vec3(0.0, 0.0, 0.0)
-    forceT[0] = -depth
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
+      forceT[0] = -depth
+    else:
+      kMaxDepth = 0.05
+      pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
+      force = wp.mul(normal, pressure)
+      forceT[0] = wp.dot(force, normal)
+
     if has_frame:
       forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
       forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
 
     dim = sensor_dim[sensor_id] // 3
-    wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    if wp.static(TACTILE_DEPTH_SEMANTICS):
+      wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
+    else:
+      wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
     wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
 
@@ -2705,6 +2725,10 @@ def sensor_acc(m: Model, d: Data):
     ],
   )
 
+  # apply sensor delay/interval for acceleration sensors
+  history.apply_sensor_delay(m, d, m.sensor_acc_adr)
+  history.apply_sensor_delay(m, d, m.sensor_limitfrc_adr)
+
   if m.callback.sensor:
     m.callback.sensor(m, d, Stage.ACC)
 
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
index bda36c7c..4f35eab0 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
@@ -823,37 +823,38 @@ def _crb_accumulate(
 
 
 @wp.kernel
-def _qM_sparse(
+def _M_sparse(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   dof_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   # Data in:
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
-  madr_ij = dof_Madr[dofid]  # dof_Madr is not batched
   bodyid = dof_bodyid[dofid]
+  madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # init M(i,i) with armature inertia
-  qM_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
+  M_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
 
   # precompute buf = crb_body_i * cdof_i
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
 
   # sparse backward pass over ancestors
   while dofid >= 0:
-    qM_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
-    madr_ij += 1
+    M_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
+    madr_ij -= 1
     dofid = dof_parentid[dofid]
 
 
 @wp.kernel
-def _qM_dense(
+def _M_dense(
   # Model:
   dof_bodyid: wp.array[int],
   dof_parentid: wp.array[int],
@@ -862,7 +863,7 @@ def _qM_dense(
   cdof_in: wp.array2d[wp.spatial_vector],
   crb_in: wp.array2d[vec10],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, dofid = wp.tid()
   bodyid = dof_bodyid[dofid]
@@ -873,15 +874,15 @@ def _qM_dense(
   buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
   M += wp.dot(cdof_in[worldid, dofid], buf)
 
-  qM_out[worldid, dofid, dofid] = M
+  M_out[worldid, dofid, dofid] = M
 
   # sparse backward pass over ancestors
   dofidi = dofid
   dofid = dof_parentid[dofid]
   while dofid >= 0:
-    qMij = wp.dot(cdof_in[worldid, dofid], buf)
-    qM_out[worldid, dofidi, dofid] += qMij
-    qM_out[worldid, dofid, dofidi] += qMij
+    Mij = wp.dot(cdof_in[worldid, dofid], buf)
+    M_out[worldid, dofidi, dofid] += Mij
+    M_out[worldid, dofid, dofidi] += Mij
     dofid = dof_parentid[dofid]
 
 
@@ -898,17 +899,17 @@ def crb(m: Model, d: Data):
     body_tree = m.body_tree[i]
     wp.launch(_crb_accumulate, dim=(d.nworld, body_tree.size), inputs=[m.body_parentid, d.crb, body_tree], outputs=[d.crb])
 
-  d.qM.zero_()
+  d.M.zero_()
   if m.is_sparse:
     wp.launch(
-      _qM_sparse,
+      _M_sparse,
       dim=(d.nworld, m.nv),
-      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_Madr, m.dof_armature, d.cdof, d.crb],
-      outputs=[d.qM],
+      inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, m.M_rownnz, m.M_rowadr, d.cdof, d.crb],
+      outputs=[d.M],
     )
   else:
     wp.launch(
-      _qM_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.qM]
+      _M_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.M]
     )
 
 
@@ -916,16 +917,17 @@ def crb(m: Model, d: Data):
 def _tendon_armature(
   # Model:
   dof_parentid: wp.array[int],
-  dof_Madr: wp.array[int],
   ten_J_rownnz: wp.array[int],
   ten_J_rowadr: wp.array[int],
   ten_J_colind: wp.array[int],
   tendon_armature: wp.array2d[float],
+  M_rownnz: wp.array[int],
+  M_rowadr: wp.array[int],
   is_sparse: bool,
   # Data in:
   ten_J_in: wp.array2d[float],
   # Data out:
-  qM_out: wp.array3d[float],
+  M_out: wp.array3d[float],
 ):
   worldid, tenid, dofid = wp.tid()
 
@@ -947,7 +949,7 @@ def _tendon_armature(
     return
 
   if is_sparse:
-    madr_ij = dof_Madr[dofid]
+    madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
 
   # sparse backward pass over ancestors
   dofidi = dofid
@@ -967,52 +969,40 @@ def _tendon_armature(
       else:
         ten_Jj = float(0.0)
 
-    qMij = armature * ten_Jj * ten_Ji
+    Mij = armature * ten_Jj * ten_Ji
 
     if is_sparse:
-      wp.atomic_add(qM_out[worldid, 0], madr_ij, qMij)
-      madr_ij += 1
+      wp.atomic_add(M_out[worldid, 0], madr_ij, Mij)
+      madr_ij -= 1
     else:
-      wp.atomic_add(qM_out[worldid, dofidi], dofid, qMij)
+      wp.atomic_add(M_out[worldid, dofidi], dofid, Mij)
       if dofidi != dofid:
-        wp.atomic_add(qM_out[worldid, dofid], dofidi, qMij)
+        wp.atomic_add(M_out[worldid, dofid], dofidi, Mij)
 
     dofid = dof_parentid[dofid]
 
 
 @event_scope
 def tendon_armature(m: Model, d: Data):
-  """Add tendon armature to qM."""
+  """Add tendon armature to M."""
   wp.launch(
     _tendon_armature,
     dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
     inputs=[
       m.dof_parentid,
-      m.dof_Madr,
       m.ten_J_rownnz,
       m.ten_J_rowadr,
       m.ten_J_colind,
       m.tendon_armature,
+      m.M_rownnz,
+      m.M_rowadr,
       m.is_sparse,
       d.ten_J,
     ],
-    outputs=[d.qM],
+    outputs=[d.M],
   )
 
 
-@wp.kernel
-def _copy_CSR(
-  # Model:
-  mapM2M: wp.array[int],
-  # In:
-  M_in: wp.array3d[float],
-  # Out:
-  L_out: wp.array3d[float],
-):
-  worldid, ind = wp.tid()
-  L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]]
-
-
 @wp.kernel
 def _qLD_acc(
   # Model:
@@ -1055,7 +1045,7 @@ def _qLDiag_div(
 
 def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]):
   """Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
-  wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L])
+  wp.copy(L, M)
 
   for i in reversed(range(len(m.qLD_updates))):
     qLD_updates = m.qLD_updates[i]
@@ -1071,7 +1061,7 @@ def _tile_cholesky_factorize(tile: TileSet):
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     adr: wp.array[int],
     # Out:
@@ -1081,8 +1071,8 @@ def _tile_cholesky_factorize(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    L_tile = wp.tile_cholesky(M_tile)
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L_out[worldid], L_tile, offset=(dofid, dofid))
 
   return cholesky_factorize
@@ -1090,7 +1080,7 @@ def _tile_cholesky_factorize(tile: TileSet):
 
 def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
   """Dense Cholesky factorization of inertia-like matrix M, assumed spd."""
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize(tile),
       dim=(d.nworld, tile.adr.size),
@@ -1104,9 +1094,9 @@ def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
 def factor_m(m: Model, d: Data):
   """Factorization of inertia-like matrix M, assumed spd."""
   if m.is_sparse:
-    _factor_i_sparse(m, d, d.qM, d.qLD, d.qLDiagInv)
+    _factor_i_sparse(m, d, d.M, d.qLD, d.qLDiagInv)
   else:
-    _factor_i_dense(m, d, d.qM, d.qLD)
+    _factor_i_dense(m, d, d.M, d.qLD)
 
 
 @wp.kernel
@@ -2842,15 +2832,15 @@ def _tile_cholesky_solve(tile: TileSet):
     dofid = adr[nodeid]
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
     L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_solve
 
 
 def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]):
-  """Computes dense backsubstitution: x = inv(L'*L)*y."""
-  for tile in m.qM_tiles:
+  """Computes dense backsubstitution: x = inv(U.T @ U) * y."""
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2868,9 +2858,9 @@ def solve_LD(
   x: wp.array2d[float],
   y: wp.array2d[float],
 ):
-  """Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y.
+  """Computes backsubstitution for the inertia factorization.
 
-  L and D are the factors from the Cholesky factorization of the inertia matrix.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function dispatches to either a sparse or dense solver depending on Model options.
 
@@ -2907,8 +2897,9 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
 
   @wp.kernel(module="unique", enable_backward=False)
   def cholesky_factorize_solve(
+    # Data in:
+    M_in: wp.array3d[float],
     # In:
-    M: wp.array3d[float],
     y: wp.array2d[float],
     adr: wp.array[int],
     # Out:
@@ -2919,12 +2910,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
     TILE_SIZE = wp.static(tile.size)
 
     dofid = adr[nodeid]
-    M_tile = wp.tile_load(M[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
     y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
 
-    L_tile = wp.tile_cholesky(M_tile)
+    L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
     wp.tile_store(L[worldid], L_tile, offset=(dofid, dofid))
-    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
     wp.tile_store(x[worldid], x_slice, offset=(dofid,))
 
   return cholesky_factorize_solve
@@ -2938,7 +2929,7 @@ def _factor_solve_i_dense(
   y: wp.array2d[float],
   L: wp.array3d[float],
 ):
-  for tile in m.qM_tiles:
+  for tile in m.M_tiles:
     wp.launch_tiled(
       _tile_cholesky_factorize_solve(tile),
       dim=(d.nworld, tile.adr.size),
@@ -2949,9 +2940,9 @@ def _factor_solve_i_dense(
 
 
 def factor_solve_i(m, d, M, L, D, x, y):
-  """Factorizes and solves the linear system: x = inv(L'*D*L) * y or x = inv(L'*L) * y.
+  """Factorizes and solves the inertia-like linear system.
 
-  M is an inertia-like matrix and L, D are its Cholesky-like factors.
+  Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
 
   This function first factorizes the matrix M (sparse or dense depending on model options),
   then solves the system for x given right-hand side y.
@@ -2960,7 +2951,7 @@ def factor_solve_i(m, d, M, L, D, x, y):
     m: The model containing factorization and sparsity information.
     d: The data object containing workspace and factorization results.
     M: The inertia-like matrix to factorize.
-    L: Output lower-triangular factor from the factorization (sparse or dense).
+    L: Output sparse factor or dense upper Cholesky factor.
     D: Output diagonal factor from the factorization (only used for sparse).
     x: Output array for the solution.
     y: Input right-hand side array.
@@ -2972,6 +2963,132 @@ def factor_solve_i(m, d, M, L, D, x, y):
     _factor_solve_i_dense(m, d, M, x, y, L)
 
 
+@cache_kernel
+def _factor_solve_lu_sparse_fused(nv: int):
+  """Fused sparse LU factorization and solve in a single kernel."""
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Model:
+    D_rownnz: wp.array[int],
+    D_rowadr: wp.array[int],
+    D_diag: wp.array[int],
+    D_colind: wp.array[int],
+    # In:
+    qfrc: wp.array2d[float],
+    # Data out:
+    qacc_out: wp.array2d[float],
+    qLU_out: wp.array3d[float],
+  ):
+    worldid = wp.tid()
+    NV = wp.static(nv)
+
+    # Phase 1: LU factorization (in-place on qLU_out)
+    for i in range(NV):
+      qacc_out[worldid, i] = float(D_rownnz[i])
+
+    # process diagonal elements from n-1 down to 0
+    for r_rev in range(NV):
+      i = NV - 1 - r_rev
+
+      rem_i = int(qacc_out[worldid, i])
+      rowadr_i = D_rowadr[i]
+      ii = rowadr_i + rem_i - 1
+      qacc_out[worldid, i] = float(rem_i - 1)
+
+      # cache diagonal element for row i
+      LUii = qLU_out[worldid, 0, ii]
+
+      # rows j above i (j < i), processed from i-1 down to 0
+      for c in range(i):
+        j = i - 1 - c
+
+        # get address of last remaining element of row j
+        rem_j = int(qacc_out[worldid, j])
+        rowadr_j = D_rowadr[j]
+        ji = rowadr_j + rem_j - 1
+
+        # process row j if (j,i) is non-zero
+        if D_colind[ji] == i:
+          # adjust remaining counter
+          rem_j = rem_j - 1
+          qacc_out[worldid, j] = float(rem_j)
+
+          # (j,i) = (j,i) / (i,i)
+          LUji = qLU_out[worldid, 0, ji] / LUii
+          qLU_out[worldid, 0, ji] = LUji
+
+          # (j,k) = (j,k) - (i,k) * (j,i) for k < i
+          icnt = rowadr_i
+          jcnt = rowadr_j
+          jend = rowadr_j + rem_j
+          while jcnt < jend:
+            col_i = D_colind[icnt]
+            col_j = D_colind[jcnt]
+            if col_i == col_j:
+              qLU_out[worldid, 0, jcnt] = qLU_out[worldid, 0, jcnt] - qLU_out[worldid, 0, icnt] * LUji
+              icnt = icnt + 1
+              jcnt = jcnt + 1
+            elif col_i > col_j:
+              jcnt = jcnt + 1
+            else:
+              icnt = icnt + 1
+
+    # Phase 2: LU solve (backward + forward substitution)
+
+    # Backward substitution: solve (U+I)*qacc = qfrc
+    for k_rev in range(NV):
+      i = NV - 1 - k_rev
+
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+      d1 = diag_i + 1
+      nnz_upper = D_rownnz[i] - d1
+
+      acc = qfrc[worldid, i]
+      for j in range(nnz_upper):
+        adr_j = rowadr_i + d1 + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+      qacc_out[worldid, i] = acc
+
+    # Forward substitution: solve L*qacc = qacc
+    for i in range(NV):
+      diag_i = D_diag[i]
+      rowadr_i = D_rowadr[i]
+
+      acc = qacc_out[worldid, i]
+      for j in range(diag_i):
+        adr_j = rowadr_i + j
+        col = D_colind[adr_j]
+        acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
+
+      qacc_out[worldid, i] = acc / qLU_out[worldid, 0, rowadr_i + diag_i]
+
+  return kernel
+
+
+@event_scope
+def factor_solve_lu(m: Model, d: Data, qLU: wp.array3d[float], qacc: wp.array2d[float], qfrc: wp.array2d[float]):
+  r"""Factorize and solve non-symmetric implicit system: qacc = A \\ qfrc.
+
+  qLU is overwritten in-place with the LU factors, then used to solve for qacc.
+
+  Args:
+    m: The model containing D-structure sparsity information.
+    d: The data object.
+    qLU: array containing the system matrix, overwritten with LU factors.
+    qacc: output array for the solution.
+    qfrc: input right-hand side.
+  """
+  wp.launch(
+    _factor_solve_lu_sparse_fused(m.nv),
+    dim=(d.nworld,),
+    inputs=[m.D_rownnz, m.D_rowadr, m.D_diag, m.D_colind, qfrc],
+    outputs=[qacc, qLU],
+  )
+
+
 @wp.kernel
 def _subtree_vel_forward(
   # Model:
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
index 0c804716..1a8ab32b 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
@@ -13,18 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 
-import dataclasses
 from math import ceil
 from math import sqrt
 
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src import island
 from mujoco.mjx.third_party.mujoco_warp._src import math
 from mujoco.mjx.third_party.mujoco_warp._src import smooth
 from mujoco.mjx.third_party.mujoco_warp._src import support
 from mujoco.mjx.third_party.mujoco_warp._src import types
-from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_func
+from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_factorize_solve_func
 from mujoco.mjx.third_party.mujoco_warp._src.block_cholesky import create_blocked_cholesky_solve_func
+from mujoco.mjx.third_party.mujoco_warp._src.types import InverseContext
+from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
+from mujoco.mjx.third_party.mujoco_warp._src.types import SolverContext
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import cache_kernel
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
 from mujoco.mjx.third_party.mujoco_warp._src.warp_util import scoped_mathdx_gemm_disabled
@@ -34,55 +37,12 @@ wp.set_module_options({"enable_backward": False})
 _BLOCK_CHOLESKY_DIM = 32
 
 
-@dataclasses.dataclass
-class InverseContext:
-  """Workspace arrays for inverse dynamics."""
-
-  Jaref: wp.array2d[float]
-  search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
-  done: wp.array[bool]
-  changed_efc_ids: wp.array2d[int]
-  changed_efc_count: wp.array[int]
-
-
-@dataclasses.dataclass
-class SolverContext:
-  """Workspace arrays for constraint solver."""
-
-  Jaref: wp.array2d[float]
-  search_dot: wp.array[float]
-  gauss: wp.array[float]
-  cost: wp.array[float]
-  prev_cost: wp.array[float]
-  done: wp.array[bool]
-  grad: wp.array2d[float]
-  grad_dot: wp.array[float]
-  Mgrad: wp.array2d[float]
-  search: wp.array2d[float]
-  mv: wp.array2d[float]
-  jv: wp.array2d[float]
-  quad: wp.array2d[wp.vec3]
-  quad_gauss: wp.array[wp.vec3]
-  alpha: wp.array[float]
-  prev_grad: wp.array2d[float]
-  prev_Mgrad: wp.array2d[float]
-  beta: wp.array[float]
-  h: wp.array3d[float]
-  hfactor: wp.array3d[float]
-  # Incremental Hessian update (Newton only)
-  changed_efc_ids: wp.array2d[int]
-  changed_efc_count: wp.array[int]
-
-
 def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
   """Create an InverseContext with allocated workspace arrays.
 
   Args:
-    m: Model containing nv, nv_pad, and solver type.
-    d: Data containing nworld and njmax.
+    m: Model.
+    d: Data.
 
   Returns:
     InverseContext with allocated arrays.
@@ -102,12 +62,55 @@ def create_inverse_context(m: types.Model, d: types.Data) -> InverseContext:
   )
 
 
+def create_island_solver_context(m: types.Model, d: types.Data) -> IslandSolverContext:
+  """Create an IslandSolverContext with allocated workspace arrays.
+
+  Args:
+    m: Model.
+    d: Data.
+
+  Returns:
+    IslandSolverContext with allocated arrays.
+  """
+  nworld = d.nworld
+  nv = m.nv
+  nv_pad = m.nv_pad
+  njmax = d.njmax
+  ntree = m.ntree
+
+  alloc_h = m.opt.solver == types.SolverType.NEWTON
+  alloc_island_cg = m.opt.solver == types.SolverType.CG
+
+  return IslandSolverContext(
+    Jaref=wp.empty((nworld, njmax), dtype=float),
+    jv=wp.empty((nworld, njmax), dtype=float),
+    search=wp.empty((nworld, nv), dtype=float),
+    mv=wp.empty((nworld, nv), dtype=float),
+    grad=wp.zeros((nworld, nv_pad), dtype=float),
+    Mgrad=wp.zeros((nworld, nv_pad), dtype=float),
+    prev_grad=wp.empty((nworld, nv), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    prev_Mgrad=wp.empty((nworld, nv), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    h=wp.zeros((nworld, nv_pad, nv_pad), dtype=float) if alloc_h else wp.empty((nworld, 0, 0), dtype=float),
+    # Per-island solver scalars
+    cost=wp.empty((nworld, ntree), dtype=float),
+    prev_cost=wp.empty((nworld, ntree), dtype=float),
+    gauss=wp.empty((nworld, ntree), dtype=float),
+    search_dot=wp.empty((nworld, ntree), dtype=float),
+    grad_dot=wp.empty((nworld, ntree), dtype=float),
+    done=wp.empty((nworld, ntree), dtype=bool),
+    solver_niter=wp.empty((nworld, ntree), dtype=int),
+    beta=wp.empty((nworld, ntree), dtype=float) if alloc_island_cg else wp.empty((nworld, 0), dtype=float),
+    alpha=wp.empty((nworld, ntree), dtype=float),
+    Ma=wp.empty((nworld, nv), dtype=float),
+  )
+
+
 def create_solver_context(m: types.Model, d: types.Data) -> SolverContext:
   """Create a SolverContext with allocated workspace arrays.
 
   Args:
-    m: Model containing nv, nv_pad, and solver type.
-    d: Data containing nworld and njmax.
+    m: Model.
+    d: Data.
 
   Returns:
     SolverContext with allocated arrays.
@@ -1669,7 +1672,7 @@ def _linesearch(m: types.Model, d: types.Data, ctx: SolverContext, cost: wp.arra
     ctx: SolverContext
     cost: Scratch array for storing costs per (world, alpha) - used for parallel mode
   """
-  # mv = qM @ search (common to both parallel and iterative)
+  # mv = M @ search (common to both parallel and iterative)
   support.mul_m(m, d, ctx.mv, ctx.search, skip=ctx.done)
 
   # Fuse jv computation in-kernel for small nv (iterative only, dense only)
@@ -2063,9 +2066,9 @@ def update_gradient_h_incremental(
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  """Incrementally update lower triangle of H for changed constraints.
+  """Incrementally update upper triangle of H for changed constraints.
 
-  Each thread handles one (i, j) element of the lower triangle.
+  Each thread handles one unique (i, j) element and writes it to the upper triangle.
   For each changed constraint, adds or subtracts D * J[i] * J[j].
   """
   worldid, elementid = wp.tid()
@@ -2074,28 +2077,28 @@ def update_gradient_h_incremental(
   if n_changes == 0:
     return
 
-  # Lower triangle index: elementid -> (i, j) where i >= j
-  i = (int(wp.sqrt(float(1 + 8 * elementid))) - 1) // 2
-  j = elementid - (i * (i + 1)) // 2
+  # Upper-triangle enumeration: elementid -> (row, col) where row <= col.
+  col = (int(wp.sqrt(float(1 + 8 * elementid))) - 1) // 2
+  row = elementid - (col * (col + 1)) // 2
 
   delta = float(0.0)
   for change_idx in range(n_changes):
     efcid = changed_ids_in[worldid, change_idx]
-    Ji = efc_J_in[worldid, efcid, i]
-    if Ji == 0.0:
+    Jrow = efc_J_in[worldid, efcid, row]
+    if Jrow == 0.0:
       continue
-    Jj = efc_J_in[worldid, efcid, j]
-    if Jj == 0.0:
+    Jcol = efc_J_in[worldid, efcid, col]
+    if Jcol == 0.0:
       continue
 
     D = efc_D_in[worldid, efcid]
     if efc_state_in[worldid, efcid] == types.ConstraintState.QUADRATIC.value:
-      delta += D * Ji * Jj
+      delta += D * Jrow * Jcol
     else:
-      delta -= D * Ji * Jj
+      delta -= D * Jrow * Jcol
 
   if delta != 0.0:
-    ctx_h_out[worldid, i, j] += delta
+    ctx_h_out[worldid, row, col] += delta
 
 
 @wp.kernel
@@ -2113,7 +2116,7 @@ def update_gradient_h_incremental_sparse(
   # Out:
   ctx_h_out: wp.array3d[float],
 ):
-  """Incrementally update lower triangle of H for changed constraints (sparse J)."""
+  """Incrementally update upper triangle of H for changed constraints (sparse J)."""
   worldid, change_idx = wp.tid()
 
   n_changes = changed_count_in[worldid]
@@ -2144,8 +2147,8 @@ def update_gradient_h_incremental_sparse(
         continue
       colindj = efc_J_colind_in[worldid, 0, sparseidj]
       h = sign * Ji * Jj
-      # Ensure lower triangle: larger index first
-      if colindi >= colindj:
+      # Ensure upper triangle: smaller index first.
+      if colindi <= colindj:
         wp.atomic_add(ctx_h_out[worldid, colindi], colindj, h)
       else:
         wp.atomic_add(ctx_h_out[worldid, colindj], colindi, h)
@@ -2257,12 +2260,13 @@ def update_gradient_grad(
 
 
 @wp.kernel
-def update_gradient_set_h_qM_lower_sparse(
+def update_gradient_set_h_M_upper_sparse(
   # Model:
-  qM_fullm_i: wp.array[int],
-  qM_fullm_j: wp.array[int],
+  M_fullm_upper_i: wp.array[int],
+  M_fullm_upper_j: wp.array[int],
+  M_fullm_upper_elemid: wp.array[int],
   # Data in:
-  qM_in: wp.array3d[float],
+  M_in: wp.array3d[float],
   # In:
   ctx_done_in: wp.array[bool],
   # Out:
@@ -2273,9 +2277,10 @@ def update_gradient_set_h_qM_lower_sparse(
   if ctx_done_in[worldid]:
     return
 
-  i = qM_fullm_i[elementid]
-  j = qM_fullm_j[elementid]
-  ctx_h_out[worldid, i, j] += qM_in[worldid, 0, elementid]
+  i = M_fullm_upper_i[elementid]
+  j = M_fullm_upper_j[elementid]
+  madr = M_fullm_upper_elemid[elementid]
+  ctx_h_out[worldid, i, j] += M_in[worldid, 0, madr]
 
 
 @wp.func
@@ -2317,12 +2322,12 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    # get lower diagonal index
-    i = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
-    j = elementid - (i * (i + 1)) // 2
+    # Upper-triangle tile index: elementid -> (row, col) where row <= col.
+    col = (int(sqrt(float(1 + 8 * elementid))) - 1) // 2
+    row = elementid - (col * (col + 1)) // 2
 
-    offset_i = i * TILE_SIZE
-    offset_j = j * TILE_SIZE
+    offset_row = row * TILE_SIZE
+    offset_col = col * TILE_SIZE
 
     sum_val = wp.tile_zeros(shape=(TILE_SIZE, TILE_SIZE), dtype=wp.float32)
 
@@ -2334,12 +2339,12 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
       # AD: leaving bounds-check disabled here because I'm not entirely sure that
       # everything always hits the fast path. The padding takes care of any
       # potential OOB accesses.
-      J_ki = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_i), bounds_check=False)
+      J_krow = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_row), bounds_check=False)
 
-      if offset_i != offset_j:
-        J_kj = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_j), bounds_check=False)
+      if offset_row != offset_col:
+        J_kcol = wp.tile_load(efc_J_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(k, offset_col), bounds_check=False)
       else:
-        wp.tile_assign(J_kj, J_ki, (0, 0))
+        wp.tile_assign(J_kcol, J_krow, (0, 0))
 
       D_k = wp.tile_load(efc_D_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
       state = wp.tile_load(efc_state_in[worldid], shape=TILE_SIZE, offset=k, bounds_check=False)
@@ -2353,13 +2358,13 @@ def update_gradient_JTDAJ_sparse_tiled(tile_size: int, njmax: int):
       active_tile = wp.tile_map(active_check, tid_tile, threshold_tile)
       D_k = wp.tile_map(wp.mul, active_tile, D_k)
 
-      J_ki = wp.tile_map(wp.mul, wp.tile_transpose(J_ki), wp.tile_broadcast(D_k, shape=(TILE_SIZE, TILE_SIZE)))
+      J_krow = wp.tile_map(wp.mul, wp.tile_transpose(J_krow), wp.tile_broadcast(D_k, shape=(TILE_SIZE, TILE_SIZE)))
 
-      sum_val += wp.tile_matmul(J_ki, J_kj)
+      sum_val += wp.tile_matmul(J_krow, J_kcol)
 
     # AD: setting bounds_check to True explicitly here because for some reason it was
     # slower to disable it.
-    wp.tile_store(ctx_h_out[worldid], sum_val, offset=(offset_i, offset_j), bounds_check=True)
+    wp.tile_store(ctx_h_out[worldid], sum_val, offset=(offset_row, offset_col), bounds_check=True)
 
   return kernel
 
@@ -2375,7 +2380,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
   def kernel(
     # Data in:
     nefc_in: wp.array[int],
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     efc_J_in: wp.array3d[float],
     efc_D_in: wp.array2d[float],
     efc_state_in: wp.array2d[int],
@@ -2391,7 +2396,7 @@ def update_gradient_JTDAJ_dense_tiled(nv_pad: int, tile_size: int, njmax: int):
 
     nefc = nefc_in[worldid]
 
-    sum_val = wp.tile_load(qM_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True)
+    sum_val = wp.tile_load(M_in[worldid], shape=(nv_pad, nv_pad), bounds_check=True)
 
     # Each tile processes one output tile by looping over all constraints
     for k in range(0, njmax, TILE_SIZE_K):
@@ -2746,9 +2751,9 @@ def update_gradient_cholesky(tile_size: int):
       return
 
     mat_tile = wp.tile_load(h_in[worldid], shape=(TILE_SIZE, TILE_SIZE))
-    fact_tile = wp.tile_cholesky(mat_tile)
+    fact_tile = wp.tile_cholesky(mat_tile, fill_mode="upper")
     input_tile = wp.tile_load(ctx_grad_in[worldid], shape=TILE_SIZE)
-    output_tile = wp.tile_cholesky_solve(fact_tile, input_tile)
+    output_tile = wp.tile_cholesky_solve(fact_tile, input_tile, fill_mode="upper")
     wp.tile_store(ctx_Mgrad_out[worldid], output_tile)
 
   return kernel
@@ -2777,14 +2782,14 @@ def update_gradient_cholesky_blocked(tile_size: int, matrix_size: int):
     # runtime input is needed for the loop bounds, otherwise warp will unroll
     # unconditionally leading to shared memory capacity issues.
 
-    wp.static(create_blocked_cholesky_func(TILE_SIZE))(ctx_h_in[worldid], matrix_size, ctx_hfactor[worldid])
-    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
-      ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
+    wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+      ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
     )
 
   return kernel
 
 
+@cache_kernel
 def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size: int):
   """Blocked Cholesky that skips factorization when no constraints changed."""
 
@@ -2806,11 +2811,13 @@ def update_gradient_cholesky_blocked_skip_unchanged(tile_size: int, matrix_size:
       return
 
     if changed_count_in[worldid] > 0:
-      wp.static(create_blocked_cholesky_func(TILE_SIZE))(ctx_h_in[worldid], matrix_size, ctx_hfactor[worldid])
-
-    wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
-      ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
-    )
+      wp.static(create_blocked_cholesky_factorize_solve_func(TILE_SIZE, matrix_size))(
+        ctx_h_in[worldid], ctx_grad_in[worldid], matrix_size, ctx_hfactor[worldid], ctx_Mgrad_out[worldid]
+      )
+    else:
+      wp.static(create_blocked_cholesky_solve_func(TILE_SIZE, matrix_size))(
+        ctx_hfactor[worldid], ctx_grad_in[worldid], matrix_size, ctx_Mgrad_out[worldid]
+      )
 
   return kernel
 
@@ -2913,9 +2920,9 @@ def _JTDAJ_sparse(
         colindj = efc_J_colind_in[worldid, 0, sparseidj]
 
       h = Ji * Jj * efc_D
-      # Store in lower triangle only: ensure row >= col
-      row = wp.max(colindi, colindj)
-      col = wp.min(colindi, colindj)
+      # Store in upper triangle only: ensure row <= col.
+      row = wp.min(colindi, colindj)
+      col = wp.max(colindi, colindj)
       wp.atomic_add(h_out[worldid, row], col, h)
 
 
@@ -2933,7 +2940,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
   if m.opt.solver == types.SolverType.CG:
     smooth.solve_m(m, d, ctx.Mgrad, ctx.grad)
   elif m.opt.solver == types.SolverType.NEWTON:
-    # h = qM + (efc_J.T * efc_D * active) @ efc_J
+    # h = M + (efc_J.T * efc_D * active) @ efc_J
     if m.is_sparse:
       ctx.h.zero_()
       wp.launch(
@@ -2944,9 +2951,9 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
       )
 
       wp.launch(
-        update_gradient_set_h_qM_lower_sparse,
-        dim=(d.nworld, m.qM_fullm_i.size),
-        inputs=[m.qM_fullm_i, m.qM_fullm_j, d.qM, ctx.done],
+        update_gradient_set_h_M_upper_sparse,
+        dim=(d.nworld, m.M_fullm_upper_i.size),
+        inputs=[m.M_fullm_upper_i, m.M_fullm_upper_j, m.M_fullm_upper_elemid, d.M, ctx.done],
         outputs=[ctx.h],
       )
     else:
@@ -2956,7 +2963,7 @@ def _update_gradient(m: types.Model, d: types.Data, ctx: SolverContext):
           dim=d.nworld,
           inputs=[
             d.nefc,
-            d.qM,
+            d.M,
             d.efc.J,
             d.efc.D,
             d.efc.state,
@@ -3066,7 +3073,7 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
     outputs=[ctx.grad, ctx.grad_dot],
   )
 
-  # Update lower triangle of H with delta from changed constraints
+  # Update upper triangle of H with delta from changed constraints.
   if m.is_sparse:
     wp.launch(
       update_gradient_h_incremental_sparse,
@@ -3084,10 +3091,10 @@ def _update_gradient_incremental(m: types.Model, d: types.Data, ctx: SolverConte
       outputs=[ctx.h],
     )
   else:
-    lower_tri_dim = m.nv * (m.nv + 1) // 2
+    tri_dim = m.nv * (m.nv + 1) // 2
     wp.launch(
       update_gradient_h_incremental,
-      dim=(d.nworld, lower_tri_dim),
+      dim=(d.nworld, tri_dim),
       inputs=[
         d.efc.J,
         d.efc.D,
@@ -3327,7 +3334,7 @@ def init_context(m: types.Model, d: types.Data, ctx: SolverContext | InverseCont
     outputs=[ctx.Jaref],
   )
 
-  # Ma = qM @ qacc
+  # Ma = M @ qacc
   support.mul_m(m, d, d.efc.Ma, d.qacc, skip=ctx.done)
 
   _update_constraint(m, d, ctx)
@@ -3342,8 +3349,17 @@ def solve(m: types.Model, d: types.Data):
     wp.copy(d.qacc, d.qacc_smooth)
     d.solver_niter.fill_(0)
   else:
-    ctx = create_solver_context(m, d)
-    _solve(m, d, ctx)
+    if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
+      ctx = create_island_solver_context(m, d)
+      island.compute_island_mapping(m, d, ctx)
+      island.gather_island_inputs(m, d, ctx)
+      _solve_islands(m, d, ctx)
+      # Ma is needed by Euler/implicit integrators for implicit damping
+      scatter_Ma = m.opt.integrator != types.IntegratorType.RK4
+      island.scatter_island_results(m, d, ctx, scatter_Ma=scatter_Ma)
+    else:
+      ctx = create_solver_context(m, d)
+      _solve(m, d, ctx)
 
 
 def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
@@ -3384,3 +3400,2104 @@ def _solve(m: types.Model, d: types.Data, ctx: SolverContext):
     # It should be removed when JAX becomes compatible.
     for _ in range(m.opt.iterations):
       _solver_iteration(m, d, ctx, step_size_cost, nsolving)
+
+
+# TODO(team): Consolidate monolithic and island solver code where possible
+@event_scope
+def _solve_islands(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Solve constraints for all islands in parallel.
+
+  All islands are processed simultaneously. Island-local arrays in ctx
+  (iacc, iefc_J, iefc_D, etc.) are indexed by idof/iefc, and each thread
+  determines its island via idof_islandid/iefc_islandid lookup tables.
+  """
+  # Initialize iacc from warmstart or smooth
+  if not (m.opt.disableflags & types.DisableBit.WARMSTART):
+    wp.launch(
+      gather_warmstart_island,
+      dim=(d.nworld, m.nv),
+      inputs=[d.nidof, d.qacc_warmstart, d.map_idof2dof],
+      outputs=[d.iqacc],
+    )
+  else:
+    wp.copy(d.iqacc, d.iqacc_smooth)
+
+  # Initialize island context
+  init_context_island(m, d, ctx)
+
+  # search = -Mgrad
+  wp.launch(
+    solve_init_search_island,
+    dim=(d.nworld, m.nv),
+    inputs=[d.nidof, ctx.Mgrad, d.dof_islandid, ctx.done],
+    outputs=[ctx.search, ctx.search_dot],
+  )
+
+  # nsolving tracks how many active islands still have unconverged globally
+  nsolving = wp.zeros((1,), dtype=int)
+  wp.launch(
+    solve_init_nsolving_island,
+    dim=d.nworld,
+    inputs=[d.nisland],
+    outputs=[nsolving],
+  )
+
+  if m.opt.iterations != 0 and m.opt.graph_conditional:
+    wp.capture_while(
+      nsolving,
+      while_body=_solver_iteration_island,
+      m=m,
+      d=d,
+      ctx=ctx,
+      nsolving=nsolving,
+    )
+  else:
+    for _ in range(m.opt.iterations):
+      _solver_iteration_island(m, d, ctx, nsolving)
+
+
+@wp.kernel
+def gather_warmstart_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  qacc_warmstart_in: wp.array2d[float],
+  map_idof2dof_in: wp.array2d[int],
+  # Out:
+  iacc_out: wp.array2d[float],
+):
+  """Gather qacc_warmstart into island-local order."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  dof = map_idof2dof_in[worldid, idofid]
+  iacc_out[worldid, idofid] = qacc_warmstart_in[worldid, dof]
+
+
+@wp.kernel
+def solve_init_efc_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  # Out:
+  island_cost_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+  island_done_out: wp.array2d[bool],
+  island_solver_niter_out: wp.array2d[int],
+):
+  """Initialize per-island solver scalars."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  island_cost_out[worldid, islandid] = 0.0
+  island_search_dot_out[worldid, islandid] = 0.0
+  island_done_out[worldid, islandid] = False
+  island_solver_niter_out[worldid, islandid] = 0
+
+
+@wp.kernel
+def solve_init_jaref_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iacc_in: wp.array2d[float],
+  iefc_aref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Jaref_out: wp.array2d[float],
+):
+  """Jaref[iefcid] = iefc_J[iefcid] @ iacc - iefc_aref[iefcid] for all island EFCs."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  acc = float(0.0)
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k in range(rownnz):
+      adr = rowadr + k
+      Ji = iefc_J_in[worldid, 0, adr]
+      idof = iefc_J_colind_in[worldid, 0, adr]
+      acc += Ji * iacc_in[worldid, idof]
+  else:
+    idofadr = island_idofadr_in[worldid, islandid]
+    inv = island_nv_in[worldid, islandid]
+    for i in range(inv):
+      idof = idofadr + i
+      acc += iefc_J_in[worldid, iefcid, idof] * iacc_in[worldid, idof]
+
+  Jaref_out[worldid, iefcid] = acc - iefc_aref_in[worldid, iefcid]
+
+
+@wp.kernel
+def solve_init_search_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  Mgrad_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  search_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+):
+  """Search = -Mgrad for all island DOFs."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  s = -Mgrad_in[worldid, idofid]
+  search_out[worldid, idofid] = s
+  wp.atomic_add(island_search_dot_out, worldid, islandid, s * s)
+
+
+@wp.kernel
+def update_constraint_init_cost_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  # In:
+  island_cost_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_gauss_out: wp.array2d[float],
+  island_cost_out: wp.array2d[float],
+  island_prev_cost_out: wp.array2d[float],
+):
+  """Save prev_cost and zero cost/gauss per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  island_prev_cost_out[worldid, islandid] = island_cost_in[worldid, islandid]
+  island_cost_out[worldid, islandid] = 0.0
+  island_gauss_out[worldid, islandid] = 0.0
+
+
+@wp.kernel
+def update_constraint_efc_island(
+  # Model:
+  opt_impratio_invsqrt: wp.array[float],
+  # Data in:
+  nefc_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  njmax_in: int,
+  nacon_in: wp.array[int],
+  # In:
+  iefc_type_in: wp.array2d[int],
+  iefc_id_in: wp.array2d[int],
+  iefc_D_in: wp.array2d[float],
+  iefc_frictionloss_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  iefc_force_out: wp.array2d[float],
+  iefc_state_out: wp.array2d[int],
+  island_cost_out: wp.array2d[float],
+):
+  """Compute force, state, and cost for each island constraint."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  # Local position within island
+  iefcadr = island_efcadr_in[worldid, islandid]
+  local_iefcid = iefcid - iefcadr
+  ine = island_ne_in[worldid, islandid]
+  inf = island_nf_in[worldid, islandid]
+
+  jaref = Jaref_in[worldid, iefcid]
+  D = iefc_D_in[worldid, iefcid]
+
+  force = float(0.0)
+  state = types.ConstraintState.SATISFIED.value
+  cost = float(0.0)
+
+  # Determine constraint category by position within island
+  if local_iefcid < ine:
+    # Equality constraint: always quadratic
+    force = -D * jaref
+    state = types.ConstraintState.QUADRATIC.value
+    cost = 0.5 * D * jaref * jaref
+  elif local_iefcid < ine + inf:
+    # Friction constraint
+    f = iefc_frictionloss_in[worldid, iefcid]
+    rf = math.safe_div(f, D)
+    if jaref <= -rf:
+      force = f
+      state = types.ConstraintState.LINEARNEG.value
+      cost = -f * (0.5 * rf + jaref)
+    elif jaref >= rf:
+      force = -f
+      state = types.ConstraintState.LINEARPOS.value
+      cost = -f * (0.5 * rf - jaref)
+    else:
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+  elif iefc_type_in[worldid, iefcid] != types.ConstraintType.CONTACT_ELLIPTIC:
+    # Limit, frictionless contact, pyramidal friction cone contact
+    if jaref >= 0.0:
+      force = 0.0
+      state = types.ConstraintState.SATISFIED.value
+      cost = 0.0
+    else:
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+  else:
+    # Elliptic friction cone contact
+    conid = iefc_id_in[worldid, iefcid]
+
+    if conid >= nacon_in[0]:
+      return
+
+    dim = contact_dim_in[conid]
+    friction = contact_friction_in[conid]
+    mu = friction[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+
+    # Get island-local index for lead row
+    efcid0_global = contact_efc_address_in[conid, 0]
+    if efcid0_global < 0:
+      return
+    ic0 = map_efc2iefc_in[worldid, efcid0_global]
+
+    N = Jaref_in[worldid, ic0] * mu
+
+    ufrictionj = float(0.0)
+    TT = float(0.0)
+    for j in range(1, dim):
+      efcidj_global = contact_efc_address_in[conid, j]
+      if efcidj_global < 0:
+        return
+      icj = map_efc2iefc_in[worldid, efcidj_global]
+      frictionj = friction[j - 1]
+      uj = Jaref_in[worldid, icj] * frictionj
+      TT += uj * uj
+      if iefcid == icj:
+        ufrictionj = uj * frictionj
+
+    if TT <= 0.0:
+      T = 0.0
+    else:
+      T = wp.sqrt(TT)
+
+    # top zone
+    if (N >= mu * T) or ((T <= 0.0) and (N >= 0.0)):
+      force = 0.0
+      state = types.ConstraintState.SATISFIED.value
+    # bottom zone
+    elif (mu * N + T <= 0.0) or ((T <= 0.0) and (N < 0.0)):
+      force = -D * jaref
+      state = types.ConstraintState.QUADRATIC.value
+      cost = 0.5 * D * jaref * jaref
+    # middle zone
+    else:
+      dm = math.safe_div(iefc_D_in[worldid, ic0], mu * mu * (1.0 + mu * mu))
+      nmt = N - mu * T
+
+      force_normal = -dm * nmt * mu
+
+      if iefcid == ic0:
+        force = force_normal
+        cost = 0.5 * dm * nmt * nmt
+      else:
+        force = -math.safe_div(force_normal, T) * ufrictionj
+
+      state = types.ConstraintState.CONE.value
+
+  iefc_force_out[worldid, iefcid] = force
+  iefc_state_out[worldid, iefcid] = state
+  wp.atomic_add(island_cost_out, worldid, islandid, cost)
+
+
+@wp.kernel
+def update_constraint_init_qfrc_constraint_dense_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  nidof_in: wp.array[int],
+  island_nefc_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  iefc_J_in: wp.array3d[float],
+  iefc_force_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ifrc_constraint_out: wp.array2d[float],
+):
+  """ifrc_constraint = iefc_J.T @ iefc_force for all island DOFs."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    ifrc_constraint_out[worldid, idofid] = 0.0
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    ifrc_constraint_out[worldid, idofid] = 0.0
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  iefcadr = island_efcadr_in[worldid, islandid]
+  inefc = island_nefc_in[worldid, islandid]
+  acc = float(0.0)
+  for iefcid in range(iefcadr, iefcadr + inefc):
+    acc += iefc_J_in[worldid, iefcid, idofid] * iefc_force_in[worldid, iefcid]
+
+  ifrc_constraint_out[worldid, idofid] = acc
+
+
+@wp.kernel
+def update_constraint_init_qfrc_constraint_sparse_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_force_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ifrc_constraint_out: wp.array2d[float],
+):
+  """ifrc_constraint += iefc_J.T @ iefc_force for all island EFCs (sparse parallel per EFC)."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+
+  rownnz = iefc_J_rownnz_in[worldid, iefcid]
+  rowadr = iefc_J_rowadr_in[worldid, iefcid]
+  force = iefc_force_in[worldid, iefcid]
+
+  for k in range(rownnz):
+    adr = rowadr + k
+    Ji = iefc_J_in[worldid, 0, adr]
+    idof = iefc_J_colind_in[worldid, 0, adr]
+    wp.atomic_add(ifrc_constraint_out, worldid, idof, Ji * force)
+
+
+@wp.kernel
+def update_constraint_gauss_cost_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  iacc_in: wp.array2d[float],
+  ifrc_smooth_in: wp.array2d[float],
+  iacc_smooth_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_gauss_out: wp.array2d[float],
+  island_cost_out: wp.array2d[float],
+):
+  """Gauss cost: 0.5 * (Ma - qfrc_smooth).T @ (qacc - qacc_smooth) per island."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  dq = iacc_in[worldid, idofid] - iacc_smooth_in[worldid, idofid]
+  df = iMa_in[worldid, idofid] - ifrc_smooth_in[worldid, idofid]
+  gauss = 0.5 * df * dq
+
+  wp.atomic_add(island_gauss_out, worldid, islandid, gauss)
+  wp.atomic_add(island_cost_out, worldid, islandid, gauss)
+
+
+@wp.kernel
+def update_gradient_grad_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  ifrc_smooth_in: wp.array2d[float],
+  ifrc_constraint_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  grad_out: wp.array2d[float],
+  island_grad_dot_out: wp.array2d[float],
+):
+  """Grad = Ma - qfrc_smooth - qfrc_constraint, grad_dot per island."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  g = iMa_in[worldid, idofid] - ifrc_smooth_in[worldid, idofid] - ifrc_constraint_in[worldid, idofid]
+  grad_out[worldid, idofid] = g
+  wp.atomic_add(island_grad_dot_out, worldid, islandid, g * g)
+
+
+@wp.kernel
+def linesearch_jv_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  nidof_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  njmax_in: int,
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  search_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  jv_out: wp.array2d[float],
+):
+  """Jv = iefc_J @ search for all island EFCs."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  acc = float(0.0)
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k in range(rownnz):
+      adr = rowadr + k
+      Ji = iefc_J_in[worldid, 0, adr]
+      idof = iefc_J_colind_in[worldid, 0, adr]
+      acc += Ji * search_in[worldid, idof]
+  else:
+    idofadr = island_idofadr_in[worldid, islandid]
+    inv = island_nv_in[worldid, islandid]
+    for i in range(inv):
+      idof = idofadr + i
+      acc += iefc_J_in[worldid, iefcid, idof] * search_in[worldid, idof]
+
+  jv_out[worldid, iefcid] = acc
+
+
+@wp.func
+def _eval_elliptic_cost_island(
+  # Model:
+  opt_impratio_invsqrt: float,  # kernel_analyzer: off
+  # Data in:
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  # In:
+  alpha: float,
+  conid: int,
+  iefc_D_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  jv_in: wp.array2d[float],
+  worldid: int,
+) -> wp.vec3:
+  dim = contact_dim_in[conid]
+  friction = contact_friction_in[conid]
+  mu = friction[0] * opt_impratio_invsqrt
+
+  ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+  D0 = iefc_D_in[worldid, ic0]
+  ja0 = Jaref_in[worldid, ic0]
+  jv0 = jv_in[worldid, ic0]
+
+  # Bottom-zone quad for the full contact (scalar quadratic over all rows)
+  quad = wp.vec3(0.5 * ja0 * ja0 * D0, jv0 * ja0 * D0, 0.5 * jv0 * jv0 * D0)
+
+  u0 = ja0 * mu
+  v0 = jv0 * mu
+  uu = float(0.0)
+  uv = float(0.0)
+  vv = float(0.0)
+
+  for j in range(1, dim):
+    icj = map_efc2iefc_in[worldid, contact_efc_address_in[conid, j]]
+    jaj = Jaref_in[worldid, icj]
+    jvj = jv_in[worldid, icj]
+    dj = iefc_D_in[worldid, icj]
+    DJj = dj * jaj
+
+    quad += wp.vec3(0.5 * jaj * DJj, jvj * DJj, 0.5 * jvj * dj * jvj)
+
+    frictionj = friction[j - 1]
+    uj = jaj * frictionj
+    vj = jvj * frictionj
+    uu += uj * uj
+    uv += uj * vj
+    vv += vj * vj
+
+  mu2 = mu * mu
+  dm = math.safe_div(D0, mu2 * (1.0 + mu2))
+
+  # Compute N, Tsqr
+  N = u0 + alpha * v0
+  Tsqr = uu + alpha * (2.0 * uv + alpha * vv)
+
+  if Tsqr <= 0.0:
+    if N < 0.0:
+      return _eval_pt(quad, alpha)
+    return wp.vec3(0.0)
+
+  T = wp.sqrt(Tsqr)
+
+  # top zone
+  if N >= mu * T:
+    return wp.vec3(0.0)
+  # bottom zone
+  if mu * N + T <= 0.0:
+    return _eval_pt(quad, alpha)
+
+  # middle zone
+  N1 = v0
+  T1 = (uv + alpha * vv) / T
+  T2 = vv / T - (uv + alpha * vv) * T1 / (T * T)
+  nmt = N - mu * T
+  return wp.vec3(
+    0.5 * dm * nmt * nmt,
+    dm * nmt * (N1 - mu * T1),
+    dm * ((N1 - mu * T1) * (N1 - mu * T1) + nmt * (-mu * T2)),
+  )
+
+
+# TODO(team): refactor linesearch_island
+@wp.kernel
+def linesearch_island(
+  # Model:
+  opt_tolerance: wp.array[float],
+  opt_ls_tolerance: wp.array[float],
+  opt_ls_iterations: int,
+  opt_impratio_invsqrt: wp.array[float],
+  stat_meaninertia: wp.array[float],
+  # Data in:
+  nefc_in: wp.array[int],
+  nisland_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  nidof_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  island_ne_in: wp.array2d[int],
+  island_nf_in: wp.array2d[int],
+  island_efcadr_in: wp.array2d[int],
+  island_nefc_in: wp.array2d[int],
+  map_efc2iefc_in: wp.array2d[int],
+  njmax_in: int,
+  nacon_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  # In:
+  iefc_type_in: wp.array2d[int],
+  iefc_id_in: wp.array2d[int],
+  iefc_D_in: wp.array2d[float],
+  iefc_frictionloss_in: wp.array2d[float],
+  Jaref_in: wp.array2d[float],
+  jv_in: wp.array2d[float],
+  mv_in: wp.array2d[float],
+  search_in: wp.array2d[float],
+  ifrc_smooth_in: wp.array2d[float],
+  iMa_in: wp.array2d[float],
+  island_search_dot_in: wp.array2d[float],
+  island_gauss_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_alpha_out: wp.array2d[float],
+):
+  """Linesearch per island."""
+  worldid, islandid = wp.tid()
+  nisland = nisland_in[worldid]
+  if islandid >= nisland:
+    island_alpha_out[worldid, islandid] = 0.0
+    return
+  nefc = wp.min(njmax_in, nefc_in[worldid])
+  tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
+  ls_tolerance = opt_ls_tolerance[worldid % opt_ls_tolerance.shape[0]]
+  meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]]
+  impratio_invsqrt = opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+
+  if island_done_in[worldid, islandid]:
+    island_alpha_out[worldid, islandid] = 0.0
+    return
+
+  iefcadr = island_efcadr_in[worldid, islandid]
+  ine = island_ne_in[worldid, islandid]
+  inf = island_nf_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+  idofadr = island_idofadr_in[worldid, islandid]
+
+  # Get island nefc
+  isle_nefc_end = iefcadr + island_nefc_in[worldid, islandid]
+
+  # Compute gauss quad: [gauss, s.T @ (Ma - frc_smooth), 0.5 * s.T @ mv]
+  quad_gauss_1 = float(0.0)
+  quad_gauss_2 = float(0.0)
+  for i in range(inv):
+    idof = idofadr + i
+    s = search_in[worldid, idof]
+    quad_gauss_1 += s * (iMa_in[worldid, idof] - ifrc_smooth_in[worldid, idof])
+    quad_gauss_2 += 0.5 * s * mv_in[worldid, idof]
+  quad_gauss = wp.vec3(island_gauss_in[worldid, islandid], quad_gauss_1, quad_gauss_2)
+
+  # gtol
+  snorm = wp.sqrt(island_search_dot_in[worldid, islandid])
+  scale = meaninertia * float(inv)
+  gtol = wp.max(tolerance * ls_tolerance * snorm * scale, 1e-6)
+
+  # p0: cost/grad/hessian at alpha=0
+  p0 = wp.vec3(quad_gauss[0], quad_gauss[1], 2.0 * quad_gauss[2])
+  for iefcid in range(iefcadr, isle_nefc_end):
+    if iefcid >= nefc:
+      break
+    local_iefcid = iefcid - iefcadr
+    D = iefc_D_in[worldid, iefcid]
+    ja = Jaref_in[worldid, iefcid]
+    jv_val = jv_in[worldid, iefcid]
+    if local_iefcid < ine:
+      # Equality: always active
+      jvD = jv_val * D
+      p0 += wp.vec3(0.5 * D * ja * ja, jvD * ja, jv_val * jvD)
+    elif local_iefcid < ine + inf:
+      # Friction
+      f = iefc_frictionloss_in[worldid, iefcid]
+      rf = math.safe_div(f, D)
+      p0 += _eval_frictionloss_pt(ja, f, rf, jv_val, D)
+    elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+      conid = iefc_id_in[worldid, iefcid]
+      if conid < nacon_in[0]:
+        ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+        if iefcid == ic0:
+          p0 += _eval_elliptic_cost_island(
+            impratio_invsqrt,
+            contact_friction_in,
+            contact_dim_in,
+            contact_efc_address_in,
+            map_efc2iefc_in,
+            0.0,
+            conid,
+            iefc_D_in,
+            Jaref_in,
+            jv_in,
+            worldid,
+          )
+    else:
+      # Inequality
+      if ja < 0.0:
+        jvD = jv_val * D
+        p0 += wp.vec3(0.5 * D * ja * ja, jvD * ja, jv_val * jvD)
+
+  # Newton step: lo_alpha_in = -p0[1] / p0[2]
+  lo_alpha_in = -math.safe_div(p0[1], p0[2])
+
+  # Evaluate at Newton step
+  lo_in = _eval_pt(quad_gauss, lo_alpha_in)
+  for iefcid in range(iefcadr, isle_nefc_end):
+    if iefcid >= nefc:
+      break
+    local_iefcid = iefcid - iefcadr
+    D = iefc_D_in[worldid, iefcid]
+    ja = Jaref_in[worldid, iefcid]
+    jv_val = jv_in[worldid, iefcid]
+    if local_iefcid < ine:
+      lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+    elif local_iefcid < ine + inf:
+      f = iefc_frictionloss_in[worldid, iefcid]
+      rf = math.safe_div(f, D)
+      x_a = ja + lo_alpha_in * jv_val
+      lo_in += _eval_frictionloss_pt(x_a, f, rf, jv_val, D)
+    elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+      conid = iefc_id_in[worldid, iefcid]
+      if conid < nacon_in[0]:
+        ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+        if iefcid == ic0:
+          lo_in += _eval_elliptic_cost_island(
+            impratio_invsqrt,
+            contact_friction_in,
+            contact_dim_in,
+            contact_efc_address_in,
+            map_efc2iefc_in,
+            lo_alpha_in,
+            conid,
+            iefc_D_in,
+            Jaref_in,
+            jv_in,
+            worldid,
+          )
+    else:
+      x_a = ja + lo_alpha_in * jv_val
+      if x_a < 0.0:
+        lo_in += _eval_pt_direct(ja, jv_val, D, lo_alpha_in)
+
+  # Accept Newton step if derivative is small and cost improved
+  initial_converged = wp.abs(lo_in[1]) < gtol and lo_in[0] < p0[0]
+
+  if initial_converged:
+    alpha = lo_alpha_in
+  else:
+    alpha = float(0.0)
+
+    # Initialize brackets
+    lo_less = int(0)
+    if lo_in[1] < p0[1]:
+      lo_less = int(1)
+    if lo_less == 1:
+      lo = lo_in
+      lo_alpha = lo_alpha_in
+      hi = p0
+      hi_alpha = float(0.0)
+    else:
+      lo = p0
+      lo_alpha = float(0.0)
+      hi = lo_in
+      hi_alpha = lo_alpha_in
+
+    for _iter in range(opt_ls_iterations):
+      lo_next_alpha = lo_alpha - math.safe_div(lo[1], lo[2])
+      hi_next_alpha = hi_alpha - math.safe_div(hi[1], hi[2])
+      mid_alpha = 0.5 * (lo_alpha + hi_alpha)
+
+      # Evaluate at 3 candidate alphas
+      lo_next = _eval_pt(quad_gauss, lo_next_alpha)
+      hi_next = _eval_pt(quad_gauss, hi_next_alpha)
+      mid = _eval_pt(quad_gauss, mid_alpha)
+
+      for iefcid in range(iefcadr, isle_nefc_end):
+        if iefcid >= nefc:
+          break
+        local_iefcid = iefcid - iefcadr
+        D = iefc_D_in[worldid, iefcid]
+        ja = Jaref_in[worldid, iefcid]
+        jv_val = jv_in[worldid, iefcid]
+        if local_iefcid < ine:
+          r_lo, r_hi, r_mid = _eval_pt_direct_3alphas(ja, jv_val, D, lo_next_alpha, hi_next_alpha, mid_alpha)
+        elif local_iefcid < ine + inf:
+          f = iefc_frictionloss_in[worldid, iefcid]
+          rf = math.safe_div(f, D)
+          x_lo = ja + lo_next_alpha * jv_val
+          x_hi = ja + hi_next_alpha * jv_val
+          x_mid = ja + mid_alpha * jv_val
+          r_lo = _eval_frictionloss_pt(x_lo, f, rf, jv_val, D)
+          r_hi = _eval_frictionloss_pt(x_hi, f, rf, jv_val, D)
+          r_mid = _eval_frictionloss_pt(x_mid, f, rf, jv_val, D)
+        elif iefc_type_in[worldid, iefcid] == types.ConstraintType.CONTACT_ELLIPTIC:
+          conid = iefc_id_in[worldid, iefcid]
+          r_lo = wp.vec3(0.0)
+          r_hi = wp.vec3(0.0)
+          r_mid = wp.vec3(0.0)
+          if conid < nacon_in[0]:
+            ic0 = map_efc2iefc_in[worldid, contact_efc_address_in[conid, 0]]
+            if iefcid == ic0:
+              r_lo = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                lo_next_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+              r_hi = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                hi_next_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+              r_mid = _eval_elliptic_cost_island(
+                impratio_invsqrt,
+                contact_friction_in,
+                contact_dim_in,
+                contact_efc_address_in,
+                map_efc2iefc_in,
+                mid_alpha,
+                conid,
+                iefc_D_in,
+                Jaref_in,
+                jv_in,
+                worldid,
+              )
+        else:
+          x_lo = ja + lo_next_alpha * jv_val
+          x_hi = ja + hi_next_alpha * jv_val
+          x_mid = ja + mid_alpha * jv_val
+          r_lo = wp.vec3(0.0)
+          r_hi = wp.vec3(0.0)
+          r_mid = wp.vec3(0.0)
+          if x_lo < 0.0:
+            r_lo = _eval_pt_direct(ja, jv_val, D, lo_next_alpha)
+          if x_hi < 0.0:
+            r_hi = _eval_pt_direct(ja, jv_val, D, hi_next_alpha)
+          if x_mid < 0.0:
+            r_mid = _eval_pt_direct(ja, jv_val, D, mid_alpha)
+        lo_next += r_lo
+        hi_next += r_hi
+        mid += r_mid
+
+      # Bracket swapping
+      swap_lo = int(0)
+      if _in_bracket(lo, lo_next):
+        lo = lo_next
+        lo_alpha = lo_next_alpha
+        swap_lo = int(1)
+      if _in_bracket(lo, mid):
+        lo = mid
+        lo_alpha = mid_alpha
+        swap_lo = int(1)
+      if _in_bracket(lo, hi_next):
+        lo = hi_next
+        lo_alpha = hi_next_alpha
+        swap_lo = int(1)
+
+      swap_hi = int(0)
+      if _in_bracket(hi, hi_next):
+        hi = hi_next
+        hi_alpha = hi_next_alpha
+        swap_hi = int(1)
+      if _in_bracket(hi, mid):
+        hi = mid
+        hi_alpha = mid_alpha
+        swap_hi = int(1)
+      if _in_bracket(hi, lo_next):
+        hi = lo_next
+        hi_alpha = lo_next_alpha
+        swap_hi = int(1)
+
+      # Done check
+      ls_done = (swap_lo == 0 and swap_hi == 0) or (lo[1] < 0.0 and lo[1] > -gtol) or (hi[1] > 0.0 and hi[1] < gtol)
+
+      # Update alpha if improved
+      if lo[0] < p0[0] or hi[0] < p0[0]:
+        if lo[0] < hi[0]:
+          alpha = lo_alpha
+        else:
+          alpha = hi_alpha
+
+      if ls_done:
+        break
+
+  island_alpha_out[worldid, islandid] = alpha
+
+
+@wp.kernel
+def linesearch_qacc_ma_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  search_in: wp.array2d[float],
+  mv_in: wp.array2d[float],
+  island_alpha_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  iacc_out: wp.array2d[float],
+  iMa_out: wp.array2d[float],
+):
+  """Update iacc and iMa after linesearch."""
+  worldid, tid = wp.tid()
+
+  # Process DOFs
+  if tid < nidof_in[worldid]:
+    idof = tid
+    islandid = idof_islandid_in[worldid, idof]
+    if islandid >= 0 and not island_done_in[worldid, islandid]:
+      alpha = island_alpha_in[worldid, islandid]
+      iacc_out[worldid, idof] += alpha * search_in[worldid, idof]
+      iMa_out[worldid, idof] += alpha * mv_in[worldid, idof]
+
+
+@wp.kernel
+def linesearch_jaref_island(
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  # In:
+  jv_in: wp.array2d[float],
+  island_alpha_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Jaref_out: wp.array2d[float],
+):
+  """Update Jaref after linesearch."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  alpha = island_alpha_in[worldid, islandid]
+  Jaref_out[worldid, iefcid] += alpha * jv_in[worldid, iefcid]
+
+
+@wp.kernel
+def solve_prev_grad_Mgrad_island(
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  grad_in: wp.array2d[float],
+  Mgrad_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  prev_grad_out: wp.array2d[float],
+  prev_Mgrad_out: wp.array2d[float],
+):
+  """Save prev_grad and prev_Mgrad per island DOF."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  prev_grad_out[worldid, idofid] = grad_in[worldid, idofid]
+  prev_Mgrad_out[worldid, idofid] = Mgrad_in[worldid, idofid]
+
+
+@wp.kernel
+def solve_beta_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  island_idofadr_in: wp.array2d[int],
+  grad_in: wp.array2d[float],
+  Mgrad_in: wp.array2d[float],
+  prev_grad_in: wp.array2d[float],
+  prev_Mgrad_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  island_beta_out: wp.array2d[float],
+):
+  """Polak-Ribière beta per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  if island_done_in[worldid, islandid]:
+    island_beta_out[worldid, islandid] = 0.0
+    return
+
+  idofadr = island_idofadr_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+
+  beta_num = float(0.0)
+  beta_den = float(0.0)
+  for i in range(inv):
+    idof = idofadr + i
+    pMg = prev_Mgrad_in[worldid, idof]
+    beta_num += grad_in[worldid, idof] * (Mgrad_in[worldid, idof] - pMg)
+    beta_den += prev_grad_in[worldid, idof] * pMg
+
+  island_beta_out[worldid, islandid] = wp.max(0.0, beta_num / wp.max(types.MJ_MINVAL, beta_den))
+
+
+@wp.kernel
+def solve_search_update_island(
+  # Model:
+  opt_solver: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  # In:
+  Mgrad_in: wp.array2d[float],
+  search_in: wp.array2d[float],
+  island_beta_in: wp.array2d[float],
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  search_out: wp.array2d[float],
+  island_search_dot_out: wp.array2d[float],
+):
+  """Update search direction per island DOF."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  s = -Mgrad_in[worldid, idofid]
+  if opt_solver == types.SolverType.CG:
+    s += island_beta_in[worldid, islandid] * search_in[worldid, idofid]
+
+  search_out[worldid, idofid] = s
+  wp.atomic_add(island_search_dot_out, worldid, islandid, s * s)
+
+
+@wp.kernel
+def solve_init_nsolving_island(
+  nisland_in: wp.array[int],
+  nsolving_out: wp.array[int],
+):
+  """Initialize active island count nsolving on device without CPU stalls."""
+  worldid = wp.tid()
+  wp.atomic_add(nsolving_out, 0, nisland_in[worldid])
+
+
+@wp.kernel
+def solve_done_island(
+  # Model:
+  opt_tolerance: wp.array[float],
+  opt_iterations: int,
+  stat_meaninertia: wp.array[float],
+  # Data in:
+  nisland_in: wp.array[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  island_grad_dot_in: wp.array2d[float],
+  island_cost_in: wp.array2d[float],
+  island_prev_cost_in: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  # Data out:
+  solver_niter_out: wp.array[int],
+  # Out:
+  island_done_out: wp.array2d[bool],
+  island_solver_niter_out: wp.array2d[int],
+  nsolving_out: wp.array[int],
+):
+  """Check convergence per island."""
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+
+  tolerance = opt_tolerance[worldid % opt_tolerance.shape[0]]
+  meaninertia = stat_meaninertia[worldid % stat_meaninertia.shape[0]]
+
+  if island_done_in[worldid, islandid]:
+    niter = island_solver_niter_out[worldid, islandid]
+    wp.atomic_max(solver_niter_out, worldid, niter)
+    return
+
+  island_solver_niter_out[worldid, islandid] += 1
+  niter = island_solver_niter_out[worldid, islandid]
+  wp.atomic_max(solver_niter_out, worldid, niter)
+
+  inv = island_nv_in[worldid, islandid]
+  improvement = _rescale(inv, meaninertia, island_prev_cost_in[worldid, islandid] - island_cost_in[worldid, islandid])
+  gradient = _rescale(inv, meaninertia, wp.sqrt(island_grad_dot_in[worldid, islandid]))
+  done = (improvement < tolerance) or (gradient < tolerance)
+  if done or niter >= opt_iterations:
+    island_done_out[worldid, islandid] = True
+    wp.atomic_sub(nsolving_out, 0, 1)
+
+
+@wp.kernel
+def update_gradient_JTDAJ_island(
+  # Model:
+  is_sparse: bool,
+  # Data in:
+  nefc_in: wp.array[int],
+  njmax_in: int,
+  island_idofadr_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_D_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Build island Hessian: ih += Jᵀ·D·J for active constraints."""
+  worldid, iefcid = wp.tid()
+
+  if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
+    return
+
+  islandid = iefc_islandid_in[worldid, iefcid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  state = iefc_state_in[worldid, iefcid]
+  if state != types.ConstraintState.QUADRATIC.value:
+    return
+  D = iefc_D_in[worldid, iefcid]
+
+  idofadr = island_idofadr_in[worldid, islandid]
+  inv = island_nv_in[worldid, islandid]
+  if is_sparse:
+    rownnz = iefc_J_rownnz_in[worldid, iefcid]
+    rowadr = iefc_J_rowadr_in[worldid, iefcid]
+    for k1 in range(rownnz):
+      adr1 = rowadr + k1
+      Ji = iefc_J_in[worldid, 0, adr1]
+      i = iefc_J_colind_in[worldid, 0, adr1]
+
+      for k2 in range(k1 + 1):
+        adr2 = rowadr + k2
+        Jj = iefc_J_in[worldid, 0, adr2]
+        j = iefc_J_colind_in[worldid, 0, adr2]
+
+        h = Ji * Jj * D
+        wp.atomic_add(ih_out[worldid, i], j, h)
+        if i != j:
+          wp.atomic_add(ih_out[worldid, j], i, h)
+  else:
+    for ii in range(inv):
+      i = idofadr + ii
+      Ji = iefc_J_in[worldid, iefcid, i]
+      if Ji == 0.0:
+        continue
+      for jj in range(ii + 1):
+        j = idofadr + jj
+        Jj = iefc_J_in[worldid, iefcid, j]
+        if Jj == 0.0:
+          continue
+        h = Ji * Jj * D
+        wp.atomic_add(ih_out[worldid, i], j, h)
+        if i != j:
+          wp.atomic_add(ih_out[worldid, j], i, h)
+
+
+@wp.kernel
+def update_gradient_set_h_M_sparse_island(
+  # Model:
+  M_fullm_i: wp.array[int],
+  M_fullm_j: wp.array[int],
+  M_elemid: wp.array2d[int],
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  dof_island_in: wp.array2d[int],
+  map_dof2idof_in: wp.array2d[int],
+  # In:
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add sparse mass matrix to island Hessian using global-to-island DOF mapping."""
+  worldid, elementid = wp.tid()
+
+  i_global = M_fullm_i[elementid]
+  j_global = M_fullm_j[elementid]
+
+  madr = M_elemid[i_global, j_global]
+  if madr < 0:
+    return
+
+  # Check both DOFs belong to an island
+  island_i = dof_island_in[worldid, i_global]
+  if island_i < 0:
+    return
+  if island_done_in[worldid, island_i]:
+    return
+
+  island_j = dof_island_in[worldid, j_global]
+  if island_j < 0:
+    return
+
+  # Both DOFs must be in the same island
+  if island_i != island_j:
+    return
+
+  idof_i = map_dof2idof_in[worldid, i_global]
+  idof_j = map_dof2idof_in[worldid, j_global]
+
+  val = M_in[worldid, 0, madr]
+  ih_out[worldid, idof_i, idof_j] += val
+  if idof_i != idof_j:
+    ih_out[worldid, idof_j, idof_i] += val
+
+
+@wp.kernel
+def update_gradient_set_h_M_dense_island(
+  # Model:
+  nv: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add dense mass matrix to island Hessian."""
+  worldid, idofid = wp.tid()
+
+  if idofid >= nidof_in[worldid]:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  dof_i = map_idof2dof_in[worldid, idofid]
+
+  # Copy row from M to ih, mapping columns
+  nid = nidof_in[worldid]
+  for jdof in range(nid):
+    dof_j = map_idof2dof_in[worldid, jdof]
+    ih_out[worldid, idofid, jdof] += M_in[worldid, dof_i, dof_j]
+
+
+@wp.kernel
+def update_gradient_JTCJ_island(
+  # Model:
+  opt_impratio_invsqrt: wp.array[float],
+  is_sparse: bool,
+  # Data in:
+  nacon_in: wp.array[int],
+  contact_friction_in: wp.array[types.vec5],
+  contact_dim_in: wp.array[int],
+  contact_efc_address_in: wp.array2d[int],
+  contact_worldid_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  naconmax_in: int,
+  nidof_in: wp.array[int],
+  map_efc2iefc_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  iefc_J_rownnz_in: wp.array2d[int],
+  iefc_J_rowadr_in: wp.array2d[int],
+  iefc_J_colind_in: wp.array3d[int],
+  iefc_J_in: wp.array3d[float],
+  iefc_D_in: wp.array2d[float],
+  iefc_state_in: wp.array2d[int],
+  Jaref_in: wp.array2d[float],
+  iefc_islandid_in: wp.array2d[int],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  ih_out: wp.array3d[float],
+):
+  """Add elliptic cone Hessian correction: Jᵀ·C·J for contacts in CONE state."""
+  conid = wp.tid()
+
+  if conid >= wp.min(naconmax_in, nacon_in[0]):
+    return
+
+  worldid = contact_worldid_in[conid]
+  condim = contact_dim_in[conid]
+
+  if condim == 1:
+    return
+
+  efcid0_global = contact_efc_address_in[conid, 0]
+  if efcid0_global < 0:
+    return
+
+  ic0 = map_efc2iefc_in[worldid, efcid0_global]
+  if iefc_state_in[worldid, ic0] != types.ConstraintState.CONE.value:
+    return
+
+  islandid = iefc_islandid_in[worldid, ic0]
+  if islandid < 0:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  inv = island_nv_in[worldid, islandid]
+  idofadr = island_idofadr_in[worldid, islandid]
+
+  fri = contact_friction_in[conid]
+  mu = fri[0] * opt_impratio_invsqrt[worldid % opt_impratio_invsqrt.shape[0]]
+  mu2 = mu * mu
+  dm = math.safe_div(iefc_D_in[worldid, ic0], mu2 * (1.0 + mu2))
+
+  if dm == 0.0:
+    return
+
+  # Compute n and u vector
+  n = Jaref_in[worldid, ic0] * mu
+  u = types.vec6(n, 0.0, 0.0, 0.0, 0.0, 0.0)
+
+  tt = float(0.0)
+  for j in range(1, condim):
+    efcidj_global = contact_efc_address_in[conid, j]
+    if efcidj_global < 0:
+      return
+    icj = map_efc2iefc_in[worldid, efcidj_global]
+    uj = Jaref_in[worldid, icj] * fri[j - 1]
+    tt += uj * uj
+    u[j] = uj
+
+  if tt <= 0.0:
+    t = 0.0
+  else:
+    t = wp.sqrt(tt)
+  t = wp.max(t, types.MJ_MINVAL)
+  ttt = wp.max(t * t * t, types.MJ_MINVAL)
+
+  # Accumulate cone correction into ih
+  for dim1id in range(condim):
+    if dim1id == 0:
+      ic1 = ic0
+    else:
+      efcid1_global = contact_efc_address_in[conid, dim1id]
+      if efcid1_global < 0:
+        return
+      ic1 = map_efc2iefc_in[worldid, efcid1_global]
+
+    ui = u[dim1id]
+
+    for dim2id in range(dim1id + 1):
+      if dim2id == 0:
+        ic2 = ic0
+      else:
+        efcid2_global = contact_efc_address_in[conid, dim2id]
+        if efcid2_global < 0:
+          return
+        ic2 = map_efc2iefc_in[worldid, efcid2_global]
+
+      uj = u[dim2id]
+
+      # Cone correction matrix
+      if dim1id == 0 and dim2id == 0:
+        hcone = 1.0
+      elif dim1id == 0:
+        hcone = -math.safe_div(mu, t) * uj
+      elif dim2id == 0:
+        hcone = -math.safe_div(mu, t) * ui
+      else:
+        hcone = mu * math.safe_div(n, ttt) * ui * uj
+        if dim1id == dim2id:
+          hcone += mu2 - mu * math.safe_div(n, t)
+
+      # Scale by dm * friction
+      if dim1id == 0:
+        fri1 = mu
+      else:
+        fri1 = fri[dim1id - 1]
+      if dim2id == 0:
+        fri2 = mu
+      else:
+        fri2 = fri[dim2id - 1]
+
+      hcone *= dm * fri1 * fri2
+
+      if hcone == 0.0:
+        continue
+
+      # Accumulate J1^T * hcone * J2 into ih (lower triangle)
+      if is_sparse:
+        if dim1id == dim2id:
+          rownnz = iefc_J_rownnz_in[worldid, ic1]
+          rowadr = iefc_J_rowadr_in[worldid, ic1]
+          for k1 in range(rownnz):
+            adr1 = rowadr + k1
+            J1 = iefc_J_in[worldid, 0, adr1]
+            i = iefc_J_colind_in[worldid, 0, adr1]
+
+            for k2 in range(k1 + 1):
+              adr2 = rowadr + k2
+              J2 = iefc_J_in[worldid, 0, adr2]
+              j = iefc_J_colind_in[worldid, 0, adr2]
+
+              val = hcone * J1 * J2
+              wp.atomic_add(ih_out[worldid, i], j, val)
+              if i != j:
+                wp.atomic_add(ih_out[worldid, j], i, val)
+        else:
+          rownnz1 = iefc_J_rownnz_in[worldid, ic1]
+          rowadr1 = iefc_J_rowadr_in[worldid, ic1]
+          rownnz2 = iefc_J_rownnz_in[worldid, ic2]
+          rowadr2 = iefc_J_rowadr_in[worldid, ic2]
+
+          for k1 in range(rownnz1):
+            adr1 = rowadr1 + k1
+            J1 = iefc_J_in[worldid, 0, adr1]
+            i = iefc_J_colind_in[worldid, 0, adr1]
+
+            for k2 in range(rownnz2):
+              adr2 = rowadr2 + k2
+              J2 = iefc_J_in[worldid, 0, adr2]
+              j = iefc_J_colind_in[worldid, 0, adr2]
+
+              val = hcone * J1 * J2
+              if i == j:
+                wp.atomic_add(ih_out[worldid, i], j, val * 2.0)
+              else:
+                wp.atomic_add(ih_out[worldid, i], j, val)
+                wp.atomic_add(ih_out[worldid, j], i, val)
+      else:
+        for i in range(inv):
+          J1i = iefc_J_in[worldid, ic1, idofadr + i]
+          if J1i == 0.0:
+            continue
+          for jj in range(i + 1):
+            J2j = iefc_J_in[worldid, ic2, idofadr + jj]
+            if J2j == 0.0:
+              continue
+            val = hcone * J1i * J2j
+            wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val)
+            if i != jj:
+              wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val)
+
+            if dim1id != dim2id:
+              J1j = iefc_J_in[worldid, ic1, idofadr + jj]
+              J2i = iefc_J_in[worldid, ic2, idofadr + i]
+              if J1j != 0.0 and J2i != 0.0:
+                val2 = hcone * J1j * J2i
+                wp.atomic_add(ih_out[worldid, idofadr + i], idofadr + jj, val2)
+                if i != jj:
+                  wp.atomic_add(ih_out[worldid, idofadr + jj], idofadr + i, val2)
+
+
+@wp.kernel
+def _cholesky_factorize_solve_island(
+  # Data in:
+  nisland_in: wp.array[int],
+  island_idofadr_in: wp.array2d[int],
+  island_nv_in: wp.array2d[int],
+  # In:
+  grad_in: wp.array2d[float],
+  ih_in: wp.array3d[float],
+  island_done_in: wp.array2d[bool],
+  # Out:
+  Mgrad_out: wp.array2d[float],
+):
+  """Per-island Cholesky factorize and solve: Mgrad = H⁻¹ @ grad.
+
+  One thread per (world, island). Performs dense in-place Cholesky factorization
+  on the island's inv x inv subblock of ih, then forward/backward substitution.
+  """
+  worldid, islandid = wp.tid()
+
+  if islandid >= nisland_in[worldid]:
+    return
+  if island_done_in[worldid, islandid]:
+    return
+
+  inv = island_nv_in[worldid, islandid]
+
+  if inv == 0:
+    return
+
+  adr = island_idofadr_in[worldid, islandid]
+  # Cholesky factorization in-place: L such that H = L @ L^T
+  for i in range(inv):
+    for j in range(i + 1):
+      s = ih_in[worldid, adr + i, adr + j]
+      for k in range(j):
+        s -= ih_in[worldid, adr + i, adr + k] * ih_in[worldid, adr + j, adr + k]
+      if i == j:
+        if s <= 1e-6:
+          s = 1e-6
+        ih_in[worldid, adr + i, adr + j] = wp.sqrt(s)
+      else:
+        div = ih_in[worldid, adr + j, adr + j]
+        ih_in[worldid, adr + i, adr + j] = s / wp.max(1e-6, div)
+
+  # Forward substitution: L @ y = grad  =>  y
+  for i in range(inv):
+    s = grad_in[worldid, adr + i]
+    for k in range(i):
+      s -= ih_in[worldid, adr + i, adr + k] * Mgrad_out[worldid, adr + k]
+    Mgrad_out[worldid, adr + i] = s / wp.max(1e-6, ih_in[worldid, adr + i, adr + i])
+
+  # Backward substitution: L^T @ x = y  =>  x = Mgrad
+  for i_rev in range(inv):
+    i = inv - 1 - i_rev
+    s = Mgrad_out[worldid, adr + i]
+    for k in range(i + 1, inv):
+      s -= ih_in[worldid, adr + k, adr + i] * Mgrad_out[worldid, adr + k]
+    Mgrad_out[worldid, adr + i] = s / wp.max(types.MJ_MINVAL, ih_in[worldid, adr + i, adr + i])
+
+
+def init_context_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Initialize island solver context."""
+  # Init per-island scalars
+  d.solver_niter.zero_()
+  wp.launch(
+    solve_init_efc_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.nisland],
+    outputs=[ctx.cost, ctx.search_dot, ctx.done, ctx.solver_niter],
+  )
+
+  # Jaref = iefc_J @ iacc - iefc_aref
+  wp.launch(
+    solve_init_jaref_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.island_nv,
+      d.njmax,
+      d.island_dofadr,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      d.iqacc,
+      d.efc.iaref,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.Jaref],
+  )
+
+  # iMa = M @ iacc (all islands in parallel)
+  support.mul_m_island(
+    m,
+    d,
+    ctx.Ma,
+    d.iqacc,
+    d.nidof,
+    d.map_idof2dof,
+    d.map_dof2idof,
+    d.dof_islandid,
+  )
+
+  # Update constraint
+  _update_constraint_island(m, d, ctx)
+
+  # Update gradient
+  if m.opt.solver == types.SolverType.NEWTON:
+    _update_gradient_incremental_island(m, d, ctx)
+  else:
+    _update_gradient_island(m, d, ctx)
+
+
+@event_scope
+def _update_constraint_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Update constraint arrays for island solver."""
+  # Save prev cost, zero cost/gauss
+  wp.launch(
+    update_constraint_init_cost_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[d.nisland, ctx.cost, ctx.done],
+    outputs=[ctx.gauss, ctx.cost, ctx.prev_cost],
+  )
+
+  # Compute force, state, cost per EFC
+  wp.launch(
+    update_constraint_efc_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.opt.impratio_invsqrt,
+      d.nefc,
+      d.contact.friction,
+      d.contact.dim,
+      d.contact.efc_address,
+      d.island_nefc,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.map_efc2iefc,
+      d.njmax,
+      d.nacon,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.iD,
+      d.efc.ifrictionloss,
+      ctx.Jaref,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[d.efc.iforce, d.efc.istate, ctx.cost],
+  )
+
+  # qfrc_constraint = J^T @ force
+  if m.is_sparse:
+    d.iqfrc_constraint.zero_()
+    wp.launch(
+      update_constraint_init_qfrc_constraint_sparse_island,
+      dim=(d.nworld, d.njmax),
+      inputs=[
+        d.nefc,
+        d.njmax,
+        d.efc.iJ_rownnz,
+        d.efc.iJ_rowadr,
+        d.efc.iJ_colind,
+        d.efc.iJ,
+        d.efc.iforce,
+        d.efc_islandid,
+        ctx.done,
+      ],
+      outputs=[d.iqfrc_constraint],
+    )
+  else:
+    wp.launch(
+      update_constraint_init_qfrc_constraint_dense_island,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        d.nefc,
+        d.nidof,
+        d.island_nefc,
+        d.island_efcadr,
+        d.njmax,
+        d.efc.iJ,
+        d.efc.iforce,
+        d.dof_islandid,
+        ctx.done,
+      ],
+      outputs=[d.iqfrc_constraint],
+    )
+
+  # Gauss cost
+  wp.launch(
+    update_constraint_gauss_cost_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqacc,
+      d.iqfrc_smooth,
+      d.iqacc_smooth,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.gauss, ctx.cost],
+  )
+
+
+@event_scope
+def _update_gradient_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Update gradient for island solver."""
+  # Zero grad_dot per island
+  ctx.grad_dot.zero_()
+
+  # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
+  wp.launch(
+    update_gradient_grad_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqfrc_smooth,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.grad, ctx.grad_dot],
+  )
+
+  # CG preconditioner: Mgrad = M^{-1} @ grad (direct solve)
+  support.solve_m_island(
+    m,
+    d,
+    ctx.Mgrad,
+    ctx.grad,
+    d.nidof,
+    d.map_idof2dof,
+  )
+
+
+@event_scope
+def _update_gradient_incremental_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Full Newton gradient update for islands: build H, factorize, solve."""
+  # Zero grad_dot per island
+  ctx.grad_dot.zero_()
+
+  # grad = Ma - frc_smooth - frc_constraint, accumulate grad_dot
+  wp.launch(
+    update_gradient_grad_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      d.iqfrc_smooth,
+      d.iqfrc_constraint,
+      ctx.Ma,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.grad, ctx.grad_dot],
+  )
+
+  # Build H = qM + Jᵀ·D·J
+  ctx.h.zero_()
+
+  # JTDAJ
+  wp.launch(
+    update_gradient_JTDAJ_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.njmax,
+      d.island_dofadr,
+      d.island_nv,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      d.efc.iD,
+      d.efc.istate,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.h],
+  )
+
+  # Add mass matrix
+  if m.is_sparse:
+    wp.launch(
+      update_gradient_set_h_M_sparse_island,
+      dim=(d.nworld, m.M_fullm_i.shape[0]),
+      inputs=[
+        m.M_fullm_i,
+        m.M_fullm_j,
+        m.M_elemid,
+        d.nidof,
+        d.M,
+        d.dof_island,
+        d.map_dof2idof,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+  else:
+    wp.launch(
+      update_gradient_set_h_M_dense_island,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.nv,
+        d.nidof,
+        d.M,
+        d.map_idof2dof,
+        d.dof_islandid,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+
+  # Elliptic cone correction: JTCJ
+  if m.opt.cone == types.ConeType.ELLIPTIC and d.naconmax > 0:
+    wp.launch(
+      update_gradient_JTCJ_island,
+      dim=d.naconmax,
+      inputs=[
+        m.opt.impratio_invsqrt,
+        m.is_sparse,
+        d.nacon,
+        d.contact.friction,
+        d.contact.dim,
+        d.contact.efc_address,
+        d.contact.worldid,
+        d.island_dofadr,
+        d.naconmax,
+        d.nidof,
+        d.map_efc2iefc,
+        d.island_nv,
+        d.efc.iJ_rownnz,
+        d.efc.iJ_rowadr,
+        d.efc.iJ_colind,
+        d.efc.iJ,
+        d.efc.iD,
+        d.efc.istate,
+        ctx.Jaref,
+        d.efc_islandid,
+        ctx.done,
+      ],
+      outputs=[ctx.h],
+    )
+
+  # Cholesky factorize and solve: Mgrad = H⁻¹ @ grad
+  wp.launch(
+    _cholesky_factorize_solve_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      d.nisland,
+      d.island_dofadr,
+      d.island_nv,
+      ctx.grad,
+      ctx.h,
+      ctx.done,
+    ],
+    outputs=[ctx.Mgrad],
+  )
+
+
+@event_scope
+def _linesearch_island(m: types.Model, d: types.Data, ctx: IslandSolverContext):
+  """Linesearch for island solver."""
+  # mv = M @ search (all islands)
+  support.mul_m_island(
+    m,
+    d,
+    ctx.mv,
+    ctx.search,
+    d.nidof,
+    d.map_idof2dof,
+    d.map_dof2idof,
+    d.dof_islandid,
+    island_done=ctx.done,
+  )
+
+  # jv = J @ search
+  wp.launch(
+    linesearch_jv_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      m.is_sparse,
+      d.nefc,
+      d.nidof,
+      d.island_nv,
+      d.njmax,
+      d.island_dofadr,
+      d.efc.iJ_rownnz,
+      d.efc.iJ_rowadr,
+      d.efc.iJ_colind,
+      d.efc.iJ,
+      ctx.search,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.jv],
+  )
+
+  # linesearch
+  wp.launch(
+    linesearch_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.opt.tolerance,
+      m.opt.ls_tolerance,
+      m.opt.ls_iterations,
+      m.opt.impratio_invsqrt,
+      m.stat.meaninertia,
+      d.nefc,
+      d.nisland,
+      d.contact.friction,
+      d.contact.dim,
+      d.contact.efc_address,
+      d.nidof,
+      d.island_nv,
+      d.island_ne,
+      d.island_nf,
+      d.island_efcadr,
+      d.island_nefc,
+      d.map_efc2iefc,
+      d.njmax,
+      d.nacon,
+      d.island_dofadr,
+      d.efc.itype,
+      d.efc.iid,
+      d.efc.iD,
+      d.efc.ifrictionloss,
+      ctx.Jaref,
+      ctx.jv,
+      ctx.mv,
+      ctx.search,
+      d.iqfrc_smooth,
+      ctx.Ma,
+      ctx.search_dot,
+      ctx.gauss,
+      ctx.done,
+    ],
+    outputs=[ctx.alpha],
+  )
+
+  # Update iacc, iMa
+  wp.launch(
+    linesearch_qacc_ma_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      d.nidof,
+      ctx.search,
+      ctx.mv,
+      ctx.alpha,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[d.iqacc, ctx.Ma],
+  )
+
+  # Update Jaref
+  wp.launch(
+    linesearch_jaref_island,
+    dim=(d.nworld, d.njmax),
+    inputs=[
+      d.nefc,
+      d.njmax,
+      ctx.jv,
+      ctx.alpha,
+      d.efc_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.Jaref],
+  )
+
+
+@event_scope
+def _solver_iteration_island(
+  m: types.Model,
+  d: types.Data,
+  ctx: IslandSolverContext,
+  nsolving: wp.array[int],
+):
+  """One iteration of island solver for all islands in parallel."""
+  _linesearch_island(m, d, ctx)
+
+  is_newton = m.opt.solver == types.SolverType.NEWTON
+  is_cg = not is_newton
+
+  # Save prev_grad, prev_Mgrad for CG
+  if is_cg:
+    wp.launch(
+      solve_prev_grad_Mgrad_island,
+      dim=(d.nworld, m.nv),
+      inputs=[d.nidof, ctx.grad, ctx.Mgrad, d.dof_islandid, ctx.done],
+      outputs=[ctx.prev_grad, ctx.prev_Mgrad],
+    )
+
+  # Update constraint
+  _update_constraint_island(m, d, ctx)
+
+  # Update gradient
+  if is_newton:
+    _update_gradient_incremental_island(m, d, ctx)
+  else:
+    _update_gradient_island(m, d, ctx)
+
+  # Polak-Ribière beta (CG only)
+  if is_cg:
+    wp.launch(
+      solve_beta_island,
+      dim=(d.nworld, m.ntree),
+      inputs=[
+        d.nisland,
+        d.island_nv,
+        d.island_dofadr,
+        ctx.grad,
+        ctx.Mgrad,
+        ctx.prev_grad,
+        ctx.prev_Mgrad,
+        ctx.done,
+      ],
+      outputs=[ctx.beta],
+    )
+
+  # Zero search_dot
+  ctx.search_dot.zero_()
+
+  # Search update
+  wp.launch(
+    solve_search_update_island,
+    dim=(d.nworld, m.nv),
+    inputs=[
+      m.opt.solver,
+      d.nidof,
+      ctx.Mgrad,
+      ctx.search,
+      ctx.beta,
+      d.dof_islandid,
+      ctx.done,
+    ],
+    outputs=[ctx.search, ctx.search_dot],
+  )
+
+  # Convergence check
+  d.solver_niter.zero_()
+  wp.launch(
+    solve_done_island,
+    dim=(d.nworld, m.ntree),
+    inputs=[
+      m.opt.tolerance,
+      m.opt.iterations,
+      m.stat.meaninertia,
+      d.nisland,
+      d.island_nv,
+      ctx.grad_dot,
+      ctx.cost,
+      ctx.prev_cost,
+      ctx.done,
+    ],
+    outputs=[d.solver_niter, ctx.done, ctx.solver_niter, nsolving],
+  )
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
index b42e9808..61e5b219 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
@@ -69,11 +69,11 @@ def mul_m_sparse(check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_sparse(
     # Model:
-    qM_mulm_rowadr: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -89,12 +89,12 @@ def mul_m_sparse(check_skip: bool):
 
     # Gather all contributions (diagonal + off-diagonal)
     acc = float(0.0)
-    start = qM_mulm_rowadr[dofid]
-    end = qM_mulm_rowadr[dofid + 1]
+    start = M_mulm_rowadr[dofid]
+    end = M_mulm_rowadr[dofid + 1]
     for k in range(start, end):
-      col = qM_mulm_col[k]
-      madr = qM_mulm_madr[k]
-      acc += qM_in[worldid, 0, madr] * vec[worldid, col]
+      col = M_mulm_col[k]
+      madr = M_mulm_madr[k]
+      acc += M_in[worldid, 0, madr] * vec[worldid, col]
 
     res[worldid, dofid] = acc
 
@@ -108,7 +108,7 @@ def mul_m_dense(nv: int, check_skip: bool):
   @wp.kernel(module="unique")
   def _mul_m_dense(
     # Data in:
-    qM_in: wp.array3d[float],
+    M_in: wp.array3d[float],
     # In:
     vec: wp.array2d[float],
     skip: wp.array[bool],
@@ -123,7 +123,7 @@ def mul_m_dense(nv: int, check_skip: bool):
 
     acc = float(0.0)
     for j in range(wp.static(nv)):
-      acc += qM_in[worldid, i, j] * vec[worldid, j]
+      acc += M_in[worldid, i, j] * vec[worldid, j]
     res[worldid, i] = acc
 
   return _mul_m_dense
@@ -143,8 +143,8 @@ def mul_m(
   Args:
     m: The model containing kinematic and dynamic information (device).
     d: The data object containing the current state and output arrays (device).
-    res: Result: qM @ vec.
-    vec: Input vector to multiply by qM.
+    res: Result: M @ vec.
+    vec: Input vector to multiply by M.
     skip: Per-world bitmask to skip computing output.
     M: Input matrix: M @ vec.
   """
@@ -152,13 +152,13 @@ def mul_m(
   skip = skip or wp.empty(0, dtype=bool)
 
   if M is None:
-    M = d.qM
+    M = d.M
 
   if m.is_sparse:
     wp.launch(
       mul_m_sparse(check_skip),
       dim=(d.nworld, m.nv),
-      inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
+      inputs=[m.M_mulm_rowadr, m.M_mulm_col, m.M_mulm_madr, M, vec, skip],
       outputs=[res],
     )
 
@@ -171,6 +171,340 @@ def mul_m(
     )
 
 
+@wp.kernel
+def _mul_m_island_sparse(
+  # Model:
+  M_mulm_rowadr: wp.array[int],
+  M_mulm_col: wp.array[int],
+  M_mulm_madr: wp.array[int],
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Sparse island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  start = M_mulm_rowadr[dof]
+  end = M_mulm_rowadr[dof + 1]
+  for k in range(start, end):
+    col = M_mulm_col[k]
+    madr = M_mulm_madr[k]
+    col_idof = map_dof2idof_in[worldid, col]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, 0, madr] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@wp.kernel
+def _mul_m_island_dense(
+  # Model:
+  nv: int,
+  # Data in:
+  nidof_in: wp.array[int],
+  M_in: wp.array3d[float],
+  map_dof2idof_in: wp.array2d[int],
+  map_idof2dof_in: wp.array2d[int],
+  # In:
+  idof_islandid_in: wp.array2d[int],
+  vec: wp.array2d[float],
+  island_done_in: wp.array2d[bool],
+  check_skip: int,
+  # Out:
+  res: wp.array2d[float],
+):
+  """Dense island mul_m for ALL islands in parallel."""
+  worldid, idofid = wp.tid()
+
+  nidof = nidof_in[worldid]
+  if idofid >= nidof:
+    return
+
+  islandid = idof_islandid_in[worldid, idofid]
+  if islandid < 0:
+    return
+
+  if check_skip:
+    if island_done_in[worldid, islandid]:
+      return
+
+  dof = map_idof2dof_in[worldid, idofid]
+
+  acc = float(0.0)
+  for j in range(nv):
+    col_idof = map_dof2idof_in[worldid, j]
+    # skip unconstrained DOFs
+    if col_idof < nidof:
+      acc += M_in[worldid, dof, j] * vec[worldid, col_idof]
+
+  res[worldid, idofid] = acc
+
+
+@event_scope
+def mul_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+  map_dof2idof: wp.array2d[int],
+  idof_islandid: wp.array2d[int],
+  island_done: Optional[wp.array] = None,
+  M: Optional[wp.array] = None,
+):
+  """Multiply island-local vectors by inertia matrix for all islands in parallel.
+
+  Args:
+    m: The model containing kinematic and dynamic information.
+    d: The data object containing the current state and output arrays.
+    res: Result: qM @ vec (island-local DOF order).
+    vec: Input vector (island-local DOF order).
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF → global DOF map.
+    map_dof2idof: Global DOF → island-local DOF map.
+    idof_islandid: Island ID per island-local DOF.
+    island_done: Per-island done flags (nworld, ntree).
+    M: Optional mass matrix override.
+  """
+  check_skip = int(island_done is not None)
+  island_done = island_done or wp.empty((0, 0), dtype=bool)
+
+  if M is None:
+    M = d.M
+
+  if m.is_sparse:
+    wp.launch(
+      _mul_m_island_sparse,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.M_mulm_rowadr,
+        m.M_mulm_col,
+        m.M_mulm_madr,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+  else:
+    wp.launch(
+      _mul_m_island_dense,
+      dim=(d.nworld, m.nv),
+      inputs=[
+        m.nv,
+        nidof,
+        M,
+        map_dof2idof,
+        map_idof2dof,
+        idof_islandid,
+        vec,
+        island_done,
+        check_skip,
+      ],
+      outputs=[res],
+    )
+
+
+@cache_kernel
+def _solve_LD_sparse_island(nv: int, nlevels: int):
+  """Sparse backsubstitution with island-local index remapping.
+
+  Same algorithm as _solve_LD_sparse_fused, but reads/writes x/y in
+  island-local DOF order. The L/D factorization stays in global DOF order.
+  """
+
+  @wp.func_native(snippet="WP_TILE_SYNC();")
+  def _syncthreads():
+    pass
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def kernel(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    map_idof2dof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    D: wp.array2d[float],
+    all_updates: wp.array[wp.vec3i],
+    level_offsets: wp.array[int],
+    y: wp.array2d[float],
+    # Out:
+    x_out: wp.array2d[float],
+  ):
+    worldid, tid = wp.tid()
+    NLEVELS = wp.static(nlevels)
+    BLOCK_DIM = wp.block_dim()
+    nid = nidof_in[worldid]
+
+    # Copy y to x_out (island-local, only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      x_out[worldid, idof] = y[worldid, idof]
+    _syncthreads()
+
+    # Forward substitution
+    for level in range(NLEVELS):
+      level_idx = NLEVELS - 1 - level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_i = map_dof2idof_in[worldid, i]
+        if idof_i < nid:
+          idof_k = map_dof2idof_in[worldid, k]
+          wp.atomic_sub(x_out[worldid], idof_i, L[worldid, 0, Madr_ki] * x_out[worldid, idof_k])
+      _syncthreads()
+
+    # Diagonal multiply (only iterate up to nid)
+    for idof in range(tid, nid, BLOCK_DIM):
+      dofid = map_idof2dof_in[worldid, idof]
+      x_out[worldid, idof] *= D[worldid, dofid]
+    _syncthreads()
+
+    # Backward substitution
+    for level in range(NLEVELS):
+      level_idx = level
+      level_offset = level_offsets[level_idx]
+      level_size = level_offsets[level_idx + 1] - level_offset
+
+      for u in range(tid, level_size, BLOCK_DIM):
+        update = all_updates[level_offset + u]
+        i, k, Madr_ki = update[0], update[1], update[2]
+        idof_k = map_dof2idof_in[worldid, k]
+        if idof_k < nid:
+          idof_i = map_dof2idof_in[worldid, i]
+          wp.atomic_sub(x_out[worldid], idof_k, L[worldid, 0, Madr_ki] * x_out[worldid, idof_i])
+      _syncthreads()
+
+  return kernel
+
+
+@cache_kernel
+def _tile_cholesky_solve_island(tile):
+  """Dense Cholesky backsubstitution with island-local index remapping.
+
+  L is loaded from global DOF offsets (factorization unchanged).
+  y/x are loaded/stored at island-local DOF offsets via map_dof2idof.
+  """
+
+  @wp.kernel(module="unique", enable_backward=False)
+  def cholesky_solve(
+    # Data in:
+    nidof_in: wp.array[int],
+    map_dof2idof_in: wp.array2d[int],
+    # In:
+    L: wp.array3d[float],
+    y: wp.array2d[float],
+    adr: wp.array[int],
+    # Out:
+    x: wp.array2d[float],
+  ):
+    worldid, nodeid = wp.tid()
+    TILE_SIZE = wp.static(tile.size)
+
+    dofid = adr[nodeid]
+    idofid = map_dof2idof_in[worldid, dofid]
+
+    # Skip unconstrained trees (uniform branch — all threads in block agree)
+    if idofid >= nidof_in[worldid]:
+      return
+
+    # L stays in global order
+    L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
+    # y and x use island-local offsets
+    y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(idofid,))
+    x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
+    wp.tile_store(x[worldid], x_slice, offset=(idofid,))
+
+  return cholesky_solve
+
+
+@event_scope
+def solve_m_island(
+  m: Model,
+  d: Data,
+  res: wp.array2d[float],
+  vec: wp.array2d[float],
+  nidof: wp.array[int],
+  map_idof2dof: wp.array2d[int],
+):
+  """Compute res = M^{-1} @ vec for island-local DOFs.
+
+  Args:
+    m: Model.
+    d: Data.
+    res: Output in island-local DOF order.
+    vec: Input in island-local DOF order.
+    nidof: Number of island DOFs per world.
+    map_idof2dof: Island-local DOF -> global DOF map.
+  """
+  if m.is_sparse:
+    nlevels = len(m.qLD_updates)
+    if wp.get_device().is_cuda:
+      dim_block = m.block_dim.solve_LD_sparse_fused
+    else:
+      dim_block = 1
+
+    wp.launch(
+      _solve_LD_sparse_island(m.nv, nlevels),
+      dim=(d.nworld, dim_block),
+      inputs=[
+        d.nidof,
+        d.map_dof2idof,
+        map_idof2dof,
+        d.qLD,
+        d.qLDiagInv,
+        m.qLD_all_updates,
+        m.qLD_level_offsets,
+        vec,
+      ],
+      outputs=[res],
+      block_dim=dim_block,
+    )
+  else:
+    for tile in m.M_tiles:
+      wp.launch_tiled(
+        _tile_cholesky_solve_island(tile),
+        dim=(d.nworld, tile.adr.size),
+        inputs=[d.nidof, d.map_dof2idof, d.qLD, vec, tile.adr],
+        outputs=[res],
+        block_dim=m.block_dim.cholesky_solve,
+      )
+
+
 @wp.kernel
 def _apply_ft(
   # Model:
@@ -604,11 +938,13 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # Data in:
     time_in: wp.array[float],
     qpos_in: wp.array2d[float],
     qvel_in: wp.array2d[float],
     act_in: wp.array2d[float],
+    history_in: wp.array2d[float],
     qacc_warmstart_in: wp.array2d[float],
     ctrl_in: wp.array2d[float],
     qfrc_applied_in: wp.array2d[float],
@@ -647,6 +983,10 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             state_out[worldid, adr + j] = act_in[worldid, j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            state_out[worldid, adr + j] = history_in[worldid, j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             state_out[worldid, adr + j] = qacc_warmstart_in[worldid, j]
@@ -700,10 +1040,12 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       d.time,
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,
@@ -743,6 +1085,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     nbody: int,
     neq: int,
     nmocap: int,
+    nhistory: int,
     # In:
     sig_in: int,
     active_in: wp.array[bool],
@@ -752,6 +1095,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
     qpos_out: wp.array2d[float],
     qvel_out: wp.array2d[float],
     act_out: wp.array2d[float],
+    history_out: wp.array2d[float],
     qacc_warmstart_out: wp.array2d[float],
     ctrl_out: wp.array2d[float],
     qfrc_applied_out: wp.array2d[float],
@@ -785,6 +1129,10 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
           for j in range(na):
             act_out[worldid, j] = state_in[worldid, adr + j]
           adr += na
+        elif element == State.HISTORY:
+          for j in range(nhistory):
+            history_out[worldid, j] = state_in[worldid, adr + j]
+          adr += nhistory
         elif element == State.WARMSTART:
           for j in range(nv):
             qacc_warmstart_out[worldid, j] = state_in[worldid, adr + j]
@@ -844,6 +1192,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       m.nbody,
       m.neq,
       m.nmocap,
+      m.nhistory,
       int(sig),
       active or wp.ones(d.nworld, dtype=bool),
       state,
@@ -853,6 +1202,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
       d.qpos,
       d.qvel,
       d.act,
+      d.history,
       d.qacc_warmstart,
       d.ctrl,
       d.qfrc_applied,
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
index 01c974c0..77c6cca2 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
@@ -17,18 +17,19 @@ import enum
 from typing import Callable
 
 import mujoco
-from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
 import numpy as np
 import warp as wp
 
+from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
+
 MJ_MINVAL = mujoco.mjMINVAL
 MJ_MAXVAL = mujoco.mjMAXVAL
 MJ_MINIMP = mujoco.mjMINIMP  # minimum constraint impedance
 MJ_MAXIMP = mujoco.mjMAXIMP  # maximum constraint impedance
 MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
 MJ_MINMU = mujoco.mjMINMU  # minimum friction
-# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
-_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
+NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
+TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
 # maximum size (by number of edges) of an horizon in EPA algorithm
 MJ_MAX_EPAHORIZON = 24
 # maximum average number of trianglarfaces EPA can insert at each iteration
@@ -47,6 +48,25 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   # collision_driver
@@ -131,13 +151,8 @@ class ProjectionType(enum.IntEnum):
     ORTHOGRAPHIC: orthographic projection
   """
 
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if hasattr(mujoco, "mjtProjection"):
-    PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
-    ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
-  else:
-    PERSPECTIVE = 0
-    ORTHOGRAPHIC = 1
+  PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
+  ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
 
 
 class Stage(enum.IntEnum):
@@ -188,6 +203,7 @@ class DisableBit(enum.IntFlag):
     EULERDAMP:    implicit damping for Euler integration
     NATIVECCD:    native convex collision detection (ignored in MJWarp)
     ISLAND:       constraint islands
+    MULTICCD:     disable multiple contacts with CCD
   """
 
   CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -217,7 +233,6 @@ class EnableBit(enum.IntFlag):
   Attributes:
     ENERGY: energy computation
     INVDISCRETE: discrete-time inverse dynamics
-    MULTICCD: multiple contacts with CCD
   """
 
   ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
@@ -338,12 +353,13 @@ class IntegratorType(enum.IntEnum):
     EULER: semi-implicit Euler
     RK4: 4th-order Runge Kutta
     IMPLICITFAST: implicit in velocity, no rne derivative
+    IMPLICIT: implicit in velocity, with rne derivative
   """
 
   EULER = mujoco.mjtIntegrator.mjINT_EULER
   RK4 = mujoco.mjtIntegrator.mjINT_RK4
   IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
-  # unsupported: IMPLICIT
+  IMPLICIT = mujoco.mjtIntegrator.mjINT_IMPLICIT
 
 
 class GeomType(enum.IntEnum):
@@ -617,6 +633,7 @@ class State(enum.IntEnum):
     QPOS: position
     QVEL: velocity
     ACT: actuator activation
+    HISTORY: delay/interval history buffers
     WARMSTART: acceleration used for warmstart
     CTRL: control
     QFRC_APPLIED: applied generalized force
@@ -625,7 +642,7 @@ class State(enum.IntEnum):
     MOCAP_POS: positions of mocap bodies
     MOCAP_QUAT: orientations of mocap bodies
     NSTATE: number of state elements
-    PHYSICS: QPOS | QVEL | ACT
+    PHYSICS: TIME | QPOS | QVEL | ACT | HISTORY
     FULLPHYSICS: TIME | PHYSICS | PLUGIN
     USER: CTRL | QFRC_APPLIED | XFRC_APPLIED | EQ_ACTIVE | MOCAP_POS | MOCAP_QUAT | USERDATA
     INTEGRATION: FULLPHYSICS | USER | WARMSTART
@@ -635,6 +652,7 @@ class State(enum.IntEnum):
   QPOS = mujoco.mjtState.mjSTATE_QPOS
   QVEL = mujoco.mjtState.mjSTATE_QVEL
   ACT = mujoco.mjtState.mjSTATE_ACT
+  HISTORY = mujoco.mjtState.mjSTATE_HISTORY
   WARMSTART = mujoco.mjtState.mjSTATE_WARMSTART
   CTRL = mujoco.mjtState.mjSTATE_CTRL
   QFRC_APPLIED = mujoco.mjtState.mjSTATE_QFRC_APPLIED
@@ -643,7 +661,7 @@ class State(enum.IntEnum):
   MOCAP_POS = mujoco.mjtState.mjSTATE_MOCAP_POS
   MOCAP_QUAT = mujoco.mjtState.mjSTATE_MOCAP_QUAT
   NSTATE = mujoco.mjtState.mjNSTATE
-  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS
+  PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS  # includes HISTORY
   FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS
   USER = mujoco.mjtState.mjSTATE_USER
   INTEGRATION = mujoco.mjtState.mjSTATE_INTEGRATION
@@ -861,6 +879,7 @@ class Model:
     ntree: number of kinematic trees
     nM: number of non-zeros in sparse inertia matrix
     nC: number of non-zeros in sparse body-dof matrix
+    nD: number of non-zeros in sparse derivative matrix
     ngeom: number of geoms
     nsite: number of sites
     ncam: number of cameras
@@ -898,6 +917,7 @@ class Model:
     nJmom: number of non-zeros in actuator_moment
     ngravcomp: number of bodies with nonzero gravcomp
     nsensordata: number of elements in sensor data vector
+    nhistory: number of history buffer entries
     opt: physics options
     stat: model statistics
     qpos0: qpos values at default pose                       (*, nq)
@@ -1129,6 +1149,9 @@ class Model:
     actuator_trnid: transmission id: joint, tendon, site     (nu, 2)
     actuator_actadr: first activation address; -1: stateless (nu,)
     actuator_actnum: number of activation variables          (nu,)
+    actuator_history: history buffer sizes                   (nu, 2)
+    actuator_historyadr: history buffer address              (nu,)
+    actuator_delay: delay in seconds                         (nu,)
     actuator_ctrllimited: is control limited                 (nu,)
     actuator_forcelimited: is force limited                  (nu,)
     actuator_actlimited: is activation limited               (nu,)
@@ -1153,12 +1176,22 @@ class Model:
     sensor_dim: number of scalar outputs                     (nsensor,)
     sensor_adr: address in sensor array                      (nsensor,)
     sensor_cutoff: cutoff for real and positive; 0: ignore   (nsensor,)
+    sensor_history: history buffer sizes                     (nsensor, 2)
+    sensor_historyadr: history buffer address                (nsensor,)
+    sensor_delay: delay in seconds                           (nsensor,)
+    sensor_interval: sensor interval and phase               (nsensor, 2)
     plugin: globally registered plugin slot number           (nplugin,)
     plugin_attr: config attributes of geom plugin            (nplugin, _NPLUGINATTR)
-    M_rownnz: number of non-zeros in each row of qM          (nv,)
-    M_rowadr: index of each row in qM                        (nv,)
-    M_colind: column indices of non-zeros in qM              (nM,)
+    M_rownnz: number of non-zeros in each row of M           (nv,)
+    M_rowadr: index of each row in M                         (nv,)
+    M_colind: column indices of non-zeros in M               (nC,)
     mapM2M: index mapping from M (legacy) to M (CSR)         (nC)
+    D_rownnz: non-zeros per row in D-structure               (nv,)
+    D_rowadr: row start addresses in D-structure             (nv,)
+    D_diag: diagonal element index within each row           (nv,)
+    D_colind: column indices in D-structure                  (nD,)
+    mapM2D: index mapping from M to D                        (nD,)
+    mapD2M: index mapping from D to M                        (nC,)
     flex_vertflexid: flex id for each flex vertex            (nflexvert,)
 
   warp only fields:
@@ -1187,8 +1220,8 @@ class Model:
     jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
     jnt_limited_ball_adr: limited/ball jntadr
     body_isdofancestor: precomputed mask of which DOFs affect each body
-    dof_tri_row: dof lower triangle row (used in solver)
-    dof_tri_col: dof lower triangle col (used in solver)
+    dof_tri_row: dof upper triangle row (used in solver)
+    dof_tri_col: dof upper triangle col (used in solver)
     nxn_geom_pair: collision pair geom ids [-2, ngeom-1]
     nxn_geom_pair_filtered: valid collision pair geom ids
                             [-1, ngeom - 1]
@@ -1243,15 +1276,21 @@ class Model:
     sensor_rangefinder_bodyid: bodyid for rangefinder        (nrangefinder,)
     taxel_vertadr: tactile sensor vertex address             (nsensortaxel,)
     taxel_sensorid: address for tactile sensors
-    qM_tiles: tiling configuration
+    M_tiles: tiling configuration
     qLD_updates: tuple of index triples for sparse factorization
     qLD_all_updates: tuple of all levels concatenated
     qLD_level_offsets: tuple of start offsets for each level
-    qM_fullm_i: sparse mass matrix addressing
-    qM_fullm_j: sparse mass matrix addressing
-    qM_mulm_rowadr: sparse matmul row pointers
-    qM_mulm_col: sparse matmul column indices
-    qM_mulm_madr: sparse matmul matrix addresses
+    M_fullm_i: sparse mass matrix addressing
+    M_fullm_j: sparse mass matrix addressing
+    M_elemid: (row, col) -> CSR madr addresses; -1 if not a chain ancestor
+    M_fullm_upper_i: upper-triangle row indices for solver h seeding
+    M_fullm_upper_j: upper-triangle column indices for solver h seeding
+    M_fullm_upper_elemid: source elemid into M_fullm_i/M_fullm_j
+    qD_fullm_i: D-structure row indices for RNE derivatives
+    qD_fullm_j: D-structure column indices for RNE derivatives
+    M_mulm_rowadr: sparse matmul row pointers
+    M_mulm_col: sparse matmul column indices
+    M_mulm_madr: sparse matmul matrix addresses
   """
 
   nq: int
@@ -1264,6 +1303,7 @@ class Model:
   ntree: int
   nM: int
   nC: int
+  nD: int
   ngeom: int
   nsite: int
   ncam: int
@@ -1301,6 +1341,7 @@ class Model:
   nJmom: int
   ngravcomp: int
   nsensordata: int
+  nhistory: int
   opt: Option
   stat: Statistic
   qpos0: array("*", "nq", float)
@@ -1532,6 +1573,9 @@ class Model:
   actuator_trnid: array("nu", wp.vec2i)
   actuator_actadr: array("nu", int)
   actuator_actnum: array("nu", int)
+  actuator_history: array("nu", wp.vec2i)
+  actuator_historyadr: array("nu", int)
+  actuator_delay: array("nu", float)
   actuator_ctrllimited: array("nu", bool)
   actuator_forcelimited: array("nu", bool)
   actuator_actlimited: array("nu", bool)
@@ -1556,12 +1600,22 @@ class Model:
   sensor_dim: array("nsensor", int)
   sensor_adr: array("nsensor", int)
   sensor_cutoff: array("nsensor", float)
+  sensor_history: array("nsensor", wp.vec2i)
+  sensor_historyadr: array("nsensor", int)
+  sensor_delay: array("nsensor", float)
+  sensor_interval: array("nsensor", wp.vec2)
   plugin: array("nplugin", int)
   plugin_attr: array("nplugin", vec_pluginattr)
   M_rownnz: array("nv", int)
   M_rowadr: array("nv", int)
   M_colind: array("nC", int)
   mapM2M: array("nC", int)
+  D_rownnz: array("nv", int)
+  D_rowadr: array("nv", int)
+  D_diag: array("nv", int)
+  D_colind: array("nD", int)
+  mapM2D: array("nD", int)
+  mapD2M: array("nC", int)
   flex_vertflexid: array("nflexvert", int)
   # warp only fields:
   callback: Callback
@@ -1635,23 +1689,32 @@ class Model:
   sensor_rangefinder_bodyid: array("nrangefinder", int)
   taxel_vertadr: array("nsensortaxel", int)
   taxel_sensorid: wp.array[int]
-  qM_tiles: tuple[TileSet, ...]
+  M_tiles: tuple[TileSet, ...]
   qLD_updates: tuple[wp.array[wp.vec3i], ...]
   qLD_all_updates: wp.array[wp.vec3i]
   qLD_level_offsets: wp.array[int]
-  qM_fullm_i: wp.array[int]
-  qM_fullm_j: wp.array[int]
+  # TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
+  # directly in the solver/derivative kernels
+  M_fullm_i: wp.array[int]
+  M_fullm_j: wp.array[int]
+  M_elemid: wp.array2d[int]  # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
+  M_fullm_upper_i: wp.array[int]
+  M_fullm_upper_j: wp.array[int]
+  M_fullm_upper_elemid: wp.array[int]
+  qD_fullm_i: wp.array[int]  # D-structure (full square) row indices for RNE derivatives
+  qD_fullm_j: wp.array[int]  # D-structure (full square) column indices for RNE derivatives
   # Gather-based sparse mul_m indices (thread per DOF, no atomics)
-  qM_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
-  qM_mulm_col: wp.array[int]  # column index to gather from
-  qM_mulm_madr: wp.array[int]  # matrix address to read
+  M_mulm_rowadr: wp.array[int]  # start address for each row [nv+1]
+  M_mulm_col: wp.array[int]  # column index to gather from
+  M_mulm_madr: wp.array[int]  # matrix address to read
 
 
 class ContactType(enum.IntFlag):
   """Type of contact.
 
-  CONSTRAINT: contact for constraint solver.
-  SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO).
+  Attributes:
+    CONSTRAINT: contact for constraint solver
+    SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO)
   """
 
   CONSTRAINT = 1
@@ -1673,6 +1736,8 @@ class Contact:
     solimp: constraint solver impedance                              (naconmax, 5)
     dim: contact space dimensionality: 1, 3, 4 or 6                  (naconmax,)
     geom: geom ids; -1 for flex                                      (naconmax, 2)
+    flex: flex ids; -1 for geom                                      (naconmax, 2)
+    vert: vertex ids for flex/mesh contact                           (naconmax, 2)
     efc_address: address in efc; -1: not included                    (naconmax, nmaxpyramid)
     worldid: world id                                                (naconmax,)
     type: ContactType                                                (naconmax,)
@@ -1711,9 +1776,9 @@ class Constraint:
     J_rowadr: row start address in colind array       (nworld, 0) dense
                                                       (nworld, njmax) sparse
     J_colind: column indices in J                     (nworld, 0, 0) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     J: constraint Jacobian                            (nworld, njmax_pad, nv_pad) dense
-                                                      (nworld, 1, njmax * nv) sparse
+                                                      (nworld, 1, njmax_nnz) sparse
     pos: constraint position (equality, contact)      (nworld, njmax)
     margin: inclusion margin (contact)                (nworld, njmax)
     D: constraint mass                                (nworld, njmax_pad)
@@ -1722,6 +1787,20 @@ class Constraint:
     frictionloss: frictionloss (friction)             (nworld, njmax)
     force: constraint force in constraint space       (nworld, njmax)
     state: constraint state                           (nworld, njmax_pad)
+    island: island ID per constraint                  (nworld, njmax)
+    itype: island constraint type                     (nworld, njmax)
+    iid: island constraint id                         (nworld, njmax)
+    iJ_rownnz: island J_rownnz                        (nworld, njmax)
+    iJ_rowadr: island J_rowadr                        (nworld, njmax)
+    iJ_colind: island J_colind                        (nworld, 0, 0) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iJ: island J                                      (nworld, njmax, nv) dense
+                                                      (nworld, 1, njmax_nnz) sparse
+    iD: island constraint mass                        (nworld, njmax_pad)
+    iaref: island aref                                (nworld, njmax)
+    ifrictionloss: island frictionloss                (nworld, njmax)
+    iforce: island force                              (nworld, njmax)
+    istate: island state                              (nworld, njmax_pad)
   warp only fields:
     Ma: M*qacc                                        (nworld, nv)
     Jqvel: J*qvel                                     (nworld, njmax)
@@ -1729,8 +1808,8 @@ class Constraint:
 
   type: array("nworld", "njmax", int)
   id: array("nworld", "njmax", int)
-  J_rownnz: wp.array2d[int]
-  J_rowadr: wp.array2d[int]
+  J_rownnz: array("nworld", "njmax", int)
+  J_rowadr: array("nworld", "njmax", int)
   J_colind: wp.array3d[int]
   J: wp.array3d[float]
   pos: array("nworld", "njmax", float)
@@ -1741,9 +1820,22 @@ class Constraint:
   frictionloss: array("nworld", "njmax", float)
   force: array("nworld", "njmax", float)
   state: array("nworld", "njmax_pad", int)
+  island: array("nworld", "njmax", int)
   Ma: array("nworld", "nv", float)
   Jqvel: array("nworld", "njmax", float)
 
+  itype: array("nworld", "njmax", int)
+  iid: array("nworld", "njmax", int)
+  iJ_rownnz: array("nworld", "njmax", int)
+  iJ_rowadr: array("nworld", "njmax", int)
+  iJ_colind: wp.array3d[int]
+  iJ: wp.array3d[float]
+  iD: array("nworld", "njmax_pad", float)
+  iaref: array("nworld", "njmax", float)
+  ifrictionloss: array("nworld", "njmax", float)
+  iforce: array("nworld", "njmax", float)
+  istate: array("nworld", "njmax_pad", int)
+
 
 @dataclasses.dataclass
 class Data:
@@ -1756,11 +1848,13 @@ class Data:
     nl: number of limit constraints                             (nworld,)
     nefc: number of constraints                                 (nworld,)
     nisland: number of constraint islands                       (nworld,)
+    nidof: total DOFs in islands                                (nworld,)
     time: simulation time                                       (nworld,)
     energy: potential, kinetic energy                           (nworld, 2)
     qpos: position                                              (nworld, nq)
     qvel: velocity                                              (nworld, nv)
     act: actuator activation                                    (nworld, na)
+    history: history buffer for delays                          (nworld, nhistory)
     qacc_warmstart: acceleration used for warmstart             (nworld, nv)
     ctrl: control                                               (nworld, nu)
     qfrc_applied: applied generalized force                     (nworld, nv)
@@ -1791,7 +1885,7 @@ class Data:
     cinert: com-based body inertia and mass                     (nworld, nbody, 10)
     flexvert_xpos: cartesian flex vertex positions              (nworld, nflexvert, 3)
     flexedge_J: edge length Jacobian                            (nworld, nJfe)
-    flexedge_length: flex edge lengths                          (nworld, nflexedge, 1)
+    flexedge_length: flex edge lengths                          (nworld, nflexedge)
     ten_wrapadr: start address of tendon's path                 (nworld, ntendon)
     ten_wrapnum: number of wrap points in path                  (nworld, ntendon)
     ten_J: tendon Jacobian                                      (nworld, nJten)
@@ -1804,10 +1898,10 @@ class Data:
     moment_colind: column indices in sparse actuator_moment     (nworld, nJmom)
     actuator_moment: actuator moments                           (nworld, nJmom)
     crb: com-based composite inertia and mass                   (nworld, nbody, 10)
-    qM: total inertia                                           (nworld, nv, nv) if dense
-                                                                (nworld, 1, nM) if sparse
-    qLD: L'*D*L factorization of M                              (nworld, nv, nv) if dense
+    M: total inertia                                            (nworld, nv_pad, nv_pad) if dense
                                                                 (nworld, 1, nC) if sparse
+    qLD: upper Cholesky factorization                           (nworld, nv, nv) if dense
+         L'*D*L factorization of M                              (nworld, 1, nC) if sparse
     qLDiagInv: 1/diag(D)                                        (nworld, nv)
     flexedge_velocity: flex edge velocities                     (nworld, nflexedge)
     ten_velocity: tendon velocities                             (nworld, ntendon)
@@ -1822,6 +1916,7 @@ class Data:
     qfrc_passive: total passive force                           (nworld, nv)
     subtree_linvel: linear velocity of subtree com              (nworld, nbody, 3)
     subtree_angmom: angular momentum about subtree com          (nworld, nbody, 3)
+    qLU: sparse LU factorization of (M - dt*qDeriv)             (nworld, 1, nD)
     actuator_force: actuator force in actuation space           (nworld, nu)
     qfrc_actuator: actuator force                               (nworld, nv)
     qfrc_smooth: net unconstrained force                        (nworld, nv)
@@ -1836,6 +1931,23 @@ class Data:
     contact: contact data
     efc: constraint data
     tree_island: island ID per tree (-1 if unconstrained)       (nworld, ntree)
+    dof_island: island ID per DOF (-1 if unconstrained)         (nworld, nv)
+    island_dofadr: island start address in dof vector           (nworld, ntree)
+    island_nv: DOFs per island                                  (nworld, ntree)
+    island_nefc: constraints per island                         (nworld, ntree)
+    island_ne: equality constraints per island                  (nworld, ntree)
+    island_nf: friction constraints per island                  (nworld, ntree)
+    island_efcadr: island start address in efc vector           (nworld, ntree)
+    map_dof2idof: global DOF -> island-local DOF                (nworld, nv)
+    map_idof2dof: island-local DOF -> global DOF                (nworld, nv)
+    map_efc2iefc: global EFC -> island-local EFC                (nworld, njmax)
+    map_iefc2efc: island-local EFC -> global EFC                (nworld, njmax)
+    dof_islandid: island ID per island-DOF                      (nworld, nv)
+    efc_islandid: island ID per island-EFC                      (nworld, njmax)
+    iqacc: island-local qacc                                    (nworld, nv)
+    iqacc_smooth: island-local qacc_smooth                      (nworld, nv)
+    iqfrc_smooth: island-local qfrc_smooth                      (nworld, nv)
+    iqfrc_constraint: island-local qfrc_constraint              (nworld, nv)
 
   warp only fields:
     nworld: number of worlds
@@ -1854,11 +1966,13 @@ class Data:
   nl: array("nworld", int)
   nefc: array("nworld", int)
   nisland: array("nworld", int)
+  nidof: array("nworld", int)
   time: array("nworld", float)
   energy: array("nworld", wp.vec2)
   qpos: array("nworld", "nq", float)
   qvel: array("nworld", "nv", float)
   act: array("nworld", "na", float)
+  history: array("nworld", "nhistory", float)
   qacc_warmstart: array("nworld", "nv", float)
   ctrl: array("nworld", "nu", float)
   qfrc_applied: array("nworld", "nv", float)
@@ -1902,7 +2016,7 @@ class Data:
   moment_colind: array("nworld", "nJmom", int)
   actuator_moment: array("nworld", "nJmom", float)
   crb: array("nworld", "nbody", vec10)
-  qM: wp.array3d[float]
+  M: wp.array3d[float]
   qLD: wp.array3d[float]
   qLDiagInv: array("nworld", "nv", float)
   flexedge_velocity: array("nworld", "nflexedge", float)
@@ -1918,6 +2032,7 @@ class Data:
   qfrc_passive: array("nworld", "nv", float)
   subtree_linvel: array("nworld", "nbody", wp.vec3)
   subtree_angmom: array("nworld", "nbody", wp.vec3)
+  qLU: array("nworld", 1, "nD", float)
   actuator_force: array("nworld", "nu", float)
   qfrc_actuator: array("nworld", "nv", float)
   qfrc_smooth: array("nworld", "nv", float)
@@ -1930,6 +2045,23 @@ class Data:
   contact: Contact
   efc: Constraint
   tree_island: array("nworld", "ntree", int)
+  dof_island: array("nworld", "nv", int)
+  island_dofadr: array("nworld", "ntree", int)
+  island_nv: array("nworld", "ntree", int)
+  island_nefc: array("nworld", "ntree", int)
+  island_ne: array("nworld", "ntree", int)
+  island_nf: array("nworld", "ntree", int)
+  island_efcadr: array("nworld", "ntree", int)
+  map_dof2idof: array("nworld", "nv", int)
+  map_idof2dof: array("nworld", "nv", int)
+  map_efc2iefc: array("nworld", "njmax", int)
+  map_iefc2efc: array("nworld", "njmax", int)
+  dof_islandid: array("nworld", "nv", int)
+  efc_islandid: array("nworld", "njmax", int)
+  iqacc: wp.array2d[float]
+  iqacc_smooth: wp.array2d[float]
+  iqfrc_smooth: wp.array2d[float]
+  iqfrc_constraint: wp.array2d[float]
 
   # warp only fields:
   nworld: int
@@ -1942,6 +2074,77 @@ class Data:
   ncollision: array(1, int)
 
 
+@dataclasses.dataclass
+class InverseContext:
+  """Workspace arrays for inverse dynamics."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
+@dataclasses.dataclass
+class IslandSolverContext:
+  """Workspace arrays for island constraint solver."""
+
+  # Re-ordered workspace arrays (sized per-dof / per-constraint)
+  Jaref: wp.array2d[float]
+  jv: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  grad: wp.array2d[float]
+  Mgrad: wp.array2d[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  h: wp.array3d[float]
+
+  # Per-island solver scalars (nworld, ntree)
+  cost: wp.array2d[float]
+  prev_cost: wp.array2d[float]
+  gauss: wp.array2d[float]
+  search_dot: wp.array2d[float]
+  grad_dot: wp.array2d[float]
+  done: wp.array2d[bool]  # per-island convergence
+  solver_niter: wp.array2d[int]  # iterations per island
+  beta: wp.array2d[float]
+  alpha: wp.array2d[float]
+  Ma: wp.array2d[float]  # island-local Ma (nworld, nv)
+
+
+@dataclasses.dataclass
+class SolverContext:
+  """Workspace arrays for constraint solver."""
+
+  Jaref: wp.array2d[float]
+  search_dot: wp.array[float]
+  gauss: wp.array[float]
+  cost: wp.array[float]
+  prev_cost: wp.array[float]
+  done: wp.array[bool]
+  grad: wp.array2d[float]
+  grad_dot: wp.array[float]
+  Mgrad: wp.array2d[float]
+  search: wp.array2d[float]
+  mv: wp.array2d[float]
+  jv: wp.array2d[float]
+  quad: wp.array2d[wp.vec3]
+  quad_gauss: wp.array[wp.vec3]
+  alpha: wp.array[float]
+  prev_grad: wp.array2d[float]
+  prev_Mgrad: wp.array2d[float]
+  beta: wp.array[float]
+  h: wp.array3d[float]
+  hfactor: wp.array3d[float]
+  # Incremental Hessian update (Newton only)
+  changed_efc_ids: wp.array2d[int]
+  changed_efc_count: wp.array[int]
+
+
 @dataclasses.dataclass
 class RenderContext:
   """Context for rendering.
@@ -1952,6 +2155,8 @@ class RenderContext:
     cam_id_map: camera id map
     use_textures: whether to use textures
     use_shadows: whether to use shadows
+    use_ambient_lighting: whether to use ambient lighting
+    background_color: background color
     use_precomputed_rays: whether to use precomputed rays
     bvh_ngeom: number of geometries in the BVH
     enabled_geom_ids: enabled geometry ids
@@ -1971,7 +2176,10 @@ class RenderContext:
     flex_bvh_id: per-flex BVH ids
     flex_group_root: per-flex group roots (nworld x n_flex_bvh)
     flex_render_smooth: whether to render flex meshes smoothly
-    flex_dim: flex dimension per flex (1D/2D/3D)
+    bvh_nflexgeom: number of flex geometries in the BVH
+    flex_dim_np: flex dimension per flex (1D/2D/3D)
+    flex_geom_flexid: map from flex geom ID to flex ID
+    flex_geom_edgeid: map from flex geom ID to flex edge ID
     bvh: scene BVH
     bvh_id: scene BVH id
     lower: lower bounds
@@ -1993,6 +2201,11 @@ class RenderContext:
     render_skybox: whether to shade missed rays with the MuJoCo skybox texture
     skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
     skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
+    enable_backface_culling: drop primitive ray hits whose normal faces away
+      from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
+      mesh-ray rule. When False, the renderer reports inner-surface hits, which
+      is faster but causes a camera placed inside a geom to render that geom's
+      back wall.
   """
 
   nrender: int
@@ -2000,6 +2213,7 @@ class RenderContext:
   cam_id_map: array("ncam", int)
   use_textures: bool
   use_shadows: bool
+  use_ambient_lighting: bool
   background_color: wp.uint32
   use_precomputed_rays: bool
   render_skybox: bool
@@ -2045,3 +2259,4 @@ class RenderContext:
   render_seg: array("ncam", bool)
   znear: float
   total_rays: int
+  enable_backface_culling: bool
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
index 9510b8fd..d5860af2 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
@@ -23,7 +23,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
   """Parse a version string into comparable components.
 
   Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
-  "-foo3") are treated as local build identifiers and stripped before
+  "-newton") are treated as local build identifiers and stripped before
   parsing. Each component is wrapped in a tuple: (0, int) for numeric parts,
   (-1, str) for non-numeric. A (0, 0) sentinel is appended so that stable
   releases sort above pre-release suffixes during Python tuple comparison
@@ -32,13 +32,13 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
 
   Args:
     version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
-      "3.9.0-foo3".
+      "3.9.0-newton".
 
   Returns:
     Tuple of (type_order, value) pairs for comparison, where type_order is 0
     for integers and -1 for strings, followed by a (0, 0) sentinel.
   """
-  # Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
+  # Strip local build identifier (e.g., "3.9.0-newton" -> "3.9.0")
   version_str = version_str.split("-", 1)[0]
   parts = version_str.split(".")
   return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
diff --git a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
index 92750398..eac2e265 100644
--- a/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
+++ b/mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-warp"
-version = "3.8.0.1"
+version = "3.9.0.1"
 # TODO(team): create a distribution list
 authors = [
     {name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -31,7 +31,7 @@ dependencies = [
     "etils[epath]",
     "mujoco>=3.8.0",
     "numpy",
-    "warp-lang>=1.12",
+    "warp-lang>=1.13",
 ]
 
 [[tool.uv.index]]
diff --git a/mjx/mujoco/mjx/warp/__init__.py b/mjx/mujoco/mjx/warp/__init__.py
index cc582e4a..28ebd35f 100644
--- a/mjx/mujoco/mjx/warp/__init__.py
+++ b/mjx/mujoco/mjx/warp/__init__.py
@@ -23,6 +23,7 @@ if not typing.TYPE_CHECKING:
   warp: Any = None
   mujoco_warp: Any = None
   mjwp_types: Any = None
+  mjwp_io: Any = None
   WARP_INSTALLED: bool = False
 
   # pylint: disable=g-import-not-at-top
@@ -35,6 +36,7 @@ if not typing.TYPE_CHECKING:
   try:
     from mujoco.mjx.third_party import mujoco_warp
     from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
+    from mujoco.mjx.third_party.mujoco_warp._src import io as mjwp_io
   except (ImportError, RuntimeError) as e:
     print('Failed to import mujoco_warp:', e)
     pass
@@ -67,7 +69,11 @@ else:
     def BlockDim(self, *args, **kwargs):  # pylint: disable=invalid-name
       pass
 
+  class _MjwpIoStub:
+    ENABLE_ISLANDS = True
+
   WARP_INSTALLED: bool = True
   warp: Any = _WpStub()
   mujoco_warp: Any = _MjwpStub()
   mjwp_types: Any = _MjwpTypesStub()
+  mjwp_io: Any = _MjwpIoStub()
diff --git a/mjx/mujoco/mjx/warp/forward.py b/mjx/mujoco/mjx/warp/forward.py
index 52a84999..8bae7621 100644
--- a/mjx/mujoco/mjx/warp/forward.py
+++ b/mjx/mujoco/mjx/warp/forward.py
@@ -53,8 +53,18 @@ _cb = mjwp_types.Callback(
 def _forward_shim(
     # Model
     nworld: int,
+    M_elemid: wp.array2d[int],
+    M_fullm_i: wp.array[int],
+    M_fullm_j: wp.array[int],
+    M_fullm_upper_elemid: wp.array[int],
+    M_fullm_upper_i: wp.array[int],
+    M_fullm_upper_j: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
     M_rowadr: wp.array[int],
     M_rownnz: wp.array[int],
+    M_tiles: tuple[mjwp_types.TileSet, ...],
     actuator_acc0: wp.array2d[float],
     actuator_actadr: wp.array[int],
     actuator_actearly: wp.array[bool],
@@ -66,6 +76,7 @@ def _forward_shim(
     actuator_cranklength: wp.array2d[float],
     actuator_ctrllimited: wp.array[bool],
     actuator_ctrlrange: wp.array2d[wp.vec2],
+    actuator_delay: wp.array[float],
     actuator_dynprm: wp.array2d[mjwp_types.vec10f],
     actuator_dyntype: wp.array[int],
     actuator_forcelimited: wp.array[bool],
@@ -73,6 +84,8 @@ def _forward_shim(
     actuator_gainprm: wp.array2d[mjwp_types.vec10f],
     actuator_gaintype: wp.array[int],
     actuator_gear: wp.array2d[wp.spatial_vector],
+    actuator_history: wp.array[wp.vec2i],
+    actuator_historyadr: wp.array[int],
     actuator_lengthrange: wp.array2d[wp.vec2],
     actuator_trnid: wp.array[wp.vec2i],
     actuator_trntype: wp.array[int],
@@ -115,7 +128,6 @@ def _forward_shim(
     cam_resolution: wp.array[wp.vec2i],
     cam_sensorsize: wp.array[wp.vec2],
     cam_targetbodyid: wp.array[int],
-    dof_Madr: wp.array[int],
     dof_armature: wp.array2d[float],
     dof_bodyid: wp.array[int],
     dof_damping: wp.array2d[float],
@@ -239,7 +251,6 @@ def _forward_shim(
     light_pos0: wp.array2d[wp.vec3],
     light_poscom0: wp.array2d[wp.vec3],
     light_targetbodyid: wp.array[int],
-    mapM2M: wp.array[int],
     mat_rgba: wp.array2d[wp.vec4],
     max_ten_J_rownnz: int,
     mesh_face: wp.array[wp.vec3i],
@@ -263,7 +274,6 @@ def _forward_shim(
     mesh_vert: wp.array[wp.vec3],
     mesh_vertadr: wp.array[int],
     mesh_vertnum: wp.array[int],
-    nC: int,
     nJten: int,
     na: int,
     nacttrnbody: int,
@@ -278,6 +288,7 @@ def _forward_shim(
     nflexvert: int,
     ngeom: int,
     ngravcomp: int,
+    nhistory: int,
     njnt: int,
     nlight: int,
     nmaxcondim: int,
@@ -314,1988 +325,6 @@ def _forward_shim(
     qLD_all_updates: wp.array[wp.vec3i],
     qLD_level_offsets: wp.array[int],
     qLD_updates: tuple[wp.array[wp.vec3i], ...],
-    qM_fullm_i: wp.array[int],
-    qM_fullm_j: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
-    qM_mulm_rowadr: wp.array[int],
-    qM_tiles: tuple[mjwp_types.TileSet, ...],
-    qpos0: wp.array2d[float],
-    qpos_spring: wp.array2d[float],
-    rangefinder_sensor_adr: wp.array[int],
-    sensor_acc_adr: wp.array[int],
-    sensor_adr: wp.array[int],
-    sensor_adr_to_contact_adr: wp.array[int],
-    sensor_contact_adr: wp.array[int],
-    sensor_cutoff: wp.array[float],
-    sensor_datatype: wp.array[int],
-    sensor_dim: wp.array[int],
-    sensor_e_kinetic: bool,
-    sensor_e_potential: bool,
-    sensor_intprm: wp.array2d[int],
-    sensor_limitfrc_adr: wp.array[int],
-    sensor_limitpos_adr: wp.array[int],
-    sensor_limitvel_adr: wp.array[int],
-    sensor_objid: wp.array[int],
-    sensor_objtype: wp.array[int],
-    sensor_pos_adr: wp.array[int],
-    sensor_rangefinder_adr: wp.array[int],
-    sensor_rangefinder_bodyid: wp.array[int],
-    sensor_refid: wp.array[int],
-    sensor_reftype: wp.array[int],
-    sensor_rne_postconstraint: bool,
-    sensor_subtree_vel: bool,
-    sensor_tendonactfrc_adr: wp.array[int],
-    sensor_touch_adr: wp.array[int],
-    sensor_type: wp.array[int],
-    sensor_vel_adr: wp.array[int],
-    site_bodyid: wp.array[int],
-    site_pos: wp.array2d[wp.vec3],
-    site_quat: wp.array2d[wp.quat],
-    site_size: wp.array[wp.vec3],
-    site_type: wp.array[int],
-    taxel_sensorid: wp.array[int],
-    taxel_vertadr: wp.array[int],
-    ten_J_colind: wp.array[int],
-    ten_J_rowadr: wp.array[int],
-    ten_J_rownnz: wp.array[int],
-    tendon_actfrclimited: wp.array[bool],
-    tendon_actfrcrange: wp.array2d[wp.vec2],
-    tendon_adr: wp.array[int],
-    tendon_armature: wp.array2d[float],
-    tendon_damping: wp.array2d[float],
-    tendon_dampingpoly: wp.array2d[wp.vec2],
-    tendon_frictionloss: wp.array2d[float],
-    tendon_geom_adr: wp.array[int],
-    tendon_invweight0: wp.array2d[float],
-    tendon_jnt_adr: wp.array[int],
-    tendon_length0: wp.array2d[float],
-    tendon_lengthspring: wp.array2d[wp.vec2],
-    tendon_limited_adr: wp.array[int],
-    tendon_margin: wp.array2d[float],
-    tendon_num: wp.array[int],
-    tendon_range: wp.array2d[wp.vec2],
-    tendon_site_pair_adr: wp.array[int],
-    tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
-    tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
-    tendon_solref_fri: wp.array2d[wp.vec2],
-    tendon_solref_lim: wp.array2d[wp.vec2],
-    tendon_stiffness: wp.array2d[float],
-    tendon_stiffnesspoly: wp.array2d[wp.vec2],
-    wrap_geom_adr: wp.array[int],
-    wrap_jnt_adr: wp.array[int],
-    wrap_objid: wp.array[int],
-    wrap_prm: wp.array[float],
-    wrap_pulley_scale: wp.array[float],
-    wrap_site_pair_adr: wp.array[int],
-    wrap_type: wp.array[int],
-    opt__broadphase: int,
-    opt__broadphase_filter: int,
-    opt__ccd_iterations: int,
-    opt__ccd_tolerance: wp.array[float],
-    opt__cone: int,
-    opt__contact_sensor_maxmatch: int,
-    opt__density: wp.array[float],
-    opt__disableflags: int,
-    opt__enableflags: int,
-    opt__graph_conditional: bool,
-    opt__gravity: wp.array[wp.vec3],
-    opt__impratio_invsqrt: wp.array[float],
-    opt__iterations: int,
-    opt__ls_iterations: int,
-    opt__ls_parallel: bool,
-    opt__ls_parallel_min_step: float,
-    opt__ls_tolerance: wp.array[float],
-    opt__magnetic: wp.array[wp.vec3],
-    opt__run_collision_detection: bool,
-    opt__sdf_initpoints: int,
-    opt__sdf_iterations: int,
-    opt__solver: int,
-    opt__timestep: wp.array[float],
-    opt__tolerance: wp.array[float],
-    opt__viscosity: wp.array[float],
-    opt__wind: wp.array[wp.vec3],
-    stat__meaninertia: wp.array[float],
-    # Data
-    naccdmax: int,
-    naconmax: int,
-    njmax: int,
-    njmax_nnz: int,
-    act: wp.array2d[float],
-    act_dot: wp.array2d[float],
-    actuator_force: wp.array2d[float],
-    actuator_length: wp.array2d[float],
-    actuator_moment: wp.array2d[float],
-    actuator_velocity: wp.array2d[float],
-    cacc: wp.array2d[wp.spatial_vector],
-    cam_xmat: wp.array2d[wp.mat33],
-    cam_xpos: wp.array2d[wp.vec3],
-    cdof: wp.array2d[wp.spatial_vector],
-    cdof_dot: wp.array2d[wp.spatial_vector],
-    cfrc_ext: wp.array2d[wp.spatial_vector],
-    cfrc_int: wp.array2d[wp.spatial_vector],
-    cinert: wp.array2d[mjwp_types.vec10],
-    crb: wp.array2d[mjwp_types.vec10],
-    ctrl: wp.array2d[float],
-    cvel: wp.array2d[wp.spatial_vector],
-    energy: wp.array[wp.vec2],
-    eq_active: wp.array2d[bool],
-    flexedge_J: wp.array2d[float],
-    flexedge_length: wp.array2d[float],
-    flexedge_velocity: wp.array2d[float],
-    flexvert_xpos: wp.array2d[wp.vec3],
-    geom_xmat: wp.array2d[wp.mat33],
-    geom_xpos: wp.array2d[wp.vec3],
-    light_xdir: wp.array2d[wp.vec3],
-    light_xpos: wp.array2d[wp.vec3],
-    mocap_pos: wp.array2d[wp.vec3],
-    mocap_quat: wp.array2d[wp.quat],
-    moment_colind: wp.array2d[int],
-    moment_rowadr: wp.array2d[int],
-    moment_rownnz: wp.array2d[int],
-    nacon: wp.array[int],
-    ncollision: wp.array[int],
-    ne: wp.array[int],
-    nefc: wp.array[int],
-    nf: wp.array[int],
-    nisland: wp.array[int],
-    nl: wp.array[int],
-    qLD: wp.array3d[float],
-    qLDiagInv: wp.array2d[float],
-    qM: wp.array3d[float],
-    qacc: wp.array2d[float],
-    qacc_smooth: wp.array2d[float],
-    qacc_warmstart: wp.array2d[float],
-    qfrc_actuator: wp.array2d[float],
-    qfrc_applied: wp.array2d[float],
-    qfrc_bias: wp.array2d[float],
-    qfrc_constraint: wp.array2d[float],
-    qfrc_damper: wp.array2d[float],
-    qfrc_fluid: wp.array2d[float],
-    qfrc_gravcomp: wp.array2d[float],
-    qfrc_passive: wp.array2d[float],
-    qfrc_smooth: wp.array2d[float],
-    qfrc_spring: wp.array2d[float],
-    qpos: wp.array2d[float],
-    qvel: wp.array2d[float],
-    sensordata: wp.array2d[float],
-    site_xmat: wp.array2d[wp.mat33],
-    site_xpos: wp.array2d[wp.vec3],
-    solver_niter: wp.array[int],
-    subtree_angmom: wp.array2d[wp.vec3],
-    subtree_com: wp.array2d[wp.vec3],
-    subtree_linvel: wp.array2d[wp.vec3],
-    ten_J: wp.array2d[float],
-    ten_length: wp.array2d[float],
-    ten_velocity: wp.array2d[float],
-    ten_wrapadr: wp.array2d[int],
-    ten_wrapnum: wp.array2d[int],
-    time: wp.array[float],
-    tree_island: wp.array2d[int],
-    wrap_obj: wp.array2d[wp.vec2i],
-    wrap_xpos: wp.array2d[wp.spatial_vector],
-    xanchor: wp.array2d[wp.vec3],
-    xaxis: wp.array2d[wp.vec3],
-    xfrc_applied: wp.array2d[wp.spatial_vector],
-    ximat: wp.array2d[wp.mat33],
-    xipos: wp.array2d[wp.vec3],
-    xmat: wp.array2d[wp.mat33],
-    xpos: wp.array2d[wp.vec3],
-    xquat: wp.array2d[wp.quat],
-    contact__dim: wp.array[int],
-    contact__dist: wp.array[float],
-    contact__efc_address: wp.array2d[int],
-    contact__flex: wp.array[wp.vec2i],
-    contact__frame: wp.array[wp.mat33],
-    contact__friction: wp.array[mjwp_types.vec5],
-    contact__geom: wp.array[wp.vec2i],
-    contact__geomcollisionid: wp.array[int],
-    contact__includemargin: wp.array[float],
-    contact__pos: wp.array[wp.vec3],
-    contact__solimp: wp.array[mjwp_types.vec5],
-    contact__solref: wp.array[wp.vec2],
-    contact__solreffriction: wp.array[wp.vec2],
-    contact__type: wp.array[int],
-    contact__vert: wp.array[wp.vec2i],
-    contact__worldid: wp.array[int],
-    efc__D: wp.array2d[float],
-    efc__J: wp.array3d[float],
-    efc__J_colind: wp.array3d[int],
-    efc__J_rowadr: wp.array2d[int],
-    efc__J_rownnz: wp.array2d[int],
-    efc__Jqvel: wp.array2d[float],
-    efc__Ma: wp.array2d[float],
-    efc__aref: wp.array2d[float],
-    efc__force: wp.array2d[float],
-    efc__frictionloss: wp.array2d[float],
-    efc__id: wp.array2d[int],
-    efc__margin: wp.array2d[float],
-    efc__pos: wp.array2d[float],
-    efc__state: wp.array2d[int],
-    efc__type: wp.array2d[int],
-    efc__vel: wp.array2d[float],
-):
-  _m.stat = _s
-  _m.opt = _o
-  _m.callback = _cb
-  _d.efc = _e
-  _d.contact = _c
-  _m.M_rowadr = M_rowadr
-  _m.M_rownnz = M_rownnz
-  _m.actuator_acc0 = actuator_acc0
-  _m.actuator_actadr = actuator_actadr
-  _m.actuator_actearly = actuator_actearly
-  _m.actuator_actlimited = actuator_actlimited
-  _m.actuator_actnum = actuator_actnum
-  _m.actuator_actrange = actuator_actrange
-  _m.actuator_biasprm = actuator_biasprm
-  _m.actuator_biastype = actuator_biastype
-  _m.actuator_cranklength = actuator_cranklength
-  _m.actuator_ctrllimited = actuator_ctrllimited
-  _m.actuator_ctrlrange = actuator_ctrlrange
-  _m.actuator_dynprm = actuator_dynprm
-  _m.actuator_dyntype = actuator_dyntype
-  _m.actuator_forcelimited = actuator_forcelimited
-  _m.actuator_forcerange = actuator_forcerange
-  _m.actuator_gainprm = actuator_gainprm
-  _m.actuator_gaintype = actuator_gaintype
-  _m.actuator_gear = actuator_gear
-  _m.actuator_lengthrange = actuator_lengthrange
-  _m.actuator_trnid = actuator_trnid
-  _m.actuator_trntype = actuator_trntype
-  _m.actuator_trntype_body_adr = actuator_trntype_body_adr
-  _m.block_dim = block_dim
-  _m.body_branch_start = body_branch_start
-  _m.body_branches = body_branches
-  _m.body_dofadr = body_dofadr
-  _m.body_dofnum = body_dofnum
-  _m.body_fluid_ellipsoid = body_fluid_ellipsoid
-  _m.body_geomadr = body_geomadr
-  _m.body_geomnum = body_geomnum
-  _m.body_gravcomp = body_gravcomp
-  _m.body_inertia = body_inertia
-  _m.body_invweight0 = body_invweight0
-  _m.body_ipos = body_ipos
-  _m.body_iquat = body_iquat
-  _m.body_isdofancestor = body_isdofancestor
-  _m.body_jntadr = body_jntadr
-  _m.body_jntnum = body_jntnum
-  _m.body_mass = body_mass
-  _m.body_mocapid = body_mocapid
-  _m.body_parentid = body_parentid
-  _m.body_pos = body_pos
-  _m.body_quat = body_quat
-  _m.body_rootid = body_rootid
-  _m.body_subtreemass = body_subtreemass
-  _m.body_tree = body_tree
-  _m.body_treeid = body_treeid
-  _m.body_weldid = body_weldid
-  _m.cam_bodyid = cam_bodyid
-  _m.cam_fovy = cam_fovy
-  _m.cam_intrinsic = cam_intrinsic
-  _m.cam_mat0 = cam_mat0
-  _m.cam_mode = cam_mode
-  _m.cam_pos = cam_pos
-  _m.cam_pos0 = cam_pos0
-  _m.cam_poscom0 = cam_poscom0
-  _m.cam_quat = cam_quat
-  _m.cam_resolution = cam_resolution
-  _m.cam_sensorsize = cam_sensorsize
-  _m.cam_targetbodyid = cam_targetbodyid
-  _m.dof_Madr = dof_Madr
-  _m.dof_armature = dof_armature
-  _m.dof_bodyid = dof_bodyid
-  _m.dof_damping = dof_damping
-  _m.dof_dampingpoly = dof_dampingpoly
-  _m.dof_frictionloss = dof_frictionloss
-  _m.dof_invweight0 = dof_invweight0
-  _m.dof_jntid = dof_jntid
-  _m.dof_parentid = dof_parentid
-  _m.dof_solimp = dof_solimp
-  _m.dof_solref = dof_solref
-  _m.dof_treeid = dof_treeid
-  _m.dof_tri_col = dof_tri_col
-  _m.dof_tri_row = dof_tri_row
-  _m.eq_connect_adr = eq_connect_adr
-  _m.eq_data = eq_data
-  _m.eq_flex_adr = eq_flex_adr
-  _m.eq_jnt_adr = eq_jnt_adr
-  _m.eq_obj1id = eq_obj1id
-  _m.eq_obj2id = eq_obj2id
-  _m.eq_objtype = eq_objtype
-  _m.eq_solimp = eq_solimp
-  _m.eq_solref = eq_solref
-  _m.eq_ten_adr = eq_ten_adr
-  _m.eq_type = eq_type
-  _m.eq_wld_adr = eq_wld_adr
-  _m.flex_bending = flex_bending
-  _m.flex_bendingadr = flex_bendingadr
-  _m.flex_centered = flex_centered
-  _m.flex_conaffinity = flex_conaffinity
-  _m.flex_condim = flex_condim
-  _m.flex_contype = flex_contype
-  _m.flex_damping = flex_damping
-  _m.flex_dim = flex_dim
-  _m.flex_edge = flex_edge
-  _m.flex_edgeadr = flex_edgeadr
-  _m.flex_edgeflap = flex_edgeflap
-  _m.flex_edgenum = flex_edgenum
-  _m.flex_elem = flex_elem
-  _m.flex_elemadr = flex_elemadr
-  _m.flex_elemdataadr = flex_elemdataadr
-  _m.flex_elemedge = flex_elemedge
-  _m.flex_elemedgeadr = flex_elemedgeadr
-  _m.flex_elemnum = flex_elemnum
-  _m.flex_friction = flex_friction
-  _m.flex_gap = flex_gap
-  _m.flex_margin = flex_margin
-  _m.flex_priority = flex_priority
-  _m.flex_radius = flex_radius
-  _m.flex_shell = flex_shell
-  _m.flex_shelldataadr = flex_shelldataadr
-  _m.flex_shellnum = flex_shellnum
-  _m.flex_solimp = flex_solimp
-  _m.flex_solmix = flex_solmix
-  _m.flex_solref = flex_solref
-  _m.flex_stiffness = flex_stiffness
-  _m.flex_stiffnessadr = flex_stiffnessadr
-  _m.flex_vert = flex_vert
-  _m.flex_vertadr = flex_vertadr
-  _m.flex_vertbodyid = flex_vertbodyid
-  _m.flex_vertflexid = flex_vertflexid
-  _m.flex_vertnum = flex_vertnum
-  _m.flexedge_J_colind = flexedge_J_colind
-  _m.flexedge_J_rowadr = flexedge_J_rowadr
-  _m.flexedge_J_rownnz = flexedge_J_rownnz
-  _m.flexedge_invweight0 = flexedge_invweight0
-  _m.flexedge_length0 = flexedge_length0
-  _m.geom_aabb = geom_aabb
-  _m.geom_bodyid = geom_bodyid
-  _m.geom_conaffinity = geom_conaffinity
-  _m.geom_condim = geom_condim
-  _m.geom_contype = geom_contype
-  _m.geom_dataid = geom_dataid
-  _m.geom_fluid = geom_fluid
-  _m.geom_friction = geom_friction
-  _m.geom_gap = geom_gap
-  _m.geom_group = geom_group
-  _m.geom_margin = geom_margin
-  _m.geom_matid = geom_matid
-  _m.geom_pair_type_count = geom_pair_type_count
-  _m.geom_plugin_index = geom_plugin_index
-  _m.geom_pos = geom_pos
-  _m.geom_priority = geom_priority
-  _m.geom_quat = geom_quat
-  _m.geom_rbound = geom_rbound
-  _m.geom_rgba = geom_rgba
-  _m.geom_size = geom_size
-  _m.geom_solimp = geom_solimp
-  _m.geom_solmix = geom_solmix
-  _m.geom_solref = geom_solref
-  _m.geom_type = geom_type
-  _m.has_fluid = has_fluid
-  _m.has_sdf_geom = has_sdf_geom
-  _m.hfield_adr = hfield_adr
-  _m.hfield_data = hfield_data
-  _m.hfield_ncol = hfield_ncol
-  _m.hfield_nrow = hfield_nrow
-  _m.hfield_size = hfield_size
-  _m.is_sparse = is_sparse
-  _m.jnt_actfrclimited = jnt_actfrclimited
-  _m.jnt_actfrcrange = jnt_actfrcrange
-  _m.jnt_actgravcomp = jnt_actgravcomp
-  _m.jnt_axis = jnt_axis
-  _m.jnt_bodyid = jnt_bodyid
-  _m.jnt_dofadr = jnt_dofadr
-  _m.jnt_limited_ball_adr = jnt_limited_ball_adr
-  _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
-  _m.jnt_margin = jnt_margin
-  _m.jnt_pos = jnt_pos
-  _m.jnt_qposadr = jnt_qposadr
-  _m.jnt_range = jnt_range
-  _m.jnt_solimp = jnt_solimp
-  _m.jnt_solref = jnt_solref
-  _m.jnt_stiffness = jnt_stiffness
-  _m.jnt_stiffnesspoly = jnt_stiffnesspoly
-  _m.jnt_type = jnt_type
-  _m.light_bodyid = light_bodyid
-  _m.light_dir = light_dir
-  _m.light_dir0 = light_dir0
-  _m.light_mode = light_mode
-  _m.light_pos = light_pos
-  _m.light_pos0 = light_pos0
-  _m.light_poscom0 = light_poscom0
-  _m.light_targetbodyid = light_targetbodyid
-  _m.mapM2M = mapM2M
-  _m.mat_rgba = mat_rgba
-  _m.max_ten_J_rownnz = max_ten_J_rownnz
-  _m.mesh_face = mesh_face
-  _m.mesh_faceadr = mesh_faceadr
-  _m.mesh_graph = mesh_graph
-  _m.mesh_graphadr = mesh_graphadr
-  _m.mesh_normal = mesh_normal
-  _m.mesh_normaladr = mesh_normaladr
-  _m.mesh_normalnum = mesh_normalnum
-  _m.mesh_octadr = mesh_octadr
-  _m.mesh_polyadr = mesh_polyadr
-  _m.mesh_polymap = mesh_polymap
-  _m.mesh_polymapadr = mesh_polymapadr
-  _m.mesh_polymapnum = mesh_polymapnum
-  _m.mesh_polynormal = mesh_polynormal
-  _m.mesh_polynum = mesh_polynum
-  _m.mesh_polyvert = mesh_polyvert
-  _m.mesh_polyvertadr = mesh_polyvertadr
-  _m.mesh_polyvertnum = mesh_polyvertnum
-  _m.mesh_quat = mesh_quat
-  _m.mesh_vert = mesh_vert
-  _m.mesh_vertadr = mesh_vertadr
-  _m.mesh_vertnum = mesh_vertnum
-  _m.nC = nC
-  _m.nJten = nJten
-  _m.na = na
-  _m.nacttrnbody = nacttrnbody
-  _m.nbody = nbody
-  _m.nbranch = nbranch
-  _m.ncam = ncam
-  _m.neq = neq
-  _m.nflex = nflex
-  _m.nflexedge = nflexedge
-  _m.nflexelem = nflexelem
-  _m.nflexshelldata = nflexshelldata
-  _m.nflexvert = nflexvert
-  _m.ngeom = ngeom
-  _m.ngravcomp = ngravcomp
-  _m.njnt = njnt
-  _m.nlight = nlight
-  _m.nmaxcondim = nmaxcondim
-  _m.nmaxmeshdeg = nmaxmeshdeg
-  _m.nmaxpolygon = nmaxpolygon
-  _m.nmaxpyramid = nmaxpyramid
-  _m.nmeshface = nmeshface
-  _m.nrangefinder = nrangefinder
-  _m.nsensorcollision = nsensorcollision
-  _m.nsensorcontact = nsensorcontact
-  _m.nsensortaxel = nsensortaxel
-  _m.nsite = nsite
-  _m.ntendon = ntendon
-  _m.ntree = ntree
-  _m.nu = nu
-  _m.nv = nv
-  _m.nv_pad = nv_pad
-  _m.nwrap = nwrap
-  _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered
-  _m.nxn_pairid = nxn_pairid
-  _m.nxn_pairid_filtered = nxn_pairid_filtered
-  _m.oct_aabb = oct_aabb
-  _m.oct_child = oct_child
-  _m.oct_coeff = oct_coeff
-  _m.opt.broadphase = opt__broadphase
-  _m.opt.broadphase_filter = opt__broadphase_filter
-  _m.opt.ccd_iterations = opt__ccd_iterations
-  _m.opt.ccd_tolerance = opt__ccd_tolerance
-  _m.opt.cone = opt__cone
-  _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
-  _m.opt.density = opt__density
-  _m.opt.disableflags = opt__disableflags
-  _m.opt.enableflags = opt__enableflags
-  _m.opt.graph_conditional = opt__graph_conditional
-  _m.opt.gravity = opt__gravity
-  _m.opt.impratio_invsqrt = opt__impratio_invsqrt
-  _m.opt.iterations = opt__iterations
-  _m.opt.ls_iterations = opt__ls_iterations
-  _m.opt.ls_parallel = opt__ls_parallel
-  _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
-  _m.opt.ls_tolerance = opt__ls_tolerance
-  _m.opt.magnetic = opt__magnetic
-  _m.opt.run_collision_detection = opt__run_collision_detection
-  _m.opt.sdf_initpoints = opt__sdf_initpoints
-  _m.opt.sdf_iterations = opt__sdf_iterations
-  _m.opt.solver = opt__solver
-  _m.opt.timestep = opt__timestep
-  _m.opt.tolerance = opt__tolerance
-  _m.opt.viscosity = opt__viscosity
-  _m.opt.wind = opt__wind
-  _m.pair_dim = pair_dim
-  _m.pair_friction = pair_friction
-  _m.pair_gap = pair_gap
-  _m.pair_margin = pair_margin
-  _m.pair_solimp = pair_solimp
-  _m.pair_solref = pair_solref
-  _m.pair_solreffriction = pair_solreffriction
-  _m.plugin = plugin
-  _m.plugin_attr = plugin_attr
-  _m.qLD_all_updates = qLD_all_updates
-  _m.qLD_level_offsets = qLD_level_offsets
-  _m.qLD_updates = qLD_updates
-  _m.qM_fullm_i = qM_fullm_i
-  _m.qM_fullm_j = qM_fullm_j
-  _m.qM_mulm_col = qM_mulm_col
-  _m.qM_mulm_madr = qM_mulm_madr
-  _m.qM_mulm_rowadr = qM_mulm_rowadr
-  _m.qM_tiles = qM_tiles
-  _m.qpos0 = qpos0
-  _m.qpos_spring = qpos_spring
-  _m.rangefinder_sensor_adr = rangefinder_sensor_adr
-  _m.sensor_acc_adr = sensor_acc_adr
-  _m.sensor_adr = sensor_adr
-  _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
-  _m.sensor_contact_adr = sensor_contact_adr
-  _m.sensor_cutoff = sensor_cutoff
-  _m.sensor_datatype = sensor_datatype
-  _m.sensor_dim = sensor_dim
-  _m.sensor_e_kinetic = sensor_e_kinetic
-  _m.sensor_e_potential = sensor_e_potential
-  _m.sensor_intprm = sensor_intprm
-  _m.sensor_limitfrc_adr = sensor_limitfrc_adr
-  _m.sensor_limitpos_adr = sensor_limitpos_adr
-  _m.sensor_limitvel_adr = sensor_limitvel_adr
-  _m.sensor_objid = sensor_objid
-  _m.sensor_objtype = sensor_objtype
-  _m.sensor_pos_adr = sensor_pos_adr
-  _m.sensor_rangefinder_adr = sensor_rangefinder_adr
-  _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
-  _m.sensor_refid = sensor_refid
-  _m.sensor_reftype = sensor_reftype
-  _m.sensor_rne_postconstraint = sensor_rne_postconstraint
-  _m.sensor_subtree_vel = sensor_subtree_vel
-  _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
-  _m.sensor_touch_adr = sensor_touch_adr
-  _m.sensor_type = sensor_type
-  _m.sensor_vel_adr = sensor_vel_adr
-  _m.site_bodyid = site_bodyid
-  _m.site_pos = site_pos
-  _m.site_quat = site_quat
-  _m.site_size = site_size
-  _m.site_type = site_type
-  _m.stat.meaninertia = stat__meaninertia
-  _m.taxel_sensorid = taxel_sensorid
-  _m.taxel_vertadr = taxel_vertadr
-  _m.ten_J_colind = ten_J_colind
-  _m.ten_J_rowadr = ten_J_rowadr
-  _m.ten_J_rownnz = ten_J_rownnz
-  _m.tendon_actfrclimited = tendon_actfrclimited
-  _m.tendon_actfrcrange = tendon_actfrcrange
-  _m.tendon_adr = tendon_adr
-  _m.tendon_armature = tendon_armature
-  _m.tendon_damping = tendon_damping
-  _m.tendon_dampingpoly = tendon_dampingpoly
-  _m.tendon_frictionloss = tendon_frictionloss
-  _m.tendon_geom_adr = tendon_geom_adr
-  _m.tendon_invweight0 = tendon_invweight0
-  _m.tendon_jnt_adr = tendon_jnt_adr
-  _m.tendon_length0 = tendon_length0
-  _m.tendon_lengthspring = tendon_lengthspring
-  _m.tendon_limited_adr = tendon_limited_adr
-  _m.tendon_margin = tendon_margin
-  _m.tendon_num = tendon_num
-  _m.tendon_range = tendon_range
-  _m.tendon_site_pair_adr = tendon_site_pair_adr
-  _m.tendon_solimp_fri = tendon_solimp_fri
-  _m.tendon_solimp_lim = tendon_solimp_lim
-  _m.tendon_solref_fri = tendon_solref_fri
-  _m.tendon_solref_lim = tendon_solref_lim
-  _m.tendon_stiffness = tendon_stiffness
-  _m.tendon_stiffnesspoly = tendon_stiffnesspoly
-  _m.wrap_geom_adr = wrap_geom_adr
-  _m.wrap_jnt_adr = wrap_jnt_adr
-  _m.wrap_objid = wrap_objid
-  _m.wrap_prm = wrap_prm
-  _m.wrap_pulley_scale = wrap_pulley_scale
-  _m.wrap_site_pair_adr = wrap_site_pair_adr
-  _m.wrap_type = wrap_type
-  _d.act = act
-  _d.act_dot = act_dot
-  _d.actuator_force = actuator_force
-  _d.actuator_length = actuator_length
-  _d.actuator_moment = actuator_moment
-  _d.actuator_velocity = actuator_velocity
-  _d.cacc = cacc
-  _d.cam_xmat = cam_xmat
-  _d.cam_xpos = cam_xpos
-  _d.cdof = cdof
-  _d.cdof_dot = cdof_dot
-  _d.cfrc_ext = cfrc_ext
-  _d.cfrc_int = cfrc_int
-  _d.cinert = cinert
-  _d.contact.dim = contact__dim
-  _d.contact.dist = contact__dist
-  _d.contact.efc_address = contact__efc_address
-  _d.contact.flex = contact__flex
-  _d.contact.frame = contact__frame
-  _d.contact.friction = contact__friction
-  _d.contact.geom = contact__geom
-  _d.contact.geomcollisionid = contact__geomcollisionid
-  _d.contact.includemargin = contact__includemargin
-  _d.contact.pos = contact__pos
-  _d.contact.solimp = contact__solimp
-  _d.contact.solref = contact__solref
-  _d.contact.solreffriction = contact__solreffriction
-  _d.contact.type = contact__type
-  _d.contact.vert = contact__vert
-  _d.contact.worldid = contact__worldid
-  _d.crb = crb
-  _d.ctrl = ctrl
-  _d.cvel = cvel
-  _d.efc.D = efc__D
-  _d.efc.J = efc__J
-  _d.efc.J_colind = efc__J_colind
-  _d.efc.J_rowadr = efc__J_rowadr
-  _d.efc.J_rownnz = efc__J_rownnz
-  _d.efc.Jqvel = efc__Jqvel
-  _d.efc.Ma = efc__Ma
-  _d.efc.aref = efc__aref
-  _d.efc.force = efc__force
-  _d.efc.frictionloss = efc__frictionloss
-  _d.efc.id = efc__id
-  _d.efc.margin = efc__margin
-  _d.efc.pos = efc__pos
-  _d.efc.state = efc__state
-  _d.efc.type = efc__type
-  _d.efc.vel = efc__vel
-  _d.energy = energy
-  _d.eq_active = eq_active
-  _d.flexedge_J = flexedge_J
-  _d.flexedge_length = flexedge_length
-  _d.flexedge_velocity = flexedge_velocity
-  _d.flexvert_xpos = flexvert_xpos
-  _d.geom_xmat = geom_xmat
-  _d.geom_xpos = geom_xpos
-  _d.light_xdir = light_xdir
-  _d.light_xpos = light_xpos
-  _d.mocap_pos = mocap_pos
-  _d.mocap_quat = mocap_quat
-  _d.moment_colind = moment_colind
-  _d.moment_rowadr = moment_rowadr
-  _d.moment_rownnz = moment_rownnz
-  _d.naccdmax = naccdmax
-  _d.nacon = nacon
-  _d.naconmax = naconmax
-  _d.ncollision = ncollision
-  _d.ne = ne
-  _d.nefc = nefc
-  _d.nf = nf
-  _d.nisland = nisland
-  _d.njmax = njmax
-  _d.njmax_nnz = njmax_nnz
-  _d.nl = nl
-  _d.qLD = qLD
-  _d.qLDiagInv = qLDiagInv
-  _d.qM = qM
-  _d.qacc = qacc
-  _d.qacc_smooth = qacc_smooth
-  _d.qacc_warmstart = qacc_warmstart
-  _d.qfrc_actuator = qfrc_actuator
-  _d.qfrc_applied = qfrc_applied
-  _d.qfrc_bias = qfrc_bias
-  _d.qfrc_constraint = qfrc_constraint
-  _d.qfrc_damper = qfrc_damper
-  _d.qfrc_fluid = qfrc_fluid
-  _d.qfrc_gravcomp = qfrc_gravcomp
-  _d.qfrc_passive = qfrc_passive
-  _d.qfrc_smooth = qfrc_smooth
-  _d.qfrc_spring = qfrc_spring
-  _d.qpos = qpos
-  _d.qvel = qvel
-  _d.sensordata = sensordata
-  _d.site_xmat = site_xmat
-  _d.site_xpos = site_xpos
-  _d.solver_niter = solver_niter
-  _d.subtree_angmom = subtree_angmom
-  _d.subtree_com = subtree_com
-  _d.subtree_linvel = subtree_linvel
-  _d.ten_J = ten_J
-  _d.ten_length = ten_length
-  _d.ten_velocity = ten_velocity
-  _d.ten_wrapadr = ten_wrapadr
-  _d.ten_wrapnum = ten_wrapnum
-  _d.time = time
-  _d.tree_island = tree_island
-  _d.wrap_obj = wrap_obj
-  _d.wrap_xpos = wrap_xpos
-  _d.xanchor = xanchor
-  _d.xaxis = xaxis
-  _d.xfrc_applied = xfrc_applied
-  _d.ximat = ximat
-  _d.xipos = xipos
-  _d.xmat = xmat
-  _d.xpos = xpos
-  _d.xquat = xquat
-  _d.nworld = nworld
-  mjwarp.forward(_m, _d)
-
-
-def _forward_jax_impl(m: types.Model, d: types.Data):
-  output_dims = {
-      'act_dot': d.act_dot.shape,
-      'actuator_force': d.actuator_force.shape,
-      'actuator_length': d.actuator_length.shape,
-      'actuator_moment': d._impl.actuator_moment.shape,
-      'actuator_velocity': d._impl.actuator_velocity.shape,
-      'cacc': d._impl.cacc.shape,
-      'cam_xmat': d.cam_xmat.shape,
-      'cam_xpos': d.cam_xpos.shape,
-      'cdof': d.cdof.shape,
-      'cdof_dot': d.cdof_dot.shape,
-      'cfrc_ext': d._impl.cfrc_ext.shape,
-      'cfrc_int': d._impl.cfrc_int.shape,
-      'cinert': d._impl.cinert.shape,
-      'crb': d._impl.crb.shape,
-      'cvel': d.cvel.shape,
-      'energy': d._impl.energy.shape,
-      'flexedge_J': d._impl.flexedge_J.shape,
-      'flexedge_length': d._impl.flexedge_length.shape,
-      'flexedge_velocity': d._impl.flexedge_velocity.shape,
-      'flexvert_xpos': d._impl.flexvert_xpos.shape,
-      'geom_xmat': d.geom_xmat.shape,
-      'geom_xpos': d.geom_xpos.shape,
-      'light_xdir': d._impl.light_xdir.shape,
-      'light_xpos': d._impl.light_xpos.shape,
-      'moment_colind': d._impl.moment_colind.shape,
-      'moment_rowadr': d._impl.moment_rowadr.shape,
-      'moment_rownnz': d._impl.moment_rownnz.shape,
-      'nacon': d._impl.nacon.shape,
-      'ncollision': d._impl.ncollision.shape,
-      'ne': d._impl.ne.shape,
-      'nefc': d._impl.nefc.shape,
-      'nf': d._impl.nf.shape,
-      'nisland': d._impl.nisland.shape,
-      'nl': d._impl.nl.shape,
-      'qLD': d._impl.qLD.shape,
-      'qLDiagInv': d._impl.qLDiagInv.shape,
-      'qM': d._impl.qM.shape,
-      'qacc': d.qacc.shape,
-      'qacc_smooth': d.qacc_smooth.shape,
-      'qfrc_actuator': d.qfrc_actuator.shape,
-      'qfrc_bias': d.qfrc_bias.shape,
-      'qfrc_constraint': d.qfrc_constraint.shape,
-      'qfrc_damper': d._impl.qfrc_damper.shape,
-      'qfrc_fluid': d.qfrc_fluid.shape,
-      'qfrc_gravcomp': d.qfrc_gravcomp.shape,
-      'qfrc_passive': d.qfrc_passive.shape,
-      'qfrc_smooth': d.qfrc_smooth.shape,
-      'qfrc_spring': d._impl.qfrc_spring.shape,
-      'qvel': d.qvel.shape,
-      'sensordata': d.sensordata.shape,
-      'site_xmat': d.site_xmat.shape,
-      'site_xpos': d.site_xpos.shape,
-      'solver_niter': d._impl.solver_niter.shape,
-      'subtree_angmom': d._impl.subtree_angmom.shape,
-      'subtree_com': d.subtree_com.shape,
-      'subtree_linvel': d._impl.subtree_linvel.shape,
-      'ten_J': d._impl.ten_J.shape,
-      'ten_length': d.ten_length.shape,
-      'ten_velocity': d._impl.ten_velocity.shape,
-      'ten_wrapadr': d._impl.ten_wrapadr.shape,
-      'ten_wrapnum': d._impl.ten_wrapnum.shape,
-      'tree_island': d._impl.tree_island.shape,
-      'wrap_obj': d._impl.wrap_obj.shape,
-      'wrap_xpos': d._impl.wrap_xpos.shape,
-      'xanchor': d.xanchor.shape,
-      'xaxis': d.xaxis.shape,
-      'ximat': d.ximat.shape,
-      'xipos': d.xipos.shape,
-      'xmat': d.xmat.shape,
-      'xpos': d.xpos.shape,
-      'xquat': d.xquat.shape,
-      'contact__dim': d._impl.contact__dim.shape,
-      'contact__dist': d._impl.contact__dist.shape,
-      'contact__efc_address': d._impl.contact__efc_address.shape,
-      'contact__flex': d._impl.contact__flex.shape,
-      'contact__frame': d._impl.contact__frame.shape,
-      'contact__friction': d._impl.contact__friction.shape,
-      'contact__geom': d._impl.contact__geom.shape,
-      'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape,
-      'contact__includemargin': d._impl.contact__includemargin.shape,
-      'contact__pos': d._impl.contact__pos.shape,
-      'contact__solimp': d._impl.contact__solimp.shape,
-      'contact__solref': d._impl.contact__solref.shape,
-      'contact__solreffriction': d._impl.contact__solreffriction.shape,
-      'contact__type': d._impl.contact__type.shape,
-      'contact__vert': d._impl.contact__vert.shape,
-      'contact__worldid': d._impl.contact__worldid.shape,
-      'efc__D': d._impl.efc__D.shape,
-      'efc__J': d._impl.efc__J.shape,
-      'efc__J_colind': d._impl.efc__J_colind.shape,
-      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
-      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
-      'efc__Jqvel': d._impl.efc__Jqvel.shape,
-      'efc__Ma': d._impl.efc__Ma.shape,
-      'efc__aref': d._impl.efc__aref.shape,
-      'efc__force': d._impl.efc__force.shape,
-      'efc__frictionloss': d._impl.efc__frictionloss.shape,
-      'efc__id': d._impl.efc__id.shape,
-      'efc__margin': d._impl.efc__margin.shape,
-      'efc__pos': d._impl.efc__pos.shape,
-      'efc__state': d._impl.efc__state.shape,
-      'efc__type': d._impl.efc__type.shape,
-      'efc__vel': d._impl.efc__vel.shape,
-  }
-  jf = ffi.jax_callable_variadic_tuple(
-      _forward_shim,
-      num_outputs=103,
-      output_dims=output_dims,
-      vmap_method=None,
-      in_out_argnames=set([
-          'act_dot',
-          'actuator_force',
-          'actuator_length',
-          'actuator_moment',
-          'actuator_velocity',
-          'cacc',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'cfrc_ext',
-          'cfrc_int',
-          'cinert',
-          'crb',
-          'cvel',
-          'energy',
-          'flexedge_J',
-          'flexedge_length',
-          'flexedge_velocity',
-          'flexvert_xpos',
-          'geom_xmat',
-          'geom_xpos',
-          'light_xdir',
-          'light_xpos',
-          'moment_colind',
-          'moment_rowadr',
-          'moment_rownnz',
-          'nacon',
-          'ncollision',
-          'ne',
-          'nefc',
-          'nf',
-          'nisland',
-          'nl',
-          'qLD',
-          'qLDiagInv',
-          'qM',
-          'qacc',
-          'qacc_smooth',
-          'qfrc_actuator',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_damper',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qfrc_spring',
-          'qvel',
-          'sensordata',
-          'site_xmat',
-          'site_xpos',
-          'solver_niter',
-          'subtree_angmom',
-          'subtree_com',
-          'subtree_linvel',
-          'ten_J',
-          'ten_length',
-          'ten_velocity',
-          'ten_wrapadr',
-          'ten_wrapnum',
-          'tree_island',
-          'wrap_obj',
-          'wrap_xpos',
-          'xanchor',
-          'xaxis',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-          'contact__dim',
-          'contact__dist',
-          'contact__efc_address',
-          'contact__flex',
-          'contact__frame',
-          'contact__friction',
-          'contact__geom',
-          'contact__geomcollisionid',
-          'contact__includemargin',
-          'contact__pos',
-          'contact__solimp',
-          'contact__solref',
-          'contact__solreffriction',
-          'contact__type',
-          'contact__vert',
-          'contact__worldid',
-          'efc__D',
-          'efc__J',
-          'efc__J_colind',
-          'efc__J_rowadr',
-          'efc__J_rownnz',
-          'efc__Jqvel',
-          'efc__Ma',
-          'efc__aref',
-          'efc__force',
-          'efc__frictionloss',
-          'efc__id',
-          'efc__margin',
-          'efc__pos',
-          'efc__state',
-          'efc__type',
-          'efc__vel',
-      ]),
-      stage_in_argnames=set([
-          'act',
-          'act_dot',
-          'actuator_acc0',
-          'actuator_actrange',
-          'actuator_biasprm',
-          'actuator_cranklength',
-          'actuator_ctrlrange',
-          'actuator_dynprm',
-          'actuator_force',
-          'actuator_forcerange',
-          'actuator_gainprm',
-          'actuator_gear',
-          'actuator_length',
-          'actuator_lengthrange',
-          'body_gravcomp',
-          'body_inertia',
-          'body_invweight0',
-          'body_ipos',
-          'body_iquat',
-          'body_mass',
-          'body_pos',
-          'body_quat',
-          'body_subtreemass',
-          'cam_fovy',
-          'cam_intrinsic',
-          'cam_mat0',
-          'cam_pos',
-          'cam_pos0',
-          'cam_poscom0',
-          'cam_quat',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'ctrl',
-          'cvel',
-          'dof_armature',
-          'dof_damping',
-          'dof_dampingpoly',
-          'dof_frictionloss',
-          'dof_invweight0',
-          'dof_solimp',
-          'dof_solref',
-          'eq_active',
-          'eq_data',
-          'eq_solimp',
-          'eq_solref',
-          'geom_aabb',
-          'geom_friction',
-          'geom_gap',
-          'geom_margin',
-          'geom_matid',
-          'geom_pos',
-          'geom_quat',
-          'geom_rbound',
-          'geom_rgba',
-          'geom_size',
-          'geom_solimp',
-          'geom_solmix',
-          'geom_solref',
-          'geom_xmat',
-          'geom_xpos',
-          'hfield_data',
-          'jnt_actfrcrange',
-          'jnt_axis',
-          'jnt_margin',
-          'jnt_pos',
-          'jnt_range',
-          'jnt_solimp',
-          'jnt_solref',
-          'jnt_stiffness',
-          'jnt_stiffnesspoly',
-          'light_dir',
-          'light_dir0',
-          'light_pos',
-          'light_pos0',
-          'light_poscom0',
-          'mat_rgba',
-          'mocap_pos',
-          'mocap_quat',
-          'opt__density',
-          'opt__gravity',
-          'opt__ls_tolerance',
-          'opt__magnetic',
-          'opt__timestep',
-          'opt__tolerance',
-          'opt__viscosity',
-          'opt__wind',
-          'pair_friction',
-          'pair_gap',
-          'pair_margin',
-          'pair_solimp',
-          'pair_solref',
-          'pair_solreffriction',
-          'qacc',
-          'qacc_smooth',
-          'qacc_warmstart',
-          'qfrc_actuator',
-          'qfrc_applied',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qpos',
-          'qpos0',
-          'qpos_spring',
-          'qvel',
-          'sensordata',
-          'site_pos',
-          'site_quat',
-          'site_xmat',
-          'site_xpos',
-          'subtree_com',
-          'ten_length',
-          'tendon_actfrcrange',
-          'tendon_armature',
-          'tendon_damping',
-          'tendon_dampingpoly',
-          'tendon_frictionloss',
-          'tendon_invweight0',
-          'tendon_length0',
-          'tendon_lengthspring',
-          'tendon_margin',
-          'tendon_range',
-          'tendon_solimp_fri',
-          'tendon_solimp_lim',
-          'tendon_solref_fri',
-          'tendon_solref_lim',
-          'tendon_stiffness',
-          'tendon_stiffnesspoly',
-          'time',
-          'xanchor',
-          'xaxis',
-          'xfrc_applied',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-      ]),
-      stage_out_argnames=set([
-          'act_dot',
-          'actuator_force',
-          'actuator_length',
-          'cam_xmat',
-          'cam_xpos',
-          'cdof',
-          'cdof_dot',
-          'cvel',
-          'geom_xmat',
-          'geom_xpos',
-          'qacc',
-          'qacc_smooth',
-          'qfrc_actuator',
-          'qfrc_bias',
-          'qfrc_constraint',
-          'qfrc_fluid',
-          'qfrc_gravcomp',
-          'qfrc_passive',
-          'qfrc_smooth',
-          'qvel',
-          'sensordata',
-          'site_xmat',
-          'site_xpos',
-          'subtree_com',
-          'ten_length',
-          'xanchor',
-          'xaxis',
-          'ximat',
-          'xipos',
-          'xmat',
-          'xpos',
-          'xquat',
-      ]),
-      graph_mode=m.opt._impl.graph_mode,
-      has_side_effect=False,
-  )
-  out = jf(
-      d.qpos.shape[0],
-      m._impl.M_rowadr,
-      m._impl.M_rownnz,
-      m.actuator_acc0,
-      m.actuator_actadr,
-      m.actuator_actearly,
-      m.actuator_actlimited,
-      m.actuator_actnum,
-      m.actuator_actrange,
-      m.actuator_biasprm,
-      m.actuator_biastype,
-      m.actuator_cranklength,
-      m.actuator_ctrllimited,
-      m.actuator_ctrlrange,
-      m.actuator_dynprm,
-      m.actuator_dyntype,
-      m.actuator_forcelimited,
-      m.actuator_forcerange,
-      m.actuator_gainprm,
-      m.actuator_gaintype,
-      m.actuator_gear,
-      m.actuator_lengthrange,
-      m.actuator_trnid,
-      m.actuator_trntype,
-      m._impl.actuator_trntype_body_adr,
-      m._impl.block_dim,
-      m._impl.body_branch_start,
-      m._impl.body_branches,
-      m.body_dofadr,
-      m.body_dofnum,
-      m._impl.body_fluid_ellipsoid,
-      m.body_geomadr,
-      m.body_geomnum,
-      m.body_gravcomp,
-      m.body_inertia,
-      m.body_invweight0,
-      m.body_ipos,
-      m.body_iquat,
-      m._impl.body_isdofancestor,
-      m.body_jntadr,
-      m.body_jntnum,
-      m.body_mass,
-      m.body_mocapid,
-      m.body_parentid,
-      m.body_pos,
-      m.body_quat,
-      m.body_rootid,
-      m.body_subtreemass,
-      m._impl.body_tree,
-      m.body_treeid,
-      m.body_weldid,
-      m.cam_bodyid,
-      m.cam_fovy,
-      m.cam_intrinsic,
-      m.cam_mat0,
-      m.cam_mode,
-      m.cam_pos,
-      m.cam_pos0,
-      m.cam_poscom0,
-      m.cam_quat,
-      m.cam_resolution,
-      m.cam_sensorsize,
-      m.cam_targetbodyid,
-      m.dof_Madr,
-      m.dof_armature,
-      m.dof_bodyid,
-      m.dof_damping,
-      m.dof_dampingpoly,
-      m.dof_frictionloss,
-      m.dof_invweight0,
-      m.dof_jntid,
-      m.dof_parentid,
-      m.dof_solimp,
-      m.dof_solref,
-      m.dof_treeid,
-      m._impl.dof_tri_col,
-      m._impl.dof_tri_row,
-      m._impl.eq_connect_adr,
-      m.eq_data,
-      m._impl.eq_flex_adr,
-      m._impl.eq_jnt_adr,
-      m.eq_obj1id,
-      m.eq_obj2id,
-      m.eq_objtype,
-      m.eq_solimp,
-      m.eq_solref,
-      m._impl.eq_ten_adr,
-      m.eq_type,
-      m._impl.eq_wld_adr,
-      m._impl.flex_bending,
-      m._impl.flex_bendingadr,
-      m._impl.flex_centered,
-      m._impl.flex_conaffinity,
-      m._impl.flex_condim,
-      m._impl.flex_contype,
-      m._impl.flex_damping,
-      m._impl.flex_dim,
-      m._impl.flex_edge,
-      m._impl.flex_edgeadr,
-      m._impl.flex_edgeflap,
-      m._impl.flex_edgenum,
-      m._impl.flex_elem,
-      m._impl.flex_elemadr,
-      m._impl.flex_elemdataadr,
-      m._impl.flex_elemedge,
-      m._impl.flex_elemedgeadr,
-      m._impl.flex_elemnum,
-      m._impl.flex_friction,
-      m._impl.flex_gap,
-      m._impl.flex_margin,
-      m._impl.flex_priority,
-      m._impl.flex_radius,
-      m._impl.flex_shell,
-      m._impl.flex_shelldataadr,
-      m._impl.flex_shellnum,
-      m._impl.flex_solimp,
-      m._impl.flex_solmix,
-      m._impl.flex_solref,
-      m._impl.flex_stiffness,
-      m._impl.flex_stiffnessadr,
-      m._impl.flex_vert,
-      m.flex_vertadr,
-      m._impl.flex_vertbodyid,
-      m._impl.flex_vertflexid,
-      m.flex_vertnum,
-      m._impl.flexedge_J_colind,
-      m._impl.flexedge_J_rowadr,
-      m._impl.flexedge_J_rownnz,
-      m._impl.flexedge_invweight0,
-      m._impl.flexedge_length0,
-      m.geom_aabb,
-      m.geom_bodyid,
-      m.geom_conaffinity,
-      m.geom_condim,
-      m.geom_contype,
-      jax.numpy.expand_dims(m.geom_dataid, 0),
-      m.geom_fluid,
-      m.geom_friction,
-      m.geom_gap,
-      m.geom_group,
-      m.geom_margin,
-      m.geom_matid,
-      m._impl.geom_pair_type_count,
-      m._impl.geom_plugin_index,
-      m.geom_pos,
-      m.geom_priority,
-      m.geom_quat,
-      m.geom_rbound,
-      m.geom_rgba,
-      m.geom_size,
-      m.geom_solimp,
-      m.geom_solmix,
-      m.geom_solref,
-      m.geom_type,
-      m._impl.has_fluid,
-      m._impl.has_sdf_geom,
-      m.hfield_adr,
-      m.hfield_data,
-      m.hfield_ncol,
-      m.hfield_nrow,
-      m.hfield_size,
-      m._impl.is_sparse,
-      m.jnt_actfrclimited,
-      m.jnt_actfrcrange,
-      m.jnt_actgravcomp,
-      m.jnt_axis,
-      m.jnt_bodyid,
-      m.jnt_dofadr,
-      m._impl.jnt_limited_ball_adr,
-      m._impl.jnt_limited_slide_hinge_adr,
-      m.jnt_margin,
-      m.jnt_pos,
-      m.jnt_qposadr,
-      m.jnt_range,
-      m.jnt_solimp,
-      m.jnt_solref,
-      m.jnt_stiffness,
-      m.jnt_stiffnesspoly,
-      m.jnt_type,
-      m._impl.light_bodyid,
-      m.light_dir,
-      m.light_dir0,
-      m.light_mode,
-      m.light_pos,
-      m.light_pos0,
-      m.light_poscom0,
-      m._impl.light_targetbodyid,
-      m._impl.mapM2M,
-      m.mat_rgba,
-      m._impl.max_ten_J_rownnz,
-      m.mesh_face,
-      m.mesh_faceadr,
-      m.mesh_graph,
-      m.mesh_graphadr,
-      m.mesh_normal,
-      m.mesh_normaladr,
-      m.mesh_normalnum,
-      m.mesh_octadr,
-      m._impl.mesh_polyadr,
-      m._impl.mesh_polymap,
-      m._impl.mesh_polymapadr,
-      m._impl.mesh_polymapnum,
-      m._impl.mesh_polynormal,
-      m._impl.mesh_polynum,
-      m._impl.mesh_polyvert,
-      m._impl.mesh_polyvertadr,
-      m._impl.mesh_polyvertnum,
-      m.mesh_quat,
-      m.mesh_vert,
-      m.mesh_vertadr,
-      m.mesh_vertnum,
-      m.nC,
-      m.nJten,
-      m.na,
-      m._impl.nacttrnbody,
-      m.nbody,
-      m._impl.nbranch,
-      m.ncam,
-      m.neq,
-      m.nflex,
-      m._impl.nflexedge,
-      m._impl.nflexelem,
-      m._impl.nflexshelldata,
-      m._impl.nflexvert,
-      m.ngeom,
-      m.ngravcomp,
-      m.njnt,
-      m.nlight,
-      m._impl.nmaxcondim,
-      m._impl.nmaxmeshdeg,
-      m._impl.nmaxpolygon,
-      m._impl.nmaxpyramid,
-      m.nmeshface,
-      m._impl.nrangefinder,
-      m._impl.nsensorcollision,
-      m._impl.nsensorcontact,
-      m._impl.nsensortaxel,
-      m.nsite,
-      m.ntendon,
-      m._impl.ntree,
-      m.nu,
-      m.nv,
-      m._impl.nv_pad,
-      m.nwrap,
-      m._impl.nxn_geom_pair_filtered,
-      m._impl.nxn_pairid,
-      m._impl.nxn_pairid_filtered,
-      m._impl.oct_aabb,
-      m._impl.oct_child,
-      m._impl.oct_coeff,
-      m.pair_dim,
-      m.pair_friction,
-      m.pair_gap,
-      m.pair_margin,
-      m.pair_solimp,
-      m.pair_solref,
-      m.pair_solreffriction,
-      m._impl.plugin,
-      m._impl.plugin_attr,
-      m._impl.qLD_all_updates,
-      m._impl.qLD_level_offsets,
-      m._impl.qLD_updates,
-      m._impl.qM_fullm_i,
-      m._impl.qM_fullm_j,
-      m._impl.qM_mulm_col,
-      m._impl.qM_mulm_madr,
-      m._impl.qM_mulm_rowadr,
-      m._impl.qM_tiles,
-      m.qpos0,
-      m.qpos_spring,
-      m._impl.rangefinder_sensor_adr,
-      m._impl.sensor_acc_adr,
-      m.sensor_adr,
-      m._impl.sensor_adr_to_contact_adr,
-      m._impl.sensor_contact_adr,
-      m.sensor_cutoff,
-      m.sensor_datatype,
-      m.sensor_dim,
-      m._impl.sensor_e_kinetic,
-      m._impl.sensor_e_potential,
-      m.sensor_intprm,
-      m._impl.sensor_limitfrc_adr,
-      m._impl.sensor_limitpos_adr,
-      m._impl.sensor_limitvel_adr,
-      m.sensor_objid,
-      m.sensor_objtype,
-      m._impl.sensor_pos_adr,
-      m._impl.sensor_rangefinder_adr,
-      m._impl.sensor_rangefinder_bodyid,
-      m.sensor_refid,
-      m.sensor_reftype,
-      m._impl.sensor_rne_postconstraint,
-      m._impl.sensor_subtree_vel,
-      m._impl.sensor_tendonactfrc_adr,
-      m._impl.sensor_touch_adr,
-      m.sensor_type,
-      m._impl.sensor_vel_adr,
-      m.site_bodyid,
-      m.site_pos,
-      m.site_quat,
-      m.site_size,
-      m.site_type,
-      m._impl.taxel_sensorid,
-      m._impl.taxel_vertadr,
-      m._impl.ten_J_colind,
-      m._impl.ten_J_rowadr,
-      m._impl.ten_J_rownnz,
-      m.tendon_actfrclimited,
-      m.tendon_actfrcrange,
-      m.tendon_adr,
-      m.tendon_armature,
-      m.tendon_damping,
-      m.tendon_dampingpoly,
-      m.tendon_frictionloss,
-      m._impl.tendon_geom_adr,
-      m.tendon_invweight0,
-      m._impl.tendon_jnt_adr,
-      m.tendon_length0,
-      m.tendon_lengthspring,
-      m._impl.tendon_limited_adr,
-      m.tendon_margin,
-      m.tendon_num,
-      m.tendon_range,
-      m._impl.tendon_site_pair_adr,
-      m.tendon_solimp_fri,
-      m.tendon_solimp_lim,
-      m.tendon_solref_fri,
-      m.tendon_solref_lim,
-      m.tendon_stiffness,
-      m.tendon_stiffnesspoly,
-      m._impl.wrap_geom_adr,
-      m._impl.wrap_jnt_adr,
-      m.wrap_objid,
-      m.wrap_prm,
-      m._impl.wrap_pulley_scale,
-      m._impl.wrap_site_pair_adr,
-      m.wrap_type,
-      m.opt._impl.broadphase,
-      m.opt._impl.broadphase_filter,
-      m.opt._impl.ccd_iterations,
-      m.opt._impl.ccd_tolerance,
-      m.opt.cone,
-      m.opt._impl.contact_sensor_maxmatch,
-      m.opt.density,
-      m.opt.disableflags,
-      m.opt.enableflags,
-      m.opt._impl.graph_conditional,
-      m.opt.gravity,
-      m.opt._impl.impratio_invsqrt,
-      m.opt.iterations,
-      m.opt.ls_iterations,
-      m.opt._impl.ls_parallel,
-      m.opt._impl.ls_parallel_min_step,
-      m.opt.ls_tolerance,
-      m.opt.magnetic,
-      m.opt._impl.run_collision_detection,
-      m.opt._impl.sdf_initpoints,
-      m.opt._impl.sdf_iterations,
-      m.opt.solver,
-      m.opt.timestep,
-      m.opt.tolerance,
-      m.opt.viscosity,
-      m.opt.wind,
-      m.stat.meaninertia,
-      d._impl.naccdmax,
-      d._impl.naconmax,
-      d._impl.njmax,
-      d._impl.njmax_nnz,
-      d.act,
-      d.act_dot,
-      d.actuator_force,
-      d.actuator_length,
-      d._impl.actuator_moment,
-      d._impl.actuator_velocity,
-      d._impl.cacc,
-      d.cam_xmat,
-      d.cam_xpos,
-      d.cdof,
-      d.cdof_dot,
-      d._impl.cfrc_ext,
-      d._impl.cfrc_int,
-      d._impl.cinert,
-      d._impl.crb,
-      d.ctrl,
-      d.cvel,
-      d._impl.energy,
-      d.eq_active,
-      d._impl.flexedge_J,
-      d._impl.flexedge_length,
-      d._impl.flexedge_velocity,
-      d._impl.flexvert_xpos,
-      d.geom_xmat,
-      d.geom_xpos,
-      d._impl.light_xdir,
-      d._impl.light_xpos,
-      d.mocap_pos,
-      d.mocap_quat,
-      d._impl.moment_colind,
-      d._impl.moment_rowadr,
-      d._impl.moment_rownnz,
-      d._impl.nacon,
-      d._impl.ncollision,
-      d._impl.ne,
-      d._impl.nefc,
-      d._impl.nf,
-      d._impl.nisland,
-      d._impl.nl,
-      d._impl.qLD,
-      d._impl.qLDiagInv,
-      d._impl.qM,
-      d.qacc,
-      d.qacc_smooth,
-      d.qacc_warmstart,
-      d.qfrc_actuator,
-      d.qfrc_applied,
-      d.qfrc_bias,
-      d.qfrc_constraint,
-      d._impl.qfrc_damper,
-      d.qfrc_fluid,
-      d.qfrc_gravcomp,
-      d.qfrc_passive,
-      d.qfrc_smooth,
-      d._impl.qfrc_spring,
-      d.qpos,
-      d.qvel,
-      d.sensordata,
-      d.site_xmat,
-      d.site_xpos,
-      d._impl.solver_niter,
-      d._impl.subtree_angmom,
-      d.subtree_com,
-      d._impl.subtree_linvel,
-      d._impl.ten_J,
-      d.ten_length,
-      d._impl.ten_velocity,
-      d._impl.ten_wrapadr,
-      d._impl.ten_wrapnum,
-      d.time,
-      d._impl.tree_island,
-      d._impl.wrap_obj,
-      d._impl.wrap_xpos,
-      d.xanchor,
-      d.xaxis,
-      d.xfrc_applied,
-      d.ximat,
-      d.xipos,
-      d.xmat,
-      d.xpos,
-      d.xquat,
-      d._impl.contact__dim,
-      d._impl.contact__dist,
-      d._impl.contact__efc_address,
-      d._impl.contact__flex,
-      d._impl.contact__frame,
-      d._impl.contact__friction,
-      d._impl.contact__geom,
-      d._impl.contact__geomcollisionid,
-      d._impl.contact__includemargin,
-      d._impl.contact__pos,
-      d._impl.contact__solimp,
-      d._impl.contact__solref,
-      d._impl.contact__solreffriction,
-      d._impl.contact__type,
-      d._impl.contact__vert,
-      d._impl.contact__worldid,
-      d._impl.efc__D,
-      d._impl.efc__J,
-      d._impl.efc__J_colind,
-      d._impl.efc__J_rowadr,
-      d._impl.efc__J_rownnz,
-      d._impl.efc__Jqvel,
-      d._impl.efc__Ma,
-      d._impl.efc__aref,
-      d._impl.efc__force,
-      d._impl.efc__frictionloss,
-      d._impl.efc__id,
-      d._impl.efc__margin,
-      d._impl.efc__pos,
-      d._impl.efc__state,
-      d._impl.efc__type,
-      d._impl.efc__vel,
-  )
-  d = d.tree_replace({
-      'act_dot': out[0],
-      'actuator_force': out[1],
-      'actuator_length': out[2],
-      '_impl.actuator_moment': out[3],
-      '_impl.actuator_velocity': out[4],
-      '_impl.cacc': out[5],
-      'cam_xmat': out[6],
-      'cam_xpos': out[7],
-      'cdof': out[8],
-      'cdof_dot': out[9],
-      '_impl.cfrc_ext': out[10],
-      '_impl.cfrc_int': out[11],
-      '_impl.cinert': out[12],
-      '_impl.crb': out[13],
-      'cvel': out[14],
-      '_impl.energy': out[15],
-      '_impl.flexedge_J': out[16],
-      '_impl.flexedge_length': out[17],
-      '_impl.flexedge_velocity': out[18],
-      '_impl.flexvert_xpos': out[19],
-      'geom_xmat': out[20],
-      'geom_xpos': out[21],
-      '_impl.light_xdir': out[22],
-      '_impl.light_xpos': out[23],
-      '_impl.moment_colind': out[24],
-      '_impl.moment_rowadr': out[25],
-      '_impl.moment_rownnz': out[26],
-      '_impl.nacon': out[27],
-      '_impl.ncollision': out[28],
-      '_impl.ne': out[29],
-      '_impl.nefc': out[30],
-      '_impl.nf': out[31],
-      '_impl.nisland': out[32],
-      '_impl.nl': out[33],
-      '_impl.qLD': out[34],
-      '_impl.qLDiagInv': out[35],
-      '_impl.qM': out[36],
-      'qacc': out[37],
-      'qacc_smooth': out[38],
-      'qfrc_actuator': out[39],
-      'qfrc_bias': out[40],
-      'qfrc_constraint': out[41],
-      '_impl.qfrc_damper': out[42],
-      'qfrc_fluid': out[43],
-      'qfrc_gravcomp': out[44],
-      'qfrc_passive': out[45],
-      'qfrc_smooth': out[46],
-      '_impl.qfrc_spring': out[47],
-      'qvel': out[48],
-      'sensordata': out[49],
-      'site_xmat': out[50],
-      'site_xpos': out[51],
-      '_impl.solver_niter': out[52],
-      '_impl.subtree_angmom': out[53],
-      'subtree_com': out[54],
-      '_impl.subtree_linvel': out[55],
-      '_impl.ten_J': out[56],
-      'ten_length': out[57],
-      '_impl.ten_velocity': out[58],
-      '_impl.ten_wrapadr': out[59],
-      '_impl.ten_wrapnum': out[60],
-      '_impl.tree_island': out[61],
-      '_impl.wrap_obj': out[62],
-      '_impl.wrap_xpos': out[63],
-      'xanchor': out[64],
-      'xaxis': out[65],
-      'ximat': out[66],
-      'xipos': out[67],
-      'xmat': out[68],
-      'xpos': out[69],
-      'xquat': out[70],
-      '_impl.contact__dim': out[71],
-      '_impl.contact__dist': out[72],
-      '_impl.contact__efc_address': out[73],
-      '_impl.contact__flex': out[74],
-      '_impl.contact__frame': out[75],
-      '_impl.contact__friction': out[76],
-      '_impl.contact__geom': out[77],
-      '_impl.contact__geomcollisionid': out[78],
-      '_impl.contact__includemargin': out[79],
-      '_impl.contact__pos': out[80],
-      '_impl.contact__solimp': out[81],
-      '_impl.contact__solref': out[82],
-      '_impl.contact__solreffriction': out[83],
-      '_impl.contact__type': out[84],
-      '_impl.contact__vert': out[85],
-      '_impl.contact__worldid': out[86],
-      '_impl.efc__D': out[87],
-      '_impl.efc__J': out[88],
-      '_impl.efc__J_colind': out[89],
-      '_impl.efc__J_rowadr': out[90],
-      '_impl.efc__J_rownnz': out[91],
-      '_impl.efc__Jqvel': out[92],
-      '_impl.efc__Ma': out[93],
-      '_impl.efc__aref': out[94],
-      '_impl.efc__force': out[95],
-      '_impl.efc__frictionloss': out[96],
-      '_impl.efc__id': out[97],
-      '_impl.efc__margin': out[98],
-      '_impl.efc__pos': out[99],
-      '_impl.efc__state': out[100],
-      '_impl.efc__type': out[101],
-      '_impl.efc__vel': out[102],
-  })
-  return d
-
-
-@jax.custom_batching.custom_vmap
-@ffi.marshal_jax_warp_callable
-def forward(m: types.Model, d: types.Data):
-  return _forward_jax_impl(m, d)
-
-
-@forward.def_vmap
-@ffi.marshal_custom_vmap
-def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
-  d = forward(m, d)
-  return d, is_batched[1]
-
-
-@ffi.format_args_for_warp
-def _step_shim(
-    # Model
-    nworld: int,
-    M_rowadr: wp.array[int],
-    M_rownnz: wp.array[int],
-    actuator_acc0: wp.array2d[float],
-    actuator_actadr: wp.array[int],
-    actuator_actearly: wp.array[bool],
-    actuator_actlimited: wp.array[bool],
-    actuator_actnum: wp.array[int],
-    actuator_actrange: wp.array2d[wp.vec2],
-    actuator_biasprm: wp.array2d[mjwp_types.vec10f],
-    actuator_biastype: wp.array[int],
-    actuator_cranklength: wp.array2d[float],
-    actuator_ctrllimited: wp.array[bool],
-    actuator_ctrlrange: wp.array2d[wp.vec2],
-    actuator_dynprm: wp.array2d[mjwp_types.vec10f],
-    actuator_dyntype: wp.array[int],
-    actuator_forcelimited: wp.array[bool],
-    actuator_forcerange: wp.array2d[wp.vec2],
-    actuator_gainprm: wp.array2d[mjwp_types.vec10f],
-    actuator_gaintype: wp.array[int],
-    actuator_gear: wp.array2d[wp.spatial_vector],
-    actuator_lengthrange: wp.array2d[wp.vec2],
-    actuator_trnid: wp.array[wp.vec2i],
-    actuator_trntype: wp.array[int],
-    actuator_trntype_body_adr: wp.array[int],
-    block_dim: mjwp_types.BlockDim,
-    body_branch_start: wp.array[int],
-    body_branches: wp.array[int],
-    body_dofadr: wp.array[int],
-    body_dofnum: wp.array[int],
-    body_fluid_ellipsoid: wp.array[bool],
-    body_geomadr: wp.array[int],
-    body_geomnum: wp.array[int],
-    body_gravcomp: wp.array2d[float],
-    body_inertia: wp.array2d[wp.vec3],
-    body_invweight0: wp.array2d[wp.vec2],
-    body_ipos: wp.array2d[wp.vec3],
-    body_iquat: wp.array2d[wp.quat],
-    body_isdofancestor: wp.array2d[int],
-    body_jntadr: wp.array[int],
-    body_jntnum: wp.array[int],
-    body_mass: wp.array2d[float],
-    body_mocapid: wp.array[int],
-    body_parentid: wp.array[int],
-    body_pos: wp.array2d[wp.vec3],
-    body_quat: wp.array2d[wp.quat],
-    body_rootid: wp.array[int],
-    body_subtreemass: wp.array2d[float],
-    body_tree: tuple[wp.array[int], ...],
-    body_treeid: wp.array[int],
-    body_weldid: wp.array[int],
-    cam_bodyid: wp.array[int],
-    cam_fovy: wp.array2d[float],
-    cam_intrinsic: wp.array2d[wp.vec4],
-    cam_mat0: wp.array2d[wp.mat33],
-    cam_mode: wp.array[int],
-    cam_pos: wp.array2d[wp.vec3],
-    cam_pos0: wp.array2d[wp.vec3],
-    cam_poscom0: wp.array2d[wp.vec3],
-    cam_quat: wp.array2d[wp.quat],
-    cam_resolution: wp.array[wp.vec2i],
-    cam_sensorsize: wp.array[wp.vec2],
-    cam_targetbodyid: wp.array[int],
-    dof_Madr: wp.array[int],
-    dof_armature: wp.array2d[float],
-    dof_bodyid: wp.array[int],
-    dof_damping: wp.array2d[float],
-    dof_dampingpoly: wp.array2d[wp.vec2],
-    dof_frictionloss: wp.array2d[float],
-    dof_invweight0: wp.array2d[float],
-    dof_jntid: wp.array[int],
-    dof_parentid: wp.array[int],
-    dof_solimp: wp.array2d[mjwp_types.vec5],
-    dof_solref: wp.array2d[wp.vec2],
-    dof_treeid: wp.array[int],
-    dof_tri_col: wp.array[int],
-    dof_tri_row: wp.array[int],
-    eq_connect_adr: wp.array[int],
-    eq_data: wp.array2d[mjwp_types.vec11],
-    eq_flex_adr: wp.array[int],
-    eq_jnt_adr: wp.array[int],
-    eq_obj1id: wp.array[int],
-    eq_obj2id: wp.array[int],
-    eq_objtype: wp.array[int],
-    eq_solimp: wp.array2d[mjwp_types.vec5],
-    eq_solref: wp.array2d[wp.vec2],
-    eq_ten_adr: wp.array[int],
-    eq_type: wp.array[int],
-    eq_wld_adr: wp.array[int],
-    flex_bending: wp.array[float],
-    flex_bendingadr: wp.array[int],
-    flex_centered: wp.array[bool],
-    flex_conaffinity: wp.array[int],
-    flex_condim: wp.array[int],
-    flex_contype: wp.array[int],
-    flex_damping: wp.array[float],
-    flex_dim: wp.array[int],
-    flex_edge: wp.array[wp.vec2i],
-    flex_edgeadr: wp.array[int],
-    flex_edgeflap: wp.array[wp.vec2i],
-    flex_edgenum: wp.array[int],
-    flex_elem: wp.array[int],
-    flex_elemadr: wp.array[int],
-    flex_elemdataadr: wp.array[int],
-    flex_elemedge: wp.array[int],
-    flex_elemedgeadr: wp.array[int],
-    flex_elemnum: wp.array[int],
-    flex_friction: wp.array[wp.vec3],
-    flex_gap: wp.array[float],
-    flex_margin: wp.array[float],
-    flex_priority: wp.array[int],
-    flex_radius: wp.array[float],
-    flex_shell: wp.array[int],
-    flex_shelldataadr: wp.array[int],
-    flex_shellnum: wp.array[int],
-    flex_solimp: wp.array[mjwp_types.vec5],
-    flex_solmix: wp.array[float],
-    flex_solref: wp.array[wp.vec2],
-    flex_stiffness: wp.array[float],
-    flex_stiffnessadr: wp.array[int],
-    flex_vert: wp.array[wp.vec3],
-    flex_vertadr: wp.array[int],
-    flex_vertbodyid: wp.array[int],
-    flex_vertflexid: wp.array[int],
-    flex_vertnum: wp.array[int],
-    flexedge_J_colind: wp.array[int],
-    flexedge_J_rowadr: wp.array[int],
-    flexedge_J_rownnz: wp.array[int],
-    flexedge_invweight0: wp.array[float],
-    flexedge_length0: wp.array[float],
-    geom_aabb: wp.array3d[wp.vec3],
-    geom_bodyid: wp.array[int],
-    geom_conaffinity: wp.array[int],
-    geom_condim: wp.array[int],
-    geom_contype: wp.array[int],
-    geom_dataid: wp.array2d[int],
-    geom_fluid: wp.array2d[float],
-    geom_friction: wp.array2d[wp.vec3],
-    geom_gap: wp.array2d[float],
-    geom_group: wp.array[int],
-    geom_margin: wp.array2d[float],
-    geom_matid: wp.array2d[int],
-    geom_pair_type_count: tuple[int, ...],
-    geom_plugin_index: wp.array[int],
-    geom_pos: wp.array2d[wp.vec3],
-    geom_priority: wp.array[int],
-    geom_quat: wp.array2d[wp.quat],
-    geom_rbound: wp.array2d[float],
-    geom_rgba: wp.array2d[wp.vec4],
-    geom_size: wp.array2d[wp.vec3],
-    geom_solimp: wp.array2d[mjwp_types.vec5],
-    geom_solmix: wp.array2d[float],
-    geom_solref: wp.array2d[wp.vec2],
-    geom_type: wp.array[int],
-    has_fluid: bool,
-    has_sdf_geom: bool,
-    hfield_adr: wp.array[int],
-    hfield_data: wp.array[float],
-    hfield_ncol: wp.array[int],
-    hfield_nrow: wp.array[int],
-    hfield_size: wp.array[wp.vec4],
-    is_sparse: bool,
-    jnt_actfrclimited: wp.array[bool],
-    jnt_actfrcrange: wp.array2d[wp.vec2],
-    jnt_actgravcomp: wp.array[int],
-    jnt_axis: wp.array2d[wp.vec3],
-    jnt_bodyid: wp.array[int],
-    jnt_dofadr: wp.array[int],
-    jnt_limited_ball_adr: wp.array[int],
-    jnt_limited_slide_hinge_adr: wp.array[int],
-    jnt_margin: wp.array2d[float],
-    jnt_pos: wp.array2d[wp.vec3],
-    jnt_qposadr: wp.array[int],
-    jnt_range: wp.array2d[wp.vec2],
-    jnt_solimp: wp.array2d[mjwp_types.vec5],
-    jnt_solref: wp.array2d[wp.vec2],
-    jnt_stiffness: wp.array2d[float],
-    jnt_stiffnesspoly: wp.array2d[wp.vec2],
-    jnt_type: wp.array[int],
-    light_bodyid: wp.array[int],
-    light_dir: wp.array2d[wp.vec3],
-    light_dir0: wp.array2d[wp.vec3],
-    light_mode: wp.array[int],
-    light_pos: wp.array2d[wp.vec3],
-    light_pos0: wp.array2d[wp.vec3],
-    light_poscom0: wp.array2d[wp.vec3],
-    light_targetbodyid: wp.array[int],
-    mapM2M: wp.array[int],
-    mat_rgba: wp.array2d[wp.vec4],
-    max_ten_J_rownnz: int,
-    mesh_face: wp.array[wp.vec3i],
-    mesh_faceadr: wp.array[int],
-    mesh_graph: wp.array[int],
-    mesh_graphadr: wp.array[int],
-    mesh_normal: wp.array[wp.vec3],
-    mesh_normaladr: wp.array[int],
-    mesh_normalnum: wp.array[int],
-    mesh_octadr: wp.array[int],
-    mesh_polyadr: wp.array[int],
-    mesh_polymap: wp.array[int],
-    mesh_polymapadr: wp.array[int],
-    mesh_polymapnum: wp.array[int],
-    mesh_polynormal: wp.array[wp.vec3],
-    mesh_polynum: wp.array[int],
-    mesh_polyvert: wp.array[int],
-    mesh_polyvertadr: wp.array[int],
-    mesh_polyvertnum: wp.array[int],
-    mesh_quat: wp.array[wp.quat],
-    mesh_vert: wp.array[wp.vec3],
-    mesh_vertadr: wp.array[int],
-    mesh_vertnum: wp.array[int],
-    nC: int,
-    nJten: int,
-    nM: int,
-    na: int,
-    nacttrnbody: int,
-    nbody: int,
-    nbranch: int,
-    ncam: int,
-    neq: int,
-    nflex: int,
-    nflexedge: int,
-    nflexelem: int,
-    nflexshelldata: int,
-    nflexvert: int,
-    ngeom: int,
-    ngravcomp: int,
-    njnt: int,
-    nlight: int,
-    nmaxcondim: int,
-    nmaxmeshdeg: int,
-    nmaxpolygon: int,
-    nmaxpyramid: int,
-    nmeshface: int,
-    nrangefinder: int,
-    nsensorcollision: int,
-    nsensorcontact: int,
-    nsensortaxel: int,
-    nsite: int,
-    ntendon: int,
-    ntree: int,
-    nu: int,
-    nv: int,
-    nv_pad: int,
-    nwrap: int,
-    nxn_geom_pair_filtered: wp.array[wp.vec2i],
-    nxn_pairid: wp.array[wp.vec2i],
-    nxn_pairid_filtered: wp.array[wp.vec2i],
-    oct_aabb: wp.array2d[wp.vec3],
-    oct_child: wp.array[mjwp_types.vec8i],
-    oct_coeff: wp.array[mjwp_types.vec8],
-    pair_dim: wp.array[int],
-    pair_friction: wp.array2d[mjwp_types.vec5],
-    pair_gap: wp.array2d[float],
-    pair_margin: wp.array2d[float],
-    pair_solimp: wp.array2d[mjwp_types.vec5],
-    pair_solref: wp.array2d[wp.vec2],
-    pair_solreffriction: wp.array2d[wp.vec2],
-    plugin: wp.array[int],
-    plugin_attr: wp.array[mjwp_types.vec_pluginattr],
-    qLD_all_updates: wp.array[wp.vec3i],
-    qLD_level_offsets: wp.array[int],
-    qLD_updates: tuple[wp.array[wp.vec3i], ...],
-    qM_fullm_i: wp.array[int],
-    qM_fullm_j: wp.array[int],
-    qM_mulm_col: wp.array[int],
-    qM_mulm_madr: wp.array[int],
-    qM_mulm_rowadr: wp.array[int],
-    qM_tiles: tuple[mjwp_types.TileSet, ...],
     qpos0: wp.array2d[float],
     qpos_spring: wp.array2d[float],
     rangefinder_sensor_adr: wp.array[int],
@@ -2305,9 +334,13 @@ def _step_shim(
     sensor_contact_adr: wp.array[int],
     sensor_cutoff: wp.array[float],
     sensor_datatype: wp.array[int],
+    sensor_delay: wp.array[float],
     sensor_dim: wp.array[int],
     sensor_e_kinetic: bool,
     sensor_e_potential: bool,
+    sensor_history: wp.array[wp.vec2i],
+    sensor_historyadr: wp.array[int],
+    sensor_interval: wp.array[wp.vec2],
     sensor_intprm: wp.array2d[int],
     sensor_limitfrc_adr: wp.array[int],
     sensor_limitpos_adr: wp.array[int],
@@ -2398,6 +431,7 @@ def _step_shim(
     naconmax: int,
     njmax: int,
     njmax_nnz: int,
+    M: wp.array3d[float],
     act: wp.array2d[float],
     act_dot: wp.array2d[float],
     actuator_force: wp.array2d[float],
@@ -2415,6 +449,9 @@ def _step_shim(
     crb: wp.array2d[mjwp_types.vec10],
     ctrl: wp.array2d[float],
     cvel: wp.array2d[wp.spatial_vector],
+    dof_island: wp.array2d[int],
+    dof_islandid: wp.array2d[int],
+    efc_islandid: wp.array2d[int],
     energy: wp.array[wp.vec2],
     eq_active: wp.array2d[bool],
     flexedge_J: wp.array2d[float],
@@ -2423,8 +460,23 @@ def _step_shim(
     flexvert_xpos: wp.array2d[wp.vec3],
     geom_xmat: wp.array2d[wp.mat33],
     geom_xpos: wp.array2d[wp.vec3],
+    history: wp.array2d[float],
+    iqacc: wp.array2d[float],
+    iqacc_smooth: wp.array2d[float],
+    iqfrc_constraint: wp.array2d[float],
+    iqfrc_smooth: wp.array2d[float],
+    island_dofadr: wp.array2d[int],
+    island_efcadr: wp.array2d[int],
+    island_ne: wp.array2d[int],
+    island_nefc: wp.array2d[int],
+    island_nf: wp.array2d[int],
+    island_nv: wp.array2d[int],
     light_xdir: wp.array2d[wp.vec3],
     light_xpos: wp.array2d[wp.vec3],
+    map_dof2idof: wp.array2d[int],
+    map_efc2iefc: wp.array2d[int],
+    map_idof2dof: wp.array2d[int],
+    map_iefc2efc: wp.array2d[int],
     mocap_pos: wp.array2d[wp.vec3],
     mocap_quat: wp.array2d[wp.quat],
     moment_colind: wp.array2d[int],
@@ -2435,11 +487,11 @@ def _step_shim(
     ne: wp.array[int],
     nefc: wp.array[int],
     nf: wp.array[int],
+    nidof: wp.array[int],
     nisland: wp.array[int],
     nl: wp.array[int],
     qLD: wp.array3d[float],
     qLDiagInv: wp.array2d[float],
-    qM: wp.array3d[float],
     qacc: wp.array2d[float],
     qacc_smooth: wp.array2d[float],
     qacc_warmstart: wp.array2d[float],
@@ -2505,7 +557,19 @@ def _step_shim(
     efc__aref: wp.array2d[float],
     efc__force: wp.array2d[float],
     efc__frictionloss: wp.array2d[float],
+    efc__iD: wp.array2d[float],
+    efc__iJ: wp.array3d[float],
+    efc__iJ_colind: wp.array3d[int],
+    efc__iJ_rowadr: wp.array2d[int],
+    efc__iJ_rownnz: wp.array2d[int],
+    efc__iaref: wp.array2d[float],
     efc__id: wp.array2d[int],
+    efc__iforce: wp.array2d[float],
+    efc__ifrictionloss: wp.array2d[float],
+    efc__iid: wp.array2d[int],
+    efc__island: wp.array2d[int],
+    efc__istate: wp.array2d[int],
+    efc__itype: wp.array2d[int],
     efc__margin: wp.array2d[float],
     efc__pos: wp.array2d[float],
     efc__state: wp.array2d[int],
@@ -2517,8 +581,18 @@ def _step_shim(
   _m.callback = _cb
   _d.efc = _e
   _d.contact = _c
+  _m.M_elemid = M_elemid
+  _m.M_fullm_i = M_fullm_i
+  _m.M_fullm_j = M_fullm_j
+  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
+  _m.M_fullm_upper_i = M_fullm_upper_i
+  _m.M_fullm_upper_j = M_fullm_upper_j
+  _m.M_mulm_col = M_mulm_col
+  _m.M_mulm_madr = M_mulm_madr
+  _m.M_mulm_rowadr = M_mulm_rowadr
   _m.M_rowadr = M_rowadr
   _m.M_rownnz = M_rownnz
+  _m.M_tiles = M_tiles
   _m.actuator_acc0 = actuator_acc0
   _m.actuator_actadr = actuator_actadr
   _m.actuator_actearly = actuator_actearly
@@ -2530,6 +604,7 @@ def _step_shim(
   _m.actuator_cranklength = actuator_cranklength
   _m.actuator_ctrllimited = actuator_ctrllimited
   _m.actuator_ctrlrange = actuator_ctrlrange
+  _m.actuator_delay = actuator_delay
   _m.actuator_dynprm = actuator_dynprm
   _m.actuator_dyntype = actuator_dyntype
   _m.actuator_forcelimited = actuator_forcelimited
@@ -2537,6 +612,8 @@ def _step_shim(
   _m.actuator_gainprm = actuator_gainprm
   _m.actuator_gaintype = actuator_gaintype
   _m.actuator_gear = actuator_gear
+  _m.actuator_history = actuator_history
+  _m.actuator_historyadr = actuator_historyadr
   _m.actuator_lengthrange = actuator_lengthrange
   _m.actuator_trnid = actuator_trnid
   _m.actuator_trntype = actuator_trntype
@@ -2579,7 +656,6 @@ def _step_shim(
   _m.cam_resolution = cam_resolution
   _m.cam_sensorsize = cam_sensorsize
   _m.cam_targetbodyid = cam_targetbodyid
-  _m.dof_Madr = dof_Madr
   _m.dof_armature = dof_armature
   _m.dof_bodyid = dof_bodyid
   _m.dof_damping = dof_damping
@@ -2703,7 +779,6 @@ def _step_shim(
   _m.light_pos0 = light_pos0
   _m.light_poscom0 = light_poscom0
   _m.light_targetbodyid = light_targetbodyid
-  _m.mapM2M = mapM2M
   _m.mat_rgba = mat_rgba
   _m.max_ten_J_rownnz = max_ten_J_rownnz
   _m.mesh_face = mesh_face
@@ -2727,9 +802,7 @@ def _step_shim(
   _m.mesh_vert = mesh_vert
   _m.mesh_vertadr = mesh_vertadr
   _m.mesh_vertnum = mesh_vertnum
-  _m.nC = nC
   _m.nJten = nJten
-  _m.nM = nM
   _m.na = na
   _m.nacttrnbody = nacttrnbody
   _m.nbody = nbody
@@ -2743,6 +816,7 @@ def _step_shim(
   _m.nflexvert = nflexvert
   _m.ngeom = ngeom
   _m.ngravcomp = ngravcomp
+  _m.nhistory = nhistory
   _m.njnt = njnt
   _m.nlight = nlight
   _m.nmaxcondim = nmaxcondim
@@ -2806,12 +880,6 @@ def _step_shim(
   _m.qLD_all_updates = qLD_all_updates
   _m.qLD_level_offsets = qLD_level_offsets
   _m.qLD_updates = qLD_updates
-  _m.qM_fullm_i = qM_fullm_i
-  _m.qM_fullm_j = qM_fullm_j
-  _m.qM_mulm_col = qM_mulm_col
-  _m.qM_mulm_madr = qM_mulm_madr
-  _m.qM_mulm_rowadr = qM_mulm_rowadr
-  _m.qM_tiles = qM_tiles
   _m.qpos0 = qpos0
   _m.qpos_spring = qpos_spring
   _m.rangefinder_sensor_adr = rangefinder_sensor_adr
@@ -2821,9 +889,13 @@ def _step_shim(
   _m.sensor_contact_adr = sensor_contact_adr
   _m.sensor_cutoff = sensor_cutoff
   _m.sensor_datatype = sensor_datatype
+  _m.sensor_delay = sensor_delay
   _m.sensor_dim = sensor_dim
   _m.sensor_e_kinetic = sensor_e_kinetic
   _m.sensor_e_potential = sensor_e_potential
+  _m.sensor_history = sensor_history
+  _m.sensor_historyadr = sensor_historyadr
+  _m.sensor_interval = sensor_interval
   _m.sensor_intprm = sensor_intprm
   _m.sensor_limitfrc_adr = sensor_limitfrc_adr
   _m.sensor_limitpos_adr = sensor_limitpos_adr
@@ -2882,6 +954,7 @@ def _step_shim(
   _m.wrap_pulley_scale = wrap_pulley_scale
   _m.wrap_site_pair_adr = wrap_site_pair_adr
   _m.wrap_type = wrap_type
+  _d.M = M
   _d.act = act
   _d.act_dot = act_dot
   _d.actuator_force = actuator_force
@@ -2915,6 +988,8 @@ def _step_shim(
   _d.crb = crb
   _d.ctrl = ctrl
   _d.cvel = cvel
+  _d.dof_island = dof_island
+  _d.dof_islandid = dof_islandid
   _d.efc.D = efc__D
   _d.efc.J = efc__J
   _d.efc.J_colind = efc__J_colind
@@ -2925,12 +1000,25 @@ def _step_shim(
   _d.efc.aref = efc__aref
   _d.efc.force = efc__force
   _d.efc.frictionloss = efc__frictionloss
+  _d.efc.iD = efc__iD
+  _d.efc.iJ = efc__iJ
+  _d.efc.iJ_colind = efc__iJ_colind
+  _d.efc.iJ_rowadr = efc__iJ_rowadr
+  _d.efc.iJ_rownnz = efc__iJ_rownnz
+  _d.efc.iaref = efc__iaref
   _d.efc.id = efc__id
+  _d.efc.iforce = efc__iforce
+  _d.efc.ifrictionloss = efc__ifrictionloss
+  _d.efc.iid = efc__iid
+  _d.efc.island = efc__island
+  _d.efc.istate = efc__istate
+  _d.efc.itype = efc__itype
   _d.efc.margin = efc__margin
   _d.efc.pos = efc__pos
   _d.efc.state = efc__state
   _d.efc.type = efc__type
   _d.efc.vel = efc__vel
+  _d.efc_islandid = efc_islandid
   _d.energy = energy
   _d.eq_active = eq_active
   _d.flexedge_J = flexedge_J
@@ -2939,8 +1027,23 @@ def _step_shim(
   _d.flexvert_xpos = flexvert_xpos
   _d.geom_xmat = geom_xmat
   _d.geom_xpos = geom_xpos
+  _d.history = history
+  _d.iqacc = iqacc
+  _d.iqacc_smooth = iqacc_smooth
+  _d.iqfrc_constraint = iqfrc_constraint
+  _d.iqfrc_smooth = iqfrc_smooth
+  _d.island_dofadr = island_dofadr
+  _d.island_efcadr = island_efcadr
+  _d.island_ne = island_ne
+  _d.island_nefc = island_nefc
+  _d.island_nf = island_nf
+  _d.island_nv = island_nv
   _d.light_xdir = light_xdir
   _d.light_xpos = light_xpos
+  _d.map_dof2idof = map_dof2idof
+  _d.map_efc2iefc = map_efc2iefc
+  _d.map_idof2dof = map_idof2dof
+  _d.map_iefc2efc = map_iefc2efc
   _d.mocap_pos = mocap_pos
   _d.mocap_quat = mocap_quat
   _d.moment_colind = moment_colind
@@ -2953,13 +1056,13 @@ def _step_shim(
   _d.ne = ne
   _d.nefc = nefc
   _d.nf = nf
+  _d.nidof = nidof
   _d.nisland = nisland
   _d.njmax = njmax
   _d.njmax_nnz = njmax_nnz
   _d.nl = nl
   _d.qLD = qLD
   _d.qLDiagInv = qLDiagInv
-  _d.qM = qM
   _d.qacc = qacc
   _d.qacc_smooth = qacc_smooth
   _d.qacc_warmstart = qacc_warmstart
@@ -3000,12 +1103,12 @@ def _step_shim(
   _d.xpos = xpos
   _d.xquat = xquat
   _d.nworld = nworld
-  mjwarp.step(_m, _d)
+  mjwarp.forward(_m, _d)
 
 
-def _step_jax_impl(m: types.Model, d: types.Data):
+def _forward_jax_impl(m: types.Model, d: types.Data):
   output_dims = {
-      'act': d.act.shape,
+      'M': d._impl.M.shape,
       'act_dot': d.act_dot.shape,
       'actuator_force': d.actuator_force.shape,
       'actuator_length': d.actuator_length.shape,
@@ -3021,6 +1124,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'cinert': d._impl.cinert.shape,
       'crb': d._impl.crb.shape,
       'cvel': d.cvel.shape,
+      'dof_island': d._impl.dof_island.shape,
+      'dof_islandid': d._impl.dof_islandid.shape,
+      'efc_islandid': d._impl.efc_islandid.shape,
       'energy': d._impl.energy.shape,
       'flexedge_J': d._impl.flexedge_J.shape,
       'flexedge_length': d._impl.flexedge_length.shape,
@@ -3028,8 +1134,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'flexvert_xpos': d._impl.flexvert_xpos.shape,
       'geom_xmat': d.geom_xmat.shape,
       'geom_xpos': d.geom_xpos.shape,
+      'history': d.history.shape,
+      'iqacc': d._impl.iqacc.shape,
+      'iqacc_smooth': d._impl.iqacc_smooth.shape,
+      'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
+      'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
+      'island_dofadr': d._impl.island_dofadr.shape,
+      'island_efcadr': d._impl.island_efcadr.shape,
+      'island_ne': d._impl.island_ne.shape,
+      'island_nefc': d._impl.island_nefc.shape,
+      'island_nf': d._impl.island_nf.shape,
+      'island_nv': d._impl.island_nv.shape,
       'light_xdir': d._impl.light_xdir.shape,
       'light_xpos': d._impl.light_xpos.shape,
+      'map_dof2idof': d._impl.map_dof2idof.shape,
+      'map_efc2iefc': d._impl.map_efc2iefc.shape,
+      'map_idof2dof': d._impl.map_idof2dof.shape,
+      'map_iefc2efc': d._impl.map_iefc2efc.shape,
       'moment_colind': d._impl.moment_colind.shape,
       'moment_rowadr': d._impl.moment_rowadr.shape,
       'moment_rownnz': d._impl.moment_rownnz.shape,
@@ -3038,14 +1159,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'ne': d._impl.ne.shape,
       'nefc': d._impl.nefc.shape,
       'nf': d._impl.nf.shape,
+      'nidof': d._impl.nidof.shape,
       'nisland': d._impl.nisland.shape,
       'nl': d._impl.nl.shape,
       'qLD': d._impl.qLD.shape,
       'qLDiagInv': d._impl.qLDiagInv.shape,
-      'qM': d._impl.qM.shape,
       'qacc': d.qacc.shape,
       'qacc_smooth': d.qacc_smooth.shape,
-      'qacc_warmstart': d.qacc_warmstart.shape,
       'qfrc_actuator': d.qfrc_actuator.shape,
       'qfrc_bias': d.qfrc_bias.shape,
       'qfrc_constraint': d.qfrc_constraint.shape,
@@ -3055,7 +1175,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'qfrc_passive': d.qfrc_passive.shape,
       'qfrc_smooth': d.qfrc_smooth.shape,
       'qfrc_spring': d._impl.qfrc_spring.shape,
-      'qpos': d.qpos.shape,
       'qvel': d.qvel.shape,
       'sensordata': d.sensordata.shape,
       'site_xmat': d.site_xmat.shape,
@@ -3069,7 +1188,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'ten_velocity': d._impl.ten_velocity.shape,
       'ten_wrapadr': d._impl.ten_wrapadr.shape,
       'ten_wrapnum': d._impl.ten_wrapnum.shape,
-      'time': d.time.shape,
       'tree_island': d._impl.tree_island.shape,
       'wrap_obj': d._impl.wrap_obj.shape,
       'wrap_xpos': d._impl.wrap_xpos.shape,
@@ -3106,7 +1224,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'efc__aref': d._impl.efc__aref.shape,
       'efc__force': d._impl.efc__force.shape,
       'efc__frictionloss': d._impl.efc__frictionloss.shape,
+      'efc__iD': d._impl.efc__iD.shape,
+      'efc__iJ': d._impl.efc__iJ.shape,
+      'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
+      'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
+      'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
+      'efc__iaref': d._impl.efc__iaref.shape,
       'efc__id': d._impl.efc__id.shape,
+      'efc__iforce': d._impl.efc__iforce.shape,
+      'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
+      'efc__iid': d._impl.efc__iid.shape,
+      'efc__island': d._impl.efc__island.shape,
+      'efc__istate': d._impl.efc__istate.shape,
+      'efc__itype': d._impl.efc__itype.shape,
       'efc__margin': d._impl.efc__margin.shape,
       'efc__pos': d._impl.efc__pos.shape,
       'efc__state': d._impl.efc__state.shape,
@@ -3114,12 +1244,12 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       'efc__vel': d._impl.efc__vel.shape,
   }
   jf = ffi.jax_callable_variadic_tuple(
-      _step_shim,
-      num_outputs=107,
+      _forward_shim,
+      num_outputs=134,
       output_dims=output_dims,
       vmap_method=None,
       in_out_argnames=set([
-          'act',
+          'M',
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -3135,6 +1265,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'cinert',
           'crb',
           'cvel',
+          'dof_island',
+          'dof_islandid',
+          'efc_islandid',
           'energy',
           'flexedge_J',
           'flexedge_length',
@@ -3142,8 +1275,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'flexvert_xpos',
           'geom_xmat',
           'geom_xpos',
+          'history',
+          'iqacc',
+          'iqacc_smooth',
+          'iqfrc_constraint',
+          'iqfrc_smooth',
+          'island_dofadr',
+          'island_efcadr',
+          'island_ne',
+          'island_nefc',
+          'island_nf',
+          'island_nv',
           'light_xdir',
           'light_xpos',
+          'map_dof2idof',
+          'map_efc2iefc',
+          'map_idof2dof',
+          'map_iefc2efc',
           'moment_colind',
           'moment_rowadr',
           'moment_rownnz',
@@ -3152,14 +1300,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ne',
           'nefc',
           'nf',
+          'nidof',
           'nisland',
           'nl',
           'qLD',
           'qLDiagInv',
-          'qM',
           'qacc',
           'qacc_smooth',
-          'qacc_warmstart',
           'qfrc_actuator',
           'qfrc_bias',
           'qfrc_constraint',
@@ -3169,7 +1316,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'qfrc_passive',
           'qfrc_smooth',
           'qfrc_spring',
-          'qpos',
           'qvel',
           'sensordata',
           'site_xmat',
@@ -3183,7 +1329,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'ten_velocity',
           'ten_wrapadr',
           'ten_wrapnum',
-          'time',
           'tree_island',
           'wrap_obj',
           'wrap_xpos',
@@ -3220,7 +1365,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'efc__aref',
           'efc__force',
           'efc__frictionloss',
+          'efc__iD',
+          'efc__iJ',
+          'efc__iJ_colind',
+          'efc__iJ_rowadr',
+          'efc__iJ_rownnz',
+          'efc__iaref',
           'efc__id',
+          'efc__iforce',
+          'efc__ifrictionloss',
+          'efc__iid',
+          'efc__island',
+          'efc__istate',
+          'efc__itype',
           'efc__margin',
           'efc__pos',
           'efc__state',
@@ -3291,6 +1448,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'geom_xmat',
           'geom_xpos',
           'hfield_data',
+          'history',
           'jnt_actfrcrange',
           'jnt_axis',
           'jnt_margin',
@@ -3371,7 +1529,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'xquat',
       ]),
       stage_out_argnames=set([
-          'act',
           'act_dot',
           'actuator_force',
           'actuator_length',
@@ -3382,9 +1539,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'cvel',
           'geom_xmat',
           'geom_xpos',
+          'history',
           'qacc',
           'qacc_smooth',
-          'qacc_warmstart',
           'qfrc_actuator',
           'qfrc_bias',
           'qfrc_constraint',
@@ -3392,14 +1549,12 @@ def _step_jax_impl(m: types.Model, d: types.Data):
           'qfrc_gravcomp',
           'qfrc_passive',
           'qfrc_smooth',
-          'qpos',
           'qvel',
           'sensordata',
           'site_xmat',
           'site_xpos',
           'subtree_com',
           'ten_length',
-          'time',
           'xanchor',
           'xaxis',
           'ximat',
@@ -3413,8 +1568,18 @@ def _step_jax_impl(m: types.Model, d: types.Data):
   )
   out = jf(
       d.qpos.shape[0],
+      m._impl.M_elemid,
+      m._impl.M_fullm_i,
+      m._impl.M_fullm_j,
+      m._impl.M_fullm_upper_elemid,
+      m._impl.M_fullm_upper_i,
+      m._impl.M_fullm_upper_j,
+      m._impl.M_mulm_col,
+      m._impl.M_mulm_madr,
+      m._impl.M_mulm_rowadr,
       m._impl.M_rowadr,
       m._impl.M_rownnz,
+      m._impl.M_tiles,
       m.actuator_acc0,
       m.actuator_actadr,
       m.actuator_actearly,
@@ -3426,6 +1591,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.actuator_cranklength,
       m.actuator_ctrllimited,
       m.actuator_ctrlrange,
+      m._impl.actuator_delay,
       m.actuator_dynprm,
       m.actuator_dyntype,
       m.actuator_forcelimited,
@@ -3433,6 +1599,8 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.actuator_gainprm,
       m.actuator_gaintype,
       m.actuator_gear,
+      m._impl.actuator_history,
+      m._impl.actuator_historyadr,
       m.actuator_lengthrange,
       m.actuator_trnid,
       m.actuator_trntype,
@@ -3475,7 +1643,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.cam_resolution,
       m.cam_sensorsize,
       m.cam_targetbodyid,
-      m.dof_Madr,
       m.dof_armature,
       m.dof_bodyid,
       m.dof_damping,
@@ -3599,7 +1766,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.light_pos0,
       m.light_poscom0,
       m._impl.light_targetbodyid,
-      m._impl.mapM2M,
       m.mat_rgba,
       m._impl.max_ten_J_rownnz,
       m.mesh_face,
@@ -3623,9 +1789,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m.mesh_vert,
       m.mesh_vertadr,
       m.mesh_vertnum,
-      m.nC,
       m.nJten,
-      m.nM,
       m.na,
       m._impl.nacttrnbody,
       m.nbody,
@@ -3639,6 +1803,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.nflexvert,
       m.ngeom,
       m.ngravcomp,
+      m.nhistory,
       m.njnt,
       m.nlight,
       m._impl.nmaxcondim,
@@ -3675,12 +1840,6 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.qLD_all_updates,
       m._impl.qLD_level_offsets,
       m._impl.qLD_updates,
-      m._impl.qM_fullm_i,
-      m._impl.qM_fullm_j,
-      m._impl.qM_mulm_col,
-      m._impl.qM_mulm_madr,
-      m._impl.qM_mulm_rowadr,
-      m._impl.qM_tiles,
       m.qpos0,
       m.qpos_spring,
       m._impl.rangefinder_sensor_adr,
@@ -3690,9 +1849,13 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       m._impl.sensor_contact_adr,
       m.sensor_cutoff,
       m.sensor_datatype,
+      m._impl.sensor_delay,
       m.sensor_dim,
       m._impl.sensor_e_kinetic,
       m._impl.sensor_e_potential,
+      m._impl.sensor_history,
+      m._impl.sensor_historyadr,
+      m._impl.sensor_interval,
       m.sensor_intprm,
       m._impl.sensor_limitfrc_adr,
       m._impl.sensor_limitpos_adr,
@@ -3782,6 +1945,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.naconmax,
       d._impl.njmax,
       d._impl.njmax_nnz,
+      d._impl.M,
       d.act,
       d.act_dot,
       d.actuator_force,
@@ -3799,6 +1963,9 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.crb,
       d.ctrl,
       d.cvel,
+      d._impl.dof_island,
+      d._impl.dof_islandid,
+      d._impl.efc_islandid,
       d._impl.energy,
       d.eq_active,
       d._impl.flexedge_J,
@@ -3807,8 +1974,23 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.flexvert_xpos,
       d.geom_xmat,
       d.geom_xpos,
+      d.history,
+      d._impl.iqacc,
+      d._impl.iqacc_smooth,
+      d._impl.iqfrc_constraint,
+      d._impl.iqfrc_smooth,
+      d._impl.island_dofadr,
+      d._impl.island_efcadr,
+      d._impl.island_ne,
+      d._impl.island_nefc,
+      d._impl.island_nf,
+      d._impl.island_nv,
       d._impl.light_xdir,
       d._impl.light_xpos,
+      d._impl.map_dof2idof,
+      d._impl.map_efc2iefc,
+      d._impl.map_idof2dof,
+      d._impl.map_iefc2efc,
       d.mocap_pos,
       d.mocap_quat,
       d._impl.moment_colind,
@@ -3819,11 +2001,11 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.ne,
       d._impl.nefc,
       d._impl.nf,
+      d._impl.nidof,
       d._impl.nisland,
       d._impl.nl,
       d._impl.qLD,
       d._impl.qLDiagInv,
-      d._impl.qM,
       d.qacc,
       d.qacc_smooth,
       d.qacc_warmstart,
@@ -3889,7 +2071,19 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.efc__aref,
       d._impl.efc__force,
       d._impl.efc__frictionloss,
+      d._impl.efc__iD,
+      d._impl.efc__iJ,
+      d._impl.efc__iJ_colind,
+      d._impl.efc__iJ_rowadr,
+      d._impl.efc__iJ_rownnz,
+      d._impl.efc__iaref,
       d._impl.efc__id,
+      d._impl.efc__iforce,
+      d._impl.efc__ifrictionloss,
+      d._impl.efc__iid,
+      d._impl.efc__island,
+      d._impl.efc__istate,
+      d._impl.efc__itype,
       d._impl.efc__margin,
       d._impl.efc__pos,
       d._impl.efc__state,
@@ -3897,7 +2091,7 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       d._impl.efc__vel,
   )
   d = d.tree_replace({
-      'act': out[0],
+      '_impl.M': out[0],
       'act_dot': out[1],
       'actuator_force': out[2],
       'actuator_length': out[3],
@@ -3913,97 +2107,2366 @@ def _step_jax_impl(m: types.Model, d: types.Data):
       '_impl.cinert': out[13],
       '_impl.crb': out[14],
       'cvel': out[15],
-      '_impl.energy': out[16],
-      '_impl.flexedge_J': out[17],
-      '_impl.flexedge_length': out[18],
-      '_impl.flexedge_velocity': out[19],
-      '_impl.flexvert_xpos': out[20],
-      'geom_xmat': out[21],
-      'geom_xpos': out[22],
-      '_impl.light_xdir': out[23],
-      '_impl.light_xpos': out[24],
-      '_impl.moment_colind': out[25],
-      '_impl.moment_rowadr': out[26],
-      '_impl.moment_rownnz': out[27],
-      '_impl.nacon': out[28],
-      '_impl.ncollision': out[29],
-      '_impl.ne': out[30],
-      '_impl.nefc': out[31],
-      '_impl.nf': out[32],
-      '_impl.nisland': out[33],
-      '_impl.nl': out[34],
-      '_impl.qLD': out[35],
-      '_impl.qLDiagInv': out[36],
-      '_impl.qM': out[37],
-      'qacc': out[38],
-      'qacc_smooth': out[39],
-      'qacc_warmstart': out[40],
-      'qfrc_actuator': out[41],
-      'qfrc_bias': out[42],
-      'qfrc_constraint': out[43],
-      '_impl.qfrc_damper': out[44],
-      'qfrc_fluid': out[45],
-      'qfrc_gravcomp': out[46],
-      'qfrc_passive': out[47],
-      'qfrc_smooth': out[48],
-      '_impl.qfrc_spring': out[49],
-      'qpos': out[50],
-      'qvel': out[51],
-      'sensordata': out[52],
-      'site_xmat': out[53],
-      'site_xpos': out[54],
-      '_impl.solver_niter': out[55],
-      '_impl.subtree_angmom': out[56],
-      'subtree_com': out[57],
-      '_impl.subtree_linvel': out[58],
-      '_impl.ten_J': out[59],
-      'ten_length': out[60],
-      '_impl.ten_velocity': out[61],
-      '_impl.ten_wrapadr': out[62],
-      '_impl.ten_wrapnum': out[63],
-      'time': out[64],
-      '_impl.tree_island': out[65],
-      '_impl.wrap_obj': out[66],
-      '_impl.wrap_xpos': out[67],
-      'xanchor': out[68],
-      'xaxis': out[69],
-      'ximat': out[70],
-      'xipos': out[71],
-      'xmat': out[72],
-      'xpos': out[73],
-      'xquat': out[74],
-      '_impl.contact__dim': out[75],
-      '_impl.contact__dist': out[76],
-      '_impl.contact__efc_address': out[77],
-      '_impl.contact__flex': out[78],
-      '_impl.contact__frame': out[79],
-      '_impl.contact__friction': out[80],
-      '_impl.contact__geom': out[81],
-      '_impl.contact__geomcollisionid': out[82],
-      '_impl.contact__includemargin': out[83],
-      '_impl.contact__pos': out[84],
-      '_impl.contact__solimp': out[85],
-      '_impl.contact__solref': out[86],
-      '_impl.contact__solreffriction': out[87],
-      '_impl.contact__type': out[88],
-      '_impl.contact__vert': out[89],
-      '_impl.contact__worldid': out[90],
-      '_impl.efc__D': out[91],
-      '_impl.efc__J': out[92],
-      '_impl.efc__J_colind': out[93],
-      '_impl.efc__J_rowadr': out[94],
-      '_impl.efc__J_rownnz': out[95],
-      '_impl.efc__Jqvel': out[96],
-      '_impl.efc__Ma': out[97],
-      '_impl.efc__aref': out[98],
-      '_impl.efc__force': out[99],
-      '_impl.efc__frictionloss': out[100],
-      '_impl.efc__id': out[101],
-      '_impl.efc__margin': out[102],
-      '_impl.efc__pos': out[103],
-      '_impl.efc__state': out[104],
-      '_impl.efc__type': out[105],
-      '_impl.efc__vel': out[106],
+      '_impl.dof_island': out[16],
+      '_impl.dof_islandid': out[17],
+      '_impl.efc_islandid': out[18],
+      '_impl.energy': out[19],
+      '_impl.flexedge_J': out[20],
+      '_impl.flexedge_length': out[21],
+      '_impl.flexedge_velocity': out[22],
+      '_impl.flexvert_xpos': out[23],
+      'geom_xmat': out[24],
+      'geom_xpos': out[25],
+      'history': out[26],
+      '_impl.iqacc': out[27],
+      '_impl.iqacc_smooth': out[28],
+      '_impl.iqfrc_constraint': out[29],
+      '_impl.iqfrc_smooth': out[30],
+      '_impl.island_dofadr': out[31],
+      '_impl.island_efcadr': out[32],
+      '_impl.island_ne': out[33],
+      '_impl.island_nefc': out[34],
+      '_impl.island_nf': out[35],
+      '_impl.island_nv': out[36],
+      '_impl.light_xdir': out[37],
+      '_impl.light_xpos': out[38],
+      '_impl.map_dof2idof': out[39],
+      '_impl.map_efc2iefc': out[40],
+      '_impl.map_idof2dof': out[41],
+      '_impl.map_iefc2efc': out[42],
+      '_impl.moment_colind': out[43],
+      '_impl.moment_rowadr': out[44],
+      '_impl.moment_rownnz': out[45],
+      '_impl.nacon': out[46],
+      '_impl.ncollision': out[47],
+      '_impl.ne': out[48],
+      '_impl.nefc': out[49],
+      '_impl.nf': out[50],
+      '_impl.nidof': out[51],
+      '_impl.nisland': out[52],
+      '_impl.nl': out[53],
+      '_impl.qLD': out[54],
+      '_impl.qLDiagInv': out[55],
+      'qacc': out[56],
+      'qacc_smooth': out[57],
+      'qfrc_actuator': out[58],
+      'qfrc_bias': out[59],
+      'qfrc_constraint': out[60],
+      '_impl.qfrc_damper': out[61],
+      'qfrc_fluid': out[62],
+      'qfrc_gravcomp': out[63],
+      'qfrc_passive': out[64],
+      'qfrc_smooth': out[65],
+      '_impl.qfrc_spring': out[66],
+      'qvel': out[67],
+      'sensordata': out[68],
+      'site_xmat': out[69],
+      'site_xpos': out[70],
+      '_impl.solver_niter': out[71],
+      '_impl.subtree_angmom': out[72],
+      'subtree_com': out[73],
+      '_impl.subtree_linvel': out[74],
+      '_impl.ten_J': out[75],
+      'ten_length': out[76],
+      '_impl.ten_velocity': out[77],
+      '_impl.ten_wrapadr': out[78],
+      '_impl.ten_wrapnum': out[79],
+      '_impl.tree_island': out[80],
+      '_impl.wrap_obj': out[81],
+      '_impl.wrap_xpos': out[82],
+      'xanchor': out[83],
+      'xaxis': out[84],
+      'ximat': out[85],
+      'xipos': out[86],
+      'xmat': out[87],
+      'xpos': out[88],
+      'xquat': out[89],
+      '_impl.contact__dim': out[90],
+      '_impl.contact__dist': out[91],
+      '_impl.contact__efc_address': out[92],
+      '_impl.contact__flex': out[93],
+      '_impl.contact__frame': out[94],
+      '_impl.contact__friction': out[95],
+      '_impl.contact__geom': out[96],
+      '_impl.contact__geomcollisionid': out[97],
+      '_impl.contact__includemargin': out[98],
+      '_impl.contact__pos': out[99],
+      '_impl.contact__solimp': out[100],
+      '_impl.contact__solref': out[101],
+      '_impl.contact__solreffriction': out[102],
+      '_impl.contact__type': out[103],
+      '_impl.contact__vert': out[104],
+      '_impl.contact__worldid': out[105],
+      '_impl.efc__D': out[106],
+      '_impl.efc__J': out[107],
+      '_impl.efc__J_colind': out[108],
+      '_impl.efc__J_rowadr': out[109],
+      '_impl.efc__J_rownnz': out[110],
+      '_impl.efc__Jqvel': out[111],
+      '_impl.efc__Ma': out[112],
+      '_impl.efc__aref': out[113],
+      '_impl.efc__force': out[114],
+      '_impl.efc__frictionloss': out[115],
+      '_impl.efc__iD': out[116],
+      '_impl.efc__iJ': out[117],
+      '_impl.efc__iJ_colind': out[118],
+      '_impl.efc__iJ_rowadr': out[119],
+      '_impl.efc__iJ_rownnz': out[120],
+      '_impl.efc__iaref': out[121],
+      '_impl.efc__id': out[122],
+      '_impl.efc__iforce': out[123],
+      '_impl.efc__ifrictionloss': out[124],
+      '_impl.efc__iid': out[125],
+      '_impl.efc__island': out[126],
+      '_impl.efc__istate': out[127],
+      '_impl.efc__itype': out[128],
+      '_impl.efc__margin': out[129],
+      '_impl.efc__pos': out[130],
+      '_impl.efc__state': out[131],
+      '_impl.efc__type': out[132],
+      '_impl.efc__vel': out[133],
+  })
+  return d
+
+
+@jax.custom_batching.custom_vmap
+@ffi.marshal_jax_warp_callable
+def forward(m: types.Model, d: types.Data):
+  return _forward_jax_impl(m, d)
+
+
+@forward.def_vmap
+@ffi.marshal_custom_vmap
+def forward_vmap(unused_axis_size, is_batched, m: types.Model, d: types.Data):
+  d = forward(m, d)
+  return d, is_batched[1]
+
+
+@ffi.format_args_for_warp
+def _step_shim(
+    # Model
+    nworld: int,
+    D_colind: wp.array[int],
+    D_diag: wp.array[int],
+    D_rowadr: wp.array[int],
+    D_rownnz: wp.array[int],
+    M_elemid: wp.array2d[int],
+    M_fullm_i: wp.array[int],
+    M_fullm_j: wp.array[int],
+    M_fullm_upper_elemid: wp.array[int],
+    M_fullm_upper_i: wp.array[int],
+    M_fullm_upper_j: wp.array[int],
+    M_mulm_col: wp.array[int],
+    M_mulm_madr: wp.array[int],
+    M_mulm_rowadr: wp.array[int],
+    M_rowadr: wp.array[int],
+    M_rownnz: wp.array[int],
+    M_tiles: tuple[mjwp_types.TileSet, ...],
+    actuator_acc0: wp.array2d[float],
+    actuator_actadr: wp.array[int],
+    actuator_actearly: wp.array[bool],
+    actuator_actlimited: wp.array[bool],
+    actuator_actnum: wp.array[int],
+    actuator_actrange: wp.array2d[wp.vec2],
+    actuator_biasprm: wp.array2d[mjwp_types.vec10f],
+    actuator_biastype: wp.array[int],
+    actuator_cranklength: wp.array2d[float],
+    actuator_ctrllimited: wp.array[bool],
+    actuator_ctrlrange: wp.array2d[wp.vec2],
+    actuator_delay: wp.array[float],
+    actuator_dynprm: wp.array2d[mjwp_types.vec10f],
+    actuator_dyntype: wp.array[int],
+    actuator_forcelimited: wp.array[bool],
+    actuator_forcerange: wp.array2d[wp.vec2],
+    actuator_gainprm: wp.array2d[mjwp_types.vec10f],
+    actuator_gaintype: wp.array[int],
+    actuator_gear: wp.array2d[wp.spatial_vector],
+    actuator_history: wp.array[wp.vec2i],
+    actuator_historyadr: wp.array[int],
+    actuator_lengthrange: wp.array2d[wp.vec2],
+    actuator_trnid: wp.array[wp.vec2i],
+    actuator_trntype: wp.array[int],
+    actuator_trntype_body_adr: wp.array[int],
+    block_dim: mjwp_types.BlockDim,
+    body_branch_start: wp.array[int],
+    body_branches: wp.array[int],
+    body_dofadr: wp.array[int],
+    body_dofnum: wp.array[int],
+    body_fluid_ellipsoid: wp.array[bool],
+    body_geomadr: wp.array[int],
+    body_geomnum: wp.array[int],
+    body_gravcomp: wp.array2d[float],
+    body_inertia: wp.array2d[wp.vec3],
+    body_invweight0: wp.array2d[wp.vec2],
+    body_ipos: wp.array2d[wp.vec3],
+    body_iquat: wp.array2d[wp.quat],
+    body_isdofancestor: wp.array2d[int],
+    body_jntadr: wp.array[int],
+    body_jntnum: wp.array[int],
+    body_mass: wp.array2d[float],
+    body_mocapid: wp.array[int],
+    body_parentid: wp.array[int],
+    body_pos: wp.array2d[wp.vec3],
+    body_quat: wp.array2d[wp.quat],
+    body_rootid: wp.array[int],
+    body_subtreemass: wp.array2d[float],
+    body_tree: tuple[wp.array[int], ...],
+    body_treeid: wp.array[int],
+    body_weldid: wp.array[int],
+    cam_bodyid: wp.array[int],
+    cam_fovy: wp.array2d[float],
+    cam_intrinsic: wp.array2d[wp.vec4],
+    cam_mat0: wp.array2d[wp.mat33],
+    cam_mode: wp.array[int],
+    cam_pos: wp.array2d[wp.vec3],
+    cam_pos0: wp.array2d[wp.vec3],
+    cam_poscom0: wp.array2d[wp.vec3],
+    cam_quat: wp.array2d[wp.quat],
+    cam_resolution: wp.array[wp.vec2i],
+    cam_sensorsize: wp.array[wp.vec2],
+    cam_targetbodyid: wp.array[int],
+    dof_armature: wp.array2d[float],
+    dof_bodyid: wp.array[int],
+    dof_damping: wp.array2d[float],
+    dof_dampingpoly: wp.array2d[wp.vec2],
+    dof_frictionloss: wp.array2d[float],
+    dof_invweight0: wp.array2d[float],
+    dof_jntid: wp.array[int],
+    dof_parentid: wp.array[int],
+    dof_solimp: wp.array2d[mjwp_types.vec5],
+    dof_solref: wp.array2d[wp.vec2],
+    dof_treeid: wp.array[int],
+    dof_tri_col: wp.array[int],
+    dof_tri_row: wp.array[int],
+    eq_connect_adr: wp.array[int],
+    eq_data: wp.array2d[mjwp_types.vec11],
+    eq_flex_adr: wp.array[int],
+    eq_jnt_adr: wp.array[int],
+    eq_obj1id: wp.array[int],
+    eq_obj2id: wp.array[int],
+    eq_objtype: wp.array[int],
+    eq_solimp: wp.array2d[mjwp_types.vec5],
+    eq_solref: wp.array2d[wp.vec2],
+    eq_ten_adr: wp.array[int],
+    eq_type: wp.array[int],
+    eq_wld_adr: wp.array[int],
+    flex_bending: wp.array[float],
+    flex_bendingadr: wp.array[int],
+    flex_centered: wp.array[bool],
+    flex_conaffinity: wp.array[int],
+    flex_condim: wp.array[int],
+    flex_contype: wp.array[int],
+    flex_damping: wp.array[float],
+    flex_dim: wp.array[int],
+    flex_edge: wp.array[wp.vec2i],
+    flex_edgeadr: wp.array[int],
+    flex_edgeflap: wp.array[wp.vec2i],
+    flex_edgenum: wp.array[int],
+    flex_elem: wp.array[int],
+    flex_elemadr: wp.array[int],
+    flex_elemdataadr: wp.array[int],
+    flex_elemedge: wp.array[int],
+    flex_elemedgeadr: wp.array[int],
+    flex_elemnum: wp.array[int],
+    flex_friction: wp.array[wp.vec3],
+    flex_gap: wp.array[float],
+    flex_margin: wp.array[float],
+    flex_priority: wp.array[int],
+    flex_radius: wp.array[float],
+    flex_shell: wp.array[int],
+    flex_shelldataadr: wp.array[int],
+    flex_shellnum: wp.array[int],
+    flex_solimp: wp.array[mjwp_types.vec5],
+    flex_solmix: wp.array[float],
+    flex_solref: wp.array[wp.vec2],
+    flex_stiffness: wp.array[float],
+    flex_stiffnessadr: wp.array[int],
+    flex_vert: wp.array[wp.vec3],
+    flex_vertadr: wp.array[int],
+    flex_vertbodyid: wp.array[int],
+    flex_vertflexid: wp.array[int],
+    flex_vertnum: wp.array[int],
+    flexedge_J_colind: wp.array[int],
+    flexedge_J_rowadr: wp.array[int],
+    flexedge_J_rownnz: wp.array[int],
+    flexedge_invweight0: wp.array[float],
+    flexedge_length0: wp.array[float],
+    geom_aabb: wp.array3d[wp.vec3],
+    geom_bodyid: wp.array[int],
+    geom_conaffinity: wp.array[int],
+    geom_condim: wp.array[int],
+    geom_contype: wp.array[int],
+    geom_dataid: wp.array2d[int],
+    geom_fluid: wp.array2d[float],
+    geom_friction: wp.array2d[wp.vec3],
+    geom_gap: wp.array2d[float],
+    geom_group: wp.array[int],
+    geom_margin: wp.array2d[float],
+    geom_matid: wp.array2d[int],
+    geom_pair_type_count: tuple[int, ...],
+    geom_plugin_index: wp.array[int],
+    geom_pos: wp.array2d[wp.vec3],
+    geom_priority: wp.array[int],
+    geom_quat: wp.array2d[wp.quat],
+    geom_rbound: wp.array2d[float],
+    geom_rgba: wp.array2d[wp.vec4],
+    geom_size: wp.array2d[wp.vec3],
+    geom_solimp: wp.array2d[mjwp_types.vec5],
+    geom_solmix: wp.array2d[float],
+    geom_solref: wp.array2d[wp.vec2],
+    geom_type: wp.array[int],
+    has_fluid: bool,
+    has_sdf_geom: bool,
+    hfield_adr: wp.array[int],
+    hfield_data: wp.array[float],
+    hfield_ncol: wp.array[int],
+    hfield_nrow: wp.array[int],
+    hfield_size: wp.array[wp.vec4],
+    is_sparse: bool,
+    jnt_actfrclimited: wp.array[bool],
+    jnt_actfrcrange: wp.array2d[wp.vec2],
+    jnt_actgravcomp: wp.array[int],
+    jnt_axis: wp.array2d[wp.vec3],
+    jnt_bodyid: wp.array[int],
+    jnt_dofadr: wp.array[int],
+    jnt_limited_ball_adr: wp.array[int],
+    jnt_limited_slide_hinge_adr: wp.array[int],
+    jnt_margin: wp.array2d[float],
+    jnt_pos: wp.array2d[wp.vec3],
+    jnt_qposadr: wp.array[int],
+    jnt_range: wp.array2d[wp.vec2],
+    jnt_solimp: wp.array2d[mjwp_types.vec5],
+    jnt_solref: wp.array2d[wp.vec2],
+    jnt_stiffness: wp.array2d[float],
+    jnt_stiffnesspoly: wp.array2d[wp.vec2],
+    jnt_type: wp.array[int],
+    light_bodyid: wp.array[int],
+    light_dir: wp.array2d[wp.vec3],
+    light_dir0: wp.array2d[wp.vec3],
+    light_mode: wp.array[int],
+    light_pos: wp.array2d[wp.vec3],
+    light_pos0: wp.array2d[wp.vec3],
+    light_poscom0: wp.array2d[wp.vec3],
+    light_targetbodyid: wp.array[int],
+    mapM2D: wp.array[int],
+    mat_rgba: wp.array2d[wp.vec4],
+    max_ten_J_rownnz: int,
+    mesh_face: wp.array[wp.vec3i],
+    mesh_faceadr: wp.array[int],
+    mesh_graph: wp.array[int],
+    mesh_graphadr: wp.array[int],
+    mesh_normal: wp.array[wp.vec3],
+    mesh_normaladr: wp.array[int],
+    mesh_normalnum: wp.array[int],
+    mesh_octadr: wp.array[int],
+    mesh_polyadr: wp.array[int],
+    mesh_polymap: wp.array[int],
+    mesh_polymapadr: wp.array[int],
+    mesh_polymapnum: wp.array[int],
+    mesh_polynormal: wp.array[wp.vec3],
+    mesh_polynum: wp.array[int],
+    mesh_polyvert: wp.array[int],
+    mesh_polyvertadr: wp.array[int],
+    mesh_polyvertnum: wp.array[int],
+    mesh_quat: wp.array[wp.quat],
+    mesh_vert: wp.array[wp.vec3],
+    mesh_vertadr: wp.array[int],
+    mesh_vertnum: wp.array[int],
+    nC: int,
+    nD: int,
+    nJten: int,
+    na: int,
+    nacttrnbody: int,
+    nbody: int,
+    nbranch: int,
+    ncam: int,
+    neq: int,
+    nflex: int,
+    nflexedge: int,
+    nflexelem: int,
+    nflexshelldata: int,
+    nflexvert: int,
+    ngeom: int,
+    ngravcomp: int,
+    nhistory: int,
+    njnt: int,
+    nlight: int,
+    nmaxcondim: int,
+    nmaxmeshdeg: int,
+    nmaxpolygon: int,
+    nmaxpyramid: int,
+    nmeshface: int,
+    nrangefinder: int,
+    nsensorcollision: int,
+    nsensorcontact: int,
+    nsensortaxel: int,
+    nsite: int,
+    ntendon: int,
+    ntree: int,
+    nu: int,
+    nv: int,
+    nv_pad: int,
+    nwrap: int,
+    nxn_geom_pair_filtered: wp.array[wp.vec2i],
+    nxn_pairid: wp.array[wp.vec2i],
+    nxn_pairid_filtered: wp.array[wp.vec2i],
+    oct_aabb: wp.array2d[wp.vec3],
+    oct_child: wp.array[mjwp_types.vec8i],
+    oct_coeff: wp.array[mjwp_types.vec8],
+    pair_dim: wp.array[int],
+    pair_friction: wp.array2d[mjwp_types.vec5],
+    pair_gap: wp.array2d[float],
+    pair_margin: wp.array2d[float],
+    pair_solimp: wp.array2d[mjwp_types.vec5],
+    pair_solref: wp.array2d[wp.vec2],
+    pair_solreffriction: wp.array2d[wp.vec2],
+    plugin: wp.array[int],
+    plugin_attr: wp.array[mjwp_types.vec_pluginattr],
+    qD_fullm_i: wp.array[int],
+    qD_fullm_j: wp.array[int],
+    qLD_all_updates: wp.array[wp.vec3i],
+    qLD_level_offsets: wp.array[int],
+    qLD_updates: tuple[wp.array[wp.vec3i], ...],
+    qpos0: wp.array2d[float],
+    qpos_spring: wp.array2d[float],
+    rangefinder_sensor_adr: wp.array[int],
+    sensor_acc_adr: wp.array[int],
+    sensor_adr: wp.array[int],
+    sensor_adr_to_contact_adr: wp.array[int],
+    sensor_contact_adr: wp.array[int],
+    sensor_cutoff: wp.array[float],
+    sensor_datatype: wp.array[int],
+    sensor_delay: wp.array[float],
+    sensor_dim: wp.array[int],
+    sensor_e_kinetic: bool,
+    sensor_e_potential: bool,
+    sensor_history: wp.array[wp.vec2i],
+    sensor_historyadr: wp.array[int],
+    sensor_interval: wp.array[wp.vec2],
+    sensor_intprm: wp.array2d[int],
+    sensor_limitfrc_adr: wp.array[int],
+    sensor_limitpos_adr: wp.array[int],
+    sensor_limitvel_adr: wp.array[int],
+    sensor_objid: wp.array[int],
+    sensor_objtype: wp.array[int],
+    sensor_pos_adr: wp.array[int],
+    sensor_rangefinder_adr: wp.array[int],
+    sensor_rangefinder_bodyid: wp.array[int],
+    sensor_refid: wp.array[int],
+    sensor_reftype: wp.array[int],
+    sensor_rne_postconstraint: bool,
+    sensor_subtree_vel: bool,
+    sensor_tendonactfrc_adr: wp.array[int],
+    sensor_touch_adr: wp.array[int],
+    sensor_type: wp.array[int],
+    sensor_vel_adr: wp.array[int],
+    site_bodyid: wp.array[int],
+    site_pos: wp.array2d[wp.vec3],
+    site_quat: wp.array2d[wp.quat],
+    site_size: wp.array[wp.vec3],
+    site_type: wp.array[int],
+    taxel_sensorid: wp.array[int],
+    taxel_vertadr: wp.array[int],
+    ten_J_colind: wp.array[int],
+    ten_J_rowadr: wp.array[int],
+    ten_J_rownnz: wp.array[int],
+    tendon_actfrclimited: wp.array[bool],
+    tendon_actfrcrange: wp.array2d[wp.vec2],
+    tendon_adr: wp.array[int],
+    tendon_armature: wp.array2d[float],
+    tendon_damping: wp.array2d[float],
+    tendon_dampingpoly: wp.array2d[wp.vec2],
+    tendon_frictionloss: wp.array2d[float],
+    tendon_geom_adr: wp.array[int],
+    tendon_invweight0: wp.array2d[float],
+    tendon_jnt_adr: wp.array[int],
+    tendon_length0: wp.array2d[float],
+    tendon_lengthspring: wp.array2d[wp.vec2],
+    tendon_limited_adr: wp.array[int],
+    tendon_margin: wp.array2d[float],
+    tendon_num: wp.array[int],
+    tendon_range: wp.array2d[wp.vec2],
+    tendon_site_pair_adr: wp.array[int],
+    tendon_solimp_fri: wp.array2d[mjwp_types.vec5],
+    tendon_solimp_lim: wp.array2d[mjwp_types.vec5],
+    tendon_solref_fri: wp.array2d[wp.vec2],
+    tendon_solref_lim: wp.array2d[wp.vec2],
+    tendon_stiffness: wp.array2d[float],
+    tendon_stiffnesspoly: wp.array2d[wp.vec2],
+    wrap_geom_adr: wp.array[int],
+    wrap_jnt_adr: wp.array[int],
+    wrap_objid: wp.array[int],
+    wrap_prm: wp.array[float],
+    wrap_pulley_scale: wp.array[float],
+    wrap_site_pair_adr: wp.array[int],
+    wrap_type: wp.array[int],
+    opt__broadphase: int,
+    opt__broadphase_filter: int,
+    opt__ccd_iterations: int,
+    opt__ccd_tolerance: wp.array[float],
+    opt__cone: int,
+    opt__contact_sensor_maxmatch: int,
+    opt__density: wp.array[float],
+    opt__disableflags: int,
+    opt__enableflags: int,
+    opt__graph_conditional: bool,
+    opt__gravity: wp.array[wp.vec3],
+    opt__impratio_invsqrt: wp.array[float],
+    opt__integrator: int,
+    opt__iterations: int,
+    opt__ls_iterations: int,
+    opt__ls_parallel: bool,
+    opt__ls_parallel_min_step: float,
+    opt__ls_tolerance: wp.array[float],
+    opt__magnetic: wp.array[wp.vec3],
+    opt__run_collision_detection: bool,
+    opt__sdf_initpoints: int,
+    opt__sdf_iterations: int,
+    opt__solver: int,
+    opt__timestep: wp.array[float],
+    opt__tolerance: wp.array[float],
+    opt__viscosity: wp.array[float],
+    opt__wind: wp.array[wp.vec3],
+    stat__meaninertia: wp.array[float],
+    # Data
+    naccdmax: int,
+    naconmax: int,
+    njmax: int,
+    njmax_nnz: int,
+    M: wp.array3d[float],
+    act: wp.array2d[float],
+    act_dot: wp.array2d[float],
+    actuator_force: wp.array2d[float],
+    actuator_length: wp.array2d[float],
+    actuator_moment: wp.array2d[float],
+    actuator_velocity: wp.array2d[float],
+    cacc: wp.array2d[wp.spatial_vector],
+    cam_xmat: wp.array2d[wp.mat33],
+    cam_xpos: wp.array2d[wp.vec3],
+    cdof: wp.array2d[wp.spatial_vector],
+    cdof_dot: wp.array2d[wp.spatial_vector],
+    cfrc_ext: wp.array2d[wp.spatial_vector],
+    cfrc_int: wp.array2d[wp.spatial_vector],
+    cinert: wp.array2d[mjwp_types.vec10],
+    crb: wp.array2d[mjwp_types.vec10],
+    ctrl: wp.array2d[float],
+    cvel: wp.array2d[wp.spatial_vector],
+    dof_island: wp.array2d[int],
+    dof_islandid: wp.array2d[int],
+    efc_islandid: wp.array2d[int],
+    energy: wp.array[wp.vec2],
+    eq_active: wp.array2d[bool],
+    flexedge_J: wp.array2d[float],
+    flexedge_length: wp.array2d[float],
+    flexedge_velocity: wp.array2d[float],
+    flexvert_xpos: wp.array2d[wp.vec3],
+    geom_xmat: wp.array2d[wp.mat33],
+    geom_xpos: wp.array2d[wp.vec3],
+    history: wp.array2d[float],
+    iqacc: wp.array2d[float],
+    iqacc_smooth: wp.array2d[float],
+    iqfrc_constraint: wp.array2d[float],
+    iqfrc_smooth: wp.array2d[float],
+    island_dofadr: wp.array2d[int],
+    island_efcadr: wp.array2d[int],
+    island_ne: wp.array2d[int],
+    island_nefc: wp.array2d[int],
+    island_nf: wp.array2d[int],
+    island_nv: wp.array2d[int],
+    light_xdir: wp.array2d[wp.vec3],
+    light_xpos: wp.array2d[wp.vec3],
+    map_dof2idof: wp.array2d[int],
+    map_efc2iefc: wp.array2d[int],
+    map_idof2dof: wp.array2d[int],
+    map_iefc2efc: wp.array2d[int],
+    mocap_pos: wp.array2d[wp.vec3],
+    mocap_quat: wp.array2d[wp.quat],
+    moment_colind: wp.array2d[int],
+    moment_rowadr: wp.array2d[int],
+    moment_rownnz: wp.array2d[int],
+    nacon: wp.array[int],
+    ncollision: wp.array[int],
+    ne: wp.array[int],
+    nefc: wp.array[int],
+    nf: wp.array[int],
+    nidof: wp.array[int],
+    nisland: wp.array[int],
+    nl: wp.array[int],
+    qLD: wp.array3d[float],
+    qLDiagInv: wp.array2d[float],
+    qLU: wp.array3d[float],
+    qacc: wp.array2d[float],
+    qacc_smooth: wp.array2d[float],
+    qacc_warmstart: wp.array2d[float],
+    qfrc_actuator: wp.array2d[float],
+    qfrc_applied: wp.array2d[float],
+    qfrc_bias: wp.array2d[float],
+    qfrc_constraint: wp.array2d[float],
+    qfrc_damper: wp.array2d[float],
+    qfrc_fluid: wp.array2d[float],
+    qfrc_gravcomp: wp.array2d[float],
+    qfrc_passive: wp.array2d[float],
+    qfrc_smooth: wp.array2d[float],
+    qfrc_spring: wp.array2d[float],
+    qpos: wp.array2d[float],
+    qvel: wp.array2d[float],
+    sensordata: wp.array2d[float],
+    site_xmat: wp.array2d[wp.mat33],
+    site_xpos: wp.array2d[wp.vec3],
+    solver_niter: wp.array[int],
+    subtree_angmom: wp.array2d[wp.vec3],
+    subtree_com: wp.array2d[wp.vec3],
+    subtree_linvel: wp.array2d[wp.vec3],
+    ten_J: wp.array2d[float],
+    ten_length: wp.array2d[float],
+    ten_velocity: wp.array2d[float],
+    ten_wrapadr: wp.array2d[int],
+    ten_wrapnum: wp.array2d[int],
+    time: wp.array[float],
+    tree_island: wp.array2d[int],
+    wrap_obj: wp.array2d[wp.vec2i],
+    wrap_xpos: wp.array2d[wp.spatial_vector],
+    xanchor: wp.array2d[wp.vec3],
+    xaxis: wp.array2d[wp.vec3],
+    xfrc_applied: wp.array2d[wp.spatial_vector],
+    ximat: wp.array2d[wp.mat33],
+    xipos: wp.array2d[wp.vec3],
+    xmat: wp.array2d[wp.mat33],
+    xpos: wp.array2d[wp.vec3],
+    xquat: wp.array2d[wp.quat],
+    contact__dim: wp.array[int],
+    contact__dist: wp.array[float],
+    contact__efc_address: wp.array2d[int],
+    contact__flex: wp.array[wp.vec2i],
+    contact__frame: wp.array[wp.mat33],
+    contact__friction: wp.array[mjwp_types.vec5],
+    contact__geom: wp.array[wp.vec2i],
+    contact__geomcollisionid: wp.array[int],
+    contact__includemargin: wp.array[float],
+    contact__pos: wp.array[wp.vec3],
+    contact__solimp: wp.array[mjwp_types.vec5],
+    contact__solref: wp.array[wp.vec2],
+    contact__solreffriction: wp.array[wp.vec2],
+    contact__type: wp.array[int],
+    contact__vert: wp.array[wp.vec2i],
+    contact__worldid: wp.array[int],
+    efc__D: wp.array2d[float],
+    efc__J: wp.array3d[float],
+    efc__J_colind: wp.array3d[int],
+    efc__J_rowadr: wp.array2d[int],
+    efc__J_rownnz: wp.array2d[int],
+    efc__Jqvel: wp.array2d[float],
+    efc__Ma: wp.array2d[float],
+    efc__aref: wp.array2d[float],
+    efc__force: wp.array2d[float],
+    efc__frictionloss: wp.array2d[float],
+    efc__iD: wp.array2d[float],
+    efc__iJ: wp.array3d[float],
+    efc__iJ_colind: wp.array3d[int],
+    efc__iJ_rowadr: wp.array2d[int],
+    efc__iJ_rownnz: wp.array2d[int],
+    efc__iaref: wp.array2d[float],
+    efc__id: wp.array2d[int],
+    efc__iforce: wp.array2d[float],
+    efc__ifrictionloss: wp.array2d[float],
+    efc__iid: wp.array2d[int],
+    efc__island: wp.array2d[int],
+    efc__istate: wp.array2d[int],
+    efc__itype: wp.array2d[int],
+    efc__margin: wp.array2d[float],
+    efc__pos: wp.array2d[float],
+    efc__state: wp.array2d[int],
+    efc__type: wp.array2d[int],
+    efc__vel: wp.array2d[float],
+):
+  _m.stat = _s
+  _m.opt = _o
+  _m.callback = _cb
+  _d.efc = _e
+  _d.contact = _c
+  _m.D_colind = D_colind
+  _m.D_diag = D_diag
+  _m.D_rowadr = D_rowadr
+  _m.D_rownnz = D_rownnz
+  _m.M_elemid = M_elemid
+  _m.M_fullm_i = M_fullm_i
+  _m.M_fullm_j = M_fullm_j
+  _m.M_fullm_upper_elemid = M_fullm_upper_elemid
+  _m.M_fullm_upper_i = M_fullm_upper_i
+  _m.M_fullm_upper_j = M_fullm_upper_j
+  _m.M_mulm_col = M_mulm_col
+  _m.M_mulm_madr = M_mulm_madr
+  _m.M_mulm_rowadr = M_mulm_rowadr
+  _m.M_rowadr = M_rowadr
+  _m.M_rownnz = M_rownnz
+  _m.M_tiles = M_tiles
+  _m.actuator_acc0 = actuator_acc0
+  _m.actuator_actadr = actuator_actadr
+  _m.actuator_actearly = actuator_actearly
+  _m.actuator_actlimited = actuator_actlimited
+  _m.actuator_actnum = actuator_actnum
+  _m.actuator_actrange = actuator_actrange
+  _m.actuator_biasprm = actuator_biasprm
+  _m.actuator_biastype = actuator_biastype
+  _m.actuator_cranklength = actuator_cranklength
+  _m.actuator_ctrllimited = actuator_ctrllimited
+  _m.actuator_ctrlrange = actuator_ctrlrange
+  _m.actuator_delay = actuator_delay
+  _m.actuator_dynprm = actuator_dynprm
+  _m.actuator_dyntype = actuator_dyntype
+  _m.actuator_forcelimited = actuator_forcelimited
+  _m.actuator_forcerange = actuator_forcerange
+  _m.actuator_gainprm = actuator_gainprm
+  _m.actuator_gaintype = actuator_gaintype
+  _m.actuator_gear = actuator_gear
+  _m.actuator_history = actuator_history
+  _m.actuator_historyadr = actuator_historyadr
+  _m.actuator_lengthrange = actuator_lengthrange
+  _m.actuator_trnid = actuator_trnid
+  _m.actuator_trntype = actuator_trntype
+  _m.actuator_trntype_body_adr = actuator_trntype_body_adr
+  _m.block_dim = block_dim
+  _m.body_branch_start = body_branch_start
+  _m.body_branches = body_branches
+  _m.body_dofadr = body_dofadr
+  _m.body_dofnum = body_dofnum
+  _m.body_fluid_ellipsoid = body_fluid_ellipsoid
+  _m.body_geomadr = body_geomadr
+  _m.body_geomnum = body_geomnum
+  _m.body_gravcomp = body_gravcomp
+  _m.body_inertia = body_inertia
+  _m.body_invweight0 = body_invweight0
+  _m.body_ipos = body_ipos
+  _m.body_iquat = body_iquat
+  _m.body_isdofancestor = body_isdofancestor
+  _m.body_jntadr = body_jntadr
+  _m.body_jntnum = body_jntnum
+  _m.body_mass = body_mass
+  _m.body_mocapid = body_mocapid
+  _m.body_parentid = body_parentid
+  _m.body_pos = body_pos
+  _m.body_quat = body_quat
+  _m.body_rootid = body_rootid
+  _m.body_subtreemass = body_subtreemass
+  _m.body_tree = body_tree
+  _m.body_treeid = body_treeid
+  _m.body_weldid = body_weldid
+  _m.cam_bodyid = cam_bodyid
+  _m.cam_fovy = cam_fovy
+  _m.cam_intrinsic = cam_intrinsic
+  _m.cam_mat0 = cam_mat0
+  _m.cam_mode = cam_mode
+  _m.cam_pos = cam_pos
+  _m.cam_pos0 = cam_pos0
+  _m.cam_poscom0 = cam_poscom0
+  _m.cam_quat = cam_quat
+  _m.cam_resolution = cam_resolution
+  _m.cam_sensorsize = cam_sensorsize
+  _m.cam_targetbodyid = cam_targetbodyid
+  _m.dof_armature = dof_armature
+  _m.dof_bodyid = dof_bodyid
+  _m.dof_damping = dof_damping
+  _m.dof_dampingpoly = dof_dampingpoly
+  _m.dof_frictionloss = dof_frictionloss
+  _m.dof_invweight0 = dof_invweight0
+  _m.dof_jntid = dof_jntid
+  _m.dof_parentid = dof_parentid
+  _m.dof_solimp = dof_solimp
+  _m.dof_solref = dof_solref
+  _m.dof_treeid = dof_treeid
+  _m.dof_tri_col = dof_tri_col
+  _m.dof_tri_row = dof_tri_row
+  _m.eq_connect_adr = eq_connect_adr
+  _m.eq_data = eq_data
+  _m.eq_flex_adr = eq_flex_adr
+  _m.eq_jnt_adr = eq_jnt_adr
+  _m.eq_obj1id = eq_obj1id
+  _m.eq_obj2id = eq_obj2id
+  _m.eq_objtype = eq_objtype
+  _m.eq_solimp = eq_solimp
+  _m.eq_solref = eq_solref
+  _m.eq_ten_adr = eq_ten_adr
+  _m.eq_type = eq_type
+  _m.eq_wld_adr = eq_wld_adr
+  _m.flex_bending = flex_bending
+  _m.flex_bendingadr = flex_bendingadr
+  _m.flex_centered = flex_centered
+  _m.flex_conaffinity = flex_conaffinity
+  _m.flex_condim = flex_condim
+  _m.flex_contype = flex_contype
+  _m.flex_damping = flex_damping
+  _m.flex_dim = flex_dim
+  _m.flex_edge = flex_edge
+  _m.flex_edgeadr = flex_edgeadr
+  _m.flex_edgeflap = flex_edgeflap
+  _m.flex_edgenum = flex_edgenum
+  _m.flex_elem = flex_elem
+  _m.flex_elemadr = flex_elemadr
+  _m.flex_elemdataadr = flex_elemdataadr
+  _m.flex_elemedge = flex_elemedge
+  _m.flex_elemedgeadr = flex_elemedgeadr
+  _m.flex_elemnum = flex_elemnum
+  _m.flex_friction = flex_friction
+  _m.flex_gap = flex_gap
+  _m.flex_margin = flex_margin
+  _m.flex_priority = flex_priority
+  _m.flex_radius = flex_radius
+  _m.flex_shell = flex_shell
+  _m.flex_shelldataadr = flex_shelldataadr
+  _m.flex_shellnum = flex_shellnum
+  _m.flex_solimp = flex_solimp
+  _m.flex_solmix = flex_solmix
+  _m.flex_solref = flex_solref
+  _m.flex_stiffness = flex_stiffness
+  _m.flex_stiffnessadr = flex_stiffnessadr
+  _m.flex_vert = flex_vert
+  _m.flex_vertadr = flex_vertadr
+  _m.flex_vertbodyid = flex_vertbodyid
+  _m.flex_vertflexid = flex_vertflexid
+  _m.flex_vertnum = flex_vertnum
+  _m.flexedge_J_colind = flexedge_J_colind
+  _m.flexedge_J_rowadr = flexedge_J_rowadr
+  _m.flexedge_J_rownnz = flexedge_J_rownnz
+  _m.flexedge_invweight0 = flexedge_invweight0
+  _m.flexedge_length0 = flexedge_length0
+  _m.geom_aabb = geom_aabb
+  _m.geom_bodyid = geom_bodyid
+  _m.geom_conaffinity = geom_conaffinity
+  _m.geom_condim = geom_condim
+  _m.geom_contype = geom_contype
+  _m.geom_dataid = geom_dataid
+  _m.geom_fluid = geom_fluid
+  _m.geom_friction = geom_friction
+  _m.geom_gap = geom_gap
+  _m.geom_group = geom_group
+  _m.geom_margin = geom_margin
+  _m.geom_matid = geom_matid
+  _m.geom_pair_type_count = geom_pair_type_count
+  _m.geom_plugin_index = geom_plugin_index
+  _m.geom_pos = geom_pos
+  _m.geom_priority = geom_priority
+  _m.geom_quat = geom_quat
+  _m.geom_rbound = geom_rbound
+  _m.geom_rgba = geom_rgba
+  _m.geom_size = geom_size
+  _m.geom_solimp = geom_solimp
+  _m.geom_solmix = geom_solmix
+  _m.geom_solref = geom_solref
+  _m.geom_type = geom_type
+  _m.has_fluid = has_fluid
+  _m.has_sdf_geom = has_sdf_geom
+  _m.hfield_adr = hfield_adr
+  _m.hfield_data = hfield_data
+  _m.hfield_ncol = hfield_ncol
+  _m.hfield_nrow = hfield_nrow
+  _m.hfield_size = hfield_size
+  _m.is_sparse = is_sparse
+  _m.jnt_actfrclimited = jnt_actfrclimited
+  _m.jnt_actfrcrange = jnt_actfrcrange
+  _m.jnt_actgravcomp = jnt_actgravcomp
+  _m.jnt_axis = jnt_axis
+  _m.jnt_bodyid = jnt_bodyid
+  _m.jnt_dofadr = jnt_dofadr
+  _m.jnt_limited_ball_adr = jnt_limited_ball_adr
+  _m.jnt_limited_slide_hinge_adr = jnt_limited_slide_hinge_adr
+  _m.jnt_margin = jnt_margin
+  _m.jnt_pos = jnt_pos
+  _m.jnt_qposadr = jnt_qposadr
+  _m.jnt_range = jnt_range
+  _m.jnt_solimp = jnt_solimp
+  _m.jnt_solref = jnt_solref
+  _m.jnt_stiffness = jnt_stiffness
+  _m.jnt_stiffnesspoly = jnt_stiffnesspoly
+  _m.jnt_type = jnt_type
+  _m.light_bodyid = light_bodyid
+  _m.light_dir = light_dir
+  _m.light_dir0 = light_dir0
+  _m.light_mode = light_mode
+  _m.light_pos = light_pos
+  _m.light_pos0 = light_pos0
+  _m.light_poscom0 = light_poscom0
+  _m.light_targetbodyid = light_targetbodyid
+  _m.mapM2D = mapM2D
+  _m.mat_rgba = mat_rgba
+  _m.max_ten_J_rownnz = max_ten_J_rownnz
+  _m.mesh_face = mesh_face
+  _m.mesh_faceadr = mesh_faceadr
+  _m.mesh_graph = mesh_graph
+  _m.mesh_graphadr = mesh_graphadr
+  _m.mesh_normal = mesh_normal
+  _m.mesh_normaladr = mesh_normaladr
+  _m.mesh_normalnum = mesh_normalnum
+  _m.mesh_octadr = mesh_octadr
+  _m.mesh_polyadr = mesh_polyadr
+  _m.mesh_polymap = mesh_polymap
+  _m.mesh_polymapadr = mesh_polymapadr
+  _m.mesh_polymapnum = mesh_polymapnum
+  _m.mesh_polynormal = mesh_polynormal
+  _m.mesh_polynum = mesh_polynum
+  _m.mesh_polyvert = mesh_polyvert
+  _m.mesh_polyvertadr = mesh_polyvertadr
+  _m.mesh_polyvertnum = mesh_polyvertnum
+  _m.mesh_quat = mesh_quat
+  _m.mesh_vert = mesh_vert
+  _m.mesh_vertadr = mesh_vertadr
+  _m.mesh_vertnum = mesh_vertnum
+  _m.nC = nC
+  _m.nD = nD
+  _m.nJten = nJten
+  _m.na = na
+  _m.nacttrnbody = nacttrnbody
+  _m.nbody = nbody
+  _m.nbranch = nbranch
+  _m.ncam = ncam
+  _m.neq = neq
+  _m.nflex = nflex
+  _m.nflexedge = nflexedge
+  _m.nflexelem = nflexelem
+  _m.nflexshelldata = nflexshelldata
+  _m.nflexvert = nflexvert
+  _m.ngeom = ngeom
+  _m.ngravcomp = ngravcomp
+  _m.nhistory = nhistory
+  _m.njnt = njnt
+  _m.nlight = nlight
+  _m.nmaxcondim = nmaxcondim
+  _m.nmaxmeshdeg = nmaxmeshdeg
+  _m.nmaxpolygon = nmaxpolygon
+  _m.nmaxpyramid = nmaxpyramid
+  _m.nmeshface = nmeshface
+  _m.nrangefinder = nrangefinder
+  _m.nsensorcollision = nsensorcollision
+  _m.nsensorcontact = nsensorcontact
+  _m.nsensortaxel = nsensortaxel
+  _m.nsite = nsite
+  _m.ntendon = ntendon
+  _m.ntree = ntree
+  _m.nu = nu
+  _m.nv = nv
+  _m.nv_pad = nv_pad
+  _m.nwrap = nwrap
+  _m.nxn_geom_pair_filtered = nxn_geom_pair_filtered
+  _m.nxn_pairid = nxn_pairid
+  _m.nxn_pairid_filtered = nxn_pairid_filtered
+  _m.oct_aabb = oct_aabb
+  _m.oct_child = oct_child
+  _m.oct_coeff = oct_coeff
+  _m.opt.broadphase = opt__broadphase
+  _m.opt.broadphase_filter = opt__broadphase_filter
+  _m.opt.ccd_iterations = opt__ccd_iterations
+  _m.opt.ccd_tolerance = opt__ccd_tolerance
+  _m.opt.cone = opt__cone
+  _m.opt.contact_sensor_maxmatch = opt__contact_sensor_maxmatch
+  _m.opt.density = opt__density
+  _m.opt.disableflags = opt__disableflags
+  _m.opt.enableflags = opt__enableflags
+  _m.opt.graph_conditional = opt__graph_conditional
+  _m.opt.gravity = opt__gravity
+  _m.opt.impratio_invsqrt = opt__impratio_invsqrt
+  _m.opt.integrator = opt__integrator
+  _m.opt.iterations = opt__iterations
+  _m.opt.ls_iterations = opt__ls_iterations
+  _m.opt.ls_parallel = opt__ls_parallel
+  _m.opt.ls_parallel_min_step = opt__ls_parallel_min_step
+  _m.opt.ls_tolerance = opt__ls_tolerance
+  _m.opt.magnetic = opt__magnetic
+  _m.opt.run_collision_detection = opt__run_collision_detection
+  _m.opt.sdf_initpoints = opt__sdf_initpoints
+  _m.opt.sdf_iterations = opt__sdf_iterations
+  _m.opt.solver = opt__solver
+  _m.opt.timestep = opt__timestep
+  _m.opt.tolerance = opt__tolerance
+  _m.opt.viscosity = opt__viscosity
+  _m.opt.wind = opt__wind
+  _m.pair_dim = pair_dim
+  _m.pair_friction = pair_friction
+  _m.pair_gap = pair_gap
+  _m.pair_margin = pair_margin
+  _m.pair_solimp = pair_solimp
+  _m.pair_solref = pair_solref
+  _m.pair_solreffriction = pair_solreffriction
+  _m.plugin = plugin
+  _m.plugin_attr = plugin_attr
+  _m.qD_fullm_i = qD_fullm_i
+  _m.qD_fullm_j = qD_fullm_j
+  _m.qLD_all_updates = qLD_all_updates
+  _m.qLD_level_offsets = qLD_level_offsets
+  _m.qLD_updates = qLD_updates
+  _m.qpos0 = qpos0
+  _m.qpos_spring = qpos_spring
+  _m.rangefinder_sensor_adr = rangefinder_sensor_adr
+  _m.sensor_acc_adr = sensor_acc_adr
+  _m.sensor_adr = sensor_adr
+  _m.sensor_adr_to_contact_adr = sensor_adr_to_contact_adr
+  _m.sensor_contact_adr = sensor_contact_adr
+  _m.sensor_cutoff = sensor_cutoff
+  _m.sensor_datatype = sensor_datatype
+  _m.sensor_delay = sensor_delay
+  _m.sensor_dim = sensor_dim
+  _m.sensor_e_kinetic = sensor_e_kinetic
+  _m.sensor_e_potential = sensor_e_potential
+  _m.sensor_history = sensor_history
+  _m.sensor_historyadr = sensor_historyadr
+  _m.sensor_interval = sensor_interval
+  _m.sensor_intprm = sensor_intprm
+  _m.sensor_limitfrc_adr = sensor_limitfrc_adr
+  _m.sensor_limitpos_adr = sensor_limitpos_adr
+  _m.sensor_limitvel_adr = sensor_limitvel_adr
+  _m.sensor_objid = sensor_objid
+  _m.sensor_objtype = sensor_objtype
+  _m.sensor_pos_adr = sensor_pos_adr
+  _m.sensor_rangefinder_adr = sensor_rangefinder_adr
+  _m.sensor_rangefinder_bodyid = sensor_rangefinder_bodyid
+  _m.sensor_refid = sensor_refid
+  _m.sensor_reftype = sensor_reftype
+  _m.sensor_rne_postconstraint = sensor_rne_postconstraint
+  _m.sensor_subtree_vel = sensor_subtree_vel
+  _m.sensor_tendonactfrc_adr = sensor_tendonactfrc_adr
+  _m.sensor_touch_adr = sensor_touch_adr
+  _m.sensor_type = sensor_type
+  _m.sensor_vel_adr = sensor_vel_adr
+  _m.site_bodyid = site_bodyid
+  _m.site_pos = site_pos
+  _m.site_quat = site_quat
+  _m.site_size = site_size
+  _m.site_type = site_type
+  _m.stat.meaninertia = stat__meaninertia
+  _m.taxel_sensorid = taxel_sensorid
+  _m.taxel_vertadr = taxel_vertadr
+  _m.ten_J_colind = ten_J_colind
+  _m.ten_J_rowadr = ten_J_rowadr
+  _m.ten_J_rownnz = ten_J_rownnz
+  _m.tendon_actfrclimited = tendon_actfrclimited
+  _m.tendon_actfrcrange = tendon_actfrcrange
+  _m.tendon_adr = tendon_adr
+  _m.tendon_armature = tendon_armature
+  _m.tendon_damping = tendon_damping
+  _m.tendon_dampingpoly = tendon_dampingpoly
+  _m.tendon_frictionloss = tendon_frictionloss
+  _m.tendon_geom_adr = tendon_geom_adr
+  _m.tendon_invweight0 = tendon_invweight0
+  _m.tendon_jnt_adr = tendon_jnt_adr
+  _m.tendon_length0 = tendon_length0
+  _m.tendon_lengthspring = tendon_lengthspring
+  _m.tendon_limited_adr = tendon_limited_adr
+  _m.tendon_margin = tendon_margin
+  _m.tendon_num = tendon_num
+  _m.tendon_range = tendon_range
+  _m.tendon_site_pair_adr = tendon_site_pair_adr
+  _m.tendon_solimp_fri = tendon_solimp_fri
+  _m.tendon_solimp_lim = tendon_solimp_lim
+  _m.tendon_solref_fri = tendon_solref_fri
+  _m.tendon_solref_lim = tendon_solref_lim
+  _m.tendon_stiffness = tendon_stiffness
+  _m.tendon_stiffnesspoly = tendon_stiffnesspoly
+  _m.wrap_geom_adr = wrap_geom_adr
+  _m.wrap_jnt_adr = wrap_jnt_adr
+  _m.wrap_objid = wrap_objid
+  _m.wrap_prm = wrap_prm
+  _m.wrap_pulley_scale = wrap_pulley_scale
+  _m.wrap_site_pair_adr = wrap_site_pair_adr
+  _m.wrap_type = wrap_type
+  _d.M = M
+  _d.act = act
+  _d.act_dot = act_dot
+  _d.actuator_force = actuator_force
+  _d.actuator_length = actuator_length
+  _d.actuator_moment = actuator_moment
+  _d.actuator_velocity = actuator_velocity
+  _d.cacc = cacc
+  _d.cam_xmat = cam_xmat
+  _d.cam_xpos = cam_xpos
+  _d.cdof = cdof
+  _d.cdof_dot = cdof_dot
+  _d.cfrc_ext = cfrc_ext
+  _d.cfrc_int = cfrc_int
+  _d.cinert = cinert
+  _d.contact.dim = contact__dim
+  _d.contact.dist = contact__dist
+  _d.contact.efc_address = contact__efc_address
+  _d.contact.flex = contact__flex
+  _d.contact.frame = contact__frame
+  _d.contact.friction = contact__friction
+  _d.contact.geom = contact__geom
+  _d.contact.geomcollisionid = contact__geomcollisionid
+  _d.contact.includemargin = contact__includemargin
+  _d.contact.pos = contact__pos
+  _d.contact.solimp = contact__solimp
+  _d.contact.solref = contact__solref
+  _d.contact.solreffriction = contact__solreffriction
+  _d.contact.type = contact__type
+  _d.contact.vert = contact__vert
+  _d.contact.worldid = contact__worldid
+  _d.crb = crb
+  _d.ctrl = ctrl
+  _d.cvel = cvel
+  _d.dof_island = dof_island
+  _d.dof_islandid = dof_islandid
+  _d.efc.D = efc__D
+  _d.efc.J = efc__J
+  _d.efc.J_colind = efc__J_colind
+  _d.efc.J_rowadr = efc__J_rowadr
+  _d.efc.J_rownnz = efc__J_rownnz
+  _d.efc.Jqvel = efc__Jqvel
+  _d.efc.Ma = efc__Ma
+  _d.efc.aref = efc__aref
+  _d.efc.force = efc__force
+  _d.efc.frictionloss = efc__frictionloss
+  _d.efc.iD = efc__iD
+  _d.efc.iJ = efc__iJ
+  _d.efc.iJ_colind = efc__iJ_colind
+  _d.efc.iJ_rowadr = efc__iJ_rowadr
+  _d.efc.iJ_rownnz = efc__iJ_rownnz
+  _d.efc.iaref = efc__iaref
+  _d.efc.id = efc__id
+  _d.efc.iforce = efc__iforce
+  _d.efc.ifrictionloss = efc__ifrictionloss
+  _d.efc.iid = efc__iid
+  _d.efc.island = efc__island
+  _d.efc.istate = efc__istate
+  _d.efc.itype = efc__itype
+  _d.efc.margin = efc__margin
+  _d.efc.pos = efc__pos
+  _d.efc.state = efc__state
+  _d.efc.type = efc__type
+  _d.efc.vel = efc__vel
+  _d.efc_islandid = efc_islandid
+  _d.energy = energy
+  _d.eq_active = eq_active
+  _d.flexedge_J = flexedge_J
+  _d.flexedge_length = flexedge_length
+  _d.flexedge_velocity = flexedge_velocity
+  _d.flexvert_xpos = flexvert_xpos
+  _d.geom_xmat = geom_xmat
+  _d.geom_xpos = geom_xpos
+  _d.history = history
+  _d.iqacc = iqacc
+  _d.iqacc_smooth = iqacc_smooth
+  _d.iqfrc_constraint = iqfrc_constraint
+  _d.iqfrc_smooth = iqfrc_smooth
+  _d.island_dofadr = island_dofadr
+  _d.island_efcadr = island_efcadr
+  _d.island_ne = island_ne
+  _d.island_nefc = island_nefc
+  _d.island_nf = island_nf
+  _d.island_nv = island_nv
+  _d.light_xdir = light_xdir
+  _d.light_xpos = light_xpos
+  _d.map_dof2idof = map_dof2idof
+  _d.map_efc2iefc = map_efc2iefc
+  _d.map_idof2dof = map_idof2dof
+  _d.map_iefc2efc = map_iefc2efc
+  _d.mocap_pos = mocap_pos
+  _d.mocap_quat = mocap_quat
+  _d.moment_colind = moment_colind
+  _d.moment_rowadr = moment_rowadr
+  _d.moment_rownnz = moment_rownnz
+  _d.naccdmax = naccdmax
+  _d.nacon = nacon
+  _d.naconmax = naconmax
+  _d.ncollision = ncollision
+  _d.ne = ne
+  _d.nefc = nefc
+  _d.nf = nf
+  _d.nidof = nidof
+  _d.nisland = nisland
+  _d.njmax = njmax
+  _d.njmax_nnz = njmax_nnz
+  _d.nl = nl
+  _d.qLD = qLD
+  _d.qLDiagInv = qLDiagInv
+  _d.qLU = qLU
+  _d.qacc = qacc
+  _d.qacc_smooth = qacc_smooth
+  _d.qacc_warmstart = qacc_warmstart
+  _d.qfrc_actuator = qfrc_actuator
+  _d.qfrc_applied = qfrc_applied
+  _d.qfrc_bias = qfrc_bias
+  _d.qfrc_constraint = qfrc_constraint
+  _d.qfrc_damper = qfrc_damper
+  _d.qfrc_fluid = qfrc_fluid
+  _d.qfrc_gravcomp = qfrc_gravcomp
+  _d.qfrc_passive = qfrc_passive
+  _d.qfrc_smooth = qfrc_smooth
+  _d.qfrc_spring = qfrc_spring
+  _d.qpos = qpos
+  _d.qvel = qvel
+  _d.sensordata = sensordata
+  _d.site_xmat = site_xmat
+  _d.site_xpos = site_xpos
+  _d.solver_niter = solver_niter
+  _d.subtree_angmom = subtree_angmom
+  _d.subtree_com = subtree_com
+  _d.subtree_linvel = subtree_linvel
+  _d.ten_J = ten_J
+  _d.ten_length = ten_length
+  _d.ten_velocity = ten_velocity
+  _d.ten_wrapadr = ten_wrapadr
+  _d.ten_wrapnum = ten_wrapnum
+  _d.time = time
+  _d.tree_island = tree_island
+  _d.wrap_obj = wrap_obj
+  _d.wrap_xpos = wrap_xpos
+  _d.xanchor = xanchor
+  _d.xaxis = xaxis
+  _d.xfrc_applied = xfrc_applied
+  _d.ximat = ximat
+  _d.xipos = xipos
+  _d.xmat = xmat
+  _d.xpos = xpos
+  _d.xquat = xquat
+  _d.nworld = nworld
+  mjwarp.step(_m, _d)
+
+
+def _step_jax_impl(m: types.Model, d: types.Data):
+  output_dims = {
+      'M': d._impl.M.shape,
+      'act': d.act.shape,
+      'act_dot': d.act_dot.shape,
+      'actuator_force': d.actuator_force.shape,
+      'actuator_length': d.actuator_length.shape,
+      'actuator_moment': d._impl.actuator_moment.shape,
+      'actuator_velocity': d._impl.actuator_velocity.shape,
+      'cacc': d._impl.cacc.shape,
+      'cam_xmat': d.cam_xmat.shape,
+      'cam_xpos': d.cam_xpos.shape,
+      'cdof': d.cdof.shape,
+      'cdof_dot': d.cdof_dot.shape,
+      'cfrc_ext': d._impl.cfrc_ext.shape,
+      'cfrc_int': d._impl.cfrc_int.shape,
+      'cinert': d._impl.cinert.shape,
+      'crb': d._impl.crb.shape,
+      'cvel': d.cvel.shape,
+      'dof_island': d._impl.dof_island.shape,
+      'dof_islandid': d._impl.dof_islandid.shape,
+      'efc_islandid': d._impl.efc_islandid.shape,
+      'energy': d._impl.energy.shape,
+      'flexedge_J': d._impl.flexedge_J.shape,
+      'flexedge_length': d._impl.flexedge_length.shape,
+      'flexedge_velocity': d._impl.flexedge_velocity.shape,
+      'flexvert_xpos': d._impl.flexvert_xpos.shape,
+      'geom_xmat': d.geom_xmat.shape,
+      'geom_xpos': d.geom_xpos.shape,
+      'history': d.history.shape,
+      'iqacc': d._impl.iqacc.shape,
+      'iqacc_smooth': d._impl.iqacc_smooth.shape,
+      'iqfrc_constraint': d._impl.iqfrc_constraint.shape,
+      'iqfrc_smooth': d._impl.iqfrc_smooth.shape,
+      'island_dofadr': d._impl.island_dofadr.shape,
+      'island_efcadr': d._impl.island_efcadr.shape,
+      'island_ne': d._impl.island_ne.shape,
+      'island_nefc': d._impl.island_nefc.shape,
+      'island_nf': d._impl.island_nf.shape,
+      'island_nv': d._impl.island_nv.shape,
+      'light_xdir': d._impl.light_xdir.shape,
+      'light_xpos': d._impl.light_xpos.shape,
+      'map_dof2idof': d._impl.map_dof2idof.shape,
+      'map_efc2iefc': d._impl.map_efc2iefc.shape,
+      'map_idof2dof': d._impl.map_idof2dof.shape,
+      'map_iefc2efc': d._impl.map_iefc2efc.shape,
+      'moment_colind': d._impl.moment_colind.shape,
+      'moment_rowadr': d._impl.moment_rowadr.shape,
+      'moment_rownnz': d._impl.moment_rownnz.shape,
+      'nacon': d._impl.nacon.shape,
+      'ncollision': d._impl.ncollision.shape,
+      'ne': d._impl.ne.shape,
+      'nefc': d._impl.nefc.shape,
+      'nf': d._impl.nf.shape,
+      'nidof': d._impl.nidof.shape,
+      'nisland': d._impl.nisland.shape,
+      'nl': d._impl.nl.shape,
+      'qLD': d._impl.qLD.shape,
+      'qLDiagInv': d._impl.qLDiagInv.shape,
+      'qLU': d._impl.qLU.shape,
+      'qacc': d.qacc.shape,
+      'qacc_smooth': d.qacc_smooth.shape,
+      'qacc_warmstart': d.qacc_warmstart.shape,
+      'qfrc_actuator': d.qfrc_actuator.shape,
+      'qfrc_bias': d.qfrc_bias.shape,
+      'qfrc_constraint': d.qfrc_constraint.shape,
+      'qfrc_damper': d._impl.qfrc_damper.shape,
+      'qfrc_fluid': d.qfrc_fluid.shape,
+      'qfrc_gravcomp': d.qfrc_gravcomp.shape,
+      'qfrc_passive': d.qfrc_passive.shape,
+      'qfrc_smooth': d.qfrc_smooth.shape,
+      'qfrc_spring': d._impl.qfrc_spring.shape,
+      'qpos': d.qpos.shape,
+      'qvel': d.qvel.shape,
+      'sensordata': d.sensordata.shape,
+      'site_xmat': d.site_xmat.shape,
+      'site_xpos': d.site_xpos.shape,
+      'solver_niter': d._impl.solver_niter.shape,
+      'subtree_angmom': d._impl.subtree_angmom.shape,
+      'subtree_com': d.subtree_com.shape,
+      'subtree_linvel': d._impl.subtree_linvel.shape,
+      'ten_J': d._impl.ten_J.shape,
+      'ten_length': d.ten_length.shape,
+      'ten_velocity': d._impl.ten_velocity.shape,
+      'ten_wrapadr': d._impl.ten_wrapadr.shape,
+      'ten_wrapnum': d._impl.ten_wrapnum.shape,
+      'time': d.time.shape,
+      'tree_island': d._impl.tree_island.shape,
+      'wrap_obj': d._impl.wrap_obj.shape,
+      'wrap_xpos': d._impl.wrap_xpos.shape,
+      'xanchor': d.xanchor.shape,
+      'xaxis': d.xaxis.shape,
+      'ximat': d.ximat.shape,
+      'xipos': d.xipos.shape,
+      'xmat': d.xmat.shape,
+      'xpos': d.xpos.shape,
+      'xquat': d.xquat.shape,
+      'contact__dim': d._impl.contact__dim.shape,
+      'contact__dist': d._impl.contact__dist.shape,
+      'contact__efc_address': d._impl.contact__efc_address.shape,
+      'contact__flex': d._impl.contact__flex.shape,
+      'contact__frame': d._impl.contact__frame.shape,
+      'contact__friction': d._impl.contact__friction.shape,
+      'contact__geom': d._impl.contact__geom.shape,
+      'contact__geomcollisionid': d._impl.contact__geomcollisionid.shape,
+      'contact__includemargin': d._impl.contact__includemargin.shape,
+      'contact__pos': d._impl.contact__pos.shape,
+      'contact__solimp': d._impl.contact__solimp.shape,
+      'contact__solref': d._impl.contact__solref.shape,
+      'contact__solreffriction': d._impl.contact__solreffriction.shape,
+      'contact__type': d._impl.contact__type.shape,
+      'contact__vert': d._impl.contact__vert.shape,
+      'contact__worldid': d._impl.contact__worldid.shape,
+      'efc__D': d._impl.efc__D.shape,
+      'efc__J': d._impl.efc__J.shape,
+      'efc__J_colind': d._impl.efc__J_colind.shape,
+      'efc__J_rowadr': d._impl.efc__J_rowadr.shape,
+      'efc__J_rownnz': d._impl.efc__J_rownnz.shape,
+      'efc__Jqvel': d._impl.efc__Jqvel.shape,
+      'efc__Ma': d._impl.efc__Ma.shape,
+      'efc__aref': d._impl.efc__aref.shape,
+      'efc__force': d._impl.efc__force.shape,
+      'efc__frictionloss': d._impl.efc__frictionloss.shape,
+      'efc__iD': d._impl.efc__iD.shape,
+      'efc__iJ': d._impl.efc__iJ.shape,
+      'efc__iJ_colind': d._impl.efc__iJ_colind.shape,
+      'efc__iJ_rowadr': d._impl.efc__iJ_rowadr.shape,
+      'efc__iJ_rownnz': d._impl.efc__iJ_rownnz.shape,
+      'efc__iaref': d._impl.efc__iaref.shape,
+      'efc__id': d._impl.efc__id.shape,
+      'efc__iforce': d._impl.efc__iforce.shape,
+      'efc__ifrictionloss': d._impl.efc__ifrictionloss.shape,
+      'efc__iid': d._impl.efc__iid.shape,
+      'efc__island': d._impl.efc__island.shape,
+      'efc__istate': d._impl.efc__istate.shape,
+      'efc__itype': d._impl.efc__itype.shape,
+      'efc__margin': d._impl.efc__margin.shape,
+      'efc__pos': d._impl.efc__pos.shape,
+      'efc__state': d._impl.efc__state.shape,
+      'efc__type': d._impl.efc__type.shape,
+      'efc__vel': d._impl.efc__vel.shape,
+  }
+  jf = ffi.jax_callable_variadic_tuple(
+      _step_shim,
+      num_outputs=139,
+      output_dims=output_dims,
+      vmap_method=None,
+      in_out_argnames=set([
+          'M',
+          'act',
+          'act_dot',
+          'actuator_force',
+          'actuator_length',
+          'actuator_moment',
+          'actuator_velocity',
+          'cacc',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'cfrc_ext',
+          'cfrc_int',
+          'cinert',
+          'crb',
+          'cvel',
+          'dof_island',
+          'dof_islandid',
+          'efc_islandid',
+          'energy',
+          'flexedge_J',
+          'flexedge_length',
+          'flexedge_velocity',
+          'flexvert_xpos',
+          'geom_xmat',
+          'geom_xpos',
+          'history',
+          'iqacc',
+          'iqacc_smooth',
+          'iqfrc_constraint',
+          'iqfrc_smooth',
+          'island_dofadr',
+          'island_efcadr',
+          'island_ne',
+          'island_nefc',
+          'island_nf',
+          'island_nv',
+          'light_xdir',
+          'light_xpos',
+          'map_dof2idof',
+          'map_efc2iefc',
+          'map_idof2dof',
+          'map_iefc2efc',
+          'moment_colind',
+          'moment_rowadr',
+          'moment_rownnz',
+          'nacon',
+          'ncollision',
+          'ne',
+          'nefc',
+          'nf',
+          'nidof',
+          'nisland',
+          'nl',
+          'qLD',
+          'qLDiagInv',
+          'qLU',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_damper',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qfrc_spring',
+          'qpos',
+          'qvel',
+          'sensordata',
+          'site_xmat',
+          'site_xpos',
+          'solver_niter',
+          'subtree_angmom',
+          'subtree_com',
+          'subtree_linvel',
+          'ten_J',
+          'ten_length',
+          'ten_velocity',
+          'ten_wrapadr',
+          'ten_wrapnum',
+          'time',
+          'tree_island',
+          'wrap_obj',
+          'wrap_xpos',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+          'contact__dim',
+          'contact__dist',
+          'contact__efc_address',
+          'contact__flex',
+          'contact__frame',
+          'contact__friction',
+          'contact__geom',
+          'contact__geomcollisionid',
+          'contact__includemargin',
+          'contact__pos',
+          'contact__solimp',
+          'contact__solref',
+          'contact__solreffriction',
+          'contact__type',
+          'contact__vert',
+          'contact__worldid',
+          'efc__D',
+          'efc__J',
+          'efc__J_colind',
+          'efc__J_rowadr',
+          'efc__J_rownnz',
+          'efc__Jqvel',
+          'efc__Ma',
+          'efc__aref',
+          'efc__force',
+          'efc__frictionloss',
+          'efc__iD',
+          'efc__iJ',
+          'efc__iJ_colind',
+          'efc__iJ_rowadr',
+          'efc__iJ_rownnz',
+          'efc__iaref',
+          'efc__id',
+          'efc__iforce',
+          'efc__ifrictionloss',
+          'efc__iid',
+          'efc__island',
+          'efc__istate',
+          'efc__itype',
+          'efc__margin',
+          'efc__pos',
+          'efc__state',
+          'efc__type',
+          'efc__vel',
+      ]),
+      stage_in_argnames=set([
+          'act',
+          'act_dot',
+          'actuator_acc0',
+          'actuator_actrange',
+          'actuator_biasprm',
+          'actuator_cranklength',
+          'actuator_ctrlrange',
+          'actuator_dynprm',
+          'actuator_force',
+          'actuator_forcerange',
+          'actuator_gainprm',
+          'actuator_gear',
+          'actuator_length',
+          'actuator_lengthrange',
+          'body_gravcomp',
+          'body_inertia',
+          'body_invweight0',
+          'body_ipos',
+          'body_iquat',
+          'body_mass',
+          'body_pos',
+          'body_quat',
+          'body_subtreemass',
+          'cam_fovy',
+          'cam_intrinsic',
+          'cam_mat0',
+          'cam_pos',
+          'cam_pos0',
+          'cam_poscom0',
+          'cam_quat',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'ctrl',
+          'cvel',
+          'dof_armature',
+          'dof_damping',
+          'dof_dampingpoly',
+          'dof_frictionloss',
+          'dof_invweight0',
+          'dof_solimp',
+          'dof_solref',
+          'eq_active',
+          'eq_data',
+          'eq_solimp',
+          'eq_solref',
+          'geom_aabb',
+          'geom_friction',
+          'geom_gap',
+          'geom_margin',
+          'geom_matid',
+          'geom_pos',
+          'geom_quat',
+          'geom_rbound',
+          'geom_rgba',
+          'geom_size',
+          'geom_solimp',
+          'geom_solmix',
+          'geom_solref',
+          'geom_xmat',
+          'geom_xpos',
+          'hfield_data',
+          'history',
+          'jnt_actfrcrange',
+          'jnt_axis',
+          'jnt_margin',
+          'jnt_pos',
+          'jnt_range',
+          'jnt_solimp',
+          'jnt_solref',
+          'jnt_stiffness',
+          'jnt_stiffnesspoly',
+          'light_dir',
+          'light_dir0',
+          'light_pos',
+          'light_pos0',
+          'light_poscom0',
+          'mat_rgba',
+          'mocap_pos',
+          'mocap_quat',
+          'opt__density',
+          'opt__gravity',
+          'opt__ls_tolerance',
+          'opt__magnetic',
+          'opt__timestep',
+          'opt__tolerance',
+          'opt__viscosity',
+          'opt__wind',
+          'pair_friction',
+          'pair_gap',
+          'pair_margin',
+          'pair_solimp',
+          'pair_solref',
+          'pair_solreffriction',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_applied',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qpos',
+          'qpos0',
+          'qpos_spring',
+          'qvel',
+          'sensordata',
+          'site_pos',
+          'site_quat',
+          'site_xmat',
+          'site_xpos',
+          'subtree_com',
+          'ten_length',
+          'tendon_actfrcrange',
+          'tendon_armature',
+          'tendon_damping',
+          'tendon_dampingpoly',
+          'tendon_frictionloss',
+          'tendon_invweight0',
+          'tendon_length0',
+          'tendon_lengthspring',
+          'tendon_margin',
+          'tendon_range',
+          'tendon_solimp_fri',
+          'tendon_solimp_lim',
+          'tendon_solref_fri',
+          'tendon_solref_lim',
+          'tendon_stiffness',
+          'tendon_stiffnesspoly',
+          'time',
+          'xanchor',
+          'xaxis',
+          'xfrc_applied',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      ]),
+      stage_out_argnames=set([
+          'act',
+          'act_dot',
+          'actuator_force',
+          'actuator_length',
+          'cam_xmat',
+          'cam_xpos',
+          'cdof',
+          'cdof_dot',
+          'cvel',
+          'geom_xmat',
+          'geom_xpos',
+          'history',
+          'qacc',
+          'qacc_smooth',
+          'qacc_warmstart',
+          'qfrc_actuator',
+          'qfrc_bias',
+          'qfrc_constraint',
+          'qfrc_fluid',
+          'qfrc_gravcomp',
+          'qfrc_passive',
+          'qfrc_smooth',
+          'qpos',
+          'qvel',
+          'sensordata',
+          'site_xmat',
+          'site_xpos',
+          'subtree_com',
+          'ten_length',
+          'time',
+          'xanchor',
+          'xaxis',
+          'ximat',
+          'xipos',
+          'xmat',
+          'xpos',
+          'xquat',
+      ]),
+      graph_mode=m.opt._impl.graph_mode,
+      has_side_effect=False,
+  )
+  out = jf(
+      d.qpos.shape[0],
+      m._impl.D_colind,
+      m._impl.D_diag,
+      m._impl.D_rowadr,
+      m._impl.D_rownnz,
+      m._impl.M_elemid,
+      m._impl.M_fullm_i,
+      m._impl.M_fullm_j,
+      m._impl.M_fullm_upper_elemid,
+      m._impl.M_fullm_upper_i,
+      m._impl.M_fullm_upper_j,
+      m._impl.M_mulm_col,
+      m._impl.M_mulm_madr,
+      m._impl.M_mulm_rowadr,
+      m._impl.M_rowadr,
+      m._impl.M_rownnz,
+      m._impl.M_tiles,
+      m.actuator_acc0,
+      m.actuator_actadr,
+      m.actuator_actearly,
+      m.actuator_actlimited,
+      m.actuator_actnum,
+      m.actuator_actrange,
+      m.actuator_biasprm,
+      m.actuator_biastype,
+      m.actuator_cranklength,
+      m.actuator_ctrllimited,
+      m.actuator_ctrlrange,
+      m._impl.actuator_delay,
+      m.actuator_dynprm,
+      m.actuator_dyntype,
+      m.actuator_forcelimited,
+      m.actuator_forcerange,
+      m.actuator_gainprm,
+      m.actuator_gaintype,
+      m.actuator_gear,
+      m._impl.actuator_history,
+      m._impl.actuator_historyadr,
+      m.actuator_lengthrange,
+      m.actuator_trnid,
+      m.actuator_trntype,
+      m._impl.actuator_trntype_body_adr,
+      m._impl.block_dim,
+      m._impl.body_branch_start,
+      m._impl.body_branches,
+      m.body_dofadr,
+      m.body_dofnum,
+      m._impl.body_fluid_ellipsoid,
+      m.body_geomadr,
+      m.body_geomnum,
+      m.body_gravcomp,
+      m.body_inertia,
+      m.body_invweight0,
+      m.body_ipos,
+      m.body_iquat,
+      m._impl.body_isdofancestor,
+      m.body_jntadr,
+      m.body_jntnum,
+      m.body_mass,
+      m.body_mocapid,
+      m.body_parentid,
+      m.body_pos,
+      m.body_quat,
+      m.body_rootid,
+      m.body_subtreemass,
+      m._impl.body_tree,
+      m.body_treeid,
+      m.body_weldid,
+      m.cam_bodyid,
+      m.cam_fovy,
+      m.cam_intrinsic,
+      m.cam_mat0,
+      m.cam_mode,
+      m.cam_pos,
+      m.cam_pos0,
+      m.cam_poscom0,
+      m.cam_quat,
+      m.cam_resolution,
+      m.cam_sensorsize,
+      m.cam_targetbodyid,
+      m.dof_armature,
+      m.dof_bodyid,
+      m.dof_damping,
+      m.dof_dampingpoly,
+      m.dof_frictionloss,
+      m.dof_invweight0,
+      m.dof_jntid,
+      m.dof_parentid,
+      m.dof_solimp,
+      m.dof_solref,
+      m.dof_treeid,
+      m._impl.dof_tri_col,
+      m._impl.dof_tri_row,
+      m._impl.eq_connect_adr,
+      m.eq_data,
+      m._impl.eq_flex_adr,
+      m._impl.eq_jnt_adr,
+      m.eq_obj1id,
+      m.eq_obj2id,
+      m.eq_objtype,
+      m.eq_solimp,
+      m.eq_solref,
+      m._impl.eq_ten_adr,
+      m.eq_type,
+      m._impl.eq_wld_adr,
+      m._impl.flex_bending,
+      m._impl.flex_bendingadr,
+      m._impl.flex_centered,
+      m._impl.flex_conaffinity,
+      m._impl.flex_condim,
+      m._impl.flex_contype,
+      m._impl.flex_damping,
+      m._impl.flex_dim,
+      m._impl.flex_edge,
+      m._impl.flex_edgeadr,
+      m._impl.flex_edgeflap,
+      m._impl.flex_edgenum,
+      m._impl.flex_elem,
+      m._impl.flex_elemadr,
+      m._impl.flex_elemdataadr,
+      m._impl.flex_elemedge,
+      m._impl.flex_elemedgeadr,
+      m._impl.flex_elemnum,
+      m._impl.flex_friction,
+      m._impl.flex_gap,
+      m._impl.flex_margin,
+      m._impl.flex_priority,
+      m._impl.flex_radius,
+      m._impl.flex_shell,
+      m._impl.flex_shelldataadr,
+      m._impl.flex_shellnum,
+      m._impl.flex_solimp,
+      m._impl.flex_solmix,
+      m._impl.flex_solref,
+      m._impl.flex_stiffness,
+      m._impl.flex_stiffnessadr,
+      m._impl.flex_vert,
+      m.flex_vertadr,
+      m._impl.flex_vertbodyid,
+      m._impl.flex_vertflexid,
+      m.flex_vertnum,
+      m._impl.flexedge_J_colind,
+      m._impl.flexedge_J_rowadr,
+      m._impl.flexedge_J_rownnz,
+      m._impl.flexedge_invweight0,
+      m._impl.flexedge_length0,
+      m.geom_aabb,
+      m.geom_bodyid,
+      m.geom_conaffinity,
+      m.geom_condim,
+      m.geom_contype,
+      jax.numpy.expand_dims(m.geom_dataid, 0),
+      m.geom_fluid,
+      m.geom_friction,
+      m.geom_gap,
+      m.geom_group,
+      m.geom_margin,
+      m.geom_matid,
+      m._impl.geom_pair_type_count,
+      m._impl.geom_plugin_index,
+      m.geom_pos,
+      m.geom_priority,
+      m.geom_quat,
+      m.geom_rbound,
+      m.geom_rgba,
+      m.geom_size,
+      m.geom_solimp,
+      m.geom_solmix,
+      m.geom_solref,
+      m.geom_type,
+      m._impl.has_fluid,
+      m._impl.has_sdf_geom,
+      m.hfield_adr,
+      m.hfield_data,
+      m.hfield_ncol,
+      m.hfield_nrow,
+      m.hfield_size,
+      m._impl.is_sparse,
+      m.jnt_actfrclimited,
+      m.jnt_actfrcrange,
+      m.jnt_actgravcomp,
+      m.jnt_axis,
+      m.jnt_bodyid,
+      m.jnt_dofadr,
+      m._impl.jnt_limited_ball_adr,
+      m._impl.jnt_limited_slide_hinge_adr,
+      m.jnt_margin,
+      m.jnt_pos,
+      m.jnt_qposadr,
+      m.jnt_range,
+      m.jnt_solimp,
+      m.jnt_solref,
+      m.jnt_stiffness,
+      m.jnt_stiffnesspoly,
+      m.jnt_type,
+      m._impl.light_bodyid,
+      m.light_dir,
+      m.light_dir0,
+      m.light_mode,
+      m.light_pos,
+      m.light_pos0,
+      m.light_poscom0,
+      m._impl.light_targetbodyid,
+      m._impl.mapM2D,
+      m.mat_rgba,
+      m._impl.max_ten_J_rownnz,
+      m.mesh_face,
+      m.mesh_faceadr,
+      m.mesh_graph,
+      m.mesh_graphadr,
+      m.mesh_normal,
+      m.mesh_normaladr,
+      m.mesh_normalnum,
+      m.mesh_octadr,
+      m._impl.mesh_polyadr,
+      m._impl.mesh_polymap,
+      m._impl.mesh_polymapadr,
+      m._impl.mesh_polymapnum,
+      m._impl.mesh_polynormal,
+      m._impl.mesh_polynum,
+      m._impl.mesh_polyvert,
+      m._impl.mesh_polyvertadr,
+      m._impl.mesh_polyvertnum,
+      m.mesh_quat,
+      m.mesh_vert,
+      m.mesh_vertadr,
+      m.mesh_vertnum,
+      m.nC,
+      m.nD,
+      m.nJten,
+      m.na,
+      m._impl.nacttrnbody,
+      m.nbody,
+      m._impl.nbranch,
+      m.ncam,
+      m.neq,
+      m.nflex,
+      m._impl.nflexedge,
+      m._impl.nflexelem,
+      m._impl.nflexshelldata,
+      m._impl.nflexvert,
+      m.ngeom,
+      m.ngravcomp,
+      m.nhistory,
+      m.njnt,
+      m.nlight,
+      m._impl.nmaxcondim,
+      m._impl.nmaxmeshdeg,
+      m._impl.nmaxpolygon,
+      m._impl.nmaxpyramid,
+      m.nmeshface,
+      m._impl.nrangefinder,
+      m._impl.nsensorcollision,
+      m._impl.nsensorcontact,
+      m._impl.nsensortaxel,
+      m.nsite,
+      m.ntendon,
+      m._impl.ntree,
+      m.nu,
+      m.nv,
+      m._impl.nv_pad,
+      m.nwrap,
+      m._impl.nxn_geom_pair_filtered,
+      m._impl.nxn_pairid,
+      m._impl.nxn_pairid_filtered,
+      m._impl.oct_aabb,
+      m._impl.oct_child,
+      m._impl.oct_coeff,
+      m.pair_dim,
+      m.pair_friction,
+      m.pair_gap,
+      m.pair_margin,
+      m.pair_solimp,
+      m.pair_solref,
+      m.pair_solreffriction,
+      m._impl.plugin,
+      m._impl.plugin_attr,
+      m._impl.qD_fullm_i,
+      m._impl.qD_fullm_j,
+      m._impl.qLD_all_updates,
+      m._impl.qLD_level_offsets,
+      m._impl.qLD_updates,
+      m.qpos0,
+      m.qpos_spring,
+      m._impl.rangefinder_sensor_adr,
+      m._impl.sensor_acc_adr,
+      m.sensor_adr,
+      m._impl.sensor_adr_to_contact_adr,
+      m._impl.sensor_contact_adr,
+      m.sensor_cutoff,
+      m.sensor_datatype,
+      m._impl.sensor_delay,
+      m.sensor_dim,
+      m._impl.sensor_e_kinetic,
+      m._impl.sensor_e_potential,
+      m._impl.sensor_history,
+      m._impl.sensor_historyadr,
+      m._impl.sensor_interval,
+      m.sensor_intprm,
+      m._impl.sensor_limitfrc_adr,
+      m._impl.sensor_limitpos_adr,
+      m._impl.sensor_limitvel_adr,
+      m.sensor_objid,
+      m.sensor_objtype,
+      m._impl.sensor_pos_adr,
+      m._impl.sensor_rangefinder_adr,
+      m._impl.sensor_rangefinder_bodyid,
+      m.sensor_refid,
+      m.sensor_reftype,
+      m._impl.sensor_rne_postconstraint,
+      m._impl.sensor_subtree_vel,
+      m._impl.sensor_tendonactfrc_adr,
+      m._impl.sensor_touch_adr,
+      m.sensor_type,
+      m._impl.sensor_vel_adr,
+      m.site_bodyid,
+      m.site_pos,
+      m.site_quat,
+      m.site_size,
+      m.site_type,
+      m._impl.taxel_sensorid,
+      m._impl.taxel_vertadr,
+      m._impl.ten_J_colind,
+      m._impl.ten_J_rowadr,
+      m._impl.ten_J_rownnz,
+      m.tendon_actfrclimited,
+      m.tendon_actfrcrange,
+      m.tendon_adr,
+      m.tendon_armature,
+      m.tendon_damping,
+      m.tendon_dampingpoly,
+      m.tendon_frictionloss,
+      m._impl.tendon_geom_adr,
+      m.tendon_invweight0,
+      m._impl.tendon_jnt_adr,
+      m.tendon_length0,
+      m.tendon_lengthspring,
+      m._impl.tendon_limited_adr,
+      m.tendon_margin,
+      m.tendon_num,
+      m.tendon_range,
+      m._impl.tendon_site_pair_adr,
+      m.tendon_solimp_fri,
+      m.tendon_solimp_lim,
+      m.tendon_solref_fri,
+      m.tendon_solref_lim,
+      m.tendon_stiffness,
+      m.tendon_stiffnesspoly,
+      m._impl.wrap_geom_adr,
+      m._impl.wrap_jnt_adr,
+      m.wrap_objid,
+      m.wrap_prm,
+      m._impl.wrap_pulley_scale,
+      m._impl.wrap_site_pair_adr,
+      m.wrap_type,
+      m.opt._impl.broadphase,
+      m.opt._impl.broadphase_filter,
+      m.opt._impl.ccd_iterations,
+      m.opt._impl.ccd_tolerance,
+      m.opt.cone,
+      m.opt._impl.contact_sensor_maxmatch,
+      m.opt.density,
+      m.opt.disableflags,
+      m.opt.enableflags,
+      m.opt._impl.graph_conditional,
+      m.opt.gravity,
+      m.opt._impl.impratio_invsqrt,
+      m.opt.integrator,
+      m.opt.iterations,
+      m.opt.ls_iterations,
+      m.opt._impl.ls_parallel,
+      m.opt._impl.ls_parallel_min_step,
+      m.opt.ls_tolerance,
+      m.opt.magnetic,
+      m.opt._impl.run_collision_detection,
+      m.opt._impl.sdf_initpoints,
+      m.opt._impl.sdf_iterations,
+      m.opt.solver,
+      m.opt.timestep,
+      m.opt.tolerance,
+      m.opt.viscosity,
+      m.opt.wind,
+      m.stat.meaninertia,
+      d._impl.naccdmax,
+      d._impl.naconmax,
+      d._impl.njmax,
+      d._impl.njmax_nnz,
+      d._impl.M,
+      d.act,
+      d.act_dot,
+      d.actuator_force,
+      d.actuator_length,
+      d._impl.actuator_moment,
+      d._impl.actuator_velocity,
+      d._impl.cacc,
+      d.cam_xmat,
+      d.cam_xpos,
+      d.cdof,
+      d.cdof_dot,
+      d._impl.cfrc_ext,
+      d._impl.cfrc_int,
+      d._impl.cinert,
+      d._impl.crb,
+      d.ctrl,
+      d.cvel,
+      d._impl.dof_island,
+      d._impl.dof_islandid,
+      d._impl.efc_islandid,
+      d._impl.energy,
+      d.eq_active,
+      d._impl.flexedge_J,
+      d._impl.flexedge_length,
+      d._impl.flexedge_velocity,
+      d._impl.flexvert_xpos,
+      d.geom_xmat,
+      d.geom_xpos,
+      d.history,
+      d._impl.iqacc,
+      d._impl.iqacc_smooth,
+      d._impl.iqfrc_constraint,
+      d._impl.iqfrc_smooth,
+      d._impl.island_dofadr,
+      d._impl.island_efcadr,
+      d._impl.island_ne,
+      d._impl.island_nefc,
+      d._impl.island_nf,
+      d._impl.island_nv,
+      d._impl.light_xdir,
+      d._impl.light_xpos,
+      d._impl.map_dof2idof,
+      d._impl.map_efc2iefc,
+      d._impl.map_idof2dof,
+      d._impl.map_iefc2efc,
+      d.mocap_pos,
+      d.mocap_quat,
+      d._impl.moment_colind,
+      d._impl.moment_rowadr,
+      d._impl.moment_rownnz,
+      d._impl.nacon,
+      d._impl.ncollision,
+      d._impl.ne,
+      d._impl.nefc,
+      d._impl.nf,
+      d._impl.nidof,
+      d._impl.nisland,
+      d._impl.nl,
+      d._impl.qLD,
+      d._impl.qLDiagInv,
+      d._impl.qLU,
+      d.qacc,
+      d.qacc_smooth,
+      d.qacc_warmstart,
+      d.qfrc_actuator,
+      d.qfrc_applied,
+      d.qfrc_bias,
+      d.qfrc_constraint,
+      d._impl.qfrc_damper,
+      d.qfrc_fluid,
+      d.qfrc_gravcomp,
+      d.qfrc_passive,
+      d.qfrc_smooth,
+      d._impl.qfrc_spring,
+      d.qpos,
+      d.qvel,
+      d.sensordata,
+      d.site_xmat,
+      d.site_xpos,
+      d._impl.solver_niter,
+      d._impl.subtree_angmom,
+      d.subtree_com,
+      d._impl.subtree_linvel,
+      d._impl.ten_J,
+      d.ten_length,
+      d._impl.ten_velocity,
+      d._impl.ten_wrapadr,
+      d._impl.ten_wrapnum,
+      d.time,
+      d._impl.tree_island,
+      d._impl.wrap_obj,
+      d._impl.wrap_xpos,
+      d.xanchor,
+      d.xaxis,
+      d.xfrc_applied,
+      d.ximat,
+      d.xipos,
+      d.xmat,
+      d.xpos,
+      d.xquat,
+      d._impl.contact__dim,
+      d._impl.contact__dist,
+      d._impl.contact__efc_address,
+      d._impl.contact__flex,
+      d._impl.contact__frame,
+      d._impl.contact__friction,
+      d._impl.contact__geom,
+      d._impl.contact__geomcollisionid,
+      d._impl.contact__includemargin,
+      d._impl.contact__pos,
+      d._impl.contact__solimp,
+      d._impl.contact__solref,
+      d._impl.contact__solreffriction,
+      d._impl.contact__type,
+      d._impl.contact__vert,
+      d._impl.contact__worldid,
+      d._impl.efc__D,
+      d._impl.efc__J,
+      d._impl.efc__J_colind,
+      d._impl.efc__J_rowadr,
+      d._impl.efc__J_rownnz,
+      d._impl.efc__Jqvel,
+      d._impl.efc__Ma,
+      d._impl.efc__aref,
+      d._impl.efc__force,
+      d._impl.efc__frictionloss,
+      d._impl.efc__iD,
+      d._impl.efc__iJ,
+      d._impl.efc__iJ_colind,
+      d._impl.efc__iJ_rowadr,
+      d._impl.efc__iJ_rownnz,
+      d._impl.efc__iaref,
+      d._impl.efc__id,
+      d._impl.efc__iforce,
+      d._impl.efc__ifrictionloss,
+      d._impl.efc__iid,
+      d._impl.efc__island,
+      d._impl.efc__istate,
+      d._impl.efc__itype,
+      d._impl.efc__margin,
+      d._impl.efc__pos,
+      d._impl.efc__state,
+      d._impl.efc__type,
+      d._impl.efc__vel,
+  )
+  d = d.tree_replace({
+      '_impl.M': out[0],
+      'act': out[1],
+      'act_dot': out[2],
+      'actuator_force': out[3],
+      'actuator_length': out[4],
+      '_impl.actuator_moment': out[5],
+      '_impl.actuator_velocity': out[6],
+      '_impl.cacc': out[7],
+      'cam_xmat': out[8],
+      'cam_xpos': out[9],
+      'cdof': out[10],
+      'cdof_dot': out[11],
+      '_impl.cfrc_ext': out[12],
+      '_impl.cfrc_int': out[13],
+      '_impl.cinert': out[14],
+      '_impl.crb': out[15],
+      'cvel': out[16],
+      '_impl.dof_island': out[17],
+      '_impl.dof_islandid': out[18],
+      '_impl.efc_islandid': out[19],
+      '_impl.energy': out[20],
+      '_impl.flexedge_J': out[21],
+      '_impl.flexedge_length': out[22],
+      '_impl.flexedge_velocity': out[23],
+      '_impl.flexvert_xpos': out[24],
+      'geom_xmat': out[25],
+      'geom_xpos': out[26],
+      'history': out[27],
+      '_impl.iqacc': out[28],
+      '_impl.iqacc_smooth': out[29],
+      '_impl.iqfrc_constraint': out[30],
+      '_impl.iqfrc_smooth': out[31],
+      '_impl.island_dofadr': out[32],
+      '_impl.island_efcadr': out[33],
+      '_impl.island_ne': out[34],
+      '_impl.island_nefc': out[35],
+      '_impl.island_nf': out[36],
+      '_impl.island_nv': out[37],
+      '_impl.light_xdir': out[38],
+      '_impl.light_xpos': out[39],
+      '_impl.map_dof2idof': out[40],
+      '_impl.map_efc2iefc': out[41],
+      '_impl.map_idof2dof': out[42],
+      '_impl.map_iefc2efc': out[43],
+      '_impl.moment_colind': out[44],
+      '_impl.moment_rowadr': out[45],
+      '_impl.moment_rownnz': out[46],
+      '_impl.nacon': out[47],
+      '_impl.ncollision': out[48],
+      '_impl.ne': out[49],
+      '_impl.nefc': out[50],
+      '_impl.nf': out[51],
+      '_impl.nidof': out[52],
+      '_impl.nisland': out[53],
+      '_impl.nl': out[54],
+      '_impl.qLD': out[55],
+      '_impl.qLDiagInv': out[56],
+      '_impl.qLU': out[57],
+      'qacc': out[58],
+      'qacc_smooth': out[59],
+      'qacc_warmstart': out[60],
+      'qfrc_actuator': out[61],
+      'qfrc_bias': out[62],
+      'qfrc_constraint': out[63],
+      '_impl.qfrc_damper': out[64],
+      'qfrc_fluid': out[65],
+      'qfrc_gravcomp': out[66],
+      'qfrc_passive': out[67],
+      'qfrc_smooth': out[68],
+      '_impl.qfrc_spring': out[69],
+      'qpos': out[70],
+      'qvel': out[71],
+      'sensordata': out[72],
+      'site_xmat': out[73],
+      'site_xpos': out[74],
+      '_impl.solver_niter': out[75],
+      '_impl.subtree_angmom': out[76],
+      'subtree_com': out[77],
+      '_impl.subtree_linvel': out[78],
+      '_impl.ten_J': out[79],
+      'ten_length': out[80],
+      '_impl.ten_velocity': out[81],
+      '_impl.ten_wrapadr': out[82],
+      '_impl.ten_wrapnum': out[83],
+      'time': out[84],
+      '_impl.tree_island': out[85],
+      '_impl.wrap_obj': out[86],
+      '_impl.wrap_xpos': out[87],
+      'xanchor': out[88],
+      'xaxis': out[89],
+      'ximat': out[90],
+      'xipos': out[91],
+      'xmat': out[92],
+      'xpos': out[93],
+      'xquat': out[94],
+      '_impl.contact__dim': out[95],
+      '_impl.contact__dist': out[96],
+      '_impl.contact__efc_address': out[97],
+      '_impl.contact__flex': out[98],
+      '_impl.contact__frame': out[99],
+      '_impl.contact__friction': out[100],
+      '_impl.contact__geom': out[101],
+      '_impl.contact__geomcollisionid': out[102],
+      '_impl.contact__includemargin': out[103],
+      '_impl.contact__pos': out[104],
+      '_impl.contact__solimp': out[105],
+      '_impl.contact__solref': out[106],
+      '_impl.contact__solreffriction': out[107],
+      '_impl.contact__type': out[108],
+      '_impl.contact__vert': out[109],
+      '_impl.contact__worldid': out[110],
+      '_impl.efc__D': out[111],
+      '_impl.efc__J': out[112],
+      '_impl.efc__J_colind': out[113],
+      '_impl.efc__J_rowadr': out[114],
+      '_impl.efc__J_rownnz': out[115],
+      '_impl.efc__Jqvel': out[116],
+      '_impl.efc__Ma': out[117],
+      '_impl.efc__aref': out[118],
+      '_impl.efc__force': out[119],
+      '_impl.efc__frictionloss': out[120],
+      '_impl.efc__iD': out[121],
+      '_impl.efc__iJ': out[122],
+      '_impl.efc__iJ_colind': out[123],
+      '_impl.efc__iJ_rowadr': out[124],
+      '_impl.efc__iJ_rownnz': out[125],
+      '_impl.efc__iaref': out[126],
+      '_impl.efc__id': out[127],
+      '_impl.efc__iforce': out[128],
+      '_impl.efc__ifrictionloss': out[129],
+      '_impl.efc__iid': out[130],
+      '_impl.efc__island': out[131],
+      '_impl.efc__istate': out[132],
+      '_impl.efc__itype': out[133],
+      '_impl.efc__margin': out[134],
+      '_impl.efc__pos': out[135],
+      '_impl.efc__state': out[136],
+      '_impl.efc__type': out[137],
+      '_impl.efc__vel': out[138],
   })
   return d
 
diff --git a/mjx/mujoco/mjx/warp/forward_test.py b/mjx/mujoco/mjx/warp/forward_test.py
index 2c38891e..95d49e0f 100644
--- a/mjx/mujoco/mjx/warp/forward_test.py
+++ b/mjx/mujoco/mjx/warp/forward_test.py
@@ -176,8 +176,8 @@ class ForwardTest(parameterized.TestCase):
 
       qm = np.zeros((m.nv, m.nv))
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      # mjwarp adds padding to qM
-      tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
+      # mjwarp adds padding to M
+      tu.assert_eq(qm, dx._impl.M[: m.nv, : m.nv], 'M')
       # qLD is fused in a cholesky factorize and solve, and not written to.
 
       tu.assert_contact_eq(d, dx, worldid=i)
diff --git a/mjx/mujoco/mjx/warp/smooth_test.py b/mjx/mujoco/mjx/warp/smooth_test.py
index 7c8bb368..442461cb 100644
--- a/mjx/mujoco/mjx/warp/smooth_test.py
+++ b/mjx/mujoco/mjx/warp/smooth_test.py
@@ -112,7 +112,7 @@ class SmoothTest(parameterized.TestCase):
 
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     rng = jax.random.PRNGKey(0)
     dx = mjx.make_data(m, impl='warp')
@@ -188,7 +188,7 @@ class SmoothTest(parameterized.TestCase):
     batch_size = 7
     d = mujoco.MjData(m)
     mx = mjx.put_model(m, impl='warp')
-    mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
+    mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
 
     worldids = jp.arange(batch_size)
     dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
diff --git a/mjx/mujoco/mjx/warp/types.py b/mjx/mujoco/mjx/warp/types.py
index df4098d2..5e7ef094 100644
--- a/mjx/mujoco/mjx/warp/types.py
+++ b/mjx/mujoco/mjx/warp/types.py
@@ -78,6 +78,29 @@ class BlockDim:
   """Block dimension 'block_dim' settings for wp.launch_tiled.
 
   TODO(team): experimental and may be removed
+
+  Attributes:
+    segmented_sort: segmented sort block dimension (collision_driver)
+    euler_dense: Euler dense block dimension (forward)
+    actuator_velocity: actuator velocity block dimension (forward)
+    ray: ray block dimension (ray)
+    contact_sort: contact sort block dimension (sensor)
+    energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
+    cholesky_factorize: Cholesky factorize block dimension (smooth)
+    cholesky_solve: Cholesky solve block dimension (smooth)
+    cholesky_factorize_solve: Cholesky factorize and solve block dimension
+      (smooth)
+    solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
+    update_gradient_cholesky: update gradient Cholesky block dimension (solver)
+    update_gradient_cholesky_blocked: update gradient Cholesky blocked block
+      dimension (solver)
+    update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension
+      (solver)
+    update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension
+      (solver)
+    linesearch_iterative: linesearch iterative block dimension (solver)
+    contact_jac_tiled: contact Jacobian tiled block dimension (solver)
+    qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
   """
 
   actuator_velocity: int
@@ -134,10 +157,26 @@ class OptionWarp(PyTreeNode):
 
 class ModelWarp(PyTreeNode):
   """Derived fields from Model."""
-
+  D_colind: np.ndarray
+  D_diag: np.ndarray
+  D_rowadr: np.ndarray
+  D_rownnz: np.ndarray
   M_colind: np.ndarray
+  M_elemid: np.ndarray
+  M_fullm_i: np.ndarray
+  M_fullm_j: np.ndarray
+  M_fullm_upper_elemid: np.ndarray
+  M_fullm_upper_i: np.ndarray
+  M_fullm_upper_j: np.ndarray
+  M_mulm_col: np.ndarray
+  M_mulm_madr: np.ndarray
+  M_mulm_rowadr: np.ndarray
   M_rowadr: np.ndarray
   M_rownnz: np.ndarray
+  M_tiles: Tuple[TileSet, ...]
+  actuator_delay: np.ndarray
+  actuator_history: np.ndarray
+  actuator_historyadr: np.ndarray
   actuator_trntype_body_adr: np.ndarray
   block_dim: BlockDim
   body_branch_start: np.ndarray
@@ -204,6 +243,8 @@ class ModelWarp(PyTreeNode):
   light_active: jax.Array
   light_bodyid: np.ndarray
   light_targetbodyid: np.ndarray
+  mapD2M: np.ndarray
+  mapM2D: np.ndarray
   mapM2M: np.ndarray
   mat_texrepeat: jax.Array
   max_ten_J_rownnz: int
@@ -249,22 +290,22 @@ class ModelWarp(PyTreeNode):
   oct_coeff: np.ndarray
   plugin: np.ndarray
   plugin_attr: np.ndarray
+  qD_fullm_i: np.ndarray
+  qD_fullm_j: np.ndarray
   qLD_all_updates: np.ndarray
   qLD_level_offsets: np.ndarray
   qLD_updates: Tuple[np.ndarray, ...]
-  qM_fullm_i: np.ndarray
-  qM_fullm_j: np.ndarray
-  qM_mulm_col: np.ndarray
-  qM_mulm_madr: np.ndarray
-  qM_mulm_rowadr: np.ndarray
-  qM_tiles: Tuple[TileSet, ...]
   rangefinder_sensor_adr: np.ndarray
   sensor_acc_adr: np.ndarray
   sensor_adr_to_contact_adr: np.ndarray
   sensor_collision_start_adr: np.ndarray
   sensor_contact_adr: np.ndarray
+  sensor_delay: np.ndarray
   sensor_e_kinetic: bool
   sensor_e_potential: bool
+  sensor_history: np.ndarray
+  sensor_historyadr: np.ndarray
+  sensor_interval: np.ndarray
   sensor_limitfrc_adr: np.ndarray
   sensor_limitpos_adr: np.ndarray
   sensor_limitvel_adr: np.ndarray
@@ -299,7 +340,7 @@ class ModelWarp(PyTreeNode):
 
 class DataWarp(PyTreeNode):
   """Derived fields from Data."""
-
+  M: jax.Array
   actuator_moment: jax.Array
   actuator_velocity: jax.Array
   cacc: jax.Array
@@ -323,6 +364,8 @@ class DataWarp(PyTreeNode):
   contact__vert: jax.Array
   contact__worldid: jax.Array
   crb: jax.Array
+  dof_island: jax.Array
+  dof_islandid: jax.Array
   efc__D: jax.Array
   efc__J: jax.Array
   efc__J_colind: jax.Array
@@ -333,19 +376,46 @@ class DataWarp(PyTreeNode):
   efc__aref: jax.Array
   efc__force: jax.Array
   efc__frictionloss: jax.Array
+  efc__iD: jax.Array
+  efc__iJ: jax.Array
+  efc__iJ_colind: jax.Array
+  efc__iJ_rowadr: jax.Array
+  efc__iJ_rownnz: jax.Array
+  efc__iaref: jax.Array
   efc__id: jax.Array
+  efc__iforce: jax.Array
+  efc__ifrictionloss: jax.Array
+  efc__iid: jax.Array
+  efc__island: jax.Array
+  efc__istate: jax.Array
+  efc__itype: jax.Array
   efc__margin: jax.Array
   efc__pos: jax.Array
   efc__state: jax.Array
   efc__type: jax.Array
   efc__vel: jax.Array
+  efc_islandid: jax.Array
   energy: jax.Array
   flexedge_J: jax.Array
   flexedge_length: jax.Array
   flexedge_velocity: jax.Array
   flexvert_xpos: jax.Array
+  iqacc: jax.Array
+  iqacc_smooth: jax.Array
+  iqfrc_constraint: jax.Array
+  iqfrc_smooth: jax.Array
+  island_dofadr: jax.Array
+  island_efcadr: jax.Array
+  island_ne: jax.Array
+  island_nefc: jax.Array
+  island_nf: jax.Array
+  island_nv: jax.Array
   light_xdir: jax.Array
   light_xpos: jax.Array
+  map_dof2idof: jax.Array
+  map_efc2iefc: jax.Array
+  map_idof2dof: jax.Array
+  map_iefc2efc: jax.Array
   moment_colind: jax.Array
   moment_rowadr: jax.Array
   moment_rownnz: jax.Array
@@ -356,6 +426,7 @@ class DataWarp(PyTreeNode):
   ne: jax.Array
   nefc: jax.Array
   nf: jax.Array
+  nidof: jax.Array
   nisland: jax.Array
   njmax: int
   njmax_nnz: int
@@ -364,7 +435,7 @@ class DataWarp(PyTreeNode):
   nworld: int
   qLD: jax.Array
   qLDiagInv: jax.Array
-  qM: jax.Array
+  qLU: jax.Array
   qfrc_damper: jax.Array
   qfrc_spring: jax.Array
   solver_niter: jax.Array
@@ -401,6 +472,7 @@ DATA_NON_VMAP = {
     'nacon',
     'naconmax',
     'ncollision',
+    'nidof',
     'njmax',
     'njmax_nnz',
     'njmax_pad',
@@ -434,6 +506,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
 
 _NDIM = {
     'Data': {
+        'M': 3,
         'act': 2,
         'act_dot': 2,
         'actuator_force': 2,
@@ -467,6 +540,8 @@ _NDIM = {
         'crb': 3,
         'ctrl': 2,
         'cvel': 3,
+        'dof_island': 2,
+        'dof_islandid': 2,
         'efc__D': 2,
         'efc__J': 3,
         'efc__J_colind': 3,
@@ -477,12 +552,25 @@ _NDIM = {
         'efc__aref': 2,
         'efc__force': 2,
         'efc__frictionloss': 2,
+        'efc__iD': 2,
+        'efc__iJ': 3,
+        'efc__iJ_colind': 3,
+        'efc__iJ_rowadr': 2,
+        'efc__iJ_rownnz': 2,
+        'efc__iaref': 2,
         'efc__id': 2,
+        'efc__iforce': 2,
+        'efc__ifrictionloss': 2,
+        'efc__iid': 2,
+        'efc__island': 2,
+        'efc__istate': 2,
+        'efc__itype': 2,
         'efc__margin': 2,
         'efc__pos': 2,
         'efc__state': 2,
         'efc__type': 2,
         'efc__vel': 2,
+        'efc_islandid': 2,
         'energy': 2,
         'eq_active': 2,
         'flexedge_J': 2,
@@ -491,8 +579,23 @@ _NDIM = {
         'flexvert_xpos': 3,
         'geom_xmat': 4,
         'geom_xpos': 3,
+        'history': 2,
+        'iqacc': 2,
+        'iqacc_smooth': 2,
+        'iqfrc_constraint': 2,
+        'iqfrc_smooth': 2,
+        'island_dofadr': 2,
+        'island_efcadr': 2,
+        'island_ne': 2,
+        'island_nefc': 2,
+        'island_nf': 2,
+        'island_nv': 2,
         'light_xdir': 3,
         'light_xpos': 3,
+        'map_dof2idof': 2,
+        'map_efc2iefc': 2,
+        'map_idof2dof': 2,
+        'map_iefc2efc': 2,
         'mocap_pos': 3,
         'mocap_quat': 3,
         'moment_colind': 2,
@@ -505,6 +608,7 @@ _NDIM = {
         'ne': 1,
         'nefc': 1,
         'nf': 1,
+        'nidof': 1,
         'nisland': 1,
         'njmax': 0,
         'njmax_nnz': 0,
@@ -513,7 +617,7 @@ _NDIM = {
         'nworld': 0,
         'qLD': 3,
         'qLDiagInv': 2,
-        'qM': 3,
+        'qLU': 3,
         'qacc': 2,
         'qacc_smooth': 2,
         'qacc_warmstart': 2,
@@ -556,9 +660,23 @@ _NDIM = {
         'xquat': 3,
     },
     'Model': {
+        'D_colind': 1,
+        'D_diag': 1,
+        'D_rowadr': 1,
+        'D_rownnz': 1,
         'M_colind': 1,
+        'M_elemid': 2,
+        'M_fullm_i': 1,
+        'M_fullm_j': 1,
+        'M_fullm_upper_elemid': 1,
+        'M_fullm_upper_i': 1,
+        'M_fullm_upper_j': 1,
+        'M_mulm_col': 1,
+        'M_mulm_madr': 1,
+        'M_mulm_rowadr': 1,
         'M_rowadr': 1,
         'M_rownnz': 1,
+        'M_tiles': -1,
         'actuator_acc0': 2,
         'actuator_actadr': 1,
         'actuator_actearly': 1,
@@ -570,6 +688,7 @@ _NDIM = {
         'actuator_cranklength': 2,
         'actuator_ctrllimited': 1,
         'actuator_ctrlrange': 3,
+        'actuator_delay': 1,
         'actuator_dynprm': 3,
         'actuator_dyntype': 1,
         'actuator_forcelimited': 1,
@@ -577,6 +696,8 @@ _NDIM = {
         'actuator_gainprm': 3,
         'actuator_gaintype': 1,
         'actuator_gear': 3,
+        'actuator_history': 2,
+        'actuator_historyadr': 1,
         'actuator_lengthrange': 3,
         'actuator_trnid': 2,
         'actuator_trntype': 1,
@@ -769,6 +890,8 @@ _NDIM = {
         'light_poscom0': 3,
         'light_targetbodyid': 1,
         'light_type': 2,
+        'mapD2M': 1,
+        'mapM2D': 1,
         'mapM2M': 1,
         'mat_rgba': 3,
         'mat_texid': 3,
@@ -797,6 +920,7 @@ _NDIM = {
         'mesh_vertnum': 1,
         'mocap_bodyid': 1,
         'nC': 0,
+        'nD': 0,
         'nJfe': 0,
         'nJmom': 0,
         'nJten': 0,
@@ -821,6 +945,7 @@ _NDIM = {
         'ngravcomp': 0,
         'nhfield': 0,
         'nhfielddata': 0,
+        'nhistory': 0,
         'njnt': 0,
         'nlight': 0,
         'nmat': 0,
@@ -897,15 +1022,11 @@ _NDIM = {
         'pair_solreffriction': 3,
         'plugin': 1,
         'plugin_attr': 2,
+        'qD_fullm_i': 1,
+        'qD_fullm_j': 1,
         'qLD_all_updates': 2,
         'qLD_level_offsets': 1,
         'qLD_updates': -1,
-        'qM_fullm_i': 1,
-        'qM_fullm_j': 1,
-        'qM_mulm_col': 1,
-        'qM_mulm_madr': 1,
-        'qM_mulm_rowadr': 1,
-        'qM_tiles': -1,
         'qpos0': 2,
         'qpos_spring': 2,
         'rangefinder_sensor_adr': 1,
@@ -916,9 +1037,13 @@ _NDIM = {
         'sensor_contact_adr': 1,
         'sensor_cutoff': 1,
         'sensor_datatype': 1,
+        'sensor_delay': 1,
         'sensor_dim': 1,
         'sensor_e_kinetic': 0,
         'sensor_e_potential': 0,
+        'sensor_history': 2,
+        'sensor_historyadr': 1,
+        'sensor_interval': 2,
         'sensor_intprm': 2,
         'sensor_limitfrc_adr': 1,
         'sensor_limitpos_adr': 1,
@@ -1016,6 +1141,7 @@ _NDIM = {
 }
 _BATCH_DIM = {
     'Data': {
+        'M': True,
         'act': True,
         'act_dot': True,
         'actuator_force': True,
@@ -1049,6 +1175,8 @@ _BATCH_DIM = {
         'crb': True,
         'ctrl': True,
         'cvel': True,
+        'dof_island': True,
+        'dof_islandid': True,
         'efc__D': True,
         'efc__J': True,
         'efc__J_colind': True,
@@ -1059,12 +1187,25 @@ _BATCH_DIM = {
         'efc__aref': True,
         'efc__force': True,
         'efc__frictionloss': True,
+        'efc__iD': True,
+        'efc__iJ': True,
+        'efc__iJ_colind': True,
+        'efc__iJ_rowadr': True,
+        'efc__iJ_rownnz': True,
+        'efc__iaref': True,
         'efc__id': True,
+        'efc__iforce': True,
+        'efc__ifrictionloss': True,
+        'efc__iid': True,
+        'efc__island': True,
+        'efc__istate': True,
+        'efc__itype': True,
         'efc__margin': True,
         'efc__pos': True,
         'efc__state': True,
         'efc__type': True,
         'efc__vel': True,
+        'efc_islandid': True,
         'energy': True,
         'eq_active': True,
         'flexedge_J': True,
@@ -1073,8 +1214,23 @@ _BATCH_DIM = {
         'flexvert_xpos': True,
         'geom_xmat': True,
         'geom_xpos': True,
+        'history': True,
+        'iqacc': True,
+        'iqacc_smooth': True,
+        'iqfrc_constraint': True,
+        'iqfrc_smooth': True,
+        'island_dofadr': True,
+        'island_efcadr': True,
+        'island_ne': True,
+        'island_nefc': True,
+        'island_nf': True,
+        'island_nv': True,
         'light_xdir': True,
         'light_xpos': True,
+        'map_dof2idof': True,
+        'map_efc2iefc': True,
+        'map_idof2dof': True,
+        'map_iefc2efc': True,
         'mocap_pos': True,
         'mocap_quat': True,
         'moment_colind': True,
@@ -1087,6 +1243,7 @@ _BATCH_DIM = {
         'ne': True,
         'nefc': True,
         'nf': True,
+        'nidof': False,
         'nisland': True,
         'njmax': False,
         'njmax_nnz': False,
@@ -1095,7 +1252,7 @@ _BATCH_DIM = {
         'nworld': False,
         'qLD': True,
         'qLDiagInv': True,
-        'qM': True,
+        'qLU': True,
         'qacc': True,
         'qacc_smooth': True,
         'qacc_warmstart': True,
@@ -1138,9 +1295,23 @@ _BATCH_DIM = {
         'xquat': True,
     },
     'Model': {
+        'D_colind': False,
+        'D_diag': False,
+        'D_rowadr': False,
+        'D_rownnz': False,
         'M_colind': False,
+        'M_elemid': False,
+        'M_fullm_i': False,
+        'M_fullm_j': False,
+        'M_fullm_upper_elemid': False,
+        'M_fullm_upper_i': False,
+        'M_fullm_upper_j': False,
+        'M_mulm_col': False,
+        'M_mulm_madr': False,
+        'M_mulm_rowadr': False,
         'M_rowadr': False,
         'M_rownnz': False,
+        'M_tiles': False,
         'actuator_acc0': True,
         'actuator_actadr': False,
         'actuator_actearly': False,
@@ -1152,6 +1323,7 @@ _BATCH_DIM = {
         'actuator_cranklength': True,
         'actuator_ctrllimited': False,
         'actuator_ctrlrange': True,
+        'actuator_delay': False,
         'actuator_dynprm': True,
         'actuator_dyntype': False,
         'actuator_forcelimited': False,
@@ -1159,6 +1331,8 @@ _BATCH_DIM = {
         'actuator_gainprm': True,
         'actuator_gaintype': False,
         'actuator_gear': True,
+        'actuator_history': False,
+        'actuator_historyadr': False,
         'actuator_lengthrange': True,
         'actuator_trnid': False,
         'actuator_trntype': False,
@@ -1351,6 +1525,8 @@ _BATCH_DIM = {
         'light_poscom0': True,
         'light_targetbodyid': False,
         'light_type': True,
+        'mapD2M': False,
+        'mapM2D': False,
         'mapM2M': False,
         'mat_rgba': True,
         'mat_texid': True,
@@ -1379,6 +1555,7 @@ _BATCH_DIM = {
         'mesh_vertnum': False,
         'mocap_bodyid': False,
         'nC': False,
+        'nD': False,
         'nJfe': False,
         'nJmom': False,
         'nJten': False,
@@ -1403,6 +1580,7 @@ _BATCH_DIM = {
         'ngravcomp': False,
         'nhfield': False,
         'nhfielddata': False,
+        'nhistory': False,
         'njnt': False,
         'nlight': False,
         'nmat': False,
@@ -1479,15 +1657,11 @@ _BATCH_DIM = {
         'pair_solreffriction': True,
         'plugin': False,
         'plugin_attr': False,
+        'qD_fullm_i': False,
+        'qD_fullm_j': False,
         'qLD_all_updates': False,
         'qLD_level_offsets': False,
         'qLD_updates': False,
-        'qM_fullm_i': False,
-        'qM_fullm_j': False,
-        'qM_mulm_col': False,
-        'qM_mulm_madr': False,
-        'qM_mulm_rowadr': False,
-        'qM_tiles': False,
         'qpos0': True,
         'qpos_spring': True,
         'rangefinder_sensor_adr': False,
@@ -1498,9 +1672,13 @@ _BATCH_DIM = {
         'sensor_contact_adr': False,
         'sensor_cutoff': False,
         'sensor_datatype': False,
+        'sensor_delay': False,
         'sensor_dim': False,
         'sensor_e_kinetic': False,
         'sensor_e_potential': False,
+        'sensor_history': False,
+        'sensor_historyadr': False,
+        'sensor_interval': False,
         'sensor_intprm': False,
         'sensor_limitfrc_adr': False,
         'sensor_limitpos_adr': False,