Merge branch 'main' of https://github.com/google-deepmind/mujoco into mjx-warp-segmentation

.github/workflows/build_steps.sh

@@ -212,25 +212,28 @@ build_test_wasm() {
     echo "Building and testing WASM bindings..."
     source emsdk/emsdk_env.sh
     export PATH="$(pwd)/node_modules/.bin:$PATH"
-
-    echo "Building Multi-Threaded version..."
+    echo "Build MuJoCo with Emscripten (Multi-Threaded)..."
     emcmake cmake -B build_wasm_mt \
         -DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL=OFF \
         -DMUJOCO_WASM_THREADS=ON \
         $WASM_CMAKE_ARGS
     cmake --build build_wasm_mt --parallel $(nproc)
 
+    echo "Run bindings tests for Multi-Threaded version..."
+    npm run test --prefix ./wasm
+
     echo "Moving Multi-Thread version under mt subfolder..."
     mkdir -p wasm/dist/mt
     mv wasm/dist/mujoco.* wasm/dist/mt/
 
-    echo "Building Single-Threaded version..."
+    echo "Build MuJoCo with Emscripten (Single-Threaded)..."
     emcmake cmake -B build_wasm_st \
         -DCMAKE_INTERPROCEDURAL_OPTIMIZATION:BOOL=OFF \
         -DMUJOCO_WASM_THREADS=OFF \
         $WASM_CMAKE_ARGS
     cmake --build build_wasm_st --parallel $(nproc)
 
+    echo "Run bindings tests for Single-Threaded version..."
     npm run test --prefix ./wasm
 }
 

CMakeLists.txt

@@ -132,9 +132,6 @@ endif()
 
 if(EMSCRIPTEN)
   add_subdirectory(wasm)
-  if(MUJOCO_BUILD_TESTS_WASM)
-    add_subdirectory(wasm/tests)
-  endif()
 endif()
 
 

doc/XMLreference.rst

@@ -3700,8 +3700,14 @@ saving the XML:
 .. _body-flexcomp-count:
 
 :at:`count`: :at-val:`int(3), "10 10 10"`
-   The number of automatically generated points in each dimension. This and the next attribute only apply to types grid,
-   box, cylinder, ellipsoid.
+   Specifies the number of automatically generated points in each dimension for types **grid**, **box**, **cylinder**,
+   and **ellipsoid**.
+
+.. _body-flexcomp-cellcount:
+
+:at:`cellcount`: :at-val:`int(3), "1 1 1"`
+   Specifies the number of cells in each dimension for the background interpolation grid when using **trilinear** or
+   **quadratic** dofs.
 
 .. _body-flexcomp-spacing:
 
@@ -4242,6 +4248,17 @@ cases, the user will specify a :el:`flexcomp` which will then automatically cons
    An array of MuJoCo body names (separated by white space) to which each node belongs. The number of body names
    should equal the number of nodes (nnode). See the flexcomp :ref:`dof<body-flexcomp-dof>` attribute for more details.
 
+.. _deformable-flex-cellcount:
+
+:at:`cellcount`: :at-val:`int(3), optional`
+   When using **trilinear** or **quadratic** dofs, this specifies the number of cells in each dimension for the
+   background interpolation grid.
+
+.. _deformable-flex-dof:
+
+:at:`dof`: :at-val:`[trilinear, quadratic], optional`
+   Interpolation order for the flex.
+
 .. _flex-edge:
 
 :el-prefix:`flex/` |-| **edge** |?|

doc/XMLschema.rst

@@ -1417,6 +1417,9 @@
             .. grid-item::
                :ref:`count<body-flexcomp-count>`
 
+            .. grid-item::
+               :ref:`cellcount<body-flexcomp-cellcount>`
+
             .. grid-item::
                :ref:`spacing<body-flexcomp-spacing>`
 
@@ -1657,6 +1660,12 @@
             .. grid-item::
                :ref:`node<deformable-flex-node>`
 
+            .. grid-item::
+               :ref:`cellcount<deformable-flex-cellcount>`
+
+            .. grid-item::
+               :ref:`dof<deformable-flex-dof>`
+
 
          .. dropdown:: :ref:`contact<flex-contact>` :octicon:`dot`
 

doc/changelog.rst

@@ -2,6 +2,15 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+General
+^^^^^^^
+
+- Added :ref:`multi-cell support<body-flexcomp-cellnum>` for trilinear and quadratic flexes. Note that the implicit
+  integrator uses a dense solver for the flex degrees of freedom, which can be slow for multi-cell flexes.
+
 Version 3.7.0 (April 14, 2026)
 ------------------------------
 

doc/includes/references.h

@@ -1315,6 +1315,8 @@ struct mjModel_ {
   int*      flex_matid;           // material id for rendering                (nflex x 1)
   int*      flex_group;           // group for visibility                     (nflex x 1)
   int*      flex_interp;          // interpolation (0: vertex, 1: nodes)      (nflex x 1)
+  int*      flex_bandwidth;       // precomputed solver bandwidth             (nflex x 1)
+  int*      flex_cellnum;         // finite cell num per dimension            (nflex x 3)
   int*      flex_nodeadr;         // first node address                       (nflex x 1)
   int*      flex_nodenum;         // number of nodes                          (nflex x 1)
   int*      flex_vertadr;         // first vertex address                     (nflex x 1)
@@ -2256,6 +2258,8 @@ typedef struct mjsFlex_ {          // flex specification
   double damping;                  // Rayleigh's damping
   double thickness;                // thickness (2D only)
   int elastic2d;                   // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
+  int cellcount[3];                // grid cell count for finite cell method
+  int order;                       // interpolation order (1: trilinear, 2: quadratic)
 
   // mesh properties
   mjStringVec* nodebody;           // node body names

doc/modeling.rst

@@ -1442,6 +1442,16 @@ instead to specify shear and volumetric stiffnesses separately using the `Poisso
 <https://en.wikipedia.org/wiki/Poisson%27s_ratio>`__ of the material. For more details, see the `Saint Venant-Kirchhoff
 <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__ hyperelastic model.
 
+**Parametrization types**.
+
+While the default behavior of :el:`flexcomp` produces a "full" flex where every node corresponds to a MuJoCo body, it
+also supports specialized :ref:`parametrizations<body-flexcomp-dof>` for volumetric objects: **trilinear** and
+**quadratic**. Instead of directly simulating all nodes, these options define a background grid of cells. The positions
+of the interior vertices are computed by interpolating the positions of the cell corners. Trilinear flexes use 8-node
+hexahedral cells with linear interpolation along each axis, while quadratic flexes use 27-node cells with quadratic
+interpolation, allowing for curved deformation modes. These grid-based parametrizations require fewer degrees of freedom
+than full flexes and can result in significantly faster simulation times, especially for large volumetric soft bodies.
+
 **Creation and visualization**.
 
 .. code-block:: xml

doc/models.rst

@@ -151,3 +151,5 @@ Biomechanical
      - Preview
    * - `Fruitfly <https://github.com/google-deepmind/mujoco_menagerie/tree/main/flybody>`_
      - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/flybody/flybody.png
+   * - `MS-Human-700 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/ms_human_700>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/ms_human_700/ms_human_700.png

include/mujoco/mjmodel.h

@@ -978,6 +978,8 @@ struct mjModel_ {
   int*      flex_matid;           // material id for rendering                (nflex x 1)
   int*      flex_group;           // group for visibility                     (nflex x 1)
   int*      flex_interp;          // interpolation (0: vertex, 1: nodes)      (nflex x 1)
+  int*      flex_bandwidth;       // precomputed solver bandwidth             (nflex x 1)
+  int*      flex_cellnum;         // finite cell num per dimension            (nflex x 3)
   int*      flex_nodeadr;         // first node address                       (nflex x 1)
   int*      flex_nodenum;         // number of nodes                          (nflex x 1)
   int*      flex_vertadr;         // first vertex address                     (nflex x 1)

include/mujoco/mjspec.h

@@ -454,6 +454,8 @@ typedef struct mjsFlex_ {          // flex specification
   double damping;                  // Rayleigh's damping
   double thickness;                // thickness (2D only)
   int elastic2d;                   // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
+  int cellcount[3];                // grid cell count for finite cell method
+  int order;                       // interpolation order (1: trilinear, 2: quadratic)
 
   // mesh properties
   mjStringVec* nodebody;           // node body names

include/mujoco/mjxmacro.h

@@ -454,6 +454,8 @@
     X   ( int,     flex_matid,            nflex,         1                    ) \
     X   ( int,     flex_group,            nflex,         1                    ) \
     X   ( int,     flex_interp,           nflex,         1                    ) \
+    X   ( int,     flex_bandwidth,        nflex,         1                    ) \
+    X   ( int,     flex_cellnum,          nflex,         3                    ) \
     X   ( int,     flex_nodeadr,          nflex,         1                    ) \
     X   ( int,     flex_nodenum,          nflex,         1                    ) \
     X   ( int,     flex_vertadr,          nflex,         1                    ) \

mjx/cuda_requirements.txt

@@ -16,8 +16,8 @@ jax-cuda12-pjrt==0.5.3; python_version >= '3.10' \
 jax-cuda12-pjrt==0.4.30; python_version == '3.9' \
     --hash=sha256:895d0198ad99638fcaf976c47592e2a543eef79ea15fabd24a402d055390c328 \
     --hash=sha256:c36fb1e0c236563bf3a87e70f4d1ab28a31d7cf5d722c9ede30c4172116e8bcb
-warp-lang==1.12.0 \
-    --hash=sha256:c78c3701d5cad86c30ef5017410d294ec46a396bb0d502ee1c98743494f3a62f \
-    --hash=sha256:a1436f60a1881cd94f787e751a83fc0987626be2d3e2b4e74c64a6947c6d1266 \
-    --hash=sha256:a2d6decba693aba5b828573c4414fd6a3f4c4a934db9c322736ef2b3fa99fe76 \
-    --hash=sha256:697248edd2f1e2952f50e3db33b214af76173641a8894aacc467bed6dc247f8a
+warp-lang==1.12.1 \
+    --hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
+    --hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
+    --hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
+    --hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/__init__.py

@@ -1,14 +1,2 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-#
-# 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.

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py

@@ -1,17 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-#
-# 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 ctypes
 from functools import reduce

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py

@@ -1,17 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-#
-# 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.
 
 from __future__ import annotations
 
@@ -223,7 +211,7 @@ class FfiKernel:
 
         # register the callback
         FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
-        self.callback_func = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame))
+        self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
         ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
         ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
         jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")
@@ -606,7 +594,7 @@ class FfiCallable:
 
         # register the callback
         FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
-        self.callback_func = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame))
+        self.callback_func = FFI_CCALLFUNC(self.ffi_callback)
         ffi_ccall_address = ctypes.cast(self.callback_func, ctypes.c_void_p)
         ffi_capsule = jax.ffi.pycapsule(ffi_ccall_address.value)
         jax.ffi.register_ffi_target(self.name, ffi_capsule, platform="CUDA")

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/xla_ffi.py

@@ -1,17 +1,5 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
-#
-# 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 ctypes
 import enum

mjx/pyproject.toml

@@ -36,7 +36,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.12.0",
+    "warp-lang==1.12.1",
 ]
 
 [project.scripts]

model/flex/bunny_multicell.xml

@@ -0,0 +1,42 @@
+<!-- Copyright 2024 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.
+-->
+
+<mujoco model="Trilinear">
+  <include file="scene.xml"/>
+
+  <option solver="CG" tolerance="1e-6" timestep=".001" integrator="Euler"/>
+
+  <size memory="100M"/>
+
+  <visual>
+    <map stiffness="100"/>
+  </visual>
+
+  <worldbody>
+    <body>
+      <joint name="press" type="slide" axis="0 0 1" damping="500"/>
+      <geom type="box" size=".02 .2 .2" pos="0 0 .5"/>
+    </body>
+    <flexcomp type="mesh" file="bunny.obj" pos="0 0 0" dim="2" euler="90 0 0" cellcount="3 3 3"
+              radius=".001" rgba="0 .7 .7 1" mass=".05" name="softbody" dof="trilinear">
+      <elasticity young="1e3" poisson="0.1" damping="0.01" elastic2d="stretch"/>
+      <contact selfcollide="none" internal="false"/>
+    </flexcomp>
+  </worldbody>
+
+  <actuator>
+    <position name="press" joint="press" gear="-1 0 0 0 0 0" ctrlrange="-1 1" kp="1000"/>
+  </actuator>
+</mujoco>

python/mujoco/introspect/structs.py

@@ -2668,6 +2668,22 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  doc='interpolation (0: vertex, 1: nodes)',
                  array_extent=('nflex',),
              ),
+             StructFieldDecl(
+                 name='flex_bandwidth',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='precomputed solver bandwidth',
+                 array_extent=('nflex',),
+             ),
+             StructFieldDecl(
+                 name='flex_cellnum',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='finite cell num per dimension',
+                 array_extent=('nflex', 3),
+             ),
              StructFieldDecl(
                  name='flex_nodeadr',
                  type=PointerType(
@@ -8237,6 +8253,19 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='2D passive forces; 0: none, 1: bending, 2: stretching, 3: both',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='cellcount',
+                 type=ArrayType(
+                     inner_type=ValueType(name='int'),
+                     extents=(3,),
+                 ),
+                 doc='grid cell count for finite cell method',
+             ),
+             StructFieldDecl(
+                 name='order',
+                 type=ValueType(name='int'),
+                 doc='interpolation order (1: trilinear, 2: quadratic)',
+             ),
              StructFieldDecl(
                  name='nodebody',
                  type=PointerType(

src/engine/engine_core_constraint.c

@@ -432,21 +432,31 @@ static int mj_vertBodyWeight(const mjModel* m, const mjData* d, int f, int* v,
     return 0;
   }
 
+  // compute parametric coordinates of the vertex in [0, 1]^3
   mjtNum coord[3] = {0, 0, 0};
   for (int i = 0; i < nw; i++) {
-    mju_addToScl3(coord,  m->flex_vert0 + 3*v[i], vweight[i]);
+    mju_addToScl3(coord, m->flex_vert0 + 3*v[i], vweight[i]);
   }
+
+  int order = m->flex_interp[f];
+  int npc = (order+1)*(order+1)*(order+1);  // number of nodes per cell
+
+  // cell lookup: get local coords and node indices
+  mjtNum local[3];
+  int nodeindices[27];  // max npc for quadratic: 3^3 = 27
+  mju_cellLookup(coord, m->flex_cellnum+3*f, order, local, nodeindices);
+
+  // evaluate basis functions for this cell's local nodes
   int nstart = m->flex_nodeadr[f];
-  int nend = m->flex_nodeadr[f] + m->flex_nodenum[f];
   int nb = 0;
 
-  for (int i = nstart; i < nend; i++) {
-    mjtNum w = mju_evalBasis(coord, i-nstart, m->flex_interp[f]);
+  for (int j = 0; j < npc; j++) {
+    mjtNum w = mju_evalBasis(local, j, order);
     if (w < 1e-5) {
       continue;
     }
     if (bweight) bweight[nb] = w;
-    body[nb++] = m->flex_nodebodyid[i];
+    body[nb++] = m->flex_nodebodyid[nstart + nodeindices[j]];
   }
 
   return nb;
@@ -871,6 +881,11 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
         break;
       }
 
+      int npc = (order+1)*(order+1)*(order+1);
+      int cx = m->flex_cellnum[3*f+0];
+      int cy = m->flex_cellnum[3*f+1];
+      int cz = m->flex_cellnum[3*f+2];
+
       // allocate stack for node positions and Jacobians
       mj_markStack(d);
       mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
@@ -910,152 +925,164 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
         }
       }
 
-      // loop over Gauss points
-      // get reference positions from m->flex_node0 (Cartesian positions at qpos0)
+      // reference positions for all nodes
       int nstart = m->flex_nodeadr[f];
       mjtNum* refpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
       for (int n = 0; n < nodenum; n++) {
         mju_copy3(refpos + 3*n, m->flex_node0 + 3*(n + nstart));
       }
 
-      // B-bar: precompute center-point values for volumetric constraint (trilinear only)
-      if (order == 1) {
-        mjtNum center[3] = {0.5, 0.5, 0.5};
-        mjtNum Fcur_c[9], Fref_c[9], Fref_inv_center[9], F_center[9];
-
-        // compute deformation gradient at center
-        mju_defGradient(Fcur_c, center, xpos, order);
-        mju_defGradient(Fref_c, center, refpos, order);
-        mat3_inverse(Fref_c, Fref_inv_center);
-        mju_mulMatMat3(F_center, Fcur_c, Fref_inv_center);
-
-        // compute C and E at center
-        mjtNum C_c[9], E_c[9];
-        mju_mulMatTMat3(C_c, F_center, F_center);
-        mju_scl(E_c, C_c, 0.5, 9);
-        E_c[0] -= 0.5;
-        E_c[4] -= 0.5;
-        E_c[8] -= 0.5;
-
-        // J = det(F) at center
-        mjtNum I1_center = E_c[0] + E_c[4] + E_c[8];
-        mjtNum J_center = mat3_det(F_center);
-
-        // compute shape function gradients at center (8 nodes for trilinear)
-        mjtNum grad_center[8][3];
-        shape_gradients(order, center, grad_center);
-
-        // add I1 and J-1 constraints at center (reduced integration for volumetric)
-        mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-        for (int inv = 0; inv < 2; inv++) {
-          if (inv == 0) {
-            // I1 = tr(E), dI1/dE = I
-            cpos[0] = I1_center;
-          } else {
-            // J - 1 = det(F) - 1, dJ/dF = cofactor(F)
-            cpos[0] = J_center - 1.0;
-          }
-
-          volumetric_dSdx(inv, nodenum, grad_center, F_center, Fref_inv_center, dSdx);
-          strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
-
-          if (issparse) {
-            mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum);
-            for (int k = 0; k < combined_nnz; k++) {
-              sparse_jac[k] = strain_jac[combined_chain[k]];
-            }
-            mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
-                             combined_nnz, combined_chain);
-          } else {
-            mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
-          }
-        }
-      }
-
-      // add I1 and J-1 constraints at center (reduced integration for volumetric)
+      // per-cell arrays
+      mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* refpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* dSdx_local = mjSTACKALLOC(d, 3*npc, mjtNum);
       mjtNum* dSdx = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-      for (int g = 0; g < ngauss; g++) {
-        mjtNum* p = gauss[g];
+      int gindices[125];  // max npc = 125 for quadratic
 
-        // F = Fcur * Fref_inv
-        mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9];
-        mju_defGradient(Fcur, p, xpos, order);
-        mju_defGradient(Fref, p, refpos, order);
-        mat3_inverse(Fref, Fref_inv);
-        mju_mulMatMat3(F, Fcur, Fref_inv);
+      // loop over cells
+      for (int ci = 0; ci < cx; ci++) {
+        for (int cj = 0; cj < cy; cj++) {
+          for (int ck = 0; ck < cz; ck++) {
+            // gather cell-local node positions
+            mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, refpos, xpos_c, NULL,
+                                    refpos_c, gindices, NULL);
 
-        // compute Green-Lagrange strain E = 0.5*(C - I)
-        mjtNum C[9], E[9];
-        mju_mulMatTMat3(C, F, F);
-        for (int j = 0; j < 9; j++) {
-          E[j] = 0.5 * C[j];
-        }
-        E[0] -= 0.5;
-        E[4] -= 0.5;
-        E[8] -= 0.5;
+            // B-bar: center-point volumetric constraints (trilinear)
+            if (order == 1) {
+              mjtNum center[3] = {0.5, 0.5, 0.5};
+              mjtNum Fcur_c[9], Fref_c[9], Fref_inv_c[9], F_c[9];
 
-        // compute 3 invariants of E
-        mjtNum I1 = E[0] + E[4] + E[8];
-        mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6] +
-                      E[3]*E[1] + E[4]*E[4] + E[5]*E[7] +
-                      E[6]*E[2] + E[7]*E[5] + E[8]*E[8];
-        mjtNum I2 = 0.5 * (I1*I1 - trE2);
-        mjtNum I3 = mat3_det(E);
+              mju_defGradient(Fcur_c, center, xpos_c, order);
+              mju_defGradient(Fref_c, center, refpos_c, order);
+              mat3_inverse(Fref_c, Fref_inv_c);
+              mju_mulMatMat3(F_c, Fcur_c, Fref_inv_c);
 
-        // compute shape function gradients at this Gauss point
-        mjtNum grad[27][3];
-        shape_gradients(order, p, grad);
+              mjtNum C_c[9], E_c[9];
+              mju_mulMatTMat3(C_c, F_c, F_c);
+              mju_scl(E_c, C_c, 0.5, 9);
+              E_c[0] -= 0.5; E_c[4] -= 0.5; E_c[8] -= 0.5;
 
-        // trilinear: 3 constraints per Gauss point (I1, I2, I3 skipped - only shear)
-        // quadratic: 6 constraints per Gauss point
-        for (int s = 0; s < 6; s++) {
-          // skip I1, I2, I3 for trilinear (I1, J-1 at center; I2 is small for small strain)
-          if (order == 1 && (s == 0 || s == 1 || s == 2)) {
-            continue;
-          }
+              mjtNum I1_c = E_c[0] + E_c[4] + E_c[8];
+              mjtNum J_c = mat3_det(F_c);
 
-          mjtNum dSdE[9];
-          mju_zero(dSdE, 9);
+              mjtNum grad_c[8][3];
+              shape_gradients(order, center, grad_c);
 
-          if (s == 0) {
-            // I1 = tr(E), dI1/dE = I (only for quadratic)
-            cpos[0] = I1;
-            dSdE[0] = dSdE[4] = dSdE[8] = 1.0;
-          } else if (s == 1) {
-            // I2 = 0.5*(tr(E)^2 - tr(E^2)), dI2/dE = tr(E)*I - E
-            cpos[0] = I2;
-            dSdE[0] = I1 - E[0];
-            dSdE[4] = I1 - E[4];
-            dSdE[8] = I1 - E[8];
-            dSdE[1] = -E[1]; dSdE[3] = -E[3];
-            dSdE[2] = -E[2]; dSdE[6] = -E[6];
-            dSdE[5] = -E[5]; dSdE[7] = -E[7];
-          } else if (s == 2) {
-            // I3 = det(E), dI3/dE = cofactor(E)
-            cpos[0] = I3;
-            mat3_cofactor(E, dSdE);
-          } else {
-            // off-diagonal entries: s=3->E12, s=4->E13, s=5->E23
-            int offdiag_idx[3] = {1, 2, 5};
-            int ij = offdiag_idx[s - 3];
-            cpos[0] = E[ij];
-            dSdE[ij] = 1.0;
-          }
+              for (int inv = 0; inv < 2; inv++) {
+                cpos[0] = (inv == 0) ? I1_c : J_c - 1.0;
 
-          // compute dS/dx for all nodes
-          invariant_dSdx(nodenum, grad, F, Fref_inv, dSdE, dSdx);
-          strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
+                // compute local dSdx
+                volumetric_dSdx(inv, npc, grad_c, F_c, Fref_inv_c, dSdx_local);
 
-          // add constraint
-          if (issparse) {
-            mjtNum* sparse_jac = mjSTACKALLOC(d, combined_nnz, mjtNum);
-            for (int k = 0; k < combined_nnz; k++) {
-              sparse_jac[k] = strain_jac[combined_chain[k]];
+                // scatter to global dSdx
+                mju_zero(dSdx, 3*nodenum);
+                for (int n = 0; n < npc; n++) {
+                  mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n);
+                }
+
+                strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
+
+                if (issparse) {
+                  mj_markStack(d);
+                  mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum);
+                  for (int k = 0; k < combined_nnz; k++) {
+                    sj[k] = strain_jac[combined_chain[k]];
+                  }
+                  mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
+                                   combined_nnz, combined_chain);
+                  mj_freeStack(d);
+                } else {
+                  mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
+                }
+              }
+            }
+
+            // Gauss integration per cell
+            for (int g = 0; g < ngauss; g++) {
+              mjtNum* p = gauss[g];
+
+              // F = Fcur * Fref_inv
+              mjtNum Fcur[9], Fref[9], Fref_inv[9], F[9];
+              mju_defGradient(Fcur, p, xpos_c, order);
+              mju_defGradient(Fref, p, refpos_c, order);
+              mat3_inverse(Fref, Fref_inv);
+              mju_mulMatMat3(F, Fcur, Fref_inv);
+
+              // Green-Lagrange strain E = 0.5*(C - I)
+              mjtNum C[9], E[9];
+              mju_mulMatTMat3(C, F, F);
+              for (int j = 0; j < 9; j++) {
+                E[j] = 0.5 * C[j];
+              }
+              E[0] -= 0.5; E[4] -= 0.5; E[8] -= 0.5;
+
+              // 3 invariants of E
+              mjtNum I1 = E[0] + E[4] + E[8];
+              mjtNum trE2 = E[0]*E[0] + E[1]*E[3] + E[2]*E[6]
+                          + E[3]*E[1] + E[4]*E[4] + E[5]*E[7]
+                          + E[6]*E[2] + E[7]*E[5] + E[8]*E[8];
+              mjtNum I2 = 0.5 * (I1*I1 - trE2);
+              mjtNum I3 = mat3_det(E);
+
+              // shape function gradients at Gauss point
+              mjtNum grad[27][3];
+              shape_gradients(order, p, grad);
+
+              for (int s = 0; s < 6; s++) {
+                // skip I1,I2,I3 for trilinear (B-bar handles vol)
+                if (order == 1 && (s == 0 || s == 1 || s == 2)) {
+                  continue;
+                }
+
+                mjtNum dSdE[9];
+                mju_zero(dSdE, 9);
+
+                if (s == 0) {
+                  cpos[0] = I1;
+                  dSdE[0] = dSdE[4] = dSdE[8] = 1.0;
+                } else if (s == 1) {
+                  cpos[0] = I2;
+                  dSdE[0] = I1-E[0]; dSdE[4] = I1-E[4];
+                  dSdE[8] = I1-E[8];
+                  dSdE[1] = -E[1]; dSdE[3] = -E[3];
+                  dSdE[2] = -E[2]; dSdE[6] = -E[6];
+                  dSdE[5] = -E[5]; dSdE[7] = -E[7];
+                } else if (s == 2) {
+                  cpos[0] = I3;
+                  mat3_cofactor(E, dSdE);
+                } else {
+                  int offdiag_idx[3] = {1, 2, 5};
+                  int ij = offdiag_idx[s - 3];
+                  cpos[0] = E[ij];
+                  dSdE[ij] = 1.0;
+                }
+
+                // compute local dS/dx for cell nodes
+                invariant_dSdx(npc, grad, F, Fref_inv, dSdE,
+                               dSdx_local);
+
+                // scatter to global dSdx
+                mju_zero(dSdx, 3*nodenum);
+                for (int n = 0; n < npc; n++) {
+                  mju_addTo3(dSdx + 3*gindices[n], dSdx_local + 3*n);
+                }
+
+                strain_jacobian(nodenum, nv, dSdx, node_jac, strain_jac);
+
+                if (issparse) {
+                  mj_markStack(d);
+                  mjtNum* sj = mjSTACKALLOC(d, combined_nnz, mjtNum);
+                  for (int k = 0; k < combined_nnz; k++) {
+                    sj[k] = strain_jac[combined_chain[k]];
+                  }
+                  mj_addConstraint(m, d, sj, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
+                                   combined_nnz, combined_chain);
+                  mj_freeStack(d);
+                } else {
+                  mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
+                }
+              }
             }
-            mj_addConstraint(m, d, sparse_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i,
-                             combined_nnz, combined_chain);
-          } else {
-            mj_addConstraint(m, d, strain_jac, cpos, 0, 0, 1, mjCNSTR_EQUALITY, i, 0, NULL);
           }
         }
       }
@@ -1899,10 +1926,13 @@ void mj_diagApprox(const mjModel* m, mjData* d) {
         int nstart = m->flex_nodeadr[flex_id];
         int order = m->flex_interp[flex_id];
 
-        // compute constraint count: trilinear (2 + 3*8 = 26), quadratic (6*27 = 162)
+        // compute constraint count per cell, then multiply by ncells
         int nquad = order + 1;
         int ngauss = nquad * nquad * nquad;
-        int nconstraint = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss);
+        int ncells = m->flex_cellnum[3*flex_id+0]
+                   * m->flex_cellnum[3*flex_id+1]
+                   * m->flex_cellnum[3*flex_id+2];
+        int nconstraint = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
 
         mjtNum avg_invweight = 0;
         for (int n = 0; n < nodenum; n++) {
@@ -2510,7 +2540,10 @@ static int mj_ne(const mjModel* m, mjData* d, int* nnz) {
       }
       int nquad = order + 1;  // 2 for order=1, 3 for order=2
       int ngauss = nquad * nquad * nquad;  // 8 or 27
-      size = (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss);  // 26 or 162
+      int ncells = m->flex_cellnum[3*id[0]+0]
+                 * m->flex_cellnum[3*id[0]+1]
+                 * m->flex_cellnum[3*id[0]+2];
+      size = ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
 
       if (nnz) {
         // Count unique DOFs across all node bodies (matching instantiation)

src/engine/engine_core_smooth.c

@@ -580,9 +580,11 @@ void mj_flex(const mjModel* m, mjData* d) {
       }
     }
 
-    // trilinear interpolation
+    // trilinear/quadratic interpolation
     else {
-      mjtNum nodexpos[3*mjMAXFLEXNODES];
+      int nodenum = nend - nstart;
+      mj_markStack(d);
+      mjtNum* nodexpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
       if (m->flex_centered[f]) {
         for (int i=nstart; i < nend; i++) {
           mji_copy3(nodexpos + 3*(i-nstart), d->xpos + 3*m->flex_nodebodyid[i]);
@@ -596,14 +598,26 @@ void mj_flex(const mjModel* m, mjData* d) {
       }
 
       int order = m->flex_interp[f];
-      if (nend - nstart != (order + 1) * (order + 1) * (order + 1)) {
+      int cx = m->flex_cellnum[3*f+0];
+      int cy = m->flex_cellnum[3*f+1];
+      int cz = m->flex_cellnum[3*f+2];
+      int nx_g = cx * order + 1;
+      int ny_g = cy * order + 1;
+      int nz_g = cz * order + 1;
+      if (nend - nstart != nx_g * ny_g * nz_g) {
         mjERROR("flex_interp_order mismatch");
       }
 
       for (int i=vstart; i < vend; i++) {
         mju_zero3(d->flexvert_xpos+3*i);
-        mju_interpolate3D(d->flexvert_xpos+3*i, m->flex_vert0 + 3*i, nodexpos, order);
+
+        // cell lookup: get local coords and node indices
+        mjtNum local[3];
+        int nodeindices[27];  // max npc for quadratic: 3^3 = 27
+        mju_cellLookup(m->flex_vert0 + 3*i, m->flex_cellnum+3*f, order, local, nodeindices);
+        mju_interpolate3D(d->flexvert_xpos+3*i, local, nodexpos, order, nodeindices);
       }
+      mj_freeStack(d);
     }
   }
 
@@ -2617,7 +2631,10 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
       if (order && nodenum) {
         int nquad = order + 1;
         int ngauss = nquad * nquad * nquad;
-        i += (order == 1) ? (2 + 3 * ngauss) : (6 * ngauss);
+        int ncells = m->flex_cellnum[3*k+0]
+                   * m->flex_cellnum[3*k+1]
+                   * m->flex_cellnum[3*k+2];
+        i += ncells * ((order == 1) ? (2 + 3 * ngauss) : (6 * ngauss));
       }
       break;
     }

src/engine/engine_core_util.c

@@ -987,6 +987,34 @@ void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
 
 //-------------------------- miscellaneous utilities -----------------------------------------------
 
+// gather global node positions and velocities
+void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel) {
+  int nodenum = m->flex_nodenum[f];
+  int nstart = m->flex_nodeadr[f];
+  int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
+
+  // compute positions
+  if (m->flex_centered[f]) {
+    for (int i=0; i < nodenum; i++) {
+      mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
+      if (vel) {
+        mju_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
+      }
+    }
+  } else {
+    mjtNum screw[6];
+    for (int i=0; i < nodenum; i++) {
+      mju_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
+      mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
+      if (vel) {
+        mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
+        mju_copy3(vel + 3*i, screw + 3);
+      }
+    }
+  }
+}
+
+
 // extract 6D force:torque for one contact, in contact frame
 void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]) {
   mjContact* con;

src/engine/engine_core_util.h

@@ -129,6 +129,9 @@ MJAPI void mj_local2Global(mjData* d, mjtNum xpos[3], mjtNum xmat[9],
 
 //-------------------------- miscellaneous ---------------------------------------------------------
 
+// gather global node positions and velocities
+MJAPI void mju_flexGatherState(const mjModel* m, mjData* d, int f, mjtNum* xpos, mjtNum* vel);
+
 // extract 6D force:torque for one contact, in contact frame
 MJAPI void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum result[6]);
 

src/engine/engine_derivative.c

@@ -872,12 +872,12 @@ typedef enum {
 
 // shared kernel for flex interpolation derivatives, scale = s1 + s2*damping
 //  op: operation type (VEC, or ADDH)
-//  res: output vector (VEC) or dense H matrix (ADDH)
+//  res: output vector (VEC) or banded H matrix (ADDH)
 //  vec: input vector for VEC operation, NULL otherwise
-//  dof_indices, ndof: DOF mapping for ADDH, ignored otherwise
+//  dof_indices, ndof, nband: DOF mapping and band width for ADDH, ignored otherwise
 static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
                                   mjtNum* res, const mjtNum* vec, mjtNum s1, mjtNum s2,
-                                  const int* dof_indices, int ndof) {
+                                  const int* dof_indices, int ndof, int nband) {
   int nv = m->nv;
 
   // build global2local map for ADDH
@@ -915,127 +915,146 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
       continue;
     }
 
+    int order = m->flex_interp[f];
+    int npc = (order+1)*(order+1)*(order+1);
+    int cx = m->flex_cellnum[3*f+0];
+    int cy = m->flex_cellnum[3*f+1];
+    int cz = m->flex_cellnum[3*f+2];
+
     int nodenum = m->flex_nodenum[f];
     int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
 
     // standard stack allocation
     mj_markStack(d);
     mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-    mjtNum* K_rot = mjSTACKALLOC(d, 9*nodenum*nodenum, mjtNum);
 
-    // sparse Jacobian allocations
-    int dim = 3 * nodenum;
-    int* rownnz = mjSTACKALLOC(d, dim, int);
-    int* rowadr = mjSTACKALLOC(d, dim, int);
-    mjtNum* J_val = mjSTACKALLOC(d, dim*nv, mjtNum);
-    int* J_colind = mjSTACKALLOC(d, dim*nv, int);
+    // per-cell arrays
+    int dim_c = 3 * npc;
+    mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+    mjtNum* K_rot_cell = mjSTACKALLOC(d, dim_c*dim_c, mjtNum);
+
+    // sparse Jacobian for one cell
+    int* J_rownnz = mjSTACKALLOC(d, dim_c, int);
+    int* J_rowadr = mjSTACKALLOC(d, dim_c, int);
+    mjtNum* J_val = mjSTACKALLOC(d, dim_c*nv, mjtNum);
+    int* J_colind = mjSTACKALLOC(d, dim_c*nv, int);
 
     // temp allocations for chain
     int* chain_colind = mjSTACKALLOC(d, nv, int);
     mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum);
 
-    // compute positions, rotation and Jacobian
-    mjtNum quat[4] = {1, 0, 0, 0};
-    mj_flexInterpState(m, d, f, xpos, NULL, quat);
+    // gather raw node positions (unrotated)
+    mju_flexGatherState(m, d, f, xpos, NULL);
 
-    // compute generalized stiffness in global frame: K_rot = R * K * R^T
-    mjtNum R[9];
-    mju_quat2Mat(R, quat);       // R = R_global2local
-    mjtNum RT[9];
-    mju_transpose(RT, R, 3, 3);  // RT = R_local2global
+    // loop over cells
+    int cell_idx = 0;
+    for (int ci = 0; ci < cx; ci++) {
+      for (int cj = 0; cj < cy; cj++) {
+        for (int ck = 0; ck < cz; ck++) {
+          // gather cell-local node positions
+          int gindices[125];  // max npc = 125 for quadratic
+          mjtNum quat[4];
+          mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos, NULL, NULL,
+                                  xpos_c, NULL, NULL, gindices, quat);
 
-    // blockwise rotation: K_rot(i,j) = scale * RT * K_local(i,j) * R
-    // note: k stores -K, so K_rot = scale * (-K_phys)
-    for (int i=0; i < nodenum; i++) {
-      for (int j=0; j < nodenum; j++) {
-        mjtNum blk[9], tmp[9];
+          // R = R_global2local, RT = R_local2global
+          mjtNum R[9], RT[9];
+          mju_quat2Mat(R, quat);
+          mju_transpose(RT, R, 3, 3);
 
-        // get K_local(i,j)
-        int adr = (3*i)*(3*nodenum) + 3*j;
-        for (int r=0; r < 3; r++) {
-          for (int c=0; c < 3; c++) {
-            blk[3*r+c] = k[adr + r*(3*nodenum) + c];
+          // get cell stiffness
+          mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc;
+
+          // compute K_rot_cell = RT * K_cell * R (block-wise)
+          mju_zero(K_rot_cell, dim_c*dim_c);
+          for (int a = 0; a < npc; a++) {
+            for (int b = 0; b < npc; b++) {
+              mjtNum blk[9], tmp[9];
+
+              // get K_cell(a,b) 3x3 block
+              int adr_cell = (3*a)*(3*npc) + 3*b;
+              for (int r = 0; r < 3; r++) {
+                for (int c = 0; c < 3; c++) {
+                  blk[3*r+c] = k_cell[adr_cell + r*(3*npc) + c];
+                }
+              }
+
+              // tmp = K * R
+              mju_mulMatMat3(tmp, blk, R);
+              // blk = RT * tmp = RT * K * R
+              mju_mulMatMat3(blk, RT, tmp);
+
+              // store in K_rot_cell at (a, b)
+              int adr_out = (3*a)*dim_c + 3*b;
+              for (int r = 0; r < 3; r++) {
+                for (int c = 0; c < 3; c++) {
+                  K_rot_cell[adr_out + r*dim_c + c] = scale * blk[3*r+c];
+                }
+              }
+            }
           }
-        }
 
-        // tmp = K * R
-        mju_mulMatMat3(tmp, blk, R);
+          // construct sparse Jacobian for this cell's nodes
+          int current_adr = 0;
+          for (int n = 0; n < npc; n++) {
+            int bid = bodyid[gindices[n]];
+            int chain_nnz = mj_bodyChain(m, bid, chain_colind);
+            mj_jacSparse(m, d, blk_jac, NULL, xpos+3*gindices[n], bid,
+                         chain_nnz, chain_colind, /*flg_skipcommon=*/0);
 
-        // blk = RT * tmp = RT * K * R
-        mju_mulMatMat3(blk, RT, tmp);
+            for (int r = 0; r < 3; r++) {
+              int row_idx = 3*n + r;
+              J_rownnz[row_idx] = chain_nnz;
+              J_rowadr[row_idx] = current_adr;
 
-        // store scaled into K_rot
-        for (int r=0; r < 3; r++) {
-          for (int c=0; c < 3; c++) {
-            K_rot[adr + r*(3*nodenum) + c] = scale * blk[3*r+c];
-          }
-        }
-      }
-    }
-
-    // construct sparse Jacobian J_val
-    int current_adr = 0;
-    for (int i=0; i < nodenum; i++) {
-        // get chain for this node
-        int chain_nnz = mj_bodyChain(m, bodyid[i], chain_colind);
-
-        // compute sparse Jacobian for this node (3 rows)
-        mj_jacSparse(m, d, blk_jac, NULL, xpos+3*i, bodyid[i], chain_nnz, chain_colind,
-                     /*flg_skipcommon=*/0);
-
-        // copy to sparse structure
-        for (int r=0; r<3; r++) {
-            int row_idx = 3*i + r;
-            rownnz[row_idx] = chain_nnz;
-            rowadr[row_idx] = current_adr;
-
-            for (int idx=0; idx<chain_nnz; idx++) {
+              for (int idx = 0; idx < chain_nnz; idx++) {
                 J_colind[current_adr] = chain_colind[idx];
                 J_val[current_adr] = blk_jac[r*chain_nnz + idx];
                 current_adr++;
-            }
-        }
-    }
-
-    // perform operation
-    if (op == mjFLEXOP_VEC) {
-      // res += J^T * K_rot * J * vec
-      addJTBJ_mulSparse(m, d, res, vec, rownnz, rowadr, J_colind, J_val, K_rot, dim);
-    } else if (op == mjFLEXOP_ADDH) {
-      // H += -J^T * K_rot * J
-      // H is dense ndof x ndof
-
-      // reuse stack for J_reduced (but now we extract from sparse J)
-      mjtNum* J_reduced = mjSTACKALLOC(d, dim*ndof, mjtNum);
-      mju_zero(J_reduced, dim*ndof);
-
-      // extract columns of J into J_reduced
-      for (int i=0; i<dim; i++) {
-          int nnz = rownnz[i];
-          int adr = rowadr[i];
-          for (int idx=0; idx<nnz; idx++) {
-              int global_col = J_colind[adr + idx];
-              int local_idx = global2local[global_col];
-              if (local_idx >= 0) {
-                  J_reduced[i*ndof + local_idx] = J_val[adr + idx];
               }
+            }
           }
-      }
 
-      // H -= J_reduced^T * K_rot * J_reduced
-      // K_rot * J_reduced (dim x ndof)
-      mjtNum* KJ = mjSTACKALLOC(d, dim*ndof, mjtNum);
-      mju_mulMatMat(KJ, K_rot, J_reduced, dim, dim, ndof);
+          // apply operation with cell's K_rot and J
+          if (op == mjFLEXOP_VEC) {
+            addJTBJ_mulSparse(m, d, res, vec, J_rownnz, J_rowadr, J_colind,
+                              J_val, K_rot_cell, dim_c);
+          } else if (op == mjFLEXOP_ADDH) {
+            mj_markStack(d);
+            // H -= J_cell^T * K_rot_cell * J_cell (banded format)
+            mjtNum* J_reduced = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
+            mju_zero(J_reduced, dim_c*ndof);
 
-      // H[i, j] -= sum_k J_reduced[k, i] * KJ[k, j]
-      for (int i=0; i<ndof; i++) {
-        for (int j=0; j<ndof; j++) {
-          mjtNum val = 0;
-          for (int dim_idx=0; dim_idx<dim; dim_idx++) {
-            val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
+            for (int i = 0; i < dim_c; i++) {
+              int nnz = J_rownnz[i];
+              int adr = J_rowadr[i];
+              for (int idx = 0; idx < nnz; idx++) {
+                int global_col = J_colind[adr + idx];
+                int local_idx = global2local[global_col];
+                if (local_idx >= 0) {
+                  J_reduced[i*ndof + local_idx] = J_val[adr + idx];
+                }
+              }
+            }
+
+            // KJ = K_rot_cell * J_reduced  (dim_c x ndof)
+            mjtNum* KJ = mjSTACKALLOC(d, dim_c*ndof, mjtNum);
+            mju_mulMatMat(KJ, K_rot_cell, J_reduced, dim_c, dim_c, ndof);
+
+            // H[i,j] -= J_reduced[k,i] * KJ[k,j], store lower triangle in banded format
+            for (int i = 0; i < ndof; i++) {
+              for (int j = mjMAX(0, i-nband+1); j <= i; j++) {
+                mjtNum val = 0;
+                for (int dim_idx = 0; dim_idx < dim_c; dim_idx++) {
+                  val += J_reduced[dim_idx*ndof + i] * KJ[dim_idx*ndof + j];
+                }
+                res[i*nband + nband-1-(i-j)] -= val;
+              }
+            }
+            mj_freeStack(d);
           }
-          // res is H
-          res[i*ndof + j] -= val;
+
+          cell_idx++;
         }
       }
     }
@@ -1053,15 +1072,16 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d, mjtFlexOp op,
 // compute res += (h^2 + h*damping) * J'*K*J * vec, for all interpolated flexes
 void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec, mjtNum h) {
   // s1=h*h, s2=h => scale = h*h + h*damping
-  mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0);
+  mjd_flexInterp_kernel(m, d, mjFLEXOP_VEC, res, vec, h * h, h, NULL, 0, 0);
 }
 
 
-// add (h^2 + h*damping) * J'*K*J to dense matrix H, for all interpolated flexes
-//  H: dense ndof x ndof matrix
+// add (h^2 + h*damping) * J'*K*J to banded matrix H, for all interpolated flexes
+//  H: banded ndof x nband matrix (lower triangle, band storage)
 //  dof_indices: maps local indices to global DOFs
-void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h) {
-  mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof);
+void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices,
+                         int ndof, int nband, mjtNum h) {
+  mjd_flexInterp_kernel(m, d, mjFLEXOP_ADDH, H, NULL, h * h, h, dof_indices, ndof, nband);
 }
 
 

src/engine/engine_derivative.h

@@ -49,7 +49,7 @@ MJAPI void mjd_flexInterp_mulKD(const mjModel* m, mjData* d, mjtNum* res, const
 
 // assemble flex stiffness matrix H_flex: H += h*h*K + h*D
 //  H is a dense matrix of size ndof x ndof, dof_indices maps local rows/cols to global DOFs
-MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, mjtNum h);
+MJAPI void mjd_flexInterp_addH(const mjModel* m, mjData* d, mjtNum* H, const int* dof_indices, int ndof, int nband, mjtNum h);
 
 
 #ifdef __cplusplus

src/engine/engine_forward.c

@@ -1350,11 +1350,12 @@ void mj_RungeKutta(const mjModel* m, mjData* d, int N) {
 }
 
 
-// context for flex interp reduced dense factorization/solve
+// context for flex interp reduced banded factorization/solve
 typedef struct {
-  mjtNum* H;              // dense Cholesky-factored matrix (ndof x ndof)
+  mjtNum* H;              // banded Cholesky-factored matrix (ndof x nband)
   int* dof_indices;       // global DOF index for each local flex DOF
   int ndof;               // number of flex DOFs
+  int nband;              // half-bandwidth + 1 (number of band columns)
   int ncoupling;          // number of off-diagonal coupling terms
   mjtNum* coupling_val;   // coupling coefficient values
   int* coupling_row;      // local flex row index for each coupling term
@@ -1391,7 +1392,7 @@ static void flexInterp_collect(const mjModel* m, int f,
 }
 
 
-// build and factor the reduced dense matrix for flex interp DOFs
+// build and factor the reduced banded matrix for flex interp DOFs
 //   mark/free stack handled by caller
 static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv) {
   FlexInterpContext ctx = {0};
@@ -1456,19 +1457,36 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
   const int* colind = implicit ? m->D_colind : m->M_colind;
   const mjtNum* source = implicit ? d->qLU : d->qH;
 
-  // count coupling terms (off-diagonal: flex row, non-flex col)
+  // get precomputed bandwidth
+  int bandwidth = 0;
+  for (int f=0; f < m->nflex; f++) {
+    if (m->flex_interp[f]) {
+      if (m->flex_bandwidth[f] > bandwidth) {
+        bandwidth = m->flex_bandwidth[f];
+      }
+    }
+  }
+
+  // compute ncoupling from sparse matrix entries
   int ncoupling = 0;
   for (int i=0; i < ndof; i++) {
     int row = dof_indices[i];
     int start = rowadr[row];
     int end = start + rownnz[row];
     for (int k=start; k < end; k++) {
-      if (global2local[colind[k]] < 0) {
+      int local_j = global2local[colind[k]];
+      if (local_j < 0) {
         ncoupling++;
       }
     }
   }
 
+  // nband = bandwidth + 1 (includes diagonal)
+  int nband = bandwidth + 1;
+
+  // cap nband at ndof (dense fallback for small systems)
+  if (nband > ndof) nband = ndof;
+
   // allocate coupling storage
   mjtNum* coupling_val = NULL;
   int* coupling_row = NULL;
@@ -1479,9 +1497,9 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
     coupling_col = mjSTACKALLOC(d, ncoupling, int);
   }
 
-  // build H_flex (dense) from qLU (implicit) or qH (implicitfast)
-  mjtNum* H = mjSTACKALLOC(d, ndof*ndof, mjtNum);
-  mju_zero(H, ndof*ndof);
+  // build H_flex (banded) from qLU (implicit) or qH (implicitfast)
+  mjtNum* H = mjSTACKALLOC(d, ndof*nband, mjtNum);
+  mju_zero(H, ndof*nband);
 
   int coup_cnt = 0;
   for (int i=0; i < ndof; i++) {
@@ -1492,7 +1510,13 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
       int col = colind[k];
       int local_j = global2local[col];
       if (local_j >= 0) {
-        H[i*ndof+local_j] = source[k];
+        // store lower triangle only: row i, col local_j, where i >= local_j
+        if (i >= local_j) {
+          H[i*nband + nband-1-(i-local_j)] = source[k];
+        } else {
+          // upper triangle entry: store symmetrically in lower triangle
+          H[local_j*nband + nband-1-(local_j-i)] = source[k];
+        }
       } else if (coup_cnt < ncoupling) {
         coupling_val[coup_cnt] = source[k];
         coupling_row[coup_cnt] = i;
@@ -1502,14 +1526,15 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
     }
   }
 
-  // add flex stiffness and factorize
-  mjd_flexInterp_addH(m, d, H, dof_indices, ndof, m->opt.timestep);
-  mju_cholFactor(H, ndof, mjMINVAL);
+  // add flex stiffness in banded format and factorize
+  mjd_flexInterp_addH(m, d, H, dof_indices, ndof, nband, m->opt.timestep);
+  mju_cholFactorBand(H, ndof, nband, 0, 0, 0);
 
   // store results in context
   ctx.H = H;
   ctx.dof_indices = dof_indices;
   ctx.ndof = ndof;
+  ctx.nband = nband;
   ctx.ncoupling = ncoupling;
   ctx.coupling_val = coupling_val;
   ctx.coupling_row = coupling_row;
@@ -1518,7 +1543,7 @@ static FlexInterpContext flexInterp_factor(const mjModel* m, mjData* d, int nv)
 }
 
 
-// solve the reduced dense system for flex interp DOFs, overwrite qacc
+// solve the reduced banded system for flex interp DOFs, overwrite qacc
 static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContext* ctx,
                              mjtNum* qacc, const mjtNum* qfrc, int nv) {
   int ndof = ctx->ndof;
@@ -1544,8 +1569,8 @@ static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContex
     qfrc_flex[ctx->coupling_row[k]] -= ctx->coupling_val[k] * qacc[ctx->coupling_col[k]];
   }
 
-  // solve and scatter back
-  mju_cholSolve(qfrc_flex, ctx->H, qfrc_flex, ndof);
+  // solve with banded Cholesky and scatter back
+  mju_cholSolveBand(qfrc_flex, ctx->H, qfrc_flex, ndof, ctx->nband, 0);
   mju_scatter(qacc, qfrc_flex, ctx->dof_indices, ndof);
 }
 

src/engine/engine_passive.c

@@ -58,61 +58,7 @@ static void inline GradSquaredLengths(mjtNum gradient[6][2][3],
   }
 }
 
-// compute interpolated flex state: xpos, vel, quat
-//  f: flex index
-//  xpos: (output) 3*nodenum
-//  vel: (output) 3*nodenum, can be NULL
-//  quat: (output) 4, rotation from global to local
-void mj_flexInterpState(const mjModel* m, mjData* d, int f,
-                        mjtNum* xpos, mjtNum* vel, mjtNum* quat) {
-  int nodenum = m->flex_nodenum[f];
-  int nstart = m->flex_nodeadr[f];
-  int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
-  mjtNum com[3] = {0};
 
-  // compute positions
-  if (m->flex_centered[f]) {
-    for (int i=0; i < nodenum; i++) {
-      mji_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
-      if (vel) {
-        mji_copy3(vel + 3*i, d->qvel + m->body_dofadr[bodyid[i]]);
-      }
-    }
-  } else {
-    mjtNum screw[6];
-    for (int i=0; i < nodenum; i++) {
-      mji_mulMatVec3(xpos + 3*i, d->xmat + 9*bodyid[i], m->flex_node + 3*(i+nstart));
-      mji_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
-      if (vel) {
-        mj_objectVelocity(m, d, mjOBJ_BODY, bodyid[i], screw, 0);
-        mji_copy3(vel + 3*i, screw + 3);
-      }
-    }
-  }
-
-  // compute center of mass
-  for (int i = 0; i < nodenum; i++) {
-    mji_addToScl3(com, xpos+3*i, 1.0/nodenum);
-  }
-
-  // compute the Jacobian at the center of mass
-  mjtNum mat[9] = {0};
-  mjtNum p[3] = {.5, .5, .5};
-  mju_defGradient(mat, p, xpos, m->flex_interp[f]);
-
-  // find rotation
-  mju_mat2Rot(quat, mat);
-  mju_negQuat(quat, quat);
-
-  // rotate vertices to quat and add reference center of mass
-  for (int i = 0; i < nodenum; i++) {
-    mju_rotVecQuat(xpos+3*i, xpos+3*i, quat);
-    mji_addTo3(xpos+3*i, p);
-    if (vel) {
-      mju_rotVecQuat(vel+3*i, vel+3*i, quat);
-    }
-  }
-}
 
 // spring and damper forces
 static void mj_springdamper(const mjModel* m, mjData* d) {
@@ -284,42 +230,111 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
     }
 
     if (m->flex_interp[f]) {
+      int order = m->flex_interp[f];
+      int npc = (order+1)*(order+1)*(order+1);  // nodes per cell
+      int cx = m->flex_cellnum[3*f+0];
+      int cy = m->flex_cellnum[3*f+1];
+      int cz = m->flex_cellnum[3*f+2];
+      int ny_g = cy * order + 1;
+      int nz_g = cz * order + 1;
+
       mj_markStack(d);
-      mjtNum* xpos = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-      mjtNum* displ = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-      mjtNum* vel = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-      mjtNum* frc = mjSTACKALLOC(d, 3*nodenum, mjtNum);
-      mjtNum* dmp = mjSTACKALLOC(d, 3*nodenum, mjtNum);
+
+      // allocate global arrays
+      mjtNum* xpos_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
+      mjtNum* vel_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
+      mjtNum* frc_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
+      mjtNum* dmp_g = mjSTACKALLOC(d, 3*nodenum, mjtNum);
       mjtNum* xpos0 = m->flex_node0 + 3*m->flex_nodeadr[f];
       int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
 
-      mjtNum quat[4] = {1, 0, 0, 0};
-      mj_flexInterpState(m, d, f, xpos, vel, quat);
+      // gather global node positions and velocities (unrotated)
+      mju_flexGatherState(m, d, f, xpos_g, vel_g);
 
-      // compute displacement
-      for (int i = 0; i < nodenum; i++) {
-        mji_addScl3(displ+3*i, xpos+3*i, xpos0+3*i, -1);
+      // zero global force accumulators
+      mju_zero(frc_g, 3*nodenum);
+      mju_zero(dmp_g, 3*nodenum);
+
+      // per-cell arrays
+      mjtNum* xpos_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* vel_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* xpos0_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* displ_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* frc_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+      mjtNum* dmp_c = mjSTACKALLOC(d, 3*npc, mjtNum);
+
+      // loop over cells
+      int cell_idx = 0;
+      for (int ci = 0; ci < cx; ci++) {
+        for (int cj = 0; cj < cy; cj++) {
+          for (int ck = 0; ck < cz; ck++) {
+            // gather cell-local node data
+            mjtNum quat[4];
+            mjtNum p[3] = {.5, .5, .5};
+            mju_flexGatherCellState(order, cy, cz, ci, cj, ck, xpos_g, vel_g, xpos0,
+                                    xpos_c, vel_c, xpos0_c, NULL, quat);
+
+            // rotate to corotational frame
+            for (int n = 0; n < npc; n++) {
+              mju_rotVecQuat(xpos_c+3*n, xpos_c+3*n, quat);
+              mji_addTo3(xpos_c+3*n, p);
+              mju_rotVecQuat(vel_c+3*n, vel_c+3*n, quat);
+            }
+
+            // compute displacement
+            for (int n = 0; n < npc; n++) {
+              mji_addScl3(displ_c+3*n, xpos_c+3*n, xpos0_c+3*n, -1);
+            }
+
+            // get cell stiffness matrix
+            mjtNum* k_cell = k + cell_idx * 3*npc * 3*npc;
+
+            // compute force in corotational frame
+            if (enbl_spring) {
+              mju_mulMatVec(frc_c, k_cell, displ_c, 3*npc, 3*npc);
+            }
+            if (enbl_damper) {
+              mju_mulMatVec(dmp_c, k_cell, vel_c, 3*npc, 3*npc);
+            }
+
+            // rotate back to global frame and scatter
+            mju_negQuat(quat, quat);
+            int local = 0;
+            for (int li = 0; li <= order; li++) {
+              for (int lj = 0; lj <= order; lj++) {
+                for (int lk = 0; lk <= order; lk++) {
+                  int gi = ci*order + li;
+                  int gj = cj*order + lj;
+                  int gk = ck*order + lk;
+                  int gidx = gi*ny_g*nz_g + gj*nz_g + gk;
+                  mjtNum qfrc[3], qdmp[3];
+                  mji_rotVecQuat(qfrc, frc_c+3*local, quat);
+                  mji_rotVecQuat(qdmp, dmp_c+3*local, quat);
+                  if (enbl_spring) {
+                    mji_addTo3(frc_g + 3*gidx, qfrc);
+                  }
+                  if (enbl_damper) {
+                    mji_addTo3(dmp_g + 3*gidx, qdmp);
+                  }
+                  local++;
+                }
+              }
+            }
+
+            cell_idx++;
+          }
+        }
       }
 
-      // compute force in the stretch frame
-      if (enbl_spring) mju_mulMatVec(frc, k, displ, 3*nodenum, 3*nodenum);
-
-      // compute damping force in stretch frame
-      if (enbl_damper) mju_mulMatVec(dmp, k, vel, 3*nodenum, 3*nodenum);
-
-      // rotate forces to global frame and add to qfrc
-      mju_negQuat(quat, quat);
+      // apply accumulated forces to bodies
       for (int i = 0; i < nodenum; i++) {
-        mjtNum qfrc[3], qdmp[3];
-        mji_rotVecQuat(qfrc, frc+3*i, quat);
-        mji_rotVecQuat(qdmp, dmp+3*i, quat);
-        mju_scl3(qdmp, qdmp, m->flex_damping[f]);
+        mju_scl3(dmp_g+3*i, dmp_g+3*i, m->flex_damping[f]);
         if (m->flex_centered[f]) {
-          if (enbl_spring) mji_addTo3(d->qfrc_spring+m->body_dofadr[bodyid[i]], qfrc);
-          if (enbl_damper) mji_addTo3(d->qfrc_damper+m->body_dofadr[bodyid[i]], qdmp);
+          if (enbl_spring) mji_addTo3(d->qfrc_spring + m->body_dofadr[bodyid[i]], frc_g+3*i);
+          if (enbl_damper) mji_addTo3(d->qfrc_damper + m->body_dofadr[bodyid[i]], dmp_g+3*i);
         } else {
-          if (enbl_spring) mj_applyFT(m, d, qfrc, 0, xpos+3*i, bodyid[i], d->qfrc_spring);
-          if (enbl_damper) mj_applyFT(m, d, qdmp, 0, xpos+3*i, bodyid[i], d->qfrc_damper);
+          if (enbl_spring) mj_applyFT(m, d, frc_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_spring);
+          if (enbl_damper) mj_applyFT(m, d, dmp_g+3*i, 0, xpos_g+3*i, bodyid[i], d->qfrc_damper);
         }
       }
 

src/engine/engine_passive.h

@@ -28,9 +28,7 @@ extern "C" {
 // all passive forces
 MJAPI void mj_passive(const mjModel* m, mjData* d);
 
-// compute interpolated flex state: xpos, vel, quat
-MJAPI void mj_flexInterpState(const mjModel* m, mjData* d, int f,
-                              mjtNum* xpos, mjtNum* vel, mjtNum* quat);
+
 
 
 //------------------------- fluid models -----------------------------------------------------------

src/engine/engine_setconst.c

@@ -636,6 +636,135 @@ static void makeFlexSparse(mjModel* m, mjData* d) {
   mj_freeStack(d);
 }
 
+// compute flex bandwidth for trilinear interpolation
+static void makeFlexBandwidth(mjModel* m, mjData* d) {
+  if (!m->nflex) {
+    return;
+  }
+
+  mj_markStack(d);
+  int* chain_dofs = mjSTACKALLOC(d, m->nv, int);
+  int* seen_dof = mjSTACKALLOC(d, m->nv, int);
+  int* dof_indices = mjSTACKALLOC(d, m->nv, int);
+  int* global2local = mjSTACKALLOC(d, m->nv, int);
+
+  mju_zeroInt(seen_dof, m->nv);
+  for (int i = 0; i < m->nv; i++) {
+    global2local[i] = -1;
+  }
+
+  int ndof = 0;
+  for (int f = 0; f < m->nflex; f++) {
+    if (m->flex_interp[f]) {
+      int nodenum = m->flex_nodenum[f];
+      int nodeadr = m->flex_nodeadr[f];
+      for (int n = 0; n < nodenum; n++) {
+        int b = m->flex_nodebodyid[nodeadr + n];
+        // only the body's own DOFs enter the reduced banded flex system;
+        // ancestor DOFs are solved by the global factorization and coupled
+        // via off-diagonal correction (see flexInterp_solve in engine_forward)
+        int chain_nnz;
+        if (m->body_dofnum[b] == 0) {
+          chain_nnz = mj_bodyChain(m, b, chain_dofs);
+        } else {
+          chain_nnz = m->body_dofnum[b];
+          for (int j = 0; j < chain_nnz; j++) {
+            chain_dofs[j] = m->body_dofadr[b] + j;
+          }
+        }
+        for (int i = 0; i < chain_nnz; i++) {
+          int dof = chain_dofs[i];
+          if (!seen_dof[dof]) {
+            seen_dof[dof] = 1;
+            dof_indices[ndof] = dof;
+            global2local[dof] = ndof++;
+          }
+        }
+      }
+    }
+  }
+
+  int bandwidth = 0;
+  if (ndof > 0) {
+    // check sparse matrix coupling (both D and M)
+    for (int integrator = 0; integrator < 2; integrator++) {
+      const int* rownnz = (integrator == 0) ? m->D_rownnz : m->M_rownnz;
+      const int* rowadr = (integrator == 0) ? m->D_rowadr : m->M_rowadr;
+      const int* colind = (integrator == 0) ? m->D_colind : m->M_colind;
+
+      // D arrays are only allocated for implicit integrators
+      if (!rownnz) continue;
+
+      for (int i = 0; i < ndof; i++) {
+        int row = dof_indices[i];
+        int start = rowadr[row];
+        int end = start + rownnz[row];
+        for (int k = start; k < end; k++) {
+          int local_j = global2local[colind[k]];
+          if (local_j >= 0) {
+            int diff = i - local_j;
+            if (diff < 0) diff = -diff;
+            if (diff > bandwidth) bandwidth = diff;
+          }
+        }
+      }
+    }
+
+    // check stiffness coupling
+    for (int f = 0; f < m->nflex; f++) {
+      if (!m->flex_interp[f]) continue;
+      int order = m->flex_interp[f];
+      int nodeadr = m->flex_nodeadr[f];
+      int nodenum = m->flex_nodenum[f];
+      int cx = m->flex_cellnum[3*f+0];
+      int cy = m->flex_cellnum[3*f+1];
+      int cz = m->flex_cellnum[3*f+2];
+      int ny = cy * order + 1;
+      int nz = cz * order + 1;
+
+      for (int icx = 0; icx < cx; icx++) {
+        for (int icy = 0; icy < cy; icy++) {
+          for (int icz = 0; icz < cz; icz++) {
+            int min_local = ndof, max_local = -1;
+            for (int lx = 0; lx <= order; lx++) {
+              for (int ly = 0; ly <= order; ly++) {
+                for (int lz = 0; lz <= order; lz++) {
+                  int gx = icx * order + lx;
+                  int gy = icy * order + ly;
+                  int gz = icz * order + lz;
+                  int node_idx = gx * ny * nz + gy * nz + gz;  // non-negative by construction
+                  if (node_idx < nodenum) {
+                    int b = m->flex_nodebodyid[nodeadr + node_idx];
+                    int chain_nnz = mj_bodyChain(m, b, chain_dofs);
+                    for (int i = 0; i < chain_nnz; i++) {
+                      int dof = chain_dofs[i];
+                      int local = global2local[dof];
+                      if (local >= 0) {
+                        if (local < min_local) min_local = local;
+                        if (local > max_local) max_local = local;
+                      }
+                    }
+                  }
+                }
+              }
+            }
+            if (max_local >= 0 && max_local - min_local > bandwidth) {
+              bandwidth = max_local - min_local;
+            }
+          }
+        }
+      }
+    }
+  }
+
+  // store bandwidth for all flexes (global max)
+  for (int f = 0; f < m->nflex; f++) {
+    m->flex_bandwidth[f] = bandwidth;
+  }
+
+  mj_freeStack(d);
+}
+
 // align 2D flexes to the XY plane
 static void mj_alignFlex(mjModel* m, mjData* d) {
   for (int f = 0; f < m->nflex; f++) {
@@ -687,6 +816,7 @@ static void mj_alignFlex(mjModel* m, mjData* d) {
 static void set0(mjModel* m, mjData* d) {
   makeTendonSparse(m);
   makeFlexSparse(m, d);
+  makeFlexBandwidth(m, d);
   mj_alignFlex(m, d);
   int nv = m->nv;
   mjtNum A[36] = {0}, pos[3], quat[4];

src/engine/engine_util_misc.c

@@ -612,11 +612,116 @@ mjtNum mju_evalBasis(const mjtNum x[3], int i, int 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)
+//   order: interpolation order (1=trilinear, 2=triquadratic)
+//   local: output local parametric coordinates within cell [0,1]^3
+//   nodeindices: output array of global node indices for the cell (size (order+1)^3, may be NULL)
+//   returns: number of nodes per cell (order+1)^3
+int mju_cellLookup(const mjtNum coord[3], const int cellnum[3], int order, mjtNum local[3],
+                   int* nodeindices) {
+  int cx = cellnum[0], cy = cellnum[1], cz = cellnum[2];
+
+  // find containing cell
+  int ci = (int)mju_floor(coord[0] * cx);
+  int cj = (int)mju_floor(coord[1] * cy);
+  int ck = (int)mju_floor(coord[2] * cz);
+  ci = mjMIN(ci, cx - 1); ci = mjMAX(ci, 0);
+  cj = mjMIN(cj, cy - 1); cj = mjMAX(cj, 0);
+  ck = mjMIN(ck, cz - 1); ck = mjMAX(ck, 0);
+
+  // local parametric coordinates within cell
+  local[0] = mju_clip(coord[0] * cx - ci, 0, 1);
+  local[1] = mju_clip(coord[1] * cy - cj, 0, 1);
+  local[2] = mju_clip(coord[2] * cz - ck, 0, 1);
+
+  // build node indices for this cell
+  if (nodeindices) {
+    int ny_g = cy * order + 1;
+    int nz_g = cz * order + 1;
+    int ni = 0;
+    for (int li = 0; li <= order; li++) {
+      for (int lj = 0; lj <= order; lj++) {
+        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 npc = (order + 1) * (order + 1) * (order + 1);
+  return npc;
+}
+
+
 // interpolate a function at x with given interpolation coefficients and order n
-void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order) {
+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);
   for (int j=0; j < npoint; j++) {
-    mju_addToScl3(res, coeff+3*j, mju_evalBasis(x, j, order));
+    int idx = nodeindices ? nodeindices[j] : j;
+    mju_addToScl3(res, coeff+3*idx, mju_evalBasis(x, j, order));
+  }
+}
+
+
+static void flexInterpRotation(int order, const mjtNum* xpos_c,
+                               const mjtNum local[3], mjtNum* quat) {
+  mjtNum mat[9] = {0};
+
+  if (order > 0) {
+    mju_defGradient(mat, local, xpos_c, order);
+  } else {
+    // order 0: fallback to identity matrix
+    mat[0] = 1;
+    mat[4] = 1;
+    mat[8] = 1;
+  }
+
+  // find rotation
+  quat[0] = 1;
+  quat[1] = 0;
+  quat[2] = 0;
+  quat[3] = 0;
+  mju_mat2Rot(quat, mat);
+  mju_negQuat(quat, quat);
+}
+
+
+// gather cell-local quantities and optionally compute rotation
+void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck,
+                             const mjtNum* xpos_g, const mjtNum* vel_g, const mjtNum* xpos0_g,
+                             mjtNum* xpos_c, mjtNum* vel_c, mjtNum* xpos0_c,
+                             int* nodeindices, mjtNum* quat) {
+  int ny_g = cy * order + 1;
+  int nz_g = cz * order + 1;
+
+  int local = 0;
+  for (int li = 0; li <= order; li++) {
+    for (int lj = 0; lj <= order; lj++) {
+      for (int lk = 0; lk <= order; lk++) {
+        int gi = ci*order + li;
+        int gj = cj*order + lj;
+        int gk = ck*order + lk;
+        int gidx = gi*ny_g*nz_g + gj*nz_g + gk;
+
+        if (xpos_c && xpos_g) mju_copy3(xpos_c + 3*local, xpos_g + 3*gidx);
+        if (vel_c && vel_g) mju_copy3(vel_c + 3*local, vel_g + 3*gidx);
+        if (xpos0_c && xpos0_g) mju_copy3(xpos0_c + 3*local, xpos0_g + 3*gidx);
+        if (nodeindices) nodeindices[local] = gidx;
+
+        local++;
+      }
+    }
+  }
+
+  if (quat && xpos_c) {
+    mjtNum p[3] = {.5, .5, .5};
+    flexInterpRotation(order, xpos_c, p, quat);
   }
 }
 

src/engine/engine_util_misc.h

@@ -89,8 +89,20 @@ 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);
 
+// 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);
+
 // interpolate a function at x with given interpolation coefficients and order n
-MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order);
+MJAPI void mju_interpolate3D(mjtNum res[3], const mjtNum x[3], const mjtNum* coeff, int order,
+                             const int* nodeindices);
+
+// gather cell-local quantities and optionally compute rotation
+MJAPI void mju_flexGatherCellState(int order, int cy, int cz, int ci, int cj, int ck,
+                                   const mjtNum* xpos_g, const mjtNum* vel_g,
+                                   const mjtNum* xpos0_g, mjtNum* xpos_c, mjtNum* vel_c,
+                                   mjtNum* xpos0_c, int* nodeindices, mjtNum* quat);
+
 
 // ----------------------------- Base64 ------------------------------------------------------------
 

src/engine/engine_vis_interact.c

@@ -863,24 +863,30 @@ int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
         flexdist = newdist;
         if (m->flex_interp[i]) {
           mjtNum* coord = m->flex_vert0 + 3*(m->flex_vertadr[i] + vertid);
+          int order = m->flex_interp[i];
+          int npc = (order+1)*(order+1)*(order+1);
+
+          // cell lookup: get local coords and node indices
+          mjtNum loc[3];
+          int nodeindices[27];  // max npc for quadratic: 3^3 = 27
+          mju_cellLookup(coord, m->flex_cellnum+3*i, order, loc, nodeindices);
+
+          // find node with largest weight in this cell
           int nodeid = -1;
           int nstart = m->flex_nodeadr[i];
-          int nend = nstart + m->flex_nodenum[i];
           mjtNum w = 0;
-          for (int j = nstart; j < nend; j++) {
-            if (mju_evalBasis(coord, j-nstart, m->flex_interp[i]) > w) {
-              w = mju_evalBasis(coord, j-nstart, m->flex_interp[i]);
-              nodeid = j;
+          for (int j = 0; j < npc; j++) {
+            mjtNum ww = mju_evalBasis(loc, j, order);
+            if (ww > w) {
+              w = ww;
+              nodeid = nodeindices[j];
             }
           }
-          if (nodeid < 0) {
-            mjERROR("flex %d: node closest to vertex %d not found", i, vertid);
-          }
-          flexbodyid = m->flex_nodebodyid[m->flex_nodeadr[i] + nodeid];
+          flexbodyid = m->flex_nodebodyid[nstart + nodeid];
           if (m->flex_centered[i]) {
             mju_copy3(flexpnt, d->xpos + 3*flexbodyid);
           } else {
-            mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*nodeid);
+            mju_mulMatVec3(flexpnt, d->xmat + 9*flexbodyid, m->flex_node + 3*(nstart + nodeid));
             mju_addTo3(flexpnt, d->xpos + 3*flexbodyid);
           }
         } else {

src/engine/engine_vis_visualize.c

@@ -1434,10 +1434,9 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
     }
 
     // control points box
-    mjtNum xpos[mjMAXFLEXNODES];
+    mjtNum* xpos = mjSTACKALLOC(d, 3*m->flex_nodenum[f], mjtNum);
     int nstart = m->flex_nodeadr[f];
     int* bodyid = m->flex_nodebodyid + m->flex_nodeadr[f];
-    int nnode = m->flex_interp[f]+1;
     if (m->flex_centered[f]) {
       for (int i=0; i < m->flex_nodenum[f]; i++) {
         mju_copy3(xpos + 3*i, d->xpos + 3*bodyid[i]);
@@ -1448,15 +1447,23 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
         mju_addTo3(xpos + 3*i, d->xpos + 3*bodyid[i]);
       }
     }
-    for (int i=0; i < nnode; i++) {
-      for (int j=0; j < nnode; j++) {
-        for (int k=0; k < nnode; k++) {
-          int nn = nnode*nnode;
-          int offset  = 3*(nn*(i+0) + nnode*(j+0) + k);
-          int offset1 = 3*(nn*(i+1) + nnode*(j+0) + k);
-          int offset2 = 3*(nn*(i+0) + nnode*(j+1) + k);
-          int offset3 = 3*(nn*(i+0) + nnode*(j+0) + (k+1));
-          if (i < nnode-1) {
+
+    int cx = m->flex_cellnum[3*f+0];
+    int cy = m->flex_cellnum[3*f+1];
+    int cz = m->flex_cellnum[3*f+2];
+    int order = m->flex_interp[f];
+    int NX = cx * order + 1;
+    int NY = cy * order + 1;
+    int NZ = cz * order + 1;
+
+    for (int i=0; i < NX; i++) {
+      for (int j=0; j < NY; j++) {
+        for (int k=0; k < NZ; k++) {
+          int offset  = 3*(i*NY*NZ + j*NZ + k);
+          int offset1 = 3*((i+1)*NY*NZ + j*NZ + k);
+          int offset2 = 3*(i*NY*NZ + (j+1)*NZ + k);
+          int offset3 = 3*(i*NY*NZ + j*NZ + (k+1));
+          if (i < NX-1) {
             mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
             if (!thisgeom) {
               return;
@@ -1465,7 +1472,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
             mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset1);
             releaseGeom(&thisgeom, scn);
           }
-          if (j < nnode-1) {
+          if (j < NY-1) {
             mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
             if (!thisgeom) {
               return;
@@ -1474,7 +1481,7 @@ static void addFlexBvhGeoms(const mjModel* m, mjData* d, const mjvOption* vopt,
             mjv_connector(thisgeom, mjGEOM_LINE, 3, xpos+offset, xpos+offset2);
             releaseGeom(&thisgeom, scn);
           }
-          if (k < nnode-1) {
+          if (k < NZ-1) {
             mjvGeom* thisgeom = acquireGeom(scn, i, mjCAT_DECOR, mjOBJ_UNKNOWN);
             if (!thisgeom) {
               return;

src/experimental/filament/assets/unlit_segmentation.mat

@@ -17,13 +17,13 @@ material {
     shadingModel : unlit,
     culling: none,
     parameters : [
-        { type : float4, name : BaseColorFactor }
+        { type : float4, name : SegmentationColor }
     ]
 }
 
 fragment {
     void material(inout MaterialInputs material) {
       prepareMaterial(material);
-      material.baseColor = materialParams.BaseColorFactor;
+      material.baseColor = materialParams.SegmentationColor;
     }
 }

src/experimental/filament/filament/imgui_bridge.cc

@@ -23,7 +23,6 @@
 #include <imgui.h>
 #include <math/vec4.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/renderable.h"
@@ -242,11 +241,10 @@ void ImguiBridge::Update() {
         renderable->UpdateMesh(0, mesh, index_offset, command.ElemCount);
       }
 
-      Material::Textures textures;
+      MaterialTextures textures;
       textures.color = textures_[command.GetTexID()].get();
-      renderable->GetMaterial().UpdateTextures(textures);
 
-      Material::Params properties;
+      MaterialParams properties;
       properties.scissor[0] = command.ClipRect.x;
       properties.scissor[1] = height - command.ClipRect.w;
       properties.scissor[2] = command.ClipRect.z - command.ClipRect.x;
@@ -260,7 +258,7 @@ void ImguiBridge::Update() {
         properties.scissor[2] = width;
         properties.scissor[3] = height;
       }
-      renderable->GetMaterial().UpdateParams(properties);
+      renderable->UpdateMaterial(properties, textures);
 
       index_offset += command.ElemCount;
       ++renderable_index;
@@ -271,15 +269,10 @@ void ImguiBridge::Update() {
 void ImguiBridge::PrepareRenderables(int count) {
   while (renderables_.size() < count) {
     auto& r = renderables_.emplace_back(
-        std::make_unique<Renderable>(object_mgr_));
+        std::make_unique<Renderable>(Renderable::Usage::Ux, object_mgr_));
     r->SetCastShadows(false);
     r->SetReceiveShadows(false);
     r->SetBlendOrder(static_cast<std::uint16_t>(renderables_.size()));
-
-    Material& material = r->GetMaterial();
-    DrawMode mode = DrawMode::Color;
-    material.SetMaterial(mode, object_mgr_->GetMaterial(ObjectManager::kUnlitUi));
-    r->SetMaterialInstance(material.GetMaterialInstance(mode));
     scene_view_->AddToUxScene(r.get());
   }
   while (renderables_.size() > count) {

src/experimental/filament/filament/material.cc

@@ -20,104 +20,54 @@
 #include <filament/RenderableManager.h>
 #include <filament/TextureSampler.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/filament/object_manager.h"
 
 namespace mujoco {
 
-Material::Material(ObjectManager* object_mgr)
-    : object_mgr_(object_mgr) {
-}
-
-Material::~Material() noexcept {
-  for (int i = 0; i < kNumDrawModes; ++i) {
-    if (instances_[i]) {
-      GetEngine()->destroy(instances_[i]);
-    }
-  }
-}
-
-void Material::SetMaterial(DrawMode mode, filament::Material* material) {
-  const int index = static_cast<int>(mode);
-  if (instances_[index]) {
-    const filament::Material* current_material =
-        instances_[index]->getMaterial();
-    if (current_material == material) {
-      return;
-    }
-
-    GetEngine()->destroy(instances_[index]);
-    instances_[index] = nullptr;
-  }
-  if (material) {
-    instances_[index] = material->createInstance();
-    UpdateMaterialInstances();
-  }
-}
-
-filament::MaterialInstance* Material::GetMaterialInstance(DrawMode mode) {
-  return instances_[static_cast<int>(mode)];
-}
-
-void Material::UpdateParams(const Params& params) {
-  params_ = params;
-  UpdateMaterialInstances();
-}
-
-void Material::UpdateTextures(const Textures& textures) {
-  textures_ = textures;
-  UpdateMaterialInstances();
-}
-
-void Material::UpdateMaterialInstances() {
-  filament::MaterialInstance* instance =
-      instances_[static_cast<int>(DrawMode::Color)];
-  if (instance == nullptr) {
-    return;
-  }
-
-  if (params_.scissor[2] != 0 && params_.scissor[3] != 0) {
-    instance->setScissor(params_.scissor[0], params_.scissor[1],
-                         params_.scissor[2], params_.scissor[3]);
+void UpdateMaterialInstance(filament::MaterialInstance* instance,
+                            const MaterialParams& params,
+                            const MaterialTextures& textures,
+                            ObjectManager* object_mgr) {
+  if (params.scissor[2] != 0 && params.scissor[3] != 0) {
+    instance->setScissor(params.scissor[0], params.scissor[1],
+                         params.scissor[2], params.scissor[3]);
   }
 
   const filament::Material* material = instance->getMaterial();
   if (material->hasParameter("BaseColorFactor")) {
     instance->setParameter("BaseColorFactor", filament::RgbaType::sRGB,
-                           params_.color);
+                           params.color);
+  }
+  if (material->hasParameter("SegmentationColor")) {
+    instance->setParameter("SegmentationColor", filament::RgbaType::LINEAR,
+                           params.segmentation_color);
   }
   if (material->hasParameter("EmissiveFactor")) {
-    instance->setParameter("EmissiveFactor", params_.emissive);
+    instance->setParameter("EmissiveFactor", params.emissive);
   }
   if (material->hasParameter("SpecularFactor")) {
-    instance->setParameter("SpecularFactor", params_.specular);
+    instance->setParameter("SpecularFactor", params.specular);
   }
   if (material->hasParameter("GlossinessFactor")) {
-    instance->setParameter("GlossinessFactor", params_.glossiness);
+    instance->setParameter("GlossinessFactor", params.glossiness);
   }
   if (material->hasParameter("MetallicFactor")) {
     instance->setParameter("MetallicFactor",
-                           params_.metallic >= 0 ? params_.metallic : 1.0f);
+                           params.metallic >= 0 ? params.metallic : 1.0f);
   }
   if (material->hasParameter("RoughnessFactor")) {
     instance->setParameter("RoughnessFactor",
-                           params_.roughness >= 0 ? params_.roughness : 1.0f);
+                           params.roughness >= 0 ? params.roughness : 1.0f);
   }
   if (material->hasParameter("UvScale")) {
-    instance->setParameter("UvScale", params_.uv_scale);
+    instance->setParameter("UvScale", params.uv_scale);
   }
   if (material->hasParameter("UvOffset")) {
-    instance->setParameter("UvOffset", params_.uv_offset);
+    instance->setParameter("UvOffset", params.uv_offset);
   }
   if (material->hasParameter("Reflectance")) {
-    instance->setParameter("Reflectance", params_.reflectance);
-  }
-
-  const int segmentation_index = static_cast<int>(DrawMode::Segmentation);
-  if (instances_[segmentation_index]) {
-    instances_[segmentation_index]->setParameter("BaseColorFactor",
-                                                 params_.segmentation_color);
+    instance->setParameter("Reflectance", params.reflectance);
   }
 
   // All textures use the same default sampler.
@@ -135,19 +85,20 @@ void Material::UpdateMaterialInstances() {
       if (texture != nullptr) {
         instance->setParameter(name, texture->GetFilamentTexture(), sampler);
       } else {
-        instance->setParameter(name, object_mgr_->GetFallbackTexture(role), sampler);
+        instance->setParameter(name, object_mgr->GetFallbackTexture(role),
+                               sampler);
       }
     }
   };
 
-  TrySetTexture("BaseColor", textures_.color, mjTEXROLE_RGB);
-  TrySetTexture("Normal", textures_.normal, mjTEXROLE_NORMAL);
-  TrySetTexture("Metallic", textures_.metallic, mjTEXROLE_METALLIC);
-  TrySetTexture("Roughness", textures_.roughness, mjTEXROLE_ROUGHNESS);
-  TrySetTexture("Occlusion", textures_.occlusion, mjTEXROLE_OCCLUSION);
-  TrySetTexture("ORM", textures_.orm, mjTEXROLE_ORM);
-  TrySetTexture("Emissive", textures_.emissive, mjTEXROLE_EMISSIVE);
-  TrySetTexture("Reflection", textures_.reflection, mjTEXROLE_USER);
+  TrySetTexture("BaseColor", textures.color, mjTEXROLE_RGB);
+  TrySetTexture("Normal", textures.normal, mjTEXROLE_NORMAL);
+  TrySetTexture("Metallic", textures.metallic, mjTEXROLE_METALLIC);
+  TrySetTexture("Roughness", textures.roughness, mjTEXROLE_ROUGHNESS);
+  TrySetTexture("Occlusion", textures.occlusion, mjTEXROLE_OCCLUSION);
+  TrySetTexture("ORM", textures.orm, mjTEXROLE_ORM);
+  TrySetTexture("Emissive", textures.emissive, mjTEXROLE_EMISSIVE);
+  TrySetTexture("Reflection", textures.reflection, mjTEXROLE_USER);
 }
 
 }  // namespace mujoco

src/experimental/filament/filament/material.h

@@ -20,82 +20,48 @@
 #include <math/vec2.h>
 #include <math/vec3.h>
 #include <math/vec4.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/texture.h"
 #include "experimental/filament/filament/object_manager.h"
 
 namespace mujoco {
 
-class Material {
- public:
-  // The textures that can be assigned to the drawable's material.
-  struct Textures {
-    const Texture* color = nullptr;
-    const Texture* normal = nullptr;
-    const Texture* metallic = nullptr;
-    const Texture* roughness = nullptr;
-    const Texture* occlusion = nullptr;
-    const Texture* orm = nullptr;
-    const Texture* emissive = nullptr;
-    const Texture* reflection = nullptr;
-  };
-
-  // The parameters that can be applied to the drawable's material.
-  struct Params {
-    filament::math::float4 color = {1, 1, 1, 1};
-    filament::math::float4 segmentation_color = {1, 1, 1, 1};
-    filament::math::float2 tex_repeat = {1, 1};
-    filament::math::float3 uv_scale = {1, 1, 1};
-    filament::math::float3 uv_offset = {0, 0, 0};
-    filament::math::float4 scissor = {0, 0, 0, 0};
-    float specular = -1.0f;
-    float glossiness = -1.0f;
-    float metallic = -1.0f;
-    float roughness = -1.0f;
-    float emissive = -1.0f;
-    float reflectance = 0.0f;
-    bool tex_uniform = false;
-    bool reflective = false;
-  };
-
-  explicit Material(ObjectManager* object_mgr);
-  ~Material() noexcept;
-
-  Material(const Material&) = delete;
-  Material& operator=(const Material&) = delete;
-
-  // Assigns a material to the draw mode.
-  void SetMaterial(DrawMode mode, filament::Material* material);
-
-  // Updates the parameters for the material.
-  void UpdateParams(const Params& params);
-
-  // Updates the textures for the material.
-  void UpdateTextures(const Textures& textures);
-
-  // Returns the current material parameters.
-  const Params& GetParams() const { return params_; }
-
-  // Returns the current material textures.
-  const Textures& GetTextures() const { return textures_; }
-
-  // Returns the material instance assigned to the draw mode.
-  filament::MaterialInstance* GetMaterialInstance(DrawMode mode);
-
-  // Returns the filament Engine managing the material.
-  filament::Engine* GetEngine() const { return object_mgr_->GetEngine(); }
-
- private:
-  // Updates the material instances based on the currently set parameters and
-  // textures.
-  void UpdateMaterialInstances();
-
-  ObjectManager* object_mgr_;
-  filament::MaterialInstance* instances_[kNumDrawModes] = {nullptr};
-  Params params_;
-  Textures textures_;
+// The textures that can be assigned to the drawable's material.
+struct MaterialTextures {
+  const Texture* color = nullptr;
+  const Texture* normal = nullptr;
+  const Texture* metallic = nullptr;
+  const Texture* roughness = nullptr;
+  const Texture* occlusion = nullptr;
+  const Texture* orm = nullptr;
+  const Texture* emissive = nullptr;
+  const Texture* reflection = nullptr;
 };
 
+// The parameters that can be applied to the drawable's material.
+struct MaterialParams {
+  filament::math::float4 color = {1, 1, 1, 1};
+  filament::math::float4 segmentation_color = {1, 1, 1, 1};
+  filament::math::float2 tex_repeat = {1, 1};
+  filament::math::float3 uv_scale = {1, 1, 1};
+  filament::math::float3 uv_offset = {0, 0, 0};
+  filament::math::float4 scissor = {0, 0, 0, 0};
+  float specular = -1.0f;
+  float glossiness = -1.0f;
+  float metallic = -1.0f;
+  float roughness = -1.0f;
+  float emissive = -1.0f;
+  float reflectance = 0.0f;
+  bool tex_uniform = false;
+  bool reflective = false;
+};
+
+// Updates the material instances based on the currently set parameters and
+// textures.
+void UpdateMaterialInstance(filament::MaterialInstance* instance,
+                            const MaterialParams& params,
+                            const MaterialTextures& textures,
+                            ObjectManager* object_mgr);
+
 }  // namespace mujoco
 
 #endif  // MUJOCO_SRC_EXPERIMENTAL_FILAMENT_FILAMENT_MATERIAL_H_

src/experimental/filament/filament/renderable.cc

@@ -14,26 +14,36 @@
 
 #include "experimental/filament/filament/renderable.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <utility>
 
 #include <filament/Engine.h>
+#include <filament/Material.h>
 #include <filament/RenderableManager.h>
 #include <filament/Scene.h>
 #include <utils/EntityManager.h>
 #include <mujoco/mujoco.h>
+#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
 
 namespace mujoco {
 
-Renderable::Renderable(ObjectManager* object_mgr) : material_(object_mgr) {}
+Renderable::Renderable(Usage usage, ObjectManager* object_mgr)
+    : usage_(usage), object_mgr_(object_mgr) {}
 
 Renderable::~Renderable() noexcept {
   while (!entities_.empty()) {
     RemoveLastEntity();
   }
+  for (int i = 0; i < kNumDrawModes; ++i) {
+    if (instances_[i] != nullptr) {
+      GetEngine()->destroy(instances_[i]);
+      instances_[i] = nullptr;
+    }
+  }
 }
 
 void Renderable::RemoveLastEntity() {
@@ -103,8 +113,8 @@ void Renderable::AppendEntity(const MeshInfo& mesh_info) {
   } else {
     builder.culling(false);
   }
-  if (material_instance_) {
-    builder.material(0, material_instance_);
+  if (instances_[static_cast<int>(draw_mode_)] != nullptr) {
+    builder.material(0, instances_[static_cast<int>(draw_mode_)]);
   }
   builder.castShadows(cast_shadows_);
   builder.receiveShadows(receive_shadows_);
@@ -191,15 +201,67 @@ void Renderable::RemoveFromScene(filament::Scene* scene) {
   assigned_scene_ = nullptr;
 }
 
-void Renderable::SetMaterialInstance(filament::MaterialInstance* instance) {
-  if (instance != material_instance_) {
+void Renderable::UpdateMaterial(const MaterialParams& params,
+                                const MaterialTextures& textures) {
+  params_ = params;
+  textures_ = textures;
+
+  AssignMaterial(DrawMode::Color, GetColorMaterialType());
+  if (usage_ == Usage::SceneObject) {
+    AssignMaterial(DrawMode::Depth, ObjectManager::kUnlitDepth);
+    AssignMaterial(DrawMode::Segmentation, ObjectManager::kUnlitSegmentation);
+  }
+
+  for (int i = 0; i < kNumDrawModes; ++i) {
+    if (instances_[i]) {
+      UpdateMaterialInstance(instances_[i], params_, textures_, object_mgr_);
+    }
+  }
+  SetDrawMode(draw_mode_);
+}
+
+void Renderable::AssignMaterial(DrawMode 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();
+  }
+}
+
+const MaterialParams& Renderable::GetMaterialParams() const {
+  return params_;
+}
+
+const MaterialTextures& Renderable::GetMaterialTextures() const {
+  return textures_;
+}
+
+void Renderable::SetDrawMode(DrawMode mode) {
+  // Only SceneObjects support non-color draw modes.
+  if (usage_ != Usage::SceneObject) {
+    mode = DrawMode::Color;
+  }
+
+  filament::MaterialInstance* instance = instances_[static_cast<int>(mode)];
+  if (instance) {
     filament::RenderableManager& rm = GetEngine()->getRenderableManager();
     for (utils::Entity& entity : entities_) {
       filament::RenderableManager::Instance ri = rm.getInstance(entity);
       rm.setMaterialInstanceAt(ri, 0, instance);
     }
-    material_instance_ = instance;
   }
+  draw_mode_ = mode;
 }
 
 std::uint8_t Renderable::SetLayerMask(std::uint8_t mask) {
@@ -284,8 +346,74 @@ void Renderable::SetWireframe(bool wireframe) {
   }
 }
 
-Material& Renderable::GetMaterial() { return material_; }
 
-filament::Engine* Renderable::GetEngine() { return material_.GetEngine(); }
+ObjectManager::MaterialType Renderable::GetColorMaterialType() const {
+  if (usage_ == Usage::DecorLines) {
+    return ObjectManager::kUnlitLine;
+  } else if (usage_ == Usage::Decor) {
+    return ObjectManager::kUnlitSegmentation;
+  } else if (usage_ == Usage::Ux) {
+    return ObjectManager::kUnlitUi;
+  } else if (textures_.orm) {
+    return ObjectManager::kPbrPacked;
+  } else if (textures_.metallic) {
+    return ObjectManager::kPbr;
+  } else if (textures_.roughness) {
+    return ObjectManager::kPbr;
+  } else if (params_.metallic >= 0) {
+    return ObjectManager::kPbr;
+  } else if (params_.roughness >= 0) {
+    return ObjectManager::kPbr;
+  }
+
+  // Check to see if we're dealing with a mesh with texture coordinates.
+  // `data_id` is the id of the mesh in model (i.e. the geom has mesh
+  // geometry) and `mesh_texcoordadr` stores the address of the mesh uvs if
+  // it has them.
+  bool has_texcoords = false;
+  if (!meshes_.empty()) {
+    const auto attribs = meshes_[0].mesh->GetVertexAttributes();
+    auto it = std::find(attribs.begin(), attribs.end(),
+                        filament::VertexAttribute::UV0);
+    has_texcoords = (it != attribs.end());
+  }
+
+  if (textures_.color == nullptr) {
+    if (params_.color.a < 1.0f) {
+      return ObjectManager::kPhongColorFade;
+    } else if (params_.reflective) {
+      return ObjectManager::kPhongColorReflect;
+    } else {
+      return ObjectManager::kPhongColor;
+    }
+  } else if (textures_.color->GetFilamentTexture()->getTarget() ==
+              filament::Texture::Sampler::SAMPLER_CUBEMAP) {
+    if (params_.color.a < 1.0f) {
+      return ObjectManager::kPhongCubeFade;
+    } else if (params_.reflective) {
+      return ObjectManager::kPhongCubeReflect;
+    } else {
+      return ObjectManager::kPhongCube;
+    }
+  } else if (has_texcoords) {
+    if (params_.color.a < 1.0f) {
+      return ObjectManager::kPhong2dUvFade;
+    } else if (params_.reflective) {
+      return ObjectManager::kPhong2dUvReflect;
+    } else {
+      return ObjectManager::kPhong2dUv;
+    }
+  } else {
+    if (params_.color.a < 1.0f) {
+      return ObjectManager::kPhong2dFade;
+    } else if (params_.reflective) {
+      return ObjectManager::kPhong2dReflect;
+    } else {
+      return ObjectManager::kPhong2d;
+    }
+  }
+}
+
+filament::Engine* Renderable::GetEngine() { return object_mgr_->GetEngine(); }
 
 }  // namespace mujoco

src/experimental/filament/filament/renderable.h

@@ -21,6 +21,7 @@
 #include <filament/Engine.h>
 #include <filament/Scene.h>
 #include <utils/Entity.h>
+#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/mesh.h"
 #include "experimental/filament/filament/object_manager.h"
@@ -37,11 +38,19 @@ namespace mujoco {
 // assigns the same material instance to all of them.
 class Renderable {
  public:
+  // How the material is to be used for rendering.
+  enum class Usage {
+    SceneObject,
+    Decor,
+    DecorLines,
+    Ux,
+  };
+
   // Default filament values for priority and layer mask.
   static constexpr std::uint8_t kDefaultPriority = 4;
   static constexpr std::uint8_t kDefaultLayerMask = 0x01;
 
-  explicit Renderable(ObjectManager* object_mgr);
+  Renderable(Usage usage, ObjectManager* object_mgr);
   ~Renderable() noexcept;
 
   Renderable(const Renderable&) = delete;
@@ -95,10 +104,17 @@ class Renderable {
   void RemoveFromScene(filament::Scene* scene);
 
   // Sets the material instance for all managed entities.
-  void SetMaterialInstance(filament::MaterialInstance* material_instance);
+  void SetDrawMode(DrawMode mode);
 
-  // Returns the material for the renderables.
-  Material& GetMaterial();
+  // Updates the parameters for the material.
+  void UpdateMaterial(const MaterialParams& params,
+                      const MaterialTextures& textures);
+
+  // Returns the current material parameters.
+  const MaterialParams& GetMaterialParams() const;
+
+  // Returns the current material textures.
+  const MaterialTextures& GetMaterialTextures() const;
 
   // Returns the filament Engine managing the renderables.
   filament::Engine* GetEngine();
@@ -129,9 +145,17 @@ class Renderable {
   // Removes the last filament::Entity from the renderable.
   void RemoveLastEntity();
 
-  Material material_;
+  void AssignMaterial(DrawMode mode, ObjectManager::MaterialType material_type);
+
+  ObjectManager::MaterialType GetColorMaterialType() const;
+
+  Usage usage_;
+  ObjectManager* object_mgr_;
+  filament::MaterialInstance* instances_[kNumDrawModes] = {nullptr};
+  MaterialParams params_;
+  MaterialTextures textures_;
+  DrawMode draw_mode_ = DrawMode::Color;
   filament::Scene* assigned_scene_ = nullptr;
-  filament::MaterialInstance* material_instance_ = nullptr;
   std::vector<utils::Entity> entities_;
   std::vector<MeshInfo> meshes_;
   std::uint8_t priority_ = kDefaultPriority;

src/experimental/filament/filament/scene_geom_util.cc

@@ -31,7 +31,6 @@
 #include <utils/Entity.h>
 #include <mujoco/mjvisualize.h>
 #include <mujoco/mujoco.h>
-#include "experimental/filament/filament/draw_mode.h"
 #include "experimental/filament/filament/material.h"
 #include "experimental/filament/filament/math_util.h"
 #include "experimental/filament/filament/mesh.h"
@@ -323,12 +322,10 @@ static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,
                                ObjectManager* object_mgr,
                                const float headpos[3]) {
   const mjModel* model = model_objs->GetModel();
-  Material& material = renderable.GetMaterial();
 
   const bool use_segid_color = scene->flags[mjRND_IDCOLOR];
   const bool enable_reflection = scene->flags[mjRND_REFLECTION];
-
-  Material::Params params;
+  MaterialParams params;
   params.color = ReadFloat4(geom.rgba);
   if (geom.type == mjGEOM_PLANE) {
     if (IsBehind(headpos, geom.pos, geom.mat)) {
@@ -347,7 +344,7 @@ static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,
     renderable.SetReceiveShadows(false);
   }
 
-  Material::Textures textures;
+  MaterialTextures textures;
   if (geom.matid >= 0) {
     textures.color = model_objs->GetTexture(geom.matid, mjTEXROLE_RGB);
     textures.normal = model_objs->GetTexture(geom.matid, mjTEXROLE_NORMAL);
@@ -358,79 +355,6 @@ static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,
         model_objs->GetTexture(geom.matid, mjTEXROLE_ROUGHNESS);
     textures.occlusion =
         model_objs->GetTexture(geom.matid, mjTEXROLE_OCCLUSION);
-    material.UpdateTextures(textures);
-  }
-
-  ObjectManager::MaterialType material_type = ObjectManager::kNumMaterials;
-  if (geom.type == mjGEOM_LINE || geom.type == mjGEOM_LINEBOX) {
-    material_type = ObjectManager::kUnlitLine;
-  } else if (geom.category == mjCAT_DECOR) {
-    material_type = ObjectManager::kUnlitSegmentation;
-  } else {
-    bool material_assigned = false;
-    if (geom.matid >= 0) {
-      material_assigned = true;
-      if (textures.orm) {
-        material_type = ObjectManager::kPbrPacked;
-      } else if (textures.metallic) {
-        material_type = ObjectManager::kPbr;
-      } else if (textures.roughness) {
-        material_type = ObjectManager::kPbr;
-      } else if (model->mat_metallic[geom.matid] >= 0) {
-        material_type = ObjectManager::kPbr;
-      } else if (model->mat_roughness[geom.matid] >= 0) {
-        material_type = ObjectManager::kPbr;
-      } else {
-        material_assigned = false;
-      }
-    }
-
-    if (!material_assigned) {
-      // Check to see if we're dealing with a mesh with texture coordinates.
-      // `data_id` is the id of the mesh in model (i.e. the geom has mesh
-      // geometry) and `mesh_texcoordadr` stores the address of the mesh uvs if
-      // it has them.
-      bool has_texcoords = false;
-      if ((geom.type == mjGEOM_MESH || geom.type == mjGEOM_SDF) &&
-          geom.dataid >= 0 && model->mesh_texcoordadr[geom.dataid / 2] >= 0) {
-        has_texcoords = true;
-      }
-
-      if (textures.color == nullptr) {
-        if (params.color.a < 1.0f) {
-          material_type = ObjectManager::kPhongColorFade;
-        } else if (params.reflective) {
-          material_type = ObjectManager::kPhongColorReflect;
-        } else {
-          material_type = ObjectManager::kPhongColor;
-        }
-      } else if (textures.color->GetFilamentTexture()->getTarget() ==
-                 filament::Texture::Sampler::SAMPLER_CUBEMAP) {
-        if (params.color.a < 1.0f) {
-          material_type = ObjectManager::kPhongCubeFade;
-        } else if (params.reflective) {
-          material_type = ObjectManager::kPhongCubeReflect;
-        } else {
-          material_type = ObjectManager::kPhongCube;
-        }
-      } else if (has_texcoords) {
-        if (params.color.a < 1.0f) {
-          material_type = ObjectManager::kPhong2dUvFade;
-        } else if (params.reflective) {
-          material_type = ObjectManager::kPhong2dUvReflect;
-        } else {
-          material_type = ObjectManager::kPhong2dUv;
-        }
-      } else {
-        if (params.color.a < 1.0f) {
-          material_type = ObjectManager::kPhong2dFade;
-        } else if (params.reflective) {
-          material_type = ObjectManager::kPhong2dReflect;
-        } else {
-          material_type = ObjectManager::kPhong2d;
-        }
-      }
-    }
   }
 
   params.reflectance = geom.reflectance;
@@ -532,20 +456,21 @@ static void UpdateGeomMaterial(Renderable& renderable, const mjvGeom& geom,
   params.emissive *= model_objs->GetEmissiveMultiplier();
   params.specular *= model_objs->GetSpecularMultiplier();
   params.glossiness *= model_objs->GetShininessMultiplier();
-  material.UpdateParams(params);
 
-  material.SetMaterial(DrawMode::Color, object_mgr->GetMaterial(material_type));
-  material.SetMaterial(DrawMode::Depth,
-                       object_mgr->GetMaterial(ObjectManager::kUnlitDepth));
-  material.SetMaterial(
-      DrawMode::Segmentation,
-      object_mgr->GetMaterial(ObjectManager::kUnlitSegmentation));
+  renderable.UpdateMaterial(params, textures);
 }
 
 std::unique_ptr<Renderable> CreateGeomRenderable(
     const mjvGeom& geom, const mjvScene* scene, ObjectManager* object_mgr,
     ModelObjects* model_objs, const float headpos[3]) {
-  auto renderable = std::make_unique<Renderable>(object_mgr);
+  Renderable::Usage usage = Renderable::Usage::SceneObject;
+  if (geom.type == mjGEOM_LINE || geom.type == mjGEOM_LINEBOX) {
+    usage = Renderable::Usage::DecorLines;
+  } else if (geom.category == mjCAT_DECOR) {
+    usage = Renderable::Usage::Decor;
+  }
+
+  auto renderable = std::make_unique<Renderable>(usage, object_mgr);
 
   // The order of these calls is important. e.g. We need to create the filament
   // renderable entities before we can set their transform.

src/experimental/filament/filament/scene_view.cc

@@ -208,7 +208,7 @@ void SceneView::RemoveFromScene(Light* light) {
 void SceneView::AddToScene(Renderable* renderable) {
   if (renderables_.insert(renderable).second) {
     renderable->AddToScene(scene_);
-    if (renderable->GetMaterial().GetParams().reflective) {
+    if (renderable->GetMaterialParams().reflective) {
       AddReflectiveRenderable(renderable);
     }
   }
@@ -261,8 +261,7 @@ void SceneView::Render(filament::Renderer* renderer,
   SetupCamera(request.camera, viewport, camera_);
 
   for (auto& iter : renderables_) {
-    Material& material = iter->GetMaterial();
-    iter->SetMaterialInstance(material.GetMaterialInstance(request.draw_mode));
+    iter->SetDrawMode(request.draw_mode);
   }
 
   filament::View* view = views_[static_cast<int>(request.draw_mode)];
@@ -336,10 +335,9 @@ void SceneView::AddReflectiveRenderable(Renderable* renderable) {
   auto& target = reflect_targets_[index];
   target->Prepare(viewport.width, viewport.height);
 
-  Material& material = renderable->GetMaterial();
-  Material::Textures textures = material.GetTextures();
+  MaterialTextures textures = renderable->GetMaterialTextures();
   textures.reflection = target->GetColorTexture();
-  material.UpdateTextures(textures);
+  renderable->UpdateMaterial(renderable->GetMaterialParams(), textures);
 }
 
 void SceneView::SetColorGradingOptions(const ColorGradingOptions& opts) {

src/experimental/platform/ux/picture_gui.cc

@@ -20,16 +20,14 @@
 #include <imgui.h>
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/renderer.h"
-#include "experimental/platform/hal/window.h"
 #include "experimental/platform/ux/imgui_widgets.h"
 
 namespace mujoco::platform {
 
 // Returns false if the user requests that this picture-in-picture widget be
 // removed from the GUI.
-static bool PipGuiImpl(const mjModel* model, mjData* data,
-                       platform::Window* window, platform::Renderer* renderer,
-                       PipState* pip) {
+static bool PipGuiImpl(const mjModel* model, mjData* data, float aspect_ratio,
+                       platform::Renderer* renderer, PipState* pip) {
   bool result = true;
 
   auto get_camera_name = [model](int i) -> const char* {
@@ -41,7 +39,7 @@ static bool PipGuiImpl(const mjModel* model, mjData* data,
   };
 
   const int width = ImGui::GetContentRegionAvail().x;
-  const int height = width / window->GetAspectRatio();
+  const int height = aspect_ratio != 0.f ? width / aspect_ratio : width;
   std::vector<std::byte> output(width * height * 3);
 
   const int combo_width = (width - 30) / 2;
@@ -97,7 +95,7 @@ static bool PipGuiImpl(const mjModel* model, mjData* data,
   return result;
 }
 
-void PipGui(const mjModel* model, mjData* data, platform::Window* window,
+void PipGui(const mjModel* model, mjData* data, float aspect_ratio,
             platform::Renderer* renderer, std::vector<PipState>* pips) {
   if (pips->empty()) {
     pips->emplace_back();
@@ -106,7 +104,7 @@ void PipGui(const mjModel* model, mjData* data, platform::Window* window,
   std::vector<int> to_delete;
   for (int i = 0; i < pips->size(); ++i) {
     PipState& pip = pips->at(i);
-    if (PipGuiImpl(model, data, window, renderer, &pip) == false) {
+    if (PipGuiImpl(model, data, aspect_ratio, renderer, &pip) == false) {
       to_delete.push_back(i);
     };
     ImGui::Separator();

src/experimental/platform/ux/picture_gui.h

@@ -19,7 +19,6 @@
 
 #include <mujoco/mujoco.h>
 #include "experimental/platform/hal/renderer.h"
-#include "experimental/platform/hal/window.h"
 
 namespace mujoco::platform {
 
@@ -32,7 +31,7 @@ struct PipState {
 };
 
 // Renders the GUI for a set of picture-in-picture widgets.
-void PipGui(const mjModel* model, mjData* data, platform::Window* window,
+void PipGui(const mjModel* model, mjData* data, float aspect_ratio,
             platform::Renderer* renderer, std::vector<PipState>* pips);
 
 }  // namespace mujoco::platform

src/experimental/studio/app.cc

@@ -902,7 +902,8 @@ void App::BuildGui() {
 
   if (tmp_.picture_in_picture) {
     if (ImGui::Begin("Picture-in-Picture", &tmp_.picture_in_picture)) {
-      PipGui(model(), data(), window_.get(), renderer_.get(), &tmp_.pips);
+      platform::PipGui(model(), data(), window_->GetAspectRatio(),
+                       renderer_.get(), &tmp_.pips);
     }
     ImGui::End();
   }

src/user/user_flexcomp.cc

@@ -77,6 +77,7 @@ bool IsValidElementOrNodeHeader22(const std::string& line) {
 mjCFlexcomp::mjCFlexcomp(void) {
   type = mjFCOMPTYPE_GRID;
   count[0] = count[1] = count[2] = 10;
+  cellcount[0] = cellcount[1] = cellcount[2] = -1;
   mjuu_setvec(spacing, 0.02, 0.02, 0.02);
   mjuu_setvec(scale, 1, 1, 1);
   mass = 1;
@@ -269,10 +270,19 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
 
   // construct pinned array
   int nnode = 0;
-  if (doftype == mjFCOMPDOF_TRILINEAR) {
-    nnode = 8;
-  } else if (doftype == mjFCOMPDOF_QUADRATIC) {
-    nnode = 27;
+  if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
+    int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
+    // multi-cell count for mesh/direct/gmsh, else single cell
+    int cx = 1, cy = 1, cz = 1;
+    if (type == mjFCOMPTYPE_MESH || type == mjFCOMPTYPE_DIRECT ||
+        type == mjFCOMPTYPE_GMSH) {
+      if (cellcount[0] >= 0) {
+        cx = cellcount[0];
+        cy = cellcount[1];
+        cz = cellcount[2];
+      }
+    }
+    nnode = (cx*order+1) * (cy*order+1) * (cz*order+1);
   }
   pinned = vector<bool>(std::max(npnt, nnode), rigid);
 
@@ -562,36 +572,81 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
     }
   }
 
-  // create nodal mesh for trilinear interpolation
+  // create nodal mesh for trilinear/quadratic interpolation
   if (doftype == mjFCOMPDOF_TRILINEAR || doftype == mjFCOMPDOF_QUADRATIC) {
-    int order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
-    flex->SetOrder(order);
-    std::vector<double> node(3*(order+1)*(order+1)*(order+1), 0);
+    flex->spec.order = doftype == mjFCOMPDOF_TRILINEAR ? 1 : 2;
+
+    if (cellcount[0] >= 0) {
+      flex->spec.cellcount[0] = cellcount[0];
+      flex->spec.cellcount[1] = cellcount[1];
+      flex->spec.cellcount[2] = cellcount[2];
+    }
+
+    // total number of nodes with shared boundaries
+    int nx = flex->spec.cellcount[0] * flex->spec.order + 1;
+    int ny = flex->spec.cellcount[1] * flex->spec.order + 1;
+    int nz = flex->spec.cellcount[2] * flex->spec.order + 1;
+    int nnode = nx * ny * nz;
+
+    std::vector<double> node(3 * nnode, 0);
     int idx = 0;
-    double step = 1.0 / (double)order;
+
+    // Simpson's rule weights for quadratic mass distribution
     double massP2[3] = {1. / 6., 2. / 3., 1. / 6.};
-    for (int i=0; i <= order; i++) {
-      for (int j=0; j <= order; j++) {
-        for (int k=0; k <= order; k++) {
+
+    // compute per-node mass for trilinear:
+    //   mass / nnode (uniform), or use Simpson for quadratic
+    double node_mass_uniform = mass / nnode;
+
+    for (int gi = 0; gi < nx; gi++) {
+      for (int gj = 0; gj < ny; gj++) {
+        for (int gk = 0; gk < nz; gk++) {
+          // parametric position in [0, 1]^3
+          double s = (double)gi / (flex->spec.cellcount[0] * flex->spec.order);
+          double t = (double)gj / (flex->spec.cellcount[1] * flex->spec.order);
+          double u = (double)gk / (flex->spec.cellcount[2] * flex->spec.order);
+
+          // physical position
+          double px = minmax[0] + s * (minmax[3] - minmax[0]);
+          double py = minmax[1] + t * (minmax[4] - minmax[1]);
+          double pz = minmax[2] + u * (minmax[5] - minmax[2]);
+
           if (pinned[idx]) {
-            node[3*idx+0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
-            node[3*idx+1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
-            node[3*idx+2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
-            mjs_appendString(pf->nodebody, mjs_getName(body->element)->c_str());
+            node[3*idx+0] = px;
+            node[3*idx+1] = py;
+            node[3*idx+2] = pz;
+            mjs_appendString(pf->nodebody,
+                             mjs_getName(body->element)->c_str());
             idx++;
             continue;
           }
 
           mjsBody* pb = mjs_addBody(body, 0);
-          pb->pos[0] = minmax[0] + i * step * (minmax[3] - minmax[0]);
-          pb->pos[1] = minmax[1] + j * step * (minmax[4] - minmax[1]);
-          pb->pos[2] = minmax[2] + k * step * (minmax[5] - minmax[2]);
+          pb->pos[0] = px;
+          pb->pos[1] = py;
+          pb->pos[2] = pz;
           mjuu_zerovec(pb->ipos, 3);
+
+          // mass distribution
           if (doftype == mjFCOMPDOF_TRILINEAR) {
-            pb->mass = mass / 8;
+            pb->mass = node_mass_uniform;
           } else {
-            pb->mass = mass * massP2[i] * massP2[j] * massP2[k];
+            // local index within the cell for mass computation
+            int li = gi % flex->spec.order;
+            int lj = gj % flex->spec.order;
+            int lk = gk % flex->spec.order;
+            // boundary nodes: average mass contribution
+            int ncells_i = (gi > 0 && gi < nx-1 && li == 0) ? 2 : 1;
+            int ncells_j = (gj > 0 && gj < ny-1 && lj == 0) ? 2 : 1;
+            int ncells_k = (gk > 0 && gk < nz-1 && lk == 0) ? 2 : 1;
+            // use Simpson weights scaled by cell count
+            double wi = massP2[li == 0 ? 0 : li];
+            double wj = massP2[lj == 0 ? 0 : lj];
+            double wk = massP2[lk == 0 ? 0 : lk];
+            pb->mass = mass * wi * wj * wk * ncells_i * ncells_j * ncells_k
+                       / (flex->spec.cellcount[0] * flex->spec.cellcount[1] * flex->spec.cellcount[2]);
           }
+
           pb->inertia[0] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
           pb->inertia[1] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
           pb->inertia[2] = pb->mass*(2.0*inertiabox*inertiabox)/3.0;
@@ -607,7 +662,7 @@ bool mjCFlexcomp::Make(mjsBody* body, char* error, int error_sz, const mjVFS* vf
 
           // construct node name, add to nodebody
           char txt[100];
-          mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), i, j, k);
+          mju::sprintf_arr(txt, "%s_%d_%d_%d", name.c_str(), gi, gj, gk);
           mjs_setName(pb->element, txt);
           mjs_appendString(pf->nodebody, mjs_getName(pb->element)->c_str());
 

src/user/user_flexcomp.h

@@ -78,6 +78,7 @@ class mjCFlexcomp {
   std::string name;               // flex name
   mjtFcompType type;              // flexcomp type
   int count[3];                   // grid count in each dimension
+  int cellcount[3];               // number of cells for interpolation
   double spacing[3];              // spacing between grid elements
   double scale[3];                // scaling for mesh and direct
   double origin[3];               // origin for generating a 3D mesh from a convex 2D mesh

src/user/user_init.c

@@ -224,6 +224,9 @@ void mjs_defaultFlex(mjsFlex* flex) {
   // set other defaults
   flex->dim = 2;
   flex->radius = 0.005;
+  flex->cellcount[0] = 1;
+  flex->cellcount[1] = 1;
+  flex->cellcount[2] = 1;
   flex->internal = 0;
   flex->selfcollide = mjFLEXSELF_AUTO;
   flex->activelayers = 1;

src/user/user_mesh.cc

@@ -3961,7 +3961,10 @@ std::string mjCFlex::ComputeStiffnessCacheKey() const {
 
   combine(std::hash<double>{}(young));
   combine(std::hash<double>{}(poisson));
-  combine(std::hash<int>{}(order_));
+  combine(std::hash<int>{}(spec.order));
+  combine(std::hash<int>{}(spec.cellcount[0]));
+  combine(std::hash<int>{}(spec.cellcount[1]));
+  combine(std::hash<int>{}(spec.cellcount[2]));
 
   // compute bounding box from vertex positions
   if (!vert_.empty()) {
@@ -4086,10 +4089,23 @@ void mjCFlex::Compile(const mjVFS* vfs) {
 
   // set nnode
   nnode = static_cast<int>(nodebody_.size());
-  if (nnode && !order_) {
-    order_ = std::pow(nnode, 1.0 / 3) - 1;
-    if (nnode != std::pow(order_ + 1, 3)) {
-      throw mjCError(this, "number of nodes must be %d^3 but it is %d", nullptr, order_, nnode);
+  if (nnode && !spec.order) {
+    throw mjCError(this, "Interpolation order must be explicitly specified (dof is missing)");
+  }
+
+  // check node compatibility with count and dof
+  if (spec.order > 0) {
+    if (spec.cellcount[0] == 0 || spec.cellcount[1] == 0 || spec.cellcount[2] == 0) {
+      throw mjCError(this, "cellcount cannot be 0 in any dimension when interpolation order > 0");
+    }
+    int expected_nodes = (spec.cellcount[0] * spec.order + 1) *
+                         (spec.cellcount[1] * spec.order + 1) *
+                         (spec.cellcount[2] * spec.order + 1);
+    if (nnode != expected_nodes) {
+      std::string msg = "number of nodes (" + std::to_string(nnode) +
+                        ") does not match cellcount and dof expected (" +
+                        std::to_string(expected_nodes) + ")";
+      throw mjCError(this, msg.c_str());
     }
   }
 
@@ -4329,12 +4345,48 @@ void mjCFlex::Compile(const mjVFS* vfs) {
   }
 
   if (!stiffness_cached && young > 0 && interpolated) {
-    int n = pow(order_ + 1, 3);
-    int ndof = 3 * n;
-    if (stiffness.size() < ndof * ndof) {
-      stiffness.resize(ndof * ndof, 0);
+    int npc = pow(spec.order + 1, 3);  // nodes per cell
+    int ndof_cell = 3 * npc;
+    int cx = spec.cellcount[0], cy = spec.cellcount[1], cz = spec.cellcount[2];
+    int ncells = cx * cy * cz;
+    int ny_global = cy * spec.order + 1;
+    int nz_global = cz * spec.order + 1;
+
+    // total stiffness = ncells * ndof_cell^2
+    stiffness.resize(ncells * ndof_cell * ndof_cell, 0);
+
+    // compute stiffness per cell
+    for (int ci = 0; ci < cx; ci++) {
+      for (int cj = 0; cj < cy; cj++) {
+        for (int ck = 0; ck < cz; ck++) {
+          int cell_idx = ci * cy * cz + cj * cz + ck;
+
+          // gather cell's local node positions
+          std::vector<double> cell_pos(3 * npc);
+          int local = 0;
+          for (int li = 0; li <= spec.order; li++) {
+            for (int lj = 0; lj <= spec.order; lj++) {
+              for (int lk = 0; lk <= spec.order; lk++) {
+                int gi = ci * spec.order + li;
+                int gj = cj * spec.order + lj;
+                int gk = ck * spec.order + lk;
+                int global = gi * ny_global * nz_global + gj * nz_global + gk;
+                mjuu_copyvec(cell_pos.data() + 3*local, nodexpos.data() + 3*global, 3);
+                local++;
+              }
+            }
+          }
+
+          // compute per-cell stiffness
+          std::vector<double> K_cell(ndof_cell * ndof_cell, 0);
+          ComputeLinearStiffness(K_cell, cell_pos.data(), young, poisson, spec.order);
+
+          // copy into global stiffness array
+          mjuu_copyvec(stiffness.data() + cell_idx * ndof_cell * ndof_cell,
+                       K_cell.data(), ndof_cell * ndof_cell);
+        }
+      }
     }
-    ComputeLinearStiffness(stiffness, nodexpos.data(), young, poisson, order_);
   }
 
   // create bounding volume hierarchy

src/user/user_model.cc

@@ -2190,8 +2190,13 @@ void mjCModel::SetSizes() {
     nflexshelldata += (int)flexes_[i]->shell.size();
     nflexevpair += (int)flexes_[i]->evpair.size()/2;
     nflextexcoord += (flexes_[i]->HasTexcoord() ? flexes_[i]->get_texcoord().size()/2 : 0);
-    if (flexes_[i]->order_ != 0) {
-      extra_stiffness_size += (3 * flexes_[i]->nnode) * (3 * flexes_[i]->nnode);
+    if (flexes_[i]->spec.order != 0) {
+      int npc = (int)pow(flexes_[i]->spec.order + 1, 3);
+      int ndof_cell = 3 * npc;
+      int ncells = flexes_[i]->spec.cellcount[0] *
+                   flexes_[i]->spec.cellcount[1] *
+                   flexes_[i]->spec.cellcount[2];
+      extra_stiffness_size += ncells * ndof_cell * ndof_cell;
     }
     if (flexes_[i]->interpolated || flexes_[i]->rigid) {
       continue;
@@ -3467,18 +3472,30 @@ void mjCModel::CopyObjects(mjModel* m) {
     mjuu_copyvec(m->flex_rgba + 4 * i, pfl->rgba, 4);
 
     // elasticity
-    if (pfl->order_ == 0) {
+    if (pfl->spec.order == 0) {
       m->flex_stiffnessadr[i] = 21 * elem_adr;
     } else {
       m->flex_stiffnessadr[i] = current_extra_stiffness_adr;
-      current_extra_stiffness_adr += (3 * pfl->nnode) * (3 * pfl->nnode);
+      int npc = (int)pow(pfl->spec.order + 1, 3);
+      int ndof_cell = 3 * npc;
+      int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
+      current_extra_stiffness_adr += ncells * ndof_cell * ndof_cell;
     }
 
     if (!pfl->stiffness.empty()) {
-      mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i], pfl->stiffness.data(), pfl->stiffness.size());
+      mjuu_copyvec(m->flex_stiffness + m->flex_stiffnessadr[i],
+                   pfl->stiffness.data(), pfl->stiffness.size());
     } else {
-      int size = (pfl->order_ == 0) ? 21 * pfl->nelem : (3 * pfl->nnode) * (3 * pfl->nnode);
-      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], size);
+      int stiff_size;
+      if (pfl->spec.order == 0) {
+        stiff_size = 21 * pfl->nelem;
+      } else {
+        int npc = (int)pow(pfl->spec.order + 1, 3);
+        int ndof_cell = 3 * npc;
+        int ncells = pfl->spec.cellcount[0] * pfl->spec.cellcount[1] * pfl->spec.cellcount[2];
+        stiff_size = ncells * ndof_cell * ndof_cell;
+      }
+      mjuu_zerovec(m->flex_stiffness + m->flex_stiffnessadr[i], stiff_size);
     }
     if (!pfl->bending.empty()) {
       mjuu_copyvec(m->flex_bending + 17 * edge_adr, pfl->bending.data(), pfl->bending.size());
@@ -3613,7 +3630,12 @@ void mjCModel::CopyObjects(mjModel* m) {
     }
 
     // set interpolation type, only two types for now
-    m->flex_interp[i] = pfl->order_;
+    m->flex_interp[i] = pfl->spec.order;
+
+    // set cell count for multi-cell finite cell method
+    m->flex_cellnum[3*i+0] = pfl->spec.cellcount[0];
+    m->flex_cellnum[3*i+1] = pfl->spec.cellcount[1];
+    m->flex_cellnum[3*i+2] = pfl->spec.cellcount[2];
 
     // convert edge pairs to int array, set edge rigid
     for (int k=0; k < pfl->nedge; k++) {
@@ -4967,6 +4989,20 @@ void mjCModel::ResolveKeyframes(const mjModel* m) {
 }
 
 void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
+#if defined(__EMSCRIPTEN__) && !defined(MUJOCO_WASM_THREADS)
+  // The MuJoCo compiler defaults to usethread=1, which causes it to try to
+  // create pthreads for compilation. In the single-threaded WASM build, this
+  // crashes because there is no threading support, so we disable threading on
+  // the internal compiler struct (not the spec) to avoid permanently mutating
+  // the spec (which would cause usethread="false" to appear in a saved XML).
+  struct ScopedDisableThreading {
+    mjtByte& ref;
+    mjtByte saved;
+    explicit ScopedDisableThreading(mjtByte& r) : ref(r), saved(r) { ref = 0; }
+    ~ScopedDisableThreading() { ref = saved; }
+  } disable_usethread(compiler.usethread);
+#endif
+
   // check if nan test works
   double test = mjNAN;
   if (mjuu_defined(test)) {

src/user/user_objects.h

@@ -1042,7 +1042,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
 
   static constexpr int kNumEdges[3] = {1, 3, 6};  // number of edges per element indexed by dim
 
-  void SetOrder(int order) { order_ = order; }  // set interpolation order
+
 
  private:
   void Compile(const mjVFS* vfs);         // compiler
@@ -1052,7 +1052,7 @@ class mjCFlex: public mjCFlex_, private mjsFlex {
   std::vector<double> vert0_;             // vertex positions in [0, 1]^d in the bounding box
   std::vector<double> node0_;             // node Cartesian positions
 
-  int order_ = 0;                         // interpolation order
+
 
   // stiffness caching
   std::string ComputeStiffnessCacheKey() const;

src/xml/xml_native_reader.cc

@@ -315,7 +315,7 @@ std::vector<const char*> MJCF[nMJCF] = {
             {">"},
         {">"},
         {"flexcomp", "*", "name", "type", "group", "dim", "dof",
-            "count", "spacing", "radius", "rigid", "mass", "inertiabox",
+            "count", "cellcount", "spacing", "radius", "rigid", "mass", "inertiabox",
             "scale", "file", "point", "element", "texcoord", "material", "rgba",
             "flatskin", "pos", "quat", "axisangle", "xyaxes", "zaxis", "euler", "origin"},
         {"<"},
@@ -334,8 +334,8 @@ std::vector<const char*> MJCF[nMJCF] = {
 
     {"deformable", "*"},
     {"<"},
-        {"flex", "*", "name", "group", "dim", "radius", "material",
-            "rgba", "flatskin", "body", "vertex", "element", "texcoord", "elemtexcoord", "node"},
+        {"flex", "*", "name", "group", "dim", "radius", "material", "rgba", "flatskin", "body",
+            "vertex", "element", "texcoord", "elemtexcoord", "node", "cellcount", "dof"},
         {"<"},
             {"contact", "?",  "contype", "conaffinity", "condim", "priority",
                 "friction", "solmix", "solref", "solimp", "margin", "gap",
@@ -1501,6 +1501,16 @@ void mjXReader::OneFlex(XMLElement* elem, mjsFlex* flex) {
   ReadAttrInt(elem, "dim", &flex->dim);
   ReadAttrInt(elem, "group", &flex->group);
 
+  flex->cellcount[0] = 1;
+  flex->cellcount[1] = 1;
+  flex->cellcount[2] = 1;
+  ReadAttr(elem, "cellcount", 3, flex->cellcount, text);
+
+  flex->order = 0;
+  if (MapValue(elem, "dof", &n, fdof_map, mjNFCOMPDOFS)) {
+    flex->order = (n == mjFCOMPDOF_QUADRATIC) ? 2 : (n == mjFCOMPDOF_TRILINEAR ? 1 : 0);
+  }
+
   // read data vectors
   if (ReadAttrTxt(elem, "body", text, true)) {
     mjs_setStringVec(flex->vertbody, text.c_str());
@@ -2794,6 +2804,7 @@ void mjXReader::OneFlexcomp(XMLElement* elem, mjsBody* body, const mjVFS* vfs) {
     fcomp.type = (mjtFcompType)n;
   }
   ReadAttr(elem, "count", 3, fcomp.count, text);
+  ReadAttr(elem, "cellcount", 3, fcomp.cellcount, text);
   ReadAttr(elem, "spacing", 3, fcomp.spacing, text);
   ReadAttr(elem, "scale", 3, fcomp.scale, text);
   ReadAttr(elem, "mass", 1, &fcomp.mass, text);

src/xml/xml_native_writer.cc

@@ -141,6 +141,13 @@ void mjXWriter::OneFlex(XMLElement* elem, const mjCFlex* flex) {
   WriteAttrKey(elem, "flatskin", bool_map, 2, flex->flatskin, defflex.flatskin);
   WriteAttrInt(elem, "dim", flex->dim, defflex.dim);
   WriteAttrInt(elem, "group", flex->group, defflex.group);
+  WriteAttr(elem, "cellcount", 3, flex->spec.cellcount, defflex.spec.cellcount);
+  if (flex->spec.order != defflex.spec.order) {
+    string dof_str = "full";
+    if (flex->spec.order == 1) dof_str = "trilinear";
+    else if (flex->spec.order == 2) dof_str = "quadratic";
+    WriteAttrTxt(elem, "dof", dof_str);
+  }
 
   // data vectors
   if (!flex->get_vertbody().empty()) {

test/engine/engine_core_constraint_test.cc

@@ -35,136 +35,9 @@ namespace {
 
 using ::testing::NotNull;
 using ::testing::Pointwise;
+
 using CoreConstraintTest = MujocoTest;
 
-// compute rotation residual following formula in mj_instantiateEquality
-void RotationResidual(const mjModel *model, mjData *data,
-                        const mjtNum qpos[7], const mjtNum dqpos[6],
-                        mjtNum res[3]) {
-  // copy configuration, compute required quantities with mj_step1
-  mju_copy(data->qpos, qpos, 7);
-
-  // perturb configuration if given
-  if (dqpos) {
-    mj_integratePos(model, data->qpos, dqpos, 1);
-  }
-
-  // update relevant quantities
-  mj_step1(model, data);
-
-  // compute orientation residual
-  mjtNum quat1[4], quat2[4], quat3[4];
-  mju_copy4(quat1, data->xquat+4*1);
-  mju_negQuat(quat2, data->xquat+4*2);
-  mju_mulQuat(quat3, quat2, quat1);
-  mju_copy3(res, quat3+1);
-}
-
-// validate rotational Jacobian used in welds
-TEST_F(CoreConstraintTest, WeldRotJacobian) {
-#ifdef mjUSESINGLE
-  GTEST_SKIP() << "FD Jacobian with eps=1e-6 below float32 precision";
-#endif
-  constexpr char xml[] = R"(
-  <mujoco>
-    <option jacobian="dense"/>
-    <worldbody>
-      <body>
-        <joint type="ball"/>
-        <geom size=".1"/>
-      </body>
-      <body pos=".5 0 0">
-        <joint axis="1 0 0" pos="0 0 .01"/>
-        <joint axis="0 1 0" pos=".02 0 0"/>
-        <joint axis="0 0 1" pos="0 .03 0"/>
-        <geom size=".1"/>
-      </body>
-    </worldbody>
-  </mujoco>
-  )";
-  char error[1024];
-  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
-  ASSERT_THAT(model, testing::NotNull()) << error;
-  ASSERT_EQ(model->nq, 7);
-  ASSERT_EQ(model->nv, 6);
-  static const int nv = 6;  // for increased readability
-  mjData* data = mj_makeData(model);
-
-  // arbitrary initial values for the ball and hinge joints
-  mjtNum qpos0[7] = {.5, .5, .5, .5, .7, .8, .9};
-
-  // compute required quantities using mj_step1
-  mj_step1(model, data);
-
-  // get orientation error
-  mjtNum res[3];
-  RotationResidual(model, data, qpos0, NULL, res);
-
-  // compute Jacobian with finite-differencing
-  mjtNum jacFD[3*nv];
-  mjtNum dqpos[nv] = {0};
-  mjtNum dres[3];
-  const mjtNum eps = 1e-6;
-  for (int i=0; i < nv; i++) {
-    // nudge i-th dof
-    dqpos[i] = eps;
-
-    // get nudged residual
-    RotationResidual(model, data, qpos0, dqpos, dres);
-
-    // remove nudge
-    dqpos[i] = 0.0;
-
-    // compute Jacobian column
-    for (int j=0; j < 3; j++) {
-      jacFD[nv*j + i] = (dres[j] - res[j]) / eps;
-    }
-  }
-
-  // reset mjData to qpos0
-  mju_copy(data->qpos, qpos0, 7);
-  mj_step1(model, data);
-
-  // intermediate quaternions quat1 and quat2
-  mjtNum quat1[4], negQuat2[4];
-  mju_copy4(quat1, data->xquat+4*1);
-  mju_negQuat(negQuat2, data->xquat+4*2);
-
-  // get analytical Jacobian following formula in mj_instantiateEquality
-  mjtNum jacdif[3*nv], jac0[3*nv], jac1[3*nv];
-  mjtNum point[3] = {0};
-
-  // rotational Jacobian difference
-  mj_jacDifPair(model, data, NULL, 2, 1, point, point,
-                NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model),
-                /*flg_skipcommon=*/0);
-
-  // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1
-  mjtNum axis[3], quat3[4], quat4[4];
-  for (int j=0; j < nv; j++) {
-    // axis = [jac1-jac2]_col(j)
-    axis[0] = jacdif[0*nv+j];
-    axis[1] = jacdif[1*nv+j];
-    axis[2] = jacdif[2*nv+j];
-
-    // apply formula
-    mju_mulQuatAxis(quat3, negQuat2, axis);
-    mju_mulQuat(quat4, quat3, quat1);
-
-    // correct Jacobian
-    jacdif[0*nv+j] = 0.5*quat4[1];
-    jacdif[1*nv+j] = 0.5*quat4[2];
-    jacdif[2*nv+j] = 0.5*quat4[3];
-  }
-
-  // test that analytical and finite-differenced Jacobians match
-  EXPECT_THAT(AsVector(jacFD, 3*nv),
-              Pointwise(MjNear(eps, 1e-3), AsVector(jacdif, 3*nv)));
-
-  mj_deleteData(data);
-  mj_deleteModel(model);
-}
-
 // test formulas for penetration at rest
 TEST_F(CoreConstraintTest, RestPenetration) {
   constexpr char xml[] = R"(

test/engine/engine_core_util_test.cc

@@ -0,0 +1,740 @@
+// 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 engine/engine_core_util.c.
+
+#include "src/engine/engine_core_util.h"
+
+#include <algorithm>
+#include <cstddef>
+#include <limits>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include <mujoco/mjdata.h>
+#include <mujoco/mujoco.h>
+#include "test/fixture.h"
+
+namespace mujoco {
+namespace {
+
+using ::testing::NotNull;
+using ::testing::Pointwise;
+
+using FlexGatherStateTest = MujocoTest;
+
+TEST_F(FlexGatherStateTest, mju_flexGatherState_Grid) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <flexcomp name="flex0" type="grid" count="3 3 3" spacing=".1 .1 .1"
+                dim="3" mass="1" radius="0.01" dof="trilinear">
+        <elasticity young="5e4" poisson="0.2"/>
+        <contact selfcollide="none"/>
+      </flexcomp>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  mjData* data = mj_makeData(model);
+  mj_forward(model, data);
+
+  ASSERT_EQ(model->nflex, 1);
+  int f = 0;
+  int nodenum = model->flex_nodenum[f];
+  int nstart = model->flex_nodeadr[f];
+
+  // Simulate a rotated state (90 degrees around Z axis)
+  ASSERT_TRUE(model->flex_centered[f]);
+  for (int i = 0; i < nodenum; i++) {
+    int b = model->flex_nodebodyid[nstart + i];
+    mjtNum x = data->xpos[3*b + 0];
+    mjtNum y = data->xpos[3*b + 1];
+    mjtNum z = data->xpos[3*b + 2];
+
+    // Rotate 90 degrees around Z: (x, y, z) -> (-y, x, z)
+    data->xpos[3*b + 0] = -y;
+    data->xpos[3*b + 1] = x;
+    data->xpos[3*b + 2] = z;
+  }
+
+  std::vector<mjtNum> xpos(3 * nodenum);
+  mju_flexGatherState(model, data, f, xpos.data(), NULL);
+
+  // Verify that gathered xpos matches the rotated data->xpos
+  for (int i = 0; i < nodenum; i++) {
+    int b = model->flex_nodebodyid[nstart + i];
+    EXPECT_NEAR(xpos[3*i + 0], data->xpos[3*b + 0], 1e-5);
+    EXPECT_NEAR(xpos[3*i + 1], data->xpos[3*b + 1], 1e-5);
+    EXPECT_NEAR(xpos[3*i + 2], data->xpos[3*b + 2], 1e-5);
+  }
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+
+using AngMomMatTest = MujocoTest;
+
+static constexpr char AngMomTestingModel[] = R"(
+<mujoco>
+  <option>
+    <flag gravity="disable"/>
+  </option>
+  <worldbody>
+    <body name="link1" pos="0 0 0.5">
+      <freejoint/>
+      <geom type="ellipsoid" size="0.15 0.17 0.19" quat="1 .2 .3 .4"/>
+      <body name="link2" >
+        <joint type="hinge" axis="1 0 0" />
+        <geom type="capsule" size="0.05" fromto="0 0 0  0 0.5 0"/>
+        <body pos="0 0.6 0">
+          <joint type="slide" axis="1 0 0"/>
+          <geom type="capsule" size="0.05 0.2" quat="0.707 0 0.707 0"/>
+          <body name="link3">
+            <joint type="ball" pos="0.2 0 0"/>
+            <geom type="capsule" pos="0.2 0 0" size="0.03 0.4"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <keyframe>
+    <key qvel="0 0 0 .1 .2 .3 .4 .5 .4 .3 .2"/>
+  </keyframe>
+</mujoco>
+)";
+
+// compare subtree angular momentum computed in two ways
+TEST_F(AngMomMatTest, CompareAngMom) {
+  char error[1024];
+  mjModel* model =
+      LoadModelFromString(AngMomTestingModel, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
+
+  mjData* data = mj_makeData(model);
+
+  // reset to the keyframe with some angular velocities
+  mj_resetDataKeyframe(model, data, 0);
+  mj_forward(model, data);
+
+  // get the reference value of angular momentum
+  mj_subtreeVel(model, data);
+  mjtNum angmom_ref[3];
+  mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid);
+
+  // compute angular momentum using the angular momentum matrix
+  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+  mj_angmomMat(model, data, angmom_mat, bodyid);
+  mjtNum angmom_test[3];
+  mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv);
+
+  // compare the two angular momentum values
+  for (int i = 0; i < 3; i++) {
+    EXPECT_THAT(angmom_ref[i], MjNear(angmom_test[i], 1e-8, 1e-4));
+  }
+
+  mju_free(angmom_mat);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+// compare subtree angular momentum matrix: analytical and findiff
+TEST_F(AngMomMatTest, CompareAngMomMats) {
+  char error[1024];
+  mjModel* model =
+      LoadModelFromString(AngMomTestingModel, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
+  mjData* data = mj_makeData(model);
+  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+  mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+
+  // reset to the keyframe with some angular velocities
+  mj_resetDataKeyframe(model, data, 0);
+  mj_forward(model, data);
+
+  // compute the angular momentum matrix using the analytical method
+  mj_angmomMat(model, data, angmom_mat, bodyid);
+
+  // compute the angular momentum matrix using finite differences
+  static constexpr mjtNum eps = MjTol(1e-6, 1e-3);
+  for (int i = 0; i < nv; i++) {
+    // reset vel, forward nudge i-th dof, get angmom
+    mju_copy(data->qvel, model->key_qvel, model->nv);
+    data->qvel[i] += eps;
+    mj_forward(model, data);
+    mj_subtreeVel(model, data);
+    mjtNum agmf[3];
+    mju_copy3(agmf, data->subtree_angmom+3*bodyid);
+
+    // reset vel, backward nudge i-th dof, get angmom
+    mju_copy(data->qvel, model->key_qvel, model->nv);
+    data->qvel[i] -= eps;
+    mj_forward(model, data);
+    mj_subtreeVel(model, data);
+    mjtNum agmb[3];
+    mju_copy3(agmb, data->subtree_angmom+3*bodyid);
+
+    // finite-difference the angmom matrix
+    for (int j = 0; j < 3; j++) {
+      angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps);
+    }
+  }
+
+  // compare the two matrices
+  for (int i = 0; i < 3*nv; i++) {
+    EXPECT_THAT(angmom_mat_fd[i], MjNear(angmom_mat[i], 1e-8, 2e-4));
+  }
+
+  mju_free(angmom_mat_fd);
+  mju_free(angmom_mat);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+using JacobianTest = MujocoTest;
+static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
+
+static constexpr char kJacobianTestingModel[] = R"(
+<mujoco>
+  <worldbody>
+    <body name="distractor1" pos="0 0 .3">
+      <freejoint/>
+      <geom size=".1"/>
+    </body>
+    <body name="main">
+      <freejoint/>
+      <geom size=".1"/>
+      <body pos=".1 0 0">
+        <joint axis="0 1 0"/>
+        <geom type="capsule" size=".03" fromto="0 0 0 .2 0 0"/>
+      </body>
+      <body pos="0 .1 0">
+        <joint type="ball"/>
+        <geom type="capsule" size=".03" fromto="0 0 0 0 .2 0"/>
+        <body pos="0 .2 0">
+          <joint type="slide" axis="1 1 1"/>
+          <geom size=".05"/>
+        </body>
+      </body>
+    </body>
+    <body name="distractor2" pos="0 0 -.3">
+      <freejoint/>
+      <geom size=".1"/>
+    </body>
+  </worldbody>
+</mujoco>
+)";
+
+// compare analytic and finite-differenced subtree-com Jacobian
+TEST_F(JacobianTest, SubtreeJac) {
+  char error[1024];
+  mjModel* model =
+      LoadModelFromString(kJacobianTestingModel, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "main");
+  mjData* data = mj_makeData(model);
+  mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+  mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nq);
+  mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv);
+
+  // all we need for Jacobians are kinematics and CoM-related quantities
+  mj_kinematics(model, data);
+  mj_comPos(model, data);
+
+  // get subtree CoM Jacobian of free body
+  mj_jacSubtreeCom(model, data, jac_subtree, bodyid);
+
+  // save current subtree-com and qpos, clear nudge
+  mjtNum subtree_com[3];
+  mju_copy3(subtree_com, data->subtree_com+3*bodyid);
+  mju_copy(qpos, data->qpos, model->nq);
+  mju_zero(nudge, nv);
+
+  // compare analytic Jacobian to finite-difference approximation
+  static const mjtNum eps = 1e-6;
+  for (int i=0; i < nv; i++) {
+    // reset qpos, nudge i-th dof, update data->qpos, reset nudge
+    mju_copy(data->qpos, qpos, model->nq);
+    nudge[i] = 1;
+    mj_integratePos(model, data->qpos, nudge, eps);
+    nudge[i] = 0;
+
+    // kinematics and comPos to get nudged com
+    mj_kinematics(model, data);
+    mj_comPos(model, data);
+
+    // compare finite-differenced and analytic Jacobian
+    for (int j=0; j < 3; j++) {
+      mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps;
+      EXPECT_THAT(jac_subtree[nv*j+i], MjNear(findiff, eps, 1e-2));
+    }
+  }
+
+  mju_free(nudge);
+  mju_free(qpos);
+  mju_free(jac_subtree);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+// confirm that applying linear forces via the subtree-com Jacobian only creates
+// the expected linear accelerations (no accelerations of internal joints)
+TEST_F(JacobianTest, SubtreeJacNoInternalAcc) {
+  char error[1024];
+  mjModel* model =
+      LoadModelFromString(kJacobianTestingModel, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  int nv = model->nv;
+  int bodyid = mj_name2id(model, mjOBJ_BODY, "main");
+  mjData* data = mj_makeData(model);
+  mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
+
+  // all we need for Jacobians are kinematics and CoM-related quantities
+  mj_kinematics(model, data);
+  mj_comPos(model, data);
+
+  // get subtree CoM Jacobian of free body
+  mj_jacSubtreeCom(model, data, jac_subtree, bodyid);
+
+  // uncomment for debugging
+  // mju_printMat(jac_subtree, 3, nv);
+
+  // call fwdPosition since we'll need the factorised mass matrix in the test
+  mj_fwdPosition(model, data);
+
+  // treating the subtree Jacobian as the projection of 3 axis-aligned unit
+  // forces into joint space, solve for the resulting accelerations in-place
+  mj_solveM(model, data, jac_subtree, jac_subtree, 3);
+
+  // expect to find accelerations of magnitude 1/subtreemass in the first 3
+  // coordinates of the free joint and 0s elsewhere, since applying forces to
+  // the CoM should accelerate the whole mechanism without any internal motion
+  int body_dofadr = model->body_dofadr[bodyid];
+  mjtNum invtreemass = 1.0/model->body_subtreemass[bodyid];
+  for (int r = 0; r < 3; r++) {
+    for (int c = 0; c < nv; c++) {
+      mjtNum expected = c - body_dofadr == r ? invtreemass : 0.0;
+      EXPECT_THAT(jac_subtree[nv*r+c], MjNear(expected, max_abs_err, 1e-4));
+    }
+  }
+
+  mju_free(jac_subtree);
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+static constexpr char kQuat[] = R"(
+<mujoco>
+  <worldbody>
+    <body name="query">
+      <joint type="ball"/>
+      <geom size="1"/>
+      <site name="query" pos=".1 .2 .3"/>
+    </body>
+  </worldbody>
+  <keyframe>
+    <key qvel="2 3 5"/>
+  </keyframe>
+</mujoco>
+)";
+
+static constexpr char kFreeBall[] = R"(
+<mujoco>
+  <worldbody>
+    <body name="distractor1" pos="0 0 .3">
+      <freejoint/>
+      <geom size=".1"/>
+    </body>
+    <body name="main">
+      <freejoint/>
+      <geom size=".1"/>
+      <body pos=".1 0 0">
+        <joint axis="0 1 0"/>
+        <geom type="capsule" size=".03" fromto="0 0 0 .2 0 0"/>
+        <body pos=".2 0 0">
+          <joint type="ball" stiffness="20"/>
+          <geom type="capsule" size=".03" fromto="0 0 0 0 .2 0"/>
+          <body name="query" pos="0 .2 0">
+            <joint type="slide" axis="1 1 1"/>
+            <geom size=".05"/>
+            <site name="query" pos=".1 .2 .3"/>
+          </body>
+        </body>
+      </body>
+    </body>
+    <body name="distractor2" pos="0 0 -.3">
+      <freejoint/>
+      <geom size=".1"/>
+    </body>
+  </worldbody>
+  <keyframe>
+    <key qvel="1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1"/>
+  </keyframe>
+</mujoco>
+)";
+
+static constexpr char kQuatlessPendulum[] = R"(
+<mujoco>
+  <option integrator="implicit">
+    <flag constraint="disable"/>
+  </option>
+  <worldbody>
+    <body pos="0.15 0 0">
+      <joint type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+      <body pos="0.1 0 0">
+        <joint type="slide" axis="1 0 0" stiffness="200"/>
+        <geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
+        <body pos=".1 0 0">
+          <joint axis="1 0 0"/>
+          <joint axis="0 1 0"/>
+          <joint axis="0 0 1"/>
+          <geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
+          <body name="query" pos="0 .1 0">
+            <joint axis="1 0 0"/>
+            <geom type="capsule" size="0.02" fromto="0 0 0 0 .1 0"/>
+            <site name="query" pos=".1 0 0"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+)";
+
+static constexpr char kTelescope[] = R"(
+<mujoco>
+  <worldbody>
+    <body>
+      <joint type="ball"/>
+      <geom type="capsule" size="0.02" fromto="0 .02 0 .1 .02 0"/>
+      <body pos=".1 .02 0">
+        <joint type="slide" axis="1 0 0"/>
+        <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+        <body pos=".1 .02 0">
+          <joint type="slide" axis="1 0 0"/>
+          <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+          <body pos=".1 .02 0" name="query">
+            <joint type="slide" axis="1 0 0"/>
+            <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+            <site name="query" pos=".1 0 0"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <keyframe>
+    <key qvel="1 1 1 1 1 1"/>
+  </keyframe>
+</mujoco>
+)";
+
+static constexpr char kHinge[] = R"(
+<mujoco>
+  <worldbody>
+    <body name="query">
+      <joint name="link1" axis="0 1 0"/>
+      <geom type="capsule" size=".02" fromto="0 0 0 0 0 -1"/>
+      <site name="query" pos="0 0 -1"/>
+    </body>
+  </worldbody>
+
+  <keyframe>
+    <key qpos="1" qvel="1"/>
+  </keyframe>
+</mujoco>
+)";
+
+// compare mj_jacDot with finite-differenced mj_jac
+TEST_F(JacobianTest, JacDot) {
+  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
+    char error[1024];
+    mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+    ASSERT_THAT(model, NotNull()) << error;
+    int nv = model->nv;
+    mjData* data = mj_makeData(model);
+
+    // load keyframe if present, step for a bit
+    if (model->nkey) mj_resetDataKeyframe(model, data, 0);
+    while (data->time < 0.1) {
+      mj_step(model, data);
+    }
+
+    // minimal call required for mj_jacDot outputs to be valid
+    mj_kinematics(model, data);
+    mj_comPos(model, data);
+    mj_comVel(model, data);
+
+    // get bodyid
+    int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
+    EXPECT_GT(bodyid, 0);
+
+    // get site position
+    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
+    EXPECT_GT(siteid, -1);
+    mjtNum point[3];
+    mju_copy3(point, data->site_xpos+3*siteid);
+
+    // jac, jac_dot
+    std::vector<mjtNum> jacp(3*nv);
+    std::vector<mjtNum> jacr(3*nv);
+    mj_jac(model, data, jacp.data(), jacr.data(), point, bodyid);
+    std::vector<mjtNum> jacp_dot(3*nv);
+    std::vector<mjtNum> jacr_dot(3*nv);
+    mj_jacDot(model, data, jacp_dot.data(), jacr_dot.data(), point, bodyid);
+
+    // jac_h: jacobian after integrating qpos with a timestep of h
+    constexpr mjtNum h = MjTol(1e-7, 5e-4);
+    mj_integratePos(model, data->qpos, data->qvel, h);
+    mj_kinematics(model, data);
+    mj_comPos(model, data);
+    std::vector<mjtNum> jacp_h(3*nv);
+    std::vector<mjtNum> jacr_h(3*nv);
+    mju_copy3(point, data->site_xpos+3*siteid);  // get updated site position
+    mj_jac(model, data, jacp_h.data(), jacr_h.data(), point, bodyid);
+
+    // jac_dot_h finite-difference approximation
+    std::vector<mjtNum> jacp_dot_h(3*nv);
+    mju_sub(jacp_dot_h.data(), jacp_h.data(), jacp.data(), 3*nv);
+    mju_scl(jacp_dot_h.data(), jacp_dot_h.data(), 1/h, 3*nv);
+    std::vector<mjtNum> jacr_dot_h(3*nv);
+    mju_sub(jacr_dot_h.data(), jacr_h.data(), jacr.data(), 3*nv);
+    mju_scl(jacr_dot_h.data(), jacr_dot_h.data(), 1/h, 3*nv);
+
+    // compare finite-differenced and analytic
+    mjtNum tol = 1e-5;
+    EXPECT_THAT(jacp_dot, Pointwise(MjNear(tol, 5e-2), jacp_dot_h));
+    EXPECT_THAT(jacr_dot, Pointwise(MjNear(tol, 5e-2), jacr_dot_h));
+
+    mj_deleteData(data);
+    mj_deleteModel(model);
+  }
+}
+
+// compare mj_jacDotSparse with dense mj_jacDot
+TEST_F(JacobianTest, JacDotSparse) {
+  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
+    char error[1024];
+    mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+    ASSERT_THAT(model, NotNull()) << error;
+    int nv = model->nv;
+    mjData* data = mj_makeData(model);
+
+    // load keyframe if present, step for a bit
+    if (model->nkey) mj_resetDataKeyframe(model, data, 0);
+    while (data->time < 0.1) {
+      mj_step(model, data);
+    }
+
+    // minimal call required for mj_jacDot outputs to be valid
+    mj_kinematics(model, data);
+    mj_comPos(model, data);
+    mj_comVel(model, data);
+
+    // get bodyid and site position
+    int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
+    EXPECT_GT(bodyid, 0);
+    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
+    EXPECT_GT(siteid, -1);
+    mjtNum point[3];
+    mju_copy3(point, data->site_xpos+3*siteid);
+
+    // dense jacDot
+    std::vector<mjtNum> jacp_dense(3*nv);
+    std::vector<mjtNum> jacr_dense(3*nv);
+    mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid);
+
+    // compute body chain using public mjModel fields
+    std::vector<int> chain(nv);
+    int NV = 0;
+    int weldbody = model->body_weldid[bodyid];
+    if (weldbody) {
+      int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1;
+      while (da >= 0) {
+        chain[NV++] = da;
+        da = model->dof_parentid[da];
+      }
+      std::reverse(chain.begin(), chain.begin() + NV);
+    }
+    EXPECT_GT(NV, 0);
+
+    // sparse jacDot
+    std::vector<mjtNum> jacp_sparse(3*NV);
+    std::vector<mjtNum> jacr_sparse(3*NV);
+    mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(),
+                    point, bodyid, NV, chain.data());
+
+    // expand sparse to dense and compare
+    std::vector<mjtNum> jacp_expanded(3*nv, 0);
+    std::vector<mjtNum> jacr_expanded(3*nv, 0);
+    for (int ci = 0; ci < NV; ci++) {
+      int di = chain[ci];
+      for (int r = 0; r < 3; r++) {
+        jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV];
+        jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV];
+      }
+    }
+
+    // expect bitwise equality
+    EXPECT_EQ(jacp_expanded, jacp_dense);
+    EXPECT_EQ(jacr_expanded, jacr_dense);
+
+    mj_deleteData(data);
+    mj_deleteModel(model);
+  }
+}
+
+
+// validate rotational Jacobian used in welds
+TEST_F(JacobianTest, WeldRotJacobian) {
+#ifdef mjUSESINGLE
+  GTEST_SKIP() << "FD Jacobian with eps=1e-6 below float32 precision";
+#endif
+  constexpr char xml[] = R"(
+  <mujoco>
+    <option jacobian="dense"/>
+    <worldbody>
+      <body>
+        <joint type="ball"/>
+        <geom size=".1"/>
+      </body>
+      <body pos=".5 0 0">
+        <joint axis="1 0 0" pos="0 0 .01"/>
+        <joint axis="0 1 0" pos=".02 0 0"/>
+        <joint axis="0 0 1" pos="0 .03 0"/>
+        <geom size=".1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+  char error[1024];
+  mjModel* model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, testing::NotNull()) << error;
+  ASSERT_EQ(model->nq, 7);
+  ASSERT_EQ(model->nv, 6);
+  static const int nv = 6;  // for increased readability
+  mjData* data = mj_makeData(model);
+
+  // arbitrary initial values for the ball and hinge joints
+  mjtNum qpos0[7] = {.5, .5, .5, .5, .7, .8, .9};
+
+  // compute required quantities using mj_step1
+  mj_step1(model, data);
+
+  // get orientation error
+  mjtNum res[3];
+  // compute rotation residual following formula in mj_instantiateEquality
+  auto RotationResidual = [](const mjModel *model, mjData *data,
+                             const mjtNum qpos[7], const mjtNum dqpos[6],
+                             mjtNum res[3]) {
+    // copy configuration, compute required quantities with mj_step1
+    mju_copy(data->qpos, qpos, 7);
+
+    // perturb configuration if given
+    if (dqpos) {
+      mj_integratePos(model, data->qpos, dqpos, 1);
+    }
+
+    // update relevant quantities
+    mj_step1(model, data);
+
+    // compute orientation residual
+    mjtNum quat1[4], quat2[4], quat3[4];
+    mju_copy4(quat1, data->xquat+4*1);
+    mju_negQuat(quat2, data->xquat+4*2);
+    mju_mulQuat(quat3, quat2, quat1);
+    mju_copy3(res, quat3+1);
+  };
+
+  RotationResidual(model, data, qpos0, NULL, res);
+
+  // compute Jacobian with finite-differencing
+  mjtNum jacFD[3*nv];
+  mjtNum dqpos[nv] = {0};
+  mjtNum dres[3];
+  const mjtNum eps = 1e-6;
+  for (int i=0; i < nv; i++) {
+    // nudge i-th dof
+    dqpos[i] = eps;
+
+    // get nudged residual
+    RotationResidual(model, data, qpos0, dqpos, dres);
+
+    // remove nudge
+    dqpos[i] = 0.0;
+
+    // compute Jacobian column
+    for (int j=0; j < 3; j++) {
+      jacFD[nv*j + i] = (dres[j] - res[j]) / eps;
+    }
+  }
+
+  // reset mjData to qpos0
+  mju_copy(data->qpos, qpos0, 7);
+  mj_step1(model, data);
+
+  // intermediate quaternions quat1 and quat2
+  mjtNum quat1[4], negQuat2[4];
+  mju_copy4(quat1, data->xquat+4*1);
+  mju_negQuat(negQuat2, data->xquat+4*2);
+
+  // get analytical Jacobian following formula in mj_instantiateEquality
+  mjtNum jacdif[3*nv], jac0[3*nv], jac1[3*nv];
+  mjtNum point[3] = {0};
+
+  // rotational Jacobian difference
+  mj_jacDifPair(model, data, NULL, 2, 1, point, point,
+                NULL, NULL, NULL, jac0, jac1, jacdif, mj_isSparse(model),
+                /*flg_skipcommon=*/0);
+
+  // formula: 0.5 * neg(quat2) * (jac1-jac2) * quat1
+  mjtNum axis[3], quat3[4], quat4[4];
+  for (int j=0; j < nv; j++) {
+    // axis = [jac1-jac2]_col(j)
+    axis[0] = jacdif[0*nv+j];
+    axis[1] = jacdif[1*nv+j];
+    axis[2] = jacdif[2*nv+j];
+
+    // apply formula
+    mju_mulQuatAxis(quat3, negQuat2, axis);
+    mju_mulQuat(quat4, quat3, quat1);
+
+    // correct Jacobian
+    jacdif[0*nv+j] = 0.5*quat4[1];
+    jacdif[1*nv+j] = 0.5*quat4[2];
+    jacdif[2*nv+j] = 0.5*quat4[3];
+  }
+
+  // test that analytical and finite-differenced Jacobians match
+  EXPECT_THAT(AsVector(jacFD, 3*nv),
+              Pointwise(MjNear(eps, 1e-3), AsVector(jacdif, 3*nv)));
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
+}  // namespace
+}  // namespace mujoco
+
+

test/engine/engine_derivative_test.cc

@@ -1473,6 +1473,17 @@ void RotateFlexGrid(mjModel* model, mjData* data, const char* flex_name,
   }
 }
 
+// Helper: assemble flex stiffness into dense matrix via banded addH
+// This wraps the banded API and converts to dense for test verification.
+static void addH_dense(mjModel* m, mjData* d, mjtNum* H_dense,
+                       const int* dof_indices, int ndof, mjtNum h) {
+  // use full bandwidth (ndof) for exact dense equivalence
+  std::vector<mjtNum> H_band(ndof * ndof, 0);
+  mjd_flexInterp_addH(m, d, H_band.data(), dof_indices, ndof, ndof, h);
+  // convert banded to dense (lower triangle), then symmetrize
+  mju_band2Dense(H_dense, H_band.data(), ndof, ndof, 0, 1);
+}
+
 // compare analytic and fin-diff d_qfrc_passive/d_qvel for flex interp
 // Combined test for verify mjd_flexInterp_mulK (stiffness) and damping
 TEST_F(DerivativeTest, FlexInterpDerivatives) {
@@ -1525,7 +1536,7 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) {
       for (int i = 0; i < nv; i++) dof_indices[i] = i;
 
       // assemble K into H
-      mjd_flexInterp_addH(model, data, H.data(), dof_indices.data(), nv, 1.0);
+      addH_dense(model, data, H.data(), dof_indices.data(), nv, 1.0);
 
       // restore damping
       model->flex_damping[0] = save_damping;
@@ -1618,10 +1629,10 @@ TEST_F(DerivativeTest, FlexInterpDerivatives) {
       for (int i = 0; i < nv; i++) dof_indices[i] = i;
 
       vector<mjtNum> H1(nv * nv, 0);
-      mjd_flexInterp_addH(model, data, H1.data(), dof_indices.data(), nv, 1.0);
+      addH_dense(model, data, H1.data(), dof_indices.data(), nv, 1.0);
 
       vector<mjtNum> H2(nv * nv, 0);
-      mjd_flexInterp_addH(model, data, H2.data(), dof_indices.data(), nv, 0.5);
+      addH_dense(model, data, H2.data(), dof_indices.data(), nv, 0.5);
 
       vector<mjtNum> D(nv * nv);
       for (int i = 0; i < nv * nv; i++) {
@@ -1695,8 +1706,7 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
   for (int i = 0; i < nv; i++) dof_indices[i] = i;
 
   // h=1, damping=0 => adds K to H
-  mjd_flexInterp_addH(model, data, H_approx.data(), dof_indices.data(), nv,
-                      1.0);
+  addH_dense(model, data, H_approx.data(), dof_indices.data(), nv, 1.0);
 
   // 2. Compute Finite Difference Jacobian (Ground Truth)
   // qfrc_passive = -dV/dq

test/engine/engine_forward_test.cc

@@ -3107,7 +3107,7 @@ TEST_F(ForwardTest, FlexParentCoupling) {
       <body name="parent" pos="0 0 0">
         <freejoint/>
         <geom size=".1" mass="0.1"/>
-        <flexcomp name="flex" type="grid" count="3 3 3" spacing="1 1 1"
+        <flexcomp name="flex" type="grid" count="3 3 3" cellcount="1 1 1" spacing="1 1 1"
                   radius=".01" dim="3" mass="100" dof="trilinear" pos="1 1 1">
           <contact selfcollide="none"/>
           <elasticity young="1e4" poisson="0.3" damping="50"/>
@@ -3152,7 +3152,7 @@ TEST_F(ForwardTest, FlexParentCoupling) {
     if (diff > max_diff) max_diff = diff;
   }
 
-  EXPECT_LT(max_diff, MjTol(2e-5, 5e-3))
+  EXPECT_LT(max_diff, MjTol(2e-5, 1.5e-2))
       << "Implicit integrator should match Euler at small timestep";
 
   mj_deleteData(data);

test/engine/engine_support_test.cc

@@ -14,12 +14,9 @@
 
 // Tests for engine/{engine_support.c and engine_core_util.c}
 
-#include "src/engine/engine_core_util.h"
 #include "src/engine/engine_support.h"
 
-#include <algorithm>
 #include <cstring>
-#include <limits>
 #include <random>
 #include <string>
 #include <string_view>
@@ -41,521 +38,9 @@ using ::testing::Ne;
 using ::testing::NotNull;
 using ::testing::Pointwise;
 
-using AngMomMatTest = MujocoTest;
 
-static constexpr char AngMomTestingModel[] = R"(
-<mujoco>
-  <option>
-    <flag gravity="disable"/>
-  </option>
-  <worldbody>
-    <body name="link1" pos="0 0 0.5">
-      <freejoint/>
-      <geom type="ellipsoid" size="0.15 0.17 0.19" quat="1 .2 .3 .4"/>
-      <body name="link2" >
-        <joint type="hinge" axis="1 0 0" />
-        <geom type="capsule" size="0.05" fromto="0 0 0  0 0.5 0"/>
-        <body pos="0 0.6 0">
-          <joint type="slide" axis="1 0 0"/>
-          <geom type="capsule" size="0.05 0.2" quat="0.707 0 0.707 0"/>
-          <body name="link3">
-            <joint type="ball" pos="0.2 0 0"/>
-            <geom type="capsule" pos="0.2 0 0" size="0.03 0.4"/>
-          </body>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-  <keyframe>
-    <key qvel="0 0 0 .1 .2 .3 .4 .5 .4 .3 .2"/>
-  </keyframe>
-</mujoco>
-)";
 
-// compare subtree angular momentum computed in two ways
-TEST_F(AngMomMatTest, CompareAngMom) {
-  char error[1024];
-  mjModel* model =
-      LoadModelFromString(AngMomTestingModel, error, sizeof(error));
-  ASSERT_THAT(model, NotNull()) << error;
-  int nv = model->nv;
-  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
 
-  mjData* data = mj_makeData(model);
-
-  // reset to the keyframe with some angular velocities
-  mj_resetDataKeyframe(model, data, 0);
-  mj_forward(model, data);
-
-  // get the reference value of angular momentum
-  mj_subtreeVel(model, data);
-  mjtNum angmom_ref[3];
-  mju_copy3(angmom_ref, data->subtree_angmom+3*bodyid);
-
-  // compute angular momentum using the angular momentum matrix
-  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
-  mj_angmomMat(model, data, angmom_mat, bodyid);
-  mjtNum angmom_test[3];
-  mju_mulMatVec(angmom_test, angmom_mat, data->qvel, 3, nv);
-
-  // compare the two angular momentum values
-  for (int i = 0; i < 3; i++) {
-    EXPECT_THAT(angmom_ref[i], MjNear(angmom_test[i], 1e-8, 1e-4));
-  }
-
-  mju_free(angmom_mat);
-  mj_deleteData(data);
-  mj_deleteModel(model);
-}
-
-// compare subtree angular momentum matrix: analytical and findiff
-TEST_F(AngMomMatTest, CompareAngMomMats) {
-  char error[1024];
-  mjModel* model =
-      LoadModelFromString(AngMomTestingModel, error, sizeof(error));
-  ASSERT_THAT(model, NotNull()) << error;
-  int nv = model->nv;
-  int bodyid = mj_name2id(model, mjOBJ_BODY, "link1");
-  mjData* data = mj_makeData(model);
-  mjtNum* angmom_mat = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
-  mjtNum* angmom_mat_fd = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
-
-  // reset to the keyframe with some angular velocities
-  mj_resetDataKeyframe(model, data, 0);
-  mj_forward(model, data);
-
-  // compute the angular momentum matrix using the analytical method
-  mj_angmomMat(model, data, angmom_mat, bodyid);
-
-  // compute the angular momentum matrix using finite differences
-  static constexpr mjtNum eps = MjTol(1e-6, 1e-3);
-  for (int i = 0; i < nv; i++) {
-    // reset vel, forward nudge i-th dof, get angmom
-    mju_copy(data->qvel, model->key_qvel, model->nv);
-    data->qvel[i] += eps;
-    mj_forward(model, data);
-    mj_subtreeVel(model, data);
-    mjtNum agmf[3];
-    mju_copy3(agmf, data->subtree_angmom+3*bodyid);
-
-    // reset vel, backward nudge i-th dof, get angmom
-    mju_copy(data->qvel, model->key_qvel, model->nv);
-    data->qvel[i] -= eps;
-    mj_forward(model, data);
-    mj_subtreeVel(model, data);
-    mjtNum agmb[3];
-    mju_copy3(agmb, data->subtree_angmom+3*bodyid);
-
-    // finite-difference the angmom matrix
-    for (int j = 0; j < 3; j++) {
-      angmom_mat_fd[nv*j+i] = (agmf[j] - agmb[j]) / (2 * eps);
-    }
-  }
-
-  // compare the two matrices
-  for (int i = 0; i < 3*nv; i++) {
-    EXPECT_THAT(angmom_mat_fd[i], MjNear(angmom_mat[i], 1e-8, 2e-4));
-  }
-
-  mju_free(angmom_mat_fd);
-  mju_free(angmom_mat);
-  mj_deleteData(data);
-  mj_deleteModel(model);
-}
-
-using JacobianTest = MujocoTest;
-static const mjtNum max_abs_err = std::numeric_limits<float>::epsilon();
-
-static constexpr char kJacobianTestingModel[] = R"(
-<mujoco>
-  <worldbody>
-    <body name="distractor1" pos="0 0 .3">
-      <freejoint/>
-      <geom size=".1"/>
-    </body>
-    <body name="main">
-      <freejoint/>
-      <geom size=".1"/>
-      <body pos=".1 0 0">
-        <joint axis="0 1 0"/>
-        <geom type="capsule" size=".03" fromto="0 0 0 .2 0 0"/>
-      </body>
-      <body pos="0 .1 0">
-        <joint type="ball"/>
-        <geom type="capsule" size=".03" fromto="0 0 0 0 .2 0"/>
-        <body pos="0 .2 0">
-          <joint type="slide" axis="1 1 1"/>
-          <geom size=".05"/>
-        </body>
-      </body>
-    </body>
-    <body name="distractor2" pos="0 0 -.3">
-      <freejoint/>
-      <geom size=".1"/>
-    </body>
-  </worldbody>
-</mujoco>
-)";
-
-// compare analytic and finite-differenced subtree-com Jacobian
-TEST_F(JacobianTest, SubtreeJac) {
-  char error[1024];
-  mjModel* model =
-      LoadModelFromString(kJacobianTestingModel, error, sizeof(error));
-  ASSERT_THAT(model, NotNull()) << error;
-  int nv = model->nv;
-  int bodyid = mj_name2id(model, mjOBJ_BODY, "main");
-  mjData* data = mj_makeData(model);
-  mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
-  mjtNum* qpos = (mjtNum*) mju_malloc(sizeof(mjtNum)*model->nq);
-  mjtNum* nudge = (mjtNum*) mju_malloc(sizeof(mjtNum)*nv);
-
-  // all we need for Jacobians are kinematics and CoM-related quantities
-  mj_kinematics(model, data);
-  mj_comPos(model, data);
-
-  // get subtree CoM Jacobian of free body
-  mj_jacSubtreeCom(model, data, jac_subtree, bodyid);
-
-  // save current subtree-com and qpos, clear nudge
-  mjtNum subtree_com[3];
-  mju_copy3(subtree_com, data->subtree_com+3*bodyid);
-  mju_copy(qpos, data->qpos, model->nq);
-  mju_zero(nudge, nv);
-
-  // compare analytic Jacobian to finite-difference approximation
-  static const mjtNum eps = 1e-6;
-  for (int i=0; i < nv; i++) {
-    // reset qpos, nudge i-th dof, update data->qpos, reset nudge
-    mju_copy(data->qpos, qpos, model->nq);
-    nudge[i] = 1;
-    mj_integratePos(model, data->qpos, nudge, eps);
-    nudge[i] = 0;
-
-    // kinematics and comPos to get nudged com
-    mj_kinematics(model, data);
-    mj_comPos(model, data);
-
-    // compare finite-differenced and analytic Jacobian
-    for (int j=0; j < 3; j++) {
-      mjtNum findiff = (data->subtree_com[3*bodyid+j] - subtree_com[j]) / eps;
-      EXPECT_THAT(jac_subtree[nv*j+i], MjNear(findiff, eps, 1e-2));
-    }
-  }
-
-  mju_free(nudge);
-  mju_free(qpos);
-  mju_free(jac_subtree);
-  mj_deleteData(data);
-  mj_deleteModel(model);
-}
-
-// confirm that applying linear forces via the subtree-com Jacobian only creates
-// the expected linear accelerations (no accelerations of internal joints)
-TEST_F(JacobianTest, SubtreeJacNoInternalAcc) {
-  char error[1024];
-  mjModel* model =
-      LoadModelFromString(kJacobianTestingModel, error, sizeof(error));
-  ASSERT_THAT(model, NotNull()) << error;
-  int nv = model->nv;
-  int bodyid = mj_name2id(model, mjOBJ_BODY, "main");
-  mjData* data = mj_makeData(model);
-  mjtNum* jac_subtree = (mjtNum*) mju_malloc(sizeof(mjtNum)*3*nv);
-
-  // all we need for Jacobians are kinematics and CoM-related quantities
-  mj_kinematics(model, data);
-  mj_comPos(model, data);
-
-  // get subtree CoM Jacobian of free body
-  mj_jacSubtreeCom(model, data, jac_subtree, bodyid);
-
-  // uncomment for debugging
-  // mju_printMat(jac_subtree, 3, nv);
-
-  // call fwdPosition since we'll need the factorised mass matrix in the test
-  mj_fwdPosition(model, data);
-
-  // treating the subtree Jacobian as the projection of 3 axis-aligned unit
-  // forces into joint space, solve for the resulting accelerations in-place
-  mj_solveM(model, data, jac_subtree, jac_subtree, 3);
-
-  // expect to find accelerations of magnitude 1/subtreemass in the first 3
-  // coordinates of the free joint and 0s elsewhere, since applying forces to
-  // the CoM should accelerate the whole mechanism without any internal motion
-  int body_dofadr = model->body_dofadr[bodyid];
-  mjtNum invtreemass = 1.0/model->body_subtreemass[bodyid];
-  for (int r = 0; r < 3; r++) {
-    for (int c = 0; c < nv; c++) {
-      mjtNum expected = c - body_dofadr == r ? invtreemass : 0.0;
-      EXPECT_THAT(jac_subtree[nv*r+c], MjNear(expected, max_abs_err, 1e-4));
-    }
-  }
-
-  mju_free(jac_subtree);
-  mj_deleteData(data);
-  mj_deleteModel(model);
-}
-
-static constexpr char kQuat[] = R"(
-<mujoco>
-  <worldbody>
-    <body name="query">
-      <joint type="ball"/>
-      <geom size="1"/>
-      <site name="query" pos=".1 .2 .3"/>
-    </body>
-  </worldbody>
-  <keyframe>
-    <key qvel="2 3 5"/>
-  </keyframe>
-</mujoco>
-)";
-
-static constexpr char kFreeBall[] = R"(
-<mujoco>
-  <worldbody>
-    <body name="distractor1" pos="0 0 .3">
-      <freejoint/>
-      <geom size=".1"/>
-    </body>
-    <body name="main">
-      <freejoint/>
-      <geom size=".1"/>
-      <body pos=".1 0 0">
-        <joint axis="0 1 0"/>
-        <geom type="capsule" size=".03" fromto="0 0 0 .2 0 0"/>
-        <body pos=".2 0 0">
-          <joint type="ball" stiffness="20"/>
-          <geom type="capsule" size=".03" fromto="0 0 0 0 .2 0"/>
-          <body name="query" pos="0 .2 0">
-            <joint type="slide" axis="1 1 1"/>
-            <geom size=".05"/>
-            <site name="query" pos=".1 .2 .3"/>
-          </body>
-        </body>
-      </body>
-    </body>
-    <body name="distractor2" pos="0 0 -.3">
-      <freejoint/>
-      <geom size=".1"/>
-    </body>
-  </worldbody>
-  <keyframe>
-    <key qvel="1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1"/>
-  </keyframe>
-</mujoco>
-)";
-
-static constexpr char kQuatlessPendulum[] = R"(
-<mujoco>
-  <option integrator="implicit">
-    <flag constraint="disable"/>
-  </option>
-  <worldbody>
-    <body pos="0.15 0 0">
-      <joint type="hinge" axis="0 1 0"/>
-      <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
-      <body pos="0.1 0 0">
-        <joint type="slide" axis="1 0 0" stiffness="200"/>
-        <geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
-        <body pos=".1 0 0">
-          <joint axis="1 0 0"/>
-          <joint axis="0 1 0"/>
-          <joint axis="0 0 1"/>
-          <geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
-          <body name="query" pos="0 .1 0">
-            <joint axis="1 0 0"/>
-            <geom type="capsule" size="0.02" fromto="0 0 0 0 .1 0"/>
-            <site name="query" pos=".1 0 0"/>
-          </body>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-</mujoco>
-)";
-
-static constexpr char kTelescope[] = R"(
-<mujoco>
-  <worldbody>
-    <body>
-      <joint type="ball"/>
-      <geom type="capsule" size="0.02" fromto="0 .02 0 .1 .02 0"/>
-      <body pos=".1 .02 0">
-        <joint type="slide" axis="1 0 0"/>
-        <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
-        <body pos=".1 .02 0">
-          <joint type="slide" axis="1 0 0"/>
-          <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
-          <body pos=".1 .02 0" name="query">
-            <joint type="slide" axis="1 0 0"/>
-            <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
-            <site name="query" pos=".1 0 0"/>
-          </body>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-  <keyframe>
-    <key qvel="1 1 1 1 1 1"/>
-  </keyframe>
-</mujoco>
-)";
-
-static constexpr char kHinge[] = R"(
-<mujoco>
-  <worldbody>
-    <body name="query">
-      <joint name="link1" axis="0 1 0"/>
-      <geom type="capsule" size=".02" fromto="0 0 0 0 0 -1"/>
-      <site name="query" pos="0 0 -1"/>
-    </body>
-  </worldbody>
-
-  <keyframe>
-    <key qpos="1" qvel="1"/>
-  </keyframe>
-</mujoco>
-)";
-
-// compare mj_jacDot with finite-differenced mj_jac
-TEST_F(JacobianTest, JacDot) {
-  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
-    char error[1024];
-    mjModel* model = LoadModelFromString(xml, error, sizeof(error));
-    ASSERT_THAT(model, NotNull()) << error;
-    int nv = model->nv;
-    mjData* data = mj_makeData(model);
-
-    // load keyframe if present, step for a bit
-    if (model->nkey) mj_resetDataKeyframe(model, data, 0);
-    while (data->time < 0.1) {
-      mj_step(model, data);
-    }
-
-    // minimal call required for mj_jacDot outputs to be valid
-    mj_kinematics(model, data);
-    mj_comPos(model, data);
-    mj_comVel(model, data);
-
-    // get bodyid
-    int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
-    EXPECT_GT(bodyid, 0);
-
-    // get site position
-    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
-    EXPECT_GT(siteid, -1);
-    mjtNum point[3];
-    mju_copy3(point, data->site_xpos+3*siteid);
-
-    // jac, jac_dot
-    vector<mjtNum> jacp(3*nv);
-    vector<mjtNum> jacr(3*nv);
-    mj_jac(model, data, jacp.data(), jacr.data(), point, bodyid);
-    vector<mjtNum> jacp_dot(3*nv);
-    vector<mjtNum> jacr_dot(3*nv);
-    mj_jacDot(model, data, jacp_dot.data(), jacr_dot.data(), point, bodyid);
-
-    // jac_h: jacobian after integrating qpos with a timestep of h
-    constexpr mjtNum h = MjTol(1e-7, 5e-4);
-    mj_integratePos(model, data->qpos, data->qvel, h);
-    mj_kinematics(model, data);
-    mj_comPos(model, data);
-    vector<mjtNum> jacp_h(3*nv);
-    vector<mjtNum> jacr_h(3*nv);
-    mju_copy3(point, data->site_xpos+3*siteid);  // get updated site position
-    mj_jac(model, data, jacp_h.data(), jacr_h.data(), point, bodyid);
-
-    // jac_dot_h finite-difference approximation
-    vector<mjtNum> jacp_dot_h(3*nv);
-    mju_sub(jacp_dot_h.data(), jacp_h.data(), jacp.data(), 3*nv);
-    mju_scl(jacp_dot_h.data(), jacp_dot_h.data(), 1/h, 3*nv);
-    vector<mjtNum> jacr_dot_h(3*nv);
-    mju_sub(jacr_dot_h.data(), jacr_h.data(), jacr.data(), 3*nv);
-    mju_scl(jacr_dot_h.data(), jacr_dot_h.data(), 1/h, 3*nv);
-
-    // compare finite-differenced and analytic
-    mjtNum tol = 1e-5;
-    EXPECT_THAT(jacp_dot, Pointwise(MjNear(tol, 5e-2), jacp_dot_h));
-    EXPECT_THAT(jacr_dot, Pointwise(MjNear(tol, 5e-2), jacr_dot_h));
-
-    mj_deleteData(data);
-    mj_deleteModel(model);
-  }
-}
-
-// compare mj_jacDotSparse with dense mj_jacDot
-TEST_F(JacobianTest, JacDotSparse) {
-  for (auto xml : {kHinge, kQuat, kTelescope, kFreeBall, kQuatlessPendulum}) {
-    char error[1024];
-    mjModel* model = LoadModelFromString(xml, error, sizeof(error));
-    ASSERT_THAT(model, NotNull()) << error;
-    int nv = model->nv;
-    mjData* data = mj_makeData(model);
-
-    // load keyframe if present, step for a bit
-    if (model->nkey) mj_resetDataKeyframe(model, data, 0);
-    while (data->time < 0.1) {
-      mj_step(model, data);
-    }
-
-    // minimal call required for mj_jacDot outputs to be valid
-    mj_kinematics(model, data);
-    mj_comPos(model, data);
-    mj_comVel(model, data);
-
-    // get bodyid and site position
-    int bodyid = mj_name2id(model, mjOBJ_BODY, "query");
-    EXPECT_GT(bodyid, 0);
-    int siteid = mj_name2id(model, mjOBJ_SITE, "query");
-    EXPECT_GT(siteid, -1);
-    mjtNum point[3];
-    mju_copy3(point, data->site_xpos+3*siteid);
-
-    // dense jacDot
-    vector<mjtNum> jacp_dense(3*nv);
-    vector<mjtNum> jacr_dense(3*nv);
-    mj_jacDot(model, data, jacp_dense.data(), jacr_dense.data(), point, bodyid);
-
-    // compute body chain using public mjModel fields
-    vector<int> chain(nv);
-    int NV = 0;
-    int weldbody = model->body_weldid[bodyid];
-    if (weldbody) {
-      int da = model->body_dofadr[weldbody] + model->body_dofnum[weldbody] - 1;
-      while (da >= 0) {
-        chain[NV++] = da;
-        da = model->dof_parentid[da];
-      }
-      std::reverse(chain.begin(), chain.begin() + NV);
-    }
-    EXPECT_GT(NV, 0);
-
-    // sparse jacDot
-    vector<mjtNum> jacp_sparse(3*NV);
-    vector<mjtNum> jacr_sparse(3*NV);
-    mj_jacDotSparse(model, data, jacp_sparse.data(), jacr_sparse.data(),
-                    point, bodyid, NV, chain.data());
-
-    // expand sparse to dense and compare
-    vector<mjtNum> jacp_expanded(3*nv, 0);
-    vector<mjtNum> jacr_expanded(3*nv, 0);
-    for (int ci = 0; ci < NV; ci++) {
-      int di = chain[ci];
-      for (int r = 0; r < 3; r++) {
-        jacp_expanded[di+r*nv] = jacp_sparse[ci+r*NV];
-        jacr_expanded[di+r*nv] = jacr_sparse[ci+r*NV];
-      }
-    }
-
-    // expect bitwise equality
-    EXPECT_EQ(jacp_expanded, jacp_dense);
-    EXPECT_EQ(jacr_expanded, jacr_dense);
-
-    mj_deleteData(data);
-    mj_deleteModel(model);
-  }
-}
 
 using Name2idTest = MujocoTest;
 

test/engine/engine_util_misc_test.cc

@@ -430,13 +430,90 @@ TEST_F(InterpolationTest, mju_interpolate3D) {
     expected[0] = quadratic_function_1(sample[0], sample[1], sample[2]);
     expected[1] = quadratic_function_2(sample[0], sample[1], sample[2]);
     expected[2] = quadratic_function_3(sample[0], sample[1], sample[2]);
-    mju_interpolate3D(res, sample, coeff, order);
+    mju_interpolate3D(res, sample, coeff, order, NULL);
     EXPECT_NEAR(res[0], expected[0], MjTol(1e-10, 1e-5));
     EXPECT_NEAR(res[1], expected[1], MjTol(1e-10, 1e-5));
     EXPECT_NEAR(res[2], expected[2], MjTol(1e-10, 1e-5));
   }
 }
 
+TEST_F(InterpolationTest, mju_cellLookup_SingleCell) {
+  // single cell (1x1x1): local coords should equal global coords
+  int cellnum[3] = {1, 1, 1};
+  mjtNum coord[3] = {0.3, 0.7, 0.5};
+  mjtNum local[3];
+  int nodeindices[8];
+
+  int npc = mju_cellLookup(coord, cellnum, 1, local, nodeindices);
+  EXPECT_EQ(npc, 8);
+  EXPECT_NEAR(local[0], 0.3, MjTol(1e-12, 1e-6));
+  EXPECT_NEAR(local[1], 0.7, MjTol(1e-12, 1e-6));
+  EXPECT_NEAR(local[2], 0.5, MjTol(1e-12, 1e-6));
+
+  // for trilinear 1x1x1: nodes are 0..7 in lexicographic order
+  for (int i = 0; i < 8; i++) {
+    EXPECT_EQ(nodeindices[i], i);
+  }
+}
+
+TEST_F(InterpolationTest, mju_cellLookup_MultiCell) {
+  // 2x3x4 grid, trilinear: 3x4x5 = 60 nodes
+  int cellnum[3] = {2, 3, 4};
+  int order = 1;
+  int ny_g = 3*1 + 1;  // 4
+  int nz_g = 4*1 + 1;  // 5
+
+  // point at (0.75, 0.5, 0.125) -> cell (1, 1, 0)
+  mjtNum coord[3] = {0.75, 0.5, 0.125};
+  mjtNum local[3];
+  int nodeindices[8];
+
+  int npc = mju_cellLookup(coord, cellnum, order, local, nodeindices);
+  EXPECT_EQ(npc, 8);
+
+  // cell (1,1,0): local = (0.75*2 - 1, 0.5*3 - 1, 0.125*4 - 0)
+  EXPECT_NEAR(local[0], 0.5, 1e-12);
+  EXPECT_NEAR(local[1], 0.5, 1e-12);
+  EXPECT_NEAR(local[2], 0.5, 1e-12);
+
+  // expected node indices for cell (1,1,0), trilinear:
+  //   (gi, gj, gk) for li,lj,lk in {0,1}
+  //   gi = 1+li, gj = 1+lj, gk = 0+lk
+  //   gidx = gi*ny_g*nz_g + gj*nz_g + gk
+  int expected[8];
+  int ni = 0;
+  for (int li = 0; li <= 1; li++) {
+    for (int lj = 0; lj <= 1; lj++) {
+      for (int lk = 0; lk <= 1; lk++) {
+        expected[ni++] = (1+li)*ny_g*nz_g + (1+lj)*nz_g + lk;
+      }
+    }
+  }
+  for (int i = 0; i < 8; i++) {
+    EXPECT_EQ(nodeindices[i], expected[i]);
+  }
+}
+
+TEST_F(InterpolationTest, mju_cellLookup_Boundary) {
+  // point exactly at coord=1.0 should clamp to last cell
+  int cellnum[3] = {3, 3, 3};
+  mjtNum coord[3] = {1.0, 1.0, 1.0};
+  mjtNum local[3];
+
+  mju_cellLookup(coord, cellnum, 1, local, NULL);
+  // cell (2,2,2), local = (1*3 - 2, 1*3 - 2, 1*3 - 2) = (1, 1, 1)
+  EXPECT_NEAR(local[0], 1.0, 1e-12);
+  EXPECT_NEAR(local[1], 1.0, 1e-12);
+  EXPECT_NEAR(local[2], 1.0, 1e-12);
+
+  // point at coord=0.0 should map to first cell
+  mjtNum coord0[3] = {0.0, 0.0, 0.0};
+  mju_cellLookup(coord0, cellnum, 1, local, NULL);
+  EXPECT_NEAR(local[0], 0.0, 1e-12);
+  EXPECT_NEAR(local[1], 0.0, 1e-12);
+  EXPECT_NEAR(local[2], 0.0, 1e-12);
+}
+
 TEST_F(InterpolationTest, mju_defGradient) {
   int order = 1;
   mjtNum mat[9];
@@ -521,7 +598,48 @@ TEST_F(InterpolationTest, mju_defGradient) {
   EXPECT_THAT(mat, Pointwise(MjNear(1e-8, 1e-6), rot7));
 }
 
-// --------------------------------- Base64 ------------------------------------
+TEST_F(InterpolationTest, mju_flexInterpState_MultiCell) {
+  int order = 1;  // trilinear
+  int cy = 2;
+  int cz = 2;
+  int nodenum = 27;  // 3x3x3
+
+  std::vector<mjtNum> xpos(3 * nodenum);
+  mjtNum quat[4];
+
+  // Populate xpos directly for a grid centered at origin, rotated 90 deg around
+  // Z Original grid points: {-0.1, 0.0, 0.1}^3 Rotated: (x, y, z) -> (-y, x, z)
+  int idx = 0;
+  for (int i = 0; i < 3; i++) {
+    for (int j = 0; j < 3; j++) {
+      for (int k = 0; k < 3; k++) {
+        mjtNum x = (i - 1) * 0.1;
+        mjtNum y = (j - 1) * 0.1;
+        mjtNum z = (k - 1) * 0.1;
+
+        // Apply rotation
+        xpos[3*idx + 0] = -y;
+        xpos[3*idx + 1] = x;
+        xpos[3*idx + 2] = z;
+        idx++;
+      }
+    }
+  }
+
+  int npc = (order+1)*(order+1)*(order+1);
+  std::vector<mjtNum> xpos_c(3 * npc);
+
+  mju_flexGatherCellState(order, cy, cz, 0, 0, 0, xpos.data(), NULL, NULL,
+                          xpos_c.data(), NULL, NULL, NULL, quat);
+
+  // Expected quaternion for -90 deg around Z (global to local):
+  // [sqrt(0.5), 0, 0, -sqrt(0.5)]
+  mjtNum expected_val = mju_sqrt(0.5);
+  EXPECT_NEAR(quat[0], expected_val, 1e-5);
+  EXPECT_NEAR(quat[1], 0.0, 1e-5);
+  EXPECT_NEAR(quat[2], 0.0, 1e-5);
+  EXPECT_NEAR(quat[3], -expected_val, 1e-5);
+}
 
 using Base64Test = MujocoTest;
 

test/user/user_flex_test.cc

@@ -79,6 +79,27 @@ TEST_F(UserFlexTest, CountTooSmall) {
   EXPECT_THAT(error.data(), HasSubstr("Count too small"));
 }
 
+TEST_F(UserFlexTest, CellnumZeroInterpolated) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+  <worldbody>
+    <body name="b0"/>
+    <body name="b1"/>
+    <body name="b2"/>
+    <body name="b3"/>
+  </worldbody>
+  <deformable>
+    <flex name="test" cellcount="2 2 0" dof="trilinear"
+          dim="3" body="b0 b1 b2 b3" element="0 1 2 3"/>
+  </deformable>
+  </mujoco>
+  )";
+  std::array<char, 1024> error;
+  mjModel* m = LoadModelFromString(xml, error.data(), error.size());
+  EXPECT_THAT(m, IsNull());
+  EXPECT_THAT(error.data(), HasSubstr("cellcount cannot be 0"));
+}
+
 TEST_F(UserFlexTest, SpacingGreaterThanGeometry) {
   static constexpr char xml[] = R"(
   <mujoco>

unity/Runtime/Bindings/MjBindings.cs

@@ -1200,6 +1200,8 @@ public unsafe struct mjModel_ {
   public int* flex_matid;
   public int* flex_group;
   public int* flex_interp;
+  public int* flex_bandwidth;
+  public int* flex_cellnum;
   public int* flex_nodeadr;
   public int* flex_nodenum;
   public int* flex_vertadr;

wasm/CMakeLists.txt

@@ -22,16 +22,14 @@ include_directories(${PROJECT_SOURCE_DIR})
 
 link_directories(${CMAKE_BINARY_DIR}/lib)
 
-file(GLOB MUJOCO_WASM_FILES
-    "codegen/generated/*.cc"
-    "unpack.cc"
-)
+# Set Emscripten compile flags.
+# -fexceptions is required for val::throw_() (used by ThrowMujocoErrorToJS) to
+# actually throw a JS exception. Without it, Emscripten compiles throw as a
+# no-op. In the MT build, -pthread implicitly enables exception support, but
+# the ST build needs it set explicitly here.
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fexceptions")
 
-if(NOT MUJOCO_WASM_FILES)
-    message(FATAL_ERROR "No source files found in codegen/generated/")
-endif()
-
-# Set Emscripten linker flags
+# Set Emscripten linker flags shared by all WASM targets.
 set(EMCC_LINKER_FLAGS
     "--bind"
     "-s ASSERTIONS=1"
@@ -42,9 +40,9 @@ set(EMCC_LINKER_FLAGS
     "-s EXPORTED_RUNTIME_METHODS=['ccall','cwrap','FS','MEMFS']"
     "-s EXPORT_NAME=loadMujoco"
     "-s DISABLE_EXCEPTION_CATCHING=0"
+    "-fexceptions"
     "-gsource-map"
     "-g"
-    "--emit-tsd mujoco.d.ts"
 )
 if(MUJOCO_WASM_THREADS)
     list(APPEND EMCC_LINKER_FLAGS
@@ -53,7 +51,28 @@ if(MUJOCO_WASM_THREADS)
     )
     add_definitions(-DMUJOCO_WASM_THREADS)
 endif()
-string (REPLACE ";" " " EMCC_LINKER_FLAGS_STR "${EMCC_LINKER_FLAGS}")
+
+# Common link libraries for all WASM targets. The mujoco library is linked as a
+# whole archive to avoid losing plugin registration such as obj_decoder and
+# stl_decoder.
+set(MUJOCO_WASM_LINK_LIBRARIES
+    -Wl,--whole-archive mujoco -Wl,--no-whole-archive
+    ccd lodepng tinyxml2 qhullstatic_r
+)
+
+# --- Main WASM bindings target ---
+
+file(GLOB MUJOCO_WASM_FILES
+    "codegen/generated/*.cc"
+    "unpack.cc"
+)
+
+if(NOT MUJOCO_WASM_FILES)
+    message(FATAL_ERROR "No source files found in codegen/generated/")
+endif()
+
+set(MUJOCO_WASM_LINKER_FLAGS ${EMCC_LINKER_FLAGS} "--emit-tsd mujoco.d.ts")
+string(REPLACE ";" " " MUJOCO_WASM_LINKER_FLAGS_STR "${MUJOCO_WASM_LINKER_FLAGS}")
 
 add_executable(mujoco_wasm ${MUJOCO_WASM_FILES})
 
@@ -61,15 +80,36 @@ add_executable(mujoco_wasm ${MUJOCO_WASM_FILES})
 # `mujoco` library target, but emit artifacts named `mujoco.*` by setting the
 # output name. Also apply the emscripten linker flags to the wasm target.
 set_target_properties(mujoco_wasm PROPERTIES
-    LINK_FLAGS "${EMCC_LINKER_FLAGS_STR}"
+    LINK_FLAGS "${MUJOCO_WASM_LINKER_FLAGS_STR}"
     OUTPUT_NAME "mujoco"
 )
 
-# Link the mujoco library as a whole archive to avoid losing plugin
-# registration such as obj_decoder and stl_decoder.
-target_link_libraries(mujoco_wasm PRIVATE
-    -Wl,--whole-archive mujoco -Wl,--no-whole-archive
-    ccd lodepng tinyxml2 qhullstatic_r
-)
+target_link_libraries(mujoco_wasm PRIVATE ${MUJOCO_WASM_LINK_LIBRARIES})
 
 install(TARGETS mujoco_wasm DESTINATION ${DIVISIBLE_INSTALL_BIN_DIR})
+
+# --- Benchmark target ---
+
+if(MUJOCO_BUILD_TESTS_WASM)
+    file(GLOB MUJOCO_WASM_BENCHMARK_FILES
+        "tests/benchmark_test.cc"
+        "unpack.cc"
+    )
+
+    if(NOT MUJOCO_WASM_BENCHMARK_FILES)
+        message(FATAL_ERROR "No benchmark source files found")
+    endif()
+
+    set(BENCHMARK_LINKER_FLAGS ${EMCC_LINKER_FLAGS} "--emit-tsd mujoco_wasm_benchmark.d.ts")
+    string(REPLACE ";" " " BENCHMARK_LINKER_FLAGS_STR "${BENCHMARK_LINKER_FLAGS}")
+
+    add_executable(mujoco_wasm_benchmark ${MUJOCO_WASM_BENCHMARK_FILES})
+
+    set_target_properties(mujoco_wasm_benchmark PROPERTIES
+        LINK_FLAGS "${BENCHMARK_LINKER_FLAGS_STR}"
+    )
+
+    target_link_libraries(mujoco_wasm_benchmark PRIVATE ${MUJOCO_WASM_LINK_LIBRARIES})
+
+    install(TARGETS mujoco_wasm_benchmark DESTINATION ${DIVISIBLE_INSTALL_BIN_DIR})
+endif()

wasm/codegen/generated/bindings.cc

@@ -2446,6 +2446,15 @@ struct MjsFlex {
   void set_elastic2d(int value) {
     ptr_->elastic2d = value;
   }
+  emscripten::val cellcount() const {
+    return emscripten::val(emscripten::typed_memory_view(3, ptr_->cellcount));
+  }
+  int order() const {
+    return ptr_->order;
+  }
+  void set_order(int value) {
+    ptr_->order = value;
+  }
   mjStringVec &nodebody() const {
     return *(ptr_->nodebody);
   }
@@ -4640,6 +4649,12 @@ struct MjModel {
   emscripten::val flex_interp() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_interp));
   }
+  emscripten::val flex_bandwidth() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_bandwidth));
+  }
+  emscripten::val flex_cellnum() const {
+    return emscripten::val(emscripten::typed_memory_view(ptr_->nflex * 3, ptr_->flex_cellnum));
+  }
   emscripten::val flex_nodeadr() const {
     return emscripten::val(emscripten::typed_memory_view(ptr_->nflex, ptr_->flex_nodeadr));
   }
@@ -11803,9 +11818,11 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("eq_type", &MjModel::eq_type)
     .property("exclude_signature", &MjModel::exclude_signature)
     .property("flex_activelayers", &MjModel::flex_activelayers)
+    .property("flex_bandwidth", &MjModel::flex_bandwidth)
     .property("flex_bending", &MjModel::flex_bending)
     .property("flex_bvhadr", &MjModel::flex_bvhadr)
     .property("flex_bvhnum", &MjModel::flex_bvhnum)
+    .property("flex_cellnum", &MjModel::flex_cellnum)
     .property("flex_centered", &MjModel::flex_centered)
     .property("flex_conaffinity", &MjModel::flex_conaffinity)
     .property("flex_condim", &MjModel::flex_condim)
@@ -12569,6 +12586,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("info", &MjsExclude::info, &MjsExclude::set_info, reference());
   emscripten::class_<MjsFlex>("MjsFlex")
     .property("activelayers", &MjsFlex::activelayers, &MjsFlex::set_activelayers, reference())
+    .property("cellcount", &MjsFlex::cellcount)
     .property("conaffinity", &MjsFlex::conaffinity, &MjsFlex::set_conaffinity, reference())
     .property("condim", &MjsFlex::condim, &MjsFlex::set_condim, reference())
     .property("contype", &MjsFlex::contype, &MjsFlex::set_contype, reference())
@@ -12590,6 +12608,7 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("material", &MjsFlex::material, &MjsFlex::set_material, reference())
     .property("node", &MjsFlex::node, reference())
     .property("nodebody", &MjsFlex::nodebody, reference())
+    .property("order", &MjsFlex::order, &MjsFlex::set_order, reference())
     .property("passive", &MjsFlex::passive, &MjsFlex::set_passive, reference())
     .property("poisson", &MjsFlex::poisson, &MjsFlex::set_poisson, reference())
     .property("priority", &MjsFlex::priority, &MjsFlex::set_priority, reference())

wasm/tests/CMakeLists.txt

@@ -1,63 +0,0 @@
-# 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.
-
-set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "${CMAKE_SOURCE_DIR}/wasm/dist")
-
-set(CMAKE_INSTALL_PREFIX ${PROJECT_SOURCE_DIR}/wasm)
-
-include_directories(${PROJECT_SOURCE_DIR}/include)
-include_directories(${PROJECT_SOURCE_DIR}/src)
-include_directories(${PROJECT_SOURCE_DIR})
-
-link_directories(${CMAKE_BINARY_DIR}/lib)
-
-file(GLOB MUJOCO_WASM_FILES
-    "benchmark_test.cc"
-    "../unpack.cc"
-)
-
-if(NOT MUJOCO_WASM_FILES)
-    message(FATAL_ERROR "No source files found")
-endif()
-
-# Set Emscripten linker flags
-set(EMCC_LINKER_FLAGS
-    "--bind"
-    "-pthread"
-    "-s PTHREAD_POOL_SIZE=navigator.hardwareConcurrency"
-    "-s ASSERTIONS=1"
-    "-s ALLOW_MEMORY_GROWTH=1"
-    "-s EXPORT_ES6=1"
-    "-s MODULARIZE=1"
-    "-s FORCE_FILESYSTEM=1"
-    "-s EXPORTED_RUNTIME_METHODS=['ccall','cwrap','FS','MEMFS']"
-    "-s EXPORT_NAME=loadMujoco"
-    "-gsource-map"
-    "-g"
-    "--emit-tsd mujoco_wasm_benchmark.d.ts"
-)
-string (REPLACE ";" " " EMCC_LINKER_FLAGS_STR "${EMCC_LINKER_FLAGS}")
-
-add_executable(mujoco_wasm_benchmark ${MUJOCO_WASM_FILES})
-
-set_target_properties(mujoco_wasm_benchmark PROPERTIES LINK_FLAGS "${EMCC_LINKER_FLAGS_STR}")
-
-# Link the mujoco library as a whole archive to avoid losing plugin
-# registration such as obj_decoder and stl_decoder.
-target_link_libraries(mujoco_wasm_benchmark PRIVATE
-    -Wl,--whole-archive mujoco -Wl,--no-whole-archive
-    ccd lodepng tinyxml2 qhullstatic_r
-)
-
-install(TARGETS mujoco_wasm_benchmark DESTINATION ${DIVISIBLE_INSTALL_BIN_DIR})