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

CMakeLists.txt

@@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco
-  VERSION 3.1.2
+  VERSION 3.1.3
   DESCRIPTION "MuJoCo Physics Simulator"
   HOMEPAGE_URL "https://mujoco.org"
 )

cmake/MujocoDependencies.cmake

@@ -39,12 +39,12 @@ set(MUJOCO_DEP_VERSION_qhull
     CACHE STRING "Version of `qhull` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_Eigen3
-    454f89af9d6f3525b1df5f9ef9c86df58bf2d4d3
+    7fd7a3f946e5ac152d28dad388cff8bfa1026925
     CACHE STRING "Version of `Eigen3` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_abseil
-    fb3621f4f897824c0dbe0615fa94543df6192f30 # LTS 20230802.1
+    4a2c63365eff8823a5221db86ef490e828306f9d # LTS 20240116.0
     CACHE STRING "Version of `abseil` to be fetched."
 )
 

dist/mujoco.rc

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

dist/simulate.rc

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

doc/APIreference/APIglobals.rst

@@ -489,7 +489,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - 1000
      - Maximum number of textures allowed.
        Defined in `mjrender.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjrender.h>`_.
-   * - ``mjMAXTHREADS``
+   * - ``mjMAXTHREAD``
      - 128
      - Maximum number OS threads that can be used in a thread pool.
        Defined in `mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`_.
@@ -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``
-     - 312
+     - 313
      - 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/XMLreference.rst

@@ -311,15 +311,16 @@ any effect. The settings here are global and apply to the entire model.
 
    - All materials are discarded.
    - All textures are discarded.
-   - All geoms with :ref:`contype<body-geom-contype>`=:ref:`conaffinity<body-geom-conaffinity>`=0 are discarded, if they
+   - All geoms with :ref:`contype<body-geom-contype>` |-| = |-| :ref:`conaffinity<body-geom-conaffinity>` |-| =0 are
-     are not referenced in another MJCF element. If a discarded geom was used for inferring body inertia, an explicit
+     discarded, if they are not referenced in another MJCF element. If a discarded geom was used for inferring body
-     :ref:`inertial<body-inertial>` element is added to the body.
+     inertia, an explicit :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
+   The resulting compiled model will have exactly the same dynamics as the original model. The only engine-level
-   raycasting, as used for example by :ref:`rangefinder<sensor-rangefinder>`, since raycasting reports distances to
+   computation which might change is the output of :ref:`raycasting<mj_ray>` computations, as used for example by
-   visual geoms. When visualizing models compiled with this flag, it is important to remember that colliding geoms are
+   :ref:`rangefinder<sensor-rangefinder>` sensors, since raycasting reports distances to visual geoms. When visualizing
-   often placed in a :ref:`group<body-geom-group>` which is invisible by default.
+   models compiled with this flag, it is important to remember that collision geoms are often placed in a
+   :ref:`group<body-geom-group>` which is invisible by default.
 
 .. _compiler-convexhull:
 
@@ -670,14 +671,18 @@ is effectively a miscellaneous subsection.
 :at:`azimuth`: :at-val:`real, "90"`
    This attribute specifies the initial azimuth of the free camera around the vertical z-axis, in degrees. A value of 0
    corresponds to looking in the positive x direction, while the default value of 90 corresponds to looking in the
-   positive y direction.
+   positive y direction. The look-at point itself is specified by the :ref:`statistic/center<statistic-center>`
+   attribute, while the distance from the look-at point is controlled by the :ref:`statistic/extent<statistic-extent>`
+   attribute.
 
 .. _visual-global-elevation:
 
 :at:`elevation`: :at-val:`real, "-45"`
    This attribute specifies the initial elevation of the free camera with respect to the lookat point. Note that since
    this is a rotation around a vector parallel to the camera's X-axis (right in pixel space), *negative* numbers
-   correspond to moving the camera *up* from the horizontal plane, and vice-versa.
+   correspond to moving the camera *up* from the horizontal plane, and vice-versa. The look-at point itself is specified
+   by the :ref:`statistic/center<statistic-center>` attribute, while the distance from the look-at point is controlled
+   by the :ref:`statistic/extent<statistic-extent>` attribute.
 
 .. _visual-global-linewidth:
 
@@ -714,6 +719,12 @@ is effectively a miscellaneous subsection.
    This attribute specifies how the equivalent inertia is visualized. "false":
    use box, "true": use ellipsoid.
 
+.. _visual-global-bvactive:
+
+:at:`bvactive`: :at-val:`[false, true], "true"`
+   This attribute specifies whether collision and raycasting code should mark elements of Bounding Volume Hierarchies
+   as intersecting, for the purpose of visualization. Setting this attribute to "false" can speed up simulation for
+   models with high-resolution meshes.
 
 .. _visual-quality:
 
@@ -1121,6 +1132,17 @@ disables the rendering of the corresponding object.
 :at:`frustum`: :at-val:`real(4), "1 1 0 0.2"`
    Color used to render the camera frustum.
 
+.. _visual-rgba-bv:
+
+:at:`bv`: :at-val:`real(4), "0 1 0 0.5"`
+   Color used to render bounding volumes.
+
+.. _visual-rgba-bvactive:
+
+:at:`bvactive`: :at-val:`real(4), "1 0 0 0.5"`
+   Color used to render active bounding volumes, if the :ref:`bvactive<visual-global-bvactive>` flag is "true".
+
+
 
 .. _asset:
 
@@ -1135,6 +1157,339 @@ extension if no ``content_type`` attribute is specified. The content type is ign
 file.
 
 
+.. _asset-mesh:
+
+:el-prefix:`asset/` |-| **mesh** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This element creates a mesh asset, which can then be referenced from geoms. If the referencing geom type is
+:at-val:`mesh` the mesh is instantiated in the model, otherwise a geometric primitive is automatically fitted to it; see
+the :ref:`geom <body-geom>` element below.
+
+MuJoCo works with triangulated meshes. They can be loaded from binary STL files, OBJ files or MSH files with custom
+format described below, or vertex and face data specified directly in the XML. Software such as MeshLab can be used to
+convert from other mesh formats to STL or OBJ. While any collection of triangles can be loaded as a mesh and rendered,
+collision detection works with the convex hull of the mesh as explained in :ref:`Collision`. See also the convexhull
+attribute of the :ref:`compiler <compiler>` element which controls the automatic generation of convex hulls. The mesh
+appearance (including texture mapping) is controlled by the :at:`material` and :at:`rgba` attributes of the referencing
+geom, similarly to height fields.
+
+Meshes can have explicit texture coordinates instead of relying on the automated texture
+mapping mechanism. When provided, these explicit coordinates have priority. Note that texture coordinates can be
+specified with OBJ files and MSH files, as well as explicitly in the XML with the :at:`texcoord` attribute, but not via
+STL files. These mechanism cannot be mixed. So if you have an STL mesh, the only way to add texture coordinates to it is
+to convert to one of the other supported formats.
+
+MSH file format
+   The binary MSH file starts with 4 integers specifying the number of vertex positions (nvertex), vertex normals
+   (nnormal), vertex texture coordinates (ntexcoord), and vertex indices making up the faces (nface), followed by the
+   numeric data. nvertex must be at least 4. nnormal and ntexcoord can be zero (in which case the corresponding data is
+   not defined) or equal to nvertex. nface can also be zero, in which case faces are constructed automatically from the
+   convex hull of the vertex positions. The file size in bytes must be exactly: 16 + 12*(nvertex + nnormal + nface) +
+   8*ntexcoord. The contents of the file must be as follows:
+
+   .. code:: Text
+
+          (int32)   nvertex
+          (int32)   nnormal
+          (int32)   ntexcoord
+          (int32)   nface
+          (float)   vertex_positions[3*nvertex]
+          (float)   vertex_normals[3*nnormal]
+          (float)   vertex_texcoords[2*ntexcoord]
+          (int32)   face_vertex_indices[3*nface]
+
+Poorly designed meshes can display rendering artifacts. In particular, the shadow mapping mechanism relies on having
+some distance between front and back-facing triangle faces. If the faces are repeated, with opposite normals as
+determined by the vertex order in each triangle, this causes shadow aliasing. The solution is to remove the repeated
+faces (which can be done in MeshLab) or use a better designed mesh. Flipped faces are checked by MuJoCo for meshes
+specified as OBJ or XML and an error message is returned.
+
+The size of the mesh is determined by the 3D coordinates of the vertex data in the mesh file, multiplied by the
+components of the :at:`scale` attribute below. Scaling is applied separately for each coordinate axis. Note that
+negative scaling values can be used to flip the mesh; this is a legitimate operation. The size parameters of the
+referening geoms are ignored, similarly to height fields. We also provide a mechanism to translate and
+rotate the 3D coordinates, using the attributes :ref:`refpos<asset-mesh-refpos>` and :ref:`refquat<asset-mesh-refquat>`.
+
+A mesh can also be defined without faces (a point cloud essentially). In that case
+the convex hull is constructed automatically, even if the compiler attribute convexhull is false. This makes it easy to
+construct simple shapes directly in the XML. For example, a pyramid can be created as:
+
+.. code-block:: xml
+
+   <asset>
+       <mesh name="tetrahedron" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
+   </asset>
+
+Positioning and orienting is complicated by the fact that vertex data are often designed relative to coordinate frames
+whose origin is not inside the mesh. In contrast, MuJoCo expects the origin of a geom's local frame to coincide with the
+geometric center of the shape. We resolve this discrepancy by pre-processing the mesh in the compiler, so that it is
+centered around (0,0,0) and its principal axes of inertia are the coordinate axes. We also save the translation and
+rotation offsets needed to achieve such alignment in :ref:`mjModel.mesh_pos<mjModel>` and
+:ref:`mjModel.mesh_quat<mjModel>`. These offsets are then applied to the referencing geom's position and orientation; see
+also :at:`mesh` attribute of :ref:`geom <body-geom>` below. Fortunately most meshes used in robot models are designed in
+a coordinate frame centered at the joint. This makes the corresponding MJCF model intuitive: we set the body frame at the
+joint, so that the joint position is (0,0,0) in the body frame, and simply reference the mesh. Below is an MJCF model
+fragment of a forearm, containing all the information needed to put the mesh where one would expect it to be. The body
+position is specified relative to the parent body, namely the upper arm (not shown). It is offset by 35 cm which is the
+typical length of the human upper arm. If the mesh vertex data were not designed in the above convention, we would have
+to use the geom position and orientation (or the new refpos, refquat mechanism) to compensate, but in practice this is
+rarely needed.
+
+.. code-block:: xml
+
+   <asset>
+       <mesh file="forearm.stl"/>
+   </asset>
+
+   <body pos="0 0 0.35"/>
+       <joint type="hinge" axis="1 0 0"/>
+       <geom type="mesh" mesh="forearm"/>
+   </body>
+
+The inertial computation mentioned above is part of an algorithm used not only to center and align the mesh, but also to
+infer the mass and inertia of the body to which it is attached. This is done by computing the centroid of the triangle
+faces, connecting each face with the centroid to form a triangular pyramid, computing the mass and signed inertia of all
+pyramids (considered solid or hollow if :at:`shellinertia` is true) and accumulating them. The sign ensures that
+pyramids on the outside of the surfaces are subtracted, as it can occur with concave geometries. This algorithm can be
+found in section 1.3.8 of Computational Geometry in C (Second Edition) by Joseph O'Rourke.
+
+The full list of processing steps applied by the compiler to each mesh is as follows:
+
+#. For STL meshes, remove any repeated vertices and re-index the faces if needed. If the mesh is not STL, we assume that
+   the desired vertices and faces have already been generated and do not apply removal or re-indexing;
+#. If vertex normals are not provided, generate normals automatically, using a weighted average of the surrounding face
+   normals. If sharp edges are encountered, the renderer uses the face normals to preserve the visual information about
+   the edge, unless smoothnormal is true. Note that normals cannot be provided with STL meshes;
+#. Scale, translate and rotate the vertices and normals, re-normalize the normals in case of scaling;
+#. Construct the convex hull if specified;
+#. Find the centroid of all triangle faces, and construct the union-of-pyramids representation. Triangles whose area is
+   too small (below the :ref:`mjMINVAL <glNumeric>` value of 1E-14) result in compile error;
+#. Compute the center of mass and inertia matrix of the union-of-pyramids. Use eigenvalue decomposition to find the
+   principal axes of inertia. Center and align the mesh, saving the translational and rotational offsets for subsequent
+   geom-related computations.
+
+.. _asset-mesh-name:
+
+:at:`name`: :at-val:`string, optional`
+   Name of the mesh, used for referencing. If omitted, the mesh name equals the file name without the path and
+   extension.
+
+.. _asset-mesh-class:
+
+:at:`class`: :at-val:`string, optional`
+   Defaults class for setting unspecified attributes (only scale in this case).
+
+.. _asset-mesh-content_type:
+
+:at:`content_type`: :at-val:`string, optional`
+   If the file attribute is specified, then this sets the
+   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`_ (formerly known as MIME type) of the
+   file to be loaded. Any filename extensions will be overloaded.  Currently ``model/vnd.mujoco.msh``, ``model/obj``,
+   and ``model/stl`` are supported.
+
+.. _asset-mesh-file:
+
+:at:`file`: :at-val:`string, optional`
+   The file from which the mesh will be loaded. The path is determined as described in the meshdir attribute of
+   :ref:`compiler <compiler>`. The file extension must be "stl", "msh", or "obj" (not case sensitive) specifying the
+   file type.  If the file name is omitted, the vertex attribute becomes required.
+
+.. _asset-mesh-scale:
+
+:at:`scale`: :at-val:`real(3), "1 1 1"`
+   This attribute specifies the scaling that will be applied to the vertex data along each coordinate axis. Negative
+   values are allowed, resulting in flipping the mesh along the corresponding axis.
+
+.. _asset-mesh-smoothnormal:
+
+:at:`smoothnormal`: :at-val:`[false, true], "false"`
+   Controls the automatic generation of vertex normals when normals are not given explicitly. If true, smooth normals
+   are generated by averaging the face normals at each vertex, with weight proportional to the face area. If false,
+   faces at large angles relative to the average normal are excluded from the average. In this way, sharp edges (as in
+   cube edges) are not smoothed.
+
+.. _asset-mesh-vertex:
+
+:at:`vertex`: :at-val:`real(3*nvert), optional`
+   Vertex 3D position data. You can specify position data in the XML using this attribute, or using a binary file, but
+   not both.
+
+.. _asset-mesh-normal:
+
+:at:`normal`: :at-val:`real(3*nvert), optional`
+   Vertex 3D normal data. If specified, the number of normals must equal the number of vertices. The model compiler
+   normalizes the normals automatically.
+
+.. _asset-mesh-texcoord:
+
+:at:`texcoord`: :at-val:`real(2*nvert), optional`
+   Vertex 2D texture coordinates, which are numbers between 0 and 1. If specified, the number of texture coordinate
+   pairs must equal the number of vertices.
+
+.. _asset-mesh-face:
+
+:at:`face`: :at-val:`int(3*nface), optional`
+   Faces of the mesh. Each face is a sequence of 3 vertex indices, in counter-clockwise order. The indices must be
+   integers between 0 and nvert-1.
+
+.. _asset-mesh-refpos:
+
+:at:`refpos`: :at-val:`real(3), "0 0 0"`
+   Reference position relative to which the 3D vertex coordinates are defined. This vector is subtracted from the
+   positions.
+
+.. _asset-mesh-refquat:
+
+:at:`refquat`: :at-val:`real(4), "1 0 0 0"`
+   Reference orientation relative to which the 3D vertex coordinates and normals are defined. The conjugate of this
+   quaternion is used to rotate the positions and normals. The model compiler normalizes the quaternion automatically.
+
+.. _mesh-plugin:
+
+:el-prefix:`mesh/` |-| **plugin** (?)
+'''''''''''''''''''''''''''''''''''''
+
+Associate this mesh with an :ref:`engine plugin<exPlugin>`. Either :at:`plugin` or :at:`instance` are required.
+
+.. _mesh-plugin-plugin:
+
+:at:`plugin`: :at-val:`string, optional`
+   Plugin identifier, used for implicit plugin instantiation.
+
+.. _mesh-plugin-instance:
+
+:at:`instance`: :at-val:`string, optional`
+   Instance name, used for explicit plugin instantiation.
+
+
+
+.. _asset-hfield:
+
+:el-prefix:`asset/` |-| **hfield** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This element creates a height field asset, which can then be referenced from geoms with type "hfield". A height field,
+also known as terrain map, is a 2D matrix of elevation data. The data can be specified in one of three ways:
+
+#. The elevation data can be loaded from a PNG file. The image is converted internally to gray scale, and the intensity
+   of each pixel is used to define elevation; white is high and black is low.
+
+#. The elevation data can be loaded from a binary file in the custom format described below. As with all other matrices
+   used in MuJoCo, the data ordering is row-major, like pixels in an image. If the data size is nrow-by-ncol, the file
+   must have 4*(2+nrow*ncol) bytes:
+
+   ::
+
+              (int32)   nrow
+              (int32)   ncol
+              (float32) data[nrow*ncol]
+
+
+#. The elevation data can be left undefined at compile time. This is done by specifying the attributes nrow and ncol.
+   The compiler allocates space for the height field data in mjModel and sets it to 0. The user can then generate a
+   custom height field at runtime, either programmatically or using sensor data.
+
+| Regardless of which method is used to specify the elevation data, the compiler always normalizes it to the range [0
+  1]. However if the data is left undefined at compile time and generated later at runtime, it is the user's
+  responsibility to normalize it.
+| The position and orientation of the height field is determined by the geom that references it. The spatial extent on
+  the other hand is specified by the height field asset itself via the size attribute, and cannot be modified by the
+  referencing geom (the geom size parameters are ignored in this case). The same approach is used for meshes below:
+  positioning is done by the geom while sizing is done by the asset. This is because height fields and meshes involve
+  sizing operations that are not common to other geoms.
+| 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 50
+  (:ref:`mjMAXCONPAIR <glNumeric>`); any contacts beyond that are discarded. To avoid penetration due to discarded
+  contacts, the spatial features of the height field should be large compared to the geoms it collides with.
+
+.. _asset-hfield-name:
+
+:at:`name`: :at-val:`string, optional`
+   Name of the height field, used for referencing. If the name is omitted and a file name is specified, the height field
+   name equals the file name without the path and extension.
+
+.. _asset-hfield-content_type:
+
+:at:`content_type`: :at-val:`string, optional`
+   If the file attribute is specified, then this sets the
+   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`__ (formerly known as MIME types) of the
+   file to be loaded. Any filename extensions will be overloaded.  Currently ``image/png`` and
+   ``image/vnd.mujoco.hfield`` are supported.
+
+.. _asset-hfield-file:
+
+:at:`file`: :at-val:`string, optional`
+   If this attribute is specified, the elevation data is loaded from the given file. If the file extension is ".png",
+   not case-sensitive, the file is treated as a PNG file. Otherwise it is treated as a binary file in the above custom
+   format. The number of rows and columns in the data are determined from the file contents. Loading data from a file
+   and setting nrow or ncol below to non-zero values results is compile error, even if these settings are consistent
+   with the file contents.
+
+.. _asset-hfield-nrow:
+
+:at:`nrow`: :at-val:`int, "0"`
+   This attribute and the next are used to allocate a height field in mjModel. If the :at:`elevation` attribute is not
+   set, the elevation data is set to 0. This attribute specifies the number of rows in the elevation data matrix. The
+   default value of 0 means that the data will be loaded from a file, which will be used to infer the size of the
+   matrix.
+
+.. _asset-hfield-ncol:
+
+:at:`ncol`: :at-val:`int, "0"`
+   This attribute specifies the number of columns in the elevation data matrix.
+
+.. _asset-hfield-elevation:
+
+:at:`elevation`: :at-val:`real(nrow*ncol), optional`
+   This attribute specifies the elevation data matrix. Values are automatically normalized to lie between 0 and 1 by
+   first subtracting the minimum value and then dividing by the (maximum-minimum) difference, if not 0. If not provided,
+   values are set to 0.
+
+.. _asset-hfield-size:
+
+:at:`size`: :at-val:`real(4), required`
+   .. figure:: images/XMLreference/peaks.png
+      :width: 350px
+      :align: right
+
+   The four numbers here are (radius_x, radius_y, elevation_z, base_z). The height field is centered at the referencing
+   geom's local frame. Elevation is in the +Z direction. The first two numbers specify the X and Y extent (or "radius")
+   of the rectangle over which the height field is defined. This may seem unnatural for rectangles, but it is natural
+   for spheres and other geom types, and we prefer to use the same convention throughout the model. The third number is
+   the maximum elevation; it scales the elevation data which is normalized to [0-1]. Thus the minimum elevation point is
+   at Z=0 and the maximum elevation point is at Z=elevation_z. The last number is the depth of a box in the -Z direction
+   serving as a "base" for the height field. Without this automatically generated box, the height field would have zero
+   thickness at places there the normalized elevation data is zero. Unlike planes which impose global unilateral
+   constraints, height fields are treated as unions of regular geoms, so there is no notion of being "under" the height
+   field. Instead a geom is either inside or outside the height field - which is why the inside part must have non-zero
+   thickness. The example on the right is the MATLAB "peaks" surface saved in our custom height field format, and loaded
+   as an asset with size = "1 1 1 0.1". The horizontal size of the box is 2, the difference between the maximum and
+   minimum elevation is 1, and the depth of the base added below the minimum elevation point is 0.1.
+
+
+.. _asset-skin:
+
+:el-prefix:`asset/` |-| **skin** (*)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. _asset-skin-name:
+.. _asset-skin-file:
+.. _asset-skin-vertex:
+.. _asset-skin-texcoord:
+.. _asset-skin-face:
+.. _asset-skin-inflate:
+.. _asset-skin-material:
+.. _asset-skin-rgba:
+.. _asset-skin-group:
+
+:ref:`Skins<deformable-skin>` have been moved under the new grouping element :ref:`deformable<deformable>`. They can
+still be specified here but this functionality is now deprecated and will be removed in the future.
+
+
+
 .. _asset-texture:
 
 :el-prefix:`asset/` |-| **texture** (*)
@@ -1347,328 +1702,6 @@ file.
    If true, images loaded from file are flipped in the vertical direction. Does not affect procedural textures.
 
 
-.. _asset-hfield:
-
-:el-prefix:`asset/` |-| **hfield** (*)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This element creates a height field asset, which can then be referenced from geoms with type "hfield". A height field,
-also known as terrain map, is a 2D matrix of elevation data. The data can be specified in one of three ways:
-
-#. The elevation data can be loaded from a PNG file. The image is converted internally to gray scale, and the intensity
-   of each pixel is used to define elevation; white is high and black is low.
-
-#. The elevation data can be loaded from a binary file in the custom format described below. As with all other matrices
-   used in MuJoCo, the data ordering is row-major, like pixels in an image. If the data size is nrow-by-ncol, the file
-   must have 4*(2+nrow*ncol) bytes:
-
-   ::
-
-              (int32)   nrow
-              (int32)   ncol
-              (float32) data[nrow*ncol]
-
-
-#. The elevation data can be left undefined at compile time. This is done by specifying the attributes nrow and ncol.
-   The compiler allocates space for the height field data in mjModel and sets it to 0. The user can then generate a
-   custom height field at runtime, either programmatically or using sensor data.
-
-| Regardless of which method is used to specify the elevation data, the compiler always normalizes it to the range [0
-  1]. However if the data is left undefined at compile time and generated later at runtime, it is the user's
-  responsibility to normalize it.
-| The position and orientation of the height field is determined by the geom that references it. The spatial extent on
-  the other hand is specified by the height field asset itself via the size attribute, and cannot be modified by the
-  referencing geom (the geom size parameters are ignored in this case). The same approach is used for meshes below:
-  positioning is done by the geom while sizing is done by the asset. This is because height fields and meshes involve
-  sizing operations that are not common to other geoms.
-| 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 50
-  (:ref:`mjMAXCONPAIR <glNumeric>`); any contacts beyond that are discarded. To avoid penetration due to discarded
-  contacts, the spatial features of the height field should be large compared to the geoms it collides with.
-
-.. _asset-hfield-name:
-
-:at:`name`: :at-val:`string, optional`
-   Name of the height field, used for referencing. If the name is omitted and a file name is specified, the height field
-   name equals the file name without the path and extension.
-
-.. _asset-hfield-content_type:
-
-:at:`content_type`: :at-val:`string, optional`
-   If the file attribute is specified, then this sets the
-   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`__ (formerly known as MIME types) of the
-   file to be loaded. Any filename extensions will be overloaded.  Currently ``image/png`` and
-   ``image/vnd.mujoco.hfield`` are supported.
-
-.. _asset-hfield-file:
-
-:at:`file`: :at-val:`string, optional`
-   If this attribute is specified, the elevation data is loaded from the given file. If the file extension is ".png",
-   not case-sensitive, the file is treated as a PNG file. Otherwise it is treated as a binary file in the above custom
-   format. The number of rows and columns in the data are determined from the file contents. Loading data from a file
-   and setting nrow or ncol below to non-zero values results is compile error, even if these settings are consistent
-   with the file contents.
-
-.. _asset-hfield-nrow:
-
-:at:`nrow`: :at-val:`int, "0"`
-   This attribute and the next are used to allocate a height field in mjModel and leave the elevation data undefined
-   (i.e., set to 0). This attribute specifies the number of rows in the elevation data matrix. The default value of 0
-   means that the data will be loaded from a file, which will be used to infer the size of the matrix.
-
-.. _asset-hfield-ncol:
-
-:at:`ncol`: :at-val:`int, "0"`
-   This attribute specifies the number of columns in the elevation data matrix.
-
-.. _asset-hfield-size:
-
-:at:`size`: :at-val:`real(4), required`
-   .. figure:: images/XMLreference/peaks.png
-      :width: 350px
-      :align: right
-
-   The four numbers here are (radius_x, radius_y, elevation_z, base_z). The height field is centered at the referencing
-   geom's local frame. Elevation is in the +Z direction. The first two numbers specify the X and Y extent (or "radius")
-   of the rectangle over which the height field is defined. This may seem unnatural for rectangles, but it is natural
-   for spheres and other geom types, and we prefer to use the same convention throughout the model. The third number is
-   the maximum elevation; it scales the elevation data which is normalized to [0-1]. Thus the minimum elevation point is
-   at Z=0 and the maximum elevation point is at Z=elevation_z. The last number is the depth of a box in the -Z direction
-   serving as a "base" for the height field. Without this automatically generated box, the height field would have zero
-   thickness at places there the normalized elevation data is zero. Unlike planes which impose global unilateral
-   constraints, height fields are treated as unions of regular geoms, so there is no notion of being "under" the height
-   field. Instead a geom is either inside or outside the height field - which is why the inside part must have non-zero
-   thickness. The example on the right is the MATLAB "peaks" surface saved in our custom height field format, and loaded
-   as an asset with size = "1 1 1 0.1". The horizontal size of the box is 2, the difference between the maximum and
-   minimum elevation is 1, and the depth of the base added below the minimum elevation point is 0.1.
-
-
-.. _asset-mesh:
-
-:el-prefix:`asset/` |-| **mesh** (*)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-This element creates a mesh asset, which can then be referenced from geoms. If the referencing geom type is
-:at-val:`mesh` the mesh is instantiated in the model, otherwise a geometric primitive is automatically fitted to it; see
-the :ref:`geom <body-geom>` element below.
-
-MuJoCo works with triangulated meshes. They can be loaded from binary STL files, OBJ files or MSH files with custom
-format described below, or vertex and face data specified directly in the XML. Software such as MeshLab can be used to
-convert from other mesh formats to STL or OBJ. While any collection of triangles can be loaded as a mesh and rendered,
-collision detection works with the convex hull of the mesh as explained in :ref:`Collision`. See also the convexhull
-attribute of the :ref:`compiler <compiler>` element which controls the automatic generation of convex hulls. The mesh
-appearance (including texture mapping) is controlled by the :at:`material` and :at:`rgba` attributes of the referencing
-geom, similarly to height fields.
-
-Meshes can have explicit texture coordinates instead of relying on the automated texture
-mapping mechanism. When provided, these explicit coordinates have priority. Note that texture coordinates can be
-specified with OBJ files and MSH files, as well as explicitly in the XML with the :at:`texcoord` attribute, but not via
-STL files. These mechanism cannot be mixed. So if you have an STL mesh, the only way to add texture coordinates to it is
-to convert to one of the other supported formats.
-
-MSH file format
-   The binary MSH file starts with 4 integers specifying the number of vertex positions (nvertex), vertex normals
-   (nnormal), vertex texture coordinates (ntexcoord), and vertex indices making up the faces (nface), followed by the
-   numeric data. nvertex must be at least 4. nnormal and ntexcoord can be zero (in which case the corresponding data is
-   not defined) or equal to nvertex. nface can also be zero, in which case faces are constructed automatically from the
-   convex hull of the vertex positions. The file size in bytes must be exactly: 16 + 12*(nvertex + nnormal + nface) +
-   8*ntexcoord. The contents of the file must be as follows:
-
-   .. code:: Text
-
-          (int32)   nvertex
-          (int32)   nnormal
-          (int32)   ntexcoord
-          (int32)   nface
-          (float)   vertex_positions[3*nvertex]
-          (float)   vertex_normals[3*nnormal]
-          (float)   vertex_texcoords[2*ntexcoord]
-          (int32)   face_vertex_indices[3*nface]
-
-Poorly designed meshes can display rendering artifacts. In particular, the shadow mapping mechanism relies on having
-some distance between front and back-facing triangle faces. If the faces are repeated, with opposite normals as
-determined by the vertex order in each triangle, this causes shadow aliasing. The solution is to remove the repeated
-faces (which can be done in MeshLab) or use a better designed mesh. Flipped faces are checked by MuJoCo for meshes
-specified as OBJ or XML and an error message is returned.
-
-The size of the mesh is determined by the 3D coordinates of the vertex data in the mesh file, multiplied by the
-components of the :at:`scale` attribute below. Scaling is applied separately for each coordinate axis. Note that
-negative scaling values can be used to flip the mesh; this is a legitimate operation. The size parameters of the
-referening geoms are ignored, similarly to height fields. We also provide a mechanism to translate and
-rotate the 3D coordinates, using the attributes :ref:`refpos<asset-mesh-refpos>` and :ref:`refquat<asset-mesh-refquat>`.
-
-A mesh can also be defined without faces (a point cloud essentially). In that case
-the convex hull is constructed automatically, even if the compiler attribute convexhull is false. This makes it easy to
-construct simple shapes directly in the XML. For example, a pyramid can be created as:
-
-.. code-block:: xml
-
-   <asset>
-       <mesh name="tetrahedron" vertex="0 0 0  1 0 0  0 1 0  0 0 1"/>
-   </asset>
-
-Positioning and orienting is complicated by the fact that vertex data are often designed relative to coordinate frames
-whose origin is not inside the mesh. In contrast, MuJoCo expects the origin of a geom's local frame to coincide with the
-geometric center of the shape. We resolve this discrepancy by pre-processing the mesh in the compiler, so that it is
-centered around (0,0,0) and its principal axes of inertia are the coordinate axes. We also save the translation and
-rotation offsets needed to achieve such alignment in :ref:`mjModel.mesh_pos<mjModel>` and
-:ref:`mjModel.mesh_quat<mjModel>`. These offsets are then applied to the referencing geom's position and orientation; see
-also :at:`mesh` attribute of :ref:`geom <body-geom>` below. Fortunately most meshes used in robot models are designed in
-a coordinate frame centered at the joint. This makes the corresponding MJCF model intuitive: we set the body frame at the
-joint, so that the joint position is (0,0,0) in the body frame, and simply reference the mesh. Below is an MJCF model
-fragment of a forearm, containing all the information needed to put the mesh where one would expect it to be. The body
-position is specified relative to the parent body, namely the upper arm (not shown). It is offset by 35 cm which is the
-typical length of the human upper arm. If the mesh vertex data were not designed in the above convention, we would have
-to use the geom position and orientation (or the new refpos, refquat mechanism) to compensate, but in practice this is
-rarely needed.
-
-.. code-block:: xml
-
-   <asset>
-       <mesh file="forearm.stl"/>
-   </asset>
-
-   <body pos="0 0 0.35"/>
-       <joint type="hinge" axis="1 0 0"/>
-       <geom type="mesh" mesh="forearm"/>
-   </body>
-
-The inertial computation mentioned above is part of an algorithm used not only to center and align the mesh, but also to
-infer the mass and inertia of the body to which it is attached. This is done by computing the centroid of the triangle
-faces, connecting each face with the centroid to form a triangular pyramid, computing the mass and signed inertia of all
-pyramids (considered solid or hollow if :at:`shellinertia` is true) and accumulating them. The sign ensures that
-pyramids on the outside of the surfaces are subtracted, as it can occur with concave geometries. This algorithm can be
-found in section 1.3.8 of Computational Geometry in C (Second Edition) by Joseph O'Rourke.
-
-The full list of processing steps applied by the compiler to each mesh is as follows:
-
-#. For STL meshes, remove any repeated vertices and re-index the faces if needed. If the mesh is not STL, we assume that
-   the desired vertices and faces have already been generated and do not apply removal or re-indexing;
-#. If vertex normals are not provided, generate normals automatically, using a weighted average of the surrounding face
-   normals. If sharp edges are encountered, the renderer uses the face normals to preserve the visual information about
-   the edge, unless smoothnormal is true. Note that normals cannot be provided with STL meshes;
-#. Scale, translate and rotate the vertices and normals, re-normalize the normals in case of scaling;
-#. Construct the convex hull if specified;
-#. Find the centroid of all triangle faces, and construct the union-of-pyramids representation. Triangles whose area is
-   too small (below the :ref:`mjMINVAL <glNumeric>` value of 1E-14) result in compile error;
-#. Compute the center of mass and inertia matrix of the union-of-pyramids. Use eigenvalue decomposition to find the
-   principal axes of inertia. Center and align the mesh, saving the translational and rotational offsets for subsequent
-   geom-related computations.
-
-.. _asset-mesh-name:
-
-:at:`name`: :at-val:`string, optional`
-   Name of the mesh, used for referencing. If omitted, the mesh name equals the file name without the path and
-   extension.
-
-.. _asset-mesh-class:
-
-:at:`class`: :at-val:`string, optional`
-   Defaults class for setting unspecified attributes (only scale in this case).
-
-.. _asset-mesh-content_type:
-
-:at:`content_type`: :at-val:`string, optional`
-   If the file attribute is specified, then this sets the
-   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`_ (formerly known as MIME type) of the
-   file to be loaded. Any filename extensions will be overloaded.  Currently ``model/vnd.mujoco.msh``, ``model/obj``,
-   and ``model/stl`` are supported.
-
-.. _asset-mesh-file:
-
-:at:`file`: :at-val:`string, optional`
-   The file from which the mesh will be loaded. The path is determined as described in the meshdir attribute of
-   :ref:`compiler <compiler>`. The file extension must be "stl", "msh", or "obj" (not case sensitive) specifying the
-   file type.  If the file name is omitted, the vertex attribute becomes required.
-
-.. _asset-mesh-scale:
-
-:at:`scale`: :at-val:`real(3), "1 1 1"`
-   This attribute specifies the scaling that will be applied to the vertex data along each coordinate axis. Negative
-   values are allowed, resulting in flipping the mesh along the corresponding axis.
-
-.. _asset-mesh-smoothnormal:
-
-:at:`smoothnormal`: :at-val:`[false, true], "false"`
-   Controls the automatic generation of vertex normals when normals are not given explicitly. If true, smooth normals
-   are generated by averaging the face normals at each vertex, with weight proportional to the face area. If false,
-   faces at large angles relative to the average normal are excluded from the average. In this way, sharp edges (as in
-   cube edges) are not smoothed.
-
-.. _asset-mesh-vertex:
-
-:at:`vertex`: :at-val:`real(3*nvert), optional`
-   Vertex 3D position data. You can specify position data in the XML using this attribute, or using a binary file, but
-   not both.
-
-.. _asset-mesh-normal:
-
-:at:`normal`: :at-val:`real(3*nvert), optional`
-   Vertex 3D normal data. If specified, the number of normals must equal the number of vertices. The model compiler
-   normalizes the normals automatically.
-
-.. _asset-mesh-texcoord:
-
-:at:`texcoord`: :at-val:`real(2*nvert), optional`
-   Vertex 2D texture coordinates, which are numbers between 0 and 1. If specified, the number of texture coordinate
-   pairs must equal the number of vertices.
-
-.. _asset-mesh-face:
-
-:at:`face`: :at-val:`int(3*nface), optional`
-   Faces of the mesh. Each face is a sequence of 3 vertex indices, in counter-clockwise order. The indices must be
-   integers between 0 and nvert-1.
-
-.. _asset-mesh-refpos:
-
-:at:`refpos`: :at-val:`real(3), "0 0 0"`
-   Reference position relative to which the 3D vertex coordinates are defined. This vector is subtracted from the
-   positions.
-
-.. _asset-mesh-refquat:
-
-:at:`refquat`: :at-val:`real(4), "1 0 0 0"`
-   Reference orientation relative to which the 3D vertex coordinates and normals are defined. The conjugate of this
-   quaternion is used to rotate the positions and normals. The model compiler normalizes the quaternion automatically.
-
-.. _mesh-plugin:
-
-:el-prefix:`mesh/` |-| **plugin** (?)
-'''''''''''''''''''''''''''''''''''''
-
-Associate this mesh with an :ref:`engine plugin<exPlugin>`. Either :at:`plugin` or :at:`instance` are required.
-
-.. _mesh-plugin-plugin:
-
-:at:`plugin`: :at-val:`string, optional`
-   Plugin identifier, used for implicit plugin instantiation.
-
-.. _mesh-plugin-instance:
-
-:at:`instance`: :at-val:`string, optional`
-   Instance name, used for explicit plugin instantiation.
-
-
-.. _asset-skin:
-
-:el-prefix:`asset/` |-| **skin** (*)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-.. _asset-skin-name:
-.. _asset-skin-file:
-.. _asset-skin-vertex:
-.. _asset-skin-texcoord:
-.. _asset-skin-face:
-.. _asset-skin-inflate:
-.. _asset-skin-material:
-.. _asset-skin-rgba:
-.. _asset-skin-group:
-
-:ref:`Skins<deformable-skin>` have been moved under the new grouping element :ref:`deformable<deformable>`. They can
-still be specified here but this functionality is now deprecated and will be removed in the future.
-
 
 .. _asset-material:
 
@@ -1944,12 +1977,19 @@ adjust it properly through the XML.
 :at:`sdf_initpoints`: :at-val:`int, "40"`
    Number of starting points used for finding contacts with Signed Distance Field collisions.
 
+.. youtube:: H9qG9Zf2W44
+   :align: right
+   :width: 240px
+
 .. _option-actuatorgroupdisable:
 
-:at:`actuatorgroupdisable`: :at-val:`int(30), ""`
+:at:`actuatorgroupdisable`: :at-val:`int(31), optional`
    List of actuator groups to disable. Actuators whose :ref:`group<actuator-general-group>` is in this list will produce
    no force. If they are stateful, their activation states will not be integrated. Internally this list is
-   implemented as an integer bitfield, so values must be in the range ``0 <= group <= 30``.
+   implemented as an integer bitfield, so values must be in the range ``0 <= group <= 30``. If not set, all actuator
+   groups are enabled. See `example model
+   <https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
+   and associated screen-capture on the right.
 
 .. _option-flag:
 
@@ -2626,7 +2666,7 @@ helps clarify the role of bodies and geoms in MuJoCo.
    | capsule | 1 or 2 | Radius of the capsule; half-length of the cylinder part when not using the :at:`fromto`        |
    |         |        | specification.                                                                                 |
    +---------+--------+------------------------------------------------------------------------------------------------+
-   |ellipsoid| 1      | X radius; Y radius; Z radius.                                                                  |
+   |ellipsoid| 3      | X radius; Y radius; Z radius.                                                                  |
    +---------+--------+------------------------------------------------------------------------------------------------+
    |cylinder | 1 or 2 | Radius of the cylinder; half-length of the cylinder when not using the :at:`fromto`            |
    |         |        | specification.                                                                                 |
@@ -4561,7 +4601,7 @@ of the other body, without any joint elements in the child body.
 
 :at:`body2`: :at-val:`string, optional`
    Name of the second body. If this attribute is omitted, the second body is the world body. Welding a body to the world
-   and changing the corresponding component of mjModel.eq_active at runtime can be used to fix the body temporarily.
+   and changing the corresponding component of mjData.eq_active at runtime can be used to fix the body temporarily.
 
 .. _equality-weld-relpose:
 
@@ -7647,8 +7687,8 @@ if omitted.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 | This element sets the attributes of the dummy :ref:`general <actuator-general>` element of the defaults class.
-| All general attributes are available here except: name, class, joint, jointinparent, site, tendon, slidersite,
+| All general attributes are available here except: name, class, joint, jointinparent, site, refsite, tendon,
-  cranksite.
+  slidersite, cranksite.
 
 
 .. _default-motor:
@@ -7675,8 +7715,8 @@ if omitted.
 This and the next three elements set the attributes of the :ref:`general <actuator-general>` element using
 :ref:`Actuator shortcuts <CActShortcuts>`. It does not make sense to use more than one such shortcut in the same defaults
 class, because they set the same underlying attributes, replacing any previous settings. All
-:ref:`motor <actuator-motor>` attributes are available here except: name, class, joint, jointinparent, site, tendon,
+:ref:`motor <actuator-motor>` attributes are available here except: name, class, joint, jointinparent, site, refsite,
-slidersite, cranksite.
+tendon, slidersite, cranksite.
 
 
 .. _default-position:
@@ -7705,7 +7745,7 @@ slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`position <actuator-position>` attributes are available here except: name, class, joint, jointinparent, site,
-tendon, slidersite, cranksite.
+refsite, tendon, slidersite, cranksite.
 
 
 .. _default-velocity:
@@ -7732,7 +7772,7 @@ tendon, slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`velocity <actuator-velocity>` attributes are available here except: name, class, joint, jointinparent, site,
-tendon, slidersite, cranksite.
+refsite, tendon, slidersite, cranksite.
 
 
 .. _default-intvelocity:
@@ -7763,7 +7803,7 @@ tendon, slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`intvelocity <actuator-intvelocity>` attributes are available here except: name, class, joint, jointinparent,
-site, tendon, slidersite, cranksite.
+site, refsite, tendon, slidersite, cranksite.
 
 
 .. _default-damper:
@@ -7788,7 +7828,7 @@ site, tendon, slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`damper <actuator-damper>` attributes are available here except: name, class, joint, jointinparent, site,
-tendon, slidersite, cranksite.
+refsite, tendon, slidersite, cranksite.
 
 
 .. _default-cylinder:
@@ -7821,7 +7861,7 @@ tendon, slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`cylinder <actuator-cylinder>` attributes are available here except: name, class, joint, jointinparent, site,
-tendon, slidersite, cranksite.
+refsite, tendon, slidersite, cranksite.
 
 
 .. _default-muscle:
@@ -7864,7 +7904,7 @@ tendon, slidersite, cranksite.
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 All :ref:`muscle <actuator-muscle>` attributes are available here except: name, class, joint, jointinparent, site,
-tendon, slidersite, cranksite.
+refsite, tendon, slidersite, cranksite.
 
 
 .. _default-adhesion:

doc/XMLschema.rst

@@ -65,7 +65,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`linewidth<visual-global-linewidth>`                       | :ref:`glow<visual-global-glow>`                                 | :ref:`offwidth<visual-global-offwidth>`                         | :ref:`offheight<visual-global-offheight>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`realtime<visual-global-realtime>`                         | :ref:`ellipsoidinertia<visual-global-ellipsoidinertia>`         |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`realtime<visual-global-realtime>`                         | :ref:`ellipsoidinertia<visual-global-ellipsoidinertia>`         | :ref:`bvactive<visual-global-bvactive>`                         |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| visual |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -126,7 +126,9 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`contacttorque<visual-rgba-contacttorque>`                 | :ref:`contactgap<visual-rgba-contactgap>`                       | :ref:`rangefinder<visual-rgba-rangefinder>`                     | :ref:`constraint<visual-rgba-constraint>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`slidercrank<visual-rgba-slidercrank>`                     | :ref:`crankbroken<visual-rgba-crankbroken>`                     | :ref:`frustum<visual-rgba-frustum>`                             |                                                                 | |
+|                                    |    |    | :ref:`slidercrank<visual-rgba-slidercrank>`                     | :ref:`crankbroken<visual-rgba-crankbroken>`                     | :ref:`frustum<visual-rgba-frustum>`                             | :ref:`bv<visual-rgba-bv>`                                       | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`bvactive<visual-rgba-bvactive>`                           |                                                                 |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | mujoco |br| |L|                    |    | *no attributes*                                                                                                                                                                                                                                                              |
@@ -155,7 +157,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<asset-hfield-name>`                                  | :ref:`content_type<asset-hfield-content_type>`                  | :ref:`file<asset-hfield-file>`                                  | :ref:`nrow<asset-hfield-nrow>`                                  | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`ncol<asset-hfield-ncol>`                                  | :ref:`size<asset-hfield-size>`                                  |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`ncol<asset-hfield-ncol>`                                  | :ref:`size<asset-hfield-size>`                                  | :ref:`elevation<asset-hfield-elevation>`                        |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| asset |br| |_| |L|             |    | .. table::                                                                                                                                                                                                                                                                   |

doc/changelog.rst

@@ -5,22 +5,66 @@ Changelog
 Upcoming version (not yet released)
 -----------------------------------
 
+MJX
+^^^
+
+1. Improved performance of getting and putting device data.
+
+   - Use ``tobytes()`` for numpy array serialization, which is orders of magnitude faster than converting to tuples.
+   - Avoid reallocating host ``mjData`` arrays when array shapes are unchanged.
+   - Speed up calculation of ``mjx.ncon`` for models with many geoms.
+   - Avoid calling ``mjx.ncon`` in ``mjx.get_data_into`` when ``nc`` can be derived from ``mjx.Data``.
+2. Fixed a bug in ``mjx-viewer`` that prevented it from running.  Updated ``mjx-viewer`` to use newer
+   ``mjx.get_data_into`` function call.
+
+Version 3.1.2 (February 05, 2024)
+-----------------------------------
+
 General
 ^^^^^^^
-1. Improved the :ref:discardvisual<compiler-discardvisual> compiler flag, which now discards all visual-only assets. See
+1. Improved the :ref:`discardvisual<compiler-discardvisual>` compiler flag, which now discards all visual-only assets.
-   :ref:discardvisual<compiler-discardvisual> for details.
+   See :ref:`discardvisual<compiler-discardvisual>` for details.
+2. Removed the :ref:`timer<mjtTimer>` for midphase colllision detection, it is now folded in with the narrowphase
+   timer. This is because timing the two phases seperately required fine-grained timers inside the collision
+   functions; these functions are so small and fast that the timer itself was incurring a measurable cost.
+3. Added the flag :ref:`bvactive<visual-global-bvactive>` to ``visual/global``, allowing users to turn off
+   visualisation of active bounding volumes (the red/green boxes in this :ref:`this changelog item<midphase>`). For
+   models with very high-resolution meshes, the computation required for this visualization can slow down simulation
+   speed. Fixes :github:issue:`1279`.
+
+   - Added color of :ref:`bounding volumes<visual-rgba-bv>` and :ref:`active bounding volumes<visual-rgba-bvactive>`
+     to :ref:`visual/rgba<visual-rgba>`.
+4. Height-field elevation data can now be specified directly in XML with the :ref:`elevation<asset-hfield-elevation>`
+   attribute (and not only with PNG files).
+   See `example model <https://github.com/google-deepmind/mujoco/blob/main/test/user/testdata/hfield_xml.xml>`__.
 
 MJX
 ^^^
-2. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
+5. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
-3. Added :at:`site` transmission.
+6. Added :at:`site` transmission.
-4. Updated MJX colab tutorial with more stable quadruped environment.
+7. Updated MJX colab tutorial with more stable quadruped environment.
-5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes.
+8. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes.
+9. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`.
+10. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`.
+11. Raise a not implemented error when nonzero frictionloss is present. Fixes :github:issue:`1344`.
+
+Python bindings
+^^^^^^^^^^^^^^^
+12. Improved the implementation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
+    code will require modification.
+
+    - Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
+    - Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
+    - Outputs are no longer squeezed and always have dim=3.
+13. The ``sync`` function for the :ref:`passive viewer<PyViewerPassive>` can now pick up changes to rendering flags in
+    ``user_scn``, as requested in :github:issue:`1190`.
 
 Bug fixes
 ^^^^^^^^^
-6. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
+14. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
-   :github:issue:`1270`.
+    :github:issue:`1270`.
+15. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
+    bound of the length range. Fixes :github:issue:`1342`.
 
 
 Version 3.1.1 (December 18, 2023)
@@ -28,7 +72,8 @@ 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.
+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``
@@ -41,8 +86,8 @@ 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
+   for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is
-   robust for very small or large geom sizes.
+   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>`
@@ -52,17 +97,14 @@ General
 
 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
@@ -74,8 +116,8 @@ Bug fixes
    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
+11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where ``SYSTEM_VERSION_COMPAT`` is set.
-    supported version is still 11.0, but we release these wheels to fix compatibility for those users. See
+    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)
@@ -272,7 +314,7 @@ General
 
        - For models which have ``<option collision="all"/>``, delete the attribute.
        - For models which have ``<option collision="dynamic"/>``, delete all :ref:`pair<contact-pair>` elements.
-       - For models which have ``<option collision="pair"/>``, disable all dynamic collisions (determined
+       - For models which have ``<option collision="predefined"/>``, disable all dynamic collisions (determined
          via contype/conaffinity) by first deleting all :ref:`contype<body-geom-contype>` and
          :ref:`conaffinity<body-geom-conaffinity>` attributes in the model and then setting them globally to ``0`` using
          |br| ``<default> <geom contype="0" conaffinity="0"/> </default>``.
@@ -619,6 +661,8 @@ General
    :align: right
    :width: 350px
 
+.. _midphase:
+
 2. Added a collision mid-phase for pruning geoms in body pairs, see :ref:`documentation<coSelection>` for more details.
    This is based on static AABB bounding volume hierarchy (a BVH binary tree) in the body inertial frame. The GIF on
    the right is cut from `this longer video <https://youtu.be/e0babIM8hBo>`__.

doc/computation/index.rst

@@ -296,9 +296,9 @@ is attached; the possible attachment object types are :at:`joint`, :at:`tendon`,
 :at:`slider-crank`
    :at:`slider-crank` `transmissions <https://en.wikipedia.org/wiki/Slider-crank_linkage>`_ transform a linear force to
    a torque, as in a piston-driven combustion engine. `This model
-   <https://github.com/google-deepmind/mujoco/tree/main/model/slider_crank>`_ contains pedagogical examples.
+   <https://github.com/google-deepmind/mujoco/blob/main/model/slider_crank/slider_crank.xml>`__ contains pedagogical
-   Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality constraints to
+   examples. Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality
-   the rest of the system, but that would be less efficient.
+   constraints, but that is both less efficient and less stable.
 
 :at:`body`
    :el:`body` transmission corresponds to applying forces at contact points belonging to a body, in

doc/includes/references.h

@@ -75,7 +75,6 @@ typedef enum mjtTimer_ {     // internal timers
 
   // breakdown of mj_collision
   mjTIMER_COL_BROAD,         // broadphase
-  mjTIMER_COL_MID,           // midphase
   mjTIMER_COL_NARROW,        // narrowphase
 
   mjNTIMER                   // number of timers
@@ -148,7 +147,7 @@ struct mjData_ {
 
   // memory utilization stats
   size_t  maxuse_stack;                      // maximum stack allocation in bytes
-  size_t  maxuse_threadstack[mjMAXTHREADS];  // maximum stack allocation per thread in bytes
+  size_t  maxuse_threadstack[mjMAXTHREAD];   // maximum stack allocation per thread in bytes
   size_t  maxuse_arena;                      // maximum arena allocation in bytes
   int     maxuse_con;                        // maximum number of contacts
   int     maxuse_efc;                        // maximum number of scalar constraints
@@ -754,6 +753,7 @@ struct mjVisual_ {                // visualization options
     int   offwidth;               // width of offscreen buffer
     int   offheight;              // height of offscreen buffer
     int   ellipsoidinertia;       // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   bvactive;               // visualize active bounding volumes (0: no, 1: yes)
   } global;
 
   struct {                        // rendering quality
@@ -831,6 +831,8 @@ struct mjVisual_ {                // visualization options
     float slidercrank[4];         // slidercrank
     float crankbroken[4];         // used when crank must be stretched/broken
     float frustum[4];             // camera frustum
+    float bv[4];                  // bounding volume
+    float bvactive[4];            // active bounding volume
   } rgba;
 };
 typedef struct mjVisual_ mjVisual;

doc/mjx.rst

@@ -230,7 +230,7 @@ The following features are **in development** and coming soon:
    * - :ref:`Geom <mjtGeom>`
      - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
    * - :ref:`Constraint <mjtConstraint>`
-     - ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
+     - :ref:`Frictionloss <coFriction>`, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
    * - :ref:`Integrator <mjtIntegrator>`
      - ``IMPLICIT``, ``IMPLICITFAST``
    * - :ref:`Cone <mjtCone>`
@@ -349,3 +349,21 @@ For MJX to perform well, some configuration parameters should be adjusted from t
 
 :ref:`option-flag` element
   Disabling ``eulerdamp`` can help performance and is often not needed for stability.
+
+:ref:`option-jacobian` element
+  Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized
+  hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as
+  they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or
+  greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the
+  Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10%
+  to 20%, as long as the dense matrices can fit on the device.
+
+GPU performance tuning
+----------------------
+
+The following environment variables should be set:
+
+``XLA_FLAGS=--xla_gpu_triton_gemm_any=true``
+  This enables the Triton-based GEMM (matmul) emitter for any GEMM that it supports.  This can yield a 30% speedup on
+  NVIDIA GPUs.  If you have multiple GPUs, you may also benefit from enabling flags related to
+  `communciation between GPUs <https://jax.readthedocs.io/en/latest/gpu_performance_tips.html>`__.

doc/models.rst

@@ -7,8 +7,11 @@ Model Gallery
 MuJoCo Menagerie
 ----------------
 
-Menagerie is a collection of high-quality simulation models for MuJoCo, curated
+A physics simulator is only as good as the model it is simulating, and in a
-by Google DeepMind.
+powerful simulator like MuJoCo with many modeling options, it is easy to create
+"bad" models which do not behave as expected. The goal of **MuJoCo Menagerie**
+is to provide the community with a curated collection of well-designed models
+that work well right out of the box.
 
 For more information, visit the `Menagerie repository <https://github.com/google-deepmind/mujoco_menagerie>`__.
 
@@ -21,9 +24,20 @@ Bipeds
    * - Model
      - Preview
    * - `Agility Cassie <https://github.com/google-deepmind/mujoco_menagerie/tree/main/agility_cassie>`_
-     - .. youtube:: rcdsAdwNhtc
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/agility_cassie/cassie.png
+
+Humanoids
+^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+
+   * - Model
+     - Preview
+   * - `Unitree H1 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/unitree_h1>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/unitree_h1/h1.png
    * - `Robotis OP3 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/robotis_op3>`_
-     - .. youtube:: jLZ3sdkyz_w
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/robotis_op3/op3.png
 
 Mobile Manipulators
 ^^^^^^^^^^^^^^^^^^^
@@ -33,8 +47,10 @@ Mobile Manipulators
 
    * - Model
      - Preview
+   * - `Google Robot <https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_robot>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/google_robot/robot.png
    * - `Hello Robot Stretch 2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/hello_robot_stretch>`_
-     - .. youtube:: w_NUKO61wIc
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/hello_robot_stretch/stretch.png
 
 Quadrupeds
 ^^^^^^^^^^
@@ -45,15 +61,19 @@ Quadrupeds
    * - Model
      - Preview
    * - `Unitree A1 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/unitree_a1>`_
-     - .. youtube:: paQMrMtnTtc
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/unitree_a1/a1.png
    * - `Unitree Go1 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/unitree_go1>`_
-     - .. youtube:: 4d7I67BzDJg
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/unitree_go1/go1.png
+   * - `Unitree Go2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/unitree_go2>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/unitree_go2/go2.png
    * - `Anybotics ANYmal B <https://github.com/google-deepmind/mujoco_menagerie/tree/main/anybotics_anymal_b>`_
-     - .. youtube:: fRHau-PMGgM
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/anybotics_anymal_b/anymal_b.png
    * - `Anybotics ANYmal C <https://github.com/google-deepmind/mujoco_menagerie/tree/main/anybotics_anymal_c>`_
-     - .. youtube:: v04uJWBLwFQ
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/anybotics_anymal_c/anymal_c.png
    * - `Google Barkour v0 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0>`_
-     - .. youtube:: w9EA0joEAMo
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/google_barkour_v0/barkour_v0.png
+   * - `Google Barkour vB <https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_vb>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/google_barkour_vb/barkour_vb.png
 
 Grippers & Hands
 ^^^^^^^^^^^^^^^^
@@ -64,11 +84,11 @@ Grippers & Hands
    * - Model
      - Preview
    * - `Shadow E3M5 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/shadow_hand>`_
-     - .. youtube:: wi_zJzRm8Ic
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/shadow_hand/shadow_hand.png
    * - `Robotiq 2F-85 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/robotiq_2f85>`_
-     - .. youtube:: yYm9fLj32Xw
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/robotiq_2f85/2f85.png
    * - `Wonik Allegro <https://github.com/google-deepmind/mujoco_menagerie/tree/main/wonik_allegro>`_
-     - .. youtube:: jDWko1WTRXc
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/wonik_allegro/allegro_hand.png
 
 Arms
 ^^^^
@@ -79,15 +99,25 @@ Arms
    * - Model
      - Preview
    * - `Franka Emika Panda <https://github.com/google-deepmind/mujoco_menagerie/tree/main/franka_emika_panda>`_
-     - .. youtube:: H5zSrWcJlGs
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/franka_emika_panda/panda.png
    * - `Universal Robots UR5e <https://github.com/google-deepmind/mujoco_menagerie/tree/main/universal_robots_ur5e>`_
-     - .. youtube:: gAqwNeY0juo
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/universal_robots_ur5e/ur5e.png
+   * - `Universal Robots UR10e <https://github.com/google-deepmind/mujoco_menagerie/tree/main/universal_robots_ur10e>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/universal_robots_ur10e/ur10e.png
    * - `LBR iiwa14 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/kuka_iiwa_14>`_
-     - .. youtube:: 4Z44nkNXkwo
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/kuka_iiwa_14/iiwa_14.png
    * - `UFACTORY xArm7 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/ufactory_xarm7>`_
-     - .. youtube:: mMDisja5ark
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/ufactory_xarm7/xarm7.png
    * - `Rethink Robotics Sawyer <https://github.com/google-deepmind/mujoco_menagerie/tree/main/rethink_robotics_sawyer>`_
-     - .. youtube:: sZ41oklVvBg
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/rethink_robotics_sawyer/sawyer.png
+   * - `UFactory Lite 6 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/ufactory_lite6>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/ufactory_lite6/lite6.png
+   * - `ViperX 300 6DOF <https://github.com/google-deepmind/mujoco_menagerie/tree/main/trossen_vx300s>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/trossen_vx300s/vx300s.png
+   * - `ALOHA 2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/aloha>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/aloha/aloha.png
+   * - `Unitree Z1 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/unitree_z1>`_
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/unitree_z1/z1.png
 
 Drones
 ^^^^^^
@@ -98,4 +128,4 @@ Drones
    * - Model
      - Preview
    * - `Skydio X2 <https://github.com/google-deepmind/mujoco_menagerie/tree/main/skydio_x2>`_
-     - .. youtube:: LBsvsgnSvoM
+     - .. image:: https://raw.githubusercontent.com/google-deepmind/mujoco_menagerie/main/skydio_x2/x2.png

doc/python.rst

@@ -124,15 +124,22 @@ attributes:
 - ``is_running()``: returns ``True`` if the viewer window is running and ``False`` if it is closed.
   This method can be safely called without locking.
 
-- ``user_scn``: an :ref:`mjvScene` object that allows users to add custom visualization geoms to the rendered scene.
+- ``user_scn``: an :ref:`mjvScene` object that allows users to add change rendering flags and add custom
-  This is separate from the ``mjvScene`` that the viewer uses internally to render the final scene, and is entirely
+  visualization geoms to the rendered scene. This is separate from the ``mjvScene`` that the viewer uses internally to
-  under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or :ref:`mjv_makeConnector` to add
+  render the final scene, and is entirely under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or
-  visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer will incorporate
+  :ref:`mjv_makeConnector` to add visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer
-  these geoms to future rendered images. For example:
+  will incorporate these geoms to future rendered images. Similarly, user scripts can make changes to ``user_scn.flags``
+  which would be picked up at the next call to ``sync()``. The ``sync()`` call also copies changes to rendering flags
+  made via the GUI back into ``user_scn`` to preserve consistency. For example:
 
   .. code-block:: python
 
     with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
+
+      # Enable wireframe rendering of the entire scene.
+      viewer.user_scn.flags[mujoco.mjtRndFlag.mjRND_WIREFRAME] = 1
+      viewer.sync()
+
       while viewer.is_running():
         ...
         # Step the physics.
@@ -465,28 +472,32 @@ Open-loop rollouts
 ==================
 
 We include a code sample showing how to add additional C/C++ functionality, exposed as a Python module via pybind11. The
-sample, implemented in ``rollout.cc`` and wrapped in ``rollout.py``, implements a common use case where tight loops
+sample, implemented in `rollout.cc <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__
-implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling ``mj_step()`` in a loop), given an
+and wrapped in `rollout.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__,
-intial state and sequence of controls, and returning subsequent states and sensor values. The canonical usage form is
+implements a common use case where tight loops implemented outside of Python are beneficial: rolling out a trajectory
+(i.e., calling ``mj_step()`` in a loop), given an intial state and sequence of controls, and returning subsequent states
+and sensor values. The basic usage form is
 
 .. code-block:: python
 
-   state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+   state, sensordata = rollout.rollout(model, data, initial_state, control)
 
-``initial_state`` is a ``nstate x nqva`` array, with ``nstate`` initial states of length ``nqva``, where ``nqva =
+``initial_state`` is a ``nroll x nstate`` array, with ``nroll`` initial states of size ``nstate``, where
-model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities
+``nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`` is the size of the
-(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control
+:ref:`full physics state<geFullPhysics>`. ``control`` is a ``nroll x nstep x ncontrol`` array of controls. Controls are
-sequences.
+by default the ``mjModel.nu`` standard actuators, but any combination of :ref:`user input<geInput>` arrays can be
+specified by passing an optional ``control_spec`` bitflag.
+
+If a rollout diverges, the current state and sensor values are used to fill the remainder of the trajectory.
+Therefore, non-increasing time values can be used to detect diverged rollouts.
 
 The ``rollout`` function is designed to be completely stateless, so all inputs of the stepping pipeline are set and any
-values already present in the given ``MjData`` instance will have no effect on the output. In order to facilitate this,
+values already present in the given ``MjData`` instance will have no effect on the output.
-all inputs including ``time`` and ``qacc_warmstart`` are set to default values, as are auxillary controls
-(``qfrc_applied``, ``xfrc_applied`` and ``mocap_{pos,quat}``). These can also be optionally set by the user.
 
 Since the Global Interpreter Lock can be released, this function can be efficiently threaded using Python threads. See
 the ``test_threading`` function in
-`rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`_ for an example of
+`rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`__ for an example
-threaded operation.
+of threaded operation (and more generally for usage examples).
 
 .. _PyMjpy_migration:
 

doc/requirements.txt

@@ -1,4 +1,4 @@
-Sphinx==4.5.0
+Sphinx==5.3.0
 furo==2022.9.29
 sphinxcontrib-bibtex==2.6.1
 sphinxcontrib-katex==0.9.4
@@ -11,7 +11,7 @@ nbsphinx==0.9.1
 pandoc==1.1.0
 pygments==2.15.0
 jq==1.4.1
-Jinja2==2.11.3
+Jinja2~=3.0
 wheel
 # see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe
 markupsafe==2.0.1

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.1.2.dylib`` (it can be
+``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.1.3.dylib`` (it can be
 found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``.
 
 Linux
 _____
 
 Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from
-``~/.mujoco/mujoco-3.1.2/lib/libmujoco.so.3.1.2`` and rename it as ``libmujoco.so``.
+``~/.mujoco/mujoco-3.1.3/lib/libmujoco.so.3.1.3`` and rename it as ``libmujoco.so``.
 
 Windows
 _______

include/mujoco/mjdata.h

@@ -87,7 +87,6 @@ typedef enum mjtTimer_ {     // internal timers
 
   // breakdown of mj_collision
   mjTIMER_COL_BROAD,         // broadphase
-  mjTIMER_COL_MID,           // midphase
   mjTIMER_COL_NARROW,        // narrowphase
 
   mjNTIMER                   // number of timers
@@ -176,7 +175,7 @@ struct mjData_ {
 
   // memory utilization stats
   size_t  maxuse_stack;                      // maximum stack allocation in bytes
-  size_t  maxuse_threadstack[mjMAXTHREADS];  // maximum stack allocation per thread in bytes
+  size_t  maxuse_threadstack[mjMAXTHREAD];   // maximum stack allocation per thread in bytes
   size_t  maxuse_arena;                      // maximum arena allocation in bytes
   int     maxuse_con;                        // maximum number of contacts
   int     maxuse_efc;                        // maximum number of scalar constraints

include/mujoco/mjmodel.h

@@ -459,6 +459,7 @@ struct mjVisual_ {                // visualization options
     int   offwidth;               // width of offscreen buffer
     int   offheight;              // height of offscreen buffer
     int   ellipsoidinertia;       // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   bvactive;               // visualize active bounding volumes (0: no, 1: yes)
   } global;
 
   struct {                        // rendering quality
@@ -536,6 +537,8 @@ struct mjVisual_ {                // visualization options
     float slidercrank[4];         // slidercrank
     float crankbroken[4];         // used when crank must be stretched/broken
     float frustum[4];             // camera frustum
+    float bv[4];                  // bounding volume
+    float bvactive[4];            // active bounding volume
   } rgba;
 };
 typedef struct mjVisual_ mjVisual;

include/mujoco/mjthread.h

@@ -15,7 +15,7 @@
 #ifndef MUJOCO_INCLUDE_MJTHREAD_H_
 #define MUJOCO_INCLUDE_MJTHREAD_H_
 
-#define mjMAXTHREADS 128       // maximum number of threads in a thread pool
+#define mjMAXTHREAD 128        // maximum number of threads in a thread pool
 
 typedef enum mjtTaskStatus_ {  // status values for mjTask
   mjTASK_NEW = 0,              // newly created

include/mujoco/mjxmacro.h

@@ -739,7 +739,7 @@
 
 // vector fields of mjData
 #define MJDATA_VECTOR                                                \
-    X( size_t,         maxuse_threadstack, mjMAXTHREADS, 1         ) \
+    X( size_t,         maxuse_threadstack, mjMAXTHREAD,  1         ) \
     X( mjWarningStat,  warning,            mjNWARNING,   1         ) \
     X( mjTimerStat,    timer,              mjNTIMER,     1         ) \
     X( mjSolverStat,   solver,             mjNILSAND,    mjNSOLVER ) \

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 312
+#define mjVERSION_HEADER 313
 
 // needed to define size_t, fabs and log10
 #include <stdlib.h>

introspect/enums.py

@@ -458,9 +458,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjTIMER_POS_MAKE', 11),
              ('mjTIMER_POS_PROJECT', 12),
              ('mjTIMER_COL_BROAD', 13),
-             ('mjTIMER_COL_MID', 14),
+             ('mjTIMER_COL_NARROW', 14),
-             ('mjTIMER_COL_NARROW', 15),
+             ('mjNTIMER', 15),
-             ('mjNTIMER', 16),
          ]),
      )),
     ('mjtCatBit',

introspect/structs.py

@@ -352,6 +352,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              type=ValueType(name='int'),
                              doc='geom for inertia visualization (0: box, 1: ellipsoid)',  # pylint: disable=line-too-long
                          ),
+                         StructFieldDecl(
+                             name='bvactive',
+                             type=ValueType(name='int'),
+                             doc='visualize active bounding volumes (0: no, 1: yes)',  # pylint: disable=line-too-long
+                         ),
                      ),
                  ),
                  doc='',
@@ -780,6 +785,22 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                              ),
                              doc='camera frustum',
                          ),
+                         StructFieldDecl(
+                             name='bv',
+                             type=ArrayType(
+                                 inner_type=ValueType(name='float'),
+                                 extents=(4,),
+                             ),
+                             doc='bounding volume',
+                         ),
+                         StructFieldDecl(
+                             name='bvactive',
+                             type=ArrayType(
+                                 inner_type=ValueType(name='float'),
+                                 extents=(4,),
+                             ),
+                             doc='active bounding volume',
+                         ),
                      ),
                  ),
                  doc='',
@@ -4154,7 +4175,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  name='timer',
                  type=ArrayType(
                      inner_type=ValueType(name='mjTimerStat'),
-                     extents=(16,),
+                     extents=(15,),
                  ),
                  doc='timer statistics',
              ),

mjx/MANIFEST.in

@@ -1,2 +1 @@
 recursive-include mujoco/mjx/test_data *
-recursive-include mujoco/mjx/benchmark *.obj *.stl *.xml

mjx/mujoco/mjx/__init__.py

@@ -16,6 +16,7 @@
 
 # pylint:disable=g-importing-member
 from mujoco.mjx._src.collision_driver import collision
+from mujoco.mjx._src.collision_driver import ncon
 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
@@ -29,6 +30,7 @@ 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 get_data_into
 from mujoco.mjx._src.io import make_data
 from mujoco.mjx._src.io import put_data
 from mujoco.mjx._src.io import put_model
@@ -39,8 +41,10 @@ from mujoco.mjx._src.smooth import com_vel
 from mujoco.mjx._src.smooth import crb
 from mujoco.mjx._src.smooth import factor_m
 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.support import full_m
+from mujoco.mjx._src.support import is_sparse
+from mujoco.mjx._src.support import mul_m
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/collision_driver.py

@@ -313,8 +313,16 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
 
   body_pairs = []
   exclude_signature = set(m.exclude_signature)
+  geom_con = m.geom_contype | m.geom_conaffinity
+  b_start = m.body_geomadr
+  b_end = b_start + m.body_geomnum
+
   for b1 in range(m.nbody):
+    if not geom_con[b_start[b1]:b_end[b1]].any():
+      continue
     for b2 in range(b1, m.nbody):
+      if not geom_con[b_start[b2]:b_end[b2]].any():
+        continue
       signature = (b1 << 16) + (b2)
       if signature in exclude_signature:
         continue
@@ -323,12 +331,12 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
       body_pairs.append((b1, b2))
 
   for b1, b2 in body_pairs:
-    start1 = m.body_geomadr[b1]
+    for g1 in range(b_start[b1], b_end[b1]):
-    end1 = m.body_geomadr[b1] + m.body_geomnum[b1]
+      if not geom_con[g1]:
-    for g1 in range(start1, end1):
+        continue
-      start2 = m.body_geomadr[b2]
+      for g2 in range(b_start[b2], b_end[b2]):
-      end2 = m.body_geomadr[b2] + m.body_geomnum[b2]
+        if not geom_con[g2]:
-      for g2 in range(start2, end2):
+          continue
         mask = m.geom_contype[g1] & m.geom_conaffinity[g2]
         mask |= m.geom_contype[g2] & m.geom_conaffinity[g1]
         if mask != 0:

mjx/mujoco/mjx/_src/constraint.py

@@ -315,7 +315,7 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
 
 
 def count_constraints(
-    m: Union[Model, mujoco.MjModel]
+    m: Union[Model, mujoco.MjModel], d: Optional[Data] = None
 ) -> Tuple[int, int, int, int]:
   """Returns equality, friction, limit, and contact constraint counts."""
   if m.opt.disableflags & DisableBit.CONSTRAINT:
@@ -336,7 +336,10 @@ def count_constraints(
   else:
     nl = int(m.jnt_limited.sum())
 
-  nc = collision_driver.ncon(m) * 4
+  if d is None:
+    nc = collision_driver.ncon(m) * 4
+  else:
+    nc = d.efc_J.shape[-2] - ne - nf - nl
 
   return ne, nf, nl, nc
 

mjx/mujoco/mjx/_src/dataclasses.py

@@ -45,6 +45,7 @@ def dataclass(clz: _T) -> _T:
         field.type is jax.Array,
         dataclasses.is_dataclass(field.type),
         jax.Array in typing.get_args(field.type),
+        any(dataclasses.is_dataclass(a) for a in typing.get_args(field.type)),
     )):
       data_fields.append(field)
     else:
@@ -57,12 +58,13 @@ def dataclass(clz: _T) -> _T:
   data_clz.replace = replace
 
   def iterate_clz_with_keys(x):
-    # numpy arrays are not hashable, so convert them to tuples for jit cache
-    to_tup = lambda x: tuple(x) if len(x.shape) == 1 else tuple(map(to_tup, x))
-
     def to_meta(field, obj):
       val = getattr(obj, field.name)
-      return (to_tup(val), val.dtype) if isinstance(val, np.ndarray) else val
+      # numpy arrays are not hashable so return raw bytes instead
+      if isinstance(val, np.ndarray):
+        return (val.tobytes(), val.dtype, val.shape)
+      else:
+        return val
 
     def to_data(field, obj):
       return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name))
@@ -75,7 +77,8 @@ def dataclass(clz: _T) -> _T:
 
     def from_meta(field, meta):
       if field.type is np.ndarray:
-        return (field.name, np.array(meta[0], dtype=meta[1]))
+        arr = np.frombuffer(meta[0], dtype=meta[1]).reshape(meta[2])
+        return (field.name, arr)
       else:
         return (field.name, meta)
 

mjx/mujoco/mjx/_src/device.py

@@ -164,6 +164,9 @@ def _validate(m: mujoco.MjModel):
     if f & m.opt.enableflags:
       warnings.warn(f'Ignoring enable flag {f.name}.')
 
+  if not np.allclose(m.dof_frictionloss, 0):
+    raise NotImplementedError('dof_frictionloss is not implemented.')
+
 
 @overload
 def device_put(value: mujoco.MjData) -> types.Data:

mjx/mujoco/mjx/_src/device_test.py

@@ -59,7 +59,7 @@ def _assert_eq(testcase, a, b, attr=None, name=None):
 
 class DeviceTest(parameterized.TestCase):
 
-  @parameterized.parameters(test_util.TEST_FILES)
+  @parameterized.parameters('constraints.xml', 'pendula.xml')
   def testdevice_put(self, fname):
     """Test putting MjData and MjModel on device."""
     m = test_util.load_test_file(fname)
@@ -71,20 +71,22 @@ class DeviceTest(parameterized.TestCase):
     _assert_eq(self, mjx.device_put(d), d)
     _assert_eq(self, mjx.device_put(m), m)
 
-  @parameterized.parameters(test_util.TEST_FILES)
+  @parameterized.parameters('constraints.xml', 'pendula.xml')
   def testdevice_get(self, fname):
     """Test getting MjData from a device."""
     m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
     mx = device.device_put(m)
     dx = mjx.make_data(mx)
     d = mujoco.MjData(m)
     device.device_get_into(d, dx)
     _assert_eq(self, dx, d)
 
-  @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'})
+  @parameterized.parameters('constraints.xml', 'pendula.xml')
   def testdevice_get_batched(self, fname):
     """Test getting MjData from a device."""
     m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
     mx = device.device_put(m)
     batch_size = 32
 

mjx/mujoco/mjx/_src/forward.py

@@ -66,7 +66,7 @@ def fwd_position(m: Model, d: Data) -> Data:
   d = smooth.kinematics(m, d)
   d = smooth.com_pos(m, d)
   d = smooth.crb(m, d)
-  d = smooth.factor_m(m, d, d.qM)
+  d = smooth.factor_m(m, d)
   d = collision_driver.collision(m, d)
   d = constraint.make_constraint(m, d)
   d = smooth.transmission(m, d)
@@ -288,8 +288,8 @@ def euler(m: Model, d: Data) -> Data:
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     # TODO(robotics-simulation): can this be done with a smaller perf hit
-    mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
+    dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
-    dh = smooth.factor_m(m, d, mh)
+    dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
   return _advance(m, d, d.act_dot, qacc)

mjx/mujoco/mjx/_src/io.py

@@ -23,8 +23,10 @@ import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
 from mujoco.mjx._src import mesh
+from mujoco.mjx._src import support
 from mujoco.mjx._src import types
 import numpy as np
+import scipy
 
 
 def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
@@ -35,6 +37,9 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
   if o.cone not in set(types.ConeType):
     raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
 
+  if o.jacobian not in set(types.JacobianType):
+    raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}')
+
   if o.solver not in set(types.SolverType):
     raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
 
@@ -49,6 +54,7 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
   }
   static_fields['integrator'] = types.IntegratorType(o.integrator)
   static_fields['cone'] = types.ConeType(o.cone)
+  static_fields['jacobian'] = types.JacobianType(o.jacobian)
   static_fields['solver'] = types.SolverType(o.solver)
   static_fields['disableflags'] = types.DisableBit(o.disableflags)
 
@@ -103,6 +109,9 @@ def put_model(m: mujoco.MjModel, device=None) -> types.Model:
           f'{[mj_type(m) for m in missing]} not supported'
       )
 
+  if not np.allclose(m.dof_frictionloss, 0):
+    raise NotImplementedError('dof_frictionloss is not implemented.')
+
   opt = _put_option(m.opt, device=device)
   stat = _put_statistic(m.stat, device=device)
 
@@ -137,8 +146,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   ne, nf, nl, nc = constraint.count_constraints(m)
   nefc = ne + nf + nl + nc
 
+  zero_0 = jp.zeros(0, dtype=jp.float32)
   zero_nv = jp.zeros(m.nv, dtype=jp.float32)
   zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
+  zero_nv_nv = jp.zeros((m.nv, m.nv), 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)
@@ -180,10 +191,9 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       actuator_length=zero_nu,
       actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
       crb=zero_nbody_10,
-      qM=zero_nm,
+      qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLD=zero_nm,
+      qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLDiagInv=zero_nv,
+      qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
-      qLDiagSqrtInv=zero_nv,
       contact=types.Contact.zero(ncon),
       efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
       efc_frictionloss=zero_nefc,
@@ -226,10 +236,35 @@ 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)
+
+  if batched:
+    result = [mujoco.MjData(m) for _ in range(batch_size)]
+  else:
+    result = mujoco.MjData(m)
+
+  get_data_into(result, m, d)
+
+  return result
+
+
+def get_data_into(
+    result: Union[mujoco.MjData, List[mujoco.MjData]],
+    m: mujoco.MjModel,
+    d: types.Data,
+):
+  """Gets mjx.Data from a device into an existing mujoco.MjData or list."""
+  batched = isinstance(result, list)
+  if batched and len(d.qpos.shape) < 2:
+    raise ValueError('dst is a list, but d is not batched.')
+  if not batched and len(d.qpos.shape) >= 2:
+    raise ValueError('dst is a an MjData, but d is batched.')
+
+  d = jax.device_get(d)
+
+  batch_size = d.qpos.shape[0] if batched else 1
+  ne, nf, nl, nc = constraint.count_constraints(m, d)
   efc_type = np.array([
       mujoco.mjtConstraint.mjCNSTR_EQUALITY,
       mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
@@ -237,26 +272,34 @@ def get_data(
       mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
   ]).repeat([ne, nf, nl, nc])
 
-  ds = []
+  dof_i, dof_j = [], []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      dof_i.append(i)
+      dof_j.append(j)
+      j = m.dof_parentid[j]
+
   for i in range(batch_size):
-    dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d
+    d_i = jax.tree_map(lambda x, i=i: x[i], d) if batched else d
-    ncon = (dx_i.contact.dist <= 0).sum()
+    result_i = result[i] if batched else result
-    efc_active = (dx_i.efc_J != 0).any(axis=1)
+    ncon = (d_i.contact.dist <= 0).sum()
+    efc_active = (d_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()
+    nefc, nc = int(efc_active.sum()), int((efc_active & efc_con).sum())
-    d_i = mujoco.MjData(m)
+    result_i.nnzJ = nefc * m.nv
-    d_i.nnzJ = nefc * m.nv
+    if ncon != result_i.ncon or nefc != result_i.nefc:
-    mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
+      mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
-    d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
+    result_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
-    d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
+    result_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
-    d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
+    result_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)
+        _get_contact(result_i.contact, d_i.contact, nefc - nc)
         continue
 
-      value = getattr(dx_i, field.name)
+      value = getattr(d_i, field.name)
 
       if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
         value = value.reshape((-1, 9))
@@ -267,15 +310,21 @@ def get_data(
       if field.name == 'efc_J':
         value = value[efc_active].reshape(-1)
 
+      if field.name == 'qM' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLD' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLDiagInv' and not support.is_sparse(m):
+        value = np.ones(m.nv)
+
       if value.shape:
-        getattr(d_i, field.name)[:] = value
+        getattr(result_i, field.name)[:] = value
       else:
-        setattr(d_i, field.name, value)
+        setattr(result_i, field.name, value)
 
-    d_i.efc_type[:] = efc_type[efc_active]
+    result_i.efc_type[:] = efc_type[efc_active]
-    ds.append(d_i)
-
-  return ds if batched else ds[0]
 
 
 def _put_contact(
@@ -346,6 +395,18 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
     fields[fname] = value
 
+  # convert qM and qLD if jacobian is dense
+  if not support.is_sparse(m):
+    fields['qM'] = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, fields['qM'], d.qM)
+    # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
+    try:
+      fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM'])
+    except scipy.linalg.LinAlgError:
+      # this happens when qM is empty or unstable simulation
+      fields['qLD'] = np.zeros((m.nv, m.nv))
+    fields['qLDiagInv'] = np.zeros(0)
+
   fields = jax.device_put(fields, device=device)
   fields['contact'] = _put_contact(d.contact, ncon, device=device)
 

mjx/mujoco/mjx/_src/io_test.py

@@ -25,7 +25,7 @@ import numpy as np
 
 _MULTIPLE_CONVEX_OBJECTS = """
   <mujoco>
-    <option timestep="0.001"/>
+    <option timestep="0.001" jacobian="dense"/>
     <default>
       <geom solref=".006 1"/>
     </default>
@@ -224,7 +224,6 @@ class DataIOTest(parameterized.TestCase):
     ncon = 46
     nv = 19
     nefc = 185
-    nm = 64
 
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
@@ -251,10 +250,9 @@ class DataIOTest(parameterized.TestCase):
     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.qM.shape, (nv, nv))
-    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d.qLD.shape, (nv, nv))
-    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d.qLDiagInv.shape, (0,))
-    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))
@@ -291,7 +289,11 @@ class DataIOTest(parameterized.TestCase):
     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)
+
+    # check that qM is transformed properly
+    qm = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
 
     # 4 contacts, 2 for each capsule against the plane
     self.assertEqual(dx.contact.dist.shape, (4,))
@@ -335,8 +337,22 @@ class DataIOTest(parameterized.TestCase):
     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)
+    dx_sparse = mjx.put_data(m, d)
-    np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
+
+    # check sparse mass matrices are correct
+    np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
+
+    # check dense mass matrices are correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_from_dense = mjx.put_data(m, d)
+    qm = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
 
   def test_get_data(self):
     """Test that get_data makes correct MjData."""
@@ -396,6 +412,27 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(ds[0].ncon, 1)
     self.assertEqual(ds[1].ncon, 0)
 
+  def test_get_data_into(self):
+    """Test that get_data_into correctly populates an MjData."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2 = mujoco.MjData(m)
+    mjx.get_data_into(d_2, 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.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)
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/math.py

@@ -20,6 +20,30 @@ import jax
 from jax import numpy as jp
 
 
+def matmul_unroll(a: jax.Array, b: jax.Array) -> jax.Array:
+  """Calculates a @ b via explicit cell value operations.
+
+  This is faster than XLA matmul for small matrices (e.g. 3x3, 4x4).
+
+  Args:
+    a: left hand of matmul operand
+    b: right hand of matmul operand
+  Returns:
+    the matrix product of the inputs.
+  """
+  c = []
+  for i in range(a.shape[0]):
+    row = []
+    for j in range(b.shape[1]):
+      s = 0.0
+      for k in range(a.shape[1]):
+        s += a[i, k] * b[k, j]
+      row.append(s)
+    c.append(row)
+
+  return jp.array(c)
+
+
 def norm(
     x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None
 ) -> jax.Array:

mjx/mujoco/mjx/_src/ray.py

@@ -19,6 +19,7 @@ from typing import Sequence, Tuple
 import jax
 from jax import numpy as jp
 import mujoco
+from mujoco.mjx._src import math
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
 from mujoco.mjx._src.types import GeomType
@@ -129,14 +130,79 @@ def _ray_box(
   return jp.min(jp.where(valid, x, jp.inf))
 
 
-def _ray_mesh(
+def _ray_triangle(
-    size: jax.Array,
+    vert: jax.Array,
     pnt: jax.Array,
     vec: jax.Array,
+    b0: jax.Array,
+    b1: jax.Array,
 ) -> jax.Array:
-  """Returns the distance at which a ray intersects with a mesh."""
+  """Returns the distance at which a ray intersects with a triangle."""
-  del size, pnt, vec
+  # project difference vectors in ray normal plane
-  raise NotImplementedError("ray <> mesh not implemented yet")
+  planar = jp.dot(jp.array([b0, b1]), (vert - pnt).T)
+
+  # determine if origin is inside planar projection of triangle
+  # A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  A = jp.array(  # pylint: disable=invalid-name
+      [planar[:, 0] - planar[:, 2], planar[:, 1] - planar[:, 2]]
+  ).T.flatten()
+  b = -planar[:, 2]
+  det = A[0] * A[3] - A[1] * A[2]
+  valid = jp.abs(det) >= mujoco.mjMINVAL
+
+  t0 = (A[3] * b[0] - A[1] * b[1]) / det
+  t1 = (-A[2] * b[0] + A[0] * b[1]) / det
+  valid &= (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1)
+
+  # intersect ray with plane of triangle
+  nrm = jp.cross(vert[0] - vert[2], vert[1] - vert[2])
+  denom = jp.dot(vec, nrm)
+  valid &= jp.abs(denom) >= mujoco.mjMINVAL
+
+  dist = jp.where(valid, -jp.dot(pnt - vert[2], nrm) / denom, jp.inf)
+
+  return dist
+
+
+def _ray_mesh(
+    m: Model,
+    geom_id: np.ndarray,
+    unused_size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns the best distance and geom_id for ray mesh intersections."""
+  data_id = m.geom_dataid[geom_id]
+
+  ray_basis = lambda x: math.orthogonals(math.normalize(x))
+  b0, b1 = jax.vmap(ray_basis)(vec)
+
+  faceadr = np.append(m.mesh_faceadr, m.nmeshface)
+  vertadr = np.append(m.mesh_vertadr, m.nmeshvert)
+
+  dists = []
+  for i, id_ in enumerate(data_id):
+    face = m.mesh_face[faceadr[id_] : faceadr[id_ + 1]]
+    vert = m.mesh_vert[vertadr[id_] : vertadr[id_ + 1]]
+    dist = jax.vmap(_ray_triangle, in_axes=(0, None, None, None, None))(
+        vert[face], pnt[i], vec[i], b0[i], b1[i]
+    )
+    dists.append(dist)
+
+  # map the triangle id to data id
+  tri_id = np.append(0, (faceadr[data_id + 1] - faceadr[data_id]).cumsum())
+  tri_data_id = np.zeros(tri_id[-1], dtype=np.int32)
+  tri_data_id[tri_id[:-1]] = 1
+  tri_data_id = tri_data_id.cumsum() - 1
+
+  dists = jp.concatenate(dists)
+  min_id = jp.argmin(dists)
+  # Grab the best distance amongst all meshes, bypassing the argmin in `ray`.
+  # This avoids having to compute the best distance per mesh.
+  dist = dists[min_id, None]
+  id_ = jp.array(geom_id)[jp.array(tri_data_id)[min_id], None]
+
+  return dist, id_
 
 
 _RAY_FUNC = {
@@ -144,7 +210,7 @@ _RAY_FUNC = {
     GeomType.SPHERE: _ray_sphere,
     GeomType.CAPSULE: _ray_capsule,
     GeomType.BOX: _ray_box,
-    # GeomType.MESH: _ray_mesh,
+    GeomType.MESH: _ray_mesh,
 }
 
 
@@ -192,8 +258,13 @@ def ray(
     if id_.size == 0:
       continue
 
-    size, pnt, vec = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
+    args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
-    dist = jax.vmap(fn)(size, pnt, vec)
+
+    if geom_type == GeomType.MESH:
+      dist, id_ = fn(m, id_, *args)
+    else:
+      dist = jax.vmap(fn)(*args)
+
     dists, ids = dists + [dist], ids + [id_]
 
   if not ids:

mjx/mujoco/mjx/_src/ray_test.py

@@ -144,6 +144,43 @@ class RayTest(absltest.TestCase):
     mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
     _assert_eq(dist, mj_dist, 'dist')
 
+  def test_ray_mesh(self):
+    """Tests MJX ray<>mesh 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)
+
+    # look at the tetrahedron
+    pnt, vec = jp.array([2.0, 2.0, 2.0]), -jp.array([
+        1.0,
+        1.0,
+        1.0,
+    ])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 4, 'geom_id')
+
+    pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid)
+    _assert_eq(geomid, 4, 'geom_id')
+    _assert_eq(dist, mj_dist, 'dist-tetrahedron')
+
+    # look at the dodecahedron
+    pnt, vec = jp.array([4.0, 2.0, 2.0]), -jp.array([
+        2.0,
+        1.0,
+        1.0,
+    ])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 5, 'geom_id')
+
+    pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid)
+    _assert_eq(geomid, 5, 'geom_id')
+    _assert_eq(dist, mj_dist, 'dist-dodecahedron')
+
   def test_ray_geomgroup(self):
     """Tests ray geomgroup filter."""
     m = test_util.load_test_file('ray.xml')

mjx/mujoco/mjx/_src/smooth.py

@@ -136,20 +136,23 @@ def com_pos(m: Model, d: Data) -> Data:
   pos, mass = scan.body_tree(
       m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True
   )
-  cond = jp.tile(mass < jp.array(mujoco.mjMINVAL), (3, 1)).T
+  cond = jp.tile(mass < mujoco.mjMINVAL, (3, 1)).T
-  subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass))
+  # take maximum to avoid NaN in gradient of jp.where
+  subtree_com = jax.vmap(jp.divide)(pos, jp.maximum(mass, mujoco.mjMINVAL))
+  subtree_com = jp.where(cond, d.xipos, subtree_com)
   d = d.replace(subtree_com=subtree_com)
 
   # map inertias to frame centered at subtree_com
   @jax.vmap
   def inert_com(inert, ximat, off, mass):
     h = jp.cross(off, -jp.eye(3))
-    inert = ximat @ jp.diag(inert) @ ximat.T + h @ h.T * mass
+    inert = math.matmul_unroll((ximat * inert), ximat.T)
+    inert += math.matmul_unroll(h, h.T) * mass
     # cinert is triu(inert), mass * off, mass
-    inert = inert[(jp.array([0, 1, 2, 0, 0, 1]), jp.array([0, 1, 2, 1, 2, 2]))]
+    inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
-    return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)])
+    return jp.concatenate([inert, off * mass, mass[None]])
 
-  root_com = subtree_com[jp.array(m.body_rootid)]
+  root_com = subtree_com[m.body_rootid]
   offset = d.xipos - root_com
   cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
   d = d.replace(cinert=cinert)
@@ -208,36 +211,21 @@ def crb(m: Model, d: Data) -> Data:
   crb_body = crb_body.at[0].set(0.0)
   d = d.replace(crb=crb_body)
 
-  # TODO(erikfrey): do centralized take fn?
   crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
   crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
-
+  qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
-  dof_i, dof_j, diag = [], [], []
-  for i in range(m.nv):
-    diag.append(len(dof_i))
-    j = i
-    while j > -1:
-      dof_i, dof_j = dof_i + [i], dof_j + [j]
-      j = m.dof_parentid[j]
-
-  crb_codf_i = jp.take(crb_cdof, jp.array(dof_i), axis=0)
-  cdof_j = jp.take(d.cdof, jp.array(dof_j), axis=0)
-  qm = jax.vmap(jp.dot)(crb_codf_i, cdof_j)
-
-  # add armature to diagonal
-  qm = qm.at[jp.array(diag)].add(m.dof_armature)
-
   d = d.replace(qM=qm)
 
   return d
 
 
-def factor_m(
+def factor_m(m: Model, d: Data) -> Data:
-    m: Model,
+  """Gets factorizaton of inertia-like matrix M, assumed spd."""
-    d: Data,
+
-    qM: jax.Array,  # pylint:disable=invalid-name
+  if not support.is_sparse(m):
-) -> Data:
+    qh, _ = jax.scipy.linalg.cho_factor(d.qM)
-  """Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd."""
+    d = d.replace(qLD=qh)
+    return d
 
   # build up indices for where we will do backwards updates over qLD
   # TODO(erikfrey): do fewer updates by combining non-overlapping ranges
@@ -255,7 +243,7 @@ def factor_m(
       madr_j_range = tuple(m.dof_Madr[j : j + 2])
       updates.setdefault(madr_j_range, []).append((madr_d, madr_ij))
 
-  qld = qM
+  qld = d.qM
 
   for (out_beg, out_end), vals in sorted(updates.items(), reverse=True):
     madr_d, madr_ij = jp.array(vals).T
@@ -281,6 +269,9 @@ def factor_m(
 def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   """Computes sparse backsubstitution:  x = inv(L'*D*L)*y ."""
 
+  if not support.is_sparse(m):
+    return jax.scipy.linalg.cho_solve((d.qLD, False), x)
+
   updates_i, updates_j = {}, {}
   for i in range(m.nv):
     madr_ij, j = m.dof_Madr[i], i
@@ -307,52 +298,6 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   return x
 
 
-def dense_m(m: Model, d: Data) -> jax.Array:
-  """Reconstitute dense mass matrix from qM."""
-
-  is_, js, madr_ijs = [], [], []
-  for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
-
-    while True:
-      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
-      if j == -1:
-        break
-      is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij]
-
-  i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
-
-  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM[madr_ij])
-
-  # diagonal, upper triangular, lower triangular
-  mat = jp.diag(d.qM[jp.array(m.dof_Madr)]) + mat + mat.T
-
-  return mat
-
-
-def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
-  """Multiply vector by inertia matrix."""
-
-  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
-
-  is_, js, madr_ijs = [], [], []
-  for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
-
-    while True:
-      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
-      if j == -1:
-        break
-      is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij]
-
-  i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
-
-  out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j])
-  out = out.at[j].add(d.qM[madr_ij] * vec[i])
-
-  return out
-
-
 def com_vel(m: Model, d: Data) -> Data:
   """Computes cvel, cdof_dot."""
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -16,7 +16,6 @@
 
 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
@@ -49,6 +48,8 @@ class SmoothTest(absltest.TestCase):
     """Tests MJX smooth functions match MuJoCo smooth functions."""
 
     m = test_util.load_test_file('pendula.xml')
+    # # force MJX sparse for testing:
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     # give the system a little kick to ensure we have non-identity rotations
     d.qvel = np.random.random(m.nv)
@@ -79,8 +80,7 @@ class SmoothTest(absltest.TestCase):
     _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, 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
@@ -95,21 +95,6 @@ class SmoothTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'actuator_length')
     _assert_attr_eq(d, dx, 'actuator_moment')
 
-  def test_mul_m(self):
-    m = test_util.load_test_file('pendula.xml')
-    d = mujoco.MjData(m)
-    # give the system a little kick to ensure we have non-identity rotations
-    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')
-
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""
         <mujoco>

mjx/mujoco/mjx/_src/solver.py

@@ -14,14 +14,13 @@
 # ==============================================================================
 """Constraint solvers."""
 
-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
+from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import Data
@@ -39,7 +38,6 @@ class _Context(PyTreeNode):
     qfrc_constraint: constraint force (from Data)     (nv,)
     Jaref: Jac*qacc - aref                            (nefc,)
     efc_force: constraint force in constraint space   (nefc,)
-    M: dense mass matrix, populated for nv < 100      (nv, nv)
     Ma: M*qacc                                        (nv,)
     grad: gradient of master cost                     (nv,)
     Mgrad: M / grad                                   (nv,)
@@ -54,7 +52,6 @@ class _Context(PyTreeNode):
   qfrc_constraint: jax.Array
   Jaref: jax.Array  # pylint: disable=invalid-name
   efc_force: jax.Array
-  M: Optional[jax.Array]
   Ma: jax.Array  # pylint: disable=invalid-name
   grad: jax.Array
   Mgrad: jax.Array  # pylint: disable=invalid-name
@@ -68,15 +65,13 @@ class _Context(PyTreeNode):
   def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context':
     jaref = d.efc_J @ d.qacc - d.efc_aref
     # 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 = support.mul_m(m, d, d.qacc)
-    ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
     nv_0 = jp.zeros(m.nv)
     ctx = _Context(
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
         efc_force=d.efc_force,
-        M=M,
         Ma=ma,
         grad=nv_0,
         Mgrad=nv_0,
@@ -224,10 +219,10 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
   elif m.opt.solver == SolverType.NEWTON:
     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 = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T)
-    h = smooth.dense_m(m, d) + h
+    dh = d.replace(qM=h)
-    h_ = jax.scipy.linalg.cho_factor(h)
+    dh = smooth.factor_m(m, dh)
-    mgrad = jax.scipy.linalg.cho_solve(h_, grad)
+    mgrad = smooth.solve_m(m, dh, grad)
   else:
     raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
 
@@ -255,7 +250,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   gtol = m.opt.tolerance * m.opt.ls_tolerance * smag
 
   # compute Mv, Jv
-  mv = smooth.mul_m(m, d, ctx.search) if ctx.M is None else ctx.M @ ctx.search
+  mv = support.mul_m(m, d, ctx.search)
   jv = d.efc_J @ ctx.search
 
   # prepare quadratics

mjx/mujoco/mjx/_src/solver_test.py

@@ -39,37 +39,65 @@ def _assert_attr_eq(a, b, attr):
 
 class SolverTest(absltest.TestCase):
 
-  def test_solver(self):
+  def test_newton(self):
-    """Test solver."""
+    """Test newton solver."""
     m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(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)
-
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
     _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')
 
-    # also test normal CG
+  def test_cg(self):
+    """Test CG solver."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
     m.opt.solver = mujoco.mjtSolver.mjSOL_CG
     mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
     _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')
 
-    # without warmstart, the solution is not as close
+  def test_no_warmstart(self):
-    m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
+    """Test no warmstart."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
     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))
+    nnz = dx.efc_J.any(axis=1)
+    # without warmstart, the solution is not as close
     _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2)
 
+  def test_dense(self):
+    """Test solver works with dense mass matrices."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
+    _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')
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/support.py

@@ -14,17 +14,124 @@
 # ==============================================================================
 """Engine support functions."""
 
-from typing import Tuple
+from typing import Optional, Tuple, Union
 
 import jax
 from jax import numpy as jp
+import mujoco
 from mujoco.mjx._src import scan
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import JacobianType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
 
+def is_sparse(m: Union[mujoco.MjModel, Model]) -> bool:
+  """Return True if this model should create sparse mass matrices.
+
+  Args:
+    m: a MuJoCo or MJX model
+
+  Returns:
+    True if provided model should create sparse mass matrices
+
+  Modern TPUs have specialized hardware for rapidly operating over sparse
+  matrices, whereas GPUs tend to be faster with dense matrices as long as they
+  fit onto the device.  As such, the default behavior in MJX (via
+  ``JacobianType.AUTO``) is sparse if ``nv`` is >= 60 or MJX detects a TPU as
+  the default backend, otherwise dense.
+  """
+  # AUTO is a rough heuristic - you may see better performance for your workload
+  # and compute by explicitly setting jacobian to dense or sparse
+  if m.opt.jacobian == JacobianType.AUTO:
+    return m.nv >= 60 or jax.default_backend() == 'tpu'
+  return m.opt.jacobian == JacobianType.SPARSE
+
+
+def make_m(
+    m: Model, a: jax.Array, b: jax.Array, d: Optional[jax.Array] = None
+) -> jax.Array:
+  """Computes M = a @ b.T + diag(d)."""
+
+  ij = []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      ij.append((i, j))
+      j = m.dof_parentid[j]
+
+  i, j = (jp.array(x) for x in zip(*ij))
+
+  if not is_sparse(m):
+    qm = a @ b.T
+    if d is not None:
+      qm += jp.diag(d)
+    mask = jp.zeros((m.nv, m.nv), dtype=bool).at[(i, j)].set(True)
+    qm = qm * mask
+    qm = qm + jp.tril(qm, -1).T
+    return qm
+
+  a_i = jp.take(a, i, axis=0)
+  b_j = jp.take(b, j, axis=0)
+  qm = jax.vmap(jp.dot)(a_i, b_j)
+
+  # add diagonal
+  if d is not None:
+    qm = qm.at[m.dof_Madr].add(d)
+
+  return qm
+
+
+def full_m(m: Model, d: Data) -> jax.Array:
+  """Reconstitute dense mass matrix from qM."""
+
+  if not is_sparse(m):
+    return d.qM
+
+  ij = []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      ij.append((i, j))
+      j = m.dof_parentid[j]
+
+  i, j = (jp.array(x) for x in zip(*ij))
+
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
+
+  # also set upper triangular
+  mat = mat + jp.tril(mat, -1).T
+
+  return mat
+
+
+def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
+  """Multiply vector by inertia matrix."""
+
+  if not is_sparse(m):
+    return d.qM @ vec
+
+  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
+
+  is_, js, madr_ijs = [], [], []
+  for i in range(m.nv):
+    madr_ij, j = m.dof_Madr[i], i
+
+    while True:
+      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
+      if j == -1:
+        break
+      is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij]
+
+  i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
+
+  out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j])
+  out = out.at[j].add(d.qM[madr_ij] * vec[i])
+
+  return out
+
+
 def jac(
     m: Model, d: Data, point: jax.Array, body_id: jax.Array
 ) -> Tuple[jax.Array, jax.Array]:

mjx/mujoco/mjx/_src/support_test.py

@@ -27,7 +27,47 @@ import numpy as np
 
 class SupportTest(parameterized.TestCase):
 
-  @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'})
+  def test_mul_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # first test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    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)
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+    # also check dense
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+  def test_full_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # for the model to be sparse to exercise MJX full_M
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_full_m = jax.jit(support.full_m)(mx, dx)
+    mj_full_m = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, mj_full_m, d.qM)
+    np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5)
+
+  @parameterized.parameters('constraints.xml', 'pendula.xml')
   def test_jac(self, fname):
     np.random.seed(0)
 

mjx/mujoco/mjx/_src/test_util.py

@@ -15,20 +15,13 @@
 """Utilities for testing."""
 
 import sys
-from typing import Dict, List, Tuple
+from typing import Dict, Tuple
 from xml.etree import ElementTree as ET
 
 from etils import epath
 import mujoco
 import numpy as np
 
-TEST_FILES: List[str] = [
-    'constraints.xml',
-    'convex.xml',
-    'pendula.xml',
-    'ray.xml',
-]
-
 _ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
 _DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact']
 _DYN_PRMS = ['0.189', '2.1']

mjx/mujoco/mjx/_src/test_util_test.py

@@ -1,38 +0,0 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Tests for the test_util."""
-
-from absl.testing import absltest
-from etils import epath
-from mujoco.mjx._src import test_util
-
-
-class TestUtilTest(absltest.TestCase):
-
-  def test_files_in_test_data_match(self):
-    directory = epath.resource_path('mujoco.mjx') / 'test_data'
-    files = set([f.name for f in directory.glob('*.xml')])
-    self.assertSetEqual(
-        files,
-        set(test_util.TEST_FILES),
-        msg=(
-            '`_test_util.TEST_FILES` must match the files in the '
-            'test_data/*.xml directory'
-        ),
-    )
-
-
-if __name__ == '__main__':
-  absltest.main()

mjx/mujoco/mjx/_src/types.py

@@ -123,6 +123,19 @@ class ConeType(enum.IntEnum):
   # unsupported: ELLIPTIC
 
 
+class JacobianType(enum.IntEnum):
+  """Type of constraint Jacobian.
+
+  Attributes:
+    DENSE: dense
+    SPARSE: sparse
+    AUTO: sparse if nv>60 and device is TPU, dense otherwise
+  """
+  DENSE = mujoco.mjtJacobian.mjJAC_DENSE
+  SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
+  AUTO = mujoco.mjtJacobian.mjJAC_AUTO
+
+
 class SolverType(enum.IntEnum):
   """Constraint solver algorithm.
 
@@ -215,6 +228,10 @@ class Option(PyTreeNode):
       nonzero. Not used by mj
     integrator:       integration mode
     cone:             type of friction cone
+    jacobian:         matrix layout for mass matrices (dense or sparse)
+                      (note that this is different from MuJoCo, where jacobian
+                      specifies whether efc_J and its accompanying matrices
+                      are dense or sparse.
     solver:           solver algorithm
     iterations:       number of main solver iterations
     ls_iterations:    maximum number of CG/Newton linesearch iterations
@@ -232,7 +249,7 @@ class Option(PyTreeNode):
   # unsupported: magnetic, o_margin, o_solref, o_solimp
   integrator: IntegratorType
   cone: ConeType
-  # unsupported: jacobian
+  jacobian: JacobianType
   solver: SolverType
   iterations: int
   ls_iterations: int
@@ -264,6 +281,8 @@ class Model(PyTreeNode):
     ngeom: number of geoms
     nsite: number of sites
     nmesh: number of meshes
+    nmeshvert: number of vertices in all meshes
+    nmeshface: number of triangular faces in all meshes
     nmat: number of materials
     npair: number of predefined geom pairs
     nexclude: number of excluded geom pairs
@@ -321,6 +340,7 @@ 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_dataid: id of geom's mesh/hfield; -1: none           (ngeom,)
     geom_group: group for visibility                          (ngeom,)
     geom_matid: material id for rendering                     (ngeom,)
     geom_priority: geom contact priority                      (ngeom,)
@@ -338,6 +358,10 @@ class Model(PyTreeNode):
     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)
+    mesh_vertadr: first vertex address                        (nmesh x 1)
+    mesh_faceadr: first face address                          (nmesh x 1)
+    mesh_vert: vertex positions for all meshes                (nmeshvert, 3)
+    mesh_face: vertex face data                               (nmeshface, 3)
     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,)
@@ -390,6 +414,8 @@ class Model(PyTreeNode):
   ngeom: int
   nsite: int
   nmesh: int
+  nmeshvert: int
+  nmeshface: int
   nmat: int
   npair: int
   nexclude: int
@@ -447,6 +473,7 @@ class Model(PyTreeNode):
   geom_conaffinity: np.ndarray
   geom_condim: np.ndarray
   geom_bodyid: np.ndarray
+  geom_dataid: np.ndarray
   geom_group: np.ndarray
   geom_matid: np.ndarray
   geom_priority: np.ndarray
@@ -463,6 +490,10 @@ class Model(PyTreeNode):
   site_bodyid: np.ndarray
   site_pos: jax.Array
   site_quat: jax.Array
+  mesh_vertadr: np.ndarray
+  mesh_faceadr: np.ndarray
+  mesh_vert: np.ndarray
+  mesh_face: np.ndarray
   mat_rgba: np.ndarray
   pair_dim: np.ndarray
   pair_geom1: np.ndarray
@@ -586,10 +617,12 @@ class Data(PyTreeNode):
     actuator_length: actuator lengths                             (nu,)
     actuator_moment: actuator moments                             (nu, nv)
     crb: com-based composite inertia and mass                     (nbody, 10)
-    qM: total inertia (sparse)                                    (nM,)
+    qM: total inertia                                  if sparse: (nM,)
-    qLD: L'*D*L factorization of M (sparse)                       (nM,)
+                                                       if dense:  (nv, nv)
-    qLDiagInv: 1/diag(D)                                          (nv,)
+    qLD: L'*D*L (or Cholesky) factorization of M.      if sparse: (nM,)
-    qLDiagSqrtInv: 1/sqrt(diag(D))                                (nv,)
+                                                       if dense:  (nv, nv)
+    qLDiagInv: 1/diag(D)                               if sparse: (nv,)
+                                                       if dense:  (0,)
     contact: list of all detected contacts                        (ncon,)
     efc_J: constraint Jacobian                                    (nefc, nv)
     efc_frictionloss: frictionloss (friction)                     (nefc,)
@@ -646,7 +679,6 @@ class Data(PyTreeNode):
   qM: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
-  qLDiagSqrtInv: jax.Array  # pylint:disable=invalid-name
   contact: Contact
   efc_J: jax.Array  # pylint:disable=invalid-name
   efc_frictionloss: jax.Array

mjx/mujoco/mjx/benchmark/__init__.py

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

mjx/mujoco/mjx/benchmark/benchmark.py

@@ -1,132 +0,0 @@
-# Copyright 2023 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Run benchmarks on various devices."""
-
-import sys
-import time
-
-from absl import flags
-from etils import epath
-import google_benchmark as benchmark
-import jax
-from jax import numpy as jp
-import mujoco
-from mujoco import mjx
-
-FLAGS = flags.FLAGS
-
-_PATHS = {
-    'humanoid': 'benchmark/model/humanoid/humanoid.xml',
-    'barkour': 'benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml',
-    'shadow_hand': 'benchmark/model/shadow_hand/scene_right.xml',
-}
-
-_BATCH_SIZE = {
-    ('barkour', 'tpu_v5e'): 1024,
-    ('humanoid', 'tpu_v5e'): 1024,
-    ('shadow_hand', 'tpu_v5e'): 1024,
-    ('barkour', 'gpu_a100'): 8192,
-    ('humanoid', 'gpu_a100'): 8192,
-    ('shadow_hand', 'gpu_a100'): 4096,
-    ('barkour', 'cpu'): 64,
-    ('humanoid', 'cpu'): 64,
-    ('shadow_hand', 'cpu'): 64,
-}
-
-_SOLVER_CONFIG = {
-    ('barkour', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 4, 6),
-    ('humanoid', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 6, 6),
-    ('shadow_hand', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 8, 6),
-    ('humanoid', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-    ('barkour', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-    ('shadow_hand', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-    ('barkour', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-    ('humanoid', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-    ('shadow_hand', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
-}
-
-
-flags.DEFINE_string('model', 'humanoid', 'Model to benchmark')
-flags.DEFINE_enum('device', 'cpu', ('cpu', 'tpu_v5e', 'gpu_a100'),
-                  'Device benchmark is running on')
-
-
-def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float:
-  """Reports jit time and op time for a function."""
-
-  step_count = 100 if FLAGS.device == 'cpu' else 1000
-
-  @jax.jit
-  def run_batch(seed: jp.ndarray):
-    rngs = jax.random.split(jax.random.PRNGKey(seed), batch_size)
-    init_state = jax.vmap(init_fn)(rngs)
-
-    @jax.vmap
-    def run(state):
-      def step(state, _):
-        state = step_fn(state)
-        return state, ()
-
-      return jax.lax.scan(step, state, (), length=step_count)
-
-    return run(init_state)
-
-  # run once to jit
-  beg = time.perf_counter()
-  jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(0))
-  first_t = time.perf_counter() - beg
-
-  times = []
-  while state:
-    beg = time.perf_counter()
-    batch = run_batch(jp.array(len(times)))
-    jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch)
-    times.append(time.perf_counter() - beg)
-
-  op_time = jp.mean(jp.array(times))
-  batch_sps = batch_size * step_count / op_time
-
-  state.counters['jit_time'] = first_t - op_time
-  state.counters['batch_sps'] = batch_sps
-
-
-@benchmark.option.unit(benchmark.kSecond)
-def _run(state: benchmark.State):
-  """Benchmark a model."""
-
-  f = epath.resource_path('mujoco.mjx') / _PATHS[FLAGS.model]
-  m = mujoco.MjModel.from_xml_path(f.as_posix())
-  m.opt.solver, m.opt.iterations, m.opt.ls_iterations = _SOLVER_CONFIG[
-      (FLAGS.model, FLAGS.device)
-  ]
-  m = mjx.device_put(m)
-
-  def init(rng):
-    d = mjx.make_data(m)
-    qvel = 0.01 * jax.random.normal(rng, shape=(m.nv,))
-    d = d.replace(qvel=qvel)
-    return d
-
-  def step(d):
-    return mjx.step(m, d)
-
-  batch_size = _BATCH_SIZE[(FLAGS.model, FLAGS.device)]
-  _measure_fn(state, init, step, batch_size=batch_size)
-
-
-if __name__ == '__main__':
-  FLAGS(sys.argv)
-  benchmark.register(_run, name=FLAGS.model + '_' + FLAGS.device)
-  benchmark.main()

mjx/mujoco/mjx/test_data/ray.xml

@@ -1,6 +1,7 @@
 <mujoco model="ray">
   <asset>
       <mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.4 0.4 0.4" />
+      <mesh name="dodecahedron" file="meshes/dodecahedron.stl" scale="0.04 0.04 0.04" />
       <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>
@@ -12,5 +13,6 @@
     <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"/>
+    <geom name="mesh2" pos="2 1 1" type="mesh" mesh="dodecahedron" rgba="1 0 1 1"/>
   </worldbody>
 </mujoco>

mjx/mujoco/mjx/test_data/shadow_hand/right_hand.xml

@@ -12,7 +12,7 @@
   <default>
     <default class="right_hand">
       <mesh scale="0.001 0.001 0.001"/>
-      <joint axis="1 0 0" damping="0.05" armature="0.0002" frictionloss="0.01"/>
+      <joint axis="1 0 0" damping="0.05" armature="0.0002"/>
       <position forcerange="-1 1"/>
 
       <default class="wrist">

mjx/mujoco/mjx/testspeed.py

@@ -0,0 +1,128 @@
+# 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.
+# ==============================================================================
+"""Run benchmarks on various devices."""
+
+import os
+import time
+from typing import Sequence, Tuple
+
+from absl import app
+from absl import flags
+from etils import epath
+import jax
+import mujoco
+from mujoco import mjx
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string('mjcf', None, 'path to model', required=True)
+flags.DEFINE_integer('nstep', 1000, 'number of steps per rollout')
+flags.DEFINE_integer('batch_size', 1024, 'number of parallel rollouts')
+flags.DEFINE_integer('unroll', 1, 'loop unroll length')
+flags.DEFINE_enum('solver', 'cg', ['cg', 'newton'], 'constraint solver')
+flags.DEFINE_integer('iterations', 1, 'number of solver iterations')
+flags.DEFINE_integer('ls_iterations', 4, 'number of linesearch iterations')
+flags.DEFINE_enum('output', 'text', ['text', 'tsv'], 'format to print results')
+
+
+def _measure(fn, *args) -> Tuple[float, float]:
+  """Reports jit time and op time for a function."""
+
+  beg = time.perf_counter()
+  compiled_fn = fn.lower(*args).compile()
+  end = time.perf_counter()
+  jit_time = end - beg
+
+  beg = time.perf_counter()
+  result = compiled_fn(*args)
+  jax.block_until_ready(result)
+  end = time.perf_counter()
+  run_time = end - beg
+
+  return jit_time, run_time
+
+
+def _main(argv: Sequence[str]):
+  """Benchmark a model."""
+
+  xla_flags = os.environ.get('XLA_FLAGS', '')
+  xla_flags += ' --xla_gpu_triton_gemm_any=True'
+  os.environ['XLA_FLAGS'] = xla_flags
+
+  f = epath.resource_path('mujoco.mjx') / 'test_data' / FLAGS.mjcf
+  m = mujoco.MjModel.from_xml_path(f.as_posix())
+  m.opt.solver = {
+      'cg': mujoco.mjtSolver.mjSOL_CG,
+      'newton': mujoco.mjtSolver.mjSOL_NEWTON,
+  }[FLAGS.solver.lower()]
+  m.opt.iterations = FLAGS.iterations
+  m.opt.ls_iterations = FLAGS.ls_iterations
+  m = mjx.put_model(m)
+
+  if FLAGS.output == 'text':
+    print(f"Rolling out {FLAGS.nstep} steps at dt = {m.opt.timestep:.3f}...")
+
+  @jax.pmap
+  def init(key):
+    key = jax.random.split(key, FLAGS.batch_size // jax.device_count())
+
+    @jax.vmap
+    def random_init(key):
+      d = mjx.make_data(m)
+      qvel = 0.01 * jax.random.normal(key, shape=(m.nv,))
+      d = d.replace(qvel=qvel)
+      return d
+
+    return random_init(key)
+
+  key = jax.random.split(jax.random.key(0), jax.device_count())
+  d = init(key)
+  jax.block_until_ready(d)
+
+  @jax.pmap
+  def unroll(d):
+
+    @jax.vmap
+    def step(d, _):
+      d = mjx.step(m, d)
+      return d, None
+
+    d, _ = jax.lax.scan(step, d, None, length=FLAGS.nstep, unroll=FLAGS.unroll)
+
+    return d
+
+  jit_time, run_time = _measure(unroll, d)
+  steps = FLAGS.nstep * FLAGS.batch_size
+
+  if FLAGS.output == 'text':
+    print(f"""
+Summary for {FLAGS.batch_size} parallel rollouts
+
+ Total JIT time: {jit_time:.2f} s
+ Total simulation time: {run_time:.2f} s
+ Total steps per second: { steps / run_time:.0f}
+ Total realtime factor: { steps * m.opt.timestep / run_time:.2f} x
+ Total time per step: { 1e6 * run_time / steps:.2f} µs""")
+  elif FLAGS.output == 'tsv':
+    name = argv[0].split('/')[-1].replace('testspeed_', '')
+    print(f"{name}\tjit: {jit_time:.2f}s\tsteps/second: {steps / run_time:.0f}")
+
+
+def main():
+  app.run(_main)
+
+
+if __name__ == '__main__':
+  main()

mjx/mujoco/mjx/viewer.py

@@ -28,7 +28,8 @@ _MODEL_PATH = flags.DEFINE_string('mjcf', None, 'Path to a MuJoCo MJCF file.',
                                   required=True)
 
 
-def main(argv: Sequence[str]) -> None:
+def _main(argv: Sequence[str]) -> None:
+  """Launches MuJoCo passive viewer fed by MJX."""
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
 
@@ -37,20 +38,15 @@ def main(argv: Sequence[str]) -> None:
   print(f'Loading model from: {_MODEL_PATH.value}.')
   m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value)
   d = mujoco.MjData(m)
+  mx = mjx.put_model(m)
+  dx = mjx.put_data(m, d)
 
-  mx = mjx.device_put(m)
+  print(f'Default backend: {jax.default_backend()}')
-  dx = mjx.make_data(mx)
+  print('JIT-compiling the model physics step...')
-
-  dt = jax.device_get(mx.opt.timestep)
-  step_fn = jax.jit(mjx.step)
-
-  print(f'JAX default backend: {jax.default_backend()}')
-  print('JIT-compiling the MJX step (this may take a while)...')
   start = time.time()
-  dx = step_fn(mx, dx)
+  step_fn = jax.jit(mjx.step).lower(mx, dx).compile()
-  mjx.device_get_into(d, dx)
   elapsed = time.time() - start
-  print(f'JIT compilation took {elapsed}s.')
+  print(f'Compilation took {elapsed}s.')
 
   with mujoco.viewer.launch_passive(m, d) as v:
     while True:
@@ -68,13 +64,17 @@ def main(argv: Sequence[str]) -> None:
       })
 
       dx = step_fn(mx, dx)
-      mjx.device_get_into(d, dx)
+      mjx.get_data_into(d, m, dx)
       v.sync()
 
       elapsed = time.time() - start
-      if elapsed < dt:
+      if elapsed < m.opt.timestep:
-        time.sleep(dt - elapsed)
+        time.sleep(m.opt.timestep - elapsed)
+
+
+def main():
+  app.run(_main)
 
 
 if __name__ == '__main__':
-  app.run(main)
+  main()

mjx/pyproject.toml

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

mjx/tutorial.ipynb

@@ -102,6 +102,11 @@
         "}\n",
         "\"\"\")\n",
         "\n",
+        "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
+        "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
+        "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
+        "os.environ['XLA_FLAGS'] = xla_flags\n",
+        "\n",
         "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
         "print('Setting environment variable to use GPU rendering:')\n",
         "%env MUJOCO_GL=egl\n",
@@ -150,7 +155,6 @@
       "cell_type": "code",
       "execution_count": null,
       "metadata": {
-        "cellView": "form",
         "id": "ObF1UXrkb0Nd"
       },
       "outputs": [],
@@ -170,7 +174,7 @@
         "from brax import envs\n",
         "from brax import math\n",
         "from brax.base import Base, Motion, Transform\n",
-        "from brax.envs.base import Env, MjxEnv, State\n",
+        "from brax.envs.base import Env, PipelineEnv, State\n",
         "from brax.mjx.base import State as MjxState\n",
         "from brax.training.agents.ppo import train as ppo\n",
         "from brax.training.agents.ppo import networks as ppo_networks\n",
@@ -413,7 +417,7 @@
       "source": [
         "#@title Humanoid Env\n",
         "\n",
-        "class Humanoid(MjxEnv):\n",
+        "class Humanoid(PipelineEnv):\n",
         "\n",
         "  def __init__(\n",
         "      self,\n",
@@ -427,7 +431,7 @@
         "      **kwargs,\n",
         "  ):\n",
         "    path = epath.Path(epath.resource_path('mujoco')) / (\n",
-        "        'mjx/benchmark/model/humanoid'\n",
+        "        'mjx/test_data/humanoid'\n",
         "    )\n",
         "    mj_model = mujoco.MjModel.from_xml_path(\n",
         "        (path / 'humanoid.xml').as_posix())\n",
@@ -435,11 +439,14 @@
         "    mj_model.opt.iterations = 6\n",
         "    mj_model.opt.ls_iterations = 6\n",
         "\n",
+        "    sys = mjcf.load_model(mj_model)\n",
+        "\n",
         "    physics_steps_per_control_step = 5\n",
         "    kwargs['n_frames'] = kwargs.get(\n",
         "        'n_frames', physics_steps_per_control_step)\n",
+        "    kwargs['backend'] = 'mjx'\n",
         "\n",
-        "    super().__init__(model=mj_model, **kwargs)\n",
+        "    super().__init__(sys, **kwargs)\n",
         "\n",
         "    self._forward_reward_weight = forward_reward_weight\n",
         "    self._ctrl_cost_weight = ctrl_cost_weight\n",
@@ -465,7 +472,7 @@
         "\n",
         "    data = self.pipeline_init(qpos, qvel)\n",
         "\n",
-        "    obs = self._get_obs(data.data, jp.zeros(self.sys.nu))\n",
+        "    obs = self._get_obs(data, jp.zeros(self.sys.nu))\n",
         "    reward, done, zero = jp.zeros(3)\n",
         "    metrics = {\n",
         "        'forward_reward': zero,\n",
@@ -485,8 +492,8 @@
         "    data0 = state.pipeline_state\n",
         "    data = self.pipeline_step(data0, action)\n",
         "\n",
-        "    com_before = data0.data.subtree_com[1]\n",
+        "    com_before = data0.subtree_com[1]\n",
-        "    com_after = data.data.subtree_com[1]\n",
+        "    com_after = data.subtree_com[1]\n",
         "    velocity = (com_after - com_before) / self.dt\n",
         "    forward_reward = self._forward_reward_weight * velocity[0]\n",
         "\n",
@@ -500,7 +507,7 @@
         "\n",
         "    ctrl_cost = self._ctrl_cost_weight * jp.sum(jp.square(action))\n",
         "\n",
-        "    obs = self._get_obs(data.data, action)\n",
+        "    obs = self._get_obs(data, action)\n",
         "    reward = forward_reward + healthy_reward - ctrl_cost\n",
         "    done = 1.0 - is_healthy if self._terminate_when_unhealthy else 0.0\n",
         "    state.metrics.update(\n",
@@ -599,7 +606,7 @@
       "source": [
         "## Train Humanoid Policy\n",
         "\n",
-        "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 9-10 minutes on a Tesla A100 GPU."
+        "Let's now train a policy with PPO to make the Humanoid run forwards. Training takes about 8-9 minutes on a Tesla A100 GPU."
       ]
     },
     {
@@ -759,7 +766,7 @@
       },
       "outputs": [],
       "source": [
-        "mj_model = eval_env._model\n",
+        "mj_model = eval_env.sys.mj_model\n",
         "mj_data = mujoco.MjData(mj_model)\n",
         "\n",
         "renderer = mujoco.Renderer(mj_model)\n",
@@ -957,7 +964,7 @@
         "  return default_config\n",
         "\n",
         "\n",
-        "class BarkourEnv(MjxEnv):\n",
+        "class BarkourEnv(PipelineEnv):\n",
         "  \"\"\"Environment for training the barkour quadruped joystick policy in MJX.\"\"\"\n",
         "\n",
         "  def __init__(\n",
@@ -968,18 +975,19 @@
         "      **kwargs,\n",
         "  ):\n",
         "    path = epath.Path('mujoco_menagerie/google_barkour_vb/scene_mjx.xml')\n",
+        "    sys = mjcf.load(path.as_posix())\n",
         "    self._dt = 0.02  # this environment is 50 fps\n",
-        "    self.brax_sys = mjcf.load(path).replace(dt=self._dt)\n",
+        "    sys = sys.tree_replace({'opt.timestep': 0.004, 'dt': 0.004})\n",
-        "    model = self.brax_sys.get_model()\n",
-        "    model.opt.timestep = 0.004\n",
         "\n",
         "    # override menagerie params for smoother policy\n",
-        "    model.dof_damping[6:] = 0.5239\n",
+        "    sys = sys.replace(\n",
-        "    model.actuator_gainprm[:, 0] = 35.0\n",
+        "        dof_damping=sys.dof_damping.at[6:].set(0.5239),\n",
-        "    model.actuator_biasprm[:, 1] = -35.0\n",
+        "        actuator_gainprm=sys.actuator_gainprm.at[:, 0].set(35.0),\n",
+        "        actuator_biasprm=sys.actuator_biasprm.at[:, 1].set(-35.0),\n",
+        "    )\n",
         "\n",
-        "    n_frames = kwargs.pop('n_frames', int(self._dt / model.opt.timestep))\n",
+        "    n_frames = kwargs.pop('n_frames', int(self._dt / sys.opt.timestep))\n",
-        "    super().__init__(model=model, n_frames=n_frames)\n",
+        "    super().__init__(sys, backend='mjx', n_frames=n_frames)\n",
         "\n",
         "    self.reward_config = get_config()\n",
         "    # set custom from kwargs\n",
@@ -988,13 +996,13 @@
         "        self.reward_config.rewards.scales[k[:-6]] = v\n",
         "\n",
         "    self._torso_idx = mujoco.mj_name2id(\n",
-        "        model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
+        "        sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, 'torso'\n",
         "    )\n",
         "    self._action_scale = action_scale\n",
         "    self._obs_noise = obs_noise\n",
         "    self._kick_vel = kick_vel\n",
-        "    self._init_q = jp.array(model.keyframe('home').qpos)\n",
+        "    self._init_q = jp.array(sys.mj_model.keyframe('home').qpos)\n",
-        "    self._default_pose = model.keyframe('home').qpos[7:]\n",
+        "    self._default_pose = sys.mj_model.keyframe('home').qpos[7:]\n",
         "    self.lowers = jp.array([-0.7, -1.0, 0.05] * 4)\n",
         "    self.uppers = jp.array([0.52, 2.1, 2.1] * 4)\n",
         "    feet_site = [\n",
@@ -1004,7 +1012,7 @@
         "        'foot_hind_right',\n",
         "    ]\n",
         "    feet_site_id = [\n",
-        "        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n",
+        "        mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_SITE.value, f)\n",
         "        for f in feet_site\n",
         "    ]\n",
         "    assert not any(id_ == -1 for id_ in feet_site_id), 'Site not found.'\n",
@@ -1016,13 +1024,13 @@
         "        'lower_leg_hind_right',\n",
         "    ]\n",
         "    lower_leg_body_id = [\n",
-        "        mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n",
+        "        mujoco.mj_name2id(sys.mj_model, mujoco.mjtObj.mjOBJ_BODY.value, l)\n",
         "        for l in lower_leg_body\n",
         "    ]\n",
         "    assert not any(id_ == -1 for id_ in lower_leg_body_id), 'Body not found.'\n",
         "    self._lower_leg_body_id = np.array(lower_leg_body_id)\n",
         "    self._foot_radius = 0.0175\n",
-        "    self._nv = model.nv\n",
+        "    self._nv = sys.nv\n",
         "\n",
         "  def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n",
         "    lin_vel_x = [-0.6, 1.5]  # min max [m/s]\n",
@@ -1076,9 +1084,9 @@
         "    kick_theta = jax.random.uniform(kick_noise_2, maxval=2 * jp.pi)\n",
         "    kick = jp.array([jp.cos(kick_theta), jp.sin(kick_theta)])\n",
         "    kick *= jp.mod(state.info['step'], push_interval) == 0\n",
-        "    qvel = state.pipeline_state.data.qvel  # pytype: disable=attribute-error\n",
+        "    qvel = state.pipeline_state.qvel  # pytype: disable=attribute-error\n",
         "    qvel = qvel.at[:2].set(kick * self._kick_vel + qvel[:2])\n",
-        "    state = state.tree_replace({'pipeline_state.data.qvel': qvel})\n",
+        "    state = state.tree_replace({'pipeline_state.qvel': qvel})\n",
         "\n",
         "    # physics step\n",
         "    motor_targets = self._default_pose + action * self._action_scale\n",
@@ -1092,7 +1100,7 @@
         "    joint_vel = pipeline_state.qd[6:]\n",
         "\n",
         "    # foot contact data based on z-position\n",
-        "    foot_pos = pipeline_state.data.site_xpos[self._feet_site_id]  # pytype: disable=attribute-error\n",
+        "    foot_pos = pipeline_state.site_xpos[self._feet_site_id]  # pytype: disable=attribute-error\n",
         "    foot_contact_z = foot_pos[:, 2] - self._foot_radius\n",
         "    contact = foot_contact_z \u003c 1e-3  # a mm or less off the floor\n",
         "    contact_filt_mm = contact | state.info['last_contact']\n",
@@ -1118,7 +1126,7 @@
         "        'lin_vel_z': self._reward_lin_vel_z(xd),\n",
         "        'ang_vel_xy': self._reward_ang_vel_xy(xd),\n",
         "        'orientation': self._reward_orientation(x),\n",
-        "        'torques': self._reward_torques(pipeline_state.data.qfrc_actuator),  # pytype: disable=attribute-error\n",
+        "        'torques': self._reward_torques(pipeline_state.qfrc_actuator),  # pytype: disable=attribute-error\n",
         "        'action_rate': self._reward_action_rate(action, state.info['last_act']),\n",
         "        'stand_still': self._reward_stand_still(\n",
         "            state.info['command'], joint_angles,\n",
@@ -1262,8 +1270,8 @@
         "  ) -\u003e jax.Array:\n",
         "    # get velocities at feet which are offset from lower legs\n",
         "    # pytype: disable=attribute-error\n",
-        "    pos = pipeline_state.data.site_xpos[self._feet_site_id]  # feet position\n",
+        "    pos = pipeline_state.site_xpos[self._feet_site_id]  # feet position\n",
-        "    feet_offset = pos - pipeline_state.data.xpos[self._lower_leg_body_id]\n",
+        "    feet_offset = pos - pipeline_state.xpos[self._lower_leg_body_id]\n",
         "    # pytype: enable=attribute-error\n",
         "    offset = base.Transform.create(pos=feet_offset)\n",
         "    foot_indices = self._lower_leg_body_id - 1  # we got rid of the world body\n",
@@ -1279,7 +1287,7 @@
         "      self, trajectory: List[base.State], camera: str | None = None\n",
         "  ) -\u003e Sequence[np.ndarray]:\n",
         "    camera = camera or 'track'\n",
-        "    return super().render(trajectory, camera)\n",
+        "    return super().render(trajectory, camera=camera)\n",
         "\n",
         "envs.register_environment('barkour', BarkourEnv)"
       ]
@@ -1440,7 +1448,7 @@
       },
       "outputs": [],
       "source": [
-        "HTML(html.render(eval_env.brax_sys, rollout))"
+        "HTML(html.render(eval_env.sys, rollout))"
       ]
     }
   ],
@@ -1448,13 +1456,13 @@
     "accelerator": "GPU",
     "colab": {
       "gpuClass": "premium",
-      "gpuType": "A100",
+      "gpuType": "V100",
       "machine_shape": "hm",
       "private_outputs": true,
       "provenance": [
         {
-          "file_id": "11cFRVCJ8Kn71tlQFbFcw4JzQZ00F8BRG",
+          "file_id": "1A58SK07tnOzix53E68D0TQ2ePCTZA61f",
-          "timestamp": 1704355889284
+          "timestamp": 1707342610876
         }
       ],
       "toc_visible": true

plugin/elasticity/elasticity.h

@@ -86,13 +86,14 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
 }
 
 template <typename T>
-inline void ComputeForce(mjtNum* qfrc_passive,
+inline void ComputeForce(std::vector<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) {
+  mju_zero(qfrc_passive.data(), qfrc_passive.size());
+
   for (int t = 0; t < elements.size(); t++)  {
     const int* v = elements[t].vertices;
 
@@ -118,7 +119,7 @@ inline void ComputeForce(mjtNum* qfrc_passive,
       for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
         for (int i = 0; i < 2; i++) {
           for (int x = 0; x < 3; x++) {
-            force[3 * T::edge[ed2][i] + x] +=
+            force[3 * T::edge[ed2][i] + x] -=
                 elongation[ed1] * gradient[ed2][i][x] *
                 metric[offset * t + T::kNumEdges * ed1 + ed2];
           }
@@ -128,17 +129,31 @@ inline void ComputeForce(mjtNum* qfrc_passive,
 
     // insert into global force
     for (int i = 0; i < T::kNumVerts; i++) {
-      int body_dofnum = 3;
+      for (int x = 0; x < 3; x++) {
-      int body_dofadr = 3*v[i];
+        qfrc_passive[3*v[i]+x] += force[3*i+x];
-      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");
-        }
       }
+    }
+  }
+}
+
+// add flex force to degrees of freedom
+inline void AddFlexForce(mjtNum* qfrc,
+                         const std::vector<mjtNum>& force,
+                         const mjModel* m, mjData* d,
+                         const mjtNum* xpos,
+                         int f0) {
+  int* bodyid = m->flex_vertbodyid + m->flex_vertadr[f0];
+
+  for (int v = 0; v < m->flex_vertnum[f0]; v++) {
+    int bid = bodyid[v];
+    if (m->body_simple[bid] != 2) {
+      // this should only occur for pinned flex vertices
+      mj_applyFT(m, d, force.data() + 3*v, 0, xpos + 3*v, bid, qfrc);
+    } else {
+      int body_dofnum = m->body_dofnum[bid];
+      int body_dofadr = m->body_dofadr[bid];
       for (int x = 0; x < body_dofnum; x++) {
-        qfrc_passive[body_dofadr+x] -= force[3*i+x];
+        qfrc[body_dofadr+x] += force[3*v+x];
       }
     }
   }

plugin/elasticity/membrane.cc

@@ -117,6 +117,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
     for (int j = 0; j < m->flex_vertnum[i]; j++) {
       if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) {
         f0 = i;
+        nv = m->flex_vertnum[f0];
       }
     }
   }
@@ -168,6 +169,7 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
   deformed.assign(ne, 0);
   previous.assign(ne, 0);
   elongation.assign(ne, 0);
+  force.assign(3*nv, 0);
 
   // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
   UpdateSquaredLengths(reference, edges, body_pos);
@@ -198,11 +200,17 @@ void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
 
   // compute gradient of elastic energy and insert into passive force
   int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0];
-  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);
+  ComputeForce<Stencil2D>(force, elements, metric, elongation, m, xpos);
+
+  // insert into passive force
+  if (f0 < 0) {
+    mju_addTo(qfrc, force.data(), force.size());
+  } else {
+    AddFlexForce(qfrc, force, m, d, xpos, f0);
+  }
 
   // update stored lengths
   if (kD > 0) {

plugin/elasticity/membrane.h

@@ -57,6 +57,7 @@ class Membrane {
   std::vector<mjtNum> deformed;             // deformed lengths       (ne x 1)
   std::vector<mjtNum> previous;             // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
+  std::vector<mjtNum> force;                // force at all vertices  (nv x 3)
 
   mjtNum damping;
   mjtNum thickness;

plugin/elasticity/solid.cc

@@ -122,6 +122,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
     for (int j = 0; j < m->flex_vertnum[i]; j++) {
       if (m->flex_vertbodyid[m->flex_vertadr[i]+j] == i0) {
         f0 = i;
+        nv = m->flex_vertnum[f0];
       }
     }
   }
@@ -172,6 +173,7 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
   deformed.assign(ne, 0);
   previous.assign(ne, 0);
   elongation.assign(ne, 0);
+  force.assign(3*nv, 0);
 
   // compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
   UpdateSquaredLengths(reference, edges, body_pos);
@@ -202,11 +204,17 @@ void Solid::Compute(const mjModel* m, mjData* d, int instance) {
 
   // compute gradient of elastic energy and insert into passive force
   int flex_vertadr = f0 < 0 ? -1 : m->flex_vertadr[f0];
-  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);
+  ComputeForce<Stencil3D>(force, elements, metric, elongation, m, xpos);
+
+  // insert into passive force
+  if (f0 < 0) {
+    mju_addTo(qfrc, force.data(), force.size());
+  } else {
+    AddFlexForce(qfrc, force, m, d, xpos, f0);
+  }
 
   // update stored lengths
   if (kD > 0) {

plugin/elasticity/solid.h

@@ -55,6 +55,7 @@ class Solid {
   std::vector<mjtNum> deformed;             // deformed lengths       (ne x 1)
   std::vector<mjtNum> previous;             // previous-step lengths  (ne x 1)
   std::vector<mjtNum> elongation;           // edge elongation        (ne x 1)
+  std::vector<mjtNum> force;                // force at all vertices  (nv x 3)
 
   mjtNum damping;
 

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.1.2.dylib
+          ${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.1.3.dylib
       )
       target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
     endif()
@@ -92,7 +92,7 @@ if(NOT TARGET mujoco)
 
   if(NOT MUJOCO_FRAMEWORK)
     find_library(
-      MUJOCO_LIBRARY mujoco mujoco.3.1.2 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
+      MUJOCO_LIBRARY mujoco mujoco.3.1.3 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
     )
     find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
     message("MuJoCo is at ${MUJOCO_LIBRARY}")
@@ -140,7 +140,7 @@ findorfetch(
   GIT_REPO
   https://github.com/abseil/abseil-cpp
   GIT_TAG
-  fb3621f4f897824c0dbe0615fa94543df6192f30 # LTS 20230802.1
+  4a2c63365eff8823a5221db86ef490e828306f9d # LTS 20240116.0
   TARGETS
   ${MUJOCO_PYTHON_ABSL_TARGETS}
   EXCLUDE_FROM_ALL
@@ -173,7 +173,7 @@ findorfetch(
   GIT_REPO
   https://gitlab.com/libeigen/eigen
   GIT_TAG
-  454f89af9d6f3525b1df5f9ef9c86df58bf2d4d3
+  7fd7a3f946e5ac152d28dad388cff8bfa1026925
   TARGETS
   Eigen3::Eigen
   EXCLUDE_FROM_ALL

python/mujoco/mjpython/Info.plist

@@ -7,13 +7,13 @@
     <key>CFBundleIdentifier</key>
     <string>org.mujoco.mjpython</string>
     <key>CFBundleVersion</key>
-    <string>3.1.2</string>
+    <string>3.1.3</string>
     <key>CFBundleGetInfoString</key>
-    <string>3.1.2</string>
+    <string>3.1.3</string>
     <key>CFBundleLongVersionString</key>
-    <string>3.1.2</string>
+    <string>3.1.3</string>
     <key>CFBundleShortVersionString</key>
-    <string>3.1.2</string>
+    <string>3.1.3</string>
     <key>CFBundleExecutable</key>
     <string>mjpython</string>
     <key>CFBundleIconFile</key>

python/mujoco/rollout.cc

@@ -12,15 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <array>
-#include <cstdio>
 #include <iostream>
 #include <optional>
 #include <sstream>
-#include <string>
 
-#include "functions.h"
+#include <mujoco/mujoco.h>
+#include "errors.h"
 #include "raw.h"
+#include "structs.h"
 #include <pybind11/buffer_info.h>
 #include <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -32,112 +31,114 @@ namespace {
 
 namespace py = ::pybind11;
 
+// NOLINTBEGIN(whitespace/line_length)
+
 const auto rollout_doc = R"(
-Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
+Roll out open-loop trajectories from initial states, get resulting states and sensor values.
 
   input arguments (required):
-    model              an instance of MjModel
+    model              instance of MjModel
-    data               an associated instance of MjData
+    data               associated instance of MjData
-    nstate             an integer, number of initial states from which to roll out trajectories
+    nroll              integer, number of initial states from which to roll out trajectories
-    nstep              an integer, number of steps to be taken for each trajectory
+    nstep              integer, number of steps to be taken for each trajectory
+    control_spec       specification of controls, ncontrol = mj_stateSize(m, control_spec)
+    state0             (nroll x nstate) nroll initial state vectors,
+                                        nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS)
   input arguments (optional):
-    initial_state      (nstate x nqva)                 nstate initial state vectors, nqva=nq+nv+na
+    warmstart0         (nroll x nv)                   nroll qacc_warmstart vectors
-    initial_time       (nstate x 1)                    nstate initial times
+    control            (nroll x nstep x ncontrol)     nroll trajectories of nstep controls
-    initial_warmstart  (nstate x nv)                   nstate qacc_warmstart vectors
-    ctrl               (nstate x nstep x nu)           nstate length-nstep controls
-    qfrc_applied       (nstate x nstep x nv)           nstate length-nstep generalized forces
-    xfrc_applied       (nstate x nstep x nbody*6)      nstate length-nstep Cartesian wrenches
-    mocap              (nstate x nstep x nmocap*7)     nstate length-nstep mocap body poses
   output arguments (optional):
-    state              (nstate x nstep x nqva)         nstate length-nstep states
+    state              (nroll x nstep x nstate)       nroll nstep states
-    sensordata         (nstate x nstep x nsendordata)  nstate length-nstep sensordatas
+    sensordata         (nroll x nstep x nsendordata)  nroll trajectories of nstep sensordata vectors
 )";
 
 // C-style rollout function, assumes all arguments are valid
 // all input fields of d are initialised, contents at call time do not matter
 // after returning, d will contain the last step of the last rollout
-void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
+void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned int control_spec,
-                     const mjtNum* state0, const mjtNum* ctrl,
+                     const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control,
-                     const mjtNum* qfrc, const mjtNum* xfrc,
-                     const mjtNum* mocap, const mjtNum* time0,
-                     const mjtNum* warmstart0,
                      mjtNum* state, mjtNum* sensordata) {
-  // model sizes
+  // sizes
-  int nq = m->nq;
+  int nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS);
-  int nv = m->nv;
+  int ncontrol = mj_stateSize(m, control_spec);
-  int na = m->na;
+  int nv = m->nv, nbody = m->nbody, neq = m->neq;
-  int nqva = nq + nv + na;
-  int nu = m->nu;
-  int nbody = m->nbody;
-  int nmocap = m->nmocap;
   int nsensordata = m->nsensordata;
 
-  // loop over initial states
+  // clear user inputs if unspecified
-  for (int s=0; s < nstate; s++) {
+  if (!(control_spec & mjSTATE_CTRL)) {
-
+    mju_zero(d->ctrl, m->nu);
-    // set initial state
+  }
-    if (state0) {
+  if (!(control_spec & mjSTATE_QFRC_APPLIED)) {
-      mju_copy(d->qpos, state0 + s*nqva, nq);
+    mju_zero(d->qfrc_applied, nv);
-      mju_copy(d->qvel, state0 + s*nqva + nq, nv);
+  }
-      mju_copy(d->act,  state0 + s*nqva + nq + nv, na);
+  if (!(control_spec & mjSTATE_XFRC_APPLIED)) {
-    } else {
+    mju_zero(d->xfrc_applied, 6*nbody);
-      mju_copy(d->qpos, m->qpos0, nq);
+  }
-      mju_zero(d->qvel, nv);
+  if (!(control_spec & mjSTATE_MOCAP_POS)) {
-      mju_zero(d->act, na);
+    for (int i = 0; i < nbody; i++) {
+      int id = m->body_mocapid[i];
+      if (id >= 0) mju_copy3(d->mocap_pos+3*id, m->body_pos+3*i);
     }
+  }
+  if (!(control_spec & mjSTATE_MOCAP_QUAT)) {
+    for (int i = 0; i < nbody; i++) {
+      int id = m->body_mocapid[i];
+      if (id >= 0) mju_copy4(d->mocap_quat+4*id, m->body_quat+4*i);
+    }
+  }
+  if (!(control_spec & mjSTATE_EQ_ACTIVE)) {
+    for (int i = 0; i < neq; i++) {
+      d->eq_active[i] = m->eq_active0[i];
+    }
+  }
 
-    // set initial time
+  // loop over rollouts
-    d->time = time0 ? time0[s] : 0;
+  for (int r = 0; r < nroll; r++) {
+    // set initial state
+    mj_setState(m, d, state0 + r*nstate, mjSTATE_FULLPHYSICS);
 
     // set warmstart accelerations
     if (warmstart0) {
-      mju_copy(d->qacc_warmstart, warmstart0 + s*nv, nv);
+      mju_copy(d->qacc_warmstart, warmstart0 + r*nv, nv);
     } else {
       mju_zero(d->qacc_warmstart, nv);
     }
 
-    // clear control inputs if unspecified
+    // clear warning counters
-    if (s == 0) {
+    for (int i = 0; i < mjNWARNING; i++) {
-      if (!ctrl) {
+      d->warning[i].number = 0;
-        mju_zero(d->ctrl, nu);
-      }
-      if (!qfrc) {
-        mju_zero(d->qfrc_applied, nv);
-      }
-      if (!xfrc) {
-        mju_zero(d->xfrc_applied, 6*nbody);
-      }
-      if (!mocap) {
-        for (int j=0; j<nbody; j++) {
-          int id = m->body_mocapid[j];
-          if (id>=0) {
-            mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j);
-            mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j);
-          }
-        }
-      }
     }
 
-    // roll out trajectories
+    // roll out trajectory
     for (int t = 0; t < nstep; t++) {
+      // check for warnings
+      bool nwarning = false;
+      for (int i = 0; i < mjNWARNING; i++) {
+        if (d->warning[i].number) {
+          nwarning = true;
+          break;
+        }
+      }
+
+      // if any warnings, fill remaining outputs with current outputs, break
+      if (nwarning) {
+        for (; t < nstep; t++) {
+          int step = r*nstep + t;
+          if (state) {
+            mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS);
+          }
+          if (sensordata) {
+            mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata);
+          }
+        }
+        break;
+      }
+
+      int step = r*nstep + t;
+
       // controls
-      if (ctrl) {
+      if (control) {
-        mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu);
+        mj_setState(m, d, control + step*ncontrol, control_spec);
-      }
-      // generalized forces
-      if (qfrc) {
-        mju_copy(d->qfrc_applied, qfrc + s*nstep*nv + t*nv, nv);
-      }
-      // Cartesian wrenches
-      if (xfrc) {
-        mju_copy(d->xfrc_applied, xfrc + s*nstep*6*nbody + t*6*nbody, 6*nbody);
-      }
-      // mocap bodies
-      if (mocap) {
-        mju_copy(d->mocap_pos,
-                 mocap + s*nstep*7*nmocap + t*7*nmocap, 3*nmocap);
-        mju_copy(d->mocap_quat,
-                 mocap + s*nstep*7*nmocap + t*7*nmocap + 3*nmocap, 4*nmocap);
       }
 
       // step
@@ -145,23 +146,22 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
 
       // copy out new state
       if (state) {
-        mju_copy(state + s*nstep*nqva + t*nqva,           d->qpos, nq);
+        mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS);
-        mju_copy(state + s*nstep*nqva + t*nqva + nq,      d->qvel, nv);
-        mju_copy(state + s*nstep*nqva + t*nqva + nq + nv, d->act,  na);
       }
+
       // copy out sensor values
       if (sensordata) {
-        mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata,
+        mju_copy(sensordata + step*nsensordata, d->sensordata, nsensordata);
-                 d->sensordata, nsensordata);
       }
     }
   }
 }
 
+// NOLINTEND(whitespace/line_length)
 
 // check size of optional argument to rollout(), return raw pointer
 mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
-                      const char* name, int nstate, int nstep, int dim) {
+                      const char* name, int nroll, int nstep, int dim) {
   // if empty return nullptr
   if (!arg.has_value()) {
     return nullptr;
@@ -171,10 +171,10 @@ mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
   py::buffer_info info = arg->request();
 
   // check size
-  int expected_size = nstate * nstep * dim;
+  int expected_size = nroll * nstep * dim;
   if (info.size != expected_size) {
     std::ostringstream msg;
-    msg << name << ".size should be " << expected_size <<  ", got " << info.size;
+    msg << name << ".size should be " << expected_size << ", got " << info.size;
     throw py::value_error(msg.str());
   }
   return static_cast<mjtNum*>(info.ptr);
@@ -189,44 +189,35 @@ PYBIND11_MODULE(_rollout, pymodule) {
   // get subsequent states and corresponding sensor values
   pymodule.def(
       "rollout",
-      [](const MjModelWrapper& m, MjDataWrapper& d, int nstate, int nstep,
+      [](const MjModelWrapper& m, MjDataWrapper& d,
-         std::optional<const PyCArray> init_state,
+         int nroll, int nstep, unsigned int control_spec,
-         std::optional<const PyCArray> init_time,
+         const PyCArray state0,
-         std::optional<const PyCArray> init_warmstart,
+         std::optional<const PyCArray> warmstart0,
-         std::optional<const PyCArray> ctrl,
+         std::optional<const PyCArray> control,
-         std::optional<const PyCArray> qfrc,
-         std::optional<const PyCArray> xfrc,
-         std::optional<const PyCArray> mocap,
          std::optional<const PyCArray> state,
          std::optional<const PyCArray> sensordata
          ) {
-
         const raw::MjModel* model = m.get();
         raw::MjData* data = d.get();
 
         // check that some steps need to be taken, return if not
-        if (nstate < 1 || nstep < 1) {
+        if (nroll < 1 || nstep < 1) {
           return;
         }
 
+        // get sizes
+        int nstate = mj_stateSize(model, mjSTATE_FULLPHYSICS);
+        int ncontrol = mj_stateSize(model, control_spec);
+
         // get raw pointers
-        int nqva = model->nq + model->nv + model->na;
+        mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate);
-        mjtNum* init_state_ptr =
+        mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll,
-            get_array_ptr(init_state, "initial_state", nstate, 1, nqva);
+                                               1, model->nv);
-        mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu);
+        mjtNum* control_ptr = get_array_ptr(control, "control", nroll,
-        mjtNum* qfrc_ptr =
+                                            nstep, ncontrol);
-            get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv);
+        mjtNum* state_ptr = get_array_ptr(state, "state", nroll, nstep, nstate);
-        mjtNum* xfrc_ptr =
+        mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nroll,
-            get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody);
+                                               nstep, model->nsensordata);
-        mjtNum* mocap_ptr =
-            get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap);
-        mjtNum* init_time_ptr =
-            get_array_ptr(init_time, "init_time", nstate, 1, 1);
-        mjtNum* init_warmstart_ptr =
-            get_array_ptr(init_warmstart, "init_warmstart", nstate, 1, model->nv);
-        mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva);
-        mjtNum* sensordata_ptr =
-            get_array_ptr(sensordata, "sensordata", nstate, nstep, model->nsensordata);
 
         // perform rollouts
         {
@@ -235,29 +226,25 @@ PYBIND11_MODULE(_rollout, pymodule) {
 
           // call unsafe rollout function
           InterceptMjErrors(_unsafe_rollout)(
-              model, data, nstate, nstep, init_state_ptr, ctrl_ptr, qfrc_ptr,
+              model, data, nroll, nstep, control_spec, state0_ptr,
-              xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr,
+              warmstart0_ptr, control_ptr, state_ptr, sensordata_ptr);
-              sensordata_ptr);
         }
       },
       py::arg("model"),
       py::arg("data"),
-      py::arg("nstate"),
+      py::arg("nroll"),
       py::arg("nstep"),
-      py::arg("initial_state")     = py::none(),
+      py::arg("control_spec"),
-      py::arg("initial_time")      = py::none(),
+      py::arg("state0"),
-      py::arg("initial_warmstart") = py::none(),
+      py::arg("warmstart0") = py::none(),
-      py::arg("ctrl")              = py::none(),
+      py::arg("control")    = py::none(),
-      py::arg("qfrc_applied")      = py::none(),
+      py::arg("state")      = py::none(),
-      py::arg("xfrc_applied")      = py::none(),
+      py::arg("sensordata") = py::none(),
-      py::arg("mocap")             = py::none(),
-      py::arg("state")             = py::none(),
-      py::arg("sensordata")        = py::none(),
       py::doc(rollout_doc)
   );
+}
 
 }  // namespace
 
-}
+}  // namespace mujoco::python
 
-}

python/mujoco/rollout.py

@@ -14,132 +14,144 @@
 # ==============================================================================
 """Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
 
+from typing import Optional
+
+import mujoco
 from mujoco import _rollout
 import numpy as np
+from numpy import typing as npt
 
 
-def rollout(model, data, initial_state=None, ctrl=None,
+def rollout(model: mujoco.MjModel,
-            *,  # require following arguments to be named
+            data: mujoco.MjData,
-            skip_checks=False,
+            initial_state: npt.ArrayLike,
-            nstate=None,
+            control: Optional[npt.ArrayLike] = None,
-            nstep=None,
+            *,  # require subsequent arguments to be named
-            initial_time=None,
+            control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
-            initial_warmstart=None,
+            skip_checks: bool = False,
-            qfrc_applied=None,
+            nroll: Optional[int] = None,
-            xfrc_applied=None,
+            nstep: Optional[int] = None,
-            mocap=None,
+            initial_warmstart: Optional[npt.ArrayLike] = None,
-            state=None,
+            state: Optional[npt.ArrayLike] = None,
-            sensordata=None):
+            sensordata: Optional[npt.ArrayLike] = None):
-  """Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
+  """Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
 
-    This function serves as a Python wrapper for the C++ functionality in
+  Python wrapper for rollout.cc, see documentation therein.
-    `rollout.cc`, please see documentation therein. This python funtion will
+  Infers nroll and nstep.
-    infer `nstate` and `nstep`, tile input arguments with singleton dimensions,
+  Tiles inputs with singleton dimensions.
-    and allocate output arguments if none are given.
+  Allocates outputs if none are given.
+
+  Args:
+    model: An mjModel instance.
+    data: An associated mjData instance.
+    initial_state: Array of initial states from which to roll out trajectories.
+      ([nroll or 1] x nstate)
+    control: Open-loop controls array to apply during the rollouts.
+      ([nroll or 1] x [nstep or 1] x ncontrol)
+    control_spec: mjtState specification of control vectors.
+    skip_checks: Whether to skip internal shape and type checks.
+    nroll: Number of rollouts (inferred if unspecified).
+    nstep: Number of steps in rollouts (inferred if unspecified).
+    initial_warmstart: Initial qfrc_warmstart array (optional).
+      ([nroll or 1] x nv)
+    state: State output array (optional).
+      (nroll x nstep x nstate)
+    sensordata: Sensor data output array (optional).
+      (nroll x nstep x nsensordata)
+
+  Returns:
+    state:
+      State output array, (nroll x nstep x nstate).
+    sensordata:
+      Sensor data output array, (nroll x nstep x nsensordata).
+
+  Raises:
+    ValueError: bad shapes or sizes.
   """
-  # don't infer nstate/nstep, don't support singleton expansion, don't allocate
+  # skip_checks shortcut:
-  # output arrays, just call rollout
+  #   don't infer nroll/nstep
+  #   don't support singleton expansion
+  #   don't allocate output arrays
+  #   just call rollout and return
   if skip_checks:
-    _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
+    _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
-                     initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
+                     initial_warmstart, control, state, sensordata)
-                     state, sensordata)
     return state, sensordata
 
+  # check control_spec
+  if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
+    raise ValueError('control_spec can only contain bits in mjSTATE_USER')
+
   # check types
-  if nstate and not isinstance(nstate, int):
+  if nroll and not isinstance(nroll, int):
-    raise ValueError('nstate must be an integer')
+    raise ValueError('nroll must be an integer')
   if nstep and not isinstance(nstep, int):
     raise ValueError('nstep must be an integer')
   _check_must_be_numeric(
       initial_state=initial_state,
-      initial_time=initial_time,
       initial_warmstart=initial_warmstart,
-      ctrl=ctrl,
+      control=control,
-      qfrc_applied=qfrc_applied,
-      xfrc_applied=xfrc_applied,
-      mocap=mocap,
       state=state,
       sensordata=sensordata)
 
   # check number of dimensions
   _check_number_of_dimensions(2,
                               initial_state=initial_state,
-                              initial_time=initial_time,
                               initial_warmstart=initial_warmstart)
   _check_number_of_dimensions(3,
-                              ctrl=ctrl,
+                              control=control,
-                              qfrc_applied=qfrc_applied,
-                              xfrc_applied=xfrc_applied,
-                              mocap=mocap,
                               state=state,
                               sensordata=sensordata)
 
   # ensure 2D, make contiguous, row-major (C ordering)
   initial_state = _ensure_2d(initial_state)
-  initial_time = _ensure_2d(initial_time)
   initial_warmstart = _ensure_2d(initial_warmstart)
 
   # ensure 3D, make contiguous, row-major (C ordering)
-  ctrl = _ensure_3d(ctrl)
+  control = _ensure_3d(control)
-  qfrc_applied = _ensure_3d(qfrc_applied)
-  xfrc_applied = _ensure_3d(xfrc_applied)
-  mocap = _ensure_3d(mocap)
   state = _ensure_3d(state)
   sensordata = _ensure_3d(sensordata)
 
   # check trailing dimensions
-  _check_trailing_dimension(model.nq + model.nv + model.na,
+  nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
-                            initial_state=initial_state, state=state)
+  _check_trailing_dimension(nstate, initial_state=initial_state, state=state)
-  _check_trailing_dimension(1, initial_time=initial_time)
+  ncontrol = mujoco.mj_stateSize(model, control_spec)
-  _check_trailing_dimension(model.nu, ctrl=ctrl)
+  _check_trailing_dimension(ncontrol, control=control)
-  _check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied)
+  _check_trailing_dimension(model.nv, initial_warmstart=initial_warmstart)
-  _check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied)
-  _check_trailing_dimension(model.nmocap*7, mocap=mocap)
   _check_trailing_dimension(model.nsensordata, sensordata=sensordata)
 
-  # infer nstate, check for incompatibilities
+  # infer nroll, check for incompatibilities
-  nstate = _infer_dimension(0, nstate or 1,
+  nroll = _infer_dimension(0, nroll or 1,
-                            initial_state=initial_state,
+                           initial_state=initial_state,
-                            initial_time=initial_time,
+                           initial_warmstart=initial_warmstart,
-                            initial_warmstart=initial_warmstart,
+                           control=control,
-                            ctrl=ctrl,
+                           state=state,
-                            qfrc_applied=qfrc_applied,
+                           sensordata=sensordata)
-                            xfrc_applied=xfrc_applied,
-                            mocap=mocap,
-                            state=state,
-                            sensordata=sensordata)
 
   # infer nstep, check for incompatibilities
   nstep = _infer_dimension(1, nstep or 1,
-                           ctrl=ctrl,
+                           control=control,
-                           qfrc_applied=qfrc_applied,
-                           xfrc_applied=xfrc_applied,
-                           mocap=mocap,
                            state=state,
                            sensordata=sensordata)
 
   # tile input arrays if required (singleton expansion)
-  initial_state = _tile_if_required(initial_state, nstate)
+  initial_state = _tile_if_required(initial_state, nroll)
-  initial_time = _tile_if_required(initial_time, nstate)
+  initial_warmstart = _tile_if_required(initial_warmstart, nroll)
-  initial_warmstart = _tile_if_required(initial_warmstart, nstate)
+  control = _tile_if_required(control, nroll, nstep)
-  ctrl = _tile_if_required(ctrl, nstate, nstep)
-  qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep)
-  xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep)
-  mocap = _tile_if_required(mocap, nstate, nstep)
 
   # allocate output if not provided
   if state is None:
-    state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+    state = np.empty((nroll, nstep, nstate))
   if sensordata is None:
-    sensordata = np.empty((nstate, nstep, model.nsensordata))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
 
   # call rollout
-  _rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
+  _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
-                   initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
+                   initial_warmstart, control, state, sensordata)
-                   state, sensordata)
+
+  # return outputs
+  return state, sensordata
 
-  # return squeezed outputs
-  return state.squeeze(), sensordata.squeeze()
 
 def _check_must_be_numeric(**kwargs):
   for key, value in kwargs.items():
@@ -148,6 +160,7 @@ def _check_must_be_numeric(**kwargs):
     if not isinstance(value, np.ndarray) and not isinstance(value, float):
       raise ValueError(f'{key} must be a numpy array or float')
 
+
 def _check_number_of_dimensions(ndim, **kwargs):
   for key, value in kwargs.items():
     if value is None:
@@ -155,12 +168,16 @@ def _check_number_of_dimensions(ndim, **kwargs):
     if value.ndim > ndim:
       raise ValueError(f'{key} can have at most {ndim} dimensions')
 
+
 def _check_trailing_dimension(dim, **kwargs):
   for key, value in kwargs.items():
     if value is None:
       continue
     if value.shape[-1] != dim:
-      raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}')
+      raise ValueError(
+          f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}'
+      )
+
 
 def _ensure_2d(arg):
   if arg is None:
@@ -168,6 +185,7 @@ def _ensure_2d(arg):
   else:
     return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
 
+
 def _ensure_3d(arg):
   if arg is None:
     return None
@@ -181,7 +199,22 @@ def _ensure_3d(arg):
       arg = arg[np.newaxis, ...]
     return np.ascontiguousarray(arg, dtype=np.float64)
 
+
 def _infer_dimension(dim, value, **kwargs):
+  """Infers dimension `dim` given guess `value` from set of arrays.
+
+  Args:
+    dim: Dimension to be inferred.
+    value: Initial guess of inferred value (1: unknown).
+    **kwargs: List of arrays which should all have the same size (or 1)
+      along dimension dim.
+
+  Returns:
+    Inferred dimension.
+
+  Raises:
+    ValueError: If mismatch between array shapes or initial guess.
+  """
   for name, array in kwargs.items():
     if array is None:
       continue
@@ -190,10 +223,12 @@ def _infer_dimension(dim, value, **kwargs):
         value = array.shape[dim]
       elif array.shape[dim] != 1:
         raise ValueError(
-            f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}'
+            f'dimension {dim} inferred as {value} '
+            f'but {name} has {array.shape[dim]}'
         )
   return value
 
+
 def _tile_if_required(array, dim0, dim1=None):
   if array is None:
     return

python/mujoco/rollout_test.py

@@ -14,15 +14,16 @@
 # ==============================================================================
 """tests for rollout function."""
 
+import concurrent.futures
+import threading
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import mujoco
-import numpy as np
-import concurrent.futures
-import threading
 from mujoco import rollout
+import numpy as np
 
-#--------------------------- models used for testing ---------------------------
+# -------------------------- models used for testing ---------------------------
 
 TEST_XML = r"""
 <mujoco>
@@ -96,7 +97,7 @@ TEST_XML_MOCAP = r"""
   </worldbody>
   <sensor>
     <framepos objtype="xbody" objname="1"/>
-    <framequat objtype="xbody" objname="2"/>
+    <framequat objtype="xbody" objname="1"/>
   </sensor>
 </mujoco>
 """
@@ -106,12 +107,33 @@ TEST_XML_EMPTY = r"""
 </mujoco>
 """
 
+TEST_XML_DIVERGE = r"""
+<mujoco>
+  <option>
+    <flag gravity="disable"/>
+  </option>
+
+  <worldbody>
+    <geom type="plane" size="5 5 .1"/>
+    <body pos="0 0 -.3" euler="30 45 90">
+      <freejoint/>
+      <geom type="box" size=".1 .2 .4"/>
+    </body>
+  </worldbody>
+
+  <keyframe>
+    <key name="non-diverging" qpos="0 0 .5 1 0 0 0"/>
+  </keyframe>
+</mujoco>
+"""
+
 ALL_MODELS = {'TEST_XML': TEST_XML,
               'TEST_XML_NO_SENSORS': TEST_XML_NO_SENSORS,
               'TEST_XML_NO_ACTUATORS': TEST_XML_NO_ACTUATORS,
               'TEST_XML_EMPTY': TEST_XML_EMPTY}
 
-#------------------------------- tests -----------------------------------------
+# ------------------------------ tests -----------------------------------------
+
 
 class MuJoCoRolloutTest(parameterized.TestCase):
 
@@ -119,179 +141,209 @@ class MuJoCoRolloutTest(parameterized.TestCase):
     super().setUp()
     np.random.seed(42)
 
-  #----------------------------- test basic operation
+  # ----------------------------- test basic operation
 
   @parameterized.parameters(ALL_MODELS.keys())
   def test_single_step(self, model_name):
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nstate)
-    ctrl = np.random.randn(model.nu)
+    control = np.random.randn(model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
     mujoco.mj_resetData(model, data)
-    py_state, py_sensordata = step(model, data, initial_state, ctrl=ctrl)
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
-
-
   @parameterized.parameters(ALL_MODELS.keys())
-  def test_single_rollout(self, model_name):
+  def test_one_rollout(self, model_name):
     nstep = 3  # number of timesteps
 
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nstate)
-    ctrl = np.random.randn(nstep, model.nu)
+    control = np.random.randn(nstep, model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
-    py_state, py_sensordata = single_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                             ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
-    np.testing.assert_array_equal(state, np.asarray(py_state))
+    np.testing.assert_array_equal(sensordata, py_sensordata)
-    np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata))
 
   @parameterized.parameters(ALL_MODELS.keys())
   def test_multi_step(self, model_name):
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    nstate = 5  # number of initial states
+    nroll = 5  # number of rollouts
+    nstep = 1  # number of steps
 
-    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nroll, nstate)
-    ctrl = np.random.randn(nstate, 1, model.nu)
+    control = np.random.randn(nroll, nstep, model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
     mujoco.mj_resetData(model, data)
-    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                            ctrl=ctrl)
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
   @parameterized.parameters(ALL_MODELS.keys())
-  def test_single_rollout_fixed_ctrl(self, model_name):
+  def test_one_rollout_fixed_ctrl(self, model_name):
-    nstep = 3
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    nroll = 1  # number of rollouts
-    ctrl = np.random.randn(model.nu)
+    nstep = 3  # number of steps
-    state = np.empty((nstep, model.nq + model.nv + model.na))
+
-    sensordata = np.empty((nstep, model.nsensordata))
+    initial_state = np.random.randn(nstate)
-    rollout.rollout(model, data, initial_state, ctrl,
+    control = np.random.randn(model.nu)
+    state = np.empty((nroll, nstep, nstate))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
+    rollout.rollout(model, data, initial_state, control,
                     state=state, sensordata=sensordata)
 
-    ctrl = np.tile(ctrl, (nstep, 1))  # repeat??
+    control = np.tile(control, (nstep, 1))
-    py_state, py_sensordata = single_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                             ctrl=ctrl)
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
   @parameterized.parameters(ALL_MODELS.keys())
   def test_multi_rollout(self, model_name):
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    nstate = 2  # number of initial states
+    nroll = 2  # number of initial states
     nstep = 3  # number of timesteps
 
-    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nroll, nstate)
-    ctrl = np.random.randn(nstate, nstep, model.nu)
+    control = np.random.randn(nroll, nstep, model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
-    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                            ctrl=ctrl)
+    np.testing.assert_array_equal(state, py_state)
-    np.testing.assert_array_equal(py_state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
-    np.testing.assert_array_equal(py_sensordata, py_sensordata)
 
   @parameterized.parameters(ALL_MODELS.keys())
   def test_multi_rollout_fixed_ctrl_infer_from_output(self, model_name):
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    nstate = 2  # number of initial states
+    nroll = 2  # number of rollouts
     nstep = 3  # number of timesteps
 
-    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nroll, nstate)
-    ctrl = np.random.randn(nstate, 1, model.nu)  # 1 control in the time dimension
+    control = np.random.randn(nroll, 1, model.nu)
-    state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+    state = np.empty((nroll, nstep, nstate))
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+    state, sensordata = rollout.rollout(model, data, initial_state, control,
                                         state=state)
 
-    ctrl = np.repeat(ctrl, nstep, axis=1)
+    control = np.repeat(control, nstep, axis=1)
-    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                            ctrl=ctrl)
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
-  @parameterized.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]],
+  @parameterized.parameters(ALL_MODELS.keys())
-                         model_name=list(ALL_MODELS.keys()))
+  def test_py_rollout_generalized_control(self, model_name):
-  def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name):
     model = mujoco.MjModel.from_xml_string(ALL_MODELS[model_name])
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    nstate = 4  # number of initial states
+    nroll = 4  # number of rollouts
+    nstep = 3  # number of timesteps
 
-    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nroll, nstate)
 
-    # arg_nstep is the horizon for {ctrl, qfrc_applied, xfrc_applied}, respectively
+    control_spec = (mujoco.mjtState.mjSTATE_CTRL |
-    ctrl = np.random.randn(nstate, arg_nstep[0], model.nu)
+                    mujoco.mjtState.mjSTATE_QFRC_APPLIED |
-    qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv)
+                    mujoco.mjtState.mjSTATE_XFRC_APPLIED)
-    xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6)
+    ncontrol = mujoco.mj_stateSize(model, control_spec)
+    control = np.random.randn(nroll, nstep, ncontrol)
 
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+    state, sensordata = rollout.rollout(model, data, initial_state, control,
-                                        qfrc_applied=qfrc_applied,
+                                        control_spec=control_spec)
-                                        xfrc_applied=xfrc_applied)
 
-    # tile singleton arguments
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control,
-    nstep = max(arg_nstep)
+                                         control_spec=control_spec)
-    if arg_nstep[0] == 1:
-      ctrl = np.repeat(ctrl, nstep, axis=1)
-    if arg_nstep[1] == 1:
-      qfrc_applied = np.repeat(qfrc_applied, nstep, axis=1)
-    if arg_nstep[2] == 1:
-      xfrc_applied = np.repeat(xfrc_applied, nstep, axis=1)
-
-    py_state, py_sensordata = multi_rollout(model, data, initial_state,
-                                            ctrl=ctrl,
-                                            qfrc_applied=qfrc_applied,
-                                            xfrc_applied=xfrc_applied)
     np.testing.assert_array_equal(state, py_state)
     np.testing.assert_array_equal(sensordata, py_sensordata)
 
-  #----------------------------- test threaded operation
+  def test_detect_divergence(self):
+    model = mujoco.MjModel.from_xml_string(TEST_XML_DIVERGE)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    data = mujoco.MjData(model)
+
+    nroll = 4  # number of rollouts
+    initial_state = np.empty((nroll, nstate))
+
+    # get diverging (0, 2) and non-diverging (1, 3) states
+    mujoco.mj_getState(model, data, initial_state[0],
+                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(model, data, initial_state[2],
+                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_resetDataKeyframe(model, data, 0)  # keyframe 0 does not diverge
+    mujoco.mj_getState(model, data, initial_state[1],
+                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+    mujoco.mj_getState(model, data, initial_state[3],
+                       mujoco.mjtState.mjSTATE_FULLPHYSICS)
+
+    nstep = 10000  # divergence after ~15s, timestep = 2e-3
+
+    state = np.random.randn(nroll, nstep, nstate)
+
+    rollout.rollout(model, data, initial_state, state=state)
+
+    # initial_state[0,2] diverged, final timesteps are identical
+    assert state[0][-1][0] == state[0][-2][0]
+    assert state[2][-1][0] == state[2][-2][0]
+
+    # initial_state[1,3] did not diverge, final timesteps are different
+    assert state[1][-1][0] != state[1][-2][0]
+    assert state[3][-1][0] != state[3][-2][0]
+
+  # ----------------------------- test threaded operation
 
   def test_threading(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     num_workers = 32
-    nstate = 10000
+    nroll = 10000
     nstep = 5
-    initial_state = np.random.randn(nstate, model.nq+model.nv+model.na)
+    initial_state = np.random.randn(nroll, nstate)
-    state = np.zeros((nstate, nstep, model.nq+model.nv+model.na))
+    state = np.empty((nroll, nstep, nstate))
-    sensordata = np.zeros((nstate, nstep, model.nsensordata))
+    sensordata = np.empty((nroll, nstep, model.nsensordata))
-    ctrl = np.random.randn(nstate, nstep, model.nu)
+    control = np.random.randn(nroll, nstep, model.nu)
 
     thread_local = threading.local()
 
     def thread_initializer():
       thread_local.data = mujoco.MjData(model)
 
-    def call_rollout(initial_state, ctrl, state):
+    def call_rollout(initial_state, control, state, sensordata):
-      rollout.rollout(model, thread_local.data, skip_checks=True,
+      rollout.rollout(model, thread_local.data, initial_state, control,
-                      nstate=initial_state.shape[0], nstep=nstep,
+                      skip_checks=True, nroll=initial_state.shape[0],
-                      initial_state=initial_state, ctrl=ctrl, state=state)
+                      nstep=nstep, state=state, sensordata=sensordata)
 
-    n = initial_state.shape[0] // num_workers  # integer division
+    n = nroll // num_workers  # integer division
     chunks = []  # a list of tuples, one per worker
     for i in range(num_workers-1):
-      chunks.append(
+      chunks.append((initial_state[i*n:(i+1)*n],
-          (initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n]))
+                     control[i*n:(i+1)*n],
+                     state[i*n:(i+1)*n],
+                     sensordata[i*n:(i+1)*n]))
+
     # last chunk, absorbing the remainder:
-    chunks.append(
+    chunks.append((initial_state[(num_workers-1)*n:],
-        (initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:],
+                   control[(num_workers-1)*n:],
-         state[(num_workers-1)*n:]))
+                   state[(num_workers-1)*n:],
+                   sensordata[(num_workers-1)*n:]))
 
     with concurrent.futures.ThreadPoolExecutor(
         max_workers=num_workers, initializer=thread_initializer) as executor:
@@ -302,187 +354,237 @@ class MuJoCoRolloutTest(parameterized.TestCase):
         future.result()
 
     data = mujoco.MjData(model)
-    py_state, py_sensordata = multi_rollout(model, data, initial_state,
+    py_state, py_sensordata = py_rollout(model, data, initial_state, control)
-                                            ctrl=ctrl)
     np.testing.assert_array_equal(state, py_state)
+    np.testing.assert_array_equal(sensordata, py_sensordata)
 
-  #----------------------------- test advanced operation
+  # ---------------------------- test advanced operation
-
-  def test_time(self):
-    model = mujoco.MjModel.from_xml_string(TEST_XML)
-    data = mujoco.MjData(model)
-
-    nstate = 1
-    nstep = 3
-
-    initial_time = np.array([[2.]])
-    initial_state = np.random.randn(nstate, model.nq + model.nv + model.na)
-    ctrl = np.random.randn(nstate, nstep, model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
-                                        initial_time=initial_time)
-
-    self.assertAlmostEqual(data.time, 2 + nstep*model.opt.timestep)
 
   def test_warmstart(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    state0 = np.zeros(model.nq + model.nv + model.na)
+    # take one step, save the state
-    ctrl = np.zeros(model.nu)
+    state0 = np.zeros(nstate)
-    state1, _ = step(model, data, state0, ctrl=ctrl)
+    control = np.zeros(model.nu)
+    state1, _ = step(model, data, state0, control)
+
+    # save qacc_warmstart
     initial_warmstart = data.qacc_warmstart.copy()
 
-    state2, _ = step(model, data, state1, ctrl=ctrl)
+    # take one more step (uses correct warmstart)
+    state2, _ = step(model, data, state1[0], control)
 
-    state, _ = rollout.rollout(model, data, state1, ctrl)
+    # take step using rollout, don't take warmstart into account
-    assert np.linalg.norm(state-state2) > 0
+    state, _ = rollout.rollout(model, data, state1[0], control)
 
-    state, _ = rollout.rollout(model, data, state1, ctrl,
+    # assert that stepping without warmstarts is not exact
+    np.testing.assert_raises(AssertionError,
+                             np.testing.assert_array_equal, state, state2)
+
+    # take step using rollout, take warmstart into account
+    state, _ = rollout.rollout(model, data, state1, control,
                                initial_warmstart=initial_warmstart)
-    np.testing.assert_array_equal(state, state2)
+
+    # assert exact equality
+    np.testing.assert_array_equal(state, np.expand_dims(state2, axis=0))
 
   def test_mocap(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML_MOCAP)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.zeros(model.nq + model.nv + model.na)
+    initial_state = np.zeros(nstate)
+
+    control_spec = (mujoco.mjtState.mjSTATE_MOCAP_POS |
+                    mujoco.mjtState.mjSTATE_MOCAP_QUAT)
+
     pos1 = np.array((1., 2., 3.))
     quat1 = np.array((1., 2., 3., 4.))
     quat1 /= np.linalg.norm(quat1)
     pos2 = np.array((2., 3., 4.))
     quat2 = np.array((2., 3., 4., 5.))
     quat2 /= np.linalg.norm(quat2)
-    mocap = np.hstack((pos1, quat1, pos2, quat2))
+    control = np.hstack((pos1, pos2, quat1, quat2))
 
-    state, sensordata = rollout.rollout(model, data, initial_state, mocap=mocap)
+    _, sensordata = rollout.rollout(model, data, initial_state, control,
+                                    control_spec=control_spec)
 
-    np.testing.assert_array_almost_equal(sensordata[:3], pos1)
+    np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1)
-    np.testing.assert_array_almost_equal(sensordata[3:], quat2)
+    np.testing.assert_array_almost_equal(sensordata[0][0][3:], quat1)
 
-  #----------------------------- test correctness
+  # ---------------------------- test correctness
 
   def test_intercept_mj_errors(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.zeros(model.nq + model.nv + model.na)
+    nroll = 1
-    ctrl = np.zeros((3, model.nu))
+    nstep = 3
+
+    initial_state = np.zeros((nroll, nstate))
+    ctrl = np.zeros((nroll, nstep, model.nu))
 
     model.opt.solver = 10  # invalid solver type
     with self.assertRaisesWithLiteralMatch(
         mujoco.FatalError, 'mj_fwdConstraint: unknown solver type 10'):
-      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+      rollout.rollout(model, data, initial_state, ctrl)
 
   def test_invalid(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.zeros(model.nq + model.nv + model.na)
+    nroll = 1
 
-    ctrl = 'string'
+    initial_state = np.zeros((nroll, nstate))
-    with self.assertRaisesWithLiteralMatch(
-        ValueError, 'ctrl must be a numpy array or float'):
-      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
 
-    qfrc_applied = np.zeros((2, 3, 4, 5))
+    control = 'string'
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'qfrc_applied can have at most 3 dimensions'):
+        ValueError, 'control must be a numpy array or float'):
-      state, sensordata = rollout.rollout(model, data, initial_state,
+      rollout.rollout(model, data, initial_state, control)
-                                          qfrc_applied=qfrc_applied)
+
+    control = np.zeros((2, 3, 4, 5))
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'control can have at most 3 dimensions'):
+      rollout.rollout(model, data, initial_state, control)
 
   def test_bad_sizes(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    initial_state = np.random.randn(model.nq + model.nv + model.na+1)
+    nroll = 1
-    with self.assertRaisesWithLiteralMatch(
+    nstep = 3
-        ValueError, 'trailing dimension of initial_state must be 5, got 6'):
-      state, sensordata = rollout.rollout(model, data, initial_state)
 
-    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nroll, nstate + 1)
-    ctrl = np.random.randn(model.nu+1)
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'trailing dimension of ctrl must be 2, got 3'):
+        ValueError, 'trailing dimension of initial_state must be 6, got 7'):
-      state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+      rollout.rollout(model, data, initial_state)
 
-    ctrl = np.random.randn(2, model.nu)
+    initial_state = np.random.randn(nroll, nstate)
-    qfrc_applied = np.random.randn(3, model.nv)  # incompatible horizon
+    control = np.random.randn(1, nstep, model.nu + 1)
     with self.assertRaisesWithLiteralMatch(
-        ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'):
+        ValueError, 'trailing dimension of control must be 2, got 3'):
-      state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
+      rollout.rollout(model, data, initial_state, control)
-                                          qfrc_applied=qfrc_applied)
+
+    control = np.random.randn(nroll, nstep, model.nu)
+    state = np.random.randn(nroll, nstep+1, nstate)  # incompatible nstep
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'dimension 1 inferred as 3 but state has 4'):
+      rollout.rollout(model, data, initial_state, control, state=state)
+
+    initial_state = np.random.randn(nroll, nstate)
+    control = np.random.randn(nroll, nstep, model.nu)
+    bad_spec = mujoco.mjtState.mjSTATE_ACT
+    with self.assertRaisesWithLiteralMatch(
+        ValueError, 'control_spec can only contain bits in mjSTATE_USER'):
+      rollout.rollout(model, data, initial_state, control,
+                      control_spec=bad_spec)
 
   def test_stateless(self):
     model = mujoco.MjModel.from_xml_string(TEST_XML)
-    model.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART.value
+    nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
     data = mujoco.MjData(model)
 
-    # call step with a clean mjData
+    # step with a clean mjData
-    initial_state = np.random.randn(model.nq + model.nv + model.na)
+    initial_state = np.random.randn(nstate)
-    ctrl = np.random.randn(model.nu)
+    control = np.random.randn(3, 3, model.nu)
-    state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+    state, sensordata = rollout.rollout(model, data, initial_state, control)
 
-    # fill mjData with some debug value, see that we still get the same outputs
+    # fill user fields with random values
-    mujoco.mj_resetDataDebug(model, data, 255)
+    for attr in [
-    debug_state, debug_sensordata = rollout.rollout(model, data, initial_state,
+        'ctrl',
-                                                    ctrl)
+        'qfrc_applied',
+        'xfrc_applied',
+        'mocap_pos',
+        'mocap_quat',
+    ]:
+      setattr(data, attr, np.random.randn(*getattr(data, attr).shape))
 
-    np.testing.assert_array_equal(state, debug_state)
+    # roll out again
-    np.testing.assert_array_equal(sensordata, debug_sensordata)
+    state2, sensordata2 = rollout.rollout(model, data, initial_state, control)
+
+    # assert that we still get the same outputs
+    np.testing.assert_array_equal(state, state2)
+    np.testing.assert_array_equal(sensordata, sensordata2)
 
 
-#--------------- Python implementation of rollout functionality ----------------
+# -------------- Python implementation of rollout functionality ----------------
 
-def get_state(data):
-  return np.hstack((data.qpos, data.qvel, data.act))
 
-def set_state(model, data, state):
+def get_state(model, data):
-  data.qpos = state[:model.nq]
+  nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
-  data.qvel = state[model.nq:model.nq+model.nv]
+  state = np.empty(nstate)
-  data.act = state[model.nq+model.nv:model.nq+model.nv+model.na]
+  mujoco.mj_getState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+  return state.reshape((1, nstate))
 
-def step(model, data, state, **kwargs):
+
+def step(model, data, state, control,
+         control_spec=mujoco.mjtState.mjSTATE_CTRL):
   if state is not None:
-    set_state(model, data, state)
+    mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
-  for key, value in kwargs.items():
+  mujoco.mj_setState(model, data, control, control_spec)
-    if value is not None:
-      setattr(data, key, np.reshape(value, getattr(data, key).shape))
   mujoco.mj_step(model, data)
-  return (get_state(data), data.sensordata)
+  return (get_state(model, data), data.sensordata)
 
-def single_rollout(model, data, initial_state, **kwargs):
-  arg_nstep = set([a.shape[0] for a in kwargs.values()])
-  assert len(arg_nstep) == 1  # nstep dimensions must match
-  nstep = arg_nstep.pop()
 
-  state = np.empty((nstep, model.nq + model.nv + model.na))
+def one_rollout(model, data, initial_state, control,
+                control_spec=mujoco.mjtState.mjSTATE_CTRL):
+  nstep = control.shape[0]
+  nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+  state = np.empty((nstep, nstate))
   sensordata = np.empty((nstep, model.nsensordata))
 
   mujoco.mj_resetData(model, data)
   for t in range(nstep):
-    kwargs_t = {}
-    for key, value in kwargs.items():
-      kwargs_t[key] = value[0 if value.ndim == 1 else t]
     state[t], sensordata[t] = step(model, data,
-                                   initial_state if t==0 else None,
+                                   initial_state if t == 0 else None,
-                                   **kwargs_t)
+                                   control[t], control_spec)
   return state, sensordata
 
-def multi_rollout(model, data, initial_state, **kwargs):
-  nstate = initial_state.shape[0]
-  arg_nstep = set([a.shape[1] for a in kwargs.values()])
-  assert len(arg_nstep) == 1  # nstep dimensions must match
-  nstep = arg_nstep.pop()
 
-  state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
+def ensure_2d(arg):
-  sensordata = np.empty((nstate, nstep, model.nsensordata))
+  if arg is None:
-  for s in range(nstate):
+    return None
-    kwargs_s = {key : value[s] for key, value in kwargs.items()}
+  else:
-    state_s, sensordata_s = single_rollout(model, data, initial_state[s],
+    return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
-                                           **kwargs_s)
+
-    state[s] = state_s
+
-    sensordata[s] = sensordata_s
+def ensure_3d(arg):
-  return state.squeeze(), sensordata.squeeze()
+  if arg is None:
+    return None
+  else:
+    # np.atleast_3d adds both leading and trailing dims, we want only leading
+    if arg.ndim == 0:
+      arg = arg[np.newaxis, np.newaxis, np.newaxis, ...]
+    elif arg.ndim == 1:
+      arg = arg[np.newaxis, np.newaxis, ...]
+    elif arg.ndim == 2:
+      arg = arg[np.newaxis, ...]
+    return np.ascontiguousarray(arg, dtype=np.float64)
+
+
+def py_rollout(model, data, initial_state, control,
+               control_spec=mujoco.mjtState.mjSTATE_CTRL):
+  initial_state = ensure_2d(initial_state)
+  control = ensure_3d(control)
+  nroll = initial_state.shape[0]
+  nstep = control.shape[1]
+  nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
+
+  state = np.empty((nroll, nstep, nstate))
+  sensordata = np.empty((nroll, nstep, model.nsensordata))
+  for r in range(nroll):
+    state_r, sensordata_r = one_rollout(
+        model, data, initial_state[r], control[r], control_spec
+    )
+    state[r] = state_r
+    sensordata[r] = sensordata_r
+  return state, sensordata
+
 
 if __name__ == '__main__':
   absltest.main()

python/mujoco/structs.cc

@@ -1400,6 +1400,7 @@ PYBIND11_MODULE(_structs, m) {
   X(offwidth);
   X(offheight);
   X(ellipsoidinertia);
+  X(bvactive);
 #undef X
 
   py::class_<raw::MjVisualQuality> mjVisualQuality(mjVisual, "Quality");
@@ -1529,6 +1530,8 @@ PYBIND11_MODULE(_structs, m) {
   X(slidercrank);
   X(crankbroken);
   X(frustum);
+  X(bv);
+  X(bvactive);
 #undef X
 
 #define X(var)                    \