Merge branch 'google-deepmind:main' into usd-integration

.github/workflows/build.yml

@@ -164,6 +164,7 @@ jobs:
       if: ${{ runner.os != 'Windows' }}
       working-directory: build
       run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
+           cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
@@ -171,6 +172,7 @@ jobs:
       if: ${{ runner.os == 'Windows' }}
       working-directory: build
       run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
+           cp bin/Release/actuator.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp bin/Release/elasticity.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp bin/Release/sensor.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
     - name: Configure samples

CMakeLists.txt

@@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco
-  VERSION 3.0.2
+  VERSION 3.1.2
   DESCRIPTION "MuJoCo Physics Simulator"
   HOMEPAGE_URL "https://mujoco.org"
 )
@@ -83,6 +83,7 @@ target_include_directories(
 )
 
 add_subdirectory(plugin/elasticity)
+add_subdirectory(plugin/actuator)
 add_subdirectory(plugin/sensor)
 add_subdirectory(plugin/sdf)
 add_subdirectory(src/engine)

README.md

@@ -164,7 +164,7 @@ These packages give users of various languages access to MuJoCo functionality:
   by [Manoj Velmurugan](https://github.com/vmanoj1996).
 - **Swift**: [swift-mujoco](https://github.com/liuliu/swift-mujoco)
 - **Java**: [mujoco-java](https://github.com/CommonWealthRobotics/mujoco-java)
-- **Julia**: [Lyceum](https://github.com/Lyceum/MuJoCo.jl) (unmaintained)
+- **Julia**: [MuJoCo.jl](https://github.com/JamieMair/MuJoCo.jl)
 
 
 ### Converters

cmake/MujocoDependencies.cmake

@@ -39,7 +39,7 @@ set(MUJOCO_DEP_VERSION_qhull
     CACHE STRING "Version of `qhull` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_Eigen3
-    aa6964bf3a34fd607837dd8123bc42465185c4f8
+    454f89af9d6f3525b1df5f9ef9c86df58bf2d4d3
     CACHE STRING "Version of `Eigen3` to be fetched."
 )
 
@@ -54,7 +54,7 @@ set(MUJOCO_DEP_VERSION_gtest
 )
 
 set(MUJOCO_DEP_VERSION_benchmark
-    344117638c8ff7e239044fd0fa7085839fc03021 # v1.8.3
+    e45585a4b8e75c28479fa4107182c28172799640 # v1.8.3
     CACHE STRING "Version of `benchmark` to be fetched."
 )
 

dist/mujoco.rc

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

dist/simulate.rc

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

doc/APIreference/APIglobals.rst

@@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - Maximum number of UI rectangles.
        Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
    * - ``mjVERSION_HEADER``
-     - 302
+     - 312
      - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software
        version, so 210 corresponds to version 2.1. Defined in  mujoco.h. The API function :ref:`mj_version` returns a
        number with the same meaning but for the compiled library.

doc/APIreference/APItypes.rst

@@ -1219,9 +1219,9 @@ a frame at the center-of-mass of the local kinematic subtree (``mjData.subtree_c
 This choice increases the precision of kinematic computations for mechanisms that are distant from the global origin.
 
 ``cdof``:
-  These 6D motion vectors describe the instantaneous axis of a degree-of-freedom and are used by all Jacobian functions.
+  These 6D motion vectors (3 rotation, 3 translation) describe the instantaneous axis of a degree-of-freedom and are
-  Therefore, the minimal computation required for analytic Jacobians is :ref:`mj_kinematics` followed by
+  used by all Jacobian functions. The minimal computation required for analytic Jacobians is :ref:`mj_kinematics`
-  :ref:`mj_comPos`.
+  followed by :ref:`mj_comPos`.
 
 ``cinert``:
   These 10-vectors describe the inertial properties of a body in the c-frame and are used by the Composite Rigid Body

doc/APIreference/functions.rst

@@ -583,6 +583,17 @@ mj_RungeKutta
 
 Runge-Kutta explicit order-N integrator.
 
+.. _mj_implicit:
+
+mj_implicit
+~~~~~~~~~~~
+
+.. mujoco-include:: mj_implicit
+
+Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see
+:ref:`Numerical Integration<geIntegration>`), and advances simulation time. See `mjdata.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__ for fields computed by this function.
+
 .. _mj_invPosition:
 
 mj_invPosition
@@ -2974,7 +2985,7 @@ mju_eig3
 
 .. mujoco-include:: mju_eig3
 
-Eigenvalue decomposition of symmetric 3x3 matrix.
+Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.
 
 .. _mju_boxQP:
 
@@ -3011,10 +3022,10 @@ outputs (optional):
 
 notes:
   The initial value of ``res`` is used to warmstart the solver.
-  ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output.
+  ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output.
-  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output.
+  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output.
   The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures.
-  Only the lower triangles of H and R and are read from and written to, respectively.
+  Only the lower triangles of H and R are read from and written to, respectively.
 
 .. _mju_boxQPmalloc:
 

doc/APIreference/functions_override.rst

@@ -69,6 +69,12 @@ These functions can be used to print various quantities to the screen for debugg
 These are components of the simulation pipeline, called internally from :ref:`mj_step`, :ref:`mj_forward` and
 :ref:`mj_inverse`. It is unlikely that the user will need to call them.
 
+.. _mj_implicit:
+
+Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see
+:ref:`Numerical Integration<geIntegration>`), and advances simulation time. See `mjdata.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__ for fields computed by this function.
+
 .. _Subcomponents:
 
 These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
@@ -486,10 +492,10 @@ outputs (optional):
 
 notes:
   The initial value of ``res`` is used to warmstart the solver.
-  ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output.
+  ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output.
-  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output.
+  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output.
   The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures.
-  Only the lower triangles of H and R and are read from and written to, respectively.
+  Only the lower triangles of H and R are read from and written to, respectively.
 
 .. _mju_boxQPmalloc:
 

doc/XMLreference.rst

@@ -103,17 +103,76 @@ In the remainder of this chapter we describe all valid MJCF elements and their a
 multiple contexts, in which case their meaning depends on the parent element. This is why we always show the parent as a
 prefix in the documentation below.
 
+.. _meta-element:
+
+Meta elements
+~~~~~~~~~~~~~
+
+These elements are not strictly part of the low-level MJCF format definition, but rather instruct the compiler to
+perform some operation on the model. A general property of meta-elements is that they disappear from the model upon
+saving the XML. There are currently four meta-elements in MJCF:
+
+- :ref:`include<include>` and :ref:`frame<frame>`, which are outside of the schema.
+- :ref:`composite<body-composite>` and :ref:`flexcomp<body-flexcomp>` which are part of the schema, but serve to
+  procedurally generate other MJCF elements.
+
+.. _frame:
+
+**frame** (R)
+^^^^^^^^^^^^^
+
+The frame meta-element is a pure coordinate transformation that can wrap any group of elements in the kinematic tree
+(under :ref:`worldbody<body>`). After compilation, frame elements disappear and their transformation is accumulated
+in their direct children. The attributes of the frame meta-element are documented :ref:`below<body-frame>`.
+
+.. collapse:: Usage example of frame
+
+   Loading this model and saving it:
+
+   .. code-block:: xml
+
+      <mujoco>
+        <worldbody>
+          <frame quat="0 0 1 0">
+             <geom name="Alice" quat="0 1 0 0" size="1"/>
+          </frame>
+
+          <frame pos="0 1 0">
+            <geom name="Bob" pos="0 1 0" size="1"/>
+            <body name="Carl" pos="1 0 0">
+              ...
+            </body>
+          </frame>
+        </worldbody>
+      </mujoco>
+
+   Results in this model:
+
+   .. code-block:: xml
+
+      <mujoco>
+        <worldbody>
+          <geom name="Alice" quat="0 0 0 1" size="1"/>
+          <geom name="Bob" pos="0 2 0" size="1"/>
+          <body name="Carl" pos="1 1 0">
+            ...
+          </body>
+        </worldbody>
+      </mujoco>
+
+   Note that in the saved model, the frame elements have disappeared but their transformation was accumulated with those
+   of their child elements.
 
 .. _include:
 
 **include** (*)
-~~~~~~~~~~~~~~~
+^^^^^^^^^^^^^^^
 
-This element does not strictly speaking belong to MJCF. Instead it is a meta-element, used to assemble multiple XML
+This element does not strictly belong to MJCF. Instead it is a meta-element, used to assemble multiple XML
 files in a single document object model (DOM) before parsing. The included file must be a valid XML file with a unique
 top-level element. This top-level element is removed by the parser, and the elements below it are inserted at the
 location of the :el:`include` element. At least one element must be inserted as a result of this procedure. The
-:el:`include` element can be used where ever an XML element is expected in the MJFC file. Nested includes are allowed,
+:el:`include` element can be used where ever an XML element is expected in the MJCF file. Nested includes are allowed,
 however a given XML file can be included at most once in the entire model. After all the included XML files have been
 assembled into a single DOM, it must correspond to a valid MJCF model. Other than that, it is up to the user to decide
 how to use includes and how to modularize large files if desired.
@@ -216,11 +275,11 @@ any effect. The settings here are global and apply to the entire model.
 .. _compiler-eulerseq:
 
 :at:`eulerseq`: :at-val:`string, "xyz"`
-   This attribute specifies the sequence of Euler rotations for all euler attributes of elements that have spatial
+   This attribute specifies the sequence of Euler rotations for all :at:`euler` attributes of elements that have spatial
-   frames, as explained in :ref:`COrientation`. This must be a string with exactly 3
+   frames, as explained in :ref:`COrientation`. This must be a string with exactly 3 characters from the set {x, y, z,
-   characters from the set {'x', 'y', 'z', 'X', 'Y', 'Z'}. The character at position n determines the axis around which
+   X, Y, Z}. The character at position n determines the axis around which the n-th rotation is performed. Lower case
-   the n-th rotation is performed. Lower case denotes axes that rotate with the frame, while upper case denotes axes
+   letters denote axes that rotate with the frame (intrinsic), while upper case letters denote axes that remain fixed in
-   that remain fixed in the parent frame. The "rpy" convention used in URDF corresponds to the default "xyz" in MJCF.
+   the parent frame (extrinsic). The "rpy" convention used in URDF corresponds to "XYZ" in MJCF.
 
 .. _compiler-meshdir:
 
@@ -247,14 +306,20 @@ any effect. The settings here are global and apply to the entire model.
 .. _compiler-discardvisual:
 
 :at:`discardvisual`: :at-val:`[false, true], "false" for MJCF, "true" for URDF`
-   This attribute instructs the parser to discard "visual geoms", defined as geoms whose contype and conaffinity
+   This attribute instructs the compiler to discard all model elements which are purely visual and have no effect on the
-   attributes are both set to 0. This functionality is useful for models that contain two sets of geoms, one for
+   physics (with one exception, see below). This often enables smaller :ref:`mjModel` structs and faster simulation.
-   collisions and the other for visualization. Note that URDF models are usually constructed in this way. It rarely
+
-   makes sense to have two sets of geoms in the model, especially since MuJoCo uses convex hulls for collisions, so we
+   - All materials are discarded.
-   recommend using this feature to discard redundant geoms. Keep in mind however that geoms considered visual per the
+   - All textures are discarded.
-   above definition can still participate in collisions, if they appear in the explicit list of contact
+   - All geoms with :ref:`contype<body-geom-contype>`=:ref:`conaffinity<body-geom-conaffinity>`=0 are discarded, if they
-   :ref:`pairs <contact-pair>`. The parser does not check this list before discarding geoms; it relies solely on the geom
+     are not referenced in another MJCF element. If a discarded geom was used for inferring body inertia, an explicit
-   attributes to make the determination.
+     :ref:`inertial<body-inertial>` element is added to the body.
+   - All meshes which are not referenced by any geom (in particular those discarded above) are discarded.
+
+   The resulting compiled model will have exactly the same dynamics as the original model, with the exception of
+   raycasting, as used for example by :ref:`rangefinder<sensor-rangefinder>`, since raycasting reports distances to
+   visual geoms. When visualizing models compiled with this flag, it is important to remember that colliding geoms are
+   often placed in a :ref:`group<body-geom-group>` which is invisible by default.
 
 .. _compiler-convexhull:
 
@@ -1319,9 +1384,9 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe
 | For collision detection, a height field is treated as a union of triangular prisms. Collisions between height fields
   and other geoms (except for planes and other height fields which are not supported) are computed by first selecting
   the sub-grid of prisms that could collide with the geom based on its bounding box, and then using the general convex
-  collider. The number of possible contacts between a height field and a geom is limited to 9; any contacts beyond that
+  collider. The number of possible contacts between a height field and a geom is limited to 50
-  are discarded. To avoid penetration due to discarded contacts, the spatial features of the height field should be
+  (:ref:`mjMAXCONPAIR <glNumeric>`); any contacts beyond that are discarded. To avoid penetration due to discarded
-  large compared to the geoms it collides with.
+  contacts, the spatial features of the height field should be large compared to the geoms it collides with.
 
 .. _asset-hfield-name:
 
@@ -1877,7 +1942,7 @@ adjust it properly through the XML.
 .. _option-sdf_initpoints:
 
 :at:`sdf_initpoints`: :at-val:`int, "40"`
-   Number of starting points used for fining contacts with Signed Distance Field collisions.
+   Number of starting points used for finding contacts with Signed Distance Field collisions.
 
 .. _option-actuatorgroupdisable:
 
@@ -1924,7 +1989,7 @@ from its default.
 .. _option-flag-contact:
 
 :at:`contact`: :at-val:`[disable, enable], "enable"`
-   This flag disables all standard computations related to contact constraints.
+   This flag disables collision detection and all standard computations related to contact constraints.
 
 .. _option-flag-passive:
 
@@ -3913,6 +3978,35 @@ Associate this flexcomp with an :ref:`engine plugin<exPlugin>`. Either :at:`plug
 :at:`instance`: :at-val:`string, optional`
    Instance name, used for explicit plugin instantiation.
 
+
+.. _body-frame:
+
+:el-prefix:`body/` |-| **frame** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Frames specify a coordinate transformation which is applied to all child elements. They disappear during compilation
+and the transformation they encode is accumulated in their direct children. See :ref:`frame<frame>` for examples.
+
+.. _frame-pos:
+
+:at:`pos`: :at-val:`real(3), "0 0 0"`
+   The 3D position of the frame, in the parent coordinate system.
+
+.. _frame-quat:
+
+.. _frame-axisangle:
+
+.. _frame-xyaxes:
+
+.. _frame-zaxis:
+
+.. _frame-euler:
+
+:at:`quat`, :at:`axisangle`, :at:`xyaxes`, :at:`zaxis`, :at:`euler`
+   See :ref:`COrientation`.
+
+
+
 .. _contact:
 
 **contact** (*)
@@ -4532,7 +4626,7 @@ joint types (slide and hinge) can be used.
 :at:`polycoef`: :at-val:`real(5), "0 1 0 0 0"`
    Coefficients a0 ... a4 of the quartic polynomial. If the two joint values are y and x, and their reference positions
    (corresponding to the joint values in the initial model configuration) are y0 and x0, the constraint is:
-   y-y0 = a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4
+   y-y0 = a0 + a1*(x-x0) + a2*(x-x0)^2 + a3*(x-x0)^3 + a4*(x-x0)^4.
    Omitting the second joint is equivalent to setting x = x0, in which case the constraint is y = y0 + a0.
 
 
@@ -4942,21 +5036,21 @@ specify them independently.
 .. _actuator-general-ctrlrange:
 
 :at:`ctrlrange`: :at-val:`real(2), "0 0"`
-   Range for clamping the control input. The compiler expects the first value to be smaller than the second value.
+   Range for clamping the control input. The first value must be smaller than the second value.
    |br| Setting this attribute without specifying :at:`ctrllimited` is an error, unless :at:`autolimits` is set in
    :ref:`compiler <compiler>`.
 
 .. _actuator-general-forcerange:
 
 :at:`forcerange`: :at-val:`real(2), "0 0"`
-   Range for clamping the force output. The compiler expects the first value to be no greater than the second value.
+   Range for clamping the force output. The first value must be no greater than the second value.
    |br| Setting this attribute without specifying :at:`forcelimited` is an error, unless :at:`autolimits` is set in
    :ref:`compiler <compiler>`.
 
 .. _actuator-general-actrange:
 
 :at:`actrange`: :at-val:`real(2), "0 0"`
-   Range for clamping the activation state. The compiler expects the first value to be no greater than the second value.
+   Range for clamping the activation state. The first value must be no greater than the second value.
    See the :ref:`Activation clamping <CActRange>` section for more details.
    |br| Setting this attribute without specifying :at:`actlimited` is an error, unless :at:`autolimits` is set in
    :ref:`compiler <compiler>`.
@@ -5229,13 +5323,13 @@ This element does not have custom attributes. It only has common attributes, whi
 
 This element creates a position servo. The underlying :el:`general` attributes are set as follows:
 
-========= ======= ========= =======
+========= ======= ========= =========
 Attribute Setting Attribute Setting
-========= ======= ========= =======
+========= ======= ========= =========
 dyntype   none    dynprm    1 0 0
 gaintype  fixed   gainprm   kp 0 0
-biastype  affine  biasprm   0 -kp 0
+biastype  affine  biasprm   0 -kp -kv
-========= ======= ========= =======
+========= ======= ========= =========
 
 
 This element has one custom attribute in addition to the common attributes:
@@ -5289,6 +5383,11 @@ This element has one custom attribute in addition to the common attributes:
 :at:`kp`: :at-val:`real, "1"`
    Position feedback gain.
 
+.. _actuator-position-kv:
+
+:at:`kv`: :at-val:`real, "0"`
+   Damping applied by the actuator.
+   When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators<geIntegration>`.
 
 .. _actuator-velocity:
 
@@ -5297,7 +5396,9 @@ This element has one custom attribute in addition to the common attributes:
 
 This element creates a velocity servo. Note that in order create a PD controller, one has to define two actuators: a
 position servo and a velocity servo. This is because MuJoCo actuators are SISO while a PD controller takes two control
-inputs (reference position and reference velocity). The underlying :el:`general` attributes are set as follows:
+inputs (reference position and reference velocity).
+When using this actuator, it is recommended to use the implicitfast or implicit :ref:`integrators<geIntegration>`.
+The underlying :el:`general` attributes are set as follows:
 
 ========= ======= ========= =======
 Attribute Setting Attribute Setting
@@ -5368,14 +5469,14 @@ This element creates an integrated-velocity servo. For more information, see the
 :ref:`Activation clamping <CActRange>` section of the Modeling chapter. The underlying
 :el:`general` attributes are set as follows:
 
-==========   =========== ========= =======
+==========   =========== ========= =========
 Attribute    Setting     Attribute Setting
-==========   =========== ========= =======
+==========   =========== ========= =========
 dyntype      integrator  dynprm    1 0 0
 gaintype     fixed       gainprm   kp 0 0
-biastype     affine      biasprm   0 -kp 0
+biastype     affine      biasprm   0 -kp -kv
 actlimited   true
-==========   =========== ========= =======
+==========   =========== ========= =========
 
 This element has one custom attribute in addition to the common attributes:
 
@@ -5430,6 +5531,11 @@ This element has one custom attribute in addition to the common attributes:
 :at:`kp`: :at-val:`real, "1"`
    Position feedback gain.
 
+.. _actuator-intvelocity-kv:
+
+:at:`kv`: :at-val:`real, "0"`
+   Damping applied by the actuator.
+   When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators<geIntegration>`.
 
 .. _actuator-damper:
 
@@ -5437,8 +5543,9 @@ This element has one custom attribute in addition to the common attributes:
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 This element is an active damper which produces a force proportional to both velocity and control: ``F = - kv * velocity
-* control``, where ``kv`` must be nonnegative. :at:`ctrlrange` is required and must also be nonnegative. The underlying
+* control``, where ``kv`` must be nonnegative. :at:`ctrlrange` is required and must also be nonnegative.
-:el:`general` attributes are set as follows:
+When using this actuator, it is recommended to use the implicitfast or implicit :ref:`integrators<geIntegration>`.
+The underlying :el:`general` attributes are set as follows:
 
 =========== ======= ========= =======
 Attribute   Setting Attribute Setting
@@ -5784,12 +5891,29 @@ Associate this actuator with an :ref:`engine plugin<exPlugin>`. Either :at:`plug
 :at:`instance`: :at-val:`string, optional`
    Instance name, used for explicit plugin instantiation.
 
+.. _actuator-plugin-dyntype:
+
+:at:`dyntype`: :at-val:`[none, integrator, filter, filterexact, muscle, user], "none"`
+   Activation dynamics type for the actuator. The available dynamics types were already described in the :ref:`Actuation
+   model <geActuation>` section. If :ref:`dyntype<actuator-general-dyntype>` is not "none", an activation variable will
+   be added to the actuator. This variable will be added after any activation state computed by the plugin (see
+   :ref:`actuator plugin activations<exActuatorAct>`).
+
+.. _actuator-plugin-actrange:
+
+:at:`actrange`: :at-val:`real(2), "0 0"`
+   Range for clamping the activation state associated with this actuator's dyntype. The limit doesn't apply to
+   activations computed by the plugin. The first value must be no greater than the second value.
+   See the :ref:`Activation clamping <CActRange>` section for more details.
+
 .. _actuator-plugin-name:
 
 .. _actuator-plugin-class:
 
 .. _actuator-plugin-group:
 
+.. _actuator-plugin-actlimited:
+
 .. _actuator-plugin-ctrllimited:
 
 .. _actuator-plugin-forcelimited:
@@ -5818,9 +5942,14 @@ Associate this actuator with an :ref:`engine plugin<exPlugin>`. Either :at:`plug
 
 .. _actuator-plugin-user:
 
-.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`ctrllimited`,
+.. _actuator-plugin-dynprm:
+
+.. _actuator-plugin-actearly:
+
+.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`actlimited`, :at:`ctrllimited`,
    :at:`forcelimited`, :at:`ctrlrange`, :at:`forcerange`, :at:`lengthrange`, :at:`gear`, :at:`cranklength`,
-   :at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`
+   :at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`,
+   :at:`dynprm`, :at:`actearly`
 
 |actuator/plugin attrib list|
    Same as in actuator/ :ref:`general <actuator-general>`.
@@ -7570,6 +7699,8 @@ slidersite, cranksite.
 
 .. _default-position-kp:
 
+.. _default-position-kv:
+
 :el-prefix:`default/` |-| **position** (?)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -7626,6 +7757,8 @@ tendon, slidersite, cranksite.
 
 .. _default-intvelocity-kp:
 
+.. _default-intvelocity-kv:
+
 :el-prefix:`default/` |-| **intvelocity** (?)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

doc/XMLschema.rst

@@ -778,7 +778,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`jointinparent<actuator-position-jointinparent>`           | :ref:`tendon<actuator-position-tendon>`                         | :ref:`slidersite<actuator-position-slidersite>`                 | :ref:`cranksite<actuator-position-cranksite>`                   | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`site<actuator-position-site>`                             | :ref:`refsite<actuator-position-refsite>`                       | :ref:`kp<actuator-position-kp>`                                 |                                                                 | |
+|                                    |    |    | :ref:`site<actuator-position-site>`                             | :ref:`refsite<actuator-position-refsite>`                       | :ref:`kp<actuator-position-kp>`                                 | :ref:`kv<actuator-position-kv>`                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| actuator |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -810,6 +810,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`cranksite<actuator-intvelocity-cranksite>`                | :ref:`site<actuator-intvelocity-site>`                          | :ref:`refsite<actuator-intvelocity-refsite>`                    | :ref:`kp<actuator-intvelocity-kp>`                              | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`kv<actuator-intvelocity-kv>`                              |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| actuator |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`damper                       | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -879,13 +881,15 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<actuator-plugin-name>`                               | :ref:`class<actuator-plugin-class>`                             | :ref:`plugin<actuator-plugin-plugin>`                           | :ref:`instance<actuator-plugin-instance>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`group<actuator-plugin-group>`                             | :ref:`ctrllimited<actuator-plugin-ctrllimited>`                 | :ref:`forcelimited<actuator-plugin-forcelimited>`               | :ref:`ctrlrange<actuator-plugin-ctrlrange>`                     | |
+|                                    |    |    | :ref:`group<actuator-plugin-group>`                             | :ref:`ctrllimited<actuator-plugin-ctrllimited>`                 | :ref:`forcelimited<actuator-plugin-forcelimited>`               | :ref:`actlimited<actuator-plugin-actlimited>`                   | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`forcerange<actuator-plugin-forcerange>`                   | :ref:`lengthrange<actuator-plugin-lengthrange>`                 | :ref:`gear<actuator-plugin-gear>`                               | :ref:`cranklength<actuator-plugin-cranklength>`                 | |
+|                                    |    |    | :ref:`ctrlrange<actuator-plugin-ctrlrange>`                     | :ref:`forcerange<actuator-plugin-forcerange>`                   | :ref:`actrange<actuator-plugin-actrange>`                       | :ref:`lengthrange<actuator-plugin-lengthrange>`                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`joint<actuator-plugin-joint>`                             | :ref:`jointinparent<actuator-plugin-jointinparent>`             | :ref:`site<actuator-plugin-site>`                               | :ref:`tendon<actuator-plugin-tendon>`                           | |
+|                                    |    |    | :ref:`gear<actuator-plugin-gear>`                               | :ref:`cranklength<actuator-plugin-cranklength>`                 | :ref:`joint<actuator-plugin-joint>`                             | :ref:`jointinparent<actuator-plugin-jointinparent>`             | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               |                                                                 | |
+|                                    |    |    | :ref:`site<actuator-plugin-site>`                               | :ref:`dyntype<actuator-plugin-dyntype>`                         | :ref:`dynprm<actuator-plugin-dynprm>`                           | :ref:`tendon<actuator-plugin-tendon>`                           | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               | :ref:`actearly<actuator-plugin-actearly>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| plugin |br| |_2| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -1438,7 +1442,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`gear<default-position-gear>`                              | :ref:`cranklength<default-position-cranklength>`                | :ref:`user<default-position-user>`                              | :ref:`group<default-position-group>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`kp<default-position-kp>`                                  |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`kp<default-position-kp>`                                  | :ref:`kv<default-position-kv>`                                  |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| default |br| |_| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -1460,7 +1464,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`actrange<default-intvelocity-actrange>`                   | :ref:`gear<default-intvelocity-gear>`                           | :ref:`cranklength<default-intvelocity-cranklength>`             | :ref:`user<default-intvelocity-user>`                           | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`group<default-intvelocity-group>`                         | :ref:`kp<default-intvelocity-kp>`                               |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`group<default-intvelocity-group>`                         | :ref:`kp<default-intvelocity-kp>`                               | :ref:`kv<default-intvelocity-kv>`                               |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| default |br| |_| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |

doc/changelog.rst

@@ -2,6 +2,82 @@
 Changelog
 =========
 
+Upcoming version (not yet released)
+-----------------------------------
+
+General
+^^^^^^^
+1. Improved the :ref:discardvisual<compiler-discardvisual> compiler flag, which now discards all visual-only assets. See
+   :ref:discardvisual<compiler-discardvisual> for details.
+
+MJX
+^^^
+2. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
+3. Added :at:`site` transmission.
+4. Updated MJX colab tutorial with more stable quadruped environment.
+5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes.
+
+Bug fixes
+^^^^^^^^^
+6. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
+   :github:issue:`1270`.
+
+
+Version 3.1.1 (December 18, 2023)
+-----------------------------------
+
+Bug fixes
+^^^^^^^^^
+1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the other.
+2. Fixed a bug in :ref:`simulate<saSimulate>` where the "LOADING..." message was not showing correctly.
+3. Fixed a crash in the Python :ref:`passive viewer<PyViewerPassive>`, when used with models containing Flex objects.
+4. Fixed a bug in MJX where ``site_xmat`` was ignored in ``get_data`` and ``put_data``
+5. Fixed a bug in MJX where ``efc_address`` was sometimes incorrectly calculated in ``get_data``.
+
+
+Version 3.1.0 (December 12, 2023)
+---------------------------------
+
+General
+^^^^^^^
+1. Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function
+   for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is more
+   robust for very small or large geom sizes.
+2. Added :ref:`frame<frame>` to MJCF, a :ref:`meta-element<meta-element>` which defines a pure coordinate transformation
+   on its direct children, without requiring a :ref:`body<body>`.
+3. Added the :at:`kv` attribute to the :ref:`position<actuator-position>` and :ref:`intvelocity<actuator-intvelocity>`
+   actuators, for specifying actuator-applied damping. This can be used to implement a PD controller with 0 reference
+   velocity. When using this attribute, it is recommended to use the implicitfast or implicit
+   :ref:`integrators<geIntegration>`.
+
+Plugins
+^^^^^^^
+
+4. Allow actuator plugins to use activation variables in ``mjData.act`` as their internal state, rather than
+   ``mjData.plugin_state``. Actuator plugins can now specify :ref:`callbacks<mjpPlugin>` that compute activation
+   variables, and they can be used with built-in :ref:`dyntype<actuator-plugin-dyntype>` actuator dynamics.
+
+5. Added the `pid <https://github.com/deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ actuator plugin, a
+   configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators.
+
+MJX
+^^^
+
+6.  Added ``site_xpos`` and ``site_xmat`` to MJX.
+7. Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be
+   deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as
+   ``efc_J``.
+
+Bug fixes
+^^^^^^^^^
+8. Fix bug in Cartesian actuation with movable refsite, as when using body-centric Cartesian actuators on a quadruped.
+   Before this fix such actuators could lead to non-conservation of momentum.
+9. Fix bug that prevented using flex with :ref:`simulate<saSimulate>`.
+10. Fix bug that prevented the use of elasticity plugins in combination with pinned flex vertices.
+11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where SYSTEM_VERSION_COMPAT is set. The minimum
+    supported version is still 11.0, but we release these wheels to fix compatibility for those users. See
+    :github:issue:`1213`.
+
 Version 3.0.1 (November 15, 2023)
 ---------------------------------
 

doc/computation/fluid.rst

@@ -520,8 +520,14 @@ in the surrounding flow a circulation of sufficient strength to hold the rear st
 This is the Kutta condition, a fluid dynamic phenomenon that can be observed for solid bodies with sharp corners, such
 as slender bodies or the trailing edges of airfoils.
 
-.. cssclass:: caption-small
 .. figure:: ../images/computation/kutta_cond_plate.svg
+   :class: only-light
+   :figwidth: 95%
+   :align: left
+
+.. cssclass:: caption-small
+.. figure:: ../images/computation/kutta_cond_plate_dark.svg
+   :class: only-dark
    :figwidth: 95%
    :align: left
 

doc/computation/index.rst

@@ -300,19 +300,22 @@ is attached; the possible attachment object types are :at:`joint`, :at:`tendon`,
    Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality constraints to
    the rest of the system, but that would be less efficient.
 
-:at:`site`
+:at:`body`
-   :at:`site` transmission (without a :at:`refsite`, see below) and :at:`body` transmission targets have a fixed zero
+   :el:`body` transmission corresponds to applying forces at contact points belonging to a body, in
-   length :math:`l_i(q) = 0`. They can therefore not be used to maintain a desired length, but can be used to apply
-   forces. Site transmissions correspond to applying a Cartsian force/torque at the site, and are useful for modeling
-   jets and propellors. :el:`body` transmissions correspond to applying forces at contact points belonging to a body, in
    order to model vacuum grippers and biomechanical adhesive appendages. For more information about adhesion, see the
-   :ref:`adhesion<actuator-adhesion>` actuator documentation.
+   :ref:`adhesion<actuator-adhesion>` actuator documentation. These transmission targets have a fixed zero length
+   :math:`l_i(q) = 0`.
 
-   If a :at:`site` transmission target is defined with the optional :at:`refsite` attribute, forces and torques are
+:at:`site`
-   applied in the frame of the reference site rather than the site's own frame. If a reference site is defined then
+   Site transmissions correspond to applying a Cartsian force/torque in the frame of a site. When a :at:`refsite` is not
-   the length of the actuator is nonzero and corresponds to the pose difference of the two sites. This length can then
+   defined (see below), these targets have a fixed zero length :math:`l_i(q) = 0` and are useful for modeling jets and
-   be controlled with a :el:`position` actuator, enabling Cartesian end-effector control. See the
+   propellors: forces and torques which are fixed to the site frame.
-   :ref:`refsite<actuator-general-refsite>` documentation for more details.
+
+   If a :at:`site` transmission is defined with the optional :at:`refsite` attribute, forces and torques are applied in
+   the frame of the reference site rather than the site's own frame. If a reference site is defined, the length of the
+   actuator is nonzero and corresponds to the pose difference of the two sites, projected onto a chosen direction in the
+   reference frame. This length can then be controlled with a :el:`position` actuator, allowing for Cartesian
+   end-effector control. See the :ref:`refsite<actuator-general-refsite>` documentation for more details.
 
 .. _geActivation:
 
@@ -963,6 +966,12 @@ is :math:`E f`. The matrix of basis vectors is constructed as follows.
 .. image:: ../images/computation/contact_frame.svg
    :width: 700px
    :align: center
+   :class: only-light
+
+.. image:: ../images/computation/contact_frame_dark.svg
+   :width: 700px
+   :align: center
+   :class: only-dark
 
 The figure illustrates the full basis set corresponding to the case :math:`n = 6`. Otherwise we use only the first
 :math:`n` or :math:`2(n-1)` columns depending on the cone type. Elliptic cones are easier to understand. Since the
@@ -1322,6 +1331,12 @@ representations of the constraint Jacobian and related matrices.
    .. image:: ../images/computation/gPGS.svg
       :width: 500px
       :align: center
+      :class: only-light
+
+   .. image:: ../images/computation/gPGS_dark.svg
+      :width: 500px
+      :align: center
+      :class: only-dark
 
    When using pyramidal friction cones, the problem involves box constraints to which PGS has traditionally been
    applied. If we applied PGS directly to the conic constraints resulting from elliptic friction cones, it would get
@@ -1421,12 +1436,18 @@ approximations, no matter how accurate the approximation is. The figure below il
 where the pyramid is not even an approximation, but represents the same constraint set as the elliptic cone. We plot the
 contours of the penalty/shadow for the pyramidal (red) and elliptic (dashed blue) cones, for different friction
 coefficients varying from left to right. Mathematically, the penalty in the pyramidal case is a quadratic spline, while
-the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics - allowing
+the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics -- allowing
 circular contours around the tip of the cone.
 
 .. image:: ../images/computation/softcontact.png
    :width: 600px
    :align: center
+   :class: only-light
+
+.. image:: ../images/computation/softcontact_dark.png
+   :width: 600px
+   :align: center
+   :class: only-dark
 
 In summary, elliptic and pyramidal friction cones define different soft-contact dynamics (although they are usually very
 close). The elliptic model is more principled and more consistent with physical intuition, and the corresponding solvers
@@ -1455,48 +1476,45 @@ others can be pruned quickly without a detailed check. MuJoCo has flexible mecha
 checked in detail. The decision process involves two stages: generation and filtering.
 
 Generation
-   First we generate a list of candidate geom pairs in one of two ways: "pair" or "dynamic". The user can also specify
+   First we generate a list of candidate geom pairs by merging from two sources: pairs of bodies that might contain
-   "all" which merges both sources (and is the default). This is done via the setting ``mjModel.opt.collision``. "Pair"
+   colliding geoms and the explicit list of geom pairs defined with the :ref:`pair <contact-pair>` element in MJCF.
-   refers to an explicit list of geom pairs defined with the :ref:`pair <contact-pair>` element in MJCF. It gives the
-   user full control, however it is a static mechanism (independent of the spatial arrangement of the geoms at runtime)
-   and can be tedious for large models. It is normally used to supplement the output of the "dynamic" mechanism. Dynamic
-   generation works with bodies rather than geoms; when a body pair is included this means that all geoms attached to
-   one body can collide with all geoms attached to the other body.
 
    The body pairs are generated via broad-phase collision detection based on a modified sweep-and-prune algorithm. The
    modification is that the axis for sorting is chosen as the principal eigenvector of the covariance matrix of all geom
-   centers - which maximizes the spread. Then, for each body pair, a mid-phase collision detection using a static
+   centers -- which maximizes the spread. Then, for each body pair, mid-phase collision detection is performed using a
-   bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB) is performed. Each body is
+   static bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB). Each body is equipped
-   equipped with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes,
+   with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes,
    respectively.
 
-   Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude <contact-exclude>` element
+   Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude <contact-exclude>` element in
-   in MJCF. Exclusion is applied when "dynamic" or "all" are selected, but not when "pair" is selected. At the end of
+   MJCF. At the end of this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`,
-   this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`, but can still be
+   but can still be pruned further before detailed collision checking.
-   pruned further before detailed collision checking.
 
 Filtering
    Next we apply four filters to the list generated in the previous step. Filters 1 and 2 are applied to all geom pairs.
-   Filters 3 and 4 are applied only to pairs generated by the "dynamic" mechanism, thereby allowing the user to bypass
+   Filters 3 and 4 are applied only to pairs generated by the body-pair mechanism, thereby allowing the user to bypass
    those filters by specifying geom pairs explicitly.
 
-   #. The types of the two geoms must correspond to a collision function that is capable of performing the detailed
+   1. The types of the two geoms must correspond to a collision function that is capable of performing the detailed
       check. This is usually the case but there are exceptions (for example plane-plane collisions are not supported),
       and furthermore the user may override the default table of collision functions with NULL pointers, effectively
       disabling collisions between certain geom types.
-   #. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a
+   2. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a
       plane, this becomes a plane-sphere test.
-   #. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless
+   3. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless
       the parent is the world body. The motivation is to avoid permanent contacts within bodies and joints. Note that if
       several bodies are welded together in the sense that there are no joints between them, they are treated as a
       single body for the purposes of this test. The parent-filter test can be disabled by the user, while the same-body
       test cannot be disabled.
-   #. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and
+   4. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and
       ``conaffinity``. The boolean expression below must be true for the test to pass:
-      ``(contype1 & conaffinity2) || (contype2 & conaffinity1)`` This requires the ``contype`` of one geom and the
+
-      ``conaffinity`` of the other geom to have a common bit set to 1. This is a powerful mechanism borrowed from the
+      ``(contype1 & conaffinity2) || (contype2 & conaffinity1)``
-      Open Dynamics Engine. The default setting for all geoms is ``contype = conaffinity = 1`` which always passes the
+
-      test, so the user can ignore this mechanism if it is confusing at first.
+      This requires the ``contype`` of one geom and the ``conaffinity`` of the other geom to have a common bit set to 1.
+      This is a powerful mechanism borrowed from Open Dynamics Engine. The default setting for all geoms is
+      ``contype = conaffinity = 1`` which always passes the test, so the user can ignore this mechanism if it is
+      confusing at first.
 
 .. _coChecking:
 
@@ -1517,12 +1535,12 @@ convex hull implicitly, however pre-computing that hull can substantially improv
 model compiler does that by default, using the `qhull <http://www.qhull.org/>`__ library.
 
 In order to model a non-convex object other than a height field, the user must decompose it into a union of convex geoms
-(which can be primitive shapes or meshes) and attach them to the same body. Tools such as the
+(which can be primitive shapes or meshes) and attach them to the same body. Open tools like the `CoACD library
-`HACD <https://github.com/kmammou/v-hacd>`__ library can be used outside MuJoCo to automate this process. Finally, all
+<https://github.com/SarahWeiii/CoACD>`__ can be used outside MuJoCo to automate this process. Finally, all built-in
-built-in collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose
+collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose "triangle
-"triangle soup" collision detector for example. However we do not recommend such an approach. Pre-processing the
+soup" collision detector for example. However we do not recommend such an approach. Pre-processing the geometry and
-geometry and representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both
+representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both faster and more
-faster and more stable simulation.
+stable simulation.
 
 .. _Pipeline:
 
@@ -1538,44 +1556,59 @@ be used to skip default steps and to enable optional steps respectively. Callbac
 Forward dynamics
 ~~~~~~~~~~~~~~~~
 
-The top-level function :ref:`mj_step` invokes the sequence of computations below. Alternatively one can call
+The source file `engine_forward.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_forward.c>`__
-:ref:`mj_forward` which invokes only steps 2-21.
+contains the high-level forward dynamics pipeline:
 
-#. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence
+- The top-level function :ref:`mj_step` invokes the entire sequence of computations below.
-   is detected, the state is automatically reset and the corresponding warning is raised.
+- :ref:`mj_forward` invokes only stages **2-22**, computing the continuous-time forward dynamics, ending with the
-#. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites,
+  acceleration ``mjData.qacc``.
-   cameras and lights. It also normalizes all quaternions, just in case.
+- :ref:`mj_step1` invokes stages **1-18** and :ref:`mj_step2` invokes stages **19-25**, breaking :ref:`mj_step` into two
-#. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding
+  distinct phases. This allows the user to write controllers that depend on quantities derived from the positions and
-   kinematic subtrees (to improve floating-point accuracy).
+  velocities (but not forces, since those have not yet been computed). Note that the :ref:`mj_step1` → :ref:`mj_step2`
-#. Compute the actuator lengths and moment arms.
+  pipeline does not support the Runge Kutta integrator.
-#. Compute the composite rigid body inertias and construct the joint-space inertia matrix.
+
-#. Compute the sparse factorization of the joint-space inertia matrix.
+1. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence
-#. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection.
+   is detected, the state is automatically reset and the corresponding warning is raised:
-#. Construct the constraint Jacobian and compute the constraint residuals.
+   :ref:`mj_checkPos`, :ref:`mj_checkVel`
-#. Compute the matrices and vectors needed by the constraint solvers.
+2. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites,
-#. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
+   cameras and lights. It also normalizes all quaternions: :ref:`mj_kinematics`, :ref:`mj_camLight`
-   tendons.
+3. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding
-#. Compute sensor data that only depends on position, and the potential energy if enabled.
+   kinematic subtrees: :ref:`mj_comPos`
-#. Compute the tendon and actuator velocities.
+4. Compute quantities related to :ref:`flex<deformable-flex>` objects: :ref:`mj_flex`
-#. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at
+5. Compute the actuator lengths and moment arms: :ref:`mj_tendon`
-   the subtree centers of mass.
+6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb`
-#. Compute all passive forces: spring-dampers in joints and tendons, and fluid dynamics forces.
+7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM`
-#. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
+8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection:
-   If required by sensors, call :ref:`mj_subtreeVel`.
+   :ref:`mj_collision`
-#. Compute the reference constraint acceleration.
+9. Construct the constraint Jacobian and compute the constraint residuals: :ref:`mj_makeConstraint`
-#. Compute the vector of Coriolis, centrifugal and gravitational forces.
+10. Compute the matrices and vectors needed by the constraint solvers: :ref:`mj_projectConstraint`
-#. Compute the actuator forces and activation dynamics if defined.
+11. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
-#. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces.
+    tendons: :ref:`mj_transmission`
-#. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
+12. Compute sensor data that only depends on position, and the potential energy if enabled: :ref:`mj_sensorPos`,
-   constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics.
+    :ref:`mj_energyPos`
-#. Compute sensor data that depends on force and acceleration if enabled.
+13. Compute the tendon, flex edge and actuator velocities: :ref:`mj_fwdVelocity`
-   If required by sensors, call :ref:`mj_rnePostConstraint`.
+14. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at
-#. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
+    the subtree centers of mass: :ref:`mj_comVel`
-   automatically reset and the corresponding warning is raised.
+15. Compute passive forces -- spring-dampers in joints and tendons, and fluid forces: :ref:`mj_passive`
-#. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
+16. Compute sensor data that depends on velocity, and the kinetic energy if enabled
-   dynamics. This is an optional step, and is performed only when enabled.
+    (if required by sensors, call :ref:`mj_subtreeVel`): :ref:`mj_sensorVel`
-#. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
+17. Compute the reference constraint acceleration: :ref:`mj_referenceConstraint`
-   repeats the above sequence three more times, except for the optional computations which are performed only once.
+18. Compute the vector of Coriolis, centrifugal and gravitational forces: :ref:`mj_rne`
+19. Compute the actuator forces and activation dynamics if defined: :ref:`mj_fwdActuation`
+20. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces:
+    :ref:`mj_fwdAcceleration`
+21. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
+    constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics:
+    :ref:`mj_fwdConstraint`
+22. Compute sensor data that depends on force and acceleration if enabled
+    (if required by sensors, call :ref:`mj_rnePostConstraint`): :ref:`mj_sensorAcc`
+23. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
+    automatically reset and the corresponding warning is raised: :ref:`mj_checkAcc`
+24. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
+    dynamics. This is an optional step, and is performed only when enabled: :ref:`mj_compareFwdInv`
+25. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
+    repeats the above sequence three more times, except for the optional computations which are performed only once:
+    one of :ref:`mj_Euler`, :ref:`mj_RungeKutta`, :ref:`mj_implicit`
 
 .. _piInverse:
 

doc/images/computation/contact_frame_dark.svg

@@ -0,0 +1,296 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 559.62 244.73">
+  <defs>
+    <style>
+      .cls-1 {
+        font-size: 12px;
+      }
+
+      .cls-1, .cls-2, .cls-3, .cls-4, .cls-5, .cls-6 {
+        fill: #fff;
+      }
+
+      .cls-1, .cls-3, .cls-4, .cls-5 {
+        font-family: ArialMT, Arial;
+      }
+
+      .cls-1, .cls-3, .cls-4, .cls-5, .cls-7 {
+        isolation: isolate;
+      }
+
+      .cls-8 {
+        stroke-width: 2px;
+      }
+
+      .cls-8, .cls-9 {
+        stroke: #4c4c4c;
+      }
+
+      .cls-8, .cls-9, .cls-10 {
+        fill: none;
+      }
+
+      .cls-2, .cls-10 {
+        stroke: #fff;
+      }
+
+      .cls-3 {
+        font-size: 8px;
+      }
+
+      .cls-9 {
+        stroke-dasharray: 0 0 5 5;
+      }
+
+      .cls-4 {
+        font-size: 24px;
+      }
+
+      .cls-5 {
+        font-size: 9px;
+      }
+
+      .cls-6 {
+        stroke-width: 0px;
+      }
+    </style>
+  </defs>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="79.67" y2="174.04"/>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="72.1" y2="113.9"/>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="28.52" y2="69.94"/>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(25.61 10.3)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(32.28 13.3)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(126 209.44)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(132.67 212.44)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(123.42 85.94)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(130.09 88.94)"><tspan x="0" y="0">3</tspan></text>
+  </g>
+  <g>
+    <path class="cls-10" d="M56.74,118.44c0-11,6.36-20,14.2-20s14.19,9,14.19,20"/>
+    <path class="cls-6" d="M85.13,117.67l2.7-1.64.06.09-1.75,4.41c-.34,1.53-.67,3.07-1,4.61-.34-1.54-.68-3.08-1-4.61l-1.81-4.41v-.09l2.8,1.64Z"/>
+  </g>
+  <g>
+    <path class="cls-10" d="M90.33,180.04c0-11-6.35-20-14.19-20s-14.19,9-14.19,20"/>
+    <path class="cls-6" d="M61.94,179.27l2.7-1.64.06.09-1.75,4.41c-.34,1.53-.68,3.07-1,4.6l-1-4.6-1.75-4.41.05-.09,2.69,1.64Z"/>
+  </g>
+  <g>
+    <path class="cls-10" d="M21.49,83.23c14.25,0,25.8-5.49,25.8-12.26s-11.55-12.26-25.8-12.26"/>
+    <path class="cls-6" d="M22.25,58.71l1.64,2.71-.09.06-4.41-1.75-4.6-1,4.6-1,4.41-1.79h.09l-1.64,2.77Z"/>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(0 62.46)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(6.67 65.46)"><tspan x="0" y="0">4</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(58.77 199.36)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(65.44 202.36)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(81.16 136.09)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(87.83 139.09)"><tspan x="0" y="0">6</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(18.39 122.2)"><tspan x="0" y="0">x</tspan></text>
+  <text class="cls-1" transform="translate(41.29 162.53)"><tspan x="0" y="0">y</tspan></text>
+  <text class="cls-1" transform="translate(50.31 137.92)"><tspan x="0" y="0">z</tspan></text>
+  <path class="cls-2" d="M32,139.57c0,1.92-1.56,3.48-3.48,3.48h0c-1.92,0-3.48-1.55-3.49-3.47h0c0-1.93,1.56-3.49,3.48-3.49,0,0,0,0,.01,0,1.92,0,3.48,1.56,3.48,3.48h0Z"/>
+  <g>
+    <line class="cls-10" x1="28.52" y1="48.07" x2="28.52" y2="22.19"/>
+    <path class="cls-6" d="M28.52,22.94l-2.71,1.64-.06-.09,1.75-4.41,1-4.61c.33,1.54.67,3.07,1,4.61l1.75,4.41-.05.09-2.68-1.64Z"/>
+  </g>
+  <g>
+    <line class="cls-10" x1="94.7" y1="184.06" x2="115.42" y2="198"/>
+    <path class="cls-6" d="M114.79,197.57l.14-3.16h.11l2.68,3.91c1.09,1.14,2.17,2.28,3.26,3.41l-4.39-1.73-4.63-1-.05-.09,2.88-1.34Z"/>
+  </g>
+  <g>
+    <line class="cls-10" x1="88.95" y1="104.87" x2="110.74" y2="92.01"/>
+    <path class="cls-6" d="M110.08,92.4l-2.75-1.46.05-.1,4.68-.74,4.48-1.46-3.45,3.21-2.9,3.75h-.11v-3.2Z"/>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(265.54 123.6)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(272.21 126.6)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(229.69 111.4)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(236.36 114.4)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(307.66 91.78)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(314.33 94.78)"><tspan x="0" y="0">3</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(188.06 133.82)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(194.73 136.82)"><tspan x="0" y="0">4</tspan></text>
+  </g>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="303.2" y2="207.11"/>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="298.79" y2="149.1"/>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="255.21" y2="105.14"/>
+  <path class="cls-2" d="M258.7,174.77c0,1.92-1.56,3.48-3.48,3.48,0,0,0,0-.01,0-1.92,0-3.48-1.56-3.48-3.48h0c0-1.92,1.56-3.48,3.48-3.48h0c1.92,0,3.48,1.55,3.49,3.47,0,0,0,0,0,.01Z"/>
+  <g>
+    <path class="cls-9" d="M301.33,92.38c7.4,10.56-7.74,29.27-33.83,41.8-26.09,12.53-53.24,14.14-60.65,3.59s7.74-29.27,33.83-41.8,53.24-14.14,60.65-3.59Z"/>
+    <g>
+      <line class="cls-10" x1="209.77" y1="140.36" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M210.38,140.82l-.33,3.15h-.11l-2.45-4.03-3.06-3.59,4.28,2,4.57,1.27.05.09-2.95,1.11Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="267.54" y1="136.09" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M267.33,136.81l-3.07.75v-.11l3-3.66,2.36-4.09c-.14,1.57-.28,3.14-.43,4.7l.33,4.73-.07.07-2.12-2.39Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="242.89" y1="100.97" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M243.01,101.72l-2.39,2.07-.08-.08,1-4.64.24-4.71c.59,1.46,1.17,2.92,1.76,4.38l2.45,4.06v.09l-2.98-1.17Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="300.66" y1="96.7" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M300.33,97.35l-3.17.07v-.11l3.72-2.93c1.07-1.16,2.13-2.32,3.2-3.47l-1.44,4.49-.71,4.68-.08.06-1.52-2.79Z"/>
+    </g>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(70.45 14.72)"><tspan x="0" y="0">elliptic basis: E = I</tspan></text>
+    <text class="cls-5" transform="translate(166.83 17.72)"><tspan x="0" y="0">6</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(312.5 17.19)"><tspan x="0" y="0">pyramidal basis: E =</tspan></text>
+  <g>
+    <g>
+      <line class="cls-10" x1="255.97" y1="83.27" x2="255.97" y2="57.39"/>
+      <path class="cls-6" d="M255.97,58.15l-2.7,1.64-.06-.09,1.75-4.41c.34-1.53.67-3.07,1-4.6l1,4.6,1.75,4.41v.09l-2.74-1.64Z"/>
+    </g>
+    <g>
+      <path class="cls-10" d="M242.94,87.42c14.26,0,25.81-5.49,25.81-12.26s-11.55-12.22-25.81-12.22"/>
+      <path class="cls-6" d="M243.71,62.94l1.64,2.7-.09.06-4.41-1.76-4.6-1,4.6-1,4.41-1.75h.09l-1.64,2.75Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(239.43 76.17)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(246.1 79.17)"><tspan x="0" y="0">5</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <line class="cls-10" x1="255.97" y1="40.49" x2="255.97" y2="14.61"/>
+      <path class="cls-6" d="M255.97,15.37l-2.7,1.65-.06-.1,1.75-4.4c.34-1.54.67-3.07,1-4.61l1,4.61,1.75,4.4v.1l-2.74-1.65Z"/>
+    </g>
+    <g>
+      <path class="cls-10" d="M268.75,45.65c-14.25,0-25.81-5.49-25.81-12.26s11.56-12.26,25.81-12.26"/>
+      <path class="cls-6" d="M267.98,21.13l-1.64-2.7.09-.06,4.41,1.75,4.61,1-4.61,1-4.41,1.82-.09-.05,1.64-2.76Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(260.75 35.64)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(267.42 38.64)"><tspan x="0" y="0">6</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M323.24,213.31c0-11-6.36-20-14.2-20s-14.19,9-14.19,20"/>
+      <path class="cls-6" d="M294.85,212.55l2.7-1.65.07.1-1.75,4.4-1,4.61c-.34-1.54-.67-3.07-1-4.61l-1.75-4.4v-.1l2.73,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="309.01" y1="211" x2="309.01" y2="185.12"/>
+      <path class="cls-6" d="M309.01,185.88l-2.68,1.64-.06-.09,1.75-4.4,1-4.61c.33,1.54.67,3.07,1,4.61l1.75,4.4v.09l-2.76-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(301.82 221.38)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(308.49 224.38)"><tspan x="0" y="0">7</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M324.24,230.74c0-11.05,6.35-20,14.19-20s14.19,8.95,14.19,20"/>
+      <path class="cls-6" d="M352.62,229.94l2.71-1.65.06.09-1.75,4.41-1,4.61c-.33-1.54-.67-3.07-1-4.61l-1.75-4.41v-.09l2.73,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="338.39" y1="228.43" x2="338.39" y2="202.55"/>
+      <path class="cls-6" d="M338.39,203.31l-2.7,1.64-.06-.09,1.75-4.4c.34-1.54.67-3.07,1-4.61l1,4.61,1.75,4.4v.09l-2.74-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(331.21 238.81)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(337.88 241.81)"><tspan x="0" y="0">8</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M293.85,147.15c0-11,6.36-20,14.19-20s14.2,9,14.2,20"/>
+      <path class="cls-6" d="M322.24,146.39l2.7-1.65.06.09-1.75,4.41c-.34,1.54-.67,3.07-1,4.61l-1-4.61-1.75-4.41v-.09l2.74,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="308.01" y1="144.84" x2="308.01" y2="118.96"/>
+      <path class="cls-6" d="M308.01,119.72l-2.68,1.64-.06-.09,1.75-4.4,1-4.61c.33,1.53.67,3.07,1,4.61l1.75,4.4v.09l-2.76-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(300.82 155.22)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(307.49 158.22)"><tspan x="0" y="0">9</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M353.08,128.06c0-11.05-6.36-20-14.2-20s-14.19,9-14.19,20"/>
+      <path class="cls-6" d="M324.69,127.29l2.7-1.64.06.09-1.75,4.41c-.33,1.54-.67,3.07-1,4.61-.34-1.54-.67-3.07-1-4.61l-1.77-4.41v-.09l2.76,1.64Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="338.85" y1="125.75" x2="338.85" y2="99.87"/>
+      <path class="cls-6" d="M338.85,100.63l-2.71,1.64-.06-.09,1.75-4.41,1-4.6c.33,1.53.67,3.07,1,4.6l1.75,4.41-.05.09-2.68-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(331.66 136.13)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(338.33 139.13)"><tspan x="0" y="0">10</tspan></text>
+    </g>
+  </g>
+  <text class="cls-1" transform="translate(378.74 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(373.33 65.99)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(390.33 68.99)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(378.74 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(407.55 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(402.13 65.99)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(416.12 68.99)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(435.2 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(429.78 88.42)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(446.78 91.42)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(435.2 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(463 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(457.59 88.42)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(471.58 91.42)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(517.79 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(512.37 157.17)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(529.37 160.17)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(517.79 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(546.59 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(541.18 157.17)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(555.17 160.17)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <text class="cls-4" transform="translate(485.39 98.39)"><tspan x="0" y="0">...</tspan></text>
+</svg>

doc/images/computation/gPGS_dark.svg

@@ -0,0 +1,102 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 454.98 171.24">
+  <defs>
+    <style>
+      .cls-1 {
+        font-family: ArialMT, Arial;
+        font-size: 12px;
+      }
+
+      .cls-1, .cls-2, .cls-3 {
+        fill: #fff;
+      }
+
+      .cls-1, .cls-4 {
+        isolation: isolate;
+      }
+
+      .cls-5 {
+        stroke-dasharray: 0 0 4 4;
+      }
+
+      .cls-5, .cls-2, .cls-6, .cls-7, .cls-3, .cls-8 {
+        stroke: #dce0df;
+      }
+
+      .cls-5, .cls-6, .cls-7, .cls-9, .cls-8, .cls-10 {
+        fill: none;
+      }
+
+      .cls-11 {
+        fill: #0c0c0c;
+      }
+
+      .cls-11, .cls-10 {
+        stroke: #666;
+      }
+
+      .cls-6 {
+        stroke-dasharray: 0 5;
+      }
+
+      .cls-7 {
+        stroke-dasharray: 0 0 5 5;
+      }
+
+      .cls-3 {
+        stroke-miterlimit: 4;
+        stroke-width: 1px;
+      }
+
+      .cls-9 {
+        stroke: #fff;
+        stroke-width: 3px;
+      }
+
+      .cls-12 {
+        fill: #dce0df;
+        stroke-width: 0px;
+      }
+    </style>
+  </defs>
+  <line class="cls-10" x1="68" y1="151.26" x2="133.33" y2="67.93"/>
+  <line class="cls-10" x1="68" y1="151.26" x2="2.67" y2="67.93"/>
+  <circle class="cls-3" cx="122.67" cy="118.59" r="2.67"/>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(0 45.93)"><tspan x="0" y="0">cone</tspan></text>
+    <text class="cls-1" transform="translate(0 60.33)"><tspan x="0" y="0">constraint</tspan></text>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(133.33 134.6)"><tspan x="0" y="0">unconstrained</tspan></text>
+    <text class="cls-1" transform="translate(133.33 149)"><tspan x="0" y="0">minimum</tspan></text>
+  </g>
+  <line class="cls-7" x1="122.67" y1="118.6" x2="122.67" y2="81.54"/>
+  <line class="cls-6" x1="125.33" y1="118.6" x2="93.61" y2="118.6"/>
+  <line class="cls-9" x1="93.14" y1="119.2" x2="122.67" y2="81.54"/>
+  <g>
+    <line class="cls-8" x1="104" y1="53.93" x2="104" y2="88.93"/>
+    <path class="cls-12" d="M104,88.17l2.7-1.64.06.09-1.75,4.4c-.33,1.54-.67,3.07-1,4.61-.34-1.54-.67-3.07-1-4.61l-1.75-4.4v-.09l2.74,1.64Z"/>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(85.81 35.63)"><tspan x="0" y="0">continuum of</tspan></text>
+    <text class="cls-1" transform="translate(85.81 50.03)"><tspan x="0" y="0">PGS local minima</tspan></text>
+  </g>
+  <line class="cls-10" x1="387.29" y1="170.93" x2="452.62" y2="87.6"/>
+  <line class="cls-10" x1="387.29" y1="170.93" x2="321.96" y2="87.6"/>
+  <path class="cls-11" d="M454.48,77.71c0,19.83-30.19,35.9-67.44,35.9s-67.42-16.08-67.42-35.9,30.2-35.9,67.45-35.9,67.41,16.06,67.41,35.9Z"/>
+  <path class="cls-2" d="M417.85,70.44c0,1.03-.83,1.86-1.86,1.86s-1.86-.83-1.86-1.86.83-1.86,1.86-1.86h0c1.02,0,1.85.82,1.86,1.84,0,0,0,.01,0,.02Z"/>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(402.51 10.3)"><tspan x="0" y="0">search</tspan></text>
+    <text class="cls-1" transform="translate(402.51 24.7)"><tspan x="0" y="0">ray</tspan></text>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(297.84 33.33)"><tspan x="0" y="0">search</tspan></text>
+    <text class="cls-1" transform="translate(297.84 47.73)"><tspan x="0" y="0">ellipsoid</tspan></text>
+  </g>
+  <g>
+    <line class="cls-8" x1="416.79" y1="68.24" x2="429.11" y2="25.33"/>
+    <path class="cls-12" d="M428.9,26.07l-3.05.83v-.1l2.89-3.75,2.25-4.15-.3,4.71.47,4.71-.07.08-2.19-2.33Z"/>
+  </g>
+  <line class="cls-5" x1="403.95" y1="112.94" x2="416.79" y2="68.24"/>
+  <line class="cls-8" x1="387.29" y1="170.93" x2="403.95" y2="112.94"/>
+</svg>

doc/images/computation/kutta_cond_plate_dark.svg

@@ -0,0 +1,124 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 792.59 237.27">
+  <defs>
+    <style>
+      .cls-1 {
+        stroke-miterlimit: 4;
+        stroke-width: 16px;
+      }
+
+      .cls-1, .cls-2, .cls-3, .cls-4, .cls-5, .cls-6, .cls-7, .cls-8 {
+        fill: none;
+      }
+
+      .cls-1, .cls-5, .cls-6, .cls-8 {
+        stroke: #fff;
+      }
+
+      .cls-9 {
+        clip-path: url(#clippath);
+      }
+
+      .cls-2 {
+        stroke: #007b76;
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-3 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-3, .cls-7 {
+        stroke: #0076ba;
+      }
+
+      .cls-4, .cls-10, .cls-11, .cls-12 {
+        stroke-width: 0px;
+      }
+
+      .cls-5 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-6 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-7, .cls-8 {
+        stroke-width: 2px;
+      }
+
+      .cls-13 {
+        font-family: ArialMT, Arial;
+        font-size: 32px;
+      }
+
+      .cls-13, .cls-14, .cls-15 {
+        isolation: isolate;
+      }
+
+      .cls-13, .cls-15, .cls-11 {
+        fill: #fff;
+      }
+
+      .cls-15 {
+        font-family: STIXGeneral-Italic, STIXGeneral;
+        font-size: 32.22px;
+        font-style: italic;
+      }
+
+      .cls-10 {
+        fill: #99195e;
+      }
+
+      .cls-12 {
+        fill: #0076ba;
+      }
+    </style>
+    <clipPath id="clippath">
+      <rect class="cls-4" x="520.61" y="100.37" width="16.82" height="14.57"/>
+    </clipPath>
+  </defs>
+  <g>
+    <path class="cls-1" d="M153.2,197.67L659.31,46.23"/>
+    <path class="cls-2" d="M653.5,38.16c-69.23-65.09-392.9,103.98-501.22,149.95"/>
+    <path class="cls-7" d="M792.58,52.6c-39.16-.63-79.53-1.41-113.82-20.52-11.37-6.34-21.66-14.59-33.61-19.79-20.67-9-42.42-8.25-64.46-5.66C410.68,26.63,262.71,144.22,96.39,182.57c-31.65,7.3-63.9,11.55-96.37,12.26"/>
+    <path class="cls-7" d="M788.09,72.98c-34.82-2.64-69.65,1.7-103.92,8.29-25.81,4.96-51.35,11.2-76.72,18.09-154.05,41.88-302.08,108.19-461.01,125.56-47.05,5.14-94.43,5.87-141.68,8.47"/>
+    <path class="cls-12" d="M787.39,47.75l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M783.33,67.36l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M29.24,188.93l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M30.24,227.67l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M615.4,0l-9.6,4.8,9.6,4.8V0Z"/>
+    <path class="cls-12" d="M349.27,77.03l-6.8,8.3,10.72.46-3.92-8.76Z"/>
+    <path class="cls-12" d="M473.5,26.47l-7.5,7.68,10.64,1.39-3.15-9.07Z"/>
+    <path class="cls-3" d="M195.07,151.07l-1.58.58-.94.34"/>
+    <path class="cls-12" d="M191.84,147.14l-7.36,7.81,10.67,1.2-3.3-9.01Z"/>
+    <path class="cls-12" d="M613.28,92.82l-8.06,7.09,10.51,2.19-2.45-9.28Z"/>
+    <path class="cls-12" d="M442.34,145.91l-8.06,7.09,10.51,2.19-2.45-9.28Z"/>
+    <path class="cls-12" d="M212.2,210.07l-8.53,6.51,10.33,2.92-1.8-9.43Z"/>
+    <path class="cls-7" d="M553.05,59.87c11.3-6.46,36.25-12.77,22.33-20.73-17.28-9.89-181.61,49.61-169.38,59.08,8.52,6.6,46.77-6.02,72.27-14.91l.95-.33"/>
+    <path class="cls-12" d="M479.86,87.84l7.47-7.7-10.65-1.36,3.17,9.06Z"/>
+    <path class="cls-8" d="M405.92,122.4h219.45"/>
+    <path class="cls-11" d="M624.38,127.2l9.6-4.8-9.6-4.8v9.6Z"/>
+    <text class="cls-13" transform="translate(637.08 131.54)"><tspan x="0" y="0">v</tspan></text>
+    <path class="cls-6" d="M141.34,157.03l38.77-11.63.96-.29"/>
+    <path class="cls-11" d="M181.49,149.99l7.82-7.36-10.57-1.84,2.76,9.2Z"/>
+    <path class="cls-5" d="M141.31,157.86l-11.63-38.77-.29-.96"/>
+    <path class="cls-11" d="M134.28,117.7l-7.36-7.82-1.84,10.57,9.2-2.76Z"/>
+    <text class="cls-13" transform="translate(135.06 107.05)"><tspan x="0" y="0">y</tspan></text>
+    <text class="cls-13" transform="translate(181.5 131.54)"><tspan x="0" y="0">x</tspan></text>
+    <path class="cls-10" d="M672.67,54.95c2.03,1.97,2.03,5.17,0,7.14-2.03,1.97-5.32,1.97-7.35,0-2.03-1.97-2.03-5.17,0-7.14,2.03-1.97,5.32-1.97,7.35,0Z"/>
+    <path class="cls-10" d="M147.18,196.95c2.03,1.97,2.03,5.17,0,7.14s-5.32,1.97-7.35,0-2.03-5.17,0-7.14c2.03-1.97,5.32-1.97,7.35,0Z"/>
+    <path class="cls-8" d="M508.05,122.3c.72-19.73-3.57-30.4-12.89-31.99"/>
+    <text/>
+  </g>
+  <g class="cls-9">
+    <g class="cls-14">
+      <text class="cls-15" transform="translate(519.74 114.58)"><tspan x="0" y="0">α</tspan></text>
+    </g>
+  </g>
+</svg>

doc/includes/references.h

@@ -237,7 +237,7 @@ struct mjData_ {
 
   // computed by mj_fwdPosition/mj_comPos
   mjtNum* subtree_com;       // center of mass of each subtree                   (nbody x 3)
-  mjtNum* cdof;              // com-based motion axis of each dof                (nv x 6)
+  mjtNum* cdof;              // com-based motion axis of each dof (rot:lin)      (nv x 6)
   mjtNum* cinert;            // com-based body inertia and mass                  (nbody x 10)
 
   // computed by mj_fwdPosition/mj_flex
@@ -285,8 +285,8 @@ struct mjData_ {
   mjtNum* actuator_velocity; // actuator velocities                              (nu x 1)
 
   // computed by mj_fwdVelocity/mj_comVel
-  mjtNum* cvel;              // com-based velocity [3D rot; 3D tran]             (nbody x 6)
+  mjtNum* cvel;              // com-based velocity (rot:lin)                     (nbody x 6)
-  mjtNum* cdof_dot;          // time-derivative of cdof                          (nv x 6)
+  mjtNum* cdof_dot;          // time-derivative of cdof (rot:lin)                (nv x 6)
 
   // computed by mj_fwdVelocity/mj_rne (without acceleration)
   mjtNum* qfrc_bias;         // C(qpos,qvel)                                     (nv x 1)
@@ -459,10 +459,11 @@ typedef enum mjtGeom_ {           // type of geometric shape
   mjGEOM_ARROW1,                  // arrow without wedges
   mjGEOM_ARROW2,                  // arrow in both directions
   mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
   mjGEOM_FLEX,                    // flex
   mjGEOM_SKIN,                    // skin
   mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle connecting a frame
+  mjGEOM_TRIANGLE,                // triangle
 
   mjGEOM_NONE         = 1001      // missing geom type
 } mjtGeom;
@@ -570,7 +571,9 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_TEXT,                     // text
   mjOBJ_TUPLE,                    // tuple
   mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN                    // plugin instance
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT                       // number of object types
 } mjtObj;
 typedef enum mjtConstraint_ {     // type of constraint
   mjCNSTR_EQUALITY    = 0,        // equality constraint
@@ -1182,20 +1185,23 @@ struct mjModel_ {
   int*      skin_bonebodyid;      // body id of each bone                     (nskinbone x 1)
   int*      skin_bonevertid;      // mesh ids of vertices in each bone        (nskinbonevert x 1)
   float*    skin_bonevertweight;  // weights of vertices in each bone         (nskinbonevert x 1)
+  int*      skin_pathadr;         // address of asset path for skin; -1: none (nskin x 1)
 
   // height fields
-  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                  (nhfield x 4)
+  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                    (nhfield x 4)
-  int*      hfield_nrow;          // number of rows in grid                   (nhfield x 1)
+  int*      hfield_nrow;          // number of rows in grid                     (nhfield x 1)
-  int*      hfield_ncol;          // number of columns in grid                (nhfield x 1)
+  int*      hfield_ncol;          // number of columns in grid                  (nhfield x 1)
-  int*      hfield_adr;           // address in hfield_data                   (nhfield x 1)
+  int*      hfield_adr;           // address in hfield_data                     (nhfield x 1)
-  float*    hfield_data;          // elevation data                           (nhfielddata x 1)
+  float*    hfield_data;          // elevation data                             (nhfielddata x 1)
+  int*      hfield_pathadr;       // address of asset path for hfield; -1: none (nhfield x 1)
 
   // textures
-  int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
+  int*      tex_type;             // texture type (mjtTexture)                  (ntex x 1)
-  int*      tex_height;           // number of rows in texture image          (ntex x 1)
+  int*      tex_height;           // number of rows in texture image            (ntex x 1)
-  int*      tex_width;            // number of columns in texture image       (ntex x 1)
+  int*      tex_width;            // number of columns in texture image         (ntex x 1)
-  int*      tex_adr;              // address in rgb                           (ntex x 1)
+  int*      tex_adr;              // address in rgb                             (ntex x 1)
-  mjtByte*  tex_rgb;              // rgb (alpha = 1)                          (ntexdata x 1)
+  mjtByte*  tex_rgb;              // rgb (alpha = 1)                            (ntexdata x 1)
+  int*      tex_pathadr;         // address of asset path for texture; -1: none (ntex x 1)
 
   // materials
   int*      mat_texid;            // texture id; -1: none                     (nmat x 1)
@@ -1422,6 +1428,15 @@ struct mjpPlugin_ {
   // called by mjv_updateScene (optional)
   void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance);
 
+  // methods specific to actuators (optional)
+
+  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
+  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
+
+  // updates the actuator plugin's entries in act_dot
+  // called after native act_dot is computed and before the compute callback
+  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
+
   // methods specific to signed distance fields (optional)
 
   // signed distance from the surface
@@ -2155,6 +2170,7 @@ struct mjvSceneState_ {
     int nnames;
     int npaths;
     int nsensordata;
+    int narena;
 
     mjOption opt;
     mjVisual vis;
@@ -2226,6 +2242,7 @@ struct mjvSceneState_ {
     int* flex_vertadr;
     int* flex_vertnum;
     int* flex_elem;
+    int* flex_elemlayer;
     int* flex_elemadr;
     int* flex_elemnum;
     int* flex_elemdataadr;
@@ -2238,6 +2255,8 @@ struct mjvSceneState_ {
     mjtNum* flex_radius;
     float* flex_rgba;
 
+    int* hfield_pathadr;
+
     int* mesh_bvhadr;
     int* mesh_bvhnum;
     int* mesh_texcoordadr;
@@ -2264,6 +2283,9 @@ struct mjvSceneState_ {
     int* skin_bonebodyid;
     int* skin_bonevertid;
     float* skin_bonevertweight;
+    int* skin_pathadr;
+
+    int* tex_pathadr;
 
     int* mat_texid;
     mjtByte* mat_texuniform;
@@ -2362,6 +2384,7 @@ struct mjvSceneState_ {
     mjtNum* ten_length;
     mjtNum* wrap_xpos;
 
+    mjtNum* bvh_aabb_dyn;
     mjtByte* bvh_active;
     int* island_dofadr;
     int* island_dofind;
@@ -2373,6 +2396,7 @@ struct mjvSceneState_ {
 
     mjContact* contact;
     mjtNum* efc_force;
+    void* arena;
   } data;
 };
 typedef struct mjvSceneState_ mjvSceneState;
@@ -2435,6 +2459,7 @@ void mj_fwdAcceleration(const mjModel* m, mjData* d);
 void mj_fwdConstraint(const mjModel* m, mjData* d);
 void mj_Euler(const mjModel* m, mjData* d);
 void mj_RungeKutta(const mjModel* m, mjData* d, int N);
+void mj_implicit(const mjModel* m, mjData* d);
 void mj_invPosition(const mjModel* m, mjData* d);
 void mj_invVelocity(const mjModel* m, mjData* d);
 void mj_invConstraint(const mjModel* m, mjData* d);

doc/mjx.rst

@@ -181,9 +181,9 @@ The following features are **fully supported** in MJX:
    * - :ref:`Joint <mjtJoint>`
      - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
    * - :ref:`Transmission <mjtTrn>`
-     - ``TRN_JOINT``
+     - ``TRN_JOINT``, ``TRN_SITE``
    * - :ref:`Actuator Dynamics <mjtDyn>`
-     - ``NONE``, ``INTEGRATOR``, ``FILTER``
+     - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT``
    * - :ref:`Actuator Gain <mjtGain>`
      - ``FIXED``, ``AFFINE``
    * - :ref:`Actuator Bias <mjtBias>`
@@ -257,9 +257,9 @@ The following features are **unsupported**:
    * - Category
      - Feature
    * - :ref:`Transmission <mjtTrn>`
-     - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``
+     - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY``
    * - :ref:`Actuator Dynamics <mjtDyn>`
-     - ``FILTEREXACT``, ``USER``
+     - ``USER``
    * - :ref:`Actuator Gain <mjtGain>`
      - ``USER``
    * - :ref:`Actuator Bias <mjtBias>`
@@ -288,6 +288,22 @@ Single scene simulation
   Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been
   carefully optimized for CPU.  MJX works best when simulating thousands or tens of thousands of scenes in parallel.
 
+Collisions between large meshes
+  MJX supports collisions between convex mesh geometries. However the convex collision algorithms
+  in MJX are implemented differently than in MuJoCo. MJX uses a branchless version of the
+  `Separating Axis Test <https://ubm-twvideo01.s3.amazonaws.com/o1/vault/gdc2013/slides/822403Gregorius_Dirk_TheSeparatingAxisTest.pdf>`__
+  (SAT) to determine if geometries are colliding with convex meshes, while MuJoCo uses the Minkowski Portal Refinement (MPR)
+  algorithm as implemented in `libccd <https://github.com/danfis/libccd>`__.
+  SAT works well for smaller meshes but suffers in both runtime and memory for larger meshes.
+
+  For
+  collisions between convex meshes and primitives (spheres, capsules, planes), use **3000 vertices or less** for your convex meshes.
+  For collisions between convex meshes and other convex meshes, use **30 vertices or less**.
+  With careful
+  tuning, MJX can simulate scenes with mesh collisions -- see the MJX
+  `shadow hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
+  config for an example. Speeding up mesh collision detection is an active area of development for MJX.
+
 Large, complex scenes with many contacts
   Accelerators exhibit poor performance for
   `branching code <https://aschrein.github.io/jekyll/update/2019/06/13/whatsup-with-my-branches-on-gpu.html#tldr>`__.
@@ -296,7 +312,7 @@ Large, complex scenes with many contacts
   powerful as the one in MuJoCo.
 
   To see how this affects simulation, let us consider a physics scene with increasing numbers of humanoid bodies,
-  varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M1 Pro and a 64-core AMD 3995WX and time
+  varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M3 Max and a 64-core AMD 3995WX and time
   it using :ref:`testspeed<saTestspeed>`, using ``2 x numcore`` threads. We time the MJX simulation on an Nvidia
   A100 GPU using a batch size of 8192 and an 8-chip
   `v5 TPU <https://cloud.google.com/blog/products/compute/announcing-cloud-tpu-v5e-and-a3-gpus-in-ga>`__
@@ -306,18 +322,10 @@ Large, complex scenes with many contacts
      :width: 95%
      :align: center
 
-  The values for a single humanoid (leftmost datapoints) for the four timed architectures are **320K**, **1.8M**,
+  The values for a single humanoid (leftmost datapoints) for the four timed architectures are **650K**, **1.8M**,
   **950K** and **2.7M** steps per second, respectively. Note that as we increase the number of humanoids (which
   increases the number of potential contacts in a scene), MJX throughput decreases more rapidly than MuJoCo.
 
-Scenes with collisions between meshes with many vertices
-  MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of
-  `mesh collision algorithms <https://ubm-twvideo01.s3.amazonaws.com/o1/vault/gdc2013/slides/822403Gregorius_Dirk_TheSeparatingAxisTest.pdf>`__.
-  These algorithms work well for smaller meshes (with hundreds of vertices) but suffer with large meshes. With careful
-  tuning, MJX can simulate scenes with mesh collisions well -- see the MJX
-  `shadow hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
-  config for an example.
-
 .. _MjxPerformance:
 
 Performance tuning

doc/modeling.rst

@@ -204,16 +204,21 @@ cameras and lights.
 A related attribute is :ref:`compiler/angle<compiler-angle>`. It specifies whether angles in the MJCF file are expressed
 in degrees or radians (after compilation, angles are always expressed in radians).
 
+Positions are specified using
+
+:at:`pos`: :at-val:`real(3), "0 0 0"`
+   Position relative to parent.
+
 .. _COrientation:
 
 Frame orientations
-~~~~~~~~~~~~~~~~~~
+^^^^^^^^^^^^^^^^^^
 
 Several model elements have right-handed spatial frames associated with them. These are all the elements defined in the
 kinematic tree except for joints. A spatial frame is defined by its position and orientation. Specifying 3D positions is
 straightforward, but specifying 3D orientations can be challenging. This is why MJCF provides several alternative
-mechanisms. No matter which mechanism the user chooses, the frame orientation is always represented as a unit quaternion
+mechanisms. No matter which mechanism the user chooses, the frame orientation is always converted internally to a unit
-after compilation. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)`
+quaternion. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)`
 corresponds to the quaternion :math:`(\cos(a/2), \: \sin(a/2) \cdot (x, y, z))`. Also recall that every 3D orientation
 can be uniquely specified by a single 3D rotation by some angle around some axis.
 
@@ -267,6 +272,7 @@ approximately
 
 .. math::
    \ac + d \cdot (b v + k r) = (1 - d)\cdot \au
+   :label: eq:constraint
 
 Again, the parameters that are under the user's control are :math:`d, b, k`. The remaining quantities are functions of
 the system state and are computed automatically at each time step.
@@ -312,7 +318,15 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec
    units of :math:`\text{width}`. Note that when :math:`\text{power}` is 1, the function is linear regardless of the
    :math:`\text{midpoint}`.
 
-   |image0|
+   .. image:: images/modeling/impedance.png
+      :width: 600px
+      :align: center
+      :class: only-light
+
+   .. image:: images/modeling/impedance_dark.png
+      :width: 600px
+      :align: center
+      :class: only-dark
 
    These plots show the impedance :math:`d(r)` on the vertical axis, as a function of the constraint violation :math:`r`
    on the horizontal axis.
@@ -338,12 +352,9 @@ Next we explain the setting of the stiffness :math:`k` and damping :math:`b` whi
 
 .. admonition:: Intuitive description of the **reference acceleration**
 
-   The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in
+   The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in order to
-   order to rectify violation. For example, consider a contact between a motionless free body pulled down by gravity
+   rectify violation. Imagine a body dropped onto the plane. Upon impact the constraint will generate a normal force
-   onto a static plane geom. Since there is no motion, the penetration will be entirely determined by the impedance
+   which attempts to rectify the penetration using a particular motion; this motion is the reference acceleration.
-   while the reference has no effect. Now imagine that the body is dropped onto the plane. Upon impact the constraint
-   will generate a normal force which attempts to rectify the penetration using a particular motion; this motion is
-   the reference acceleration.
 
    Another way of understanding the reference acceleration is to think of the unmodeled deformation variables
    described in the :ref:`Computation chapter<soPrimal>`. Imagine two bodies pressed together, leading to deformation at
@@ -388,7 +399,12 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math:
    can go unstable. This is enforced internally, unless the :ref:`refsafe<option-flag-refsafe>` attribute of :ref:`flag
    <option-flag>` is set to false. The :math:`\text{dampratio}` parameter would normally be set to 1, corresponding to
    critical damping. Smaller values result in under-damped or bouncy constraints, while larger values result in
-   over-damped constraints.
+   over-damped constraints. Combining the above formula with :eq:`eq:constraint`, we can derive the following result.
+   If the reference acceleration is given using the positive number format and the impedance is constant
+   :math:`d = d_0 = d_\text{width}`, then the penetration depth at rest is
+
+   .. math::
+      r = \au \cdot (1 - d) \cdot \text{timeconst}^2 \cdot \text{dampratio}^2
 
    Next we describe the direct format where the two numbers are :math:`(-\text{stiffness}, -\text{damping})`. This
    allows direct control over restitution in particular. We still apply some scaling so that the same numbers can be
@@ -398,9 +414,15 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math:
    .. math::
       \begin{aligned}
       b &= \text{damping} / d_\text{width} \\
-      k &= \text{stiffness} / d_\text{width}^2 \\
+      k &= \text{stiffness} \cdot d(r) / d_\text{width}^2 \\
       \end{aligned}
 
+   Similarly to the above derivation, if the reference acceleration is given using the negative number format and the
+   impedance is constant, then the penetration depth at rest is
+
+   .. math::
+      r = \au \cdot (1 - d) \cdot \text{stiffness}
+
 .. tip::
    In the positive-value default format, the :math:`\text{timeconst}` parameter controls constraint **softness**.
    It is specified in units of time and means "how quickly is the constraint trying to resolve the violation". Larger
@@ -851,7 +873,15 @@ The advantage of the scaled quantities is that all muscles behave similarly in t
 captured by the Force-Length-Velocity (:math:`\text{\small FLV}`) function measured in many experimental papers. We
 approximate this function as follows:
 
-|image1|
+.. image:: images/modeling/musclemodel.png
+   :width: 650px
+   :align: center
+   :class: only-light
+
+.. image:: images/modeling/musclemodel_dark.png
+   :width: 650px
+   :align: center
+   :class: only-dark
 
 The function is in the form:
 
@@ -893,7 +923,15 @@ Before embarking on a mission to design more accurate :math:`\text{\small FLV}`
 operating range of the muscle has a bigger effect than the shape of the :math:`\text{\small FLV}` function, and in many
 cases this parameter is unknown. Below is a graphical illustration:
 
-|image2|
+.. image:: images/modeling/musclerange.png
+   :width: 500px
+   :align: center
+   :class: only-light
+
+.. image:: images/modeling/musclerange_dark.png
+   :width: 500px
+   :align: center
+   :class: only-dark
 
 This figure format is common in the biomechanics literature, showing the operating range of each muscle superimposed on
 the normalized :math:`\text{FL}` curve (ignore the vertical displacement). Our default range is shown in black. The blue
@@ -1297,7 +1335,9 @@ A flex is a collection of MuJoCo bodies that are connected with massless stretch
 capsules (1D flex), triangles (2D flex), or tetrahedra (3D flex). In all cases we allow a radius, which makes the
 elements smooth and also volumetric in 1D and 2D. The primitive elements are illustrated below:
 
-|flexelem|
+.. image:: images/modeling/flexelem.png
+   :width: 600px
+   :align: center
 
 Thus far these look like geoms. But the key difference is that they deform: as the bodies (vertices) move independently
 of each other, the shape of the elements changes in real time. Collisions and contact forces are now generalized to
@@ -1360,7 +1400,8 @@ In case of 3D flexes made of tetrahedra, it may be useful to examine how the fle
 a special visualization mode that peels off the outer layers. Below is an example with the Stanford Bunny. Note how it
 has smaller tetrahedra on the outside and larger ones on the inside. This mesh design makes sense, because we want the
 collision surface to be accurate, but on the inside we just need soft material properties - which require less spatial
-resolution.
+resolution. In order to convert a surface mesh to a tetrahedral mesh, we recommend open tools like the
+`fTetWild library <https://github.com/wildmeshing/fTetWild>`__.
 
 |bunny1| |bunny2|
 
@@ -1663,12 +1704,6 @@ in a visible way, and the energy fluctuates around the initial value instead of
    </worldbody>
 
 
-.. |image0| image:: images/modeling/impedance.png
-   :width: 600px
-.. |image1| image:: images/modeling/musclemodel.png
-   :width: 650px
-.. |image2| image:: images/modeling/musclerange.png
-   :width: 400px
 .. |image3| image:: images/modeling/tendonwraps.png
    :width: 500px
 .. |image4| image:: images/modeling/particle.png
@@ -1705,8 +1740,6 @@ in a visible way, and the energy fluctuates around the initial value instead of
    :height: 250px
 .. |particle| image:: images/models/particle.gif
    :width: 270px
-.. |flexelem| image:: images/modeling/flexelem.png
-   :width: 400px
 .. |bunny1| image:: images/modeling/bunny1.png
    :width: 300px
 .. |bunny2| image:: images/modeling/bunny2.png

doc/programming/extension.rst

@@ -199,6 +199,25 @@ When :ref:`mjData` is being copied via :ref:`mj_copyData`, MuJoCo will copy over
 code is responsible for setting up the plugin data for the newly copied :ref:`mjData`. To facilitate this, MuJoCo calls
 the ``copy`` callback from :ref:`mjpPlugin` for each plugin instance present.
 
+.. _exActuatorAct:
+
+Actuator activations
+""""""""""""""""""""
+
+When writing stateful actuator plugins, there are two choices for where to save the actuator state. One option is using
+``plugin_state`` as described above, and the other is to use ``mjData.act`` by implementing the ``actuator_actdim`` and
+``actuator_act_dot`` callbacks on :ref:`mjpPlugin`.
+
+When using the latter option, the actuator plugin's state will be added to ``mjData.act``, and MuJoCo will
+automatically integrate ``mjData.act_dot`` values between timesteps. One advantage of this approach is that
+finite-differencing functions like :ref:`mjd_transitionFD` will work as they do for native actuators. The
+``mjpPlugin.advance`` callback will be called after ``act_dot`` is integrated, and actuator plugins may overwrite
+the ``act`` values at that point, if Euler integration isn't appropriate.
+
+Users may specify the :ref:`dyntype<actuator-plugin-dyntype>` attribute on actuator plugins, to introduce a filter or
+an integrator between user inputs and actuator activations. When they do, the activation variable introduced by
+``dyntype`` will be placed *after* the plugin's activation variables in the ``act`` array.
+
 .. _exRegistration:
 
 Registration
@@ -248,8 +267,11 @@ A future version of this section will include:
 * Things that developers need to keep in mind in order to ensure that plugins function correctly when :ref:`mjData` is
   copied, stepped, or reset.
 
-Currently, there are three directories of first-party plugins:
+There are several first-party plugin directories:
 
+* **actuator:** The plugins in the `actuator/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/actuator>`__
+  directory implement custom actuators, so far only a PID controller. See the
+  `README <https://github.com/google-deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ for details.
 * **elasticity:** The plugins in the `elasticity/
   <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__ directory are passive forces based on
   continuum mechanics for 1-dimensional and 3-dimensional bodies. The 1D model is invariant under rotations and captures
@@ -269,11 +291,11 @@ Currently, there are three directories of first-party plugins:
   <https://github.com/google-deepmind/mujoco/blob/main/plugin/sdf/README.md>`__. The rest of this section will give more
   detail concerning the collision algorithm and the plugin engine interface.
 
-  Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent.
+  Collision points are found by minimizing the function A + B + abs(max(A, B)), where A and B are the two colliding
-  Because SDFs are non-convex, multiple starting points are required in order to converge to multiple local minima.
+  SDFs, via gradient descent. Because SDFs are non-convex, multiple starting points are required in order to converge to
-  The number of starting points is set using :ref:`sdf_initpoints<option-sdf_initpoints>`, and are
+  multiple local minima. The number of starting points is set using :ref:`sdf_initpoints<option-sdf_initpoints>`, and
-  initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes.
+  are initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes. The number of
-  The number of gradient descent iterations is set using :ref:`sdf_iterations<option-sdf_iterations>`.
+  gradient descent iterations is set using :ref:`sdf_iterations<option-sdf_iterations>`.
 
   While *exact* SDFs---encoding the precise signed distance to the surface---are preferred, collisions are possible with
   any function whose value vanishes at the surface and grows monotonically away from it, with a negative sign in the
@@ -323,7 +345,6 @@ loading functions. The :ref:`mjpResourceProvider` struct stores three types of f
 .. _Uniform Resource Identifier: https://en.wikipedia.org/wiki/Uniform_Resource_Identifier
 
 Resource prefix
-
   Resources are identified by prefixes in their name. The chosen prefix should have a valid `Uniform Resource
   Identifier`_ (URI) scheme syntax. Resource names should also have a valid URI syntax, however this isn't enforced. A
   resource name with the syntax ``{prefix}:{filename}`` will match a provider using the scheme ``prefix``.  For

doc/unity.rst

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

include/mujoco/mjdata.h

@@ -265,7 +265,7 @@ struct mjData_ {
 
   // computed by mj_fwdPosition/mj_comPos
   mjtNum* subtree_com;       // center of mass of each subtree                   (nbody x 3)
-  mjtNum* cdof;              // com-based motion axis of each dof                (nv x 6)
+  mjtNum* cdof;              // com-based motion axis of each dof (rot:lin)      (nv x 6)
   mjtNum* cinert;            // com-based body inertia and mass                  (nbody x 10)
 
   // computed by mj_fwdPosition/mj_flex
@@ -313,8 +313,8 @@ struct mjData_ {
   mjtNum* actuator_velocity; // actuator velocities                              (nu x 1)
 
   // computed by mj_fwdVelocity/mj_comVel
-  mjtNum* cvel;              // com-based velocity [3D rot; 3D tran]             (nbody x 6)
+  mjtNum* cvel;              // com-based velocity (rot:lin)                     (nbody x 6)
-  mjtNum* cdof_dot;          // time-derivative of cdof                          (nv x 6)
+  mjtNum* cdof_dot;          // time-derivative of cdof (rot:lin)                (nv x 6)
 
   // computed by mj_fwdVelocity/mj_rne (without acceleration)
   mjtNum* qfrc_bias;         // C(qpos,qvel)                                     (nv x 1)

include/mujoco/mjmodel.h

@@ -109,10 +109,11 @@ typedef enum mjtGeom_ {           // type of geometric shape
   mjGEOM_ARROW1,                  // arrow without wedges
   mjGEOM_ARROW2,                  // arrow in both directions
   mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
   mjGEOM_FLEX,                    // flex
   mjGEOM_SKIN,                    // skin
   mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle connecting a frame
+  mjGEOM_TRIANGLE,                // triangle
 
   mjGEOM_NONE         = 1001      // missing geom type
 } mjtGeom;
@@ -246,7 +247,9 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_TEXT,                     // text
   mjOBJ_TUPLE,                    // tuple
   mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN                    // plugin instance
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT                       // number of object types
 } mjtObj;
 
 
@@ -895,20 +898,23 @@ struct mjModel_ {
   int*      skin_bonebodyid;      // body id of each bone                     (nskinbone x 1)
   int*      skin_bonevertid;      // mesh ids of vertices in each bone        (nskinbonevert x 1)
   float*    skin_bonevertweight;  // weights of vertices in each bone         (nskinbonevert x 1)
+  int*      skin_pathadr;         // address of asset path for skin; -1: none (nskin x 1)
 
   // height fields
-  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                  (nhfield x 4)
+  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                    (nhfield x 4)
-  int*      hfield_nrow;          // number of rows in grid                   (nhfield x 1)
+  int*      hfield_nrow;          // number of rows in grid                     (nhfield x 1)
-  int*      hfield_ncol;          // number of columns in grid                (nhfield x 1)
+  int*      hfield_ncol;          // number of columns in grid                  (nhfield x 1)
-  int*      hfield_adr;           // address in hfield_data                   (nhfield x 1)
+  int*      hfield_adr;           // address in hfield_data                     (nhfield x 1)
-  float*    hfield_data;          // elevation data                           (nhfielddata x 1)
+  float*    hfield_data;          // elevation data                             (nhfielddata x 1)
+  int*      hfield_pathadr;       // address of asset path for hfield; -1: none (nhfield x 1)
 
   // textures
-  int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
+  int*      tex_type;             // texture type (mjtTexture)                  (ntex x 1)
-  int*      tex_height;           // number of rows in texture image          (ntex x 1)
+  int*      tex_height;           // number of rows in texture image            (ntex x 1)
-  int*      tex_width;            // number of columns in texture image       (ntex x 1)
+  int*      tex_width;            // number of columns in texture image         (ntex x 1)
-  int*      tex_adr;              // address in rgb                           (ntex x 1)
+  int*      tex_adr;              // address in rgb                             (ntex x 1)
-  mjtByte*  tex_rgb;              // rgb (alpha = 1)                          (ntexdata x 1)
+  mjtByte*  tex_rgb;              // rgb (alpha = 1)                            (ntexdata x 1)
+  int*      tex_pathadr;         // address of asset path for texture; -1: none (ntex x 1)
 
   // materials
   int*      mat_texid;            // texture id; -1: none                     (nmat x 1)

include/mujoco/mjplugin.h

@@ -107,6 +107,15 @@ struct mjpPlugin_ {
   // called by mjv_updateScene (optional)
   void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance);
 
+  // methods specific to actuators (optional)
+
+  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
+  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
+
+  // updates the actuator plugin's entries in act_dot
+  // called after native act_dot is computed and before the compute callback
+  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
+
   // methods specific to signed distance fields (optional)
 
   // signed distance from the surface

include/mujoco/mjvisualize.h

@@ -436,6 +436,7 @@ struct mjvSceneState_ {
     int nnames;
     int npaths;
     int nsensordata;
+    int narena;
 
     mjOption opt;
     mjVisual vis;
@@ -507,6 +508,7 @@ struct mjvSceneState_ {
     int* flex_vertadr;
     int* flex_vertnum;
     int* flex_elem;
+    int* flex_elemlayer;
     int* flex_elemadr;
     int* flex_elemnum;
     int* flex_elemdataadr;
@@ -519,6 +521,8 @@ struct mjvSceneState_ {
     mjtNum* flex_radius;
     float* flex_rgba;
 
+    int* hfield_pathadr;
+
     int* mesh_bvhadr;
     int* mesh_bvhnum;
     int* mesh_texcoordadr;
@@ -545,6 +549,9 @@ struct mjvSceneState_ {
     int* skin_bonebodyid;
     int* skin_bonevertid;
     float* skin_bonevertweight;
+    int* skin_pathadr;
+
+    int* tex_pathadr;
 
     int* mat_texid;
     mjtByte* mat_texuniform;
@@ -643,6 +650,7 @@ struct mjvSceneState_ {
     mjtNum* ten_length;
     mjtNum* wrap_xpos;
 
+    mjtNum* bvh_aabb_dyn;
     mjtByte* bvh_active;
     int* island_dofadr;
     int* island_dofind;
@@ -654,6 +662,7 @@ struct mjvSceneState_ {
 
     mjContact* contact;
     mjtNum* efc_force;
+    void* arena;
   } data;
 };
 typedef struct mjvSceneState_ mjvSceneState;

include/mujoco/mjxmacro.h

@@ -334,7 +334,7 @@
     X   ( int,     flex_vertbodyid,       nflexvert,     1                    ) \
     X   ( int,     flex_edge,             nflexedge,     2                    ) \
     XMJV( int,     flex_elem,             nflexelemdata, 1                    ) \
-    X   ( int,     flex_elemlayer,        nflexelem,     1                    ) \
+    XMJV( int,     flex_elemlayer,        nflexelem,     1                    ) \
     XMJV( int,     flex_shell,            nflexshelldata,1                    ) \
     X   ( int,     flex_evpair,           nflexevpair,   2                    ) \
     X   ( mjtNum,  flex_vert,             nflexvert,     3                    ) \
@@ -395,16 +395,19 @@
     XMJV( int,     skin_bonebodyid,       nskinbone,     1                    ) \
     XMJV( int,     skin_bonevertid,       nskinbonevert, 1                    ) \
     XMJV( float,   skin_bonevertweight,   nskinbonevert, 1                    ) \
+    XMJV( int,     skin_pathadr,          nskin,         1                    ) \
     X   ( mjtNum,  hfield_size,           nhfield,       4                    ) \
     X   ( int,     hfield_nrow,           nhfield,       1                    ) \
     X   ( int,     hfield_ncol,           nhfield,       1                    ) \
     X   ( int,     hfield_adr,            nhfield,       1                    ) \
     X   ( float,   hfield_data,           nhfielddata,   1                    ) \
+    XMJV( int,     hfield_pathadr,        nhfield,       1                    ) \
     X   ( int,     tex_type,              ntex,          1                    ) \
     X   ( int,     tex_height,            ntex,          1                    ) \
     X   ( int,     tex_width,             ntex,          1                    ) \
     X   ( int,     tex_adr,               ntex,          1                    ) \
     X   ( mjtByte, tex_rgb,               ntexdata,      1                    ) \
+    XMJV( int,     tex_pathadr,           ntex,          1                    ) \
     XMJV( int,     mat_texid,             nmat,          1                    ) \
     XMJV( mjtByte, mat_texuniform,        nmat,          1                    ) \
     XMJV( float,   mat_texrepeat,         nmat,          2                    ) \
@@ -611,7 +614,7 @@
     X   ( mjtNum,    qLD,               nM,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
     X   ( mjtNum,    qLDiagSqrtInv,     nv,          1           ) \
-    X   ( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
+    XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
     X   ( mjtNum,    flexedge_velocity, nflexedge,   1           ) \
     X   ( mjtNum,    ten_velocity,      ntendon,     1           ) \

include/mujoco/mujoco.h

@@ -24,7 +24,7 @@ extern "C" {
 #endif
 
 // header version; should match the library version as returned by mj_version()
-#define mjVERSION_HEADER 302
+#define mjVERSION_HEADER 312
 
 // needed to define size_t, fabs and log10
 #include <stdlib.h>
@@ -264,6 +264,9 @@ MJAPI void mj_Euler(const mjModel* m, mjData* d);
 // Runge-Kutta explicit order-N integrator.
 MJAPI void mj_RungeKutta(const mjModel* m, mjData* d, int N);
 
+// Implicit-in-velocity integrators.
+MJAPI void mj_implicit(const mjModel* m, mjData* d);
+
 // Run position-dependent computations in inverse dynamics.
 MJAPI void mj_invPosition(const mjModel* m, mjData* d);
 
@@ -1115,7 +1118,7 @@ MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 // Address of diagonal element i in band-dense matrix representation.
 MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
 
-// Eigenvalue decomposition of symmetric 3x3 matrix.
+// Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.
 MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
 
 // minimize 0.5*x'*H*x + x'*g  s.t. lower <= x <= upper, return rank or -1 if failed

introspect/enums.py

@@ -90,10 +90,11 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjGEOM_ARROW1', 101),
              ('mjGEOM_ARROW2', 102),
              ('mjGEOM_LINE', 103),
-             ('mjGEOM_FLEX', 104),
+             ('mjGEOM_LINEBOX', 104),
-             ('mjGEOM_SKIN', 105),
+             ('mjGEOM_FLEX', 105),
-             ('mjGEOM_LABEL', 106),
+             ('mjGEOM_SKIN', 106),
-             ('mjGEOM_TRIANGLE', 107),
+             ('mjGEOM_LABEL', 107),
+             ('mjGEOM_TRIANGLE', 108),
              ('mjGEOM_NONE', 1001),
          ]),
      )),
@@ -265,6 +266,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjOBJ_TUPLE', 23),
              ('mjOBJ_KEY', 24),
              ('mjOBJ_PLUGIN', 25),
+             ('mjNOBJECT', 26),
          ]),
      )),
     ('mjtConstraint',

introspect/enums_test.py

@@ -61,7 +61,7 @@ class EnumsTest(absltest.TestCase):
     self.assertEqual(enum_decl.values['mjGEOM_ARROW'], 100)
     self.assertEqual(enum_decl.values['mjGEOM_ARROW1'], 101)
     self.assertEqual(enum_decl.values['mjGEOM_ARROW2'], 102)
-    self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 107)
+    self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 108)
     # Skip a few...
 
     self.assertEqual(enum_decl.values['mjGEOM_NONE'], 1001)

introspect/functions.py

@@ -1156,6 +1156,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Runge-Kutta explicit order-N integrator.',
      )),
+    ('mj_implicit',
+     FunctionDecl(
+         name='mj_implicit',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+         ),
+         doc='Implicit-in-velocity integrators.',
+     )),
     ('mj_invPosition',
      FunctionDecl(
          name='mj_invPosition',
@@ -7348,7 +7368,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='Eigenvalue decomposition of symmetric 3x3 matrix.',
+         doc="Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.",  # pylint: disable=line-too-long
      )),
     ('mju_boxQP',
      FunctionDecl(

introspect/structs.py

@@ -2787,75 +2787,96 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  ),
                  doc='weights of vertices in each bone         (nskinbonevert x 1)',  # pylint: disable=line-too-long
              ),
+             StructFieldDecl(
+                 name='skin_pathadr',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='address of asset path for skin; -1: none (nskin x 1)',
+             ),
              StructFieldDecl(
                  name='hfield_size',
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
-                 doc='(x, y, z_top, z_bottom)                  (nhfield x 4)',
+                 doc='(x, y, z_top, z_bottom)                    (nhfield x 4)',
              ),
              StructFieldDecl(
                  name='hfield_nrow',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='number of rows in grid                   (nhfield x 1)',
+                 doc='number of rows in grid                     (nhfield x 1)',
              ),
              StructFieldDecl(
                  name='hfield_ncol',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='number of columns in grid                (nhfield x 1)',
+                 doc='number of columns in grid                  (nhfield x 1)',
              ),
              StructFieldDecl(
                  name='hfield_adr',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='address in hfield_data                   (nhfield x 1)',
+                 doc='address in hfield_data                     (nhfield x 1)',
              ),
              StructFieldDecl(
                  name='hfield_data',
                  type=PointerType(
                      inner_type=ValueType(name='float'),
                  ),
-                 doc='elevation data                           (nhfielddata x 1)',  # pylint: disable=line-too-long
+                 doc='elevation data                             (nhfielddata x 1)',  # pylint: disable=line-too-long
+             ),
+             StructFieldDecl(
+                 name='hfield_pathadr',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='address of asset path for hfield; -1: none (nhfield x 1)',
              ),
              StructFieldDecl(
                  name='tex_type',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='texture type (mjtTexture)                (ntex x 1)',
+                 doc='texture type (mjtTexture)                  (ntex x 1)',
              ),
              StructFieldDecl(
                  name='tex_height',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='number of rows in texture image          (ntex x 1)',
+                 doc='number of rows in texture image            (ntex x 1)',
              ),
              StructFieldDecl(
                  name='tex_width',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='number of columns in texture image       (ntex x 1)',
+                 doc='number of columns in texture image         (ntex x 1)',
              ),
              StructFieldDecl(
                  name='tex_adr',
                  type=PointerType(
                      inner_type=ValueType(name='int'),
                  ),
-                 doc='address in rgb                           (ntex x 1)',
+                 doc='address in rgb                             (ntex x 1)',
              ),
              StructFieldDecl(
                  name='tex_rgb',
                  type=PointerType(
                      inner_type=ValueType(name='mjtByte'),
                  ),
-                 doc='rgb (alpha = 1)                          (ntexdata x 1)',
+                 doc='rgb (alpha = 1)                            (ntexdata x 1)',  # pylint: disable=line-too-long
+             ),
+             StructFieldDecl(
+                 name='tex_pathadr',
+                 type=PointerType(
+                     inner_type=ValueType(name='int'),
+                 ),
+                 doc='address of asset path for texture; -1: none (ntex x 1)',
              ),
              StructFieldDecl(
                  name='mat_texid',
@@ -4472,7 +4493,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
-                 doc='com-based motion axis of each dof                (nv x 6)',  # pylint: disable=line-too-long
+                 doc='com-based motion axis of each dof (rot:lin)      (nv x 6)',  # pylint: disable=line-too-long
              ),
              StructFieldDecl(
                  name='cinert',
@@ -4682,14 +4703,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
-                 doc='com-based velocity [3D rot; 3D tran]             (nbody x 6)',  # pylint: disable=line-too-long
+                 doc='com-based velocity (rot:lin)                     (nbody x 6)',  # pylint: disable=line-too-long
              ),
              StructFieldDecl(
                  name='cdof_dot',
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
-                 doc='time-derivative of cdof                          (nv x 6)',  # pylint: disable=line-too-long
+                 doc='time-derivative of cdof (rot:lin)                (nv x 6)',  # pylint: disable=line-too-long
              ),
              StructFieldDecl(
                  name='qfrc_bias',
@@ -6313,6 +6334,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              type=ValueType(name='int'),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='narena',
+                             type=ValueType(name='int'),
+                             doc='',
+                         ),
                          StructFieldDecl(
                              name='opt',
                              type=ValueType(name='mjOption'),
@@ -6741,6 +6767,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='flex_elemlayer',
+                             type=PointerType(
+                                 inner_type=ValueType(name='int'),
+                             ),
+                             doc='',
+                         ),
                          StructFieldDecl(
                              name='flex_elemadr',
                              type=PointerType(
@@ -6818,6 +6851,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='hfield_pathadr',
+                             type=PointerType(
+                                 inner_type=ValueType(name='int'),
+                             ),
+                             doc='',
+                         ),
                          StructFieldDecl(
                              name='mesh_bvhadr',
                              type=PointerType(
@@ -6993,6 +7033,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='skin_pathadr',
+                             type=PointerType(
+                                 inner_type=ValueType(name='int'),
+                             ),
+                             doc='',
+                         ),
+                         StructFieldDecl(
+                             name='tex_pathadr',
+                             type=PointerType(
+                                 inner_type=ValueType(name='int'),
+                             ),
+                             doc='',
+                         ),
                          StructFieldDecl(
                              name='mat_texid',
                              type=PointerType(
@@ -7547,6 +7601,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='bvh_aabb_dyn',
+                             type=PointerType(
+                                 inner_type=ValueType(name='mjtNum'),
+                             ),
+                             doc='',
+                         ),
                          StructFieldDecl(
                              name='bvh_active',
                              type=PointerType(
@@ -7610,6 +7671,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='',
                          ),
+                         StructFieldDecl(
+                             name='arena',
+                             type=PointerType(
+                                 inner_type=ValueType(name='void'),
+                             ),
+                             doc='',
+                         ),
                      ),
                  ),
                  doc='',

mjx/mujoco/mjx/__init__.py

@@ -16,13 +16,24 @@
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
+from mujoco.mjx._src.constraint import count_constraints
 from mujoco.mjx._src.constraint import make_constraint
 from mujoco.mjx._src.device import device_get_into
 from mujoco.mjx._src.device import device_put
+from mujoco.mjx._src.forward import euler
 from mujoco.mjx._src.forward import forward
+from mujoco.mjx._src.forward import fwd_acceleration
+from mujoco.mjx._src.forward import fwd_actuation
+from mujoco.mjx._src.forward import fwd_position
+from mujoco.mjx._src.forward import fwd_velocity
+from mujoco.mjx._src.forward import rungekutta4
 from mujoco.mjx._src.forward import step
+from mujoco.mjx._src.io import get_data
 from mujoco.mjx._src.io import make_data
+from mujoco.mjx._src.io import put_data
+from mujoco.mjx._src.io import put_model
 from mujoco.mjx._src.passive import passive
+from mujoco.mjx._src.ray import ray
 from mujoco.mjx._src.smooth import com_pos
 from mujoco.mjx._src.smooth import com_vel
 from mujoco.mjx._src.smooth import crb
@@ -31,4 +42,5 @@ from mujoco.mjx._src.smooth import kinematics
 from mujoco.mjx._src.smooth import mul_m
 from mujoco.mjx._src.smooth import rne
 from mujoco.mjx._src.smooth import transmission
+from mujoco.mjx._src.solver import solve
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/collision_driver.py

@@ -40,7 +40,6 @@ from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import GeomType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
-import numpy as np
 
 
 # pair-wise collision functions
@@ -81,6 +80,12 @@ def _add_candidate(
   if t1 > t2:
     t1, t2, g1, g2 = t2, t1, g2, g1
 
+  # MuJoCo does not collide planes with other planes or hfields
+  if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
+    return
+  if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
+    return
+
   def mesh_key(i):
     convex_data = [[None] * m.ngeom] * 3
     if isinstance(m, Model):
@@ -284,13 +289,11 @@ def _collide_geoms(
       solimp=params.solimp,
       geom1=geom1,
       geom2=geom2,
-      dim=np.array([]),
-      efc_address=np.array([]),
   )
   return con
 
 
-def _max_contact_points(m: Model) -> int:
+def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> int:
   """Returns the maximum number of contact points when set as a numeric."""
   for i in range(m.nnumeric):
     name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0]
@@ -334,7 +337,7 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
   return candidate_set
 
 
-def ncon(m: Model) -> int:
+def ncon(m: Union[Model, mujoco.MjModel]) -> int:
   """Returns the number of contacts computed in MJX given a model."""
   if m.opt.disableflags & DisableBit.CONTACT:
     return 0
@@ -354,9 +357,8 @@ def ncon(m: Model) -> int:
 
 def collision(m: Model, d: Data) -> Data:
   """Collides geometries."""
-  ncon_ = ncon(m)
+  if ncon(m) == 0:
-  if ncon_ == 0:
+    return d.replace(contact=Contact.zero())
-    return d.replace(contact=Contact.zero(), ncon=0)
 
   candidate_set = collision_candidates(m)
 
@@ -376,13 +378,4 @@ def collision(m: Model, d: Data) -> Data:
     _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
     contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact)
 
-  if ncon_ != contact.dist.shape[0]:
+  return d.replace(contact=contact)
-    raise RuntimeError('Number of contacts does not match ncon.')
-
-  # TODO(robotics-simulation): move this logic to device_put
-  ns = d.ne + d.nf + d.nl
-  contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4))
-  # TODO(robotics-simulation): add support for other friction dimensions
-  contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32))
-
-  return d.replace(contact=contact, ncon=ncon_)

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -52,9 +52,9 @@ def _collide(
     mjcf: str, assets: Optional[Dict[str, str]] = None
 ) -> Tuple[mujoco.MjModel, mujoco.MjData, Model, Data]:
   m = mujoco.MjModel.from_xml_string(mjcf, assets or {})
-  mx = mjx.device_put(m)
+  mx = mjx.put_model(m)
   d = mujoco.MjData(m)
-  dx = mjx.device_put(d)
+  dx = mjx.put_data(m, d)
 
   mujoco.mj_step(m, d)
   collision_jit_fn = jax.jit(mjx.collision)
@@ -253,7 +253,6 @@ class CapsuleCollisionTest(parameterized.TestCase):
     self.assertGreater(c.dist[1], 0)
     # extract the contact point with penetration
     c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
-    c = c.replace(dim=c.dim[np.array([0])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
 
@@ -281,7 +280,6 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(-dx.contact.dist[2:], 0)
     # extract the contact points with penetration
     c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
-    c = c.replace(dim=c.dim[np.array([0, 1])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
 
@@ -339,7 +337,6 @@ class ConvexTest(absltest.TestCase):
     np.testing.assert_array_less(-dx.contact.dist[1:], 0)
     # extract the contact point with penetration
     c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
-    c = c.replace(dim=c.dim[np.array([0])])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
 
@@ -421,9 +418,9 @@ class BodyPairFilterTest(absltest.TestCase):
   def test_filter_parent_child(self):
     """Tests that parent-child collisions get filtered."""
     m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)
@@ -438,9 +435,9 @@ class BodyPairFilterTest(absltest.TestCase):
     """Tests that filterparent flag disables parent-child filtering."""
     m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
     m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_FILTERPARENT
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)
@@ -457,22 +454,14 @@ class NconTest(parameterized.TestCase):
   """Tests ncon."""
 
   def test_ncon(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    ncon = collision_driver.ncon(m)
-    d.qpos[2] = 0.0
+    self.assertEqual(ncon, 16)
-
-    mx = mjx.device_put(m)
-    ncon = collision_driver.ncon(mx)
-    self.assertEqual(ncon, 4)
 
   def test_disable_contact(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags |= DisableBit.CONTACT
-    d.qpos[2] = 0.0
+    ncon = collision_driver.ncon(m)
-
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
-    mx = mjx.device_put(m)
-    ncon = collision_driver.ncon(mx)
     self.assertEqual(ncon, 0)
 
 
@@ -503,12 +492,12 @@ class TopKContactTest(absltest.TestCase):
 
   def test_top_k_contacts(self):
     m = mujoco.MjModel.from_xml_string(self._CAPSULES)
-    mx_top_k = mjx.device_put(m)
+    mx_top_k = mjx.put_model(m)
     mx_all = mx_top_k.replace(
         nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([])
     )
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     collision_jit_fn = jax.jit(mjx.collision)
     kinematics_jit_fn = jax.jit(mjx.kinematics)
@@ -517,8 +506,8 @@ class TopKContactTest(absltest.TestCase):
     dx_all = collision_jit_fn(mx_all, dx)
     dx_top_k = collision_jit_fn(mx_top_k, dx)
 
-    self.assertEqual(dx_all.ncon, 3)
+    self.assertEqual(dx_all.contact.dist.shape, (3,))
-    self.assertEqual(dx_top_k.ncon, 2)
+    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/constraint.py

@@ -14,11 +14,12 @@
 # ==============================================================================
 """Core non-smooth constraint functions."""
 
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import math
 from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
@@ -276,7 +277,7 @@ def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
 def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
   """Calculates constraint rows for contacts."""
 
-  if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0:
+  if collision_driver.ncon(m) == 0:
     return None
 
   @jax.vmap
@@ -313,7 +314,9 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
   return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
 
 
-def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
+def count_constraints(
+    m: Union[Model, mujoco.MjModel]
+) -> Tuple[int, int, int, int]:
   """Returns equality, friction, limit, and contact constraint counts."""
   if m.opt.disableflags & DisableBit.CONSTRAINT:
     return 0, 0, 0, 0
@@ -333,10 +336,7 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
   else:
     nl = int(m.jnt_limited.sum())
 
-  if m.opt.disableflags & DisableBit.CONTACT:
+  nc = collision_driver.ncon(m) * 4
-    nc = 0
-  else:
-    nc = d.ncon * 4
 
   return ne, nf, nl, nc
 
@@ -344,10 +344,6 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
 def make_constraint(m: Model, d: Data) -> Data:
   """Creates constraint jacobians and other supporting data."""
 
-  ns = sum(count_constraints(m, d)[:-1])
-  # TODO(robotics-simulation): make device_put set nefc/efc_address instead
-  d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)})
-
   if m.opt.disableflags & DisableBit.CONSTRAINT:
     efcs = ()
   else:
@@ -364,7 +360,7 @@ def make_constraint(m: Model, d: Data) -> Data:
   if not efcs:
     z = jp.empty(0)
     d = d.replace(efc_J=jp.empty((0, m.nv)))
-    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z, nefc=0)
+    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
     return d
 
   efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -378,6 +374,6 @@ def make_constraint(m: Model, d: Data) -> Data:
 
   aref, r = fn(efc)
   d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
-  d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0])
+  d = d.replace(efc_frictionloss=efc.frictionloss)
 
   return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -15,164 +15,92 @@
 """Tests for constraint functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
-import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import constraint
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-from mujoco.mjx._src.types import SolverType
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3):
+# tolerance for difference between MuJoCo and MJX constraint calculations,
-  err_msg = f'mismatch: {name} at step {step} in {fname}'
+# mostly due to float precision
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
+_TOLERANCE = 5e-5
 
 
-class ConstraintTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10   # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(enumerate(test_util.TEST_FILES))
+
-  def test_constraints(self, seed, fname):
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ConstraintTest(absltest.TestCase):
+
+  def test_constraints(self):
     """Test constraints."""
-    np.random.seed(seed)
+    m = test_util.load_test_file('constraints.xml')
-
-    # exclude convex.xml since convex contacts are not exactly equivalent
-    if fname == 'convex.xml':
-      return
-
-    m = test_util.load_test_file(fname)
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
-    dx = mjx.make_data(mx)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
 
-    forward_jit_fn = jax.jit(mjx.forward)
+    dx = mjx.make_constraint(mx, dx)
-
+    nnz = dx.efc_J.any(axis=1)
-    # give the system a little kick to ensure we have non-identity rotations
+    _assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
-    d.qvel = np.random.random(m.nv)
+    _assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
-    for i in range(100):
+    _assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
-      dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel))
+    _assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      nnz_filter = dx.efc_J.any(axis=1)
-
-      mj_efc_j = d.efc_J.reshape((-1, m.nv))
-      mjx_efc_j = dx.efc_J[nnz_filter]
-      _assert_eq(mj_efc_j, mjx_efc_j, 'efc_J', i, fname)
-
-      mjx_efc_d = dx.efc_D[nnz_filter]
-      _assert_eq(d.efc_D, mjx_efc_d, 'efc_D', i, fname)
-
-      mjx_efc_aref = dx.efc_aref[nnz_filter]
-      _assert_eq(d.efc_aref, mjx_efc_aref, 'efc_aref', i, fname)
-
-      mjx_efc_frictionloss = dx.efc_frictionloss[nnz_filter]
-      _assert_eq(
-          d.efc_frictionloss,
-          mjx_efc_frictionloss,
-          'efc_frictionloss',
-          i,
-          fname,
-      )
-
-  _JNT_RANGE = """
-    <mujoco>
-      <worldbody>
-        <body pos="0 0 1">
-          <joint type="slide" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1"
-           damping="5e-4"/>
-          <geom type="box" size="0.2 0.15 0.1" mass="1"/>
-          <body>
-            <joint axis="0 1 0" damping="2e-6"/>
-            <geom type="capsule" fromto="0 0 0 0 0 1" size="0.045" mass=".1"/>
-          </body>
-        </body>
-      </worldbody>
-    </mujoco>
-  """
-
-  def test_jnt_range(self):
-    """Tests that mixed joint ranges are respected."""
-    # TODO(robotics-simulation): also test ball
-    m = mujoco.MjModel.from_xml_string(self._JNT_RANGE)
-    m.opt.solver = SolverType.CG.value
-    d = mujoco.MjData(m)
-    d.qpos = np.array([2.0, 15.0])
-
-    mx = mjx.device_put(m)
-    dx = mjx.device_put(d)
-    efc = jax.jit(constraint._instantiate_limit_slide_hinge)(mx, dx)
-
-    # first joint is outside the joint range
-    np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0)
-
-    # second joint has no range, so only one efc row
-    self.assertEqual(efc.J.shape[0], 1)
 
   def test_disable_refsafe(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
 
     timeconst = m.opt.timestep / 4.0  # timeconst < 2 * timestep
     solimp = jp.array([timeconst, 1.0])
     solref = jp.array([0.8, 0.99, 0.001, 0.2, 2])
     pos = jp.ones(3)
 
-    m.opt.disableflags = m.opt.disableflags | DisableBit.REFSAFE
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.REFSAFE
     mx = mjx.device_put(m)
     k, *_ = constraint._kbi(mx, solimp, solref, pos)
     self.assertEqual(k, 1 / (0.99**2 * timeconst**2))
 
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.REFSAFE
+  def test_disable_constraint(self):
-    mx = mjx.device_put(m)
+    m = test_util.load_test_file('constraints.xml')
-    k, *_ = constraint._kbi(mx, solimp, solref, pos)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT
-    self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2))
+    ne, nf, nl, nc = mjx.count_constraints(m)
-
+    self.assertEqual(ne, 0)
-  def test_disableconstraint(self):
+    self.assertEqual(nf, 0)
-    m = test_util.load_test_file('ant.xml')
+    self.assertEqual(nl, 0)
-    d = mujoco.MjData(m)
+    self.assertEqual(nc, 0)
-
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
-    dx = constraint.make_constraint(mx, dx)
-    self.assertGreater(dx.efc_J.shape[0], 1)
-
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT
-    mx = mjx.device_put(m)
-    dx = constraint.make_constraint(mx, dx)
     self.assertEqual(dx.efc_J.shape[0], 0)
 
   def test_disable_equality(self):
-    m = test_util.load_test_file('equality.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
-
+    ne, nf, nl, nc = mjx.count_constraints(m)
-    m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY
+    self.assertEqual(ne, 0)
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
+    self.assertEqual(nf, 0)
-    dx = constraint.make_constraint(mx, dx)
+    self.assertEqual(nl, 2)
-    self.assertEqual(dx.efc_J.shape[0], 0)
+    self.assertEqual(nc, 64)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 66)  # only joint range, contact
 
   def test_disable_contact(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
-    d = mujoco.MjData(m)
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
-    d.qpos[2] = 0.0
+    ne, nf, nl, nc = mjx.count_constraints(m)
-    mujoco.mj_forward(m, d)
+    self.assertEqual(ne, 10)
-
+    self.assertEqual(nf, 0)
-    m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT
+    self.assertEqual(nl, 2)
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
+    self.assertEqual(nc, 0)
-    dx = dx.tree_replace(
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
-        {'contact.frame': dx.contact.frame.reshape((-1, 3, 3))}
+    self.assertEqual(dx.efc_J.shape[0], 12)  # only joint range, limit
-    )
-    efc = constraint._instantiate_contact(mx, dx)
-    self.assertIsNotNone(efc)
-
-    m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
-    mx, dx = mjx.device_put(m), mjx.device_put(d)
-    efc = constraint._instantiate_contact(mx, dx)
-    self.assertIsNone(efc)
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/dataclasses.py

@@ -18,7 +18,7 @@ import copy
 import dataclasses
 
 import typing
-from typing import Dict, Optional, Sequence, TypeVar
+from typing import Any, Dict, Optional, Sequence, TypeVar
 import jax
 import numpy as np
 
@@ -62,7 +62,7 @@ def dataclass(clz: _T) -> _T:
 
     def to_meta(field, obj):
       val = getattr(obj, field.name)
-      return to_tup(val) if isinstance(val, np.ndarray) else val
+      return (to_tup(val), val.dtype) if isinstance(val, np.ndarray) else val
 
     def to_data(field, obj):
       return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name))
@@ -75,7 +75,7 @@ def dataclass(clz: _T) -> _T:
 
     def from_meta(field, meta):
       if field.type is np.ndarray:
-        return (field.name, np.array(meta))
+        return (field.name, np.array(meta[0], dtype=meta[1]))
       else:
         return (field.name, meta)
 
@@ -113,6 +113,10 @@ class PyTreeNode:
     # stub for pytype
     raise NotImplementedError
 
+  @classmethod
+  def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
+    return dataclasses.fields(cls)
+
   def tree_replace(
       self, params: Dict[str, Optional[jax.typing.ArrayLike]]
   ) -> 'PyTreeNode':

mjx/mujoco/mjx/_src/device.py

@@ -25,6 +25,7 @@ import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import mesh
 from mujoco.mjx._src import types
+import numpy as np
 
 _MJ_TYPE_ATTR = {
     mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,),
@@ -67,7 +68,7 @@ _TRANSFORMS = {
     (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
-    (types.Model, 'actuator_trnid'): lambda x: x[:, 0],
+    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-1] + (3, 3))  # pylint: disable=g-long-lambda
         if x is not None and x.shape[0] else jp.zeros((0, 3, 3))
@@ -78,6 +79,7 @@ _INVERSE_TRANSFORMS = {
     (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
+    (types.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)),
     (types.Contact, 'frame'): (
         lambda x: x.reshape(x.shape[:-2] + (9,))  # pylint: disable=g-long-lambda
         if x is not None and x.shape[0] else jp.zeros((0, 9))
@@ -182,6 +184,11 @@ def device_put(value):
   Returns:
     on-device MJX struct reflecting the input value
   """
+  warnings.warn(
+      'device_put is deprecated, use put_model and put_data instead',
+      category=DeprecationWarning,
+  )
+
   clz = _TYPE_MAP.get(type(value))
   if clz is None:
     raise NotImplementedError(f'{type(value)} is not supported for device_put.')
@@ -240,6 +247,10 @@ def device_get_into(result, value):
   Raises:
     RuntimeError: if result length doesn't match data batch size
   """
+  warnings.warn(
+      'device_get_into is deprecated, use get_data instead',
+      category=DeprecationWarning,
+  )
 
   value = jax.device_get(value)
 
@@ -263,9 +274,16 @@ def device_get_into(result, value):
 
   else:
     if isinstance(result, mujoco.MjData):
+      ncon = value.contact.dist.shape[0]
+      nefc = value.efc_J.shape[0]
       mujoco._functions._realloc_con_efc(  # pylint: disable=protected-access
-          result, ncon=value.ncon, nefc=value.nefc
+          result, ncon=ncon, nefc=nefc
       )
+      result.ncon = ncon
+      result.nefc = nefc
+      efc_start = nefc - ncon * 4
+      result.contact.efc_address[:] = np.arange(efc_start, nefc, 4)
+      result.contact.dim[:] = 3
 
     for f in dataclasses.fields(value):  # type: ignore
       if (type(value), f.name) in _DERIVED:

mjx/mujoco/mjx/_src/device_test.py

@@ -129,32 +129,26 @@ class ValidateInputTest(absltest.TestCase):
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
-  def test_trn(self):
-    m = test_util.load_test_file('ant.xml')
-    m.actuator_trntype[0] = mujoco.mjtTrn.mjTRN_SITE
-    with self.assertRaises(NotImplementedError):
-      mjx.device_put(m)
-
   def test_dyn(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_gain(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_bias(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
     with self.assertRaises(NotImplementedError):
       mjx.device_put(m)
 
   def test_condim(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('constraints.xml')
     for i in [1, 4, 6]:
       m.geom_condim[0] = i
       with self.assertRaises(NotImplementedError):

mjx/mujoco/mjx/_src/forward.py

@@ -60,7 +60,7 @@ def named_scope(fn, name: str = ''):
 
 
 @named_scope
-def _position(m: Model, d: Data) -> Data:
+def fwd_position(m: Model, d: Data) -> Data:
   """Position-dependent computations."""
   # TODO(robotics-simulation): tendon
   d = smooth.kinematics(m, d)
@@ -74,7 +74,7 @@ def _position(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _velocity(m: Model, d: Data) -> Data:
+def fwd_velocity(m: Model, d: Data) -> Data:
   """Velocity-dependent computations."""
   d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel)
   d = smooth.com_vel(m, d)
@@ -84,7 +84,7 @@ def _velocity(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _actuation(m: Model, d: Data) -> Data:
+def fwd_actuation(m: Model, d: Data) -> Data:
   """Actuation-dependent computations."""
   if not m.nu or m.opt.disableflags & DisableBit.ACTUATION:
     return d.replace(
@@ -107,7 +107,7 @@ def _actuation(m: Model, d: Data) -> Data:
       act_dot = jp.array(0.0)
     elif dyn_typ == DynType.INTEGRATOR:
       act_dot = ctrl
-    elif dyn_typ == DynType.FILTER:
+    elif dyn_typ in (DynType.FILTER, DynType.FILTEREXACT):
       act_dot = (ctrl - act) / jp.clip(dyn_prm[0], mujoco.mjMINVAL)
     else:
       raise NotImplementedError(f'dyntype {dyn_typ.name} not implemented.')
@@ -190,7 +190,7 @@ def _actuation(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _acceleration(m: Model, d: Data) -> Data:
+def fwd_acceleration(m: Model, d: Data) -> Data:
   """Add up all non-constraint forces, compute qacc_smooth."""
   qfrc_applied = d.qfrc_applied + support.xfrc_accumulate(m, d)
   qfrc_smooth = d.qfrc_passive - d.qfrc_bias + d.qfrc_actuator + qfrc_applied
@@ -228,6 +228,34 @@ def _integrate_pos(
   return jp.concatenate(qs) if qs else jp.empty((0,))
 
 
+def _next_activation(m: Model, d: Data, act_dot: jax.Array) -> jax.Array:
+  """Returns the next act given the current act_dot, after clamping."""
+  act = d.act
+
+  if not m.na:
+    return act
+
+  actrange = jp.where(
+      m.actuator_actlimited[:, None],
+      m.actuator_actrange,
+      jp.array([-jp.inf, jp.inf]),
+  )
+
+  def fn(dyntype, dynprm, act, act_dot, actrange):
+    if dyntype == DynType.FILTEREXACT:
+      tau = jp.clip(dynprm[0], a_min=mujoco.mjMINVAL)
+      act = act + act_dot * tau * (1 - jp.exp(-m.opt.timestep / tau))
+    else:
+      act = act + act_dot * m.opt.timestep
+    act = jp.clip(act, actrange[0], actrange[1])
+    return act
+
+  args = (m.actuator_dyntype, m.actuator_dynprm, act, act_dot, actrange)
+  act = scan.flat(m, fn, 'uuaau', 'a', *args, group_by='u')
+
+  return act.reshape(m.na)
+
+
 @named_scope
 def _advance(
     m: Model,
@@ -237,16 +265,7 @@ def _advance(
     qvel: Optional[jax.Array] = None,
 ) -> Data:
   """Advance state and time given activation derivatives and acceleration."""
-  act = d.act
+  act = _next_activation(m, d, act_dot)
-  if m.na:
-    act = d.act + act_dot * m.opt.timestep
-    actrange = jp.where(
-        m.actuator_actlimited[:, None],
-        m.actuator_actrange,
-        jp.array([-jp.inf, jp.inf]),
-    )
-    fn = lambda act, actrange: jp.clip(act, actrange[0], actrange[1])
-    act = scan.flat(m, fn, 'au', 'a', act, actrange, group_by='u')
 
   # advance velocities
   d = d.replace(qvel=d.qvel + qacc * m.opt.timestep)
@@ -263,7 +282,7 @@ def _advance(
 
 
 @named_scope
-def _euler(m: Model, d: Data) -> Data:
+def euler(m: Model, d: Data) -> Data:
   """Euler integrator, semi-implicit in velocity."""
   # integrate damping implicitly
   qacc = d.qacc
@@ -277,7 +296,7 @@ def _euler(m: Model, d: Data) -> Data:
 
 
 @named_scope
-def _rungekutta4(m: Model, d: Data) -> Data:
+def rungekutta4(m: Model, d: Data) -> Data:
   """Runge-Kutta explicit order 4 integrator."""
   d_t0 = d
   # pylint: disable=invalid-name
@@ -323,10 +342,10 @@ def _rungekutta4(m: Model, d: Data) -> Data:
 @named_scope
 def forward(m: Model, d: Data) -> Data:
   """Forward dynamics."""
-  d = _position(m, d)
+  d = fwd_position(m, d)
-  d = _velocity(m, d)
+  d = fwd_velocity(m, d)
-  d = _actuation(m, d)
+  d = fwd_actuation(m, d)
-  d = _acceleration(m, d)
+  d = fwd_acceleration(m, d)
 
   if d.efc_J.size == 0:
     d = d.replace(qacc=d.qacc_smooth)
@@ -343,9 +362,9 @@ def step(m: Model, d: Data) -> Data:
   d = forward(m, d)
 
   if m.opt.integrator == IntegratorType.EULER:
-    d = _euler(m, d)
+    d = euler(m, d)
   elif m.opt.integrator == IntegratorType.RK4:
-    d = _rungekutta4(m, d)
+    d = rungekutta4(m, d)
   else:
     raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
 

mjx/mujoco/mjx/_src/forward_test.py

@@ -15,77 +15,75 @@
 """Tests for forward functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
 import jax
-from jax import numpy as jp
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import forward
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3):
+# tolerance for difference between MuJoCo and MJX forward calculations - mostly
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+# due to float precision
-  a, b = getattr(a, attr), getattr(b, attr)
+_TOLERANCE = 1e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class ForwardTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(
-      filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
-  )
-  def test_forward(self, fname):
-    """Test mujoco mj forward function matches mujoco_mjx forward function."""
-    np.random.seed(test_util.TEST_FILES.index(fname))
 
-    m = test_util.load_test_file(fname)
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ForwardTest(absltest.TestCase):
+
+  def test_forward(self):
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    # apply some control and xfrc input
-    dx = mjx.make_data(mx)
+    d.ctrl = np.array([-18, 0.59, 0.47])
-    forward_jit_fn = jax.jit(mjx.forward)
+    d.xfrc_applied[0, 2] = 0.1  # torque
+    d.xfrc_applied[1, 4] = 0.3  # linear force
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
+    mujoco.mj_forward(m, d)
 
-    # give the system a little kick to ensure we have non-identity rotations
+    mx = mjx.put_model(m)
-    d.qvel = np.random.random(m.nv) * 0.05
-    for i in range(100):
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
 
-      _assert_attr_eq(d, dx, 'qfrc_smooth', i, fname)
+    # fwd_actuation
-      _assert_attr_eq(d, dx, 'qacc_smooth', i, fname)
+    dx = jax.jit(mjx.fwd_actuation)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'act_dot')
+    _assert_attr_eq(d, dx, 'qfrc_actuator')
 
-  @parameterized.parameters(
+    # fwd_accleration (fwd_position and fwd_velocity already tested elsewhere)
-      filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
+    dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d))
-  )
+    _assert_attr_eq(d, dx, 'qfrc_smooth')
-  def test_step(self, fname):
+    _assert_attr_eq(d, dx, 'qacc_smooth')
-    """Test mujoco mj step matches mujoco_mjx step."""
-    np.random.seed(test_util.TEST_FILES.index(fname))
-    m = test_util.load_test_file(fname)
-    step_jit_fn = jax.jit(forward.step)
 
-    mx = mjx.device_put(m)
+    # euler
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
+    mujoco.mj_Euler(m, d)
+    _assert_attr_eq(d, dx, 'act')
+    _assert_attr_eq(d, dx, 'qpos')
+    _assert_attr_eq(d, dx, 'time')
+
+  def test_step(self):
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    # give the system a little kick to ensure we have non-identity rotations
+    # apply some control and xfrc input
-    d.qvel = np.random.normal(m.nv) * 0.05
+    d.ctrl = np.array([-18, 0.59, 0.47])
-    for i in range(100):
+    d.xfrc_applied[0, 2] = 0.1  # torque
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    d.xfrc_applied[1, 4] = 0.3  # linear force
-      qpos, qvel = d.qpos, d.qvel
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel = qpos, qvel
-      dx = mjx.device_put(d)
 
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = step_jit_fn(mx, dx)
+    dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
-
+    mujoco.mj_step(m, d)
-      _assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2)
+    _assert_attr_eq(d, dx, 'act')
-      _assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2)
+    _assert_attr_eq(d, dx, 'time')
-      _assert_attr_eq(d, dx, 'act', i, fname)
+    _assert_attr_eq(d, dx, 'qvel')
-      _assert_attr_eq(d, dx, 'time', i, fname)
+    _assert_attr_eq(d, dx, 'qpos')
 
   def test_rk4(self):
     m = mujoco.MjModel.from_xml_string("""
@@ -94,7 +92,6 @@ class ForwardTest(parameterized.TestCase):
             <flag constraint="disable"/>
           </option>
           <worldbody>
-            <light pos="0 0 1"/>
             <geom type="plane" size="1 1 .01" pos="0 0 -1"/>
             <body pos="0.15 0 0">
               <joint type="hinge" axis="0 1 0"/>
@@ -107,41 +104,75 @@ class ForwardTest(parameterized.TestCase):
           </worldbody>
         </mujoco>
         """)
-    step_jit_fn = jax.jit(forward.step)
 
-    mx = mjx.device_put(m)
     d = mujoco.MjData(m)
     # give the system a little kick to ensure we have non-identity rotations
-    d.qvel = np.random.normal(m.nv) * 0.05
+    d.qvel = np.array([0.2, -0.1])
-    for i in range(100):
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    mujoco.mj_forward(m, d)
-      qpos, qvel = d.qpos, d.qvel
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel = qpos, qvel
-      dx = mjx.device_put(d)
 
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = step_jit_fn(mx, dx)
+    dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d))
+    mujoco.mj_RungeKutta(m, d, 4)
 
-      _assert_attr_eq(d, dx, 'qvel', i, 'test_rk4', atol=1e-2)
+    _assert_attr_eq(d, dx, 'qvel')
-      _assert_attr_eq(d, dx, 'qpos', i, 'test_rk4', atol=1e-2)
+    _assert_attr_eq(d, dx, 'qpos')
-      _assert_attr_eq(d, dx, 'act', i, 'test_rk4')
+    _assert_attr_eq(d, dx, 'act')
-      _assert_attr_eq(d, dx, 'time', i, 'test_rk4')
+    _assert_attr_eq(d, dx, 'time')
 
   def test_disable_eulerdamp(self):
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
-    m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP
+    self.assertTrue((m.dof_damping > 0).any())
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EULERDAMP
 
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    d.qvel[:] = 1.0
-    self.assertTrue((mx.dof_damping > 0).any())
+    d.qacc[:] = 1.0
-    dx = mjx.device_put(d)
+    mx = mjx.put_model(m)
-    dx = jax.jit(forward.forward)(mx, dx)
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
 
-    dx = dx.replace(qvel=jp.ones_like(dx.qvel), qacc=jp.ones_like(dx.qacc))
-    dx = jax.jit(forward._euler)(mx, dx)
     np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep)
 
 
+class ActuatorTest(absltest.TestCase):
+  _DYN_XML = """
+    <mujoco>
+      <compiler autolimits="true"/>
+      <worldbody>
+        <body name="box">
+          <joint name="slide1" type="slide" axis="1 0 0" />
+          <joint name="slide2" type="slide" axis="0 1 0" />
+          <joint name="slide3" type="slide" axis="0 0 1" />
+          <joint name="slide4" type="slide" axis="1 1 0" />
+          <geom type="box" size=".05 .05 .05" mass="1"/>
+        </body>
+      </worldbody>
+      <actuator>
+        <general joint="slide1" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide2" dyntype="integrator" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide3" dyntype="filter" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide4" dyntype="filterexact" dynprm="0.1" gainprm="1.1" />
+      </actuator>
+    </mujoco>
+  """
+
+  def test_dyntype(self):
+    m = mujoco.MjModel.from_xml_string(self._DYN_XML)
+    d = mujoco.MjData(m)
+    d.ctrl = np.array([1.5, 1.5, 1.5, 1.5])
+    d.act = np.array([0.5, 0.5, 0.5])
+
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_fwdActuation(m, d)
+    dx = jax.jit(mjx.fwd_actuation)(mx, dx)
+    _assert_attr_eq(d, dx, 'act_dot')
+
+    mujoco.mj_Euler(m, d)
+    dx = jax.jit(mjx.euler)(mx, dx)
+    _assert_attr_eq(d, dx, 'act')
+
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/io.py

@@ -14,92 +14,339 @@
 # ==============================================================================
 """Functions to initialize, load, or save data."""
 
+import copy
+from typing import List, Union
+
+import jax
 from jax import numpy as jp
+import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
-# pylint: disable=g-importing-member
+from mujoco.mjx._src import mesh
-from mujoco.mjx._src.types import Contact
+from mujoco.mjx._src import types
-from mujoco.mjx._src.types import Data
-from mujoco.mjx._src.types import Model
-# pylint: enable=g-importing-member
 import numpy as np
 
 
-def make_data(m: Model) -> Data:
+def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
-  """Allocate and initialize Data."""
+  """Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
+  if o.integrator not in set(types.IntegratorType):
+    raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
 
-  # create first d to get num contacts and nc
+  if o.cone not in set(types.ConeType):
-  d = Data(
+    raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
-      solver_niter=jp.array(0, dtype=jp.int32),
+
-      ne=0,
+  if o.solver not in set(types.SolverType):
-      nf=0,
+    raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
-      nl=0,
+
-      nefc=0,
+  for i in range(mujoco.mjtEnableBit.mjNENABLE):
-      ncon=0,
+    if o.enableflags & 2**i:
-      time=jp.zeros((), dtype=jp.float32),
+      raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
-      qpos=m.qpos0,
+
-      qvel=jp.zeros(m.nv, dtype=jp.float32),
+  static_fields = {
-      act=jp.zeros(m.na, dtype=jp.float32),
+      f.name: copy.copy(getattr(o, f.name))
-      qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32),
+      for f in types.Option.fields()
-      ctrl=jp.zeros(m.nu, dtype=jp.float32),
+      if f.type in (int, bytes, np.ndarray)
-      qfrc_applied=jp.zeros(m.nv, dtype=jp.float32),
+  }
-      xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32),
+  static_fields['integrator'] = types.IntegratorType(o.integrator)
-      eq_active=jp.zeros(m.neq, dtype=jp.int32),
+  static_fields['cone'] = types.ConeType(o.cone)
-      qacc=jp.zeros(m.nv, dtype=jp.float32),
+  static_fields['solver'] = types.SolverType(o.solver)
-      act_dot=jp.zeros(m.na, dtype=jp.float32),
+  static_fields['disableflags'] = types.DisableBit(o.disableflags)
-      xpos=jp.zeros((m.nbody, 3), dtype=jp.float32),
+
-      xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
+  device_fields = {
-      xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
+      f.name: copy.copy(getattr(o, f.name))
-      xipos=jp.zeros((m.nbody, 3), dtype=jp.float32),
+      for f in types.Option.fields()
-      ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32),
+      if f.type is jax.Array
-      xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32),
+  }
-      xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32),
+  device_fields = jax.device_put(device_fields, device=device)
-      geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
+
-      geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
+  has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
-      subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32),
+
-      cdof=jp.zeros((m.nv, 6), dtype=jp.float32),
+  return types.Option(
-      cinert=jp.zeros((m.nbody, 10), dtype=jp.float32),
+      has_fluid_params=has_fluid_params,
-      actuator_length=jp.zeros(m.nu, dtype=jp.float32),
+      **static_fields,
-      actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
+      **device_fields,
-      crb=jp.zeros((m.nbody, 10), dtype=jp.float32),
-      qM=jp.zeros(m.nM, dtype=jp.float32),
-      qLD=jp.zeros(m.nM, dtype=jp.float32),
-      qLDiagInv=jp.zeros(m.nv, dtype=jp.float32),
-      qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32),
-      contact=Contact.zero(),
-      efc_J=jp.zeros((), dtype=jp.float32),
-      efc_frictionloss=jp.zeros((), dtype=jp.float32),
-      efc_D=jp.zeros((), dtype=jp.float32),
-      actuator_velocity=jp.zeros(m.nu, dtype=jp.float32),
-      cvel=jp.zeros((m.nbody, 6), dtype=jp.float32),
-      cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32),
-      qfrc_bias=jp.zeros(m.nv, dtype=jp.float32),
-      qfrc_passive=jp.zeros(m.nv, dtype=jp.float32),
-      efc_aref=jp.zeros((), dtype=jp.float32),
-      actuator_force=jp.zeros(m.nu, dtype=jp.float32),
-      qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32),
-      qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32),
-      qacc_smooth=jp.zeros(m.nv, dtype=jp.float32),
-      qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32),
-      qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32),
-      efc_force=jp.zeros((), dtype=jp.float32),
   )
 
-  # get contact data with correct shapes
-  ncon = collision_driver.ncon(m)
-  d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon)
-  d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)})
 
-  ne, nf, nl, nc = constraint.count_constraints(m, d)
+def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
-  d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc)
+  """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
-  ns = ne + nf + nl
+  return types.Statistic(
-  d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)})
+      meaninertia=jax.device_put(s.meaninertia, device=device)
-  d = d.replace(
+  )
-      efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32),
+
-      efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32),
+
-      efc_D=jp.zeros(d.nefc, dtype=jp.float32),
+def put_model(m: mujoco.MjModel, device=None) -> types.Model:
-      efc_aref=jp.zeros(d.nefc, dtype=jp.float32),
+  """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
-      efc_force=jp.zeros(d.nefc, dtype=jp.float32),
+
+  if m.ntendon:
+    raise NotImplementedError('tendons are not supported')
+
+  if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
+    raise NotImplementedError('only condim=3 is supported')
+
+  # check collision geom types
+  for g1, g2, *_ in collision_driver.collision_candidates(m):
+    if collision_driver.get_collision_fn((g1, g2)) is None:
+      g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
+      raise NotImplementedError(f'({g1}, {g2}) has no collision function')
+
+  for enum_field, enum_type, mj_type in (
+      (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
+      (m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
+      (m.actuator_gaintype, types.GainType, mujoco.mjtGain),
+      (m.actuator_trntype, types.TrnType, mujoco.mjtTrn),
+      (m.eq_type, types.EqType, mujoco.mjtEq),
+  ):
+    missing = set(enum_field) - set(enum_type)
+    if missing:
+      raise NotImplementedError(
+          f'{[mj_type(m) for m in missing]} not supported'
+      )
+
+  opt = _put_option(m.opt, device=device)
+  stat = _put_statistic(m.stat, device=device)
+
+  static_fields = {
+      f.name: getattr(m, f.name)
+      for f in types.Model.fields()
+      if f.type in (int, bytes, np.ndarray)
+  }
+  static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4))
+  static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4))
+
+  device_fields = {
+      f.name: copy.copy(getattr(m, f.name))  # copy because device_put is async
+      for f in types.Model.fields()
+      if f.type is jax.Array
+  }
+  device_fields.update(mesh.get(m))
+  device_fields = jax.device_put(device_fields, device=device)
+
+  return types.Model(
+      opt=opt,
+      stat=stat,
+      **static_fields,
+      **device_fields,
+  )
+
+
+def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
+  """Allocate and initialize Data."""
+
+  ncon = collision_driver.ncon(m)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  zero_nv = jp.zeros(m.nv, dtype=jp.float32)
+  zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
+  zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32)
+  zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32)
+  zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32)
+  zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32)
+  zero_nefc = jp.zeros(nefc, dtype=jp.float32)
+  zero_na = jp.zeros(m.na, dtype=jp.float32)
+  zero_nu = jp.zeros(m.nu, dtype=jp.float32)
+  zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32)
+  zero_nm = jp.zeros(m.nM, dtype=jp.float32)
+
+  # create first d to get num contacts and nc
+  d = types.Data(
+      solver_niter=jp.array(0, dtype=jp.int32),
+      time=jp.array(0.0),
+      qpos=jp.array(m.qpos0),
+      qvel=zero_nv,
+      act=zero_na,
+      qacc_warmstart=zero_nv,
+      ctrl=zero_nu,
+      qfrc_applied=zero_nv,
+      xfrc_applied=zero_nbody_6,
+      eq_active=jp.zeros(m.neq, dtype=jp.int32),
+      qacc=zero_nv,
+      act_dot=zero_na,
+      xpos=zero_nbody_3,
+      xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
+      xmat=zero_nbody_3_3,
+      xipos=zero_nbody_3,
+      ximat=zero_nbody_3_3,
+      xanchor=zero_njnt_3,
+      xaxis=zero_njnt_3,
+      geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
+      geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
+      site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
+      site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
+      subtree_com=zero_nbody_3,
+      cdof=zero_nv_6,
+      cinert=zero_nbody_10,
+      actuator_length=zero_nu,
+      actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
+      crb=zero_nbody_10,
+      qM=zero_nm,
+      qLD=zero_nm,
+      qLDiagInv=zero_nv,
+      qLDiagSqrtInv=zero_nv,
+      contact=types.Contact.zero(ncon),
+      efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
+      efc_frictionloss=zero_nefc,
+      efc_D=zero_nefc,
+      actuator_velocity=zero_nu,
+      cvel=zero_nbody_6,
+      cdof_dot=zero_nv_6,
+      qfrc_bias=zero_nv,
+      qfrc_passive=zero_nv,
+      efc_aref=zero_nefc,
+      qfrc_actuator=zero_nv,
+      qfrc_smooth=zero_nv,
+      qacc_smooth=zero_nv,
+      qfrc_constraint=zero_nv,
+      qfrc_inverse=zero_nv,
+      efc_force=zero_nefc,
   )
 
   return d
+
+
+def _get_contact(
+    c: mujoco._structs._MjContactList,
+    cx: types.Contact,
+    efc_start: int,
+):
+  """Converts mjx.Contact to mujoco._structs._MjContactList."""
+  con_id = np.nonzero(cx.dist <= 0)[0]
+  for field in types.Contact.fields():
+    value = getattr(cx, field.name)[con_id]
+    if field.name == 'frame':
+      value = value.reshape((-1, 9))
+    getattr(c, field.name)[:] = value
+
+  ncon = cx.dist.shape[0]
+  c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
+
+
+def get_data(
+    m: mujoco.MjModel, d: types.Data
+) -> Union[mujoco.MjData, List[mujoco.MjData]]:
+  """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData]."""
+  dx = jax.device_get(d)
+  batched = len(d.qpos.shape) > 1
+  batch_size = d.qpos.shape[0] if batched else 1
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  efc_type = np.array([
+      mujoco.mjtConstraint.mjCNSTR_EQUALITY,
+      mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
+      mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT,
+      mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
+  ]).repeat([ne, nf, nl, nc])
+
+  ds = []
+  for i in range(batch_size):
+    dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d
+    ncon = (dx_i.contact.dist <= 0).sum()
+    efc_active = (dx_i.efc_J != 0).any(axis=1)
+    efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
+    nefc, nc = efc_active.sum(), (efc_active & efc_con).sum()
+    d_i = mujoco.MjData(m)
+    d_i.nnzJ = nefc * m.nv
+    mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
+    d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
+    d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
+    d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
+
+    for field in types.Data.fields():
+      if field.name == 'contact':
+        _get_contact(d_i.contact, dx_i.contact, nefc - nc)
+        continue
+
+      value = getattr(dx_i, field.name)
+
+      if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+        value = value.reshape((-1, 9))
+
+      if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+        value = value[efc_active]
+
+      if field.name == 'efc_J':
+        value = value[efc_active].reshape(-1)
+
+      if value.shape:
+        getattr(d_i, field.name)[:] = value
+      else:
+        setattr(d_i, field.name, value)
+
+    d_i.efc_type[:] = efc_type[efc_active]
+    ds.append(d_i)
+
+  return ds if batched else ds[0]
+
+
+def _put_contact(
+    c: mujoco._structs._MjContactList, ncon: int, device=None
+) -> types.Contact:
+  """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact."""
+  fields = {
+      f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields()
+  }
+  fields['frame'] = fields['frame'].reshape((-1, 3, 3))
+  pad_size = ncon - c.dist.shape[0]
+  pad_fn = lambda x: np.concatenate(
+      (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
+  )
+  fields = jax.tree_map(pad_fn, fields)
+  fields['dist'][-pad_size:] = np.inf
+  fields = jax.device_put(fields, device=device)
+
+  return types.Contact(**fields)
+
+
+def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
+  """Puts mujoco.MjData onto a device, resulting in mjx.Data."""
+  ncon = collision_driver.ncon(m)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  for d_val, val, name in (
+      (d.ncon, ncon, 'ncon'),
+      (d.ne, ne, 'ne'),
+      (d.nf, nf, 'nf'),
+      (d.nl, nl, 'nl'),
+      (d.nefc, nefc, 'nefc'),
+  ):
+    if d_val > val:
+      raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
+
+  fields = {
+      f.name: copy.copy(getattr(d, f.name))  # copy because device_put is async
+      for f in types.Data.fields()
+      if f.type is jax.Array
+  }
+
+  for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+    fields[fname] = fields[fname].reshape((-1, 3, 3))
+
+  # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
+  # efc_J is also optionally column-sparse (typically for large nv).  MJX is
+  # neither: it contains zeros for inactive constraints, and efc_J is always
+  # (nefc, nv).  this may change in the future.
+  if mujoco.mj_isSparse(m):
+    nr = d.efc_J_rownnz.shape[0]
+    efc_j = np.zeros((nr, m.nv))
+    for i in range(nr):
+      rowadr = d.efc_J_rowadr[i]
+      for j in range(d.efc_J_rownnz[i]):
+        efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
+    fields['efc_J'] = efc_j
+  else:
+    fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
+
+  for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+    value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
+    for i in range(4):
+      value_beg = sum([ne, nf, nl][:i])
+      d_beg = sum([d.ne, d.nf, d.nl][:i])
+      size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i]
+      value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
+    fields[fname] = value
+
+  fields = jax.device_put(fields, device=device)
+  fields['contact'] = _put_contact(d.contact, ncon, device=device)
+
+  return types.Data(**fields)

mjx/mujoco/mjx/_src/io_test.py

@@ -17,25 +17,384 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
+from jax import numpy as jp
+import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import test_util
+import numpy as np
 
 
-class IoTest(parameterized.TestCase):
+_MULTIPLE_CONVEX_OBJECTS = """
+  <mujoco>
+    <option timestep="0.001"/>
+    <default>
+      <geom solref=".006 1"/>
+    </default>
+    <asset>
+      <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 .1"/>
+      <mesh name="boxoid" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 .5 -1 -1 2" scale=".3 .2 .1"/>
+      <mesh name="pentaprism" vertex="1 0 0 0.309 0.951 0 -0.809 0.588 0 -0.809 -0.588 0 0.309 -0.951 0 1 0 1 0.309 0.951 1 -0.809 0.588 1 -0.809 -0.588 1 0.309 -0.951 1" scale=".2 .2 .1"/>
+    </asset>
+    <worldbody>
+      <geom type="plane" pos="0 0 -.5" size="3 3 .01"/>
+      <geom type="mesh" mesh="box" pos="0 0 -.15" euler="3 7 30"/>
+      <body pos="-.3 -.3 .3">
+        <freejoint/>
+        <geom type="mesh" mesh="boxoid" rgba=".8 0 0 1" euler="3 5 -130"/>
+      </body>
+      <body pos=".3 .3 0.3">
+        <freejoint/>
+        <geom type="box" euler="3 5 -80" size=".3 .2 .1" rgba="0 .8 0 1"/>
+      </body>
+      <body pos=".3 .3 .6">
+        <freejoint/>
+        <geom type="mesh" mesh="pentaprism" rgba="0 0 .8 1"/>
+      </body>
+      <body pos=".6 -.3 .3">
+        <joint name="joint" axis="1 0 0" type="hinge" range="-45 45"/>
+        <geom type="capsule" size=".2 .05" rgba=".6 0 .6 1"/>
+      </body>
+    </worldbody>
+    <actuator>
+      <motor joint="joint"/>
+    </actuator>
+  </mujoco>
+"""
 
-  @parameterized.parameters(test_util.TEST_FILES)
+_MULTIPLE_CONSTRAINTS = """
-  def test_make_data(self, fname):
+  <mujoco>
-    """Test that data created by make_data matches data returned by step."""
+    <worldbody>
+      <geom type="plane" size="3 3 .01"/>
+      <body name="cap1" pos="-.3 -.3 .2">
+        <freejoint/>
+        <geom type="capsule" size=".2 .05"/>
+        <body name="cap2" pos=".6 -.3 .3">
+          <joint axis="0 1 0" type="hinge" range="-45 45"/>
+          <joint axis="1 0 0" type="hinge" range="-0.001 0.001"/>
+          <geom type="capsule" size=".2 .05"/>
+          <site pos="-0.214 -0.078 0" quat="0.664 0.664 -0.242 -0.242"/>
+        </body>
+      </body>
+    </worldbody>
+    <equality>
+      <connect body1="cap2" anchor="0 0 1"/>
+    </equality>
+  </mujoco>
+"""
 
-    m = test_util.load_test_file(fname)
-    mx = mjx.device_put(m)
-    dx = mjx.make_data(mx)
-    dx_step = mjx.step(mx, dx)
 
-    _, dx_treedef = jax.tree_util.tree_flatten(dx)
+class ModelIOTest(parameterized.TestCase):
-    _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step)
+  """IO tests for mjx.Model."""
 
-    self.assertEqual(dx_treedef, dx_step_treedef)
+  def test_put_model(self):
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
+    mx = mjx.put_model(m)
+    self.assertEqual(mx.nq, m.nq)
+    self.assertEqual(mx.nv, m.nv)
+    self.assertEqual(mx.nu, m.nu)
+    self.assertEqual(mx.na, m.na)
+    self.assertEqual(mx.nbody, m.nbody)
+    self.assertEqual(mx.njnt, m.njnt)
+    self.assertEqual(mx.ngeom, m.ngeom)
+    self.assertEqual(mx.nmesh, m.nmesh)
+    self.assertEqual(mx.npair, m.npair)
+    self.assertEqual(mx.nexclude, m.nexclude)
+    self.assertEqual(mx.neq, m.neq)
+    self.assertEqual(mx.nnumeric, m.nnumeric)
+    self.assertEqual(mx.nM, m.nM)
+    self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
+
+    np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
+    np.testing.assert_allclose(mx.geom_type, m.geom_type)
+    np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
+    np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref)
+    np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
+    self.assertLen(mx.geom_convex_face, 6)
+    self.assertLen(mx.geom_convex_vert, 6)
+    self.assertLen(mx.geom_convex_edge, 6)
+    self.assertLen(mx.geom_convex_facenormal, 6)
+
+    np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
+    np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
+    np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid)
+    np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited)
+    np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis)
+
+    np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype)
+    np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype)
+    np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype)
+    np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype)
+    np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid)
+
+  def test_fluid_params(self):
+    """Test that has_fluid_params is set when fluid params are present."""
+    m = mjx.put_model(
+        mujoco.MjModel.from_xml_string(
+            '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
+        )
+    )
+    self.assertTrue(m.opt.has_fluid_params)
+
+  def test_implicit_not_implemented(self):
+    """Test that MJX guards against models with unimplemented features."""
+
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_cone_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_pgs_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option solver="PGS"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_tendon_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint name="left_hip" type="hinge"/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+          <tendon>
+            <fixed>
+              <joint coef="1" joint="left_hip"/>
+            </fixed>
+          </tendon>
+        </mujoco>"""))
+
+  def test_condim_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+  def test_cylinder_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom type="cylinder" size="0.05 0.05"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+
+class DataIOTest(parameterized.TestCase):
+  """IO tests for mjx.Data."""
+
+  def test_make_data(self):
+    """Test that make_data returns the correct shapes."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
+    d = mjx.make_data(m)
+
+    nq = 22
+    nbody = 5
+    ncon = 46
+    nv = 19
+    nefc = 185
+    nm = 64
+
+    self.assertEqual(d.qpos.shape, (nq,))
+    self.assertEqual(d.qvel.shape, (nv,))
+    self.assertEqual(d.act.shape, (0,))
+    self.assertEqual(d.qacc_warmstart.shape, (nv,))
+    self.assertEqual(d.ctrl.shape, (1,))
+    self.assertEqual(d.qfrc_applied.shape, (nv,))
+    self.assertEqual(d.xfrc_applied.shape, (nbody, 6))
+    self.assertEqual(d.eq_active.shape, (0,))
+    self.assertEqual(d.qacc.shape, (nv,))
+    self.assertEqual(d.act_dot.shape, (0,))
+    self.assertEqual(d.xpos.shape, (nbody, 3))
+    self.assertEqual(d.xquat.shape, (nbody, 4))
+    self.assertEqual(d.xmat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xipos.shape, (nbody, 3))
+    self.assertEqual(d.ximat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xanchor.shape, (4, 3))
+    self.assertEqual(d.xaxis.shape, (4, 3))
+    self.assertEqual(d.geom_xpos.shape, (6, 3))
+    self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
+    self.assertEqual(d.subtree_com.shape, (nbody, 3))
+    self.assertEqual(d.cdof.shape, (nv, 6))
+    self.assertEqual(d.cinert.shape, (nbody, 10))
+    self.assertEqual(d.crb.shape, (nbody, 10))
+    self.assertEqual(d.actuator_length.shape, (1,))
+    self.assertEqual(d.actuator_moment.shape, (1, nv))
+    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
+    self.assertEqual(d.contact.dist.shape, (ncon,))
+    self.assertEqual(d.contact.pos.shape, (ncon, 3))
+    self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
+    self.assertEqual(d.contact.solref.shape, (ncon, 2))
+    self.assertEqual(d.contact.solimp.shape, (ncon, 5))
+    self.assertEqual(d.contact.geom1.shape, (ncon,))
+    self.assertEqual(d.contact.geom2.shape, (ncon,))
+    self.assertEqual(d.efc_J.shape, (nefc, nv))
+    self.assertEqual(d.efc_frictionloss.shape, (nefc,))
+    self.assertEqual(d.efc_D.shape, (nefc,))
+    self.assertEqual(d.actuator_velocity.shape, (1,))
+    self.assertEqual(d.cvel.shape, (nbody, 6))
+    self.assertEqual(d.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d.qfrc_bias.shape, (nv,))
+    self.assertEqual(d.qfrc_passive.shape, (nv,))
+    self.assertEqual(d.efc_aref.shape, (nefc,))
+    self.assertEqual(d.qfrc_actuator.shape, (nv,))
+    self.assertEqual(d.qfrc_smooth.shape, (nv,))
+    self.assertEqual(d.qacc_smooth.shape, (nv,))
+    self.assertEqual(d.qfrc_constraint.shape, (nv,))
+    self.assertEqual(d.qfrc_inverse.shape, (nv,))
+    self.assertEqual(d.efc_force.shape, (nefc,))
+
+  def test_put_data(self):
+    """Test that put_data puts the correct data for dense and sparse."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+
+    # check a few fields
+    np.testing.assert_allclose(dx.qpos, d.qpos)
+    np.testing.assert_allclose(dx.xpos, d.xpos)
+    np.testing.assert_allclose(dx.cvel, d.cvel)
+    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(dx.qM, d.qM)
+
+    # 4 contacts, 2 for each capsule against the plane
+    self.assertEqual(dx.contact.dist.shape, (4,))
+    self.assertEqual(d.ncon, 1)  # however only 1 contact in this step
+    np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
+    self.assertTrue(np.isinf(dx.contact.dist[1:]).all())
+    self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
+    np.testing.assert_allclose(
+        dx.contact.frame[0].reshape(9), d.contact.frame[0]
+    )
+    np.testing.assert_allclose(dx.contact.frame[1:], 0)
+
+    # xmat, ximat, geom_xmat are all shape transformed
+    self.assertEqual(dx.xmat.shape, (3, 3, 3))
+    self.assertEqual(dx.ximat.shape, (3, 3, 3))
+    self.assertEqual(dx.geom_xmat.shape, (3, 3, 3))
+    self.assertEqual(dx.site_xmat.shape, (1, 3, 3))
+    np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat)
+    np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat)
+    np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
+    np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
+
+    # efc_ are also shape transformed and padded
+    self.assertEqual(dx.efc_J.shape, (21, 8))  # nefc, nv
+    d_efc_j = d.efc_J.reshape((-1, 8))
+    np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3])  # connect eq
+    np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3])  # one active limit
+    np.testing.assert_allclose(dx.efc_J[4], 0)  # one inactive limit
+    np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8])  # contact
+    np.testing.assert_allclose(dx.efc_J[9:], 0)  # no contact
+
+    # check another efc_ too
+    self.assertEqual(dx.efc_aref.shape, (21,))  # nefc
+    np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
+    np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
+    np.testing.assert_allclose(dx.efc_aref[4], 0)
+    np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8])
+    np.testing.assert_allclose(dx.efc_aref[9:], 0)
+
+    # check sparse transform is correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_from_sparse = mjx.put_data(m, d)
+    np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
+
+  def test_get_data(self):
+    """Test that get_data makes correct MjData."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2: mujoco.MjData = mjx.get_data(m, dx)
+
+    # check a few fields
+    np.testing.assert_allclose(d_2.qpos, d.qpos)
+    np.testing.assert_allclose(d_2.xpos, d.xpos)
+    np.testing.assert_allclose(d_2.cvel, d.cvel)
+    np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(d_2.qM, d.qM)
+
+    # only 1 contact active
+    self.assertEqual(d_2.contact.dist.shape, (1,))
+    self.assertEqual(d_2.ncon, 1)
+    np.testing.assert_allclose(d_2.contact.dist, d.contact.dist)
+    self.assertEqual(d_2.contact.frame.shape, (1, 9))
+    np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
+
+    # xmat, ximat, geom_xmat, site_xmat are all shape transformed
+    self.assertEqual(d_2.xmat.shape, (3, 9))
+    self.assertEqual(d_2.ximat.shape, (3, 9))
+    self.assertEqual(d_2.geom_xmat.shape, (3, 9))
+    self.assertEqual(d_2.site_xmat.shape, (1, 9))
+    np.testing.assert_allclose(d_2.xmat, d.xmat)
+    np.testing.assert_allclose(d_2.ximat, d.ximat)
+    np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat)
+    np.testing.assert_allclose(d_2.site_xmat, d.site_xmat)
+
+    # efc_* are also shape transformed and filtered
+    self.assertEqual(d_2.efc_J.shape, (64,))  # nefc * nv
+    np.testing.assert_allclose(d_2.efc_J, d.efc_J)
+    self.assertEqual(d_2.efc_aref.shape, (8,))  # nefc
+    np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
+
+    # efc_address is created on demand
+    np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
+
+  def test_get_data_batched(self):
+    """Test that get_data makes correct List[MjData] for batched Data."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    # second data in batch has contact dist > 0, disables contact
+    dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
+    ds = mjx.get_data(m, dx_b)
+    self.assertLen(ds, 2)
+    np.testing.assert_allclose(ds[0].qpos, d.qpos)
+    np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8)
+    self.assertEqual(ds[0].ncon, 1)
+    self.assertEqual(ds[1].ncon, 0)
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/passive.py

@@ -76,7 +76,7 @@ def _inertia_box_fluid_model(
 def passive(m: Model, d: Data) -> Data:
   """Adds all passive forces."""
   if m.opt.disableflags & DisableBit.PASSIVE:
-    return d
+    return d.replace(qfrc_passive=jp.zeros(m.nv))
 
   # joint-level springs
   def fn(jnt_typs, stiffness, qpos_spring, qpos):

mjx/mujoco/mjx/_src/passive_test.py

@@ -14,100 +14,65 @@
 # ==============================================================================
 """Tests passive forces."""
 
-import itertools
-
 from absl.testing import absltest
-from absl.testing import parameterized
-from etils import epath
 import jax
-import jax.numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import test_util
 import numpy as np
 
-
+# tolerance for difference between MuJoCo and MJX passive calculations - mostly
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-4, rtol=1e-4):
+# due to float precision
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+_TOLERANCE = 1e-7
-  a, b = getattr(a, attr), getattr(b, attr)
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class PassiveTest(parameterized.TestCase):
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  @parameterized.parameters(enumerate(('ant.xml', 'pendula.xml')))
-  def test_stiffness_damping(self, seed, fname):
-    """Tests stiffness and damping on Ant."""
-    np.random.seed(seed)
-    path = epath.resource_path('mujoco.mjx') / 'test_data'
-    path /= fname
-    m = mujoco.MjModel.from_xml_string(path.read_text())
 
-    # set stiffness/damping
+def _assert_attr_eq(a, b, attr):
-    m.jnt_stiffness = np.random.uniform(size=m.njnt)
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
-    m.dof_damping = np.random.uniform(size=m.nv)
+
+
+class PassiveTest(absltest.TestCase):
+
+  def test_passive(self):
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
-    d.qvel = np.random.random(m.nv)  # random kick
+    # give the system a little kick to ensure we have non-identity rotations
+    d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4])
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
 
-    mx = mjx.device_put(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    dx = mjx.make_data(mx)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-    passive_jit_fn = jax.jit(mjx.passive)
+    # test with fluid forces
+    m.opt.density = 0.01
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-    for i in range(100):
+    m.opt.viscosity = 0.02
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
+    mujoco.mj_forward(m, d)
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
-      dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-      _assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
 
-  @parameterized.parameters(
+    m.opt.wind = np.array([0.03, 0.04, 0.05])
-      itertools.product(range(3), ('pendula.xml',))
+    mujoco.mj_forward(m, d)
-  )
+    mx = mjx.put_model(m)
-  def test_fluid(self, seed, fname):
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    np.random.seed(seed)
+    _assert_attr_eq(d, dx, 'qfrc_passive')
-    path = epath.resource_path('mujoco.mjx') / 'test_data'
-    path /= fname
-    m = mujoco.MjModel.from_xml_string(path.read_text())
 
-    # set density/viscosity/wind
+    # test disable passive
-    m.opt.density = np.random.uniform()
+    mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
-    m.opt.viscosity = np.random.uniform()
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
-    m.opt.wind = np.random.uniform()
+    np.testing.assert_allclose(dx.qfrc_passive, 0)
-
-    passive_jit_fn = jax.jit(mjx.passive)
-
-    mx = mjx.device_put(m)
-    d = mujoco.MjData(m)
-    d.qvel = np.random.random(m.nv)  # random kick
-
-    for i in range(100):
-      mujoco.mj_step(m, d)
-      dx = mjx.device_put(d)
-      mujoco.mj_passive(m, d)
-      dx = passive_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
-
-  def test_disable_passive(self):
-    m = mujoco.MjModel.from_xml_string("""
-        <mujoco>
-          <option density="1" viscosity="2" wind="0.1 0.2 0.3">
-            <flag passive="disable"/>
-          </option>
-          <worldbody>
-            <body>
-              <joint damping="1" axis="1 0 0" type="ball"/>
-              <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
-            </body>
-          </worldbody>
-        </mujoco>
-        """)
-    mx = mjx.device_put(m)
-    d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
-    dx = dx.replace(qvel=jp.ones(mx.nv))
-
-    passive_jit_fn = jax.jit(mjx.passive)
-    dx = passive_jit_fn(mx, dx)
-    np.testing.assert_equal(dx.qfrc_passive, np.zeros(mx.nv))
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/ray.py

@@ -0,0 +1,208 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Functions for ray interesection testing."""
+
+from typing import Sequence, Tuple
+
+import jax
+from jax import numpy as jp
+import mujoco
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import GeomType
+from mujoco.mjx._src.types import Model
+# pylint: enable=g-importing-member
+import numpy as np
+
+
+def _ray_quad(
+    a: jax.Array, b: jax.Array, c: jax.Array
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = jp.sqrt(det)
+
+  x0, x1 = (-b - det_2) / a, (-b + det_2) / a
+  x0 = jp.where((det < mujoco.mjMINVAL) | (x0 < 0), jp.inf, x0)
+  x1 = jp.where((det < mujoco.mjMINVAL) | (x1 < 0), jp.inf, x1)
+
+  return x0, x1
+
+
+def _ray_plane(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -mujoco.mjMINVAL  # z-vec pointing towards front face
+  valid &= x >= 0
+  # only within rendered rectangle
+  p = pnt[0:2] + x * vec[0:2]
+  valid &= jp.all((size[0:2] <= 0) | (jp.abs(p) <= size[0:2]))
+
+  return jp.where(valid, x, jp.inf)
+
+
+def _ray_sphere(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a sphere."""
+  x0, x1 = _ray_quad(vec @ vec, vec @ pnt, pnt @ pnt - size[0] * size[0])
+  x = jp.where(jp.isinf(x0), x1, x0)
+
+  return x
+
+
+def _ray_capsule(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  a = vec[0:2] @ vec[0:2]
+  b = vec[0:2] @ pnt[0:2]
+  c = pnt[0:2] @ pnt[0:2] - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  x0, x1 = _ray_quad(a, b, c)
+  x = jp.where(jp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = jp.where(jp.abs(pnt[2] + x * vec[2]) <= size[1], x, jp.inf)
+
+  # top cap
+  dif = pnt - jp.array([0, 0, size[1]])
+  x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
+  # accept only top half of sphere
+  x = jp.where((pnt[2] + x0 * vec[2] >= size[1]) & (x0 < x), x0, x)
+  x = jp.where((pnt[2] + x1 * vec[2] >= size[1]) & (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + jp.array([0, 0, size[1]])
+  x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
+
+  # accept only bottom half of sphere
+  x = jp.where((pnt[2] + x0 * vec[2] <= -size[1]) & (x0 < x), x0, x)
+  x = jp.where((pnt[2] + x1 * vec[2] <= -size[1]) & (x1 < x), x1, x)
+
+  return x
+
+
+def _ray_box(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a box."""
+
+  iface = jp.array([(1, 2), (0, 2), (0, 1), (1, 2), (0, 2), (0, 1)])
+
+  # side +1, -1
+  # solution of pnt[i] + x * vec[i] = side * size[i]
+  x = jp.concatenate([(size - pnt) / vec, (-size - pnt) / vec])
+
+  # intersection with face
+  p0 = pnt[iface[:, 0]] + x * vec[iface[:, 0]]
+  p1 = pnt[iface[:, 1]] + x * vec[iface[:, 1]]
+  valid = jp.abs(p0) <= size[iface[:, 0]]
+  valid &= jp.abs(p1) <= size[iface[:, 1]]
+
+  return jp.min(jp.where(valid, x, jp.inf))
+
+
+def _ray_mesh(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a mesh."""
+  del size, pnt, vec
+  raise NotImplementedError("ray <> mesh not implemented yet")
+
+
+_RAY_FUNC = {
+    GeomType.PLANE: _ray_plane,
+    GeomType.SPHERE: _ray_sphere,
+    GeomType.CAPSULE: _ray_capsule,
+    GeomType.BOX: _ray_box,
+    # GeomType.MESH: _ray_mesh,
+}
+
+
+def ray(
+    m: Model,
+    d: Data,
+    pnt: jax.Array,
+    vec: jax.Array,
+    geomgroup: Sequence[int] = (),
+    flg_static: bool = True,
+    bodyexclude: int = -1,
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns the geom id and distance at which a ray intersects with a geom.
+
+  Args:
+    m: MJX model
+    d: MJX data
+    pnt: ray origin point (3,)
+    vec: ray direction    (3,)
+    geomgroup: group inclusion/exclusion mask, or empty to ignore
+    flg_static: if True, allows rays to intersect with static geoms
+    bodyexclude: ignore geoms on specified body id
+
+  Returns:
+    dist: distance from ray origin to geom surface (or -1.0 for no intersection)
+    id: id of intersected geom (or -1 for no intersection)
+  """
+
+  dists, ids = [], []
+  geom_filter = m.geom_bodyid != bodyexclude
+  geom_filter &= (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
+  geom_filter &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
+  geom_filter &= flg_static | (m.body_weldid[m.geom_bodyid] != 0)
+  if geomgroup:
+    geomgroup = np.array(geomgroup, dtype=bool)
+    geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)]
+
+  # map ray to local geom frames
+  geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
+  geom_vecs = jax.vmap(lambda x: x.T @ vec)(d.geom_xmat)
+
+  for geom_type, fn in _RAY_FUNC.items():
+    id_, = np.nonzero(geom_filter & (m.geom_type == geom_type))
+
+    if id_.size == 0:
+      continue
+
+    size, pnt, vec = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
+    dist = jax.vmap(fn)(size, pnt, vec)
+    dists, ids = dists + [dist], ids + [id_]
+
+  if not ids:
+    return jp.array(-1), jp.array(-1.0)
+
+  dists = jp.concatenate(dists)
+  ids = jp.concatenate(ids)
+  min_id = jp.argmin(dists)
+  dist = jp.where(jp.isinf(dists[min_id]), -1, dists[min_id])
+  id_ = jp.where(jp.isinf(dists[min_id]), -1, ids[min_id])
+
+  return dist, id_

mjx/mujoco/mjx/_src/ray_test.py

@@ -0,0 +1,220 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for ray functions."""
+
+from absl.testing import absltest
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+
+  def test_ray_nothing(self):
+    """Tests that MJX ray returns -1 when nothing is hit."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    pnt, vec = jp.array([12.146, 1.865, 3.895]), jp.array([0, 0, -1.0])
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_plane(self):
+    """Tests MJX ray<>plane matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 0, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking on wrong side of plane
+    pnt = jp.array([0, 0, -0.5])
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_sphere(self):
+    """Tests MJX ray<>sphere matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt, vec = jp.array([0, 0, 1.6]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 1, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_capsule(self):
+    """Tests MJX ray<>capsule matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt, vec = jp.array([0.5, 1, 1.6]), jp.array([0, 0.05, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking up at capsule from below
+    pnt, vec = jp.array([-0.5, 1, 0.05]), jp.array([0, 0.05, 1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking at cylinder of capsule from the side
+    pnt, vec = jp.array([0, 1, 0.75]), jp.array([1, 0, 0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_box(self):
+    """Tests MJX ray<>box matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt, vec = jp.array([1, 0, 1.6]), jp.array([0, 0.05, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 3, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking up at box from below
+    pnt, vec = jp.array([1, 0, 0.05]), jp.array([0, 0.05, 1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 3, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnums=(4,))
+
+    # hits plane with geom_group[0] = 1
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    geomgroup = (1, 0, 0, 0, 0, 0)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup)
+    _assert_eq(geomid, 0, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # nothing hit with geom_group[0] = 0
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    geomgroup = (0, 0, 0, 0, 0, 0)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnames=('flg_static',))
+
+    # nothing hit with flg_static = False
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, flg_static=False)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_bodyexclude(self):
+    """Tests ray bodyexclude filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnames=('bodyexclude',))
+
+    # nothing hit with bodyexclude = 0 (world body)
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, bodyexclude=0)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_invisible(self):
+    """Tests ray doesn't hit transparent geoms."""
+    m = test_util.load_test_file('ray.xml')
+    # nothing hit with transparent geoms:
+    m.geom_rgba = 0
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/scan.py

@@ -49,7 +49,9 @@ def _take(obj: Y, idx: np.ndarray) -> Y:
 
   def take(x):
     # TODO(erikfrey): if this helps perf, add support for striding too
-    if (
+    if not x.shape[0]:
+      return x
+    elif (
         len(idx.shape) == 1
         and idx.size > 0
         and (idx == np.arange(idx[0], idx[0] + idx.size)).all()
@@ -113,8 +115,12 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
     if isinstance(arg, np.ndarray) and not np.all(arg == arg[0]):
       raise RuntimeError(f'numpy arg elements do not match: {arg}')
 
+  # split out numpy and jax args
   np_args = [a[0] if isinstance(a, np.ndarray) else None for a in args]
   args = [a if n is None else None for n, a in zip(np_args, args)]
+
+  # remove empty args that we should not vmap over
+  args = jax.tree_map(lambda a: a if a.shape[0] else None, args)
   in_axes = [None if a is None else 0 for a in args]
 
   def outer_f(*args, np_args=np_args):
@@ -126,7 +132,15 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
 
 def _check_input(m: Model, args: Any, in_types: str) -> None:
   """Checks that scan input has the right shape."""
-  size = {'b': m.nbody, 'j': m.njnt, 'q': m.nq, 'v': m.nv, 'u': m.nu, 'a': m.na}
+  size = {
+      'b': m.nbody,
+      'j': m.njnt,
+      'q': m.nq,
+      'v': m.nv,
+      'u': m.nu,
+      'a': m.na,
+      's': m.nsite,
+  }
   for idx, (arg, typ) in enumerate(zip(args, in_types)):
     if len(arg) != size[typ]:
       raise IndexError(
@@ -162,7 +176,7 @@ def flat(
 ) -> Y:
   r"""Scan a function across bodies or actuators.
 
-  Scan group data according to type and batch shape then calls vmap(f) on it.
+  Scan group data according to type and batch shape then calls vmap(f) on it.\
 
   Args:
     m: an mjx model
@@ -206,6 +220,7 @@ def flat(
         m.actuator_dyntype[ids_u],
         m.actuator_trntype[ids_u],
         m.jnt_type[ids_j],
+        m.actuator_trnid[ids_u, 1] == -1,  # key by refsite being present
     )
 
   def type_ids_j(m, i):
@@ -221,16 +236,24 @@ def flat(
         'u': i,
         'a': m.actuator_actadr[i],
         'j': (
-            m.actuator_trnid[i]
+            m.actuator_trnid[i, 0]
             if m.actuator_trntype[i] == TrnType.JOINT
-            else np.array(-1)
+            else -1
+        ),
+        's': (
+            m.actuator_trnid[i]
+            if m.actuator_trntype[i] == TrnType.SITE
+            else np.array([-1, -1])
         ),
     }
-    # v/q associated with joint transmissions
+    v, q = np.array([-1]), np.array([-1])
-    typ_ids.update({
+    if m.actuator_trntype[i] == TrnType.JOINT:
-        'v': np.nonzero(m.dof_jntid == typ_ids['j'])[0],
+      # v/q are associated with the joint transmissions only
-        'q': np.nonzero(_q_jointid(m) == typ_ids['j'])[0],
+      v = np.nonzero(m.dof_jntid == typ_ids['j'])[0]
-    })
+      q = np.nonzero(_q_jointid(m) == typ_ids['j'])[0]
+
+    typ_ids.update({'v': v, 'q': q})
+
     return typ_ids
 
   # build up a grouping of type take-ids in body/actuator order

mjx/mujoco/mjx/_src/scan_test.py

@@ -193,7 +193,7 @@ class ScanTest(absltest.TestCase):
     </mujoco>
   """
 
-  def testscan_actuators(self):
+  def test_scan_actuators(self):
     """Tests scanning over actuators."""
     m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML)
     m = mjx.device_put(m)
@@ -210,15 +210,16 @@ class ScanTest(absltest.TestCase):
         m, fn, 'ujqva', 'ujqva', *args, group_by='u'
     )
 
+    actuator_trnid = m.actuator_trnid[:, 0]
     np.testing.assert_array_equal(gear, m.actuator_gear)
-    np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid])
+    np.testing.assert_array_equal(jnt_typ, m.jnt_type[actuator_trnid])
     np.testing.assert_array_equal(act, jp.array([1.4, 1.1]))
     expected_vadr = np.concatenate(
-        [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid]
+        [np.nonzero(m.dof_jntid == trnid)[0] for trnid in actuator_trnid]
     )
     np.testing.assert_array_equal(vadr, expected_vadr)
     expected_qadr = np.concatenate(
-        [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid]
+        [np.nonzero(scan._q_jointid(m) == i)[0] for i in actuator_trnid]
     )
     np.testing.assert_array_equal(qadr, expected_qadr)
 

mjx/mujoco/mjx/_src/smooth.py

@@ -19,12 +19,15 @@ from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import math
 from mujoco.mjx._src import scan
+from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
+from mujoco.mjx._src.types import TrnType
 # pylint: enable=g-importing-member
+import numpy as np
 
 
 def kinematics(m: Model, d: Data) -> Data:
@@ -101,13 +104,20 @@ def kinematics(m: Model, d: Data) -> Data:
 
   # TODO(erikfrey): confirm that quats are more performant for mjx than mats
   xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat)
-  geom_xpos, geom_xmat = local_to_global(
-      xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
-  )
-
   d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos)
   d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat)
-  d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
+
+  if m.ngeom:
+    geom_xpos, geom_xmat = local_to_global(
+        xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
+    )
+    d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
+
+  if m.nsite:
+    site_xpos, site_xmat = local_to_global(
+        xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat
+    )
+    d = d.replace(site_xpos=site_xpos, site_xmat=site_xmat)
 
   return d
 
@@ -423,45 +433,118 @@ def rne(m: Model, d: Data) -> Data:
   return d
 
 
+def _site_dof_mask(m: Model) -> np.ndarray:
+  """Creates a dof mask for site transmissions."""
+  mask = np.ones((m.nu, m.nv))
+  for i in np.nonzero(m.actuator_trnid[:, 1] != -1)[0]:
+    id_, refid = m.actuator_trnid[i]
+    # intialize last dof address for each body
+    b0 = m.body_weldid[m.site_bodyid[id_]]
+    b1 = m.body_weldid[m.site_bodyid[refid]]
+    dofadr0 = m.body_dofadr[b0] + m.body_dofnum[b0] - 1
+    dofadr1 = m.body_dofadr[b1] + m.body_dofnum[b1] - 1
+
+    # find common ancestral dof, if any
+    while dofadr0 != dofadr1:
+      if dofadr0 < dofadr1:
+        dofadr1 = m.dof_parentid[dofadr1]
+      else:
+        dofadr0 = m.dof_parentid[dofadr0]
+      if dofadr0 == -1 or dofadr1 == -1:
+        break
+
+    # if common ancestral dof was found, clear the columns of its parental chain
+    da = dofadr0 if dofadr0 == dofadr1 else -1
+    while da >= 0:
+      mask[i, da] = 0.0
+      da = m.dof_parentid[da]
+
+  return mask
+
+
 def transmission(m: Model, d: Data) -> Data:
   """Computes actuator/transmission lengths and moments."""
+  # TODO: consider combining transmission calculation into fwd_actuation.
   if not m.nu:
     return d
 
-  def fn(gear, jnt_typ, m_i, m_j, qpos):
+  def fn(
-    # handles joint transmissions only
+      trntype,
-    if jnt_typ == JointType.FREE:
+      trnid,
-      length = jp.zeros(1)
+      gear,
-      moment = gear
+      jnt_typ,
-      m_i = jp.repeat(m_i, 6)
+      m_j,
-      m_j = m_j + jp.arange(6)
+      qpos,
-    elif jnt_typ == JointType.BALL:
+      has_refsite,
-      axis, _ = math.quat_to_axis_angle(qpos)
+      site_dof_mask,
-      length = jp.dot(axis, gear[:3])[None]
+      site_xpos,
-      moment = gear[:3]
+      site_xmat,
-      m_i = jp.repeat(m_i, 3)
+      site_quat,
-      m_j = m_j + jp.arange(3)
+  ):
-    elif jnt_typ in (JointType.SLIDE, JointType.HINGE):
+    if trntype == TrnType.JOINT:
-      length = qpos * gear[0]
+      if jnt_typ == JointType.FREE:
-      moment = gear[:1]
+        length = jp.zeros(1)
-      m_i, m_j = m_i[None], m_j[None]
+        moment = gear
-    else:
+        m_j = m_j + jp.arange(6)
-      raise RuntimeError(f'unrecognized joint type: {jnt_typ}')
+      elif jnt_typ == JointType.BALL:
-    return length, moment, m_i, m_j
+        axis, angle = math.quat_to_axis_angle(qpos)
+        length = jp.dot(axis * angle, gear[:3])[None]
+        moment = gear[:3]
+        m_j = m_j + jp.arange(3)
+      elif jnt_typ in (JointType.SLIDE, JointType.HINGE):
+        length = qpos * gear[0]
+        moment = gear[:1]
+        m_j = m_j[None]
+      else:
+        raise RuntimeError(f'unrecognized joint type: {JointType(jnt_typ)}')
 
-  length, m_val, m_i, m_j = scan.flat(
+      moment = jp.zeros((m.nv,)).at[m_j].set(moment)
+    elif trntype == TrnType.SITE:
+      length = jp.zeros(1)
+      id_, refid = jp.array(m.site_bodyid)[trnid]
+      jacp, jacr = support.jac(m, d, site_xpos[0], id_)
+      frame_xmat = site_xmat[0]
+      if has_refsite:
+        vecp = site_xmat[1].T @ (site_xpos[0] - site_xpos[1])
+        vecr = math.quat_sub(site_quat[0], site_quat[1])
+        length += jp.dot(jp.concatenate([vecp, vecr]), gear)
+        jacrefp, jacrefr = support.jac(m, d, site_xpos[1], refid)
+        jacp, jacr = jacp - jacrefp, jacr - jacrefr
+        frame_xmat = site_xmat[1]
+
+      jac = jp.concatenate((jacp, jacr), axis=1) * site_dof_mask[:, None]
+      wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
+      moment = jac @ wrench
+    else:
+      raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
+
+    return length, moment
+
+  # pre-compute values for site transmissions
+  has_refsite = m.actuator_trnid[:, 1] != -1
+  site_dof_mask = _site_dof_mask(m)
+  site_quat = jax.vmap(math.quat_mul)(m.site_quat, d.xquat[m.site_bodyid])
+
+  length, moment = scan.flat(
       m,
       fn,
-      'ujujq',
+      'uuujjquusss',
-      'uvvv',
+      'uu',
+      m.actuator_trntype,
+      jp.array(m.actuator_trnid),
       m.actuator_gear,
       m.jnt_type,
-      jp.arange(m.nu),
       jp.array(m.jnt_dofadr),
       d.qpos,
+      has_refsite,
+      jp.array(site_dof_mask),
+      d.site_xpos,
+      d.site_xmat,
+      site_quat,
       group_by='u',
   )
-  moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val)
   length = length.reshape((m.nu,))
+  moment = moment.reshape((m.nu, m.nv))
+
   d = d.replace(actuator_length=length, actuator_moment=moment)
   return d

mjx/mujoco/mjx/_src/smooth_test.py

@@ -15,122 +15,107 @@
 """Tests for smooth dynamics functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
 import jax
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
-# pylint: disable=g-importing-member
-from mujoco.mjx._src.types import DisableBit
-# pylint: enable=g-importing-member
 import numpy as np
 
-
+# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
-def _assert_eq(a, b, name, step, fname, atol=5e-4, rtol=5e-4):
+# due to float precision
-  err_msg = f'mismatch: {name} at step {step} in {fname}'
+_TOLERANCE = 5e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=5e-4, rtol=5e-4):
+def _assert_eq(a, b, name):
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+  tol = _TOLERANCE * 10  # avoid test noise
-  a, b = getattr(a, attr), getattr(b, attr)
+  err_msg = f'mismatch: {name}'
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
-class SmoothTest(parameterized.TestCase):
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
 
-  @parameterized.parameters(enumerate(test_util.TEST_FILES))
-  def test_smooth(self, seed, fname):
-    """Tests mujoco mj smooth functions match mujoco_mjx smooth functions."""
-    if fname in ('convex.xml', 'equality.xml'):
-      return
 
-    np.random.seed(seed)
+class SmoothTest(absltest.TestCase):
 
-    m = test_util.load_test_file(fname)
+  def setUp(self):
+    super().setUp()
+    # although we already have generous padding of thresholds, it doesn't hurt
+    # to also fix the seed to reduce test flakiness
+    np.random.seed(0)
+
+  def test_smooth(self):
+    """Tests MJX smooth functions match MuJoCo smooth functions."""
+
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
-
-    kinematics_jit_fn = jax.jit(mjx.kinematics)
-    com_pos_jit_fn = jax.jit(mjx.com_pos)
-    crb_jit_fn = jax.jit(mjx.crb)
-    factor_m_fn = jax.jit(mjx.factor_m)
-    com_vel_jit_fn = jax.jit(mjx.com_vel)
-    rne_jit_fn = jax.jit(mjx.rne)
-    mul_m_jit_fn = jax.jit(mjx.mul_m)
-    transmission_jit_fn = jax.jit(mjx.transmission)
-
-    mx = mjx.device_put(m)
-    dx = mjx.make_data(mx)
-
     # give the system a little kick to ensure we have non-identity rotations
     d.qvel = np.random.random(m.nv)
-    for i in range(100):
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
-      qpos, qvel = d.qpos.copy(), d.qvel.copy()
+    mujoco.mj_forward(m, d)
-      mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
 
-      # kinematics
+    # kinematics
-      dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel))
+    dx = jax.jit(mjx.kinematics)(mx, mjx.put_data(m, d))
-      _assert_attr_eq(d, dx, 'xanchor', i, fname)
+    _assert_attr_eq(d, dx, 'xanchor')
-      _assert_attr_eq(d, dx, 'xaxis', i, fname)
+    _assert_attr_eq(d, dx, 'xaxis')
-      _assert_attr_eq(d, dx, 'xpos', i, fname)
+    _assert_attr_eq(d, dx, 'xpos')
-      _assert_attr_eq(d, dx, 'xquat', i, fname)
+    _assert_attr_eq(d, dx, 'xquat')
-      _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname)
+    _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat')
-      _assert_attr_eq(d, dx, 'xipos', i, fname)
+    _assert_attr_eq(d, dx, 'xipos')
-      _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname)
+    _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat')
-      _assert_attr_eq(d, dx, 'geom_xpos', i, fname)
+    _assert_attr_eq(d, dx, 'geom_xpos')
-      _assert_eq(
+    _assert_eq(d.geom_xmat.reshape((-1, 3, 3)), dx.geom_xmat, 'geom_xmat')
-          d.geom_xmat.reshape((-1, 3, 3)),
+    _assert_attr_eq(d, dx, 'site_xpos')
-          dx.geom_xmat,
+    _assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat')
-          'geom_xmat',
+    # com_pos
-          i,
+    dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
-          fname,
+    _assert_attr_eq(d, dx, 'subtree_com')
-      )
+    _assert_attr_eq(d, dx, 'cinert')
+    _assert_attr_eq(d, dx, 'cdof')
+    # crb
+    dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'crb')
+    _assert_attr_eq(d, dx, 'qM')
+    # factor_m
+    dx = mjx.put_data(m, d)
+    dx = jax.jit(mjx.factor_m)(mx, dx, dx.qM)
+    _assert_attr_eq(d, dx, 'qLD')
+    _assert_attr_eq(d, dx, 'qLDiagInv')
+    # com_vel
+    dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'cvel')
+    _assert_attr_eq(d, dx, 'cdof_dot')
+    # rne
+    dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_bias')
+    # transmission
+    dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_moment')
 
-      # com_pos
+  def test_mul_m(self):
-      dx = com_pos_jit_fn(mx, dx)
+    m = test_util.load_test_file('pendula.xml')
-      _assert_attr_eq(d, dx, 'subtree_com', i, fname)
+    d = mujoco.MjData(m)
-      _assert_attr_eq(d, dx, 'cinert', i, fname)
+    # give the system a little kick to ensure we have non-identity rotations
-      _assert_attr_eq(d, dx, 'cdof', i, fname)
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    vec = np.random.random(m.nv)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    mj_vec = np.zeros(m.nv)
+    mujoco.mj_mulM(m, d, mj_vec, vec)
+    _assert_eq(mj_vec, mjx_vec, 'mul_m')
 
-      # crb
+  def test_disable_gravity(self):
-      dx = crb_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'crb', i, fname)
-      _assert_attr_eq(d, dx, 'qM', i, fname)
-
-      # factor_m
-      dx = factor_m_fn(mx, dx, dx.qM)
-      _assert_attr_eq(d, dx, 'qLD', i, fname, atol=1e-3)
-      _assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3)
-
-      # com_vel
-      dx = com_vel_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'cvel', i, fname)
-      _assert_attr_eq(d, dx, 'cdof_dot', i, fname)
-
-      # rne
-      dx = rne_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'qfrc_bias', i, fname)
-
-      # mul_m (auxilliary function, not part of smooth step)
-      vec = np.random.random(m.nv)
-      mjx_vec = mul_m_jit_fn(mx, dx, jp.array(vec))
-      mj_vec = np.zeros(m.nv)
-      mujoco.mj_mulM(m, d, mj_vec, vec)
-      _assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname)
-
-      # transmission
-      dx = transmission_jit_fn(mx, dx)
-      _assert_attr_eq(d, dx, 'actuator_length', i, fname)
-      _assert_attr_eq(d, dx, 'actuator_moment', i, fname)
-
-
-class DisableGravityTest(absltest.TestCase):
-
-  def test_disabled(self):
     m = mujoco.MjModel.from_xml_string("""
         <mujoco>
-          <option timestep="0.01"/>
+          <option>
+            <flag gravity="disable"/>
+          </option>
           <worldbody>
             <body>
               <joint type="free"/>
@@ -139,27 +124,52 @@ class DisableGravityTest(absltest.TestCase):
           </worldbody>
         </mujoco>
         """)
-    mx = mjx.device_put(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
 
-    # test with gravity
+    dx = jax.jit(mjx.rne)(mx, dx)
-    step_jit_fn = jax.jit(mjx.step)
+    np.testing.assert_allclose(dx.qfrc_bias, 0)
-    dx = step_jit_fn(mx, dx)
-    np.testing.assert_array_almost_equal(
-        dx.qpos, np.array([0.0, 0.0, -9.81e-4, 1.0, 0.0, 0.0, 0.0]), decimal=7
-    )
 
-    # test with gravity disabled
+  def test_site_transmission(self):
-    mx = mx.tree_replace(
+    m = mujoco.MjModel.from_xml_string("""
-        {'opt.disableflags': mx.opt.disableflags | DisableBit.GRAVITY}
+        <mujoco>
-    )
+        <compiler autolimits="true"/>
-    dx = mjx.device_put(d)
+        <worldbody>
-    step_jit_fn = jax.jit(mjx.step)
+          <body>
-    dx = step_jit_fn(mx, dx)
+            <joint type="free"/>
-    np.testing.assert_equal(
+            <geom type="box" size=".05 .05 .05" mass="1"/>
-        dx.qpos, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0])
+            <site name="site1"/>
-    )
+            <body>
+              <joint type="hinge"/>
+              <geom size="0.1" mass="1"/>
+              <site name="site2" pos="0.1 0.2 0.3"/>
+            </body>
+          </body>
+          <body pos="1 0 0">
+            <joint name="slide" type="hinge"/>
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+          </body>
+        </worldbody>
+        <actuator>
+          <position site="site1" kv="0.1" gear="1 2 3 0 0 0"/>
+          <position site="site1" kv="0.2" gear="0 0 0 1 2 3"/>
+          <position site="site2" kv="0.3" gear="0 3 0 0 0 1"/>
+          <position joint="slide" kv="0.05" />
+          <position site="site2" refsite="site1" gear="1 2 3 0.5 0.4 0.6"/>
+        </actuator>
+        </mujoco>
+      """)
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_transmission(m, d)
+    dx = jax.jit(mjx.transmission)(mx, dx)
+    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_moment')
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/_src/solver.py

@@ -19,6 +19,7 @@ from typing import Optional
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.mjx._src import constraint
 from mujoco.mjx._src import math
 from mujoco.mjx._src import smooth
 # pylint: disable=g-importing-member
@@ -69,12 +70,12 @@ class _Context(PyTreeNode):
     # TODO(robotics-team): determine nv at which sparse mul is faster
     M = smooth.dense_m(m, d) if m.nv < 100 else None  # pylint: disable=invalid-name
     ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
-    nv_0 = jp.zeros((m.nv,))
+    nv_0 = jp.zeros(m.nv)
     ctx = _Context(
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
-        efc_force=jp.zeros(d.nefc),
+        efc_force=d.efc_force,
         M=M,
         Ma=ma,
         grad=nv_0,
@@ -111,7 +112,7 @@ class _LSPoint(PyTreeNode):
   @classmethod
   def create(
       cls,
-      d: Data,
+      m: Model,
       ctx: _Context,
       alpha: jax.Array,
       jv: jax.Array,
@@ -122,13 +123,14 @@ class _LSPoint(PyTreeNode):
     # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
 
     # TODO(robotics-team): change this to support friction constraints
-    active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
+    ne, nf, *_ = constraint.count_constraints(m)
+    active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True)
     quad = jax.vmap(jp.multiply)(quad, active)  # only active
     quad_total = quad_gauss + jp.sum(quad, axis=0)
 
     cost = alpha * alpha * quad_total[2] + alpha * quad_total[1] + quad_total[0]
     deriv_0 = 2 * alpha * quad_total[2] + quad_total[1]
-    deriv_1 = 2 * quad_total[2]
+    deriv_1 = 2 * quad_total[2] + (quad_total[2] == 0) * mujoco.mjMINVAL
     return _LSPoint(alpha=alpha, cost=cost, deriv_0=deriv_0, deriv_1=deriv_1)
 
 
@@ -177,12 +179,11 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
   Returns:
     context with new constraint force and costs
   """
-  del m
-
   # TODO(robotics-team): add friction constraints
 
   # only count active constraints
-  active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
+  ne, nf, *_ = constraint.count_constraints(m)
+  active = (ctx.Jaref < 0).at[:ne + nf].set(True)
 
   efc_force = d.efc_D * -ctx.Jaref * active
   qfrc_constraint = d.efc_J.T @ efc_force
@@ -221,7 +222,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
   if m.opt.solver == SolverType.CG:
     mgrad = smooth.solve_m(m, d, grad)
   elif m.opt.solver == SolverType.NEWTON:
-    active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
+    ne, nf, *_ = constraint.count_constraints(m)
+    active = (ctx.Jaref < 0).at[:ne + nf].set(True)
     h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
     h = smooth.dense_m(m, d) + h
     h_ = jax.scipy.linalg.cho_factor(h)
@@ -265,7 +267,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
   quad = (quad * d.efc_D).T
 
-  point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss)
+  point_fn = lambda a: _LSPoint.create(m, ctx, a, jv, quad, quad_gauss)
 
   def cond(ctx: _LSContext) -> jax.Array:
     done = ctx.ls_iter >= m.opt.ls_iterations

mjx/mujoco/mjx/_src/solver_test.py

@@ -12,119 +12,64 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for forward functions."""
+"""Tests for constraint functions."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
-from etils import epath
 import jax
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import test_util
 import numpy as np
 
 
-def _assert_attr_eq(a, b, attr, step, fname, atol=1e-2, rtol=1e-2):
+# tolerance for difference between MuJoCo and MJX constraint calculations,
-  err_msg = f'mismatch: {attr} at step {step} in {fname}'
+# mostly due to float precision
-  a, b = getattr(a, attr), getattr(b, attr)
+_TOLERANCE = 5e-5
-  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
 
 
-class Solver64Test(parameterized.TestCase):
+def _assert_eq(a, b, name, tol=_TOLERANCE):
-  """Tests solvers at 64 bit precision."""
+  tol = tol * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
-  def setUp(self):
-    super().setUp()
-    jax.config.update('jax_enable_x64', True)
 
-  def tearDown(self):
+def _assert_attr_eq(a, b, attr):
-    super().tearDown()
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
-    jax.config.update('jax_enable_x64', False)
 
-  @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml')))
+
-  def test_cg(self, seed, fname):
+class SolverTest(absltest.TestCase):
-    """Test mjx cg solver matches mujoco cg solver at 64 bit precision."""
+
-    f = epath.resource_path('mujoco.mjx') / 'test_data' / fname
+  def test_solver(self):
-    m = mujoco.MjModel.from_xml_string(f.read_text())
+    """Test solver."""
+    m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
 
-    jax.config.update('jax_enable_x64', True)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    forward_jit_fn = jax.jit(mjx.forward)
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
 
-    # give the system a little kick to ensure we have non-identity rotations
+    # also test normal CG
-    np.random.seed(seed)
+    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
-    d.qvel = 0.01 * np.random.random(m.nv)
+    mujoco.mj_forward(m, d)
-
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
-    for i in range(100):
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
-      # in order to avoid re-jitting, reuse the same mj_data shape
+    _assert_attr_eq(d, dx, 'qacc')
-      save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
-      d = mujoco.MjData(m)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
-      d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save
-      dx = mjx.device_put(d)
-
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      # at 64 bits the solutions returned by the two solvers are quite close
-      self.assertLessEqual(dx.solver_niter[0], d.solver_niter[0])
-      _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname)
-      _assert_attr_eq(d, dx, 'qacc', i, fname)
-
-
-class SolverTest(parameterized.TestCase):
-
-  @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml')))
-  def test_cg(self, seed, fname):
-    """Test mjx cg solver is close to mj at 32 bit precision.
-
-    Args:
-      seed: int
-      fname: file to test
-
-    At lower float resolution there's wiggle room in valid forces that satisfy
-    constraints.  So instead let's mainly validate that mjx is finding solutions
-    with as good cost as mujoco, even if the resulting forces/accelerations
-    are not quite the same.
-    """
-    f = epath.resource_path('mujoco.mjx') / 'test_data' / fname
-    m = mujoco.MjModel.from_xml_string(f.read_text())
-    d = mujoco.MjData(m)
-    mx = mjx.device_put(m)
-
-    forward_jit_fn = jax.jit(mjx.forward)
-
-    # give the system a little kick to ensure we have non-identity rotations
-    np.random.seed(seed)
-    d.qvel = 0.01 * np.random.random(m.nv)
-
-    for i in range(100):
-      # in order to avoid re-jitting, reuse the same mj_data shape
-      save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth
-      d = mujoco.MjData(m)
-      d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save
-      dx = mjx.device_put(d)
-
-      mujoco.mj_step(m, d)
-      dx = forward_jit_fn(mx, dx)
-
-      def cost(qacc):
-        jaref = np.zeros(d.nefc)
-        mujoco.mj_mulJacVec(m, d, jaref, qacc)
-        jaref -= d.efc_aref
-        cost = np.array([0.0])
-        mujoco.mj_constraintUpdate(m, d, jaref, cost, 0)
-        return cost[0]
-
-      cost_mj, cost_mjx = cost(d.qacc), cost(dx.qacc)
-
-      self.assertLessEqual(
-          cost_mjx,
-          cost_mj * 1.01,
-          msg=f'mismatch: {fname} at step {i}, cost too high',
-      )
-      _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname, atol=1e-1, rtol=1e-1)
-      _assert_attr_eq(d, dx, 'qacc', i, fname, atol=1e-1, rtol=1e-1)
 
+    # without warmstart, the solution is not as close
+    m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
+    m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2)
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/support_test.py

@@ -34,8 +34,8 @@ class SupportTest(parameterized.TestCase):
     m = test_util.load_test_file(fname)
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
     point = np.random.randn(3)
     body = np.random.choice(m.nbody)
     jacp, jacr = jax.jit(support.jac)(mx, dx, point, body)
@@ -49,11 +49,11 @@ class SupportTest(parameterized.TestCase):
     """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate."""
     np.random.seed(0)
 
-    m = test_util.load_test_file('ant.xml')
+    m = test_util.load_test_file('pendula.xml')
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
     self.assertFalse((dx.xipos == 0.0).all())
 
     xfrc = np.random.rand(*dx.xfrc_applied.shape)

mjx/mujoco/mjx/_src/test_util.py

@@ -23,20 +23,21 @@ import mujoco
 import numpy as np
 
 TEST_FILES: List[str] = [
-    'ant.xml',
+    'constraints.xml',
     'convex.xml',
-    'equality.xml',
-    'humanoid.xml',
     'pendula.xml',
+    'ray.xml',
 ]
 
 _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
+_DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact']
+_DYN_PRMS = ['0.189', '2.1']
 _JOINT_TYPES = ['free', 'hinge', 'slide', 'ball']
 _JOINT_AXES = ['1 0 0', '0 1 0', '0 0 1']
 _FRICTIONS = ['1.2 0.003 0.0002', '0.2 0.0001 0.0005']
 _KP_POS = ['1', '2']
 _KP_INTVEL = ['10000', '2000']
-_KV_VEL = ['123', '1']
+_KV_VEL = ['12', '1', '0', '0.1']
 _PAIR_FRICTIONS = ['1.2 0.9 0.003 0.0002 0.0001']
 _SOLREFS = ['0.04 1.01', '0.05 1.02', '0.03 1.1', '0.015 1.0']
 _SOLIMPS = [
@@ -47,7 +48,7 @@ _SOLIMPS = [
 _DIMS = ['3']
 _MARGINS = ['0.0', '0.01', '0.02']
 _GAPS = ['0.0', '0.005']
-_GEARS = ['20', '50', '100']
+_GEARS = ['2.1 0.0 3.3 0 2.3 0', '5.0 3.1 0 2.3 0.0 1.1']
 
 
 def p(pct: int) -> bool:
@@ -123,13 +124,29 @@ def _make_geom(
   return attr
 
 
-def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]:
+def _make_actuator(
+    actuator_type: str,
+    joint: str | None = None,
+    site: str | None = None,
+    refsite: str | None = None,
+) -> Dict[str, str]:
   """Returns attributes for an actuator."""
-  attr = {'joint': joint}
+  if joint:
-  if actuator_type == 'motor':
+    attr = {'joint': joint}
-    attr['gear'] = np.random.choice(_GEARS)
+  elif site:
-  elif actuator_type == 'position':
+    attr = {'site': site}
+  else:
+    raise ValueError('must provide a joint or site name')
+
+  if refsite:
+    attr['refsite'] = refsite
+
+  attr['gear'] = np.random.choice(_GEARS)
+
+  # set actuator type
+  if actuator_type == 'position':
     attr['kp'] = np.random.choice(_KP_POS)
+    attr['kv'] = np.random.choice(_KV_VEL)
   elif actuator_type == 'general':
     attr['biastype'] = 'affine'
     attr['gainprm'] = '35 0 0'
@@ -141,10 +158,18 @@ def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]:
   elif actuator_type == 'velocity':
     attr['kv'] = np.random.choice(_KV_VEL)
 
+  # set dyntype
+  if actuator_type == 'general':
+    attr['dyntype'] = np.random.choice(_DYN_TYPES)
+    if attr['dyntype'] != 'none':
+      attr['dynprm'] = np.random.choice(_DYN_PRMS)
+
+  # ctrlrange
   if p(50) and actuator_type != 'intvelocity':
     lb, ub = -np.random.uniform(), np.random.uniform()
     attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}'
 
+  # forcerange
   if p(50):
     lb, ub = -np.random.uniform(), np.random.uniform()
     attr['forcerange'] = f'{lb*10:.2f} {ub*10:.2f}'
@@ -235,6 +260,7 @@ def create_mjcf(
     pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}'
     n_bodies = len(list(mjcf.iter('body')))
     child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'})
+    ET.SubElement(child, 'site', {'name': f'site{n_bodies}'})
 
     n_joints = len(list(mjcf.iter('joint')))
     for nj in range(np.random.randint(1, max_stacked_joints + 1)):
@@ -272,17 +298,35 @@ def create_mjcf(
   for _ in range(num_trees):
     make_tree(world, 0)
 
+  bodies = list(mjcf.iter('body'))
+  n_bodies = len(bodies)
+
   # actuators
   if add_actuators:
     actuator = ET.SubElement(mjcf, 'actuator')
     n_joints = len(list(mjcf.iter('joint')))
     nu = np.random.randint(1, n_joints + 1)
     actuators = []
+
+    # joint transmission
     for i in range(nu):
       actuator_type = np.random.choice(_ACTUATOR_TYPES)
       attr = _make_actuator(actuator_type, joint=f'joint{i}')
       actuators.append((actuator_type, attr))
 
+    # site transmission
+    for i in range(np.random.randint(0, n_bodies)):
+      actuator_type = np.random.choice(_ACTUATOR_TYPES)
+      attr = _make_actuator(actuator_type, site=f'site{i}')
+      actuators.append((actuator_type, attr))
+
+    # site transmission with refsite
+    for i in range(np.random.randint(0, n_bodies)):
+      j = np.random.randint(0, n_bodies)
+      actuator_type = np.random.choice(_ACTUATOR_TYPES)
+      attr = _make_actuator(actuator_type, site=f'site{i}', refsite=f'site{j}')
+      actuators.append((actuator_type, attr))
+
     np.random.shuffle(actuators)
     for typ, attr in actuators:
       ET.SubElement(actuator, typ, attr)
@@ -310,9 +354,7 @@ def create_mjcf(
     ET.SubElement(contact, 'pair', attr)
 
   # exclude contacts
-  bodies = list(mjcf.iter('body'))
   body_names = [b.get('name') for b in bodies]
-  n_bodies = len(bodies)
   for _ in range(min(max_contact_excludes, (n_bodies * (n_bodies - 1) // 2))):
     if p(50):
       continue

mjx/mujoco/mjx/_src/types.py

@@ -20,9 +20,7 @@ from typing import Sequence
 import jax
 import jax.numpy as jp
 import mujoco
-# pylint: disable=g-importing-member
+from mujoco.mjx._src.dataclasses import PyTreeNode  # pylint: disable=g-importing-member
-from mujoco.mjx._src.dataclasses import PyTreeNode
-# pylint: enable=g-importing-member
 import numpy as np
 
 
@@ -155,9 +153,11 @@ class TrnType(enum.IntEnum):
 
   Attributes:
     JOINT: force on joint
+    SITE: force on site
   """
   JOINT = mujoco.mjtTrn.mjTRN_JOINT
-  # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, SITE, BODY
+  SITE = mujoco.mjtTrn.mjTRN_SITE
+  # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, BODY
 
 
 class DynType(enum.IntEnum):
@@ -166,11 +166,14 @@ class DynType(enum.IntEnum):
   Attributes:
     NONE: no internal dynamics; ctrl specifies force
     INTEGRATOR: integrator: da/dt = u
+    FILTER: linear filter: da/dt = (u-a) / tau
+    FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration
   """
   NONE = mujoco.mjtDyn.mjDYN_NONE
   INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
   FILTER = mujoco.mjtDyn.mjDYN_FILTER
-  # unsupported: FILTEREXACT, MUSCLE, USER
+  FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT
+  # unsupported: MUSCLE, USER
 
 
 class GainType(enum.IntEnum):
@@ -213,7 +216,6 @@ class Option(PyTreeNode):
     integrator:       integration mode
     cone:             type of friction cone
     solver:           solver algorithm
-    integrator:       integration mode
     iterations:       number of main solver iterations
     ls_iterations:    maximum number of CG/Newton linesearch iterations
     disableflags:     bit flags for disabling standard features
@@ -260,7 +262,9 @@ class Model(PyTreeNode):
     nbody: number of bodies
     njnt: number of joints
     ngeom: number of geoms
+    nsite: number of sites
     nmesh: number of meshes
+    nmat: number of materials
     npair: number of predefined geom pairs
     nexclude: number of excluded geom pairs
     neq: number of equality constraints
@@ -317,6 +321,8 @@ class Model(PyTreeNode):
     geom_conaffinity: geom contact affinity                   (ngeom,)
     geom_condim: contact dimensionality (1, 3, 4, 6)          (ngeom,)
     geom_bodyid: id of geom's body                            (ngeom,)
+    geom_group: group for visibility                          (ngeom,)
+    geom_matid: material id for rendering                     (ngeom,)
     geom_priority: geom contact priority                      (ngeom,)
     geom_solmix: mixing coef for solref/imp in geom pair      (ngeom,)
     geom_solref: constraint solver reference: contact         (ngeom, mjNREF)
@@ -327,6 +333,11 @@ class Model(PyTreeNode):
     geom_friction: friction for (slide, spin, roll)           (ngeom, 3)
     geom_margin: include in solver if dist<margin-gap         (ngeom,)
     geom_gap: include in solver if dist<margin-gap            (ngeom,)
+    geom_rgba: rgba when material is omitted                  (ngeom, 4)
+    site_bodyid: id of site's body                            (nsite,)
+    site_pos: local position offset rel. to body              (nsite, 3)
+    site_quat: local orientation offset rel. to body          (nsite, 4)
+    mat_rgba: rgba                                            (nmat, 4)
     geom_convex_face: vertex face data, MJX only              (ngeom,)
     geom_convex_vert: vertex data, MJX only                   (ngeom,)
     geom_convex_edge: unique edge data, MJX only              (ngeom,)
@@ -377,7 +388,9 @@ class Model(PyTreeNode):
   nbody: int
   njnt: int
   ngeom: int
+  nsite: int
   nmesh: int
+  nmat: int
   npair: int
   nexclude: int
   neq: int
@@ -434,6 +447,8 @@ class Model(PyTreeNode):
   geom_conaffinity: np.ndarray
   geom_condim: np.ndarray
   geom_bodyid: np.ndarray
+  geom_group: np.ndarray
+  geom_matid: np.ndarray
   geom_priority: np.ndarray
   geom_solmix: jax.Array
   geom_solref: jax.Array
@@ -444,6 +459,11 @@ class Model(PyTreeNode):
   geom_friction: jax.Array
   geom_margin: jax.Array
   geom_gap: jax.Array
+  geom_rgba: np.ndarray
+  site_bodyid: np.ndarray
+  site_pos: jax.Array
+  site_quat: jax.Array
+  mat_rgba: np.ndarray
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
   pair_geom2: np.ndarray
@@ -500,10 +520,8 @@ class Contact(PyTreeNode):
     solref: constraint solver reference, normal direction             (mjNREF,)
     solreffriction: constraint solver reference, friction directions  (mjNREF,)
     solimp: constraint solver impedance                               (mjNIMP,)
-    dim: contact space dimensionality: 1, 3, 4 or 6
     geom1: id of geom 1
     geom2: id of geom 2
-    efc_address: address in efc; -1: not included
   """
   dist: jax.Array
   pos: jax.Array
@@ -513,29 +531,25 @@ class Contact(PyTreeNode):
   solref: jax.Array
   solreffriction: jax.Array
   solimp: jax.Array
-  # unsupported: mu, H
+  # unsupported: mu, H, dim
-  dim: np.ndarray
   geom1: jax.Array
   geom2: jax.Array
-  efc_address: np.ndarray
+  # unsupported: efc_address, exclude
-  # unsupported: exclude
 
   @classmethod
-  def zero(cls, shape=(0,)) -> 'Contact':
+  def zero(cls, ncon: int = 0) -> 'Contact':
     """Returns a contact filled with zeros."""
     return Contact(
-        dist=jp.zeros(shape),
+        dist=jp.zeros(ncon),
-        pos=jp.zeros(shape + (3,)),
+        pos=jp.zeros((ncon, 3,)),
-        frame=jp.zeros(shape + (3, 3)),
+        frame=jp.zeros((ncon, 3, 3)),
-        includemargin=jp.zeros(shape),
+        includemargin=jp.zeros(ncon),
-        friction=jp.zeros(shape + (5,)),
+        friction=jp.zeros((ncon, 5)),
-        solref=jp.zeros(shape + (mujoco.mjNREF,)),
+        solref=jp.zeros((ncon, mujoco.mjNREF)),
-        solreffriction=jp.zeros(shape + (mujoco.mjNREF,)),
+        solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
-        solimp=jp.zeros(shape + (mujoco.mjNIMP,)),
+        solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
-        dim=np.zeros(shape, dtype=np.int32),
+        geom1=jp.zeros(ncon, dtype=jp.int32),
-        geom1=jp.zeros(shape, dtype=jp.int32),
+        geom2=jp.zeros(ncon, dtype=jp.int32),
-        geom2=jp.zeros(shape, dtype=jp.int32),
-        efc_address=np.zeros(shape, dtype=np.int32),
     )
 
 
@@ -544,11 +558,6 @@ class Data(PyTreeNode):
 
   Attributes:
     solver_niter: number of solver iterations, per island         (mjNISLAND,)
-    ne: number of equality constraints
-    nf: number of friction constraints
-    nl: number of limit constraints
-    nefc: number of constraints
-    ncon: nubmer of contacts
     time: simulation time
     qpos: position                                                (nq,)
     qvel: velocity                                                (nv,)
@@ -569,6 +578,8 @@ class Data(PyTreeNode):
     xaxis: Cartesian joint axis                                   (njnt, 3)
     geom_xpos: Cartesian geom position                            (ngeom, 3)
     geom_xmat: Cartesian geom orientation                         (ngeom, 3, 3)
+    site_xpos: Cartesian site position                            (nsite, 3)
+    site_xmat: Cartesian site orientation                         (nsite, 9)
     subtree_com: center of mass of each subtree                   (nbody, 3)
     cdof: com-based motion axis of each dof                       (nv, 6)
     cinert: com-based body inertia and mass                       (nbody, 10)
@@ -589,7 +600,6 @@ class Data(PyTreeNode):
     qfrc_bias: C(qpos,qvel)                                       (nv,)
     qfrc_passive: passive force                                   (nv,)
     efc_aref: reference pseudo-acceleration                       (nefc,)
-    actuator_force: actuator force in actuation space             (nu,)
     qfrc_actuator: actuator force                                 (nv,)
     qfrc_smooth: net unconstrained force                          (nv,)
     qacc_smooth: unconstrained acceleration                       (nv,)
@@ -600,12 +610,6 @@ class Data(PyTreeNode):
   """
   # solver statistics:
   solver_niter: jax.Array
-  # sizes (variable in MJ, constant in MJX)
-  ne: int
-  nf: int
-  nl: int
-  nefc: int
-  ncon: int
   # global properties:
   time: jax.Array
   # state:
@@ -631,6 +635,8 @@ class Data(PyTreeNode):
   xaxis: jax.Array
   geom_xpos: jax.Array
   geom_xmat: jax.Array
+  site_xpos: jax.Array
+  site_xmat: jax.Array
   subtree_com: jax.Array
   cdof: jax.Array
   cinert: jax.Array
@@ -653,7 +659,6 @@ class Data(PyTreeNode):
   qfrc_passive: jax.Array
   efc_aref: jax.Array
   # position, velcoity, control & acceleration dependent:
-  actuator_force: jax.Array
   qfrc_actuator: jax.Array
   qfrc_smooth: jax.Array
   qacc_smooth: jax.Array

mjx/mujoco/mjx/integration_test/collision_driver_test.py

@@ -58,9 +58,9 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
     )
 
     m = mujoco.MjModel.from_xml_string(mjcf)
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
     d = mujoco.MjData(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_step(m, d)
     collision_jit_fn = jax.jit(mjx.collision)
@@ -82,7 +82,6 @@ class CollisionDriverIntegrationTest(parameterized.TestCase):
     mjx_contact = jax.tree_map(
         lambda x: x.take(np.array(idx), axis=0), dx.contact
     )
-    mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx])
     for field in dataclasses.fields(Contact):
       _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)
 

mjx/mujoco/mjx/integration_test/forward_test.py

@@ -19,7 +19,6 @@ from absl.testing import parameterized
 import jax
 import mujoco
 from mujoco import mjx
-from mujoco.mjx._src import forward
 from mujoco.mjx._src import test_util
 import numpy as np
 
@@ -46,7 +45,7 @@ class ActuationIntegrationTest(parameterized.TestCase):
         enable_contact=False,
     )
     m = mujoco.MjModel.from_xml_string(mjcf)
-    actuation_jit_fn = jax.jit(forward._actuation)
+    actuation_jit_fn = jax.jit(mjx.fwd_actuation)
 
     # init
     d = mujoco.MjData(m)
@@ -57,8 +56,8 @@ class ActuationIntegrationTest(parameterized.TestCase):
     mujoco.mj_fwdVelocity(m, d)
 
     # put on device
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_fwdActuation(m, d)
     dx = actuation_jit_fn(mx, dx)

mjx/mujoco/mjx/integration_test/smooth_test.py

@@ -57,17 +57,21 @@ class TransmissionIntegrationTest(parameterized.TestCase):
     d = mujoco.MjData(m)
     d.ctrl = np.random.normal(scale=10, size=m.nu)
     d.act = np.random.normal(scale=10, size=m.na)
+    d.qpos = np.random.normal(m.nq)
     d.qvel = np.random.random(m.nv)
+    mujoco.mj_forward(m, d)
 
     # put on device
-    mx = mjx.device_put(m)
+    mx = mjx.put_model(m)
-    dx = mjx.device_put(d)
+    dx = mjx.put_data(m, d)
 
     mujoco.mj_transmission(m, d)
     dx = transmission_jit_fn(mx, dx)
 
     _assert_attr_eq(d, dx, 'actuator_length', seed, f'transmission{seed}')
-    _assert_attr_eq(d, dx, 'actuator_moment', seed, f'transmission{seed}')
+    _assert_attr_eq(
+        d, dx, 'actuator_moment', seed, f'transmission{seed}', atol=1e-4
+    )
 
 
 if __name__ == '__main__':

mjx/mujoco/mjx/test_data/ant.xml

@@ -1,82 +0,0 @@
-<mujoco model="ant">
-  <compiler angle="degree" coordinate="local" inertiafromgeom="true"/>
-  <option timestep="0.01" iterations="4" solver="CG" />
-  <default>
-    <joint armature="1" damping="1" limited="true"/>
-    <geom contype="0" conaffinity="0" condim="3" density="5.0" friction="1 0.5 0.5"/>
-  </default>
-  <asset>
-    <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>
-    <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
-    <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
-    <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
-    <material name="geom" texture="texgeom" texuniform="true"/>
-  </asset>
-  <worldbody>
-    <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
-    <geom conaffinity="1" condim="3" material="MatPlane" name="floor" pos="0 0 0" size="40 40 40" type="plane"/>
-    <body name="torso" pos="0 0 0.75">
-      <camera name="track" mode="trackcom" pos="0 -3 0.3" xyaxes="1 0 0 0 0 1"/>
-      <geom name="torso_geom" pos="0 0 0" size="0.25" type="sphere"/>
-      <joint armature="0" damping="0" limited="false" margin="0.01" name="root" pos="0 0 0" type="free"/>
-      <body name="front_left_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="aux_1_geom" size="0.08" type="capsule"/>
-        <body name="aux_1" pos="0.2 0.2 0">
-          <joint axis="0 0 1" name="hip_1" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 0.2 0.2 0.0" name="left_leg_geom" size="0.08" type="capsule"/>
-          <body pos="0.2 0.2 0">
-            <joint axis="-1 1 0" name="ankle_1" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 0.4 0.4 0.0" name="left_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="left_foot_geom" contype="1" pos="0.4 0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="front_right_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="aux_2_geom" size="0.08" type="capsule"/>
-        <body name="aux_2" pos="-0.2 0.2 0">
-          <joint axis="0 0 1" name="hip_2" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 -0.2 0.2 0.0" name="right_leg_geom" size="0.08" type="capsule"/>
-          <body pos="-0.2 0.2 0">
-            <joint axis="1 1 0" name="ankle_2" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 -0.4 0.4 0.0" name="right_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="right_foot_geom" contype="1" pos="-0.4 0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="back_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="aux_3_geom" size="0.08" type="capsule"/>
-        <body name="aux_3" pos="-0.2 -0.2 0">
-          <joint axis="0 0 1" name="hip_3" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 -0.2 -0.2 0.0" name="back_leg_geom" size="0.08" type="capsule"/>
-          <body pos="-0.2 -0.2 0">
-            <joint axis="-1 1 0" name="ankle_3" pos="0.0 0.0 0.0" range="-70 -30" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 -0.4 -0.4 0.0" name="third_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="third_foot_geom" contype="1" pos="-0.4 -0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-      <body name="right_back_leg" pos="0 0 0">
-        <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="aux_4_geom" size="0.08" type="capsule"/>
-        <body name="aux_4" pos="0.2 -0.2 0">
-          <joint axis="0 0 1" name="hip_4" pos="0.0 0.0 0.0" range="-30 30" type="hinge"/>
-          <geom fromto="0.0 0.0 0.0 0.2 -0.2 0.0" name="rightback_leg_geom" size="0.08" type="capsule"/>
-          <body pos="0.2 -0.2 0">
-            <joint axis="1 1 0" name="ankle_4" pos="0.0 0.0 0.0" range="30 70" type="hinge"/>
-            <geom fromto="0.0 0.0 0.0 0.4 -0.4 0.0" name="fourth_ankle_geom" size="0.08" type="capsule"/>
-            <geom name="fourth_foot_geom" contype="1" pos="0.4 -0.4 0" size="0.08" type="sphere" mass="0"/>
-          </body>
-        </body>
-      </body>
-    </body>
-  </worldbody>
-  <actuator>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_4" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_4" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_1" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_1" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_2" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_2" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="hip_3" gear="150"/>
-    <motor ctrllimited="true" ctrlrange="-1.0 1.0" joint="ankle_3" gear="150"/>
-  </actuator>
-</mujoco>

mjx/mujoco/mjx/test_data/constraints.xml

@@ -0,0 +1,53 @@
+<!-- For validating constraint dynamics:
+
+* connect, weld, joint constraints
+* collision and joint limits for ball and 1d joints
+* solref, solimp
+-->
+<mujoco>
+  <option timestep="0.02"/>
+
+  <default>
+    <default class="box">
+      <geom type="box" size=".2" fromto="0 0 0 0 -2 0" rgba=".4 .7 .6 .3"  contype="0"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <geom pos="0 0 -1" type="plane" size="10 10 .01"/>
+
+    <body name="anchor1" pos="-3 0 0"/>
+    <body name="beam1" pos="-3 0 0">
+      <joint name="joint1" type="ball" range="0 45" solreflimit="0.03 0.9" solimplimit="0.89 0.9 0.01 2.1"/>
+      <geom class="box"/>
+    </body>
+
+    <body name="anchor2" pos="-1 0 0"/>
+    <body name="beam2" pos="-1 0 0">
+      <freejoint/>
+      <geom class="box"/>
+    </body>
+
+    <body name="beam3" pos="1 0 0">
+      <joint name="joint3" axis="1 0 0" type="hinge" range="-20 20"/>
+      <geom class="box"/>
+    </body>
+
+    <body name="beam4" pos="3 0 0">
+      <joint name="joint4" axis="1 0 0" type="hinge" damping="1.0"/> <!-- tests no joint range -->
+      <geom class="box"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <connect name="connect"  body1="anchor1"    body2="beam1"   anchor="1 0 -1" />
+    <weld    name="weld"     body1="anchor2"    body2="beam2"   relpose="0 0 0 1 -.3 0 0" torquescale="0.002" anchor="0 -2 0"/>
+    <joint   name="joint"    joint1="joint3"    joint2="joint4" polycoef="0.5 -1 0.1 0.15 0.2" />
+  </equality>
+
+  <actuator>
+    <position ctrlrange="-20 20" gear="500" joint="joint1" name="act1"/>
+    <motor gear="50000" joint="joint3" name="act2"/>
+    <motor gear="75000" joint="joint4" name="act3"/>
+  </actuator>
+</mujoco>

mjx/mujoco/mjx/test_data/equality.xml

@@ -1,71 +0,0 @@
-<mujoco>
-  <option solver="CG" iterations="8">
-    <flag contact="disable"/>
-  </option>
-
-  <default>
-    <default class="static">
-      <geom size=".5 .1 .5" rgba=".5 .7 .5 .3"/>
-    </default>
-    <default class="free">
-      <geom type="box" size=".2" fromto="0 0 0 0 -2 0" rgba=".4 .7 .6 .3"/>
-    </default>
-  </default>
-
-  <worldbody>
-    <geom pos="0 0 -2" type="plane" size="10 10 .01"/>
-    <light pos="0 0 20"/>
-
-    <body name="box1" pos="-3 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam1" pos="-3 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box2" pos="-1 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam2" pos="-1 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box3" pos="1 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam3" pos="1 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box4" pos="3 0 0">
-      <geom type="box" class="static"/>
-    </body>
-    <body name="beam4" pos="3 0 0">
-      <freejoint/>
-      <geom class="free"/>
-    </body>
-
-    <body name="box5" pos="4 0 0">
-      <geom class="free"/>
-      <joint name="joint1" axis="1 0 0" type="hinge" />
-    </body>
-
-    <body name="box6" pos="4 0 0">
-      <geom class="free"/>
-      <joint name="joint2" axis="1 0 0" type="hinge" />
-    </body>
-
-
-  </worldbody>
-
-  <equality>
-    <connect name="connect anchor"           body1="box1"    body2="beam1"   anchor="0 0 -1" />
-    <weld    name="weld anchor weak torques" body1="box2"    body2="beam2"   torquescale="0.002" anchor="0 -2 0"/>
-    <weld    name="weld relpose"             body1="box3"    body2="beam3"   relpose="0 0 0 1 -.3 0 0"/>
-    <weld    name="weld relpose+anchor"      body1="box4"    body2="beam4"   relpose="0 0 0 1 -.3 0 0" anchor="0 0 -1"/>
-    <joint   name="joint"                    joint1="joint1" joint2="joint2" polycoef="0 -1 0.1 0.15 0.2" />
-  </equality>
-</mujoco>

mjx/mujoco/mjx/test_data/humanoid.xml

@@ -1,109 +0,0 @@
-<mujoco model="humanoid">
-    <compiler angle="degree" inertiafromgeom="true"/>
-    <default>
-        <joint armature="1" damping="1" limited="true"/>
-        <geom conaffinity="0" condim="3" contype="0" material="geom"/>
-        <motor ctrllimited="true" ctrlrange="-.4 .4"/>
-    </default>
-    <option iterations="8" timestep="0.003"/>
-    <size nkey="5" nuser_geom="1"/>
-    <visual>
-        <map fogend="5" fogstart="3"/>
-    </visual>
-    <asset>
-        <texture builtin="gradient" height="100" rgb1=".4 .5 .6" rgb2="0 0 0" type="skybox" width="100"/>
-        <!-- <texture builtin="gradient" height="100" rgb1="1 1 1" rgb2="0 0 0" type="skybox" width="100"/>-->
-        <texture builtin="flat" height="1278" mark="cross" markrgb="1 1 1" name="texgeom" random="0.01" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" type="cube" width="127"/>
-        <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
-        <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
-        <material name="geom" texture="texgeom" texuniform="true"/>
-    </asset>
-    <worldbody>
-        <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
-        <geom conaffinity="1" condim="3" friction="1 .1 .1" material="MatPlane" name="floor" pos="0 0 0" size="20 20 0.125" type="plane"/>
-        <!-- <geom condim="3" material="MatPlane" name="floor" pos="0 0 0" size="10 10 0.125" type="plane"/>-->
-        <body name="torso" pos="0 0 1.4">
-            <camera name="track" mode="trackcom" pos="0 -4 0" xyaxes="1 0 0 0 0 1"/>
-            <joint armature="0" damping="0" limited="false" name="root" pos="0 0 0" stiffness="0" type="free"/>
-            <geom fromto="0 -.07 0 0 .07 0" name="torso1" size="0.07" type="capsule"/>
-            <geom name="head" pos="0 0 .19" size=".09" type="sphere" user="258"/>
-            <geom fromto="-.01 -.06 -.12 -.01 .06 -.12" name="uwaist" size="0.06" type="capsule"/>
-            <body name="lwaist" pos="-.01 0 -0.260" quat="1.000 0 -0.002 0">
-                <geom fromto="0 -.06 0 0 .06 0" name="lwaist" size="0.06" type="capsule"/>
-                <joint armature="0.02" axis="0 0 1" damping="5" name="abdomen_z" pos="0 0 0.065" range="-45 45" stiffness="20" type="hinge"/>
-                <joint armature="0.02" axis="0 1 0" damping="5" name="abdomen_y" pos="0 0 0.065" range="-75 30" stiffness="10" type="hinge"/>
-                <body name="pelvis" pos="0 0 -0.165" quat="1.000 0 -0.002 0">
-                    <joint armature="0.02" axis="1 0 0" damping="5" name="abdomen_x" pos="0 0 0.1" range="-35 35" stiffness="10" type="hinge"/>
-                    <geom fromto="-.02 -.07 0 -.02 .07 0" name="butt" size="0.09" type="capsule"/>
-                    <body name="right_thigh" pos="0 -0.1 -0.04">
-                        <joint armature="0.01" axis="1 0 0" damping="5" name="right_hip_x" pos="0 0 0" range="-25 5" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 0 1" damping="5" name="right_hip_z" pos="0 0 0" range="-60 35" stiffness="10" type="hinge"/>
-                        <joint armature="0.0080" axis="0 1 0" damping="5" name="right_hip_y" pos="0 0 0" range="-110 20" stiffness="20" type="hinge"/>
-                        <geom fromto="0 0 0 0 0.01 -.34" name="right_thigh1" size="0.06" type="capsule"/>
-                        <body name="right_shin" pos="0 0.01 -0.403">
-                            <joint armature="0.0060" axis="0 -1 0" name="right_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
-                            <geom fromto="0 0 0 0 0 -.3" name="right_shin1" size="0.049" type="capsule"/>
-                            <body name="right_foot" pos="0 0 -0.45">
-                                <geom contype="1" name="right_foot" pos="0 0 0.1" size="0.075" type="sphere" user="0"/>
-                            </body>
-                        </body>
-                    </body>
-                    <body name="left_thigh" pos="0 0.1 -0.04">
-                        <joint armature="0.01" axis="-1 0 0" damping="5" name="left_hip_x" pos="0 0 0" range="-25 5" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 0 -1" damping="5" name="left_hip_z" pos="0 0 0" range="-60 35" stiffness="10" type="hinge"/>
-                        <joint armature="0.01" axis="0 1 0" damping="5" name="left_hip_y" pos="0 0 0" range="-110 20" stiffness="20" type="hinge"/>
-                        <geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
-                        <body name="left_shin" pos="0 -0.01 -0.403">
-                            <joint armature="0.0060" axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" stiffness="1" type="hinge"/>
-                            <geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
-                            <body name="left_foot" pos="0 0 -0.45">
-                                <geom contype="1" name="left_foot" type="sphere" size="0.075" pos="0 0 0.1" user="0" />
-                            </body>
-                        </body>
-                    </body>
-                </body>
-            </body>
-            <body name="right_upper_arm" pos="0 -0.17 0.06">
-                <joint armature="0.0068" axis="2 1 1" name="right_shoulder1" pos="0 0 0" range="-85 60" stiffness="1" type="hinge"/>
-                <joint armature="0.0051" axis="0 -1 1" name="right_shoulder2" pos="0 0 0" range="-85 60" stiffness="1" type="hinge"/>
-                <geom fromto="0 0 0 .16 -.16 -.16" name="right_uarm1" size="0.04 0.16" type="capsule"/>
-                <body name="right_lower_arm" pos=".18 -.18 -.18">
-                    <joint armature="0.0028" axis="0 -1 1" name="right_elbow" pos="0 0 0" range="-90 50" stiffness="0" type="hinge"/>
-                    <geom fromto="0.01 0.01 0.01 .17 .17 .17" name="right_larm" size="0.031" type="capsule"/>
-                    <geom name="right_hand" pos=".18 .18 .18" size="0.04" type="sphere"/>
-                    <camera pos="0 0 0"/>
-                </body>
-            </body>
-            <body name="left_upper_arm" pos="0 0.17 0.06">
-                <joint armature="0.0068" axis="2 -1 1" name="left_shoulder1" pos="0 0 0" range="-60 85" stiffness="1" type="hinge"/>
-                <joint armature="0.0051" axis="0 1 1" name="left_shoulder2" pos="0 0 0" range="-60 85" stiffness="1" type="hinge"/>
-                <geom fromto="0 0 0 .16 .16 -.16" name="left_uarm1" size="0.04 0.16" type="capsule"/>
-                <body name="left_lower_arm" pos=".18 .18 -.18">
-                    <joint armature="0.0028" axis="0 -1 -1" name="left_elbow" pos="0 0 0" range="-90 50" stiffness="0" type="hinge"/>
-                    <geom fromto="0.01 -0.01 0.01 .17 -.17 .17" name="left_larm" size="0.031" type="capsule"/>
-                    <geom name="left_hand" pos=".18 -.18 .18" size="0.04" type="sphere"/>
-                </body>
-            </body>
-        </body>
-    </worldbody>
-
-    <actuator>
-        <motor gear="100" joint="abdomen_y" name="abdomen_y"/>
-        <motor gear="100" joint="abdomen_z" name="abdomen_z"/>
-        <motor gear="100" joint="abdomen_x" name="abdomen_x"/>
-        <motor gear="100" joint="right_hip_x" name="right_hip_x"/>
-        <motor gear="100" joint="right_hip_z" name="right_hip_z"/>
-        <motor gear="300" joint="right_hip_y" name="right_hip_y"/>
-        <motor gear="200" joint="right_knee" name="right_knee"/>
-        <motor gear="100" joint="left_hip_x" name="left_hip_x"/>
-        <motor gear="100" joint="left_hip_z" name="left_hip_z"/>
-        <motor gear="300" joint="left_hip_y" name="left_hip_y"/>
-        <motor gear="200" joint="left_knee" name="left_knee"/>
-        <motor gear="25" joint="right_shoulder1" name="right_shoulder1"/>
-        <motor gear="25" joint="right_shoulder2" name="right_shoulder2"/>
-        <motor gear="25" joint="right_elbow" name="right_elbow"/>
-        <motor gear="25" joint="left_shoulder1" name="left_shoulder1"/>
-        <motor gear="25" joint="left_shoulder2" name="left_shoulder2"/>
-        <motor gear="25" joint="left_elbow" name="left_elbow"/>
-    </actuator>
-</mujoco>

mjx/mujoco/mjx/test_data/pendula.xml

@@ -18,6 +18,8 @@
   </default>
 
   <worldbody>
+    <site name="origin"/>
+
     <!-- a single free body -->
     <body pos="0 0 0">
       <freejoint/>
@@ -26,45 +28,49 @@
 
     <!-- a single ball joint with a limit -->
     <body pos="0.5 0 0">
-      <joint type="ball" range="0 35"/>
+      <joint name="joint1" type="ball" range="0 35"/>
+      <site name="s1" pos="1.3 0 0"/>
       <geom/>
     </body>
 
     <!-- a single slide joint with a limit -->
     <body pos="1.0 0 0">
-      <joint type="slide" axis="0.1 0.2 0.3" range="-1 1"/>
+      <joint name="joint2" type="slide" axis="0.1 0.2 0.3" range="-1 1"/>
+      <site name="s2"/>
       <geom/>
     </body>
 
     <!-- a single hinge joint with a limit -->
     <body pos="1.5 0 0">
-      <joint type="hinge" axis="0.1 0.2 0.3" range="-35 50"/>
+      <joint name="joint3" type="hinge" axis="0.1 0.2 0.3" range="-35 50"/>
+      <site name="s3" quat="1 0 1 0"/>
       <geom/>
     </body>
 
     <!-- stacked joint: hinge + slide -->
     <body pos="2.0 0 0">
-      <joint type="hinge" axis="0.1 0.2 0.3"/>
+      <joint name="joint4" type="hinge" axis="0.1 0.2 0.3"/>
-      <joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
+      <joint name="joint5" type="slide" axis="0.4 0.5 0.6" range="0 1"/>
       <geom/>
     </body>
 
     <!-- stacked joint: slide + ball -->
     <body pos="2.5 0 0">
-      <joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
+      <joint name="joint6" type="slide" axis="0.4 0.5 0.6" range="0 1"/>
       <joint type="ball"/>
       <geom/>
     </body>
 
     <!-- triple pendulum of hinges -->
     <body pos="3.0 0 0">
-      <joint axis="0.1 0.2 0.3" type="hinge"/>
+      <joint name="joint7" axis="0.1 0.2 0.3" type="hinge"/>
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
+        <joint name="joint8" axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
         <geom/>
         <body pos="0 -0.7 0">
-          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <joint name="joint9" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <site name="s4" pos="0 1.5 0" quat="1 0 1 0"/>
           <geom/>
         </body>
       </body>
@@ -72,14 +78,14 @@
 
     <!-- cherry pendulum: two bodies attached to same parent body -->
     <body pos="3.5 0 0">
-      <joint type="ball" damping="0.5" />
+      <joint name="joint10" type="ball" damping="0.5" />
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
+        <joint name="joint11" axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
         <geom/>
       </body>
       <body pos="0 -0.7 0">
-        <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+        <joint name="joint12" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
         <geom/>
       </body>
     </body>
@@ -89,14 +95,21 @@
       <freejoint/>
       <geom/>
       <body pos="0 0 -0.8">
-        <joint axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
+        <joint name="joint13" axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
         <geom/>
         <body pos="0 -0.7 0">
-          <joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
+          <joint name="joint14" axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
           <geom/>
         </body>
       </body>
     </body>
-
   </worldbody>
+
+  <actuator>
+    <motor gear="250 0 0" joint="joint1" name="act1"/>
+    <motor gear="0 275 0" joint="joint1" name="act2"/>
+    <motor gear="0 0 300" joint="joint1" name="act3"/>
+    <motor gear="275" joint="joint2" name="act4"/>
+    <motor gear="275" joint="joint3" name="act5"/>
+  </actuator>
 </mujoco>

mjx/mujoco/mjx/test_data/ray.xml

@@ -0,0 +1,16 @@
+<mujoco model="ray">
+  <asset>
+      <mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.4 0.4 0.4" />
+      <texture builtin="checker" height="100" name="texplane" rgb1="0 0 0" rgb2="0.8 0.8 0.8" type="2d" width="100"/>
+      <material name="MatPlane" reflectance="0.5" shininess="1" specular="1" texrepeat="60 60" texture="texplane"/>
+  </asset>
+  <worldbody>
+    <light cutoff="100" diffuse="1 1 1" dir="-0 0 -1.3" directional="true" exponent="1" pos="0 0 1.3" specular=".1 .1 .1"/>
+
+    <geom name="plane" pos="0 0 0" quat="1 0 0 0" size="4 4 4" type="plane" rgba="0.1 0.1 0.1 1"/>
+    <geom name="sphere" pos="0 0 1" quat="1 0 0 0" size="0.5" type="sphere" rgba="1 0 0 1"/>
+    <geom name="capsule" pos="0 1 1" quat="0 0.3826834 0 0.9238795 " size="0.25 0.5" type="capsule" rgba="0 1 0 1"/>
+    <geom name="box" pos="1 0 1" quat="0 0.3826834 0 0.9238795" size="0.5 0.25 0.3" type="box" rgba="0 0 1 1"/>
+    <geom name="mesh" pos="1 1 1" quat="0 0 0.3826834 0.9238795" type="mesh" mesh="tetrahedron" rgba="1 1 0 1"/>
+  </worldbody>
+</mujoco>

mjx/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name="mujoco-mjx"
-version = "3.0.2"
+version = "3.1.2"
 authors = [
     {name = "Google DeepMind", email = "mujoco@deepmind.com"},
 ]
@@ -31,13 +31,13 @@ dependencies = [
     "etils[epath]",
     "jax",
     "jaxlib",
-    "mujoco>=3.0.2.dev0",
+    "mujoco>=3.1.2.dev0",
     "scipy",
     "trimesh",
 ]
 
 [project.urls]
 Homepage = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
-Documentation = "https://mujoco.readthedocs.io/en/3.0.2"
+Documentation = "https://mujoco.readthedocs.io/en/3.1.2"
 Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
-Changelog = "https://mujoco.readthedocs.io/en/3.0.2/changelog.html"
+Changelog = "https://mujoco.readthedocs.io/en/3.1.2/changelog.html"

mjx/requirements.txt

@@ -65,7 +65,8 @@ scipy==1.11.3; python_version >= '3.9' \
     --hash=sha256:c77da50c9a91e23beb63c2a711ef9e9ca9a2060442757dffee34ea41847d8156 \
     --hash=sha256:9ea7f579182d83d00fed0e5c11a4aa5ffe01460444219dedc448a36adf0c3917 \
     --hash=sha256:5305792c7110e32ff155aed0df46aa60a60fc6e52cd4ee02cdeb67eaccd5356e \
-    --hash=sha256:a63d1ec9cadecce838467ce0631c17c15c7197ae61e49429434ba01d618caa83
+    --hash=sha256:a63d1ec9cadecce838467ce0631c17c15c7197ae61e49429434ba01d618caa83 \
+    --hash=sha256:715c9966eb8906bc67e450e962bd07a5254420077178f98258904da4004a172f
 setuptools==68.2.2 \
     --hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a
 trimesh==4.0.0 \

model/humanoid/README.md

@@ -1,17 +1,34 @@
-Humanoid
+# Humanoid
-========
-
-Degrees of Freedom: 27
-Actuators: 21
 
 This simplified humanoid model, introduced in [1], is designed for bipedal locomotion
 behaviours. While several variants of it exist in the wild, this version is based on the model
 in the DeepMind Control Suite [2], which has fairly realistic actuator gains.
 
-[1] [Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization]
+* Degrees of Freedom: 27
-     (https://doi.org/10.1109/IROS.2012.6386025).
+* Actuators: 21
+
+<p float="left">
+  <img src="humanoid.png" width="400">
+</p>
+
+## Changelog
+
+* 02-01-2024: Add more keyframes.
+* 27-11-2023: Move humanoid geoms to group 1.
+* 05-04-2023: Fix typo in texture size.
+* 20-09-2022: Use default class for left_upper_arm geom.
+* 17-09-2022: Increase offscreen render buffer resolution of the humanoid to 2560x1440.
+* 12-09-2022:
+  * Increased maximum hip flexion.
+  * Symmetrised shoulder and ankle joints.
+  * Added hamstring tendons which couple the hip and knee at large flexion values.
+  * Moved duplicated values into defaults.
+  * Added two keyframes.
+  * Improved lighting.
+  * Changed naming convention.
+
+## References
+
+[1] [Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization](https://doi.org/10.1109/IROS.2012.6386025).
 
 [2] [DeepMind Control Suite](https://arxiv.org/abs/1801.00690).
-
-
-![humanoid](humanoid.png)

model/humanoid/humanoid.xml

@@ -37,7 +37,7 @@
     <default class="body">
 
       <!-- geoms -->
-      <geom type="capsule" condim="1" friction=".7" solimp=".9 .99 .003" solref=".015 1" material="body"/>
+      <geom type="capsule" condim="1" friction=".7" solimp=".9 .99 .003" solref=".015 1" material="body" group="1"/>
       <default class="thigh">
         <geom size=".06"/>
       </default>
@@ -233,17 +233,35 @@
       left leg
       arms
     -->
-    <key name="squat" qpos="0 0 0.596
+    <key name="squat"
-                            0.988015 0 0.154359 0
+         qpos="0 0 0.596
-                            0 0.4 0
+               0.988015 0 0.154359 0
-                            -0.25 -0.5 -2.5 -2.65 -0.8 0.56
+               0 0.4 0
-                            -0.25 -0.5 -2.5 -2.65 -0.8 0.56
+               -0.25 -0.5 -2.5 -2.65 -0.8 0.56
-                            0 0 0 0 0 0"/>
+               -0.25 -0.5 -2.5 -2.65 -0.8 0.56
-    <key name="stand_on_left_leg" qpos="0 0 1.21948
+               0 0 0 0 0 0"/>
-                                        0.971588 -0.179973 0.135318 -0.0729076
+    <key name="stand_on_left_leg"
-                                        -0.0516 -0.202 0.23
+         qpos="0 0 1.21948
-                                        -0.24 -0.007 -0.34 -1.76 -0.466 -0.0415
+               0.971588 -0.179973 0.135318 -0.0729076
-                                        -0.08 -0.01 -0.37 -0.685 -0.35 -0.09
+               -0.0516 -0.202 0.23
-                                        0.109 -0.067 -0.7 -0.05 0.12 0.16"/>
+               -0.24 -0.007 -0.34 -1.76 -0.466 -0.0415
+               -0.08 -0.01 -0.37 -0.685 -0.35 -0.09
+               0.109 -0.067 -0.7 -0.05 0.12 0.16"/>
+    <key name="prone"
+         qpos="0.4 0 0.0757706
+               0.7325 0 0.680767 0
+               0 0.0729 0
+               0.0077 0.0019 -0.026 -0.351 -0.27 0
+               0.0077 0.0019 -0.026 -0.351 -0.27 0
+               0.56 -0.62 -1.752
+               0.56 -0.62 -1.752"/>
+    <key name="supine"
+         qpos="-0.4 0 0.08122
+               0.722788 0 -0.69107 0
+               0 -0.25 0
+               0.0182 0.0142 0.3 0.042 -0.44 -0.02
+               0.0182 0.0142 0.3 0.042 -0.44 -0.02
+               0.186 -0.73 -1.73
+               0.186 -0.73 -1.73"/>
   </keyframe>
 </mujoco>

model/plugin/actuator/pid.xml

@@ -0,0 +1,68 @@
+<mujoco model="PID Controllers">
+  <size memory="10M"/>
+
+  <visual>
+    <headlight diffuse=".2 .2 .2"/>
+    <global elevation="-30" />
+  </visual>
+  <statistic extent="0.50" />
+
+  <extension>
+    <plugin plugin="mujoco.pid">
+      <!-- pure proportional gain, bounces up and down -->
+      <instance name="pid1">
+        <config key="kp" value="4.0"/>
+      </instance>
+      <!-- PD controller - settles below setpoint -->
+      <instance name="pid2">
+        <config key="kp" value="4.0"/>
+        <config key="kd" value="0.4"/>
+      </instance>
+      <!-- PID controller - reaches setpoint -->
+      <instance name="pid3">
+        <config key="kp" value="4.0"/>
+        <config key="ki" value="4"/>
+        <config key="kd" value="0.4"/>
+      </instance>
+      <!-- PID controller with clamped I term. -->
+      <instance name="pid4">
+        <config key="kp" value="4.0"/>
+        <config key="ki" value="4"/>
+        <config key="kd" value="0.4"/>
+        <config key="imax" value=".5"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <default>
+    <geom type="sphere" size="0.01" rgba="1 1 1 1" />
+  </default>
+
+  <worldbody>
+    <light pos="0.4 -.15 .6" dir="-1 .5 -1" diffuse=".7 .7 .7"/>
+    <light pos="-.4 -.15 .6" dir="1 .5 -1" diffuse=".7 .7 .7"/>
+    <geom name="floor" type="plane" pos="0 0 -0.1" size="0 0 100" />
+    <body pos="-0.06 0 0">
+      <joint name="j1" type="slide" axis="0 0 1" />
+      <geom/>
+    </body>
+    <body pos="-0.02 0 0" >
+      <joint name="j2" type="slide" axis="0 0 1" />
+      <geom/>
+    </body>
+    <body pos="0.02 0 0" >
+      <joint name="j3" type="slide" axis="0 0 1" />
+      <geom/>
+    </body>
+    <body pos="0.06 0 0" >
+      <joint name="j4" type="slide" axis="0 0 1" />
+      <geom/>
+    </body>
+  </worldbody>
+  <actuator>
+    <plugin joint="j1" plugin="mujoco.pid" instance="pid1" ctrlrange="0 0.15" />
+    <plugin joint="j2" plugin="mujoco.pid" instance="pid2" ctrlrange="0 0.15" />
+    <plugin joint="j3" plugin="mujoco.pid" instance="pid3" ctrlrange="0 0.15" />
+    <plugin joint="j4" plugin="mujoco.pid" instance="pid4" ctrlrange="0 0.15" />
+  </actuator>
+</mujoco>

model/plugin/elasticity/trampoline_flex.xml

@@ -44,6 +44,7 @@
               radius=".001" mass="10" name="plate" dim="2">
       <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
       <edge equality="false" damping="10"/>
+      <pin id="0 15 240 255"/>
       <plugin plugin="mujoco.elasticity.membrane">
         <config key="poisson" value="0"/>
         <config key="thickness" value="1e-2"/>
@@ -52,11 +53,4 @@
       </plugin>
     </flexcomp>
   </worldbody>
-
-  <equality>
-    <connect body1="plate_0" anchor="0 0 0"/>
-    <connect body1="plate_15" anchor="0 0 0"/>
-    <connect body1="plate_240" anchor="0 0 0"/>
-    <connect body1="plate_255" anchor="0 0 0"/>
-  </equality>
 </mujoco>

model/plugin/sdf/nutbolt.xml

@@ -8,7 +8,7 @@
     </plugin>
     <plugin plugin="mujoco.sdf.bolt">
       <instance name="bolt">
-        <config key="radius" value="0.26"/>
+        <config key="radius" value="0.255"/>
       </instance>
     </plugin>
   </extension>
@@ -30,7 +30,7 @@
     </mesh>
   </asset>
 
-  <option sdf_iterations="15" sdf_initpoints="60"/>
+  <option sdf_iterations="10" sdf_initpoints="20"/>
 
   <default>
     <geom solref="0.01 1" solimp=".95 .99 .0001" friction="0.01"/>

plugin/actuator/CMakeLists.txt

@@ -0,0 +1,41 @@
+# Copyright 2023 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set(MUJOCO_ACTUATOR_INCLUDE ${CMAKE_CURRENT_SOURCE_DIR}/../..
+                            ${CMAKE_CURRENT_SOURCE_DIR}/../../src
+)
+
+set(MUJOCO_ACTUATOR_SRCS
+    pid.cc
+    pid.h
+    register.cc
+)
+
+add_library(actuator SHARED)
+target_sources(actuator PRIVATE ${MUJOCO_ACTUATOR_SRCS})
+target_include_directories(actuator PRIVATE ${MUJOCO_ACTUATOR_INCLUDE})
+target_link_libraries(actuator PRIVATE mujoco)
+target_compile_options(
+  actuator
+  PRIVATE ${AVX_COMPILE_OPTIONS}
+          ${MUJOCO_MACOS_COMPILE_OPTIONS}
+          ${EXTRA_COMPILE_OPTIONS}
+          ${MUJOCO_CXX_FLAGS}
+)
+target_link_options(
+  actuator
+  PRIVATE
+  ${MUJOCO_MACOS_LINK_OPTIONS}
+  ${EXTRA_LINK_OPTIONS}
+)

plugin/actuator/README.md

@@ -0,0 +1,47 @@
+# Actuator plugins
+
+## PID
+
+The `mujoco.pid` actuator plugin implements a configurable [PID controller](https://en.wikipedia.org/wiki/Proportional%E2%80%93integral%E2%80%93derivative_controller):
+
+$$f(t) = K_\text{p} e(t) + K_\text{i} \int_0^t e(\tau) \,\mathrm{d}\tau + K_\text{d} \frac{\mathrm{d}e(t)}{\mathrm{d}t},$$
+where $e(t) = u(t) - \ell(t)$ is the difference between the control $u$ and the actuator length $\ell$.
+
+You can use it like:
+
+```xml
+<mujoco>
+  <extension>
+    <plugin plugin="mujoco.pid">
+      <instance name="pid">
+        <config key="kp" value="40.0"/>
+        <config key="ki" value="40"/>
+        <config key="kd" value="4"/>
+        <config key="slewmax" value="3" />
+        <config key="imax" value="1"/>
+      </instance>
+    </plugin>
+  </extension>
+
+  <worldbody>
+    <body>
+      <joint name="j" type="slide" axis="0 0 1" />
+      <geom size="0.01" mass="1"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <plugin joint="j" plugin="mujoco.pid" instance="pid" />
+  </actuator>
+</mujoco>
+```
+
+The available options are:
+
+|Attribute | Default | Meaning |
+|----------|---------|---------|
+|`kp` | 0 | **P** gain for the controller. |
+|`ki` | 0 | **I** gain for the controller.<p/>If nonzero, one activation variable will be added to `mjData.act`, containing the current I term (in units of force). |
+|`kd` | 0 | **D** gain for the controller. |
+|`imax` | Optional | If specified, the force produced by the I term will be clipped to the range `[-imax, -imax]`. |
+|`slewmax` | Optional | The maximum rate at which the setpoint for the PID controller can change.<p/>If a bigger change is requested between two timesteps, it will be clipped to the range `[ctrl - slewmax * dt, ctrl + slewmax * dt]`<p/>If specified, one activation variable will be added to `mjData.act` containing the previous value of `ctrl`. |

plugin/actuator/pid.cc

@@ -0,0 +1,285 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "pid.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <optional>
+#include <utility>
+#include <vector>
+
+#include <mujoco/mujoco.h>
+
+namespace mujoco::plugin::actuator {
+namespace {
+
+constexpr char kAttrPGain[] = "kp";
+constexpr char kAttrIGain[] = "ki";
+constexpr char kAttrDGain[] = "kd";
+constexpr char kAttrIMax[] = "imax";
+constexpr char kAttrSlewMax[] = "slewmax";
+
+std::optional<mjtNum> ReadOptionalDoubleAttr(const mjModel* m, int instance,
+                                             const char* attr) {
+  const char* value = mj_getPluginConfig(m, instance, attr);
+  if (value == nullptr || value[0] == '\0') {
+    return std::nullopt;
+  }
+  return std::strtod(value, nullptr);
+}
+
+// returns the next act given the current act_dot, after clamping, for native
+// mujoco dyntypes.
+// copied from engine_forward.
+mjtNum NextActivation(const mjModel* m, const mjData* d, int actuator_id,
+                      int act_adr, mjtNum act_dot) {
+  mjtNum act = d->act[act_adr];
+
+  if (m->actuator_dyntype[actuator_id] == mjDYN_FILTEREXACT) {
+    // exact filter integration
+    // act_dot(0) = (ctrl-act(0)) / tau
+    // act(h) = act(0) + (ctrl-act(0)) (1 - exp(-h / tau))
+    //        = act(0) + act_dot(0) * tau * (1 - exp(-h / tau))
+    mjtNum tau = mju_max(mjMINVAL, m->actuator_dynprm[actuator_id * mjNDYN]);
+    act = act + act_dot * tau * (1 - mju_exp(-m->opt.timestep / tau));
+  } else {
+    // Euler integration
+    act = act + act_dot * m->opt.timestep;
+  }
+
+  // clamp to actrange
+  if (m->actuator_actlimited[actuator_id]) {
+    mjtNum* actrange = m->actuator_actrange + 2 * actuator_id;
+    act = mju_clip(act, actrange[0], actrange[1]);
+  }
+
+  return act;
+}
+
+bool HasSlew(const mjModel* m, int instance) {
+  return ReadOptionalDoubleAttr(m, instance, kAttrSlewMax).has_value();
+}
+
+}  // namespace
+
+PidConfig PidConfig::FromModel(const mjModel* m, int instance) {
+  PidConfig config;
+  config.p_gain = ReadOptionalDoubleAttr(m, instance, kAttrPGain).value_or(0);
+  config.i_gain = ReadOptionalDoubleAttr(m, instance, kAttrIGain).value_or(0);
+  config.d_gain = ReadOptionalDoubleAttr(m, instance, kAttrDGain).value_or(0);
+
+  // Clamps in the XML are specified in terms of maximum forces. Scale by i_gain
+  // to get the limits on the value of the error integral.
+  std::optional<double> i_clamp_max_force =
+      ReadOptionalDoubleAttr(m, instance, kAttrIMax);
+  if (i_clamp_max_force.has_value() && config.i_gain) {
+    config.i_max = *i_clamp_max_force / config.i_gain;
+  }
+
+  config.slew_max = ReadOptionalDoubleAttr(m, instance, kAttrSlewMax);
+
+  return config;
+}
+
+std::unique_ptr<Pid> Pid::Create(const mjModel* m, int instance) {
+  PidConfig config = PidConfig::FromModel(m, instance);
+
+  if (config.i_max.has_value() && *config.i_max < 0) {
+    mju_warning("negative imax");
+    return nullptr;
+  }
+
+  if (config.slew_max.value_or(0.0) < 0) {
+    mju_warning("maxslew must be non-negative");
+    return nullptr;
+  }
+
+  std::vector<int> actuators;
+  for (int i = 0; i < m->nu; i++) {
+    if (m->actuator_plugin[i] == instance) {
+      actuators.push_back(i);
+    }
+  }
+  if (actuators.empty()) {
+    mju_warning("actuator not found for plugin instance %d", instance);
+    return nullptr;
+  }
+  return std::unique_ptr<Pid>(new Pid(config, std::move(actuators)));
+}
+
+void Pid::Reset(mjtNum* plugin_state) {}
+
+mjtNum Pid::GetCtrl(const mjModel* m, const mjData* d, int actuator_idx,
+                    const State& state,
+                    bool actearly) const {
+  mjtNum ctrl = 0;
+  if (m->actuator_dyntype[actuator_idx] == mjDYN_NONE) {
+    ctrl = d->ctrl[actuator_idx];
+    // clamp ctrl
+    if (m->actuator_ctrllimited[actuator_idx]) {
+      ctrl = mju_clip(ctrl, m->actuator_ctrlrange[2 * actuator_idx],
+                      m->actuator_ctrlrange[2 * actuator_idx + 1]);
+    }
+  } else {
+    // Use of act instead of ctrl, to create integrated-velocity controllers or
+    // to filter the controls.
+    int actadr = m->actuator_actadr[actuator_idx] +
+                 m->actuator_actnum[actuator_idx] - 1;
+    if (actearly) {
+      ctrl = NextActivation(m, d, actuator_idx, actadr, d->act_dot[actadr]);
+    } else {
+      ctrl = d->act[actadr];
+    }
+  }
+  if (config_.slew_max.has_value() && state.previous_ctrl_exists) {
+    mjtNum ctrl_min = state.previous_ctrl - *config_.slew_max * m->opt.timestep;
+    mjtNum ctrl_max = state.previous_ctrl + *config_.slew_max * m->opt.timestep;
+    ctrl = mju_clip(ctrl, ctrl_min, ctrl_max);
+  }
+  return ctrl;
+}
+
+void Pid::ActDot(const mjModel* m, mjData* d, int instance) const {
+  for (int actuator_idx : actuators_) {
+    State state = GetState(m, d, actuator_idx);
+    mjtNum ctrl = GetCtrl(m, d, actuator_idx, state, /*actearly=*/false);
+    mjtNum error = ctrl - d->actuator_length[actuator_idx];
+
+    int state_idx = m->actuator_actadr[actuator_idx];
+    if (config_.i_gain) {
+      mjtNum integral = state.integral + error * m->opt.timestep;
+      if (config_.i_max.has_value()) {
+        integral = mju_clip(integral, -*config_.i_max, *config_.i_max);
+      }
+      d->act_dot[state_idx] = (integral - d->act[state_idx]) / m->opt.timestep;
+      ++state_idx;
+    }
+    if (config_.slew_max.has_value()) {
+      d->act_dot[state_idx] = (ctrl - d->act[state_idx]) / m->opt.timestep;
+      ++state_idx;
+    }
+  }
+}
+
+void Pid::Compute(const mjModel* m, mjData* d, int instance) {
+  for (int i = 0; i < actuators_.size(); i++) {
+    int actuator_idx = actuators_[i];
+    State state = GetState(m, d, actuator_idx);
+    mjtNum ctrl =
+        GetCtrl(m, d, actuator_idx, state, m->actuator_actearly[actuator_idx]);
+
+    mjtNum error = ctrl - d->actuator_length[actuator_idx];
+
+    mjtNum ctrl_dot = m->actuator_dyntype[actuator_idx] == mjDYN_NONE
+                          ? 0
+                          : d->act_dot[m->actuator_actadr[actuator_idx] +
+                                       m->actuator_actnum[actuator_idx] - 1];
+    mjtNum error_dot = ctrl_dot - d->actuator_velocity[actuator_idx];
+
+    mjtNum integral = 0;
+    if (config_.i_gain) {
+      integral = state.integral + error * m->opt.timestep;
+      if (config_.i_max.has_value()) {
+        integral =
+            mju_clip(integral, -*config_.i_max, *config_.i_max);
+      }
+    }
+
+    d->actuator_force[actuator_idx] = config_.p_gain * error +
+                                      config_.d_gain * error_dot +
+                                      config_.i_gain * integral;
+  }
+}
+
+void Pid::Advance(const mjModel* m, mjData* d, int instance) const {
+  // act variables already updated by MuJoCo integrating act_dot
+}
+
+int Pid::StateSize(const mjModel* m, int instance) {
+  return 0;
+}
+
+int Pid::ActDim(const mjModel* m, int instance, int actuator_id) {
+  double i_gain = ReadOptionalDoubleAttr(m, instance, kAttrIGain).value_or(0);
+  return (i_gain ? 1 : 0) + (HasSlew(m, instance) ? 1 : 0);
+}
+
+Pid::State Pid::GetState(const mjModel* m, mjData* d, int actuator_idx) const {
+  State state;
+  int state_idx = m->actuator_actadr[actuator_idx];
+  if (config_.i_gain) {
+    state.integral = d->act[state_idx++];
+  }
+  if (config_.slew_max.has_value()) {
+    state.previous_ctrl = d->act[state_idx++];
+    state.previous_ctrl_exists = d->time > 0;
+  }
+  return state;
+}
+
+void Pid::RegisterPlugin() {
+  mjpPlugin plugin;
+  mjp_defaultPlugin(&plugin);
+  plugin.name = "mujoco.pid";
+  plugin.capabilityflags |= mjPLUGIN_ACTUATOR;
+
+  std::vector<const char*> attributes = {kAttrPGain, kAttrIGain, kAttrDGain,
+                                         kAttrIMax, kAttrSlewMax};
+  plugin.nattribute = attributes.size();
+  plugin.attributes = attributes.data();
+
+  plugin.actuator_actdim = Pid::ActDim;
+  plugin.nstate = Pid::StateSize;
+
+  plugin.init = +[](const mjModel* m, mjData* d, int instance) {
+    std::unique_ptr<Pid> pid = Pid::Create(m, instance);
+    if (pid == nullptr) {
+      return -1;
+    }
+    d->plugin_data[instance] = reinterpret_cast<uintptr_t>(pid.release());
+    return 0;
+  };
+  plugin.destroy = +[](mjData* d, int instance) {
+    delete reinterpret_cast<Pid*>(d->plugin_data[instance]);
+    d->plugin_data[instance] = 0;
+  };
+  plugin.reset = +[](const mjModel* m, double* plugin_state, void* plugin_data,
+                     int instance) {
+    auto* pid = reinterpret_cast<Pid*>(plugin_data);
+    pid->Reset(plugin_state);
+  };
+  plugin.actuator_act_dot = +[](const mjModel* m, mjData* d, int instance) {
+    auto* pid = reinterpret_cast<Pid*>(d->plugin_data[instance]);
+    pid->ActDot(m, d, instance);
+  };
+  plugin.compute =
+      +[](const mjModel* m, mjData* d, int instance, int capability_bit) {
+        auto* pid = reinterpret_cast<Pid*>(d->plugin_data[instance]);
+        pid->Compute(m, d, instance);
+      };
+  plugin.advance = +[](const mjModel* m, mjData* d, int instance) {
+    auto* pid = reinterpret_cast<Pid*>(d->plugin_data[instance]);
+    pid->Advance(m, d, instance);
+  };
+  // TODO: b/303823996 - allow actuator plugins to compute their derivatives wrt
+  // qvel, for implicit integration
+  mjp_registerPlugin(&plugin);
+}
+
+Pid::Pid(PidConfig config, std::vector<int> actuators)
+    : config_(std::move(config)), actuators_(std::move(actuators)) {}
+
+}  // namespace mujoco::plugin::actuator

plugin/actuator/pid.h

@@ -0,0 +1,102 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef MUJOCO_PLUGIN_ACTUATOR_PID_H_
+#define MUJOCO_PLUGIN_ACTUATOR_PID_H_
+
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include <mujoco/mjdata.h>
+#include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
+
+namespace mujoco::plugin::actuator {
+
+struct PidConfig {
+  double p_gain = 0.0;
+  double i_gain = 0.0;
+  double d_gain = 0.0;
+
+  // Maximum value of the error integral.
+  // NOTE: In the XML definition, the clamp values are specified as limits on
+  // the *force* value, so are scaled by i_gain.
+  std::optional<double> i_max;
+
+  // Maximum speed at which the setpoint can change.
+  std::optional<double> slew_max;
+
+  // Reads plugin attributes to construct PID configuration.
+  static PidConfig FromModel(const mjModel* m, int instance);
+};
+
+// An actuator plugin which implements configurable PID control.
+class Pid {
+ public:
+  // Returns an instance of Pid. The result can be null in case of
+  // misconfiguration.
+  static std::unique_ptr<Pid> Create(const mjModel* m, int instance);
+
+  // Returns the number of state variables for the plugin instance
+  static int StateSize(const mjModel* m, int instance);
+
+  // Returns the number of activation variables for the plugin instance
+  static int ActDim(const mjModel* m, int instance, int actuator_id);
+
+  // Resets the C++ Pid instance's state.
+  // plugin_state is a C array pointer into mjData->plugin_state, with a size
+  // equal to the value returned from StateSize.
+  void Reset(mjtNum* plugin_state);
+
+  // Computes the rate of change for activation variables
+  void ActDot(const mjModel* m, mjData* d, int instance) const;
+
+  // Idempotent computation which updates d->actuator_force and the internal
+  // state of the class. Called after ActDot.
+  void Compute(const mjModel* m, mjData* d, int instance);
+
+  // Updates plugin state.
+  void Advance(const mjModel* m, mjData* d, int instance) const;
+
+  // Adds the PID plugin to the global registry of MuJoCo plugins.
+  static void RegisterPlugin();
+
+ private:
+  Pid(PidConfig config, std::vector<int> actuators);
+
+  struct State {
+    mjtNum previous_ctrl = 0;
+    // if using slew rate limits, mjData.act will contain an activation variable
+    // with the last ctrl value. If `false`, that value should be ignored,
+    // because it hasn't been set yet.
+    bool previous_ctrl_exists = false;
+    mjtNum integral = 0;
+  };
+  // Reads data from d->act and returns it as a State struct.
+  State GetState(const mjModel* m, mjData* d, int actuator_idx) const;
+
+  // Returns the PID setpoint, which is normally d->ctrl, but can be d->act for
+  // actuators with dyntype != none.
+  mjtNum GetCtrl(const mjModel* m, const mjData* d, int actuator_idx,
+                 const State& state, bool actearly) const;
+
+  PidConfig config_;
+  // set of actuator IDs controlled by this plugin instance.
+  std::vector<int> actuators_;
+};
+
+}  // namespace mujoco::plugin::actuator
+
+#endif  // MUJOCO_PLUGIN_ACTUATOR_PID_H_

plugin/actuator/register.cc

@@ -0,0 +1,22 @@
+// Copyright 2023 DeepMind Technologies Limited
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <mujoco/mjplugin.h>
+#include "pid.h"
+
+namespace mujoco::plugin::actuator {
+
+mjPLUGIN_LIB_INIT { Pid::RegisterPlugin(); }
+
+}  // namespace mujoco::plugin::actuator

plugin/elasticity/README.md

@@ -16,7 +16,7 @@ Parameters:
 
    - `twist` [Pa]: twisting stiffness.
    - `bend` [Pa]: bending stiffness.
-   - `flat` [bool]: if true, the stress-equilibrium configuration is that of a straight cable; if true, it is the configuration defined in the XML.
+   - `flat` [bool]: if true, the stress-equilibrium configuration is that of a straight cable; if false or unspecified, it is the configuration defined in the XML.
    - `vmax` [N/m^2]: If greater than zero, the cable is colored using mechanical stresses; the value represent the maximum stress in the color scale.
 
 ### Shell

plugin/elasticity/elasticity.h

@@ -90,6 +90,8 @@ inline void ComputeForce(mjtNum* qfrc_passive,
                          const std::vector<T>& elements,
                          const std::vector<mjtNum>& metric,
                          const std::vector<mjtNum>& elongationglob,
+                         const mjModel* m,
+                         const int* vertbodyid,
                          const mjtNum* xpos) {
   for (int t = 0; t < elements.size(); t++)  {
     const int* v = elements[t].vertices;
@@ -126,8 +128,17 @@ inline void ComputeForce(mjtNum* qfrc_passive,
 
     // insert into global force
     for (int i = 0; i < T::kNumVerts; i++) {
-      for (int x = 0; x < 3; x++) {
+      int body_dofnum = 3;
-        qfrc_passive[3*v[i]+x] -= force[3*i+x];
+      int body_dofadr = 3*v[i];
+      if (vertbodyid) {
+        body_dofnum = m->body_dofnum[vertbodyid[v[i]]];
+        body_dofadr = m->body_dofadr[vertbodyid[v[i]]];
+        if (body_dofnum && m->body_simple[vertbodyid[v[i]]] != 2) {
+          mju_error("Non-simple or non-static bodies are not yet supported");
+        }
+      }
+      for (int x = 0; x < body_dofnum; x++) {
+        qfrc_passive[body_dofadr+x] -= force[3*i+x];
       }
     }
   }

plugin/elasticity/membrane.cc

@@ -114,12 +114,17 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
 
   // count flexes
   for (int i = 0; i < m->nflex; i++) {
-    if (m->flex_vertbodyid[m->flex_vertadr[i]] == i0) {
+    for (int j = 0; j < m->flex_vertnum[i]; j++) {
-      f0 = i;
+      if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) {
-      break;
+        f0 = i;
+      }
     }
   }
 
+  // vertex positions
+  mjtNum* body_pos =
+      f0 < 0 ? m->body_pos + 3*i0 : m->flex_xvert0 + 3*m->flex_vertadr[f0];
+
   // generate triangles from the vertices
   nt = CreateStencils<Stencil2D>(elements, edges, simplex, edgeidx);
 
@@ -130,13 +135,14 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   for (int t = 0; t < nt; t++) {
     int* v = elements[t].vertices;
     for (int i = 0; i < kNumVerts; i++) {
-      if (m->body_plugin[i0+v[i]] != instance) {
+      int bi = f0 < 0 ? i0+v[i] : m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
-        mju_error("This body does not have the requested plugin instance");
+      if (bi && m->body_plugin[bi] != instance) {
+        mju_error("Body %d does not have plugin instance %d", bi, instance);
       }
     }
 
     // triangles area
-    mjtNum volume = ComputeVolume(m->body_pos+3*i0, v);
+    mjtNum volume = ComputeVolume(body_pos, v);
 
     // material parameters
     mjtNum mu = E / (2*(1+nu)) * mju_abs(volume) / 4 * thickness;
@@ -147,7 +153,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
 
     // compute edge basis
     for (int e = 0; e < kNumEdges; e++) {
-      ComputeBasis(basis[e], m->body_pos+3*i0, v,
+      ComputeBasis(basis[e], body_pos, v,
                    Stencil2D::edge[Stencil2D::edge[e][0]],
                    Stencil2D::edge[Stencil2D::edge[e][1]], volume);
     }
@@ -164,7 +170,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   elongation.assign(ne, 0);
 
   // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
-  UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
+  UpdateSquaredLengths(reference, edges, body_pos);
 
   // save previous lengths
   previous = reference;
@@ -191,8 +197,12 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
   }
 
   // compute gradient of elastic energy and insert into passive force
-  ComputeForce<Stencil2D>(d->qfrc_passive + m->body_dofadr[i0], elements,
+  int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0];
-                          metric, elongation, d->xpos + 3 * i0);
+  int* bodyid = f0 < 0 ? nullptr : m->flex_vertbodyid + flex_vertadr;
+  mjtNum* xpos = f0 < 0 ? d->xpos + 3*i0 : d->flexvert_xpos + 3*flex_vertadr;
+  mjtNum* qfrc = d->qfrc_passive + (f0 < 0 ? m->body_dofadr[i0] : 0);
+
+  ComputeForce<Stencil2D>(qfrc, elements, metric, elongation, m, bodyid, xpos);
 
   // update stored lengths
   if (kD > 0) {

plugin/elasticity/solid.cc

@@ -119,12 +119,17 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
 
   // count flexes
   for (int i = 0; i < m->nflex; i++) {
-    if (m->flex_vertbodyid[m->flex_vertadr[i]] == i0) {
+    for (int j = 0; j < m->flex_vertnum[i]; j++) {
-      f0 = i;
+      if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) {
-      break;
+        f0 = i;
+      }
     }
   }
 
+  // vertex positions
+  mjtNum* body_pos =
+      f0 < 0 ? m->body_pos + 3*i0 : m->flex_xvert0 + 3*m->flex_vertadr[f0];
+
   // generate tetrahedra from the vertices
   nt = CreateStencils<Stencil3D>(elements, edges, simplex, edgeidx);
 
@@ -135,20 +140,21 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   for (int t = 0; t < nt; t++) {
     int* v = elements[t].vertices;
     for (int i = 0; i < kNumVerts; i++) {
-      if (m->body_plugin[i0+v[i]] != instance) {
+      int bi = f0 < 0 ? i0+v[i] : m->flex_vertbodyid[m->flex_vertadr[f0]+v[i]];
-        mju_error("This body does not have the requested plugin instance");
+      if (bi && m->body_plugin[bi] != instance) {
+        mju_error("Body %d does not have plugin instance %d", bi, instance);
       }
     }
 
     // tetrahedron volume
-    mjtNum volume = ComputeVolume(m->body_pos+3*i0, v);
+    mjtNum volume = ComputeVolume(body_pos, v);
 
     // local geometric quantities
     mjtNum basis[kNumEdges][9] = {{0}, {0}, {0}, {0}, {0}, {0}};
 
     // compute edge basis
     for (int e = 0; e < kNumEdges; e++) {
-      ComputeBasis(basis[e], m->body_pos+3*i0, v,
+      ComputeBasis(basis[e], body_pos, v,
                    face[e2f[e][0]], face[e2f[e][1]], volume);
     }
 
@@ -168,7 +174,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   elongation.assign(ne, 0);
 
   // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
-  UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
+  UpdateSquaredLengths(reference, edges, body_pos);
 
   // save previous lengths
   previous = reference;
@@ -195,8 +201,12 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
   }
 
   // compute gradient of elastic energy and insert into passive force
-  ComputeForce<Stencil3D>(d->qfrc_passive + m->body_dofadr[i0], elements,
+  int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0];
-                          metric, elongation, d->xpos + 3 * i0);
+  int* bodyid = f0 < 0 ? nullptr : m->flex_vertbodyid + flex_vertadr;
+  mjtNum* xpos = f0 < 0 ? d->xpos + 3*i0 : d->flexvert_xpos + 3*flex_vertadr;
+  mjtNum* qfrc = d->qfrc_passive + (f0 < 0 ? m->body_dofadr[i0] : 0);
+
+  ComputeForce<Stencil3D>(qfrc, elements, metric, elongation, m, bodyid, xpos);
 
   // update stored lengths
   if (kD > 0) {

plugin/sdf/README.md

@@ -95,6 +95,6 @@ class MySDF {
   mjtNum attribute[MySDFAttribute::nattribute];
 
  private:
-  Torus(const mjModel* m, mjData* d, int instance);
+  MySDF(const mjModel* m, mjData* d, int instance);
 };
 ```

plugin/sdf/bolt.cc

@@ -166,12 +166,12 @@ void Bolt::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Bolt*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Bolt*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/bowl.cc

@@ -164,12 +164,12 @@ void Bowl::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Bowl*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Bowl*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/gear.cc

@@ -250,12 +250,12 @@ void Gear::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Gear*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Gear*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/nut.cc

@@ -166,12 +166,12 @@ void Nut::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<Nut*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<Nut*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
   plugin.sdf_staticdistance =

plugin/sdf/sdflib.cc

@@ -31,6 +31,34 @@ inline unsigned int* MakeNonConstUnsigned(const int* ptr) {
   return reinterpret_cast<unsigned int*>(const_cast<int*>(ptr));
 }
 
+mjtNum boxProjection(glm::vec3& point, const sdflib::BoundingBox& box) {
+  glm::vec3 r = point - box.getCenter();
+  glm::vec3 q = glm::abs(r) - 0.5f * box.getSize();
+  mjtNum dist_sqr = 0;
+  mjtNum eps = 1e-6;
+
+  // skip the projection if inside
+  if (q.x <= 0 && q.y <= 0 && q.z <= 0) {
+    return glm::max(q.x, glm::max(q.y, q.z));
+  }
+
+  // in-place projection inside the box if outside
+  if ( q.x >= 0 ) {
+    dist_sqr += q.x * q.x;
+    point.x -= r.x > 0 ? (q.x+eps) : -(q.x+eps);
+  }
+  if ( q.y >= 0 ) {
+    dist_sqr += q.y * q.y;
+    point.y -= r.y > 0 ? (q.y+eps) : -(q.y+eps);
+  }
+  if ( q.z >= 0 ) {
+    dist_sqr += q.z * q.z;
+    point.z -= r.z > 0 ? (q.z+eps) : -(q.z+eps);
+  }
+
+  return mju_sqrt(dist_sqr);
+}
+
 }  // namespace
 
 // factory function
@@ -65,8 +93,12 @@ std::optional<SdfLib> SdfLib::Create(const mjModel* m, mjData* d,
 
 // plugin constructor
 SdfLib::SdfLib(sdflib::Mesh&& mesh) {
+  sdflib::BoundingBox box = mesh.getBoundingBox();
+  const glm::vec3 modelBBsize = box.getSize();
+  box.addMargin(
+      0.1f * glm::max(glm::max(modelBBsize.x, modelBBsize.y), modelBBsize.z));
   sdf_func_ =
-      sdflib::OctreeSdf(mesh, mesh.getBoundingBox(), 8, 3, 1e-3,
+      sdflib::OctreeSdf(mesh, box, 8, 3, 1e-3,
                         sdflib::OctreeSdf::InitAlgorithm::CONTINUITY, 1);
 }
 
@@ -89,17 +121,37 @@ void SdfLib::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
 // sdf
 mjtNum SdfLib::Distance(const mjtNum p[3]) const {
   glm::vec3 point(p[0], p[1], p[2]);
-  return sdf_func_.getDistance(point);
+  mjtNum boxDist = boxProjection(point, sdf_func_.getGridBoundingBox());
+  return sdf_func_.getDistance(point) + (boxDist <= 0 ? 0 : boxDist);
 }
 
 // gradient of sdf
 void SdfLib::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
   glm::vec3 gradient;
   glm::vec3 p(point[0], point[1], point[2]);
-  sdf_func_.getDistance(p, gradient);
+
-  grad[0] = gradient[0];
+  // analytic in the interior
-  grad[1] = gradient[1];
+  if (boxProjection(p, sdf_func_.getGridBoundingBox()) <= 0) {
-  grad[2] = gradient[2];
+    sdf_func_.getDistance(p, gradient);
+    grad[0] = gradient[0];
+    grad[1] = gradient[1];
+    grad[2] = gradient[2];
+    return;
+  }
+
+  // finite difference in the exterior
+  mjtNum eps = 1e-8;
+  mjtNum dist0 = Distance(point);
+  mjtNum pointX[3] = {point[0]+eps, point[1], point[2]};
+  mjtNum distX = Distance(pointX);
+  mjtNum pointY[3] = {point[0], point[1]+eps, point[2]};
+  mjtNum distY = Distance(pointY);
+  mjtNum pointZ[3] = {point[0], point[1], point[2]+eps};
+  mjtNum distZ = Distance(pointZ);
+
+  grad[0] = (distX - dist0) / eps;
+  grad[1] = (distY - dist0) / eps;
+  grad[2] = (distZ - dist0) / eps;
 }
 
 // plugin registration
@@ -146,12 +198,12 @@ void SdfLib::RegisterPlugin() {
   plugin.sdf_distance =
       +[](const mjtNum point[3], const mjData* d, int instance) {
         auto* sdf = reinterpret_cast<SdfLib*>(d->plugin_data[instance]);
-        sdf->visualizer_.AddPoint(point);
         return sdf->Distance(point);
       };
   plugin.sdf_gradient = +[](mjtNum gradient[3], const mjtNum point[3],
                         const mjData* d, int instance) {
     auto* sdf = reinterpret_cast<SdfLib*>(d->plugin_data[instance]);
+    sdf->visualizer_.AddPoint(point);
     sdf->Gradient(gradient, point);
   };
 

python/mujoco/CMakeLists.txt

@@ -84,7 +84,7 @@ if(NOT TARGET mujoco)
     if(MUJOCO_FRAMEWORK)
       message("MuJoCo framework is at ${MUJOCO_FRAMEWORK}/mujoco.framework")
       set(MUJOCO_LIBRARY
-          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.0.2.dylib
+          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.1.2.dylib
       )
       target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
     endif()
@@ -92,7 +92,7 @@ if(NOT TARGET mujoco)
 
   if(NOT MUJOCO_FRAMEWORK)
     find_library(
-      MUJOCO_LIBRARY mujoco mujoco.3.0.2 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
+      MUJOCO_LIBRARY mujoco mujoco.3.1.2 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
     )
     find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
     message("MuJoCo is at ${MUJOCO_LIBRARY}")
@@ -173,16 +173,17 @@ findorfetch(
   GIT_REPO
   https://gitlab.com/libeigen/eigen
   GIT_TAG
-  aa6964bf3a34fd607837dd8123bc42465185c4f8
+  454f89af9d6f3525b1df5f9ef9c86df58bf2d4d3
   TARGETS
   Eigen3::Eigen
   EXCLUDE_FROM_ALL
 )
 
 # ==================== PYBIND11 ================================================
+option(MUJOCO_PYTHON_USE_SYSTEM_PYBIND11 "Use installed pybind11 version." OFF)
 findorfetch(
   USE_SYSTEM_PACKAGE
-  OFF
+  MUJOCO_PYTHON_USE_SYSTEM_PYBIND11
   PACKAGE_NAME
   pybind11
   LIBRARY_NAME

python/mujoco/bindings_test.py

@@ -822,7 +822,7 @@ Euler integrator, semi-implicit in velocity.
     self.assertLen(mujoco.mjRNDSTRING, mujoco.mjtRndFlag.mjNRNDFLAG)
     self.assertEqual(mujoco.mjDISABLESTRING[11], 'Refsafe')
     self.assertEqual(mujoco.mjVISSTRING[mujoco.mjtVisFlag.mjVIS_INERTIA],
-                     ('&Inertia', '0', 'I'))
+                     ('Inertia', '0', 'I'))
 
   def test_enum_values(self):
     self.assertEqual(mujoco.mjtJoint.mjJNT_FREE, 0)
@@ -840,7 +840,7 @@ Euler integrator, semi-implicit in velocity.
     self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW, 100)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW1, 101)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_ARROW2, 102)
-    self.assertEqual(mujoco.mjtGeom.mjGEOM_TRIANGLE, 107)
+    self.assertEqual(mujoco.mjtGeom.mjGEOM_TRIANGLE, 108)
     self.assertEqual(mujoco.mjtGeom.mjGEOM_NONE, 1001)
 
   def test_enum_from_int(self):

python/mujoco/cgl/__init__.py

@@ -57,15 +57,18 @@ class GLContext:
 
   def free(self):
     """Frees resources associated with this context."""
-    if self._context:
+    try:
-      cgl.CGLUnlockContext(self._context)
+      if self._context:
-      cgl.CGLSetCurrentContext(None)
+        cgl.CGLUnlockContext(self._context)
-      cgl.CGLReleaseContext(self._context)
+        cgl.CGLSetCurrentContext(None)
-    self._context = None
+        cgl.CGLReleaseContext(self._context)
+      self._context = None
 
-    if self._pix:
+      if self._pix:
-      cgl.CGLReleasePixelFormat(self._pix)
+        cgl.CGLReleasePixelFormat(self._pix)
-    self._context = None
+      self._pix = None
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
 
   def __del__(self):
     self.free()