chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Abhishek Joshi
2024-02-12 06:04:57 -08:00
committed by GitHub
parent ff6581df1c ab209c41be
commit 18edadb975
170 changed files with 7885 additions and 3961 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+1 -1
View File
@@ -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
+2 -2
View File
@@ -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
Vendored
+4 -4
View File
@@ -1,6 +1,6 @@
1 VERSIONINFO
FILEVERSION 3,1,2,0
PRODUCTVERSION 3,1,2,0
FILEVERSION 3,1,3,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
Vendored
+4 -4
View File
@@ -1,8 +1,8 @@
MUJOCO ICON "mujoco.ico"
1 VERSIONINFO
FILEVERSION 3,1,2,0
PRODUCTVERSION 3,1,2,0
FILEVERSION 3,1,3,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
+2 -2
View File
@@ -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
+385 -345
View File
@@ -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
are not referenced in another MJCF element. If a discarded geom was used for inferring body inertia, an explicit
:ref:`inertial<body-inertial>` element is added to the body.
- All geoms with :ref:`contype<body-geom-contype>` |-| = |-| :ref:`conaffinity<body-geom-conaffinity>` |-| =0 are
discarded, if they are not referenced in another MJCF element. If a discarded geom was used for inferring 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
raycasting, as used for example by :ref:`rangefinder<sensor-rangefinder>`, since raycasting reports distances to
visual geoms. When visualizing models compiled with this flag, it is important to remember that colliding geoms are
often placed in a :ref:`group<body-geom-group>` which is invisible by default.
The resulting compiled model will have exactly the same dynamics as the original model. The only engine-level
computation which might change is the output of :ref:`raycasting<mj_ray>` computations, as used for example by
:ref:`rangefinder<sensor-rangefinder>` sensors, since raycasting reports distances to visual geoms. When visualizing
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,
cranksite.
| All general attributes are available here except: name, class, joint, jointinparent, site, refsite, tendon,
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,
slidersite, cranksite.
:ref:`motor <actuator-motor>` attributes are available here except: name, class, joint, jointinparent, site, refsite,
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
+5 -3
View File
@@ -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
+61 -17
View File
@@ -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
:ref:discardvisual<compiler-discardvisual> for details.
1. Improved the :ref:`discardvisual<compiler-discardvisual>` compiler flag, which now discards all visual-only assets.
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``.
3. Added :at:`site` transmission.
4. Updated MJX colab tutorial with more stable quadruped environment.
5. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, and boxes.
5. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
6. Added :at:`site` transmission.
7. Updated MJX colab tutorial with more stable quadruped environment.
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
:github:issue:`1270`.
14. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
: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
robust for very small or large geom sizes.
for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is
more robust for very small or large geom sizes.
2. Added :ref:`frame<frame>` to MJCF, a :ref:`meta-element<meta-element>` which defines a pure coordinate transformation
on its direct children, without requiring a :ref:`body<body>`.
3. Added the :at:`kv` attribute to the :ref:`position<actuator-position>` and :ref:`intvelocity<actuator-intvelocity>`
@@ -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
supported version is still 11.0, but we release these wheels to fix compatibility for those users. See
11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where ``SYSTEM_VERSION_COMPAT`` is set.
The minimum supported version is still 11.0, but we release these wheels to fix compatibility for those users. See
:github:issue:`1213`.
Version 3.0.1 (November 15, 2023)
@@ -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
+3 -3
View File
@@ -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.
Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality constraints to
the rest of the system, but that would be less efficient.
<https://github.com/google-deepmind/mujoco/blob/main/model/slider_crank/slider_crank.xml>`__ contains pedagogical
examples. Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality
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
+4 -2
View File
@@ -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
+19 -1
View File
@@ -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
+49 -19
View File
@@ -7,8 +7,11 @@ Model Gallery
MuJoCo Menagerie
----------------
Menagerie is a collection of high-quality simulation models for MuJoCo, curated
by Google DeepMind.
A physics simulator is only as good as the model it is simulating, and in a
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
+29 -18
View File
@@ -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.
This is separate from the ``mjvScene`` that the viewer uses internally to render the final scene, and is entirely
under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or :ref:`mjv_makeConnector` to add
visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer will incorporate
these geoms to future rendered images. For example:
- ``user_scn``: an :ref:`mjvScene` object that allows users to add change rendering flags and add custom
visualization geoms to the rendered scene. This is separate from the ``mjvScene`` that the viewer uses internally to
render the final scene, and is entirely under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or
:ref:`mjv_makeConnector` to add visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer
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
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 canonical usage form is
sample, implemented in `rollout.cc <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__
and wrapped in `rollout.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__,
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 =
model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities
(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control
sequences.
``initial_state`` is a ``nroll x nstate`` array, with ``nroll`` initial states of size ``nstate``, where
``nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`` is the size of the
:ref:`full physics state<geFullPhysics>`. ``control`` is a ``nroll x nstep x ncontrol`` array of controls. Controls are
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,
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.
values already present in the given ``MjData`` instance will have no effect on the output.
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
threaded operation.
`rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`__ for an example
of threaded operation (and more generally for usage examples).
.. _PyMjpy_migration:
doc/requirements.txt
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+1 -2
View File
@@ -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
+3
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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
+2 -3
View File
@@ -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', 15),
('mjNTIMER', 16),
('mjTIMER_COL_NARROW', 14),
('mjNTIMER', 15),
]),
)),
('mjtCatBit',
introspect/structs.py
+22 -1
View File
@@ -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
View File
@@ -1,2 +1 @@
recursive-include mujoco/mjx/test_data *
recursive-include mujoco/mjx/benchmark *.obj *.stl *.xml
mjx/mujoco/mjx/__init__.py
+5 -1
View File
@@ -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
+14 -6
View File
@@ -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]
end1 = m.body_geomadr[b1] + m.body_geomnum[b1]
for g1 in range(start1, end1):
start2 = m.body_geomadr[b2]
end2 = m.body_geomadr[b2] + m.body_geomnum[b2]
for g2 in range(start2, end2):
for g1 in range(b_start[b1], b_end[b1]):
if not geom_con[g1]:
continue
for g2 in range(b_start[b2], b_end[b2]):
if not geom_con[g2]:
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
+5 -2
View File
@@ -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
+8 -5
View File
@@ -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
+3
View File
@@ -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
+5 -3
View File
@@ -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
+3 -3
View File
@@ -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 = smooth.factor_m(m, d, mh)
dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
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
+86 -25
View File
@@ -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,
qLD=zero_nm,
qLDiagInv=zero_nv,
qLDiagSqrtInv=zero_nv,
qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
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
ncon = (dx_i.contact.dist <= 0).sum()
efc_active = (dx_i.efc_J != 0).any(axis=1)
d_i = jax.tree_map(lambda x, i=i: x[i], d) if batched else d
result_i = result[i] if batched else result
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()
d_i = mujoco.MjData(m)
d_i.nnzJ = nefc * m.nv
mujoco._functions._realloc_con_efc(d_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access
d_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
d_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
d_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
nefc, nc = int(efc_active.sum()), int((efc_active & efc_con).sum())
result_i.nnzJ = nefc * m.nv
if ncon != result_i.ncon or nefc != result_i.nefc:
mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc) # pylint: disable=protected-access
result_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
result_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
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]
ds.append(d_i)
return ds if batched else ds[0]
result_i.efc_type[:] = efc_type[efc_active]
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
+46 -9
View File
@@ -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.qLD.shape, (nm,))
self.assertEqual(d.qLDiagInv.shape, (nv,))
self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
self.assertEqual(d.qM.shape, (nv, nv))
self.assertEqual(d.qLD.shape, (nv, nv))
self.assertEqual(d.qLDiagInv.shape, (0,))
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)
np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
dx_sparse = mjx.put_data(m, d)
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
+24
View File
@@ -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
+79 -8
View File
@@ -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(
size: jax.Array,
def _ray_triangle(
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."""
del size, pnt, vec
raise NotImplementedError("ray <> mesh not implemented yet")
"""Returns the distance at which a ray intersects with a triangle."""
# project difference vectors in ray normal plane
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_]
dist = jax.vmap(fn)(size, pnt, vec)
args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
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
+37
View File
@@ -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
+21 -76
View File
@@ -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
subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass))
cond = jp.tile(mass < mujoco.mjMINVAL, (3, 1)).T
# 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]))]
return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)])
inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
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)
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)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.replace(qM=qm)
return d
def factor_m(
m: Model,
d: Data,
qM: jax.Array, # pylint:disable=invalid-name
) -> Data:
"""Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd."""
def factor_m(m: Model, d: Data) -> Data:
"""Gets factorizaton of inertia-like matrix M, assumed spd."""
if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d.qM)
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
+3 -18
View File
@@ -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, dx, dx.qM)
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
_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
+7 -12
View File
@@ -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 = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
ma = support.mul_m(m, d, 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 = smooth.dense_m(m, d) + h
h_ = jax.scipy.linalg.cho_factor(h)
mgrad = jax.scipy.linalg.cho_solve(h_, grad)
h = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T)
dh = d.replace(qM=h)
dh = smooth.factor_m(m, dh)
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
+34 -6
View File
@@ -39,37 +39,65 @@ def _assert_attr_eq(a, b, attr):
class SolverTest(absltest.TestCase):
def test_solver(self):
"""Test solver."""
def test_newton(self):
"""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
m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
def test_no_warmstart(self):
"""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
+108 -1
View File
@@ -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
+41 -1
View File
@@ -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
+1 -8
View File
@@ -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
-38
View File
@@ -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
+38 -6
View File
@@ -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,)
qLD: L'*D*L factorization of M (sparse) (nM,)
qLDiagInv: 1/diag(D) (nv,)
qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,)
qM: total inertia if sparse: (nM,)
if dense: (nv, nv)
qLD: L'*D*L (or Cholesky) factorization of M. if sparse: (nM,)
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
-14
View File
@@ -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
-132
View File
@@ -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/benchmark/model/barkour_v0/README.md → mjx/mujoco/mjx/test_data/barkour_v0/README.md
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/abduction.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/abduction.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml → mjx/mujoco/mjx/test_data/barkour_v0/assets/barkour_v0_mjx.xml
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/body.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/body.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/foot.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/foot.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/handle.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/handle.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/head.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head_mount.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/head_mount.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/joystick.gif → mjx/mujoco/mjx/test_data/barkour_v0/assets/joystick.gif
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 4.0 MiB

After

Width:  |  Height:  |  Size: 4.0 MiB

mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/lower_leg_1to1.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/lower_leg_1to1.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/powercable.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/powercable.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_1.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_left_1.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_2.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_left_2.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_3.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_left_3.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_1.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_right_1.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_2.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_right_2.stl
View File
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_3.stl → mjx/mujoco/mjx/test_data/barkour_v0/assets/upper_right_3.stl
View File
mjx/mujoco/mjx/benchmark/model/humanoid/README.md → mjx/mujoco/mjx/test_data/humanoid/README.md
View File
mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.png → mjx/mujoco/mjx/test_data/humanoid/humanoid.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 388 KiB

After

Width:  |  Height:  |  Size: 388 KiB

mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.xml → mjx/mujoco/mjx/test_data/humanoid/humanoid.xml
View File
mjx/mujoco/mjx/test_data/ray.xml
+2
View File
@@ -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/benchmark/model/shadow_hand/README.md → mjx/mujoco/mjx/test_data/shadow_hand/README.md
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_distal_pst.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/f_distal_pst.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_knuckle.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/f_knuckle.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_middle.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/f_middle.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_proximal.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/f_proximal.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_0.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/forearm_0.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_1.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/forearm_1.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_collision.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/forearm_collision.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/lf_metacarpal.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/lf_metacarpal.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/mounting_plate.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/mounting_plate.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/palm.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/palm.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_distal_pst.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/th_distal_pst.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_middle.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/th_middle.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_proximal.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/th_proximal.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/wrist.obj → mjx/mujoco/mjx/test_data/shadow_hand/assets/wrist.obj
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/right_hand.xml → mjx/mujoco/mjx/test_data/shadow_hand/right_hand.xml
+1 -1
View File
@@ -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/benchmark/model/shadow_hand/scene_right.xml → mjx/mujoco/mjx/test_data/shadow_hand/scene_right.xml
View File
mjx/mujoco/mjx/benchmark/model/shadow_hand/shadow_hand.png → mjx/mujoco/mjx/test_data/shadow_hand/shadow_hand.png
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 1.3 MiB

After

Width:  |  Height:  |  Size: 1.3 MiB

mjx/mujoco/mjx/testspeed.py
+128
View File
@@ -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
+16 -16
View File
@@ -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)
dx = mjx.make_data(mx)
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)...')
print(f'Default backend: {jax.default_backend()}')
print('JIT-compiling the model physics step...')
start = time.time()
dx = step_fn(mx, dx)
mjx.device_get_into(d, dx)
step_fn = jax.jit(mjx.step).lower(mx, dx).compile()
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:
time.sleep(dt - elapsed)
if elapsed < m.opt.timestep:
time.sleep(m.opt.timestep - elapsed)
def main():
app.run(_main)
if __name__ == '__main__':
app.run(main)
main()
mjx/pyproject.toml
+8 -4
View File
@@ -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
+45 -37
View File
@@ -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_after = data.data.subtree_com[1]\n",
" com_before = data0.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",
" model = self.brax_sys.get_model()\n",
" model.opt.timestep = 0.004\n",
" sys = sys.tree_replace({'opt.timestep': 0.004, 'dt': 0.004})\n",
"\n",
" # override menagerie params for smoother policy\n",
" model.dof_damping[6:] = 0.5239\n",
" model.actuator_gainprm[:, 0] = 35.0\n",
" model.actuator_biasprm[:, 1] = -35.0\n",
" sys = sys.replace(\n",
" dof_damping=sys.dof_damping.at[6:].set(0.5239),\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",
" super().__init__(model=model, n_frames=n_frames)\n",
" n_frames = kwargs.pop('n_frames', int(self._dt / sys.opt.timestep))\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._default_pose = model.keyframe('home').qpos[7:]\n",
" self._init_q = jp.array(sys.mj_model.keyframe('home').qpos)\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",
" feet_offset = pos - pipeline_state.data.xpos[self._lower_leg_body_id]\n",
" pos = pipeline_state.site_xpos[self._feet_site_id] # feet position\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",
"timestamp": 1704355889284
"file_id": "1A58SK07tnOzix53E68D0TQ2ePCTZA61f",
"timestamp": 1707342610876
}
],
"toc_visible": true
plugin/elasticity/elasticity.h
+27 -12
View File
@@ -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;
int body_dofadr = 3*v[i];
if (vertbodyid) {
body_dofnum = m->body_dofnum[vertbodyid[v[i]]];
body_dofadr = m->body_dofadr[vertbodyid[v[i]]];
if (body_dofnum && m->body_simple[vertbodyid[v[i]]] != 2) {
mju_error("Non-simple or non-static bodies are not yet supported");
}
for (int x = 0; x < 3; x++) {
qfrc_passive[3*v[i]+x] += force[3*i+x];
}
}
}
}
// 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
+10 -2
View File
@@ -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
+1
View File
@@ -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
+10 -2
View File
@@ -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
+1
View File
@@ -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
+4 -4
View File
@@ -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
+4 -4
View File
@@ -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
+121 -134
View File
@@ -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
data an associated instance of MjData
nstate an integer, number of initial states from which to roll out trajectories
nstep an integer, number of steps to be taken for each trajectory
model instance of MjModel
data associated instance of MjData
nroll integer, number of initial states from which to roll out trajectories
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
initial_time (nstate x 1) nstate initial times
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
warmstart0 (nroll x nv) nroll qacc_warmstart vectors
control (nroll x nstep x ncontrol) nroll trajectories of nstep controls
output arguments (optional):
state (nstate x nstep x nqva) nstate length-nstep states
sensordata (nstate x nstep x nsendordata) nstate length-nstep sensordatas
state (nroll x nstep x nstate) nroll nstep states
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,
const mjtNum* state0, const mjtNum* ctrl,
const mjtNum* qfrc, const mjtNum* xfrc,
const mjtNum* mocap, const mjtNum* time0,
const mjtNum* warmstart0,
void _unsafe_rollout(const mjModel* m, mjData* d, int nroll, int nstep, unsigned int control_spec,
const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* control,
mjtNum* state, mjtNum* sensordata) {
// model sizes
int nq = m->nq;
int nv = m->nv;
int na = m->na;
int nqva = nq + nv + na;
int nu = m->nu;
int nbody = m->nbody;
int nmocap = m->nmocap;
// sizes
int nstate = mj_stateSize(m, mjSTATE_FULLPHYSICS);
int ncontrol = mj_stateSize(m, control_spec);
int nv = m->nv, nbody = m->nbody, neq = m->neq;
int nsensordata = m->nsensordata;
// loop over initial states
for (int s=0; s < nstate; s++) {
// set initial state
if (state0) {
mju_copy(d->qpos, state0 + s*nqva, nq);
mju_copy(d->qvel, state0 + s*nqva + nq, nv);
mju_copy(d->act, state0 + s*nqva + nq + nv, na);
} else {
mju_copy(d->qpos, m->qpos0, nq);
mju_zero(d->qvel, nv);
mju_zero(d->act, na);
// clear user inputs if unspecified
if (!(control_spec & mjSTATE_CTRL)) {
mju_zero(d->ctrl, m->nu);
}
if (!(control_spec & mjSTATE_QFRC_APPLIED)) {
mju_zero(d->qfrc_applied, nv);
}
if (!(control_spec & mjSTATE_XFRC_APPLIED)) {
mju_zero(d->xfrc_applied, 6*nbody);
}
if (!(control_spec & mjSTATE_MOCAP_POS)) {
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
d->time = time0 ? time0[s] : 0;
// loop over rollouts
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
if (s == 0) {
if (!ctrl) {
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);
}
}
}
// clear warning counters
for (int i = 0; i < mjNWARNING; i++) {
d->warning[i].number = 0;
}
// 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) {
mju_copy(d->ctrl, ctrl + s*nstep*nu + t*nu, nu);
}
// 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);
if (control) {
mj_setState(m, d, control + step*ncontrol, control_spec);
}
// 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);
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);
mj_getState(m, d, state + step*nstate, mjSTATE_FULLPHYSICS);
}
// copy out sensor values
if (sensordata) {
mju_copy(sensordata + s*nstep*nsensordata + t*nsensordata,
d->sensordata, nsensordata);
mju_copy(sensordata + step*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,
std::optional<const PyCArray> init_state,
std::optional<const PyCArray> init_time,
std::optional<const PyCArray> init_warmstart,
std::optional<const PyCArray> ctrl,
std::optional<const PyCArray> qfrc,
std::optional<const PyCArray> xfrc,
std::optional<const PyCArray> mocap,
[](const MjModelWrapper& m, MjDataWrapper& d,
int nroll, int nstep, unsigned int control_spec,
const PyCArray state0,
std::optional<const PyCArray> warmstart0,
std::optional<const PyCArray> control,
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* init_state_ptr =
get_array_ptr(init_state, "initial_state", nstate, 1, nqva);
mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu);
mjtNum* qfrc_ptr =
get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv);
mjtNum* xfrc_ptr =
get_array_ptr(xfrc, "xfrc_applied", nstate, nstep, 6*model->nbody);
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);
mjtNum* state0_ptr = get_array_ptr(state0, "state0", nroll, 1, nstate);
mjtNum* warmstart0_ptr = get_array_ptr(warmstart0, "warmstart0", nroll,
1, model->nv);
mjtNum* control_ptr = get_array_ptr(control, "control", nroll,
nstep, ncontrol);
mjtNum* state_ptr = get_array_ptr(state, "state", nroll, nstep, nstate);
mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nroll,
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,
xfrc_ptr, mocap_ptr, init_time_ptr, init_warmstart_ptr, state_ptr,
sensordata_ptr);
model, data, nroll, nstep, control_spec, state0_ptr,
warmstart0_ptr, control_ptr, state_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("initial_time") = py::none(),
py::arg("initial_warmstart") = py::none(),
py::arg("ctrl") = py::none(),
py::arg("qfrc_applied") = py::none(),
py::arg("xfrc_applied") = py::none(),
py::arg("mocap") = py::none(),
py::arg("state") = py::none(),
py::arg("sensordata") = py::none(),
py::arg("control_spec"),
py::arg("state0"),
py::arg("warmstart0") = py::none(),
py::arg("control") = py::none(),
py::arg("state") = py::none(),
py::arg("sensordata") = py::none(),
py::doc(rollout_doc)
);
}
} // namespace
}
} // namespace mujoco::python
}
python/mujoco/rollout.py
+112 -77
View File
@@ -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,
*, # require following arguments to be named
skip_checks=False,
nstate=None,
nstep=None,
initial_time=None,
initial_warmstart=None,
qfrc_applied=None,
xfrc_applied=None,
mocap=None,
state=None,
sensordata=None):
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
def rollout(model: mujoco.MjModel,
data: mujoco.MjData,
initial_state: npt.ArrayLike,
control: Optional[npt.ArrayLike] = None,
*, # require subsequent arguments to be named
control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
skip_checks: bool = False,
nroll: Optional[int] = None,
nstep: Optional[int] = None,
initial_warmstart: Optional[npt.ArrayLike] = None,
state: Optional[npt.ArrayLike] = None,
sensordata: Optional[npt.ArrayLike] = None):
"""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
`rollout.cc`, please see documentation therein. This python funtion will
infer `nstate` and `nstep`, tile input arguments with singleton dimensions,
and allocate output arguments if none are given.
Python wrapper for rollout.cc, see documentation therein.
Infers nroll and nstep.
Tiles inputs with singleton dimensions.
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
# output arrays, just call rollout
# skip_checks shortcut:
# 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,
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
state, sensordata)
_rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
initial_warmstart, control, 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):
raise ValueError('nstate must be an integer')
if nroll and not isinstance(nroll, int):
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,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
control=control,
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,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
control=control,
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)
qfrc_applied = _ensure_3d(qfrc_applied)
xfrc_applied = _ensure_3d(xfrc_applied)
mocap = _ensure_3d(mocap)
control = _ensure_3d(control)
state = _ensure_3d(state)
sensordata = _ensure_3d(sensordata)
# check trailing dimensions
_check_trailing_dimension(model.nq + model.nv + model.na,
initial_state=initial_state, state=state)
_check_trailing_dimension(1, initial_time=initial_time)
_check_trailing_dimension(model.nu, ctrl=ctrl)
_check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied)
_check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied)
_check_trailing_dimension(model.nmocap*7, mocap=mocap)
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
_check_trailing_dimension(nstate, initial_state=initial_state, state=state)
ncontrol = mujoco.mj_stateSize(model, control_spec)
_check_trailing_dimension(ncontrol, control=control)
_check_trailing_dimension(model.nv, initial_warmstart=initial_warmstart)
_check_trailing_dimension(model.nsensordata, sensordata=sensordata)
# infer nstate, check for incompatibilities
nstate = _infer_dimension(0, nstate or 1,
initial_state=initial_state,
initial_time=initial_time,
initial_warmstart=initial_warmstart,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
state=state,
sensordata=sensordata)
# infer nroll, check for incompatibilities
nroll = _infer_dimension(0, nroll or 1,
initial_state=initial_state,
initial_warmstart=initial_warmstart,
control=control,
state=state,
sensordata=sensordata)
# infer nstep, check for incompatibilities
nstep = _infer_dimension(1, nstep or 1,
ctrl=ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied,
mocap=mocap,
control=control,
state=state,
sensordata=sensordata)
# tile input arrays if required (singleton expansion)
initial_state = _tile_if_required(initial_state, nstate)
initial_time = _tile_if_required(initial_time, nstate)
initial_warmstart = _tile_if_required(initial_warmstart, nstate)
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)
initial_state = _tile_if_required(initial_state, nroll)
initial_warmstart = _tile_if_required(initial_warmstart, nroll)
control = _tile_if_required(control, nroll, 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,
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
state, sensordata)
_rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
initial_warmstart, control, 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
+301 -199
View File
@@ -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)
ctrl = np.random.randn(model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
initial_state = np.random.randn(nstate)
control = np.random.randn(model.nu)
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)
ctrl = np.random.randn(nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
initial_state = np.random.randn(nstate)
control = np.random.randn(nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
py_state, py_sensordata = single_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(state, np.asarray(py_state))
np.testing.assert_array_equal(sensordata, np.asarray(py_sensordata))
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_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)
ctrl = np.random.randn(nstate, 1, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nroll, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
mujoco.mj_resetData(model, data)
py_state, py_sensordata = multi_rollout(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_fixed_ctrl(self, model_name):
nstep = 3
def test_one_rollout_fixed_ctrl(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)
ctrl = np.random.randn(model.nu)
state = np.empty((nstep, model.nq + model.nv + model.na))
sensordata = np.empty((nstep, model.nsensordata))
rollout.rollout(model, data, initial_state, ctrl,
nroll = 1 # number of rollouts
nstep = 3 # number of steps
initial_state = np.random.randn(nstate)
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??
py_state, py_sensordata = single_rollout(model, data, initial_state,
ctrl=ctrl)
control = np.tile(control, (nstep, 1))
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_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)
ctrl = np.random.randn(nstate, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nroll, nstep, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
np.testing.assert_array_equal(py_state, py_state)
np.testing.assert_array_equal(py_sensordata, py_sensordata)
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_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)
ctrl = np.random.randn(nstate, 1, model.nu) # 1 control in the time dimension
state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(nroll, 1, model.nu)
state = np.empty((nroll, nstep, nstate))
state, sensordata = rollout.rollout(model, data, initial_state, control,
state=state)
ctrl = np.repeat(ctrl, nstep, axis=1)
py_state, py_sensordata = multi_rollout(model, data, initial_state,
ctrl=ctrl)
control = np.repeat(control, nstep, axis=1)
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.product(arg_nstep=[[3, 1, 1], [3, 3, 1], [3, 1, 3]],
model_name=list(ALL_MODELS.keys()))
def test_multi_rollout_multiple_inputs(self, arg_nstep, model_name):
@parameterized.parameters(ALL_MODELS.keys())
def test_py_rollout_generalized_control(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 = 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
ctrl = np.random.randn(nstate, arg_nstep[0], model.nu)
qfrc_applied = np.random.randn(nstate, arg_nstep[1], model.nv)
xfrc_applied = np.random.randn(nstate, arg_nstep[2], model.nbody*6)
control_spec = (mujoco.mjtState.mjSTATE_CTRL |
mujoco.mjtState.mjSTATE_QFRC_APPLIED |
mujoco.mjtState.mjSTATE_XFRC_APPLIED)
ncontrol = mujoco.mj_stateSize(model, control_spec)
control = np.random.randn(nroll, nstep, ncontrol)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
qfrc_applied=qfrc_applied,
xfrc_applied=xfrc_applied)
state, sensordata = rollout.rollout(model, data, initial_state, control,
control_spec=control_spec)
# tile singleton arguments
nstep = max(arg_nstep)
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)
py_state, py_sensordata = py_rollout(model, data, initial_state, control,
control_spec=control_spec)
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)
state = np.zeros((nstate, nstep, model.nq+model.nv+model.na))
sensordata = np.zeros((nstate, nstep, model.nsensordata))
ctrl = np.random.randn(nstate, nstep, model.nu)
initial_state = np.random.randn(nroll, nstate)
state = np.empty((nroll, nstep, nstate))
sensordata = np.empty((nroll, nstep, model.nsensordata))
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):
rollout.rollout(model, thread_local.data, skip_checks=True,
nstate=initial_state.shape[0], nstep=nstep,
initial_state=initial_state, ctrl=ctrl, state=state)
def call_rollout(initial_state, control, state, sensordata):
rollout.rollout(model, thread_local.data, initial_state, control,
skip_checks=True, nroll=initial_state.shape[0],
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(
(initial_state[i*n:(i+1)*n], ctrl[i*n:(i+1)*n], state[i*n:(i+1)*n]))
chunks.append((initial_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(
(initial_state[(num_workers-1)*n:], ctrl[(num_workers-1)*n:],
state[(num_workers-1)*n:]))
chunks.append((initial_state[(num_workers-1)*n:],
control[(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,
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)
#----------------------------- 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)
# ---------------------------- test advanced operation
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)
ctrl = np.zeros(model.nu)
state1, _ = step(model, data, state0, ctrl=ctrl)
# take one step, save the state
state0 = np.zeros(nstate)
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)
assert np.linalg.norm(state-state2) > 0
# take step using rollout, don't take warmstart into account
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[3:], quat2)
np.testing.assert_array_almost_equal(sensordata[0][0][:3], pos1)
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)
ctrl = np.zeros((3, model.nu))
nroll = 1
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'
with self.assertRaisesWithLiteralMatch(
ValueError, 'ctrl must be a numpy array or float'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
initial_state = np.zeros((nroll, nstate))
qfrc_applied = np.zeros((2, 3, 4, 5))
control = 'string'
with self.assertRaisesWithLiteralMatch(
ValueError, 'qfrc_applied can have at most 3 dimensions'):
state, sensordata = rollout.rollout(model, data, initial_state,
qfrc_applied=qfrc_applied)
ValueError, 'control must be a numpy array or float'):
rollout.rollout(model, data, initial_state, control)
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)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of initial_state must be 5, got 6'):
state, sensordata = rollout.rollout(model, data, initial_state)
nroll = 1
nstep = 3
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu+1)
initial_state = np.random.randn(nroll, nstate + 1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'trailing dimension of ctrl must be 2, got 3'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
ValueError, 'trailing dimension of initial_state must be 6, got 7'):
rollout.rollout(model, data, initial_state)
ctrl = np.random.randn(2, model.nu)
qfrc_applied = np.random.randn(3, model.nv) # incompatible horizon
initial_state = np.random.randn(nroll, nstate)
control = np.random.randn(1, nstep, model.nu + 1)
with self.assertRaisesWithLiteralMatch(
ValueError, 'dimension 1 inferred as 2 but qfrc_applied has 3'):
state, sensordata = rollout.rollout(model, data, initial_state, ctrl,
qfrc_applied=qfrc_applied)
ValueError, 'trailing dimension of control must be 2, got 3'):
rollout.rollout(model, data, initial_state, control)
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
initial_state = np.random.randn(model.nq + model.nv + model.na)
ctrl = np.random.randn(model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
# step with a clean mjData
initial_state = np.random.randn(nstate)
control = np.random.randn(3, 3, model.nu)
state, sensordata = rollout.rollout(model, data, initial_state, control)
# fill mjData with some debug value, see that we still get the same outputs
mujoco.mj_resetDataDebug(model, data, 255)
debug_state, debug_sensordata = rollout.rollout(model, data, initial_state,
ctrl)
# fill user fields with random values
for attr in [
'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)
np.testing.assert_array_equal(sensordata, debug_sensordata)
# roll out again
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):
data.qpos = state[:model.nq]
data.qvel = state[model.nq:model.nq+model.nv]
data.act = state[model.nq+model.nv:model.nq+model.nv+model.na]
def get_state(model, data):
nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS)
state = np.empty(nstate)
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)
for key, value in kwargs.items():
if value is not None:
setattr(data, key, np.reshape(value, getattr(data, key).shape))
mujoco.mj_setState(model, data, state, mujoco.mjtState.mjSTATE_FULLPHYSICS)
mujoco.mj_setState(model, data, control, control_spec)
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,
**kwargs_t)
initial_state if t == 0 else None,
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))
sensordata = np.empty((nstate, nstep, model.nsensordata))
for s in range(nstate):
kwargs_s = {key : value[s] for key, value in kwargs.items()}
state_s, sensordata_s = single_rollout(model, data, initial_state[s],
**kwargs_s)
state[s] = state_s
sensordata[s] = sensordata_s
return state.squeeze(), sensordata.squeeze()
def ensure_2d(arg):
if arg is None:
return None
else:
return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
def ensure_3d(arg):
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
+3
View File
@@ -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) \

Some files were not shown because too many files have changed in this diff Show More