diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst
index 04d8802a..097e675e 100644
--- a/doc/APIreference/functions.rst
+++ b/doc/APIreference/functions.rst
@@ -1109,14 +1109,14 @@ mj_addFileVFS
Add file to VFS. The directory argument is optional and can be NULL or empty. Returns 0 on success, 1 when VFS is full,
2 on name collision, or -1 when an internal error occurs.
-.. _mj_makeEmptyFileVFS:
+.. _mj_addBufferVFS:
-mj_makeEmptyFileVFS
-~~~~~~~~~~~~~~~~~~~
+mj_addBufferVFS
+~~~~~~~~~~~~~~~
-.. mujoco-include:: mj_makeEmptyFileVFS
+.. mujoco-include:: mj_addBufferVFS
-Make empty file in VFS, return 0: success, 1: full, 2: repeated name.
+Add file to VFS from buffer, return 0: success, 1: full, 2: repeated name, -1: failed to load.
.. _mj_findFileVFS:
@@ -1145,6 +1145,15 @@ mj_deleteVFS
Delete all files from VFS.
+.. _mj_makeEmptyFileVFS:
+
+mj_makeEmptyFileVFS
+~~~~~~~~~~~~~~~~~~~
+
+.. mujoco-include:: mj_makeEmptyFileVFS
+
+deprecated: use mj_copyBufferVFS.
+
.. _Initialization:
Initialization
@@ -1280,7 +1289,7 @@ mj_resetDataKeyframe
.. mujoco-include:: mj_resetDataKeyframe
-Reset data, set fields from specified keyframe.
+Reset data. If 0 <= key < nkey, set fields from specified keyframe.
.. _mj_markStack:
diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst
index eeb855d2..e463ffe6 100644
--- a/doc/XMLreference.rst
+++ b/doc/XMLreference.rst
@@ -1218,40 +1218,41 @@ construct simple shapes directly in the XML. For example, a pyramid can be creat
.. code-block:: xml
-
+
-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` and
-:ref:`mjModel.mesh_quat`. These offsets are then applied to the referencing geom's position and orientation; see
-also :at:`mesh` attribute of :ref:`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
+Positioning and orienting is complicated by the fact that vertex data in the source asset are often 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 save the translation
+and rotation offsets applied to the source asset in :ref:`mjModel.mesh_pos` and
+:ref:`mjModel.mesh_quat`; these are required if one reads vertex data from the source and needs to re-apply the
+transform. These offsets are then composed with the referencing geom's position and orientation; see also the :at:`mesh`
+attribute of :ref:`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 :at:`refpos`, :at:`refquat`` mechanism) to compensate, but in practice this is
rarely needed.
.. code-block:: xml
-
+
-
-
+
+
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
+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 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:
@@ -5429,6 +5430,20 @@ This element has one custom attribute in addition to the common attributes:
Damping applied by the actuator.
When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators`.
+.. _actuator-position-inheritrange:
+
+:at:`inheritrange`: :at-val:`real, "0"`
+ Automatically set the actuator's :at:`ctrlrange` to match the transmission target's :at:`range`. The default value
+ means "disabled". A positive value :at-val:`X` sets the :at:`ctrlrange` around the midpoint of the target range,
+ scaled by :at-val:`X`. For example if the target joint has :at:`range` of :at-val:`[0, 1]`, then a value of
+ :at-val:`1.0` will set :at:`ctrlrange` to :at-val:`[0, 1]`; values of :at-val:`0.8` and :at-val:`1.2` will set the
+ :at:`ctrlrange` to :at-val:`[0.1, 0.9]` and :at-val:`[-0.1, 1.1]`, respectively. Values smaller than 1 are useful for
+ not hitting the limits; values larger than 1 are useful for maintaining control authority at the limits (being able
+ to push on them). This attribute is exclusive with :at:`ctrlrange` and available only for joint and tendon
+ transmissions which have :at:`range` defined. Note that while :at:`inheritrange` is available both as a
+ :ref:`position` attribute and in the :ref:`default class`,
+ saved XMLs always convert it to explicit :at:`ctrlrange` at the actuator.
+
.. _actuator-velocity:
:el-prefix:`actuator/` |-| **velocity** (*)
@@ -5577,10 +5592,16 @@ This element has one custom attribute in addition to the common attributes:
Damping applied by the actuator.
When using this attribute, it is recommended to use the implicitfast or implicit :ref:`integrators`.
+.. _actuator-intvelocity-inheritrange:
+
+:at:`inheritrange`: :at-val:`real, "0"`
+ Identical to :ref:`position/inheritrange`, but sets :at:`actrange` (which has the same
+ length semantics as the transmission target) rather than :at:`ctrlrange` (which has velocity semantics).
+
.. _actuator-damper:
:el-prefix:`actuator/` |-| **damper** (*)
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This element is an active damper which produces a force proportional to both velocity and control: ``F = - kv * velocity
* control``, where ``kv`` must be nonnegative. :at:`ctrlrange` is required and must also be nonnegative.
@@ -7727,6 +7748,8 @@ tendon, slidersite, cranksite.
.. _default-position-ctrlrange:
+.. _default-position-inheritrange:
+
.. _default-position-forcerange:
.. _default-position-gear:
@@ -7787,6 +7810,8 @@ refsite, tendon, slidersite, cranksite.
.. _default-intvelocity-actrange:
+.. _default-intvelocity-inheritrange:
+
.. _default-intvelocity-gear:
.. _default-intvelocity-cranklength:
diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst
index 143fc733..03a651af 100644
--- a/doc/XMLschema.rst
+++ b/doc/XMLschema.rst
@@ -135,32 +135,6 @@
| :ref:`asset` | | |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| asset |br| |_| |L| | | .. table:: |
-| :ref:`texture | \* | :class: mjcf-attributes |
-| ` | | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`name` | :ref:`type` | :ref:`content_type` | :ref:`file` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`gridsize` | :ref:`gridlayout` | :ref:`fileright` | :ref:`fileleft` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`fileup` | :ref:`filedown` | :ref:`filefront` | :ref:`fileback` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`builtin` | :ref:`rgb1` | :ref:`rgb2` | :ref:`mark` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`markrgb` | :ref:`random` | :ref:`width` | :ref:`height` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`hflip` | :ref:`vflip` | | | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
-| |_| asset |br| |_| |L| | | .. table:: |
-| :ref:`hfield | \* | :class: mjcf-attributes |
-| ` | | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`name` | :ref:`content_type` | :ref:`file` | :ref:`nrow` | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`ncol` | :ref:`size` | :ref:`elevation` | | |
-| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
-| |_| asset |br| |_| |L| | | .. table:: |
| :ref:`mesh | \* | :class: mjcf-attributes |
| ` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
@@ -186,6 +160,15 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| asset |br| |_| |L| | | .. table:: |
+| :ref:`hfield | \* | :class: mjcf-attributes |
+| ` | | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`name` | :ref:`content_type` | :ref:`file` | :ref:`nrow` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`ncol` | :ref:`size` | :ref:`elevation` | | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| asset |br| |_| |L| | | .. table:: |
| :ref:`skin | \* | :class: mjcf-attributes |
| ` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
@@ -206,6 +189,23 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| asset |br| |_| |L| | | .. table:: |
+| :ref:`texture | \* | :class: mjcf-attributes |
+| ` | | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`name` | :ref:`type` | :ref:`content_type` | :ref:`file` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`gridsize` | :ref:`gridlayout` | :ref:`fileright` | :ref:`fileleft` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`fileup` | :ref:`filedown` | :ref:`filefront` | :ref:`fileback` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`builtin` | :ref:`rgb1` | :ref:`rgb2` | :ref:`mark` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`markrgb` | :ref:`random` | :ref:`width` | :ref:`height` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`hflip` | :ref:`vflip` | | | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| asset |br| |_| |L| | | .. table:: |
| :ref:`material | \* | :class: mjcf-attributes |
| ` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
@@ -774,13 +774,15 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name` | :ref:`class` | :ref:`group` | :ref:`ctrllimited` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`forcerange` | :ref:`lengthrange` | |
+| | | | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`inheritrange` | :ref:`forcerange` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`gear` | :ref:`cranklength` | :ref:`user` | :ref:`joint` | |
+| | | | :ref:`lengthrange` | :ref:`gear` | :ref:`cranklength` | :ref:`user` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`jointinparent` | :ref:`tendon` | :ref:`slidersite` | :ref:`cranksite` | |
+| | | | :ref:`joint` | :ref:`jointinparent` | :ref:`tendon` | :ref:`slidersite` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`site` | :ref:`refsite` | :ref:`kp` | :ref:`kv` | |
+| | | | :ref:`cranksite` | :ref:`site` | :ref:`refsite` | :ref:`kp` | |
+| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+| | | | :ref:`kv` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| actuator |br| |_| |L| | | .. table:: |
@@ -806,13 +808,13 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`forcerange` | :ref:`actrange` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`lengthrange` | :ref:`gear` | :ref:`cranklength` | :ref:`user` | |
+| | | | :ref:`inheritrange` | :ref:`lengthrange` | :ref:`gear` | :ref:`cranklength` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`joint` | :ref:`jointinparent` | :ref:`tendon` | :ref:`slidersite` | |
+| | | | :ref:`user` | :ref:`joint` | :ref:`jointinparent` | :ref:`tendon` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`cranksite` | :ref:`site` | :ref:`refsite` | :ref:`kp` | |
+| | | | :ref:`slidersite` | :ref:`cranksite` | :ref:`site` | :ref:`refsite` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`kv` | | | | |
+| | | | :ref:`kp` | :ref:`kv` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| actuator |br| |_| |L| | | .. table:: |
@@ -1440,11 +1442,11 @@
| :ref:`position | ? | :class: mjcf-attributes |
| ` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`ctrllimited` | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`forcerange` | |
+| | | | :ref:`ctrllimited` | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`inheritrange` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`gear` | :ref:`cranklength` | :ref:`user` | :ref:`group` | |
+| | | | :ref:`forcerange` | :ref:`gear` | :ref:`cranklength` | :ref:`user` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`kp` | :ref:`kv` | | | |
+| | | | :ref:`group` | :ref:`kp` | :ref:`kv` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| default |br| |_| |L| | | .. table:: |
@@ -1464,9 +1466,9 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`ctrllimited` | :ref:`forcelimited` | :ref:`ctrlrange` | :ref:`forcerange` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`actrange` | :ref:`gear` | :ref:`cranklength` | :ref:`user` | |
+| | | | :ref:`actrange` | :ref:`inheritrange` | :ref:`gear` | :ref:`cranklength` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-| | | | :ref:`group` | :ref:`kp` | :ref:`kv` | | |
+| | | | :ref:`user` | :ref:`group` | :ref:`kp` | :ref:`kv` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| default |br| |_| |L| | | .. table:: |
diff --git a/doc/changelog.rst b/doc/changelog.rst
index 55cf4d2d..30a419e3 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -5,17 +5,38 @@ Changelog
Upcoming version (not yet released)
-----------------------------------
+General
+^^^^^^^
+1. Added the :at:`inheritrange` attribute to :ref:`position` and
+ :ref:`intvelocity` actuators, allowing convenient setting of the actuator's
+ :at:`ctrlrange` or :at:`actrange` (respectively), according to the range of the transmission
+ target (joint or tendon). See :ref:`position/inheritrange` for
+ details.
+2. Deprecated :ref:`mj_makeEmptyFileVFS` in favor of :ref:`mj_addBufferVFS`. :ref:`mjVFS` now computes checksums of
+ its internal file buffers. :ref:`mj_addBufferVFS` allocates an empty buffer with a given name in an mjVFS and
+ copies the data buffer into it, combining and replacing the deprecated two-step process of calling
+ :ref:`mj_makeEmptyFileVFS` followed by a direct copy into the given mjVFS internal file buffer.
+
MJX
^^^
-1. Improved performance of getting and putting device data.
+3. 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
+4. Fixed a bug in ``mjx-viewer`` that prevented it from running. Updated ``mjx-viewer`` to use newer
``mjx.get_data_into`` function call.
+5. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
+6. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
+7. Added support for :ref:`mjOption.impratio` to ``mjx.Model``.
+
+Python bindings
+^^^^^^^^^^^^^^^
+7. Fixed incorrect data types in the bindings for the ``geom``, ``vert``, ``elem``, and ``flex`` array members
+ of the ``mjContact`` struct, and all array members of the ``mjrContext`` struct.
+
Version 3.1.2 (February 05, 2024)
-----------------------------------
diff --git a/doc/conf.py b/doc/conf.py
index bc244774..d9ca5bfb 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -118,7 +118,7 @@ SHARED_CSS_VARIABLES = {
html_theme_options = {
'light_css_variables': {
'font-stack--monospace': 'Inconsolata,Consolas,ui-monospace,monospace',
- 'at-color': '#bc103e',
+ 'at-color': '#830b2b',
'at-val-color': '#bc103e',
'body-color': '#14234b',
'color-highlight-on-target': '#e5e8ed',
@@ -131,7 +131,7 @@ html_theme_options = {
'wy-nav-side-background-color': '#0053d6',
},
'dark_css_variables': {
- 'at-color': '#ff95a6',
+ 'at-color': '#ffaab7',
'at-val-color': '#ff95a6',
'body-color': '#14234b',
'color-admonition-background': '#1e1e21',
diff --git a/doc/includes/references.h b/doc/includes/references.h
index b94d5d71..a61663d1 100644
--- a/doc/includes/references.h
+++ b/doc/includes/references.h
@@ -691,11 +691,12 @@ struct mjLROpt_ { // options for mj_setLengthRange()
mjtNum tolrange; // convergence tolerance (relative to range)
};
typedef struct mjLROpt_ mjLROpt;
-struct mjVFS_ { // virtual file system for loading from memory
- int nfile; // number of files present
- char filename[mjMAXVFS][mjMAXVFSNAME]; // file name without path
- size_t filesize[mjMAXVFS]; // file size in bytes
- void* filedata[mjMAXVFS]; // buffer with file data
+struct mjVFS_ { // virtual file system for loading from memory
+ int nfile; // number of files present
+ char filename[mjMAXVFS][mjMAXVFSNAME]; // file name without path
+ size_t filesize[mjMAXVFS]; // file size in bytes
+ void* filedata[mjMAXVFS]; // buffer with file data
+ uint64_t filestamp[mjMAXVFS]; // checksum of the file data
};
typedef struct mjVFS_ mjVFS;
struct mjOption_ { // physics options
@@ -1154,8 +1155,6 @@ struct mjModel_ {
int* mesh_texcoordadr; // texcoord data address; -1: no texcoord (nmesh x 1)
int* mesh_texcoordnum; // number of texcoord (nmesh x 1)
int* mesh_graphadr; // graph data address; -1: no graph (nmesh x 1)
- mjtNum* mesh_pos; // translation applied to asset vertices (nmesh x 3)
- mjtNum* mesh_quat; // rotation applied to asset vertices (nmesh x 4)
float* mesh_vert; // vertex positions for all meshes (nmeshvert x 3)
float* mesh_normal; // normals for all meshes (nmeshnormal x 3)
float* mesh_texcoord; // vertex texcoords for all meshes (nmeshtexcoord x 2)
@@ -1163,6 +1162,8 @@ struct mjModel_ {
int* mesh_facenormal; // normal face data (nmeshface x 3)
int* mesh_facetexcoord; // texture face data (nmeshface x 3)
int* mesh_graph; // convex graph data (nmeshgraph x 1)
+ mjtNum* mesh_pos; // translation applied to asset vertices (nmesh x 3)
+ mjtNum* mesh_quat; // rotation applied to asset vertices (nmesh x 4)
int* mesh_pathadr; // address of asset path for mesh; -1: none (nmesh x 1)
// skins
@@ -1375,6 +1376,7 @@ typedef struct mjModel_ mjModel;
struct mjResource_ {
char* name; // name of resource (filename, etc)
void* data; // opaque data pointer
+ char timestamp[512]; // timestamp of the resource
const struct mjpResourceProvider* provider; // pointer to the provider
};
typedef struct mjResource_ mjResource;
@@ -2406,10 +2408,11 @@ typedef struct mjvSceneState_ mjvSceneState;
//----------------------------- MJAPI FUNCTIONS --------------------------------
void mj_defaultVFS(mjVFS* vfs);
int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
-int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
+int mj_addBufferVFS(mjVFS* vfs, const char* name, const void* buffer, int nbuffer);
int mj_findFileVFS(const mjVFS* vfs, const char* filename);
int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
void mj_deleteVFS(mjVFS* vfs);
+int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz);
void mj_freeLastXML(void);
diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h
index b91ee594..9055dca0 100644
--- a/include/mujoco/mjmodel.h
+++ b/include/mujoco/mjmodel.h
@@ -16,6 +16,7 @@
#define MUJOCO_MJMODEL_H_
#include
+#include
#include
@@ -390,11 +391,12 @@ typedef struct mjLROpt_ mjLROpt;
//---------------------------------- mjVFS ---------------------------------------------------------
-struct mjVFS_ { // virtual file system for loading from memory
- int nfile; // number of files present
- char filename[mjMAXVFS][mjMAXVFSNAME]; // file name without path
- size_t filesize[mjMAXVFS]; // file size in bytes
- void* filedata[mjMAXVFS]; // buffer with file data
+struct mjVFS_ { // virtual file system for loading from memory
+ int nfile; // number of files present
+ char filename[mjMAXVFS][mjMAXVFSNAME]; // file name without path
+ size_t filesize[mjMAXVFS]; // file size in bytes
+ void* filedata[mjMAXVFS]; // buffer with file data
+ uint64_t filestamp[mjMAXVFS]; // checksum of the file data
};
typedef struct mjVFS_ mjVFS;
@@ -868,8 +870,6 @@ struct mjModel_ {
int* mesh_texcoordadr; // texcoord data address; -1: no texcoord (nmesh x 1)
int* mesh_texcoordnum; // number of texcoord (nmesh x 1)
int* mesh_graphadr; // graph data address; -1: no graph (nmesh x 1)
- mjtNum* mesh_pos; // translation applied to asset vertices (nmesh x 3)
- mjtNum* mesh_quat; // rotation applied to asset vertices (nmesh x 4)
float* mesh_vert; // vertex positions for all meshes (nmeshvert x 3)
float* mesh_normal; // normals for all meshes (nmeshnormal x 3)
float* mesh_texcoord; // vertex texcoords for all meshes (nmeshtexcoord x 2)
@@ -877,6 +877,8 @@ struct mjModel_ {
int* mesh_facenormal; // normal face data (nmeshface x 3)
int* mesh_facetexcoord; // texture face data (nmeshface x 3)
int* mesh_graph; // convex graph data (nmeshgraph x 1)
+ mjtNum* mesh_pos; // translation applied to asset vertices (nmesh x 3)
+ mjtNum* mesh_quat; // rotation applied to asset vertices (nmesh x 4)
int* mesh_pathadr; // address of asset path for mesh; -1: none (nmesh x 1)
// skins
diff --git a/include/mujoco/mjplugin.h b/include/mujoco/mjplugin.h
index ca850d57..b6897080 100644
--- a/include/mujoco/mjplugin.h
+++ b/include/mujoco/mjplugin.h
@@ -26,11 +26,12 @@
struct mjResource_ {
char* name; // name of resource (filename, etc)
void* data; // opaque data pointer
+ char timestamp[512]; // timestamp of the resource
const struct mjpResourceProvider* provider; // pointer to the provider
};
typedef struct mjResource_ mjResource;
-// callback for opeing a resource, returns zero on failure
+// callback for opening a resource, returns zero on failure
typedef int (*mjfOpenResource)(mjResource* resource);
// callback for reading a resource
@@ -44,10 +45,12 @@ typedef void (*mjfCloseResource)(mjResource* resource);
// sets dir to directory string with ndir being size of directory string
typedef void (*mjfGetResourceDir)(mjResource* resource, const char** dir, int* ndir);
-// callback for checking if a resource was modified since last read
-// returns > 0 if resource was modified since last open, 0 if resource was not
-// modified, and < 0 if inconclusive
-typedef int (*mjfResourceModified)(const mjResource* resource);
+// callback for checking if the current resource was modified from the time
+// specified by the timestamp
+// returns 0 if the resource's timestamp matches the provided timestamp
+// returns > 0 if the the resource is younger than the given timestamp
+// returns < 0 if the resource is older than the given timestamp
+typedef int (*mjfResourceModified)(const mjResource* resource, const char* timestamp);
// struct describing a single resource provider
struct mjpResourceProvider {
diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h
index 78bc9178..f32d6b5e 100644
--- a/include/mujoco/mujoco.h
+++ b/include/mujoco/mujoco.h
@@ -82,8 +82,8 @@ MJAPI void mj_defaultVFS(mjVFS* vfs);
// Add file to VFS, return 0: success, 1: full, 2: repeated name, -1: failed to load.
MJAPI int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
-// Make empty file in VFS, return 0: success, 1: full, 2: repeated name.
-MJAPI int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
+// Add file to VFS from buffer, return 0: success, 1: full, 2: repeated name, -1: failed to load.
+MJAPI int mj_addBufferVFS(mjVFS* vfs, const char* name, const void* buffer, int nbuffer);
// Return file index in VFS, or -1 if not found in VFS.
MJAPI int mj_findFileVFS(const mjVFS* vfs, const char* filename);
@@ -94,6 +94,8 @@ MJAPI int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
// Delete all files from VFS.
MJAPI void mj_deleteVFS(mjVFS* vfs);
+// deprecated: use mj_copyBufferVFS.
+MJAPI int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
//---------------------------------- Parse and compile ---------------------------------------------
@@ -182,7 +184,7 @@ MJAPI void mj_resetData(const mjModel* m, mjData* d);
// Reset data to defaults, fill everything else with debug_value.
MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
-// Reset data, set fields from specified keyframe.
+// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
// Mark a new frame on the mjData stack.
diff --git a/introspect/functions.py b/introspect/functions.py
index 0e9d2568..b2f7655f 100644
--- a/introspect/functions.py
+++ b/introspect/functions.py
@@ -66,9 +66,9 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Add file to VFS, return 0: success, 1: full, 2: repeated name, -1: failed to load.', # pylint: disable=line-too-long
)),
- ('mj_makeEmptyFileVFS',
+ ('mj_addBufferVFS',
FunctionDecl(
- name='mj_makeEmptyFileVFS',
+ name='mj_addBufferVFS',
return_type=ValueType(name='int'),
parameters=(
FunctionParameterDecl(
@@ -78,17 +78,23 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
),
FunctionParameterDecl(
- name='filename',
+ name='name',
type=PointerType(
inner_type=ValueType(name='char', is_const=True),
),
),
FunctionParameterDecl(
- name='filesize',
+ name='buffer',
+ type=PointerType(
+ inner_type=ValueType(name='void', is_const=True),
+ ),
+ ),
+ FunctionParameterDecl(
+ name='nbuffer',
type=ValueType(name='int'),
),
),
- doc='Make empty file in VFS, return 0: success, 1: full, 2: repeated name.', # pylint: disable=line-too-long
+ doc='Add file to VFS from buffer, return 0: success, 1: full, 2: repeated name, -1: failed to load.', # pylint: disable=line-too-long
)),
('mj_findFileVFS',
FunctionDecl(
@@ -144,6 +150,30 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
doc='Delete all files from VFS.',
)),
+ ('mj_makeEmptyFileVFS',
+ FunctionDecl(
+ name='mj_makeEmptyFileVFS',
+ return_type=ValueType(name='int'),
+ parameters=(
+ FunctionParameterDecl(
+ name='vfs',
+ type=PointerType(
+ inner_type=ValueType(name='mjVFS'),
+ ),
+ ),
+ FunctionParameterDecl(
+ name='filename',
+ type=PointerType(
+ inner_type=ValueType(name='char', is_const=True),
+ ),
+ ),
+ FunctionParameterDecl(
+ name='filesize',
+ type=ValueType(name='int'),
+ ),
+ ),
+ doc='deprecated: use mj_copyBufferVFS.',
+ )),
('mj_loadXML',
FunctionDecl(
name='mj_loadXML',
@@ -675,7 +705,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
type=ValueType(name='int'),
),
),
- doc='Reset data, set fields from specified keyframe.',
+ doc='Reset data. If 0 <= key < nkey, set fields from specified keyframe.', # pylint: disable=line-too-long
)),
('mj_markStack',
FunctionDecl(
diff --git a/introspect/structs.py b/introspect/structs.py
index 8f44927d..67c44261 100644
--- a/introspect/structs.py
+++ b/introspect/structs.py
@@ -121,6 +121,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='buffer with file data',
),
+ StructFieldDecl(
+ name='filestamp',
+ type=ArrayType(
+ inner_type=ValueType(name='uint64_t'),
+ extents=(2000,),
+ ),
+ doc='checksum of the file data',
+ ),
),
)),
('mjOption',
@@ -2591,20 +2599,6 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='graph data address; -1: no graph (nmesh x 1)',
),
- StructFieldDecl(
- name='mesh_pos',
- type=PointerType(
- inner_type=ValueType(name='mjtNum'),
- ),
- doc='translation applied to asset vertices (nmesh x 3)',
- ),
- StructFieldDecl(
- name='mesh_quat',
- type=PointerType(
- inner_type=ValueType(name='mjtNum'),
- ),
- doc='rotation applied to asset vertices (nmesh x 4)',
- ),
StructFieldDecl(
name='mesh_vert',
type=PointerType(
@@ -2654,6 +2648,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='convex graph data (nmeshgraph x 1)', # pylint: disable=line-too-long
),
+ StructFieldDecl(
+ name='mesh_pos',
+ type=PointerType(
+ inner_type=ValueType(name='mjtNum'),
+ ),
+ doc='translation applied to asset vertices (nmesh x 3)',
+ ),
+ StructFieldDecl(
+ name='mesh_quat',
+ type=PointerType(
+ inner_type=ValueType(name='mjtNum'),
+ ),
+ doc='rotation applied to asset vertices (nmesh x 4)',
+ ),
StructFieldDecl(
name='mesh_pathadr',
type=PointerType(
diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py
index 62d357f9..21197c15 100644
--- a/mjx/mujoco/mjx/_src/collision_convex.py
+++ b/mjx/mujoco/mjx/_src/collision_convex.py
@@ -173,7 +173,7 @@ def _manifold_points(
bp = b - poly
dist_bp = jp.abs(bp.dot(bc)) + dist_mask
dist_ap = jp.abs(ap.dot(ac)) + dist_mask
- d_idx = jp.concatenate([dist_bp, dist_ap]).argmax() % poly.shape[0]
+ d_idx = (dist_bp + dist_ap).argmax() % poly.shape[0]
return jp.array([a_idx, b_idx, c_idx, d_idx])
@@ -514,6 +514,7 @@ def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
frame = jp.stack([math.make_frame(n)] * 4, axis=0)
unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1
dist = jp.where(unique, -support[idx], 1)
+ pos = pos - 0.5 * dist[:, None] * n
return dist, pos, frame
diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py
index 5b58f7f8..f0bca5e9 100644
--- a/mjx/mujoco/mjx/_src/collision_driver_test.py
+++ b/mjx/mujoco/mjx/_src/collision_driver_test.py
@@ -265,7 +265,7 @@ class ConvexTest(absltest.TestCase):
-
+
@@ -281,7 +281,35 @@ class ConvexTest(absltest.TestCase):
# extract the contact points with penetration
c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
for field in dataclasses.fields(Contact):
- _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
+ _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-5)
+
+ _FLAT_BOX_PLANE = """
+
+
+
+
+
+
+
+
+
+ """
+
+ def test_flat_box_plane(self):
+ """Tests box collision with a plane."""
+ with jax.disable_jit():
+ d, dx = _collide(self._FLAT_BOX_PLANE)
+
+ np.testing.assert_array_less(dx.contact.dist, 0)
+
+ # sort positions for comparison
+ idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
+ dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
+ idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
+ d.contact.pos[:] = d.contact.pos[idx]
+
+ for field in dataclasses.fields(Contact):
+ _assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5)
_BOX_BOX = """
diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py
index 85df7e18..50aee9b3 100644
--- a/mjx/mujoco/mjx/_src/constraint.py
+++ b/mjx/mujoco/mjx/_src/constraint.py
@@ -297,7 +297,7 @@ def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
for diff_tan, friction in zip(diff_con[1:], c.friction[:2]):
for f in (friction, -friction):
js.append(diff_con[0] + diff_tan * f)
- invweights.append((t + f * f * t) * 2 * f * f)
+ invweights.append((t + f * f * t) * 2 * f * f / m.opt.impratio)
active = dist < 0
j, invweight = jp.stack(js) * active, jp.stack(invweights)
diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py
index 1c71a9cd..5a80ab4c 100644
--- a/mjx/mujoco/mjx/_src/dataclasses.py
+++ b/mjx/mujoco/mjx/_src/dataclasses.py
@@ -18,13 +18,23 @@ import copy
import dataclasses
import typing
-from typing import Any, Dict, Optional, Sequence, TypeVar
+from typing import Any, Dict, Optional, Sequence, TypeVar, Union
import jax
import numpy as np
_T = TypeVar('_T')
+def _jax_in_args(typ) -> bool:
+ if typ is jax.Array:
+ return True
+ if dataclasses.is_dataclass(typ):
+ return any(_jax_in_args(f.type) for f in dataclasses.fields(typ))
+ if typing.get_origin(typ) in (list, dict, Union, set):
+ return any(_jax_in_args(t) for t in typing.get_args(typ))
+ return False
+
+
def dataclass(clz: _T) -> _T:
"""Wraps a dataclass with metadata for which fields are pytrees.
@@ -41,12 +51,7 @@ def dataclass(clz: _T) -> _T:
data_clz = dataclasses.dataclass(frozen=True)(clz)
meta_fields, data_fields = [], []
for field in dataclasses.fields(data_clz):
- if any((
- 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)),
- )):
+ if _jax_in_args(field.type):
data_fields.append(field)
else:
meta_fields.append(field)
diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py
index 682cc88e..c6bbdbc1 100644
--- a/mjx/mujoco/mjx/_src/forward.py
+++ b/mjx/mujoco/mjx/_src/forward.py
@@ -287,8 +287,10 @@ def euler(m: Model, d: Data) -> Data:
# integrate damping implicitly
qacc = d.qacc
if not m.opt.disableflags & DisableBit.EULERDAMP:
- # TODO(robotics-simulation): can this be done with a smaller perf hit
- dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+ if support.is_sparse(m):
+ dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+ else:
+ dh = d.replace(qM=d.qM + jp.diag(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)
diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py
index cc767ae8..0a777797 100644
--- a/mjx/mujoco/mjx/_src/forward_test.py
+++ b/mjx/mujoco/mjx/_src/forward_test.py
@@ -27,14 +27,14 @@ import numpy as np
_TOLERANCE = 1e-5
-def _assert_eq(a, b, name):
- tol = _TOLERANCE * 10 # avoid test noise
+def _assert_eq(a, b, name, tol=_TOLERANCE):
+ tol = tol * 10 # avoid test noise
err_msg = f'mismatch: {name}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
-def _assert_attr_eq(a, b, attr):
- _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+def _assert_attr_eq(a, b, attr, tol=_TOLERANCE):
+ _assert_eq(getattr(a, attr), getattr(b, attr), attr, tol=tol)
class ForwardTest(absltest.TestCase):
@@ -46,7 +46,7 @@ class ForwardTest(absltest.TestCase):
d.ctrl = np.array([-18, 0.59, 0.47])
d.xfrc_applied[0, 2] = 0.1 # torque
d.xfrc_applied[1, 4] = 0.3 # linear force
- mujoco.mj_step(m, d, 100) # get some dynamics going
+ mujoco.mj_step(m, d, 20) # get some dynamics going
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
@@ -75,14 +75,13 @@ class ForwardTest(absltest.TestCase):
d.ctrl = np.array([-18, 0.59, 0.47])
d.xfrc_applied[0, 2] = 0.1 # torque
d.xfrc_applied[1, 4] = 0.3 # linear force
- mujoco.mj_step(m, d, 100) # get some dynamics going
+ mujoco.mj_step(m, d, 20) # get some dynamics going
- mx = mjx.put_model(m)
- dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
+ dx = jax.jit(mjx.step)(mjx.put_model(m), mjx.put_data(m, d))
mujoco.mj_step(m, d)
_assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time')
- _assert_attr_eq(d, dx, 'qvel')
+ _assert_attr_eq(d, dx, 'qvel', tol=5e-4)
_assert_attr_eq(d, dx, 'qpos')
def test_rk4(self):
@@ -111,8 +110,7 @@ class ForwardTest(absltest.TestCase):
mujoco.mj_step(m, d, 10) # let dynamics get state significantly non-zero
mujoco.mj_forward(m, d)
- mx = mjx.put_model(m)
- dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d))
+ dx = jax.jit(mjx.rungekutta4)(mjx.put_model(m), mjx.put_data(m, d))
mujoco.mj_RungeKutta(m, d, 4)
_assert_attr_eq(d, dx, 'qvel')
@@ -120,6 +118,30 @@ class ForwardTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time')
+ def test_eulerdamp(self):
+ m = test_util.load_test_file('pendula.xml')
+ self.assertTrue((m.dof_damping > 0).any())
+
+ d = mujoco.MjData(m)
+ d.qvel[:] = 1.0
+ d.qacc[:] = 1.0
+ mujoco.mj_forward(m, d)
+ dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
+ mujoco.mj_Euler(m, d)
+
+ _assert_attr_eq(d, dx, 'qpos')
+
+ # also test sparse
+ m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+ d = mujoco.MjData(m)
+ d.qvel[:] = 1.0
+ d.qacc[:] = 1.0
+ mujoco.mj_forward(m, d)
+ dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
+ mujoco.mj_Euler(m, d)
+
+ _assert_attr_eq(d, dx, 'qpos')
+
def test_disable_eulerdamp(self):
m = test_util.load_test_file('pendula.xml')
self.assertTrue((m.dof_damping > 0).any())
@@ -128,8 +150,7 @@ class ForwardTest(absltest.TestCase):
d = mujoco.MjData(m)
d.qvel[:] = 1.0
d.qacc[:] = 1.0
- mx = mjx.put_model(m)
- dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
+ dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep)
diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 2f979d71..d1fbbe55 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -146,23 +146,23 @@ 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)
- zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32)
- zero_nefc = jp.zeros(nefc, dtype=jp.float32)
- zero_na = jp.zeros(m.na, dtype=jp.float32)
- zero_nu = jp.zeros(m.nu, dtype=jp.float32)
- zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32)
- zero_nm = jp.zeros(m.nM, dtype=jp.float32)
+ zero_0 = jp.zeros(0, dtype=float)
+ zero_nv = jp.zeros(m.nv, dtype=float)
+ zero_nv_6 = jp.zeros((m.nv, 6), dtype=float)
+ zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=float)
+ zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=float)
+ zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=float)
+ zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=float)
+ zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=float)
+ zero_nefc = jp.zeros(nefc, dtype=float)
+ zero_na = jp.zeros(m.na, dtype=float)
+ zero_nu = jp.zeros(m.nu, dtype=float)
+ zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=float)
+ zero_nm = jp.zeros(m.nM, dtype=float)
# create first d to get num contacts and nc
d = types.Data(
- solver_niter=jp.array(0, dtype=jp.int32),
+ solver_niter=jp.array(0, dtype=int),
time=jp.array(0.0),
qpos=jp.array(m.qpos0),
qvel=zero_nv,
@@ -171,31 +171,31 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
ctrl=zero_nu,
qfrc_applied=zero_nv,
xfrc_applied=zero_nbody_6,
- eq_active=jp.zeros(m.neq, dtype=jp.int32),
+ eq_active=jp.zeros(m.neq, dtype=int),
qacc=zero_nv,
act_dot=zero_na,
xpos=zero_nbody_3,
- xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
+ xquat=jp.zeros((m.nbody, 4), dtype=float),
xmat=zero_nbody_3_3,
xipos=zero_nbody_3,
ximat=zero_nbody_3_3,
xanchor=zero_njnt_3,
xaxis=zero_njnt_3,
- geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
- geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
- site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
- site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
+ geom_xpos=jp.zeros((m.ngeom, 3), dtype=float),
+ geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=float),
+ site_xpos=jp.zeros((m.nsite, 3), dtype=float),
+ site_xmat=jp.zeros((m.nsite, 3, 3), dtype=float),
subtree_com=zero_nbody_3,
cdof=zero_nv_6,
cinert=zero_nbody_10,
actuator_length=zero_nu,
- actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
+ actuator_moment=jp.zeros((m.nu, m.nv), dtype=float),
crb=zero_nbody_10,
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_J=jp.zeros((nefc, m.nv), dtype=float),
efc_frictionloss=zero_nefc,
efc_D=zero_nefc,
actuator_velocity=zero_nu,
diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py
index eaa5d3bf..bebca9ac 100644
--- a/mjx/mujoco/mjx/_src/mesh.py
+++ b/mjx/mujoco/mjx/_src/mesh.py
@@ -212,10 +212,14 @@ def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
"""Derives geom mesh attributes for mjx.Model from MjModel."""
kwargs = {k: [] for k in _DERIVED_ARGS}
verts, faces = _get_faces_verts(m)
+ geom_con = m.geom_conaffinity | m.geom_contype
for geomid in range(m.ngeom):
dataid = m.geom_dataid[geomid]
- typ = m.geom_type[geomid]
- if typ == GeomType.BOX:
+ if not geom_con[geomid]:
+ # ignore visual-only meshes
+ kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
+ continue
+ elif m.geom_type[geomid] == GeomType.BOX:
vert, face = _box(m.geom_size[geomid])
elif dataid >= 0:
vert, face = verts[dataid], faces[dataid]
diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py
index cd0cc63c..dae5a82b 100644
--- a/mjx/mujoco/mjx/_src/solver.py
+++ b/mjx/mujoco/mjx/_src/solver.py
@@ -218,13 +218,13 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
mgrad = smooth.solve_m(m, d, grad)
elif m.opt.solver == SolverType.NEWTON:
ne, nf, *_ = constraint.count_constraints(m)
- active = (ctx.Jaref < 0).at[:ne + nf].set(True)
- 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)
+ active = (ctx.Jaref < 0).at[: ne + nf].set(True)
+ h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
+ h = support.full_m(m, d) + h
+ h_ = jax.scipy.linalg.cho_factor(h)
+ mgrad = jax.scipy.linalg.cho_solve(h_, grad)
else:
- raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
+ raise NotImplementedError(f'unsupported solver type: {m.opt.solver}')
ctx = ctx.replace(grad=grad, Mgrad=mgrad)
diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py
index 692f6619..f4fdf732 100644
--- a/mjx/mujoco/mjx/_src/solver_test.py
+++ b/mjx/mujoco/mjx/_src/solver_test.py
@@ -33,8 +33,8 @@ def _assert_eq(a, b, name, tol=_TOLERANCE):
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
-def _assert_attr_eq(a, b, attr):
- _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+def _assert_attr_eq(a, b, attr, tol=_TOLERANCE):
+ _assert_eq(getattr(a, attr), getattr(b, attr), attr, tol=tol)
class SolverTest(absltest.TestCase):
@@ -42,13 +42,20 @@ class SolverTest(absltest.TestCase):
def test_newton(self):
"""Test newton solver."""
m = test_util.load_test_file('constraints.xml')
+ # it's critical that mgrad is optimally calculated, so lower iterations
+ # to be sure that MJX is converging as quickly as MuJoCo
+ m.opt.iterations = 1
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))
+ mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps
+
+ # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
+ # at the beginning of the step so that MJX does not have a trivial solution
+ warmstart = d.qacc_warmstart.copy()
+ mujoco.mj_forward(m, d)
+ d.qacc_warmstart = warmstart
+
+ dx = jax.jit(mjx.solve)(mjx.put_model(m), 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)
@@ -58,23 +65,30 @@ class SolverTest(absltest.TestCase):
"""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))
+ mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps
+
+ # CG does not converge as quickly as Newton but is cheaper to calculate
+ m.opt.solver = mujoco.mjtSolver.mjSOL_CG
+ m.opt.iterations = 8
+
+ # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
+ # at the beginning of the step so that MJX does not have a trivial solution
+ warmstart = d.qacc_warmstart.copy()
+ mujoco.mj_forward(m, d)
+ d.qacc_warmstart = warmstart
+
+ dx = jax.jit(mjx.solve)(mjx.put_model(m), 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')
+ _assert_attr_eq(d, dx, 'qfrc_constraint', tol=8e-4)
nnz = dx.efc_J.any(axis=1)
- _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+ _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=5e-4)
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
+ mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps
m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
@@ -83,17 +97,21 @@ class SolverTest(absltest.TestCase):
# 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."""
+ def test_sparse(self):
+ """Test solver works with sparse mass matrices."""
m = test_util.load_test_file('constraints.xml')
+ m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
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))
+ mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps
+
+ # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
+ # at the beginning of the step so that MJX does not have a trivial solution
+ warmstart = d.qacc_warmstart.copy()
+ mujoco.mj_forward(m, d)
+ d.qacc_warmstart = warmstart
+
+ dx = jax.jit(mjx.solve)(mjx.put_model(m), 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)
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index 8b0c3cc9..578f78b6 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -15,7 +15,7 @@
"""Base types used in MJX."""
import enum
-from typing import Sequence
+from typing import List, Optional
import jax
import jax.numpy as jp
@@ -218,6 +218,7 @@ class Option(PyTreeNode):
Attributes:
timestep: timestep
+ impratio: ratio of friction-to-normal contact impedance
tolerance: main solver tolerance
ls_tolerance: CG/Newton linesearch tolerance
gravity: gravitational acceleration (3,)
@@ -238,9 +239,11 @@ class Option(PyTreeNode):
disableflags: bit flags for disabling standard features
"""
timestep: jax.Array
+ # unsupported: apirate
+ impratio: jax.Array
tolerance: jax.Array
ls_tolerance: jax.Array
- # unsupported: apirate, impratio, noslip_tolerance, mpr_tolerance
+ # unsupported: noslip_tolerance, mpr_tolerance
gravity: jax.Array
wind: jax.Array
density: jax.Array
@@ -498,10 +501,10 @@ class Model(PyTreeNode):
pair_dim: np.ndarray
pair_geom1: np.ndarray
pair_geom2: np.ndarray
- geom_convex_face: Sequence[jax.Array]
- geom_convex_vert: Sequence[jax.Array]
- geom_convex_edge: Sequence[jax.Array]
- geom_convex_facenormal: Sequence[jax.Array]
+ geom_convex_face: List[Optional[jax.Array]]
+ geom_convex_vert: List[Optional[jax.Array]]
+ geom_convex_edge: List[Optional[jax.Array]]
+ geom_convex_facenormal: List[Optional[jax.Array]]
pair_solref: jax.Array
pair_solreffriction: jax.Array
pair_solimp: jax.Array
@@ -579,8 +582,8 @@ class Contact(PyTreeNode):
solref=jp.zeros((ncon, mujoco.mjNREF)),
solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
- geom1=jp.zeros(ncon, dtype=jp.int32),
- geom2=jp.zeros(ncon, dtype=jp.int32),
+ geom1=jp.zeros(ncon, dtype=int),
+ geom2=jp.zeros(ncon, dtype=int),
)
diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml
index 52eb95d8..42294b50 100644
--- a/mjx/mujoco/mjx/test_data/constraints.xml
+++ b/mjx/mujoco/mjx/test_data/constraints.xml
@@ -5,7 +5,7 @@
* solref, solimp
-->
-
+
diff --git a/model/humanoid/README.md b/model/humanoid/README.md
index de4c7069..37a6a203 100644
--- a/model/humanoid/README.md
+++ b/model/humanoid/README.md
@@ -13,6 +13,7 @@ in the DeepMind Control Suite [2], which has fairly realistic actuator gains.
## Changelog
+* 20-02-2024: Sorted actuators in the same order as the joints.
* 02-01-2024: Add more keyframes.
* 27-11-2023: Move humanoid geoms to group 1.
* 05-04-2023: Fix typo in texture size.
diff --git a/model/humanoid/humanoid.xml b/model/humanoid/humanoid.xml
index 013ebe23..b4d8d131 100644
--- a/model/humanoid/humanoid.xml
+++ b/model/humanoid/humanoid.xml
@@ -200,21 +200,21 @@
-
+
-
+
-
+
diff --git a/python/mujoco/bindings_test.py b/python/mujoco/bindings_test.py
index 407557f0..d293448e 100644
--- a/python/mujoco/bindings_test.py
+++ b/python/mujoco/bindings_test.py
@@ -529,6 +529,7 @@ class MuJoCoBindingsTest(parameterized.TestCase):
expected_ncon = 4
mujoco.mj_forward(self.model, self.data)
self.assertLen(self.data.contact, expected_ncon)
+ np.testing.assert_array_equal(self.data.contact.geom, [[0, 1]] * 4)
expected_pos = []
for contact in self.data.contact:
@@ -551,6 +552,13 @@ class MuJoCoBindingsTest(parameterized.TestCase):
self.assertLen(expected_H, expected_ncon)
np.testing.assert_array_equal(self.data.contact.H, expected_H)
+ expected_geom = []
+ for i, contact in enumerate(self.data.contact):
+ expected_geom.append([i, i + 1])
+ contact.geom = expected_geom[-1]
+ self.assertLen(expected_geom, expected_ncon)
+ np.testing.assert_array_equal(self.data.contact.geom, expected_geom)
+
def test_realloc_con_efc(self):
self.assertEmpty(self.data.contact)
diff --git a/python/mujoco/minimize.py b/python/mujoco/minimize.py
new file mode 100644
index 00000000..7532dc21
--- /dev/null
+++ b/python/mujoco/minimize.py
@@ -0,0 +1,337 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Nonlinear Least Squares minimization with box bounds."""
+
+import dataclasses
+import enum
+import time
+from typing import Callable, List, Optional, TextIO, Tuple, Union
+
+import mujoco
+import numpy as np
+
+
+class Verbosity(enum.Enum):
+ SILENT = 0
+ FINAL = 1
+ ITER = 2
+ FULLITER = 3
+
+
+class Status(enum.Enum):
+ FACTORIZATION_FAILED = enum.auto()
+ NO_IMPORVEMENT = enum.auto()
+ MAX_ITER = enum.auto()
+ DX_TOL = enum.auto()
+
+
+_STATUS_MESSAGE = {
+ Status.FACTORIZATION_FAILED: 'factorization failed.',
+ Status.NO_IMPORVEMENT: 'no improvement found.',
+ Status.MAX_ITER: 'maximum iterations reached.',
+ Status.DX_TOL: 'norm(step) < tolerance.',
+}
+
+
+@dataclasses.dataclass(frozen=True)
+class IterLog:
+ """Log of a single iteration of the non-linear least-squares solver.
+
+ Attributes:
+ candidate: Value of the decision variable at the beginning this iteration.
+ objective: Value of the objective at the candidate.
+ reduction: Reduction of the objective during this iteration.
+ regularizer: Value of the regularizer used for this iteration.
+ residual: Optional value of the residual at the candidate.
+ jacobian: Optional value of the Jacobian at the candidate.
+ step: Optional change in decision variable during this iteration.
+ """
+
+ candidate: np.ndarray
+ objective: np.float64
+ reduction: np.float64
+ regularizer: np.float64
+ residual: Optional[np.ndarray] = None
+ jacobian: Optional[np.ndarray] = None
+ step: Optional[np.ndarray] = None
+
+
+def jacobian_fd(
+ residual: Callable[[np.ndarray], np.ndarray],
+ x: np.ndarray,
+ r: np.ndarray,
+ eps: float,
+ bounds: Optional[List[np.ndarray]] = None,
+):
+ """Finite-difference Jacobian of a residual function.
+
+ Args:
+ residual: function that returns the residual for a given point.
+ x: point at which to evaluate the Jacobian.
+ r: residual at x.
+ eps: finite-difference step size.
+ bounds: optional pair of lower and upper bounds of the solution.
+
+ Returns:
+ jac: Jacobian of the residual at x.
+ """
+ nx = x.size
+ nr = r.size
+ jac = np.zeros((nr, nx))
+ xh = x.copy()
+ for i in range(nx):
+ if bounds is not None:
+ # Have bounds: scale eps, don't cross bounds.
+ lower, upper = bounds
+ eps_i = eps * (upper[i] - lower[i])
+ if xh[i] < upper[i] - eps_i:
+ # Not near upper bound, use forward.
+ xh[i] += eps_i
+ rh = residual(xh)
+ jac[:, i] = (rh - r) / eps_i
+ else:
+ # Near upper bound, use backward.
+ xh[i] -= eps_i
+ rh = residual(xh)
+ jac[:, i] = (r - rh) / eps_i
+ else:
+ # No bounds, just use forward fin-diff.
+ xh[i] += eps
+ rh = residual(xh)
+ jac[:, i] = (rh - r) / eps
+ xh[i] = x[i]
+ return jac
+
+
+def least_squares(
+ x0: np.ndarray,
+ residual: Callable[[np.ndarray], np.ndarray],
+ bounds: Optional[List[np.ndarray]] = None,
+ jacobian: Optional[Callable[[np.ndarray, np.ndarray], np.ndarray]] = None,
+ eps: Optional[float] = -6,
+ mu_min: Optional[float] = -6,
+ mu_max: Optional[float] = 8,
+ mu_delta: Optional[float] = 0.5,
+ tol: Optional[float] = 1e-7,
+ max_iter: Optional[int] = 100,
+ verbose: Optional[Union[Verbosity, int]] = Verbosity.ITER,
+ output: Optional[TextIO] = None,
+) -> Tuple[np.ndarray, List[IterLog]]:
+ """Nonlinear Least Squares minimization with box bounds.
+
+ Args:
+ x0: initial guess
+ residual: function that returns the residual for a given point x.
+ bounds: optional pair of lower and upper bounds on the solution.
+ jacobian: optional function that returns Jacobian of the residual at a given
+ point and residual. If not given, `residual` will be finite-differenced.
+ eps: log10 of the perurbation used for automatic finite-differencing.
+ mu_min: log10 of the minimum value of the regularizer.
+ mu_max: log10 of the maximum value of the regularizer.
+ mu_delta: log10 of the factor increasing or decreasing the regularizer.
+ tol: termination tolerance on the step size.
+ max_iter: maximum number of iterations.
+ verbose: verbosity level.
+ output: optional file or StringIO to which to print messages.
+
+ Returns:
+ x: best solution found
+ trace: sequence of solution iterates.
+ """
+ t_start_total = time.time()
+
+ # Convert verbosity to int.
+ verbose = Verbosity(verbose).value
+
+ # Initialize locals.
+ x = x0.copy()
+ mu = -np.inf # Optimistically start with no regularization.
+ n = x.size
+ i = 0
+ trace = []
+ dx = np.zeros((n,))
+ scratch = np.zeros((n, n + 7))
+ dx_norm = 0.0
+ xnew = np.zeros((n,))
+ status = Status.MAX_ITER
+ n_res = 0
+ n_jac = 0
+ t_res = 0.0
+ t_jac = 0.0
+ t_qp = 0.0
+
+ # Regularization control functions.
+ def increase_mu(mu):
+ return min(mu_max, max(mu_min, mu_delta + mu))
+
+ def decrease_mu(mu):
+ return -np.inf if mu - mu_delta < mu_min else mu - mu_delta
+
+ if bounds is not None:
+ # Checks bounds.
+ if len(bounds) != 2:
+ raise ValueError('bounds must have 2 elements.')
+ if bounds[0].size != n or bounds[1].size != n:
+ raise ValueError('bounds must have the same size as x0.')
+ if not np.all(np.isfinite(bounds[0])) or not np.all(np.isfinite(bounds[1])):
+ raise ValueError('bounds must be finite.')
+ if not np.all(bounds[0] < bounds[1]):
+ raise ValueError('bounds[0] must be smaller than bounds[1].')
+ # Clip.
+ np.clip(x, bounds[0], bounds[1], out=x)
+
+ # Get initial residual.
+ t_start = time.time()
+ r = residual(x)
+ rnew = r
+ t_res += time.time() - t_start
+ n_res += 1
+
+ # Minimize.
+ for i in range(max_iter):
+ if status != Status.MAX_ITER:
+ break
+
+ # Get objective y.
+ y = 0.5 * r.dot(r)
+
+ # Get Jacobian jac.
+ t_start = time.time()
+ if jacobian is None:
+ jac = jacobian_fd(residual, x, r, 10**eps, bounds)
+ t_res += time.time() - t_start
+ n_res += n
+ else:
+ jac = jacobian(x, r)
+ t_jac += time.time() - t_start
+ n_jac += 1
+
+ # Get gradient, Gauss-Newton Hessian.
+ grad = jac.T @ r
+ hess = jac.T @ jac
+ gnorm = np.linalg.norm(grad)
+
+ # Bounds relative to x
+ dbounds = [None, None] if bounds is None else [bounds[0] - x, bounds[1] - x]
+
+ # Find some reduction.
+ reduction = -1
+ while reduction < 0:
+ # Increase mu until factorizabl.
+ factorizable = False
+ while not factorizable:
+ # Formula from https://arxiv.org/abs/2112.02089
+ reg = np.sqrt(gnorm * 10**mu) * np.eye(n)
+ t_start = time.time()
+ nfree = mujoco.mju_boxQP(
+ dx, scratch, None, hess + reg, grad, dbounds[0], dbounds[1]
+ )
+ t_qp += time.time() - t_start
+ if nfree > -1:
+ factorizable = True
+ elif mu >= mu_max:
+ status = Status.FACTORIZATION_FAILED
+ break
+ else:
+ mu += mu_delta
+
+ if status != Status.MAX_ITER:
+ break
+
+ # New candidate, residual.
+ xnew = x + dx
+ t_start = time.time()
+ rnew = residual(xnew)
+ t_res += time.time() - t_start
+ n_res += 1
+
+ # New objective, evaluate reduction.
+ ynew = 0.5 * rnew.dot(rnew)
+ reduction = y - ynew
+
+ if reduction < 0:
+ if mu >= mu_max:
+ status = Status.NO_IMPORVEMENT
+ break
+ mu = increase_mu(mu)
+
+ if status != Status.MAX_ITER:
+ break
+
+ # Compute reduction ratio.
+ expected_reduction = -(grad.dot(dx) + 0.5 * dx.T @ hess @ dx)
+ reduction_ratio = 0.0
+ if expected_reduction == 0:
+ print('Zero expected reduction: exact minimum found?', file=output)
+ elif expected_reduction < 0:
+ print('Negative expected reduction: should not occur.', file=output)
+ else:
+ reduction_ratio = reduction / expected_reduction
+
+ # Iteration message.
+ if verbose >= Verbosity.ITER.value:
+ message = (
+ f'iter: {i:<3d} y: {y:<8.3g} mu: {mu:>4.1f} '
+ f'ratio: {reduction_ratio:<5.2g} '
+ f'dx: {dx_norm:<8.3g} reduction: {reduction:<8.3g}'
+ )
+ print(message, file=output)
+
+ # Append log to trace.
+ log = IterLog(candidate=x, objective=y, reduction=reduction, regularizer=mu)
+ if verbose >= Verbosity.FULLITER.value:
+ log = dataclasses.replace(log, residual=r, jacobian=jac, step=dx)
+ trace.append(log)
+
+ # Check for success.
+ dx_norm = np.linalg.norm(dx)
+ if dx_norm < tol:
+ status = Status.DX_TOL
+ break
+
+ # Modify regularizer like in (Bazaraa, Sherali, and Shetty)
+ if reduction_ratio > 0.75:
+ mu = decrease_mu(mu)
+ elif reduction_ratio < 0.25:
+ mu = increase_mu(mu)
+
+ # Accept proposal.
+ x = xnew
+ r = rnew
+
+ # Print final diagnostics.
+ if verbose > Verbosity.SILENT.value:
+ message = f'Terminated after {i} iterations: '
+ message += _STATUS_MESSAGE[status]
+
+ message += f' Residual evals: {n_res:d}'
+ if n_jac > 0:
+ message += f', Jacobian evals: {n_jac:d}'
+ print(message, file=output)
+
+ time_total = time.time() - t_start_total
+ if time_total > 0:
+ qp_percent = 100 * t_qp / time_total
+ r_percent = 100 * t_res / time_total
+ time_scale = 1 if time_total > 1 else 1000
+ time_units = 's' if time_total > 1 else 'ms'
+ message = f'total time {time_scale * time_total:<.1f}{time_units}'
+ message += f' of which QP {qp_percent:<.1f}%, residual {r_percent:<.1f}%'
+ if t_jac > 0:
+ jac_percent = 100 * t_jac / time_total
+ message += f' Jacobian {jac_percent:<.1f}%'
+ print(message, file=output)
+
+ return x, trace
diff --git a/python/mujoco/minimize_test.py b/python/mujoco/minimize_test.py
new file mode 100644
index 00000000..117e4211
--- /dev/null
+++ b/python/mujoco/minimize_test.py
@@ -0,0 +1,145 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for minimize.py."""
+
+import io
+from typing import Tuple
+
+from absl.testing import absltest
+from mujoco import minimize
+import numpy as np
+
+
+class MinimizeTest(absltest.TestCase):
+
+ def test_basic(self) -> None:
+ def residual(x: np.ndarray) -> float:
+ return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+
+ out = io.StringIO()
+ x0 = np.array((0.0, 0.0))
+ x, _ = minimize.least_squares(x0, residual, output=out)
+ expected_x = np.array((1.0, 1.0))
+ np.testing.assert_array_almost_equal(x, expected_x)
+ self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
+
+ def test_start_at_minimum(self) -> None:
+ def residual(x: np.ndarray) -> float:
+ return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+
+ out = io.StringIO()
+ x0 = np.array((1.0, 1.0))
+ x, _ = minimize.least_squares(x0, residual, output=out)
+ expected_x = np.array((1.0, 1.0))
+ np.testing.assert_array_almost_equal(x, expected_x)
+ self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
+ self.assertContainsSubsequence(out.getvalue(), 'exact minimum found')
+
+ def test_jac_callback(self) -> None:
+ def residual(x: np.ndarray) -> float:
+ return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+
+ def jacobian(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
+ del r # Unused.
+ return np.array([[-1, 0], [-20 * x[0], 10]])
+
+ x0 = np.array((0.0, 0.0))
+ out = io.StringIO()
+ x, _ = minimize.least_squares(x0, residual, jacobian=jacobian, output=out)
+ expected_x = np.array((1.0, 1.0))
+ np.testing.assert_array_almost_equal(x, expected_x)
+ self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
+
+ # Try with bad Jacobian, expect no improvement.
+ def jac_bad1(x: np.ndarray, r: np.ndarray) -> Tuple[float, np.ndarray]:
+ return -jacobian(x, r)
+ out1 = io.StringIO()
+ minimize.least_squares(x0, residual, jacobian=jac_bad1, output=out1)
+ self.assertContainsSubsequence(out1.getvalue(), 'no improvement found.')
+
+ def test_max_iter(self) -> None:
+ dim = 20 # High-D Rosenbrock
+
+ def residual(x: np.ndarray) -> float:
+ res0 = [1 - x[i] for i in range(dim - 1)]
+ res1 = [10 * (x[i] - x[i + 1] ** 2) for i in range(dim - 1)]
+ return np.asarray(res0 + res1)
+
+ # Fail to reach minimum after 20 iterations.
+ x0 = np.zeros(dim)
+ out = io.StringIO()
+ minimize.least_squares(x0, residual, max_iter=20, output=out)
+ self.assertContainsSubsequence(out.getvalue(), 'maximum iterations')
+
+ # Succeed after 100 iterations (default).
+ x, _ = minimize.least_squares(x0, residual)
+ expected_x = np.ones(20)
+ np.testing.assert_array_almost_equal(x, expected_x)
+
+ def test_bounds(self) -> None:
+ def residual(x: np.ndarray) -> float:
+ return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+
+ out = io.StringIO()
+ x0 = np.array((0.0, 0.0))
+ expected_x = np.array((1.0, 1.0))
+
+ bounds_types = {'inbounds': [np.array((-2.0, -2.0)), np.array((2.0, 2.0))],
+ 'onlower': [np.array((-2.0, 2.0)), np.array((0.5, 3.0))],
+ 'onupper': [np.array((-2.0, -2.0)), np.array((0.5, 2.0))]}
+
+ # In bounds finds true minimum.
+ x, _ = minimize.least_squares(x0, residual, bounds=bounds_types['inbounds'],
+ output=out)
+ np.testing.assert_array_almost_equal(x, expected_x)
+ self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
+
+ # Test different bounds conditions.
+ verbose = minimize.Verbosity.FULLITER
+ for bounds in bounds_types.values():
+ out = io.StringIO()
+ x, trace = minimize.least_squares(x0, residual, bounds=bounds, output=out,
+ verbose=verbose)
+ self.assertContainsSubsequence(out.getvalue(), 'norm(step) < tol')
+ grad = trace[-1].jacobian.T @ trace[-1].residual
+ # If x_i is on the boundary, gradient points out, otherwise it is 0.
+ for i, xi in enumerate(x):
+ if xi == bounds[0][i]:
+ self.assertGreater(grad[i], 0)
+ elif xi == bounds[1][i]:
+ self.assertLess(grad[i], 0)
+ else:
+ self.assertAlmostEqual(grad[i], 0, places=4)
+
+ def test_bad_bounds(self) -> None:
+ def residual(x: np.ndarray) -> float:
+ return np.array([1 - x[0], 10 * (x[1] - x[0] ** 2)])
+
+ out = io.StringIO()
+ x0 = np.array((0.0, 0.0))
+
+ bad_bounds = [
+ [0, 1, 2],
+ [np.array((-2, 2, 0)), np.array((0.5, 3, 4))],
+ [np.array((-2, 2, 0)), np.array((0.5, 3, np.inf))],
+ [np.array((-2, 2, 0)), np.array((-5, 3, 6))],
+ ]
+
+ for bounds in bad_bounds:
+ with self.assertRaises(ValueError):
+ minimize.least_squares(x0, residual, bounds=bounds, output=out)
+
+if __name__ == '__main__':
+ absltest.main()
diff --git a/python/mujoco/render.cc b/python/mujoco/render.cc
index b2922efc..8be1fc69 100644
--- a/python/mujoco/render.cc
+++ b/python/mujoco/render.cc
@@ -14,6 +14,8 @@
#include
#include
+#include
+#include
#include
#include
@@ -40,7 +42,10 @@ class MjWrapper : public WrapperBase {
void Free();
-#define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(fogRGBA);
X(auxWidth);
X(auxHeight);
diff --git a/python/mujoco/structs.h b/python/mujoco/structs.h
index 5aba80f5..f27e17ff 100644
--- a/python/mujoco/structs.h
+++ b/python/mujoco/structs.h
@@ -17,12 +17,17 @@
#include
#include
+#include
#include
#include
#include
#include
+#include
+#include
#include
#include
+#include
+#include
#include
#include
@@ -36,6 +41,7 @@
namespace mujoco::python {
namespace _impl {
+
template
class WrapperBase {
public:
@@ -44,7 +50,7 @@ class WrapperBase {
T* get() { return ptr_; }
const T* get() const { return ptr_; }
- const pybind11::handle owner() const { return owner_; }
+ pybind11::handle owner() const { return owner_; }
protected:
static void DefaultCapsuleDestructor(PyObject* pyobj) {
@@ -167,7 +173,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(raw::MjOption* ptr, pybind11::handle owner);
~MjWrapper() = default;
- #define X(var, dim) py_array_or_tuple_t var;
+ #define X(var, dim) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var;
MJOPTION_VECTORS
#undef X
};
@@ -188,7 +197,10 @@ class MjWrapper
MjWrapper(raw::MjVisualHeadlight* ptr, pybind11::handle owner);
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(ambient);
X(diffuse);
X(specular);
@@ -209,7 +221,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(raw::MjVisualRgba* ptr, pybind11::handle owner);
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(fog);
X(haze);
X(force);
@@ -271,7 +286,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(raw::MjStatistic* ptr, pybind11::handle owner);
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(center);
#undef X
};
@@ -504,7 +522,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(raw::MjContact* ptr, pybind11::handle owner);
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(pos);
X(frame);
X(friction);
@@ -628,7 +649,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(refpos);
X(refquat);
X(refselpos);
@@ -650,7 +674,10 @@ class MjWrapper : public WrapperBase {
MjWrapper(MjWrapper&&) = default;
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(lookat);
#undef X
};
@@ -671,7 +698,10 @@ class MjWrapper : public WrapperBase {
explicit MjWrapper(raw::MjvGLCamera&& other);
~MjWrapper() = default;
- #define X(var) py_array_or_tuple_t var
+ #define X(var) \
+ py_array_or_tuple_t< \
+ std::remove_all_extents_t> \
+ var
X(pos);
X(forward);
X(up);
@@ -693,7 +723,10 @@ class MjWrapper