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

Merge origin/main into light_active PR

Browse Source
This commit is contained in:
Tarik Kelestemur
2026-05-30 16:24:27 -04:00
parent cf1b28d19a 50e823e91c
commit f8a12a827f
304 changed files with 17322 additions and 7702 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+10 -5
View File
@@ -29,7 +29,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
project(
mujoco
VERSION 3.9.0
VERSION 3.10.0
DESCRIPTION "MuJoCo Physics Simulator"
HOMEPAGE_URL "https://mujoco.org"
)
@@ -108,6 +108,7 @@ include(MujocoMacOS)
include(MujocoDependencies)
set(MUJOCO_HEADERS
include/mujoco/mjassert.h
include/mujoco/mjdata.h
include/mujoco/mjexport.h
include/mujoco/mjmacro.h
@@ -116,8 +117,8 @@ set(MUJOCO_HEADERS
include/mujoco/mjrender.h
include/mujoco/mjsan.h
include/mujoco/mjspec.h
include/mujoco/mjthread.h
include/mujoco/mjtnum.h
include/mujoco/mjspecmacro.h
include/mujoco/mjtype.h
include/mujoco/mjui.h
include/mujoco/mjvisualize.h
include/mujoco/mjxmacro.h
@@ -133,6 +134,7 @@ endif()
# Emscripten does not support SHARED libs
if(NOT EMSCRIPTEN)
find_package(Threads REQUIRED)
add_library(mujoco SHARED ${MUJOCO_RESOURCE_FILES})
else()
add_library(mujoco STATIC ${MUJOCO_RESOURCE_FILES})
@@ -157,7 +159,7 @@ add_subdirectory(plugin/stl_decoder)
add_subdirectory(src/engine)
add_subdirectory(src/user)
add_subdirectory(src/xml)
add_subdirectory(src/thread)
add_subdirectory(src/xml/mjz)
if(NOT EMSCRIPTEN AND NOT MUJOCO_USE_FILAMENT_MJR_COMPAT)
add_subdirectory(src/render/classic)
add_subdirectory(src/ui)
@@ -201,6 +203,10 @@ target_link_libraries(
tinyxml2
)
if(NOT EMSCRIPTEN)
target_link_libraries(mujoco PRIVATE Threads::Threads)
endif()
set_target_properties(
mujoco PROPERTIES VERSION "${mujoco_VERSION}" PUBLIC_HEADER "${MUJOCO_HEADERS}"
)
@@ -256,7 +262,6 @@ if(MUJOCO_BUILD_SIMULATE)
endif()
if(MUJOCO_BUILD_STUDIO)
add_subdirectory(src/xml/mjz)
add_subdirectory(src/experimental/platform)
add_subdirectory(src/experimental/studio)
endif()
cmake/MujocoDependencies.cmake
+2 -2
View File
@@ -39,7 +39,7 @@ set(MUJOCO_DEP_VERSION_qhull
CACHE STRING "Version of `qhull` to be fetched."
)
set(MUJOCO_DEP_VERSION_miniz
3.1.1
d10b03cc73475af673df40f06e5cefd1d5f940d9
CACHE STRING "Version of `miniz` to be fetched."
)
set(MUJOCO_DEP_VERSION_Eigen3
@@ -58,7 +58,7 @@ set(MUJOCO_DEP_VERSION_gtest
)
set(MUJOCO_DEP_VERSION_benchmark
5c55f5d4f45a1b09c5d98aa63a671993ebd42c69
834a61fc65e8b7885fcf177f1230ae4b897118fa
CACHE STRING "Version of `benchmark` to be fetched."
)
dist/mujoco.rc
Vendored
+4 -4
View File
@@ -1,6 +1,6 @@
1 VERSIONINFO
FILEVERSION 3,9,0,0
PRODUCTVERSION 3,9,0,0
FILEVERSION 3,10,0,0
PRODUCTVERSION 3,10,0,0
FILEOS 0x4
FILETYPE 0x1
{
@@ -9,9 +9,9 @@ FILETYPE 0x1
BLOCK "040904b0"
{
VALUE "ProductName", "MuJoCo"
VALUE "ProductVersion", "3.9.0"
VALUE "ProductVersion", "3.10.0"
VALUE "FileDescription", "MuJoCo"
VALUE "FileVersion", "3.9.0"
VALUE "FileVersion", "3.10.0"
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,9,0,0
PRODUCTVERSION 3,9,0,0
FILEVERSION 3,10,0,0
PRODUCTVERSION 3,10,0,0
FILEOS 0x4
FILETYPE 0x1
{
@@ -11,9 +11,9 @@ FILETYPE 0x1
BLOCK "040904b0"
{
VALUE "ProductName", "MuJoCo"
VALUE "ProductVersion", "3.9.0"
VALUE "ProductVersion", "3.10.0"
VALUE "FileDescription", "MuJoCo"
VALUE "FileVersion", "3.9.0"
VALUE "FileVersion", "3.10.0"
VALUE "InternalName", "simulate.exe"
VALUE "OriginalFilename", "simulate.exe"
VALUE "CompanyName", "Google DeepMind"
doc/APIreference/APIglobals.rst
+5 -8
View File
@@ -388,7 +388,7 @@ Defined in `mujoco.h <https://github.com/google-deepmind/mujoco/blob/main/includ
- value
- description
* - ``mjVERSION_HEADER``
- 3009000
- 3010000
- The version of the MuJoCo headers. This is an integer calculated from the version string "S.M.P"
using the formula ``(S * 1e6) + (M * 1e3) + P``. For example, version 4.2.1 is represented as 4002001.
The API function :ref:`mj_version` returns a number with the same meaning
@@ -415,7 +415,7 @@ indicated otherwise.
- 1E-15
- The minimal value allowed in any denominator, and in general any mathematical operation where 0 is not allowed.
In almost all cases, MuJoCo silently clamps smaller values to mjMINVAL.
Defined in `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`_.
Defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`_.
* - ``mjPI``
- :math:`\pi`
- The value of :math:`\pi`. This is used in various trigonometric functions, and also for conversion from degrees
@@ -452,10 +452,6 @@ indicated otherwise.
* - ``mjMINAWAKE``
- 10
- The minimum number of timesteps that must pass after a tree is awoken, before it is allowed to go back to sleep.
* - ``mjMAXTHREAD``
- 128
- Maximum number of 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>`_.
.. _glNumericSizes:
@@ -706,7 +702,8 @@ X Macros
^^^^^^^^
The X Macros are not needed in most user projects. They are used internally to allocate the model, and are also
available for users who know how to use this programming technique. See the header file `mjxmacro.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`_ for the actual definitions. They are
available for users who know how to use this programming technique. See the header files `mjxmacro.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`_ and `mjspecmacro.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspecmacro.h>`_ for the actual definitions. They are
particularly useful in writing MuJoCo wrappers for scripting languages, where dynamic structures matching the MuJoCo
data structures need to be constructed programmatically.
doc/APIreference/APItypes.rst
+31 -36
View File
@@ -46,7 +46,7 @@ MuJoCo defines a large number of types:
Primitive types
---------------
The two types below are defined in `mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__.
The three types below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
.. _mjtNum:
@@ -63,7 +63,7 @@ double-precision version will always be available. Thus it is safe to write user
However, our preference is to write code that works with either single or double precision. To this end we provide math
utility functions that are always defined with the correct floating-point type.
Note that changing ``mjUSESINGLE`` in ``mjtnum.h`` will not change how the library was compiled, and instead will
Note that changing ``mjUSESINGLE`` in ``mjtype.h`` will not change how the library was compiled, and instead will
result in numerous link errors. In general, the header files distributed with precompiled MuJoCo should never be
changed by the user.
@@ -84,13 +84,29 @@ changed by the user.
mjtByte
^^^^^^^
Byte type used to represent boolean variables.
Byte type used to represent small integers and binary data.
.. code-block:: C
typedef unsigned char mjtByte;
.. _mjtBool:
mjtBool
^^^^^^^
Boolean type used to represent true/false values.
.. code-block:: C
#ifndef __cplusplus
typedef _Bool mjtBool;
#else
typedef bool mjtBool;
#endif
.. _mjtSize:
mjtSize
@@ -115,7 +131,7 @@ All enum types use the ``mjt`` prefix.
Model
^^^^^
The enums below are defined in `mjmodel.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__.
The enums below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
.. _mjtDisableBit:
@@ -456,7 +472,7 @@ Formulas used to combine SDFs when calling mjc_distance and mjc_gradient.
Data
^^^^
The enums below are defined in `mjdata.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__.
The enums below are defined in `mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__.
@@ -936,6 +952,15 @@ MJCF element :ref:`statistic <statistic>`. One instance of it is embedded in mjM
.. mujoco-include:: mjStatistic
.. _mjPreContact:
mjPreContact
~~~~~~~~~~~~
This is the data structure holding information about one contact filled out by a narrowphase collision detector.
.. mujoco-include:: mjPreContact
.. _mjContact:
mjContact
@@ -988,36 +1013,6 @@ Asset cache used by the compiler to avoid repeated slow recompilation. See :ref:
.. mujoco-include:: mjCache
.. _mjtTaskStatus:
mjtTaskStatus
~~~~~~~~~~~~~
Status values for :ref:`mjTask`.
.. mujoco-include:: mjtTaskStatus
.. _mjTask:
mjTask
~~~~~~
This is a representation of a task to be run asynchronously inside of an :ref:`mjThreadPool` . It is created in the
:ref:`mju_threadPoolEnqueue` method of the :ref:`mjThreadPool` and is used to join the task at completion.
The ``status`` field uses values from :ref:`mjtTaskStatus`.
.. mujoco-include:: mjTask
.. _mjThreadPool:
mjThreadPool
~~~~~~~~~~~~
This is the data structure of the threadpool. It can only be constructed programmatically, and does not
have an analog in MJCF. In order to enable multi-threaded calculations, a pointer to an existing :ref:`mjThreadPool`
should be assigned to the ``mjData.threadpool``.
.. mujoco-include:: mjThreadPool
.. _tyStatStructure:
@@ -1780,7 +1775,7 @@ mjfCollision
.. code-block:: C
typedef int (*mjfCollision)(const mjModel* m, const mjData* d,
mjContact* con, int g1, int g2, mjtNum margin);
mjPreContact* con, int g1, int g2, mjtNum margin);
This is the function type of the callbacks in the collision table :ref:`mjCOLLISIONFUNC`.
doc/APIreference/functions.rst
+16 -50
View File
@@ -3375,59 +3375,14 @@ If no match, return NULL.
Threads
^^^^^^^
.. _mju_threadPoolCreate:
.. _mju_threadpool:
`mju_threadPoolCreate <#mju_threadPoolCreate>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
`mju_threadpool <#mju_threadpool>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_threadPoolCreate
.. mujoco-include:: mju_threadpool
Create a thread pool with the specified number of threads running.
.. _mju_bindThreadPool:
`mju_bindThreadPool <#mju_bindThreadPool>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_bindThreadPool
Adds a thread pool to mjData and configures it for multi-threaded use.
.. _mju_threadPoolEnqueue:
`mju_threadPoolEnqueue <#mju_threadPoolEnqueue>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_threadPoolEnqueue
Enqueue a task in a thread pool.
.. _mju_threadPoolDestroy:
`mju_threadPoolDestroy <#mju_threadPoolDestroy>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_threadPoolDestroy
Destroy a thread pool.
.. _mju_defaultTask:
`mju_defaultTask <#mju_defaultTask>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_defaultTask
Initialize an mjTask.
.. _mju_taskJoin:
`mju_taskJoin <#mju_taskJoin>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mju_taskJoin
Wait for a task to complete.
Create a thread pool with nthread worker threads.
.. _Standardmath:
@@ -4496,6 +4451,17 @@ Add sensor.
Add flex.
.. _mjs_makeFlex:
`mjs_makeFlex <#mjs_makeFlex>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_makeFlex
Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.
*Nullable:* ``type``, ``dof``, ``count``, ``cellcount``, ``spacing``, ``scale``, ``pos``, ``quat``, ``origin``, ``file``, ``vfs``
.. _mjs_addPair:
`mjs_addPair <#mjs_addPair>`__
doc/XMLreference.rst
+8 -9
View File
@@ -8514,15 +8514,14 @@ Extraction
:width: 30%
:target: https://github.com/google-deepmind/mujoco/blob/main/model/tactile/tactile.xml
The tactile sensor returns the penetration pressure and the sliding velocities in the tangent frame at given points
between the geom associated with the sensor and the SDF geoms in contact with it. We define the penetration pressure as
a function of the penetration depth :math:`p(d) = \frac{d}{d_{max}-d}`, which is zero at the surface and goes to
infinity as the maximum depth is reached. The sensor is associated with a geom and a mesh. It is activated by the
contact between its associated geom and other geoms. The vertices of the mesh, when positioned in the geom frame, are
the points at which sensor values are computed, so the dimension of the output is 3 times the number of vertices in the
mesh. The mesh must have 3 normal vectors per vertex, which are used to compute the tangent frame. If the penetration
depth is positive (no contact), then all values are 0 for the corresponding vertex. Only contacts with geoms of type SDF
contribute to the sensor output. The sensor can be visualized by enabling the visualization of contact points.
The tactile sensor returns the maximum penetration depth and the sliding velocities in the tangent frame at given points
between the geom associated with the sensor and the SDF geoms in contact with it. The sensor is associated with a geom
and a mesh. It is activated by the contact between its associated geom and other geoms. The vertices of the mesh, when
positioned in the geom frame, are the points at which sensor values are computed, so the dimension of the output is 3
times the number of vertices in the mesh. The mesh must have 3 normal vectors per vertex, which are used to compute the
tangent frame. If the penetration depth is positive (no contact), then all values are 0 for the corresponding vertex.
Only contacts with geoms of type SDF contribute to the sensor output. The sensor can be visualized by enabling the
visualization of contact points.
.. _sensor-tactile-geom:
doc/XMLschema.rst
+10
View File
@@ -1,3 +1,13 @@
..
DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
.. raw:: html
<div class="schema-controls" style="margin-bottom: 1em;">
<button onclick="document.querySelectorAll('.schema-small details').forEach(d => d.open = true)" class="sd-btn sd-btn-outline-primary sd-btn-sm">Expand All</button>
<button onclick="document.querySelectorAll('.schema-small details details').forEach(d => d.open = false)" class="sd-btn sd-btn-outline-secondary sd-btn-sm" style="margin-left: 0.5em;">Collapse All</button>
</div>
.. dropdown:: :ref:`mujoco<mujoco>` :octicon:`star`
:open:
doc/changelog.rst
+95 -50
View File
@@ -7,71 +7,116 @@ Upcoming version (not yet released)
General
^^^^^^^
- Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
- Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
:ref:`Exact diagonal <soExactDiag>` for details.
- The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
undesirable time dependence.
- Flexes are now allowed to sleep, with the exception of completely passive (constraint-free) flexes.
- Added compiler timing diagnostics via the new :ref:`mjtCTimer` enum and the :ref:`mjs_getTimer` C API. After
:ref:`mj_compile`, per-category timings (total, assets, mesh loading, convex hull, normals, inertia, BVH, octree,
textures) are available via ``mjs_getTimer(spec)``. The :ref:`compile<saCompile>` sample prints a detailed timing
breakdown when run without an output file.
- Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp<body-flexcomp>` element for
programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as
``body.make_flex()`` in Python.
- Added :ref:`mju_threadpool`, a new function for creating a thread pool on an ``mjData`` instance. When a thread pool
is initialized, parts of the simulation pipeline, such as collision detection and constraint solving across islands,
are parallelized. The thread pool is automatically destroyed when the ``mjData`` is freed.
.. admonition:: Breaking API changes
:class: attention
- The header file ``mjthread.h`` was removed along with the engine threading API.
**Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine.
- Moved island sparse matrix construction from :ref:`mj_island` (single threaded) into :ref:`mj_fwdConstraint`
(multi-threaded). The island-specific matrices ``iM, iLD, iefc_J`` were removed from the arena and are now
allocated on the stack.
Bug fixes
^^^^^^^^^
- Fixed a bug in the ``mjz`` :ref:`decoder <mjpDecoder>` where unnormalized paths would fail to be read.
Version 3.9.0 (May 27, 2026)
----------------------------
General
^^^^^^^
1. Added ``mjData.efc_Y``, the whitened constraint Jacobian :math:`Y = J M^{-1/2}`, allocated in the arena when
dual solvers (PGS or NoSlip) are used or when :ref:`diagexact<option-flag-diagexact>` is enabled.
2. Added the :ref:`diagexact<option-flag-diagexact>` enable flag, which computes the exact diagonal of the
constraint-space inertia matrix at the current configuration, replacing the default compile-time approximation.
This improves solver quality for models with anisotropic inertias or complex kinematic coupling. See
:ref:`Exact diagonal <soExactDiag>` for details.
3. The pseudo-random constraint visitation order in the :ref:`PGS solver<soAlgorithms>`, introduced in the previous
release, now uses a fixed seed. The previous implementation seeded with ``mjData.time``, which introduced subtle yet
undesirable time dependence.
4. Flexes are now allowed to sleep, with the exception of completely passive (constraint-free) flexes.
5. Added compiler timing diagnostics via the new :ref:`mjtCTimer` enum and the :ref:`mjs_getTimer` C API. After
:ref:`mj_compile`, per-category timings (total, assets, mesh loading, convex hull, normals, inertia, BVH, octree,
textures) are available via ``mjs_getTimer(spec)``. The :ref:`compile<saCompile>` sample prints a detailed timing
breakdown when run without an output file.
6. Added :ref:`mjtBool` to represent boolean variables, replacing :ref:`mjtByte` across all boolean fields in
:ref:`mjModel`, :ref:`mjData`, and public C API function signatures.
.. admonition:: Breaking API changes
:class: attention
- The semantics of the contact ``margin`` and ``gap`` parameters have been redesigned for conceptual clarity and
consistency with `Newton <https://github.com/newton-physics/newton>`__. See the new
:ref:`margin and gap<coMarginGap>` documentation section for details.
7. The semantics of the contact ``margin`` and ``gap`` parameters have been redesigned for conceptual clarity and
consistency with `Newton <https://github.com/newton-physics/newton>`__. See the new
:ref:`margin and gap<coMarginGap>` documentation section for details.
Previously, ``margin`` controlled the *detection threshold* (contacts exist when ``dist < margin``) and ``gap``
was subtracted from it to produce the *force threshold* (forces generated when ``dist < margin - gap``). This was
unintuitive: users expected ``margin`` to mean geometric inflation and ``gap`` to mean a spatial gap.
Previously, ``margin`` controlled the *detection threshold* (contacts exist when ``dist < margin``) and ``gap``
was subtracted from it to produce the *force threshold* (forces generated when ``dist < margin - gap``). This was
unintuitive: users expected ``margin`` to mean geometric inflation and ``gap`` to mean a spatial gap.
Under the new semantics, ``margin`` is the geometric inflation of the geom surface and ``gap`` is an additional
detection buffer beyond the inflated surface:
Under the new semantics, ``margin`` is the geometric inflation of the geom surface and ``gap`` is an additional
detection buffer beyond the inflated surface:
- **Detection**: contacts are created when ``dist < margin + gap``.
- **Force generation**: constraint forces are applied when ``dist < margin``.
- **Inactive contacts**: contacts with ``margin < dist ≤ margin + gap`` are included in ``mjData.contact`` but
generate no force (``efc_address = -1``). This is useful for :ref:`adhesion<actuator-adhesion>` actuators and
custom callbacks.
- **Detection**: contacts are created when ``dist < margin + gap``.
- **Force generation**: constraint forces are applied when ``dist < margin``.
- **Inactive contacts**: contacts with ``margin < dist ≤ margin + gap`` are included in ``mjData.contact`` but
generate no force (``efc_address = -1``). This is useful for :ref:`adhesion<actuator-adhesion>` actuators and
custom callbacks.
With the default values ``margin = 0``, ``gap = 0``, the behavior is unchanged.
With the default values ``margin = 0``, ``gap = 0``, the behavior is unchanged.
.. image:: images/modeling/margin_gap_light.svg
:width: 80%
:align: center
:class: only-light
.. image:: images/modeling/margin_gap_light.svg
:width: 80%
:align: center
:class: only-light
.. image:: images/modeling/margin_gap_dark.svg
:width: 80%
:align: center
:class: only-dark
.. image:: images/modeling/margin_gap_dark.svg
:width: 80%
:align: center
:class: only-dark
|
|
**Migration:** Models that use the default ``gap="0"`` (the vast majority) require no changes. For models with
``gap > 0``, apply the following transformation to preserve identical behavior:
**Migration:** Models that use the default ``gap="0"`` (the vast majority) require no changes. For models with
``gap > 0``, apply the following transformation to preserve identical behavior:
.. code-block::
.. code-block::
margin_new = margin_old - gap_old
gap_new = gap_old
margin_new = margin_old - gap_old
gap_new = gap_old
For example, a geom with the old attributes ``margin="0.1" gap="0.1"`` should be changed to
``margin="0" gap="0.1"``.
For example, a geom with the old attributes ``margin="0.1" gap="0.1"`` should be changed to
``margin="0" gap="0.1"``.
Negative ``margin`` values are now permitted (corresponding to ``gap > margin`` under the old semantics). The
constraint ``margin + gap >= 0`` should be maintained to ensure valid collision detection.
Negative ``margin`` values are now permitted (corresponding to ``gap > margin`` under the old semantics). The
constraint ``margin + gap >= 0`` should be maintained to ensure valid collision detection.
8. The :ref:`mjfCollision` functions now populate the :ref:`mjPreContact` struct instead of the :ref:`mjContact`
struct. The :ref:`mjPreContact` only contains the necessary fields needed for the narrowphase collision detection.
9. The header file ``mjtnum.h`` was renamed to
`mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>` and now includes all
enum type definitions.
10. The :ref:`tactile<sensor-tactile>` sensor now reports raw depth instead of an estimated pressure.
11. MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
``naconmax``.
12. Maybe-breaking: Added `mjassert.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__, a new header containing
compile-time assertions that verify the sizes of MuJoCo's public types for ABI stability. This is a first step
towards replacing ``int`` with strongly-typed enums in the public API. If these assertions fail on your compiler or
platform, please report the issue on GitHub.
- MJX: Removed the deprecated ``nconmax`` argument from ``mjx.make_data`` and ``mjx.put_data`` in favor of
``naconmax``.
Version 3.8.1 (May 11, 2026)
----------------------------
@@ -1191,7 +1236,7 @@ General
10. Added :ref:`mj_setKeyframe` for saving the current state into a model keyframe.
11. Added support for ``ball`` joints in the URDF parser ("spherical" in URDF).
12. Replaced ``mjUSEDOUBLE`` which was previously hard-coded in
`mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
`mjtnum.h <https://github.com/google-deepmind/mujoco/blob/3577e2cf8bf841475b489aefff52276a39f24d51/include/mjtnum.h>`__
with the build-time flag ``mjUSESINGLE``. If this symbol is not defined, MuJoCo will use double-precision floating
point, as usual. If ``mjUSESINGLE`` is defined, MuJoCo will use single-precision floating point. See :ref:`mjtNum`.
@@ -1677,7 +1722,7 @@ New features
If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
5. Added ``mjThreadPool`` and ``mjTask`` which allow for multi-threaded operations within the MuJoCo engine
pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
- Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
doc/conf.py
+2
View File
@@ -168,6 +168,7 @@ SHARED_CSS_VARIABLES = {
# font-stack--monospace used in code blocks, Inconsolata fits in 100 chars.
html_theme_options = {
'light_css_variables': {
'color-brand-visited': 'var(--color-brand-content)',
'font-stack--monospace': 'Inconsolata,Consolas,ui-monospace,monospace',
'at-color': '#830b2b',
'at-val-color': '#bc103e',
@@ -182,6 +183,7 @@ html_theme_options = {
'wy-nav-side-background-color': '#0053d6',
},
'dark_css_variables': {
'color-brand-visited': 'var(--color-brand-content)',
'at-color': '#ffaab7',
'at-val-color': '#ff95a6',
'body-color': '#14234b',
doc/css/theme_overrides.css
-5
View File
@@ -197,11 +197,6 @@ h4 {
padding: 6px;
}
/* Don't change color of visited links.
.rst-content a.reference:visited {
color: var(--rst-content-a-color);
} */
.rst-content code, .rst-content tt {
font-size: 90%;
padding: inherit;
doc/includes/references.h
+500 -518
View File
File diff suppressed because it is too large Load Diff
doc/mjwarp/index.rst
+3 -30
View File
@@ -228,36 +228,9 @@ Feature Parity
==============
MJWarp supports most of the main simulation features of MuJoCo, with a few exceptions. MJWarp will raise an exception if
asked to copy to device an :ref:`mjModel` with field values referencing unsupported features.
The following features are **not supported** in MJWarp:
.. list-table::
:width: 90%
:align: left
:widths: 2 5
:header-rows: 1
* - Category
- Feature
* - :ref:`Integrator <mjtIntegrator>`
- ``IMPLICIT``, ``IMPLICITFAST`` not supported with fluid drag
* - :ref:`Solver <mjtSolver>`
- ``PGS``, ``noslip``, :ref:`islands <soIsland>`
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Sensors <mjtSensor>`
- ``GEOMDIST``, ``GEOMNORMAL``, ``GEOMFROMTO``
* - Flex
- ``VERTCOLLIDE=false``, ``INTERNAL=true``
* - Jacobian format
- ``SPARSE``
* - Option
- :ref:`contact override <COverride>`
* - Plugins
- ``All`` except ``SDF``
* - :ref:`User parameters <CUser>`
- ``All``
asked to copy to device an :ref:`mjModel` with field values referencing unsupported features. For the most up-to-date
feature availability, please see
`MuJoCo API Compatibility <https://github.com/google-deepmind/mujoco_warp#mujoco-api-compatibility>`__.
.. _mjwPerf:
doc/mjx.rst
+2 -2
View File
@@ -473,7 +473,7 @@ The following table compares feature support between MJX-Warp and MJX-JAX compar
- All
- ``CONNECT``, ``WELD``, ``JOINT``, ``TENDON``
* - :ref:`Integrator <mjtIntegrator>`
- All except ``IMPLICIT``
- All except ``IMPLICITFAST`` midpoint integrator feature
- ``EULER``, ``RK4``, ``IMPLICITFAST`` (``IMPLICITFAST`` not supported with :doc:`fluid drag <computation/fluid>`)
* - :ref:`Cone <mjtCone>`
- All
@@ -503,8 +503,8 @@ The following table compares feature support between MJX-Warp and MJX-JAX compar
- Sparse and Dense
- Sparse and Dense
* - Jacobian format
- ``DENSE`` only
- ``DENSE`` and ``SPARSE``
- ``DENSE`` only
* - Lights
- ✓
- Positions and directions
doc/overview.rst
+2 -2
View File
@@ -875,8 +875,8 @@ The size of each array (``njnt`` in this case) is also given in ``mjModel``. The
first, followed by the limits of the second joint etc. This ordering reflects the fact that all matrices in MuJoCo have
row-major format.
The available element types are defined in `mjmodel.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h#L237>`_, in the enum type :ref:`mjtObj`.
The available element types are defined in `mjtype.h
<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__, in the enum type :ref:`mjtObj`.
These enums are mostly used internally. One exception are the functions :ref:`mj_name2id` and :ref:`mj_id2name` in the
MuJoCo API, which map element names to integer ids and vice versa. These functions take an element type as input.
doc/programming/index.rst
+15 -13
View File
@@ -31,9 +31,8 @@ OpenGL renderer
state-of-the-art rendering engines (and can be replaced with such an engine if desired) but nevertheless it provides
efficient and informative 3D rendering.
Thread
The threading framework is written in C++ and exposed in C. It provides a :ref:`mjThreadPool<mjThreadPool>` interface
to process tasks asynchronously. To enable use in MuJoCo, create a thread pool and assign it to the
``mjData.threadpool`` field.
The threading framework is written in C++ and exposed in C. It provides a thread pool interface
to process tasks asynchronously. To enable use in MuJoCo, call ``mju_threadpool``.
UI framework
The UI framework is written in C. UI elements are rendered in OpenGL. It has its own event
mechanism and abstract hooks for keyboard and mouse input. The code samples use it with GLFW, but it can also be used
@@ -151,39 +150,43 @@ links below, to make this documentation self-contained.
`mujoco.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mujoco.h>`__
This is the main header file and must be included in all programs using MuJoCo. It defines all API functions and
global variables, and includes all other header files except mjxmacro.h.
global variables, and includes all other header files except mjxmacro.h and mjspecmacro.h.
`mjmodel.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmodel.h>`__
Defines the C structure :ref:`mjModel` which is the runtime representation of the
model being simulated. It also defines a number of primitive types and other structures needed to define mjModel.
model being simulated.
`mjdata.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__
Defines the C structure :ref:`mjData` which is the workspace where all computations
read their inputs and write their outputs. It also defines primitive types and other structures needed to define
mjData.
read their inputs and write their outputs.
`mjvisualize.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjvisualize.h>`__
Defines the primitive types and structures needed by the abstract visualizer.
`mjrender.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjrender.h>`__
Defines the primitive types and structures needed by the OpenGL renderer.
`mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`__
Defines the primitive types and structures needed by the UI framework.
`mjtnum.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtnum.h>`__
Defines MuJoCo's ``mjtNum`` floating-point type to be either ``double`` or ``float``. See :ref:`mjtNum`.
`mjtype.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjtype.h>`__
Defines primitive types and enums, including the ``mjtNum`` floating-point type to be either ``double`` or ``float``
(see :ref:`mjtNum`).
`mjspec.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspec.h>`__
Defines enums and structs used for :doc:`procedural model editing <modeledit>`.
`mjplugin.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjplugin.h>`__
Defines data structures required by :ref:`engine plugins<exPlugin>`.
`mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`__
Defines data structures and functions required by :ref:`thread<Thread>`.
`mjmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjmacro.h>`__
Defines C macros that are useful in user code.
`mjxmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`__
This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
automate the mapping of mjModel and mjData into scripting languages, as well as other operations that require
accessing all fields of mjModel and mjData.
`mjspecmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjspecmacro.h>`__
This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
automate the mapping of mjSpec and its element structs into scripting languages, as well as other operations that require
accessing all fields of mjSpec during procedural model editing.
`mjexport.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjexport.h>`__
Macros used for exporting public symbols from the MuJoCo library. This header should not be used directly by client
code.
`mjsan.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjsan.h>`__
Definitions required when building with sanitizer instrumentation.
`mjassert.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjassert.h>`__
Compile-time size assertions verifying MuJoCo ABI stability across C and C++ compilers.
.. _inVersion:
@@ -228,8 +231,7 @@ to which the symbol belongs. First we list the prefixes corresponding to type de
Core simulation data structure (C struct), for example :ref:`mjModel`. If all characters
after the prefix are capital, for example :ref:`mjMIN`, this is a macro or a symbol (#define).
``mjt``
Primitive type, for example :ref:`mjtGeom`. Except for mjtByte and mjtNum, all other
definitions in this family are enums.
Primitive type, for example :ref:`mjtNum` and :ref:`mjtGeom`. Most types in this family are enums.
``mjf``
Callback function type, for example :ref:`mjfGeneric`.
``mjv``
doc/programming/samples.rst
+1 -1
View File
@@ -60,7 +60,7 @@ Where the command line arguments are
keyframe named "test" is present in the model, it is used as the initial state.
- The ``ctrlnoise`` argument prevents models from settling into a static state where, due to warmstarts, one can
measure artificially faster simulation.
- When ``npoolthread > 1`` is specified, an engine-internal :ref:`mjThreadPool` is created with the specified number of
- When ``npoolthread > 1`` is specified, an engine-internal thread pool is created with the specified number of
threads, to speed up simulation of large scenes. Note that while it is possible to use both ``nthread`` and
``npoolthread``, the scenarios for which one would want these different types of multithreading are usually mutually
exclusive.
doc/programming/simulation.rst
+3 -4
View File
@@ -562,10 +562,9 @@ external force computed by inverse dynamics.
Multi-threading
~~~~~~~~~~~~~~~
MuJoCo has experimental support for within-step multi-threading. When a :ref:`mjThreadPool` is assigned to
``mjData.threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
:ref:`islands<siSleep>` — can be distributed across worker threads. Note that within-step threading currently has
significant memory overhead and is still a work in progress.
MuJoCo has support for within-step multi-threading. When a thread pool is initialized via
``mju_threadpool``, parts of the simulation pipeline — such as collision detection and constraint solving across
:ref:`islands<siSleep>` — can be distributed across worker threads.
The more common and well-supported use of multi-threading is to speed up sampling operations that are
common in more advanced applications. Simulation is inherently serial over time (the output of one mj_step is the
doc/unity.rst
+2 -2
View File
@@ -37,14 +37,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.9.0.dylib`` (it can be
``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.10.0.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.9.0/lib/libmujoco.so.3.9.0`` and rename it as ``libmujoco.so``.
``~/.mujoco/mujoco-3.10.0/lib/libmujoco.so.3.10.0`` and rename it as ``libmujoco.so``.
Windows
_______
include/mujoco/experimental/usd/mjcPhysics/actuator.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_ACTUATOR_H
#define MJCPHYSICS_GENERATED_ACTUATOR_H
include/mujoco/experimental/usd/mjcPhysics/api.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_API_H
#define MJCPHYSICS_API_H
include/mujoco/experimental/usd/mjcPhysics/collisionAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_COLLISIONAPI_H
#define MJCPHYSICS_GENERATED_COLLISIONAPI_H
include/mujoco/experimental/usd/mjcPhysics/equalityAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_EQUALITYAPI_H
#define MJCPHYSICS_GENERATED_EQUALITYAPI_H
include/mujoco/experimental/usd/mjcPhysics/equalityConnectAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_EQUALITYCONNECTAPI_H
#define MJCPHYSICS_GENERATED_EQUALITYCONNECTAPI_H
include/mujoco/experimental/usd/mjcPhysics/equalityJointAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
#define MJCPHYSICS_GENERATED_EQUALITYJOINTAPI_H
include/mujoco/experimental/usd/mjcPhysics/equalityWeldAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_EQUALITYWELDAPI_H
#define MJCPHYSICS_GENERATED_EQUALITYWELDAPI_H
include/mujoco/experimental/usd/mjcPhysics/imageableAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_IMAGEABLEAPI_H
#define MJCPHYSICS_GENERATED_IMAGEABLEAPI_H
include/mujoco/experimental/usd/mjcPhysics/jointAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_JOINTAPI_H
#define MJCPHYSICS_GENERATED_JOINTAPI_H
include/mujoco/experimental/usd/mjcPhysics/keyframe.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_KEYFRAME_H
#define MJCPHYSICS_GENERATED_KEYFRAME_H
include/mujoco/experimental/usd/mjcPhysics/materialAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_MATERIALAPI_H
#define MJCPHYSICS_GENERATED_MATERIALAPI_H
include/mujoco/experimental/usd/mjcPhysics/meshCollisionAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H
#define MJCPHYSICS_GENERATED_MESHCOLLISIONAPI_H
include/mujoco/experimental/usd/mjcPhysics/sceneAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_SCENEAPI_H
#define MJCPHYSICS_GENERATED_SCENEAPI_H
include/mujoco/experimental/usd/mjcPhysics/siteAPI.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_SITEAPI_H
#define MJCPHYSICS_GENERATED_SITEAPI_H
include/mujoco/experimental/usd/mjcPhysics/tendon.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_GENERATED_TENDON_H
#define MJCPHYSICS_GENERATED_TENDON_H
include/mujoco/experimental/usd/mjcPhysics/tokens.h
+2
View File
@@ -12,6 +12,8 @@
// See the License for the specific language governing permissions and
// limitations under the License.
// DO NOT EDIT. THIS FILE IS AUTOMATICALLY GENERATED.
#ifndef MJCPHYSICS_TOKENS_H
#define MJCPHYSICS_TOKENS_H
include/mujoco/mjassert.h
+130
View File
@@ -0,0 +1,130 @@
// Copyright 2026 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_MJASSERT_H_
#define MUJOCO_MJASSERT_H_
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjrender.h>
#include <mujoco/mjspec.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjui.h>
#include <mujoco/mjvisualize.h>
#if defined(__cplusplus)
#define MJ_ASSERT_SIZE(type, size) \
static_assert(sizeof(type) == (size), #type " must be " #size " bytes for MuJoCo ABI stability")
#elif defined(__STDC_VERSION__) && __STDC_VERSION__ >= 201112L
#define MJ_ASSERT_SIZE(type, size) \
_Static_assert(sizeof(type) == (size), #type " must be " #size " bytes for MuJoCo ABI stability")
#else
#define MJ_ASSERT_SIZE(type, size) \
typedef char mj_assert_##type[sizeof(type) == (size) ? 1 : -1]
#endif
// primitive types
#if !defined(mjUSESINGLE)
MJ_ASSERT_SIZE(mjtNum, 8);
#else
MJ_ASSERT_SIZE(mjtNum, 4);
#endif
MJ_ASSERT_SIZE(mjtByte, 1);
MJ_ASSERT_SIZE(mjtBool, 1);
MJ_ASSERT_SIZE(mjtSize, 8);
// mjModel enums
MJ_ASSERT_SIZE(mjtDisableBit, 4);
MJ_ASSERT_SIZE(mjtEnableBit, 4);
MJ_ASSERT_SIZE(mjtJoint, 4);
MJ_ASSERT_SIZE(mjtGeom, 4);
MJ_ASSERT_SIZE(mjtProjection, 4);
MJ_ASSERT_SIZE(mjtCamLight, 4);
MJ_ASSERT_SIZE(mjtLightType, 4);
MJ_ASSERT_SIZE(mjtTexture, 4);
MJ_ASSERT_SIZE(mjtTextureRole, 4);
MJ_ASSERT_SIZE(mjtColorSpace, 4);
MJ_ASSERT_SIZE(mjtIntegrator, 4);
MJ_ASSERT_SIZE(mjtCone, 4);
MJ_ASSERT_SIZE(mjtJacobian, 4);
MJ_ASSERT_SIZE(mjtSolver, 4);
MJ_ASSERT_SIZE(mjtEq, 4);
MJ_ASSERT_SIZE(mjtWrap, 4);
MJ_ASSERT_SIZE(mjtTrn, 4);
MJ_ASSERT_SIZE(mjtDyn, 4);
MJ_ASSERT_SIZE(mjtGain, 4);
MJ_ASSERT_SIZE(mjtBias, 4);
MJ_ASSERT_SIZE(mjtObj, 4);
MJ_ASSERT_SIZE(mjtSensor, 4);
MJ_ASSERT_SIZE(mjtStage, 4);
MJ_ASSERT_SIZE(mjtDataType, 4);
MJ_ASSERT_SIZE(mjtConDataField, 4);
MJ_ASSERT_SIZE(mjtRayDataField, 4);
MJ_ASSERT_SIZE(mjtCamOutBit, 4);
MJ_ASSERT_SIZE(mjtSameFrame, 4);
MJ_ASSERT_SIZE(mjtSleepPolicy, 4);
MJ_ASSERT_SIZE(mjtLRMode, 4);
MJ_ASSERT_SIZE(mjtFlexSelf, 4);
MJ_ASSERT_SIZE(mjtSDFType, 4);
// mjData enums
MJ_ASSERT_SIZE(mjtState, 4);
MJ_ASSERT_SIZE(mjtConstraint, 4);
MJ_ASSERT_SIZE(mjtConstraintState, 4);
MJ_ASSERT_SIZE(mjtWarning, 4);
MJ_ASSERT_SIZE(mjtTimer, 4);
MJ_ASSERT_SIZE(mjtSleepState, 4);
// mjspec.h
MJ_ASSERT_SIZE(mjtGeomInertia, 4);
MJ_ASSERT_SIZE(mjtMeshInertia, 4);
MJ_ASSERT_SIZE(mjtMeshBuiltin, 4);
MJ_ASSERT_SIZE(mjtBuiltin, 4);
MJ_ASSERT_SIZE(mjtMark, 4);
MJ_ASSERT_SIZE(mjtLimited, 4);
MJ_ASSERT_SIZE(mjtAlignFree, 4);
MJ_ASSERT_SIZE(mjtInertiaFromGeom, 4);
MJ_ASSERT_SIZE(mjtOrientation, 4);
// mjvisualize.h
MJ_ASSERT_SIZE(mjtCatBit, 4);
MJ_ASSERT_SIZE(mjtMouse, 4);
MJ_ASSERT_SIZE(mjtPertBit, 4);
MJ_ASSERT_SIZE(mjtCamera, 4);
MJ_ASSERT_SIZE(mjtLabel, 4);
MJ_ASSERT_SIZE(mjtFrame, 4);
MJ_ASSERT_SIZE(mjtVisFlag, 4);
MJ_ASSERT_SIZE(mjtRndFlag, 4);
MJ_ASSERT_SIZE(mjtStereo, 4);
// mjrender.h
MJ_ASSERT_SIZE(mjtGridPos, 4);
MJ_ASSERT_SIZE(mjtFramebuffer, 4);
MJ_ASSERT_SIZE(mjtDepthMap, 4);
MJ_ASSERT_SIZE(mjtFontScale, 4);
MJ_ASSERT_SIZE(mjtFont, 4);
// mjui.h
MJ_ASSERT_SIZE(mjtButton, 4);
MJ_ASSERT_SIZE(mjtEvent, 4);
MJ_ASSERT_SIZE(mjtItem, 4);
MJ_ASSERT_SIZE(mjtSection, 4);
// mjplugin.h
MJ_ASSERT_SIZE(mjtPluginCapabilityBit, 4);
#undef MJ_ASSERT_SIZE
#endif // MUJOCO_MJASSERT_H_
include/mujoco/mjdata.h
+21 -121
View File
@@ -18,113 +18,24 @@
#include <stddef.h>
#include <stdint.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjthread.h>
//---------------------------------- primitive types (mjt) -----------------------------------------
typedef enum mjtState_ { // state elements
mjSTATE_TIME = 1<<0, // time
mjSTATE_QPOS = 1<<1, // position
mjSTATE_QVEL = 1<<2, // velocity
mjSTATE_ACT = 1<<3, // actuator activation
mjSTATE_HISTORY = 1<<4, // history buffers (control, sensor)
mjSTATE_WARMSTART = 1<<5, // acceleration used for warmstart
mjSTATE_CTRL = 1<<6, // control
mjSTATE_QFRC_APPLIED = 1<<7, // applied generalized force
mjSTATE_XFRC_APPLIED = 1<<8, // applied Cartesian force/torque
mjSTATE_EQ_ACTIVE = 1<<9, // enable/disable constraints
mjSTATE_MOCAP_POS = 1<<10, // positions of mocap bodies
mjSTATE_MOCAP_QUAT = 1<<11, // orientations of mocap bodies
mjSTATE_USERDATA = 1<<12, // user data
mjSTATE_PLUGIN = 1<<13, // plugin state
mjNSTATE = 14, // number of state elements
// convenience values for commonly used state specifications
mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
mjSTATE_FULLPHYSICS = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
mjSTATE_USERDATA,
mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
} mjtState;
typedef enum mjtConstraint_ { // type of constraint
mjCNSTR_EQUALITY = 0, // equality constraint
mjCNSTR_FRICTION_DOF, // dof friction
mjCNSTR_FRICTION_TENDON, // tendon friction
mjCNSTR_LIMIT_JOINT, // joint limit
mjCNSTR_LIMIT_TENDON, // tendon limit
mjCNSTR_CONTACT_FRICTIONLESS, // frictionless contact
mjCNSTR_CONTACT_PYRAMIDAL, // frictional contact, pyramidal friction cone
mjCNSTR_CONTACT_ELLIPTIC // frictional contact, elliptic friction cone
} mjtConstraint;
typedef enum mjtConstraintState_ { // constraint state
mjCNSTRSTATE_SATISFIED = 0, // constraint satisfied, zero cost (limit, contact)
mjCNSTRSTATE_QUADRATIC, // quadratic cost (equality, friction, limit, contact)
mjCNSTRSTATE_LINEARNEG, // linear cost, negative side (friction)
mjCNSTRSTATE_LINEARPOS, // linear cost, positive side (friction)
mjCNSTRSTATE_CONE // squared distance to cone cost (elliptic contact)
} mjtConstraintState;
//------------------------------------- Contact ----------------------------------------------------
typedef enum mjtWarning_ { // warning types
mjWARN_INERTIA = 0, // (near) singular inertia matrix
mjWARN_CONTACTFULL, // too many contacts in contact list
mjWARN_CNSTRFULL, // too many constraints
mjWARN_BADQPOS, // bad number in qpos
mjWARN_BADQVEL, // bad number in qvel
mjWARN_BADQACC, // bad number in qacc
mjWARN_BADCTRL, // bad number in ctrl
mjNWARNING // number of warnings
} mjtWarning;
typedef enum mjtTimer_ { // internal timers
// main api
mjTIMER_STEP = 0, // step
mjTIMER_FORWARD, // forward
mjTIMER_INVERSE, // inverse
// breakdown of step/forward
mjTIMER_POSITION, // fwdPosition
mjTIMER_VELOCITY, // fwdVelocity
mjTIMER_ACTUATION, // fwdActuation
mjTIMER_CONSTRAINT, // fwdConstraint
mjTIMER_ADVANCE, // mj_Euler, mj_implicit
// breakdown of fwdPosition
mjTIMER_POS_KINEMATICS, // kinematics, com, tendon, transmission
mjTIMER_POS_INERTIA, // inertia computations
mjTIMER_POS_COLLISION, // collision detection
mjTIMER_POS_MAKE, // make constraints
mjTIMER_POS_PROJECT, // project constraints
// breakdown of mj_collision
mjTIMER_COL_BROAD, // broadphase
mjTIMER_COL_NARROW, // narrowphase
mjNTIMER // number of timers
} mjtTimer;
typedef enum mjtSleepState_ { // sleep state of an object
mjS_STATIC = -1, // object is static
mjS_ASLEEP = 0, // object is asleep
mjS_AWAKE = 1 // object is awake
} mjtSleepState;
//---------------------------------- mjContact -----------------------------------------------------
struct mjPreContact_ { // contact parameters set by narrowphase collision functions
mjtNum dist;
mjtNum pos[3];
mjtNum normal[3]; // contact normal of the collision
mjtNum tangent[3]; // first tangent direction
};
typedef struct mjPreContact_ mjPreContact;
struct mjContact_ { // result of collision detection functions
// contact parameters set by near-phase collision function
// contact parameters set by narrowphase collision function
mjtNum dist; // distance between nearest points; neg: penetration
mjtNum pos[3]; // position of contact point: midpoint between geoms
mjtNum frame[9]; // normal is in [0-2], points from geom[0] to geom[1]
@@ -201,9 +112,12 @@ struct mjData_ {
// arena pointer
size_t parena; // first available byte in arena
// threading
uintptr_t threadpool; // thread pool pointer
mjtBool threadlock; // disable stack freeing during threaded execution
// memory utilization statistics
mjtSize maxuse_stack; // maximum stack allocation in bytes (mutable)
mjtSize maxuse_threadstack[mjMAXTHREAD]; // maximum stack allocation per thread in bytes
mjtSize maxuse_arena; // maximum arena allocation in bytes
int maxuse_con; // maximum number of contacts
int maxuse_efc; // maximum number of scalar constraints
@@ -235,10 +149,10 @@ struct mjData_ {
int nv_awake; // number of awake dofs
// flags marking lazily evaluated stages
mjtByte flg_energypos; // has mj_energyPos been called
mjtByte flg_energyvel; // has mj_energyVel been called
mjtByte flg_subtreevel; // has mj_subtreeVel been called
mjtByte flg_rnepost; // has mj_rnePostConstraint been called
mjtBool flg_energypos; // has mj_energyPos been called
mjtBool flg_energyvel; // has mj_energyVel been called
mjtBool flg_subtreevel; // has mj_subtreeVel been called
mjtBool flg_rnepost; // has mj_rnePostConstraint been called
// global properties
mjtNum time; // simulation time
@@ -264,7 +178,7 @@ struct mjData_ {
mjtNum* ctrl; // control (nu x 1)
mjtNum* qfrc_applied; // applied generalized force (nv x 1)
mjtNum* xfrc_applied; // applied Cartesian force/torque (nbody x 6)
mjtByte* eq_active; // enable/disable constraints (neq x 1)
mjtBool* eq_active; // enable/disable constraints (neq x 1)
// mocap data
mjtNum* mocap_pos; // positions of mocap bodies (nmocap x 3)
@@ -345,7 +259,7 @@ struct mjData_ {
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
// computed by mj_collision/mj_collideTree
mjtByte* bvh_active; // was bounding volume checked for collision (nbvh x 1)
mjtBool* bvh_active; // was bounding volume checked for collision (nbvh x 1)
// computed by mj_updateSleep
int* tree_awake; // is tree awake; 0: asleep; 1: awake (ntree x 1)
@@ -450,12 +364,6 @@ struct mjData_ {
// computed by mj_island (dofs sorted by island)
mjtNum* ifrc_smooth; // net unconstrained force (nidof x 1)
mjtNum* iacc_smooth; // unconstrained acceleration (nidof x 1)
int* iM_rownnz; // inertia: non-zeros in each row (nidof x 1)
int* iM_rowadr; // inertia: address of each row in iM_colind (nidof x 1)
int* iM_colind; // inertia: column indices of non-zeros (nC x 1)
mjtNum* iM; // total inertia (sparse) (nC x 1)
mjtNum* iLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* iLDiagInv; // 1/diag(D) (nidof x 1)
mjtNum* iacc; // acceleration (nidof x 1)
// computed by mj_island (island constraint structure)
@@ -470,11 +378,6 @@ struct mjData_ {
// computed by mj_island (constraints sorted by island)
int* iefc_type; // constraint type (mjtConstraint) (nefc x 1)
int* iefc_id; // id of object of specified type (nefc x 1)
int* iefc_J_rownnz; // number of non-zeros in constraint Jacobian row (nefc x 1)
int* iefc_J_rowadr; // row start address in colind array (nefc x 1)
int* iefc_J_rowsuper; // number of subsequent rows in supernode (nefc x 1)
int* iefc_J_colind; // column indices in constraint Jacobian (nJ x 1)
mjtNum* iefc_J; // constraint Jacobian (nJ x 1)
mjtNum* iefc_frictionloss; // frictionloss (friction) (nefc x 1)
mjtNum* iefc_D; // constraint mass (nefc x 1)
mjtNum* iefc_R; // inverse constraint mass (nefc x 1)
@@ -506,9 +409,6 @@ struct mjData_ {
mjtNum* efc_force; // constraint force in constraint space (nefc x 1)
mjtNum* ifrc_constraint; // constraint force (nidof x 1)
// thread pool pointer
uintptr_t threadpool;
// compilation signature
uint64_t signature; // also held by the mjSpec that compiled the model
};
@@ -533,7 +433,7 @@ typedef mjtNum (*mjfTime)(void);
typedef mjtNum (*mjfAct)(const mjModel* m, const mjData* d, int id);
// collision detection
typedef int (*mjfCollision)(const mjModel* m, mjData* d, mjContact* con, int g1, int g2,
typedef int (*mjfCollision)(const mjModel* m, mjData* d, mjPreContact* con, int g1, int g2,
mjtNum margin);
#endif // MUJOCO_MJDATA_H_
include/mujoco/mjmodel.h
+19 -448
View File
@@ -18,8 +18,7 @@
#include <stddef.h>
#include <stdint.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
// global constants
#define mjPI 3.14159265358979323846
@@ -48,434 +47,6 @@
#define mjNISLAND 20 // number of mjData.solver arrays
//---------------------------------- enum types (mjt) ----------------------------------------------
typedef enum mjtDisableBit_ { // disable default feature bitflags
mjDSBL_CONSTRAINT = 1<<0, // entire constraint solver
mjDSBL_EQUALITY = 1<<1, // equality constraints
mjDSBL_FRICTIONLOSS = 1<<2, // joint and tendon frictionloss constraints
mjDSBL_LIMIT = 1<<3, // joint and tendon limit constraints
mjDSBL_CONTACT = 1<<4, // contact constraints
mjDSBL_SPRING = 1<<5, // passive spring forces
mjDSBL_DAMPER = 1<<6, // passive damping forces
mjDSBL_GRAVITY = 1<<7, // gravitational forces
mjDSBL_CLAMPCTRL = 1<<8, // clamp control to specified range
mjDSBL_WARMSTART = 1<<9, // warmstart constraint solver
mjDSBL_FILTERPARENT = 1<<10, // remove collisions with parent body
mjDSBL_ACTUATION = 1<<11, // apply actuation forces
mjDSBL_REFSAFE = 1<<12, // integrator safety: make ref[0]>=2*timestep
mjDSBL_SENSOR = 1<<13, // sensors
mjDSBL_MIDPHASE = 1<<14, // mid-phase collision filtering
mjDSBL_EULERDAMP = 1<<15, // implicit integration of joint damping in Euler integrator
mjDSBL_AUTORESET = 1<<16, // automatic reset when numerical issues are detected
mjDSBL_NATIVECCD = 1<<17, // native convex collision detection
mjDSBL_ISLAND = 1<<18, // constraint island discovery
mjDSBL_MULTICCD = 1<<19, // multiple CCD contact points
mjNDISABLE = 20 // number of disable flags
} mjtDisableBit;
typedef enum mjtEnableBit_ { // enable optional feature bitflags
mjENBL_OVERRIDE = 1<<0, // override contact parameters
mjENBL_ENERGY = 1<<1, // energy computation
mjENBL_FWDINV = 1<<2, // record solver statistics
mjENBL_INVDISCRETE = 1<<3, // discrete-time inverse dynamics
mjENBL_SLEEP = 1<<4, // sleeping
mjENBL_DIAGEXACT = 1<<5, // exact diagonal of constraint inertia
mjNENABLE = 6 // number of enable flags
} mjtEnableBit;
typedef enum mjtJoint_ { // type of degree of freedom
mjJNT_FREE = 0, // global position and orientation (quat) (7)
mjJNT_BALL, // orientation (quat) relative to parent (4)
mjJNT_SLIDE, // sliding distance along body-fixed axis (1)
mjJNT_HINGE // rotation angle (rad) around body-fixed axis (1)
} mjtJoint;
typedef enum mjtGeom_ { // type of geometric shape
// regular geom types
mjGEOM_PLANE = 0, // plane
mjGEOM_HFIELD, // height field
mjGEOM_SPHERE, // sphere
mjGEOM_CAPSULE, // capsule
mjGEOM_ELLIPSOID, // ellipsoid
mjGEOM_CYLINDER, // cylinder
mjGEOM_BOX, // box
mjGEOM_MESH, // mesh
mjGEOM_SDF, // signed distance field
mjNGEOMTYPES, // number of regular geom types
// rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
mjGEOM_ARROW = 100, // arrow
mjGEOM_ARROW1, // arrow without wedges
mjGEOM_ARROW2, // arrow in both directions
mjGEOM_LINE, // line
mjGEOM_LINEBOX, // box with line edges
mjGEOM_FLEX, // flex
mjGEOM_SKIN, // skin
mjGEOM_LABEL, // text label
mjGEOM_TRIANGLE, // triangle
mjGEOM_NONE = 1001 // missing geom type
} mjtGeom;
typedef enum mjtProjection_ { // type of camera projection
mjPROJ_PERSPECTIVE = 0, // perspective
mjPROJ_ORTHOGRAPHIC // orthographic
} mjtProjection;
typedef enum mjtCamLight_ { // tracking mode for camera and light
mjCAMLIGHT_FIXED = 0, // pos and rot fixed in body
mjCAMLIGHT_TRACK, // pos tracks body, rot fixed in global
mjCAMLIGHT_TRACKCOM, // pos tracks subtree com, rot fixed in body
mjCAMLIGHT_TARGETBODY, // pos fixed in body, rot tracks target body
mjCAMLIGHT_TARGETBODYCOM // pos fixed in body, rot tracks target subtree com
} mjtCamLight;
typedef enum mjtLightType_ { // type of light
mjLIGHT_SPOT = 0, // spot
mjLIGHT_DIRECTIONAL, // directional
mjLIGHT_POINT, // point
mjLIGHT_IMAGE, // image-based
} mjtLightType;
typedef enum mjtTexture_ { // type of texture
mjTEXTURE_2D = 0, // 2d texture, suitable for planes and hfields
mjTEXTURE_CUBE, // cube texture, suitable for all other geom types
mjTEXTURE_SKYBOX // cube texture used as skybox
} mjtTexture;
typedef enum mjtTextureRole_ { // role of texture map in rendering
mjTEXROLE_USER = 0, // unspecified
mjTEXROLE_RGB, // base color (albedo)
mjTEXROLE_OCCLUSION, // ambient occlusion
mjTEXROLE_ROUGHNESS, // roughness
mjTEXROLE_METALLIC, // metallic
mjTEXROLE_NORMAL, // normal (bump) map
mjTEXROLE_OPACITY, // opacity
mjTEXROLE_EMISSIVE, // light emission
mjTEXROLE_RGBA, // base color, opacity
mjTEXROLE_ORM, // occlusion, roughness, metallic
mjNTEXROLE
} mjtTextureRole;
typedef enum mjtColorSpace_ { // type of color space encoding
mjCOLORSPACE_AUTO = 0, // attempts to autodetect color space, defaults to linear
mjCOLORSPACE_LINEAR, // linear color space
mjCOLORSPACE_SRGB // standard RGB color space
} mjtColorSpace;
typedef enum mjtIntegrator_ { // integrator mode
mjINT_EULER = 0, // semi-implicit Euler
mjINT_RK4, // 4th-order Runge Kutta
mjINT_IMPLICIT, // implicit in velocity
mjINT_IMPLICITFAST // implicit in velocity, no rne derivative
} mjtIntegrator;
typedef enum mjtCone_ { // type of friction cone
mjCONE_PYRAMIDAL = 0, // pyramidal
mjCONE_ELLIPTIC // elliptic
} mjtCone;
typedef enum mjtJacobian_ { // type of constraint Jacobian
mjJAC_DENSE = 0, // dense
mjJAC_SPARSE, // sparse
mjJAC_AUTO // dense if nv<60, sparse otherwise
} mjtJacobian;
typedef enum mjtSolver_ { // constraint solver algorithm
mjSOL_PGS = 0, // PGS (dual)
mjSOL_CG, // CG (primal)
mjSOL_NEWTON // Newton (primal)
} mjtSolver;
typedef enum mjtEq_ { // type of equality constraint
mjEQ_CONNECT = 0, // connect two bodies at a point (ball joint)
mjEQ_WELD, // fix relative position and orientation of two bodies
mjEQ_JOINT, // couple the values of two scalar joints with cubic
mjEQ_TENDON, // couple the lengths of two tendons with cubic
mjEQ_FLEX, // fix all edge lengths of a flex
mjEQ_FLEXVERT, // fix all vertex lengths of a flex
mjEQ_FLEXSTRAIN, // constrain strain of a trilinear/quadratic flex (B-bar)
mjEQ_DISTANCE // unsupported, will cause an error if used
} mjtEq;
typedef enum mjtWrap_ { // type of tendon wrap object
mjWRAP_NONE = 0, // null object
mjWRAP_JOINT, // constant moment arm
mjWRAP_PULLEY, // pulley used to split tendon
mjWRAP_SITE, // pass through site
mjWRAP_SPHERE, // wrap around sphere
mjWRAP_CYLINDER // wrap around (infinite) cylinder
} mjtWrap;
typedef enum mjtTrn_ { // type of actuator transmission
mjTRN_JOINT = 0, // force on joint
mjTRN_JOINTINPARENT, // force on joint, expressed in parent frame
mjTRN_SLIDERCRANK, // force via slider-crank linkage
mjTRN_TENDON, // force on tendon
mjTRN_SITE, // force on site
mjTRN_BODY, // adhesion force on a body's geoms
mjTRN_UNDEFINED = 1000 // undefined transmission type
} mjtTrn;
typedef enum mjtDyn_ { // type of actuator dynamics
mjDYN_NONE = 0, // no internal dynamics; ctrl specifies force
mjDYN_INTEGRATOR, // integrator: da/dt = u
mjDYN_FILTER, // linear filter: da/dt = (u-a) / tau
mjDYN_FILTEREXACT, // linear filter: da/dt = (u-a) / tau, with exact integration
mjDYN_MUSCLE, // piecewise linear filter with two time constants
mjDYN_DCMOTOR, // DC motor electrical dynamics
mjDYN_USER // user-defined dynamics type
} mjtDyn;
typedef enum mjtGain_ { // type of actuator gain
mjGAIN_FIXED = 0, // fixed gain
mjGAIN_AFFINE, // const + kp*length + kv*velocity
mjGAIN_MUSCLE, // muscle FLV curve computed by mju_muscleGain()
mjGAIN_DCMOTOR, // DC motor gain: K or K/R
mjGAIN_USER // user-defined gain type
} mjtGain;
typedef enum mjtBias_ { // type of actuator bias
mjBIAS_NONE = 0, // no bias
mjBIAS_AFFINE, // const + kp*length + kv*velocity
mjBIAS_MUSCLE, // muscle passive force computed by mju_muscleBias()
mjBIAS_DCMOTOR, // DC motor bias: back-EMF, cogging, LuGre friction
mjBIAS_USER // user-defined bias type
} mjtBias;
typedef enum mjtObj_ { // type of MujoCo object
mjOBJ_UNKNOWN = 0, // unknown object type
mjOBJ_BODY, // body
mjOBJ_XBODY, // body, used to access regular frame instead of i-frame
mjOBJ_JOINT, // joint
mjOBJ_DOF, // dof
mjOBJ_GEOM, // geom
mjOBJ_SITE, // site
mjOBJ_CAMERA, // camera
mjOBJ_LIGHT, // light
mjOBJ_FLEX, // flex
mjOBJ_MESH, // mesh
mjOBJ_SKIN, // skin
mjOBJ_HFIELD, // heightfield
mjOBJ_TEXTURE, // texture
mjOBJ_MATERIAL, // material for rendering
mjOBJ_PAIR, // geom pair to include
mjOBJ_EXCLUDE, // body pair to exclude
mjOBJ_EQUALITY, // equality constraint
mjOBJ_TENDON, // tendon
mjOBJ_ACTUATOR, // actuator
mjOBJ_SENSOR, // sensor
mjOBJ_NUMERIC, // numeric
mjOBJ_TEXT, // text
mjOBJ_TUPLE, // tuple
mjOBJ_KEY, // keyframe
mjOBJ_PLUGIN, // plugin instance
mjNOBJECT, // number of object types
// meta elements, do not appear in mjModel
mjOBJ_FRAME = 100, // frame
mjOBJ_DEFAULT, // default
mjOBJ_MODEL // entire model
} mjtObj;
typedef enum mjtSensor_ { // type of sensor
// common robotic sensors, attached to a site
mjSENS_TOUCH = 0, // scalar contact normal forces summed over sensor zone
mjSENS_ACCELEROMETER, // 3D linear acceleration, in local frame
mjSENS_VELOCIMETER, // 3D linear velocity, in local frame
mjSENS_GYRO, // 3D angular velocity, in local frame
mjSENS_FORCE, // 3D force between site's body and its parent body
mjSENS_TORQUE, // 3D torque between site's body and its parent body
mjSENS_MAGNETOMETER, // 3D magnetometer
mjSENS_RANGEFINDER, // scalar distance to nearest geom along z-axis
mjSENS_CAMPROJECTION, // pixel coordinates of a site in the camera image
// sensors related to scalar joints, tendons, actuators
mjSENS_JOINTPOS, // scalar joint position (hinge and slide only)
mjSENS_JOINTVEL, // scalar joint velocity (hinge and slide only)
mjSENS_TENDONPOS, // scalar tendon position
mjSENS_TENDONVEL, // scalar tendon velocity
mjSENS_ACTUATORPOS, // scalar actuator position
mjSENS_ACTUATORVEL, // scalar actuator velocity
mjSENS_ACTUATORFRC, // scalar actuator force
mjSENS_JOINTACTFRC, // scalar actuator force, measured at the joint
mjSENS_TENDONACTFRC, // scalar actuator force, measured at the tendon
// sensors related to ball joints
mjSENS_BALLQUAT, // 4D ball joint quaternion
mjSENS_BALLANGVEL, // 3D ball joint angular velocity
// joint and tendon limit sensors, in constraint space
mjSENS_JOINTLIMITPOS, // joint limit distance-margin
mjSENS_JOINTLIMITVEL, // joint limit velocity
mjSENS_JOINTLIMITFRC, // joint limit force
mjSENS_TENDONLIMITPOS, // tendon limit distance-margin
mjSENS_TENDONLIMITVEL, // tendon limit velocity
mjSENS_TENDONLIMITFRC, // tendon limit force
// sensors attached to an object with spatial frame: (x)body, geom, site, camera
mjSENS_FRAMEPOS, // 3D position
mjSENS_FRAMEQUAT, // 4D unit quaternion orientation
mjSENS_FRAMEXAXIS, // 3D unit vector: x-axis of object's frame
mjSENS_FRAMEYAXIS, // 3D unit vector: y-axis of object's frame
mjSENS_FRAMEZAXIS, // 3D unit vector: z-axis of object's frame
mjSENS_FRAMELINVEL, // 3D linear velocity
mjSENS_FRAMEANGVEL, // 3D angular velocity
mjSENS_FRAMELINACC, // 3D linear acceleration
mjSENS_FRAMEANGACC, // 3D angular acceleration
// sensors related to kinematic subtrees; attached to a body (which is the subtree root)
mjSENS_SUBTREECOM, // 3D center of mass of subtree
mjSENS_SUBTREELINVEL, // 3D linear velocity of subtree
mjSENS_SUBTREEANGMOM, // 3D angular momentum of subtree
// sensors of geometric relationships
mjSENS_INSIDESITE, // 1 if object is inside a site, 0 otherwise
mjSENS_GEOMDIST, // signed distance between two geoms
mjSENS_GEOMNORMAL, // normal direction between two geoms
mjSENS_GEOMFROMTO, // segment between two geoms
// sensors for reporting contacts which occurred during the simulation
mjSENS_CONTACT, // contacts which occurred during the simulation
// global sensors
mjSENS_E_POTENTIAL, // potential energy
mjSENS_E_KINETIC, // kinetic energy
mjSENS_CLOCK, // simulation time
// sensors related to SDFs
mjSENS_TACTILE, // tactile sensor
// plugin-controlled sensors
mjSENS_PLUGIN, // plugin-controlled
// user-defined sensor
mjSENS_USER // sensor data provided by mjcb_sensor callback
} mjtSensor;
typedef enum mjtStage_ { // computation stage
mjSTAGE_NONE = 0, // no computations
mjSTAGE_POS, // position-dependent computations
mjSTAGE_VEL, // velocity-dependent computations
mjSTAGE_ACC // acceleration/force-dependent computations
} mjtStage;
typedef enum mjtDataType_ { // data type for sensors
mjDATATYPE_REAL = 0, // real values, no constraints
mjDATATYPE_POSITIVE, // positive values; 0 or negative: inactive
mjDATATYPE_AXIS, // 3D unit vector
mjDATATYPE_QUATERNION // unit quaternion
} mjtDataType;
typedef enum mjtConDataField_ { // data fields returned by contact sensors
mjCONDATA_FOUND = 0, // whether a contact was found
mjCONDATA_FORCE, // contact force
mjCONDATA_TORQUE, // contact torque
mjCONDATA_DIST, // contact penetration distance
mjCONDATA_POS, // contact position
mjCONDATA_NORMAL, // contact frame normal
mjCONDATA_TANGENT, // contact frame first tangent
mjNCONDATA // number of contact sensor data fields
} mjtConDataField;
typedef enum mjtRayDataField_ { // data fields returned by rangefinder sensors
mjRAYDATA_DIST = 0, // distance from ray origin to nearest surface
mjRAYDATA_DIR, // normalized ray direction
mjRAYDATA_ORIGIN, // ray origin
mjRAYDATA_POINT, // point at which ray intersects nearest surface
mjRAYDATA_NORMAL, // surface normal at intersection point
mjRAYDATA_DEPTH, // depth along z-axis
mjNRAYDATA // number of rangefinder sensor data fields
} mjtRayDataField;
typedef enum mjtCamOutBit_ { // camera output type bitflags
mjCAMOUT_RGB = 1<<0, // RGB image
mjCAMOUT_DEPTH = 1<<1, // depth image (distance from camera plane)
mjCAMOUT_DIST = 1<<2, // distance image (distance from camera origin)
mjCAMOUT_NORMAL = 1<<3, // normal image
mjCAMOUT_SEG = 1<<4, // segmentation image
mjNCAMOUT = 5 // number of camera output types
} mjtCamOutBit;
typedef enum mjtSameFrame_ { // frame alignment of bodies with their children
mjSAMEFRAME_NONE = 0, // no alignment
mjSAMEFRAME_BODY, // frame is same as body frame
mjSAMEFRAME_INERTIA, // frame is same as inertial frame
mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation
mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation
} mjtSameFrame;
typedef enum mjtSleepPolicy_ { // per-tree sleep policy
mjSLEEP_AUTO = 0, // compiler chooses sleep policy
mjSLEEP_AUTO_NEVER, // compiler sleep policy: never
mjSLEEP_AUTO_ALLOWED, // compiler sleep policy: allowed
mjSLEEP_NEVER, // user sleep policy: never
mjSLEEP_ALLOWED, // user sleep policy: allowed
mjSLEEP_INIT, // user sleep policy: initialized asleep
} mjtSleepPolicy;
typedef enum mjtLRMode_ { // mode for actuator length range computation
mjLRMODE_NONE = 0, // do not process any actuators
mjLRMODE_MUSCLE, // process muscle actuators
mjLRMODE_MUSCLEUSER, // process muscle and user actuators
mjLRMODE_ALL // process all actuators
} mjtLRMode;
typedef enum mjtFlexSelf_ { // mode for flex selfcollide
mjFLEXSELF_NONE = 0, // no self-collisions
mjFLEXSELF_NARROW, // skip midphase, go directly to narrowphase
mjFLEXSELF_BVH, // use BVH in midphase (if midphase enabled)
mjFLEXSELF_SAP, // use SAP in midphase
mjFLEXSELF_AUTO // choose between BVH and SAP automatically
} mjtFlexSelf;
typedef enum mjtSDFType_ { // signed distance function (SDF) type
mjSDFTYPE_SINGLE = 0, // single SDF
mjSDFTYPE_INTERSECTION, // max(A, B)
mjSDFTYPE_MIDSURFACE, // A - B
mjSDFTYPE_COLLISION, // A + B + abs(max(A, B))
} mjtSDFType;
//---------------------------------- mjLROpt -------------------------------------------------------
@@ -840,9 +411,9 @@ struct mjModel_ {
int* jnt_bodyid; // id of joint's body (njnt x 1)
int* jnt_actuatorid; // actuator contributing damping / armature (njnt x 1)
int* jnt_group; // group for visibility (njnt x 1)
mjtByte* jnt_limited; // does joint have limits (njnt x 1)
mjtByte* jnt_actfrclimited; // does joint have actuator force limits (njnt x 1)
mjtByte* jnt_actgravcomp; // is gravcomp force applied via actuators (njnt x 1)
mjtBool* jnt_limited; // does joint have limits (njnt x 1)
mjtBool* jnt_actfrclimited; // does joint have actuator force limits (njnt x 1)
mjtBool* jnt_actgravcomp; // is gravcomp force applied via actuators (njnt x 1)
mjtNum* jnt_solref; // constraint solver reference: limit (njnt x mjNREF)
mjtNum* jnt_solimp; // constraint solver impedance: limit (njnt x mjNIMP)
mjtNum* jnt_pos; // local anchor position (njnt x 3)
@@ -941,11 +512,11 @@ struct mjModel_ {
int* light_targetbodyid; // id of targeted body; -1: none (nlight x 1)
int* light_type; // spot, directional, etc. (mjtLightType) (nlight x 1)
int* light_texid; // texture id for image lights (nlight x 1)
mjtByte* light_castshadow; // does light cast shadows (nlight x 1)
mjtBool* light_castshadow; // does light cast shadows (nlight x 1)
float* light_bulbradius; // light radius for soft shadows (nlight x 1)
float* light_intensity; // intensity, in candela (nlight x 1)
float* light_range; // range of effectiveness (nlight x 1)
mjtByte* light_active; // is light on (nlight x 1)
mjtBool* light_active; // is light on (nlight x 1)
mjtNum* light_pos; // position rel. to body frame (nlight x 3)
mjtNum* light_dir; // direction rel. to body frame (nlight x 3)
mjtNum* light_poscom0; // global position rel. to sub-com in qpos0 (nlight x 3)
@@ -969,7 +540,7 @@ struct mjModel_ {
mjtNum* flex_friction; // friction for (slide, spin, roll) (nflex x 3)
mjtNum* flex_margin; // geometric inflation for contact (nflex x 1)
mjtNum* flex_gap; // additional contact detection buffer (nflex x 1)
mjtByte* flex_internal; // internal flex collision enabled (nflex x 1)
mjtBool* flex_internal; // internal flex collision enabled (nflex x 1)
int* flex_selfcollide; // self collision mode (mjtFlexSelf) (nflex x 1)
int* flex_activelayers; // number of active element layers, 3D only (nflex x 1)
int* flex_passive; // passive collisions enabled (nflex x 1)
@@ -1025,10 +596,10 @@ struct mjModel_ {
mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1)
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
int* flex_edgeequality; // 0:none, 1:edges, 2:vertices, 3:strain (nflex x 1)
mjtByte* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
mjtBool* flex_rigid; // are all vertices in the same body (nflex x 1)
mjtBool* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtBool* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtBool* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
int* flex_bvhnum; // number of bounding volumes (nflex x 1)
int* flexedge_J_rownnz; // number of non-zeros in Jacobian row (nflexedge x 1)
@@ -1119,7 +690,7 @@ struct mjModel_ {
// materials
int* mat_texid; // indices of textures; -1: none (nmat x mjNTEXROLE)
mjtByte* mat_texuniform; // make texture cube uniform (nmat x 1)
mjtBool* mat_texuniform; // make texture cube uniform (nmat x 1)
float* mat_texrepeat; // texture repetition for 2d mapping (nmat x 2)
float* mat_emission; // emission (x rgb) (nmat x 1)
float* mat_specular; // specular (x white) (nmat x 1)
@@ -1149,7 +720,7 @@ struct mjModel_ {
int* eq_obj1id; // id of object 1 (neq x 1)
int* eq_obj2id; // id of object 2 (neq x 1)
int* eq_objtype; // type of both objects (mjtObj) (neq x 1)
mjtByte* eq_active0; // initial enable/disable constraint state (neq x 1)
mjtBool* eq_active0; // initial enable/disable constraint state (neq x 1)
mjtNum* eq_solref; // constraint solver reference (neq x mjNREF)
mjtNum* eq_solimp; // constraint solver impedance (neq x mjNIMP)
mjtNum* eq_data; // numeric data for constraint (neq x mjNEQDATA)
@@ -1165,8 +736,8 @@ struct mjModel_ {
int* ten_J_rownnz; // number of non-zeros in Jacobian row (ntendon x 1)
int* ten_J_rowadr; // row start address in colind array (ntendon x 1)
int* ten_J_colind; // column indices in sparse Jacobian (nJten x 1)
mjtByte* tendon_limited; // does tendon have length limits (ntendon x 1)
mjtByte* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1)
mjtBool* tendon_limited; // does tendon have length limits (ntendon x 1)
mjtBool* tendon_actfrclimited; // does tendon have actuator force limits (ntendon x 1)
mjtNum* tendon_width; // width for rendering (ntendon x 1)
mjtNum* tendon_solref_lim; // constraint solver reference: limit (ntendon x mjNREF)
mjtNum* tendon_solimp_lim; // constraint solver impedance: limit (ntendon x mjNIMP)
@@ -1207,13 +778,13 @@ struct mjModel_ {
int* actuator_history; // history buffer: [nsample, interp] (nu x 2)
int* actuator_historyadr; // address in history buffer; -1: none (nu x 1)
mjtNum* actuator_delay; // delay time in seconds; 0: no delay (nu x 1)
mjtByte* actuator_ctrllimited; // is control limited (nu x 1)
mjtByte* actuator_forcelimited;// is force limited (nu x 1)
mjtByte* actuator_actlimited; // is activation limited (nu x 1)
mjtBool* actuator_ctrllimited; // is control limited (nu x 1)
mjtBool* actuator_forcelimited;// is force limited (nu x 1)
mjtBool* actuator_actlimited; // is activation limited (nu x 1)
mjtNum* actuator_dynprm; // dynamics parameters (nu x mjNDYN)
mjtNum* actuator_gainprm; // gain parameters (nu x mjNGAIN)
mjtNum* actuator_biasprm; // bias parameters (nu x mjNBIAS)
mjtByte* actuator_actearly; // step activation before force (nu x 1)
mjtBool* actuator_actearly; // step activation before force (nu x 1)
mjtNum* actuator_ctrlrange; // range of controls (nu x 2)
mjtNum* actuator_forcerange; // range of forces (nu x 2)
mjtNum* actuator_actrange; // range of activations (nu x 2)
include/mujoco/mjplugin.h
+1 -1
View File
@@ -18,7 +18,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjspec.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
include/mujoco/mjspec.h
+1 -1
View File
@@ -17,7 +17,7 @@
#include <stddef.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
// this is a C-API
include/mujoco/mjspecmacro.h
+603
View File
@@ -0,0 +1,603 @@
// Copyright 2026 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_MJSPECMACRO_H_
#define MUJOCO_MJSPECMACRO_H_
//-------------------------------- mjsElement ------------------------------------------------------
#define MJSELEMENT_FIELDS \
X( mjtObj, elemtype, 1 ) \
X( uint64_t, signature, 1 )
//-------------------------------- mjsCompiler -----------------------------------------------------
#define MJSCOMPILER_FIELDS \
X ( mjtByte, autolimits, 1 ) \
X ( double, boundmass, 1 ) \
X ( double, boundinertia, 1 ) \
X ( double, settotalmass, 1 ) \
X ( mjtByte, balanceinertia, 1 ) \
X ( mjtByte, fitaabb, 1 ) \
X ( mjtByte, degree, 1 ) \
XVEC( char, eulerseq, 3 ) \
X ( mjtByte, discardvisual, 1 ) \
X ( mjtByte, usethread, 1 ) \
X ( mjtByte, fusestatic, 1 ) \
X ( int, inertiafromgeom, 1 ) \
XVEC( int, inertiagrouprange, 2 ) \
X ( mjtByte, saveinertial, 1 ) \
X ( int, alignfree, 1 ) \
X ( mjLROpt, LRopt, 1 ) \
X ( mjString*, meshdir, 1 ) \
X ( mjString*, texturedir, 1 )
//-------------------------------- mjSpec ----------------------------------------------------------
#define MJSPEC_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjString*, modelname, 1 ) \
X( mjsCompiler, compiler, 1 ) \
X( mjtByte, strippath, 1 ) \
X( mjOption, option, 1 ) \
X( mjVisual, visual, 1 ) \
X( mjStatistic, stat, 1 ) \
X( mjtSize, memory, 1 ) \
X( int, nemax, 1 ) \
X( int, nuserdata, 1 ) \
X( int, nuser_body, 1 ) \
X( int, nuser_jnt, 1 ) \
X( int, nuser_geom, 1 ) \
X( int, nuser_site, 1 ) \
X( int, nuser_cam, 1 ) \
X( int, nuser_tendon, 1 ) \
X( int, nuser_actuator, 1 ) \
X( int, nuser_sensor, 1 ) \
X( int, nkey, 1 ) \
X( int, njmax, 1 ) \
X( int, nconmax, 1 ) \
X( mjtSize, nstack, 1 ) \
X( mjString*, comment, 1 ) \
X( mjString*, modelfiledir, 1 ) \
X( mjtByte, hasImplicitPluginElem, 1 )
//-------------------------------- mjsOrientation --------------------------------------------------
#define MJSORIENTATION_FIELDS \
X ( mjtOrientation, type, 1 ) \
XVEC( double, axisangle, 4 ) \
XVEC( double, xyaxes, 6 ) \
XVEC( double, zaxis, 3 ) \
XVEC( double, euler, 3 )
//-------------------------------- mjsPlugin -------------------------------------------------------
#define MJSPLUGIN_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjString*, name, 1 ) \
X( mjString*, plugin_name, 1 ) \
X( mjtByte, active, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsBody ---------------------------------------------------------
#define MJSBODY_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, childclass, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, quat, 4 ) \
X ( mjsOrientation, alt, 1 ) \
X ( double, mass, 1 ) \
XVEC( double, ipos, 3 ) \
XVEC( double, iquat, 4 ) \
XVEC( double, inertia, 3 ) \
X ( mjsOrientation, ialt, 1 ) \
XVEC( double, fullinertia, 6 ) \
X ( mjtByte, mocap, 1 ) \
X ( double, gravcomp, 1 ) \
X ( mjtSleepPolicy, sleep, 1 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjtByte, explicitinertial, 1 ) \
X ( mjsPlugin, plugin, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsFrame --------------------------------------------------------
#define MJSFRAME_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, childclass, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, quat, 4 ) \
X ( mjsOrientation, alt, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsJoint --------------------------------------------------------
#define MJSJOINT_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtJoint, type, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, axis, 3 ) \
X ( double, ref, 1 ) \
X ( int, align, 1 ) \
XVEC( double, stiffness, mjNPOLY+1 ) \
X ( double, springref, 1 ) \
XVEC( double, springdamper, 2 ) \
X ( int, limited, 1 ) \
XVEC( double, range, 2 ) \
X ( double, margin, 1 ) \
XVEC( mjtNum, solref_limit, mjNREF ) \
XVEC( mjtNum, solimp_limit, mjNIMP ) \
X ( int, actfrclimited, 1 ) \
XVEC( double, actfrcrange, 2 ) \
X ( double, armature, 1 ) \
XVEC( double, damping, mjNPOLY+1 ) \
X ( double, frictionloss, 1 ) \
XVEC( mjtNum, solref_friction, mjNREF ) \
XVEC( mjtNum, solimp_friction, mjNIMP ) \
X ( int, group, 1 ) \
X ( mjtByte, actgravcomp, 1 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsGeom ---------------------------------------------------------
#define MJSGEOM_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtGeom, type, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, quat, 4 ) \
X ( mjsOrientation, alt, 1 ) \
XVEC( double, fromto, 6 ) \
XVEC( double, size, 3 ) \
X ( int, contype, 1 ) \
X ( int, conaffinity, 1 ) \
X ( int, condim, 1 ) \
X ( int, priority, 1 ) \
XVEC( double, friction, 3 ) \
X ( double, solmix, 1 ) \
XVEC( mjtNum, solref, mjNREF ) \
XVEC( mjtNum, solimp, mjNIMP ) \
X ( double, margin, 1 ) \
X ( double, gap, 1 ) \
X ( double, mass, 1 ) \
X ( double, density, 1 ) \
X ( mjtGeomInertia, typeinertia, 1 ) \
X ( mjtNum, fluid_ellipsoid, 1 ) \
XVEC( mjtNum, fluid_coefs, 5 ) \
X ( mjString*, material, 1 ) \
XVEC( float, rgba, 4 ) \
X ( int, group, 1 ) \
X ( mjString*, hfieldname, 1 ) \
X ( mjString*, meshname, 1 ) \
X ( double, fitscale, 1 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjsPlugin, plugin, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsSite ---------------------------------------------------------
#define MJSSITE_FIELDS \
X ( mjsElement*, element, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, quat, 4 ) \
X ( mjsOrientation, alt, 1 ) \
XVEC( double, fromto, 6 ) \
XVEC( double, size, 3 ) \
X ( mjtGeom, type, 1 ) \
X ( mjString*, material, 1 ) \
X ( int, group, 1 ) \
XVEC( float, rgba, 4 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsCamera -------------------------------------------------------
#define MJSCAMERA_FIELDS \
X ( mjsElement*, element, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, quat, 4 ) \
X ( mjsOrientation, alt, 1 ) \
X ( mjtCamLight, mode, 1 ) \
X ( mjString*, targetbody, 1 ) \
X ( mjtProjection, proj, 1 ) \
XVEC( int, resolution, 2 ) \
X ( int, output, 1 ) \
X ( double, fovy, 1 ) \
X ( double, ipd, 1 ) \
XVEC( float, intrinsic, 4 ) \
XVEC( float, sensor_size, 2 ) \
XVEC( float, focal_length, 2 ) \
XVEC( float, focal_pixel, 2 ) \
XVEC( float, principal_length, 2 ) \
XVEC( float, principal_pixel, 2 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsLight --------------------------------------------------------
#define MJSLIGHT_FIELDS \
X ( mjsElement*, element, 1 ) \
XVEC( double, pos, 3 ) \
XVEC( double, dir, 3 ) \
X ( mjtCamLight, mode, 1 ) \
X ( mjString*, targetbody, 1 ) \
X ( mjtByte, active, 1 ) \
X ( mjtLightType, type, 1 ) \
X ( mjString*, texture, 1 ) \
X ( mjtByte, castshadow, 1 ) \
X ( float, bulbradius, 1 ) \
X ( float, intensity, 1 ) \
X ( float, range, 1 ) \
XVEC( float, attenuation, 3 ) \
X ( float, cutoff, 1 ) \
X ( float, exponent, 1 ) \
XVEC( float, ambient, 3 ) \
XVEC( float, diffuse, 3 ) \
XVEC( float, specular, 3 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsFlex ---------------------------------------------------------
#define MJSFLEX_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( int, contype, 1 ) \
X ( int, conaffinity, 1 ) \
X ( int, condim, 1 ) \
X ( int, priority, 1 ) \
XVEC( double, friction, 3 ) \
X ( double, solmix, 1 ) \
XVEC( mjtNum, solref, mjNREF ) \
XVEC( mjtNum, solimp, mjNIMP ) \
X ( double, margin, 1 ) \
X ( double, gap, 1 ) \
X ( int, dim, 1 ) \
X ( double, radius, 1 ) \
XVEC( double, size, 3 ) \
X ( mjtByte, internal, 1 ) \
X ( mjtByte, flatskin, 1 ) \
X ( int, selfcollide, 1 ) \
X ( int, passive, 1 ) \
X ( int, activelayers, 1 ) \
X ( int, group, 1 ) \
X ( double, edgestiffness, 1 ) \
X ( double, edgedamping, 1 ) \
XVEC( float, rgba, 4 ) \
X ( mjString*, material, 1 ) \
X ( double, young, 1 ) \
X ( double, poisson, 1 ) \
X ( double, damping, 1 ) \
X ( double, thickness, 1 ) \
X ( int, elastic2d, 1 ) \
XVEC( int, cellcount, 3 ) \
X ( int, order, 1 ) \
X ( mjStringVec*, nodebody, 1 ) \
X ( mjStringVec*, vertbody, 1 ) \
X ( mjDoubleVec*, node, 1 ) \
X ( mjDoubleVec*, vert, 1 ) \
X ( mjIntVec*, elem, 1 ) \
X ( mjFloatVec*, texcoord, 1 ) \
X ( mjIntVec*, elemtexcoord, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsMesh ---------------------------------------------------------
#define MJSMESH_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, content_type, 1 ) \
X ( mjString*, file, 1 ) \
XVEC( double, refpos, 3 ) \
XVEC( double, refquat, 4 ) \
XVEC( double, scale, 3 ) \
X ( mjtMeshInertia, inertia, 1 ) \
X ( mjtByte, smoothnormal, 1 ) \
X ( mjtByte, needsdf, 1 ) \
X ( int, maxhullvert, 1 ) \
X ( mjFloatVec*, uservert, 1 ) \
X ( mjFloatVec*, usernormal, 1 ) \
X ( mjFloatVec*, usertexcoord, 1 ) \
X ( mjIntVec*, userface, 1 ) \
X ( mjIntVec*, userfacenormal, 1 ) \
X ( mjIntVec*, userfacetexcoord, 1 ) \
X ( mjsPlugin, plugin, 1 ) \
X ( mjString*, material, 1 ) \
X ( int, octree_maxdepth, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsHField -------------------------------------------------------
#define MJSHFIELD_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, content_type, 1 ) \
X ( mjString*, file, 1 ) \
XVEC( double, size, 4 ) \
X ( int, nrow, 1 ) \
X ( int, ncol, 1 ) \
X ( mjFloatVec*, userdata, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsSkin ---------------------------------------------------------
#define MJSSKIN_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, file, 1 ) \
X ( mjString*, material, 1 ) \
XVEC( float, rgba, 4 ) \
X ( float, inflate, 1 ) \
X ( int, group, 1 ) \
X ( mjFloatVec*, vert, 1 ) \
X ( mjFloatVec*, texcoord, 1 ) \
X ( mjIntVec*, face, 1 ) \
X ( mjStringVec*, bodyname, 1 ) \
X ( mjFloatVec*, bindpos, 1 ) \
X ( mjFloatVec*, bindquat, 1 ) \
X ( mjIntVecVec*, vertid, 1 ) \
X ( mjFloatVecVec*, vertweight, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsTexture ------------------------------------------------------
#define MJSTEXTURE_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtTexture, type, 1 ) \
X ( mjtColorSpace, colorspace, 1 ) \
X ( int, builtin, 1 ) \
X ( int, mark, 1 ) \
XVEC( double, rgb1, 3 ) \
XVEC( double, rgb2, 3 ) \
XVEC( double, markrgb, 3 ) \
X ( double, random, 1 ) \
X ( int, height, 1 ) \
X ( int, width, 1 ) \
X ( int, nchannel, 1 ) \
X ( mjString*, content_type, 1 ) \
X ( mjString*, file, 1 ) \
XVEC( int, gridsize, 2 ) \
XVEC( char, gridlayout, 12 ) \
X ( mjStringVec*, cubefiles, 1 ) \
X ( mjByteVec*, data, 1 ) \
X ( mjtByte, hflip, 1 ) \
X ( mjtByte, vflip, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsMaterial -----------------------------------------------------
#define MJSMATERIAL_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjStringVec*, textures, 1 ) \
X ( mjtByte, texuniform, 1 ) \
XVEC( float, texrepeat, 2 ) \
X ( float, emission, 1 ) \
X ( float, specular, 1 ) \
X ( float, shininess, 1 ) \
X ( float, reflectance, 1 ) \
X ( float, metallic, 1 ) \
X ( float, roughness, 1 ) \
XVEC( float, rgba, 4 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsPair ---------------------------------------------------------
#define MJSPAIR_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjString*, geomname1, 1 ) \
X ( mjString*, geomname2, 1 ) \
X ( int, condim, 1 ) \
XVEC( mjtNum, solref, mjNREF ) \
XVEC( mjtNum, solreffriction, mjNREF ) \
XVEC( mjtNum, solimp, mjNIMP ) \
X ( double, margin, 1 ) \
X ( double, gap, 1 ) \
XVEC( double, friction, 5 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsExclude ------------------------------------------------------
#define MJSEXCLUDE_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjString*, bodyname1, 1 ) \
X( mjString*, bodyname2, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsEquality -----------------------------------------------------
#define MJSEQUALITY_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtEq, type, 1 ) \
XVEC( double, data, mjNEQDATA ) \
X ( mjtByte, active, 1 ) \
X ( mjString*, name1, 1 ) \
X ( mjString*, name2, 1 ) \
X ( mjtObj, objtype, 1 ) \
XVEC( mjtNum, solref, mjNREF ) \
XVEC( mjtNum, solimp, mjNIMP ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsTendon -------------------------------------------------------
#define MJSTENDON_FIELDS \
X ( mjsElement*, element, 1 ) \
XVEC( double, stiffness, mjNPOLY+1 ) \
XVEC( double, springlength, 2 ) \
XVEC( double, damping, mjNPOLY+1 ) \
X ( double, frictionloss, 1 ) \
XVEC( mjtNum, solref_friction, mjNREF ) \
XVEC( mjtNum, solimp_friction, mjNIMP ) \
X ( double, armature, 1 ) \
X ( int, limited, 1 ) \
X ( int, actfrclimited, 1 ) \
XVEC( double, range, 2 ) \
XVEC( double, actfrcrange, 2 ) \
X ( double, margin, 1 ) \
XVEC( mjtNum, solref_limit, mjNREF ) \
XVEC( mjtNum, solimp_limit, mjNIMP ) \
X ( mjString*, material, 1 ) \
X ( double, width, 1 ) \
XVEC( float, rgba, 4 ) \
X ( int, group, 1 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsWrap ---------------------------------------------------------
#define MJSWRAP_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjtWrap, type, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsActuator -----------------------------------------------------
#define MJSACTUATOR_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtGain, gaintype, 1 ) \
XVEC( double, gainprm, mjNGAIN ) \
X ( mjtBias, biastype, 1 ) \
XVEC( double, biasprm, mjNGAIN ) \
X ( mjtDyn, dyntype, 1 ) \
XVEC( double, dynprm, mjNDYN ) \
X ( int, actdim, 1 ) \
X ( mjtByte, actearly, 1 ) \
X ( mjtTrn, trntype, 1 ) \
XVEC( double, gear, 6 ) \
X ( mjString*, target, 1 ) \
X ( mjString*, refsite, 1 ) \
X ( mjString*, slidersite, 1 ) \
X ( double, cranklength, 1 ) \
XVEC( double, lengthrange, 2 ) \
X ( double, inheritrange, 1 ) \
XVEC( double, damping, mjNPOLY+1 ) \
X ( double, armature, 1 ) \
X ( int, ctrllimited, 1 ) \
XVEC( double, ctrlrange, 2 ) \
X ( int, forcelimited, 1 ) \
XVEC( double, forcerange, 2 ) \
X ( int, actlimited, 1 ) \
XVEC( double, actrange, 2 ) \
X ( int, group, 1 ) \
X ( int, nsample, 1 ) \
X ( int, interp, 1 ) \
X ( double, delay, 1 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjsPlugin, plugin, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsSensor -------------------------------------------------------
#define MJSSENSOR_FIELDS \
X ( mjsElement*, element, 1 ) \
X ( mjtSensor, type, 1 ) \
X ( mjtObj, objtype, 1 ) \
X ( mjString*, objname, 1 ) \
X ( mjtObj, reftype, 1 ) \
X ( mjString*, refname, 1 ) \
XVEC( int, intprm, mjNSENS ) \
X ( mjtDataType, datatype, 1 ) \
X ( mjtStage, needstage, 1 ) \
X ( int, dim, 1 ) \
X ( double, cutoff, 1 ) \
X ( double, noise, 1 ) \
X ( int, nsample, 1 ) \
X ( int, interp, 1 ) \
X ( double, delay, 1 ) \
XVEC( double, interval, 2 ) \
X ( mjDoubleVec*, userdata, 1 ) \
X ( mjsPlugin, plugin, 1 ) \
X ( mjString*, info, 1 )
//-------------------------------- mjsNumeric ------------------------------------------------------
#define MJSNUMERIC_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjDoubleVec*, data, 1 ) \
X( int, size, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsText ---------------------------------------------------------
#define MJSTEXT_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjString*, data, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsTuple --------------------------------------------------------
#define MJSTUPLE_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjIntVec*, objtype, 1 ) \
X( mjStringVec*, objname, 1 ) \
X( mjDoubleVec*, objprm, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsKey ----------------------------------------------------------
#define MJSKEY_FIELDS \
X( mjsElement*, element, 1 ) \
X( double, time, 1 ) \
X( mjDoubleVec*, qpos, 1 ) \
X( mjDoubleVec*, qvel, 1 ) \
X( mjDoubleVec*, act, 1 ) \
X( mjDoubleVec*, mpos, 1 ) \
X( mjDoubleVec*, mquat, 1 ) \
X( mjDoubleVec*, ctrl, 1 ) \
X( mjString*, info, 1 )
//-------------------------------- mjsDefault ------------------------------------------------------
#define MJSDEFAULT_FIELDS \
X( mjsElement*, element, 1 ) \
X( mjsJoint*, joint, 1 ) \
X( mjsGeom*, geom, 1 ) \
X( mjsSite*, site, 1 ) \
X( mjsCamera*, camera, 1 ) \
X( mjsLight*, light, 1 ) \
X( mjsFlex*, flex, 1 ) \
X( mjsMesh*, mesh, 1 ) \
X( mjsMaterial*, material, 1 ) \
X( mjsPair*, pair, 1 ) \
X( mjsEquality*, equality, 1 ) \
X( mjsTendon*, tendon, 1 ) \
X( mjsActuator*, actuator, 1 )
#endif // MUJOCO_MJSPECMACRO_H_
include/mujoco/mjthread.h
-42
View File
@@ -1,42 +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.
#ifndef MUJOCO_INCLUDE_MJTHREAD_H_
#define MUJOCO_INCLUDE_MJTHREAD_H_
#define mjMAXTHREAD 128 // maximum number of threads in a thread pool
typedef enum mjtTaskStatus_ { // status values for mjTask
mjTASK_NEW = 0, // newly created
mjTASK_QUEUED, // enqueued in a thread pool
mjTASK_COMPLETED // completed execution
} mjtTaskStatus;
// function pointer type for mjTask
typedef void* (*mjfTask)(void*);
// An opaque type representing a thread pool.
struct mjThreadPool_ {
int nworker; // number of workers in the pool
};
typedef struct mjThreadPool_ mjThreadPool;
struct mjTask_ { // a task that can be executed by a thread pool.
mjfTask func; // pointer to the function that implements the task
void* args; // arguments to func
volatile int status; // status of the task
};
typedef struct mjTask_ mjTask;
#endif // MUJOCO_INCLUDE_MJTHREAD_H_
include/mujoco/mjtnum.h
-46
View File
@@ -1,46 +0,0 @@
// Copyright 2021 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_INCLUDE_MJTNUM_H_
#define MUJOCO_INCLUDE_MJTNUM_H_
#include <stdint.h>
//---------------------------------- floating-point definition -------------------------------------
// floating point data type and minval
#ifndef mjUSESINGLE
typedef double mjtNum;
#define mjMINVAL 1E-15 // minimum value in any denominator
#else
typedef float mjtNum;
#define mjMINVAL 1E-15f
#endif
//---------------------------------- byte definition -----------------------------------------------
typedef unsigned char mjtByte; // used for true/false
//---------------------------------- size definition -----------------------------------------------
typedef int64_t mjtSize; // used for buffer sizes
#endif // MUJOCO_INCLUDE_MJTNUM_H_
include/mujoco/mjtype.h
+579
View File
@@ -0,0 +1,579 @@
// Copyright 2021 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_INCLUDE_MJTYPE_H_
#define MUJOCO_INCLUDE_MJTYPE_H_
#include <stdbool.h>
#include <stdint.h>
//---------------------------------- floating-point definition -------------------------------------
// floating point data type and minval
#ifndef mjUSESINGLE
typedef double mjtNum;
#define mjMINVAL 1E-15 // minimum value in any denominator
#else
typedef float mjtNum;
#define mjMINVAL 1E-15f
#endif
//---------------------------------- byte definition -----------------------------------------------
typedef unsigned char mjtByte; // used for small integers and binary data
#ifndef __cplusplus
typedef _Bool mjtBool; // used for boolean values
#else
typedef bool mjtBool; // used for boolean values
#endif
//---------------------------------- size definition -----------------------------------------------
typedef int64_t mjtSize; // used for buffer sizes
//---------------------------------- enum types (mjModel) ------------------------------------------
typedef enum mjtDisableBit_ { // disable default feature bitflags
mjDSBL_CONSTRAINT = 1<<0, // entire constraint solver
mjDSBL_EQUALITY = 1<<1, // equality constraints
mjDSBL_FRICTIONLOSS = 1<<2, // joint and tendon frictionloss constraints
mjDSBL_LIMIT = 1<<3, // joint and tendon limit constraints
mjDSBL_CONTACT = 1<<4, // contact constraints
mjDSBL_SPRING = 1<<5, // passive spring forces
mjDSBL_DAMPER = 1<<6, // passive damping forces
mjDSBL_GRAVITY = 1<<7, // gravitational forces
mjDSBL_CLAMPCTRL = 1<<8, // clamp control to specified range
mjDSBL_WARMSTART = 1<<9, // warmstart constraint solver
mjDSBL_FILTERPARENT = 1<<10, // remove collisions with parent body
mjDSBL_ACTUATION = 1<<11, // apply actuation forces
mjDSBL_REFSAFE = 1<<12, // integrator safety: make ref[0]>=2*timestep
mjDSBL_SENSOR = 1<<13, // sensors
mjDSBL_MIDPHASE = 1<<14, // mid-phase collision filtering
mjDSBL_EULERDAMP = 1<<15, // implicit integration of joint damping in Euler integrator
mjDSBL_AUTORESET = 1<<16, // automatic reset when numerical issues are detected
mjDSBL_NATIVECCD = 1<<17, // native convex collision detection
mjDSBL_ISLAND = 1<<18, // constraint island discovery
mjDSBL_MULTICCD = 1<<19, // multiple CCD contact points
mjNDISABLE = 20 // number of disable flags
} mjtDisableBit;
typedef enum mjtEnableBit_ { // enable optional feature bitflags
mjENBL_OVERRIDE = 1<<0, // override contact parameters
mjENBL_ENERGY = 1<<1, // energy computation
mjENBL_FWDINV = 1<<2, // record solver statistics
mjENBL_INVDISCRETE = 1<<3, // discrete-time inverse dynamics
mjENBL_SLEEP = 1<<4, // sleeping
mjENBL_DIAGEXACT = 1<<5, // exact diagonal of constraint inertia
mjNENABLE = 6 // number of enable flags
} mjtEnableBit;
typedef enum mjtJoint_ { // type of degree of freedom
mjJNT_FREE = 0, // global position and orientation (quat) (7)
mjJNT_BALL, // orientation (quat) relative to parent (4)
mjJNT_SLIDE, // sliding distance along body-fixed axis (1)
mjJNT_HINGE // rotation angle (rad) around body-fixed axis (1)
} mjtJoint;
typedef enum mjtGeom_ { // type of geometric shape
// regular geom types
mjGEOM_PLANE = 0, // plane
mjGEOM_HFIELD, // height field
mjGEOM_SPHERE, // sphere
mjGEOM_CAPSULE, // capsule
mjGEOM_ELLIPSOID, // ellipsoid
mjGEOM_CYLINDER, // cylinder
mjGEOM_BOX, // box
mjGEOM_MESH, // mesh
mjGEOM_SDF, // signed distance field
mjNGEOMTYPES, // number of regular geom types
// rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
mjGEOM_ARROW = 100, // arrow
mjGEOM_ARROW1, // arrow without wedges
mjGEOM_ARROW2, // arrow in both directions
mjGEOM_LINE, // line
mjGEOM_LINEBOX, // box with line edges
mjGEOM_FLEX, // flex
mjGEOM_SKIN, // skin
mjGEOM_LABEL, // text label
mjGEOM_TRIANGLE, // triangle
mjGEOM_NONE = 1001 // missing geom type
} mjtGeom;
typedef enum mjtProjection_ { // type of camera projection
mjPROJ_PERSPECTIVE = 0, // perspective
mjPROJ_ORTHOGRAPHIC // orthographic
} mjtProjection;
typedef enum mjtCamLight_ { // tracking mode for camera and light
mjCAMLIGHT_FIXED = 0, // pos and rot fixed in body
mjCAMLIGHT_TRACK, // pos tracks body, rot fixed in global
mjCAMLIGHT_TRACKCOM, // pos tracks subtree com, rot fixed in body
mjCAMLIGHT_TARGETBODY, // pos fixed in body, rot tracks target body
mjCAMLIGHT_TARGETBODYCOM // pos fixed in body, rot tracks target subtree com
} mjtCamLight;
typedef enum mjtLightType_ { // type of light
mjLIGHT_SPOT = 0, // spot
mjLIGHT_DIRECTIONAL, // directional
mjLIGHT_POINT, // point
mjLIGHT_IMAGE, // image-based
} mjtLightType;
typedef enum mjtTexture_ { // type of texture
mjTEXTURE_2D = 0, // 2d texture, suitable for planes and hfields
mjTEXTURE_CUBE, // cube texture, suitable for all other geom types
mjTEXTURE_SKYBOX // cube texture used as skybox
} mjtTexture;
typedef enum mjtTextureRole_ { // role of texture map in rendering
mjTEXROLE_USER = 0, // unspecified
mjTEXROLE_RGB, // base color (albedo)
mjTEXROLE_OCCLUSION, // ambient occlusion
mjTEXROLE_ROUGHNESS, // roughness
mjTEXROLE_METALLIC, // metallic
mjTEXROLE_NORMAL, // normal (bump) map
mjTEXROLE_OPACITY, // opacity
mjTEXROLE_EMISSIVE, // light emission
mjTEXROLE_RGBA, // base color, opacity
mjTEXROLE_ORM, // occlusion, roughness, metallic
mjNTEXROLE
} mjtTextureRole;
typedef enum mjtColorSpace_ { // type of color space encoding
mjCOLORSPACE_AUTO = 0, // attempts to autodetect color space, defaults to linear
mjCOLORSPACE_LINEAR, // linear color space
mjCOLORSPACE_SRGB // standard RGB color space
} mjtColorSpace;
typedef enum mjtIntegrator_ { // integrator mode
mjINT_EULER = 0, // semi-implicit Euler
mjINT_RK4, // 4th-order Runge Kutta
mjINT_IMPLICIT, // implicit in velocity
mjINT_IMPLICITFAST // implicit in velocity, no rne derivative
} mjtIntegrator;
typedef enum mjtCone_ { // type of friction cone
mjCONE_PYRAMIDAL = 0, // pyramidal
mjCONE_ELLIPTIC // elliptic
} mjtCone;
typedef enum mjtJacobian_ { // type of constraint Jacobian
mjJAC_DENSE = 0, // dense
mjJAC_SPARSE, // sparse
mjJAC_AUTO // dense if nv<60, sparse otherwise
} mjtJacobian;
typedef enum mjtSolver_ { // constraint solver algorithm
mjSOL_PGS = 0, // PGS (dual)
mjSOL_CG, // CG (primal)
mjSOL_NEWTON // Newton (primal)
} mjtSolver;
typedef enum mjtEq_ { // type of equality constraint
mjEQ_CONNECT = 0, // connect two bodies at a point (ball joint)
mjEQ_WELD, // fix relative position and orientation of two bodies
mjEQ_JOINT, // couple the values of two scalar joints with cubic
mjEQ_TENDON, // couple the lengths of two tendons with cubic
mjEQ_FLEX, // fix all edge lengths of a flex
mjEQ_FLEXVERT, // fix all vertex lengths of a flex
mjEQ_FLEXSTRAIN, // constrain strain of a trilinear/quadratic flex (B-bar)
mjEQ_DISTANCE // unsupported, will cause an error if used
} mjtEq;
typedef enum mjtWrap_ { // type of tendon wrap object
mjWRAP_NONE = 0, // null object
mjWRAP_JOINT, // constant moment arm
mjWRAP_PULLEY, // pulley used to split tendon
mjWRAP_SITE, // pass through site
mjWRAP_SPHERE, // wrap around sphere
mjWRAP_CYLINDER // wrap around (infinite) cylinder
} mjtWrap;
typedef enum mjtTrn_ { // type of actuator transmission
mjTRN_JOINT = 0, // force on joint
mjTRN_JOINTINPARENT, // force on joint, expressed in parent frame
mjTRN_SLIDERCRANK, // force via slider-crank linkage
mjTRN_TENDON, // force on tendon
mjTRN_SITE, // force on site
mjTRN_BODY, // adhesion force on a body's geoms
mjTRN_UNDEFINED = 1000 // undefined transmission type
} mjtTrn;
typedef enum mjtDyn_ { // type of actuator dynamics
mjDYN_NONE = 0, // no internal dynamics; ctrl specifies force
mjDYN_INTEGRATOR, // integrator: da/dt = u
mjDYN_FILTER, // linear filter: da/dt = (u-a) / tau
mjDYN_FILTEREXACT, // linear filter: da/dt = (u-a) / tau, with exact integration
mjDYN_MUSCLE, // piecewise linear filter with two time constants
mjDYN_DCMOTOR, // DC motor electrical dynamics
mjDYN_USER // user-defined dynamics type
} mjtDyn;
typedef enum mjtGain_ { // type of actuator gain
mjGAIN_FIXED = 0, // fixed gain
mjGAIN_AFFINE, // const + kp*length + kv*velocity
mjGAIN_MUSCLE, // muscle FLV curve computed by mju_muscleGain()
mjGAIN_DCMOTOR, // DC motor gain: K or K/R
mjGAIN_USER // user-defined gain type
} mjtGain;
typedef enum mjtBias_ { // type of actuator bias
mjBIAS_NONE = 0, // no bias
mjBIAS_AFFINE, // const + kp*length + kv*velocity
mjBIAS_MUSCLE, // muscle passive force computed by mju_muscleBias()
mjBIAS_DCMOTOR, // DC motor bias: back-EMF, cogging, LuGre friction
mjBIAS_USER // user-defined bias type
} mjtBias;
typedef enum mjtObj_ { // type of MujoCo object
mjOBJ_UNKNOWN = 0, // unknown object type
mjOBJ_BODY, // body
mjOBJ_XBODY, // body, used to access regular frame instead of i-frame
mjOBJ_JOINT, // joint
mjOBJ_DOF, // dof
mjOBJ_GEOM, // geom
mjOBJ_SITE, // site
mjOBJ_CAMERA, // camera
mjOBJ_LIGHT, // light
mjOBJ_FLEX, // flex
mjOBJ_MESH, // mesh
mjOBJ_SKIN, // skin
mjOBJ_HFIELD, // heightfield
mjOBJ_TEXTURE, // texture
mjOBJ_MATERIAL, // material for rendering
mjOBJ_PAIR, // geom pair to include
mjOBJ_EXCLUDE, // body pair to exclude
mjOBJ_EQUALITY, // equality constraint
mjOBJ_TENDON, // tendon
mjOBJ_ACTUATOR, // actuator
mjOBJ_SENSOR, // sensor
mjOBJ_NUMERIC, // numeric
mjOBJ_TEXT, // text
mjOBJ_TUPLE, // tuple
mjOBJ_KEY, // keyframe
mjOBJ_PLUGIN, // plugin instance
mjNOBJECT, // number of object types
// meta elements, do not appear in mjModel
mjOBJ_FRAME = 100, // frame
mjOBJ_DEFAULT, // default
mjOBJ_MODEL // entire model
} mjtObj;
typedef enum mjtSensor_ { // type of sensor
// common robotic sensors, attached to a site
mjSENS_TOUCH = 0, // scalar contact normal forces summed over sensor zone
mjSENS_ACCELEROMETER, // 3D linear acceleration, in local frame
mjSENS_VELOCIMETER, // 3D linear velocity, in local frame
mjSENS_GYRO, // 3D angular velocity, in local frame
mjSENS_FORCE, // 3D force between site's body and its parent body
mjSENS_TORQUE, // 3D torque between site's body and its parent body
mjSENS_MAGNETOMETER, // 3D magnetometer
mjSENS_RANGEFINDER, // scalar distance to nearest geom along z-axis
mjSENS_CAMPROJECTION, // pixel coordinates of a site in the camera image
// sensors related to scalar joints, tendons, actuators
mjSENS_JOINTPOS, // scalar joint position (hinge and slide only)
mjSENS_JOINTVEL, // scalar joint velocity (hinge and slide only)
mjSENS_TENDONPOS, // scalar tendon position
mjSENS_TENDONVEL, // scalar tendon velocity
mjSENS_ACTUATORPOS, // scalar actuator position
mjSENS_ACTUATORVEL, // scalar actuator velocity
mjSENS_ACTUATORFRC, // scalar actuator force
mjSENS_JOINTACTFRC, // scalar actuator force, measured at the joint
mjSENS_TENDONACTFRC, // scalar actuator force, measured at the tendon
// sensors related to ball joints
mjSENS_BALLQUAT, // 4D ball joint quaternion
mjSENS_BALLANGVEL, // 3D ball joint angular velocity
// joint and tendon limit sensors, in constraint space
mjSENS_JOINTLIMITPOS, // joint limit distance-margin
mjSENS_JOINTLIMITVEL, // joint limit velocity
mjSENS_JOINTLIMITFRC, // joint limit force
mjSENS_TENDONLIMITPOS, // tendon limit distance-margin
mjSENS_TENDONLIMITVEL, // tendon limit velocity
mjSENS_TENDONLIMITFRC, // tendon limit force
// sensors attached to an object with spatial frame: (x)body, geom, site, camera
mjSENS_FRAMEPOS, // 3D position
mjSENS_FRAMEQUAT, // 4D unit quaternion orientation
mjSENS_FRAMEXAXIS, // 3D unit vector: x-axis of object's frame
mjSENS_FRAMEYAXIS, // 3D unit vector: y-axis of object's frame
mjSENS_FRAMEZAXIS, // 3D unit vector: z-axis of object's frame
mjSENS_FRAMELINVEL, // 3D linear velocity
mjSENS_FRAMEANGVEL, // 3D angular velocity
mjSENS_FRAMELINACC, // 3D linear acceleration
mjSENS_FRAMEANGACC, // 3D angular acceleration
// sensors related to kinematic subtrees; attached to a body (which is the subtree root)
mjSENS_SUBTREECOM, // 3D center of mass of subtree
mjSENS_SUBTREELINVEL, // 3D linear velocity of subtree
mjSENS_SUBTREEANGMOM, // 3D angular momentum of subtree
// sensors of geometric relationships
mjSENS_INSIDESITE, // 1 if object is inside a site, 0 otherwise
mjSENS_GEOMDIST, // signed distance between two geoms
mjSENS_GEOMNORMAL, // normal direction between two geoms
mjSENS_GEOMFROMTO, // segment between two geoms
// sensors for reporting contacts which occurred during the simulation
mjSENS_CONTACT, // contacts which occurred during the simulation
// global sensors
mjSENS_E_POTENTIAL, // potential energy
mjSENS_E_KINETIC, // kinetic energy
mjSENS_CLOCK, // simulation time
// sensors related to SDFs
mjSENS_TACTILE, // tactile sensor
// plugin-controlled sensors
mjSENS_PLUGIN, // plugin-controlled
// user-defined sensor
mjSENS_USER // sensor data provided by mjcb_sensor callback
} mjtSensor;
typedef enum mjtStage_ { // computation stage
mjSTAGE_NONE = 0, // no computations
mjSTAGE_POS, // position-dependent computations
mjSTAGE_VEL, // velocity-dependent computations
mjSTAGE_ACC // acceleration/force-dependent computations
} mjtStage;
typedef enum mjtDataType_ { // data type for sensors
mjDATATYPE_REAL = 0, // real values, no constraints
mjDATATYPE_POSITIVE, // positive values; 0 or negative: inactive
mjDATATYPE_AXIS, // 3D unit vector
mjDATATYPE_QUATERNION // unit quaternion
} mjtDataType;
typedef enum mjtConDataField_ { // data fields returned by contact sensors
mjCONDATA_FOUND = 0, // whether a contact was found
mjCONDATA_FORCE, // contact force
mjCONDATA_TORQUE, // contact torque
mjCONDATA_DIST, // contact penetration distance
mjCONDATA_POS, // contact position
mjCONDATA_NORMAL, // contact frame normal
mjCONDATA_TANGENT, // contact frame first tangent
mjNCONDATA // number of contact sensor data fields
} mjtConDataField;
typedef enum mjtRayDataField_ { // data fields returned by rangefinder sensors
mjRAYDATA_DIST = 0, // distance from ray origin to nearest surface
mjRAYDATA_DIR, // normalized ray direction
mjRAYDATA_ORIGIN, // ray origin
mjRAYDATA_POINT, // point at which ray intersects nearest surface
mjRAYDATA_NORMAL, // surface normal at intersection point
mjRAYDATA_DEPTH, // depth along z-axis
mjNRAYDATA // number of rangefinder sensor data fields
} mjtRayDataField;
typedef enum mjtCamOutBit_ { // camera output type bitflags
mjCAMOUT_RGB = 1<<0, // RGB image
mjCAMOUT_DEPTH = 1<<1, // depth image (distance from camera plane)
mjCAMOUT_DIST = 1<<2, // distance image (distance from camera origin)
mjCAMOUT_NORMAL = 1<<3, // normal image
mjCAMOUT_SEG = 1<<4, // segmentation image
mjNCAMOUT = 5 // number of camera output types
} mjtCamOutBit;
typedef enum mjtSameFrame_ { // frame alignment of bodies with their children
mjSAMEFRAME_NONE = 0, // no alignment
mjSAMEFRAME_BODY, // frame is same as body frame
mjSAMEFRAME_INERTIA, // frame is same as inertial frame
mjSAMEFRAME_BODYROT, // frame orientation is same as body orientation
mjSAMEFRAME_INERTIAROT // frame orientation is same as inertia orientation
} mjtSameFrame;
typedef enum mjtSleepPolicy_ { // per-tree sleep policy
mjSLEEP_AUTO = 0, // compiler chooses sleep policy
mjSLEEP_AUTO_NEVER, // compiler sleep policy: never
mjSLEEP_AUTO_ALLOWED, // compiler sleep policy: allowed
mjSLEEP_NEVER, // user sleep policy: never
mjSLEEP_ALLOWED, // user sleep policy: allowed
mjSLEEP_INIT, // user sleep policy: initialized asleep
} mjtSleepPolicy;
typedef enum mjtLRMode_ { // mode for actuator length range computation
mjLRMODE_NONE = 0, // do not process any actuators
mjLRMODE_MUSCLE, // process muscle actuators
mjLRMODE_MUSCLEUSER, // process muscle and user actuators
mjLRMODE_ALL // process all actuators
} mjtLRMode;
typedef enum mjtFlexSelf_ { // mode for flex selfcollide
mjFLEXSELF_NONE = 0, // no self-collisions
mjFLEXSELF_NARROW, // skip midphase, go directly to narrowphase
mjFLEXSELF_BVH, // use BVH in midphase (if midphase enabled)
mjFLEXSELF_SAP, // use SAP in midphase
mjFLEXSELF_AUTO // choose between BVH and SAP automatically
} mjtFlexSelf;
typedef enum mjtSDFType_ { // signed distance function (SDF) type
mjSDFTYPE_SINGLE = 0, // single SDF
mjSDFTYPE_INTERSECTION, // max(A, B)
mjSDFTYPE_MIDSURFACE, // A - B
mjSDFTYPE_COLLISION, // A + B + abs(max(A, B))
} mjtSDFType;
//---------------------------------- enum types (mjData) -------------------------------------------
typedef enum mjtState_ { // state elements
mjSTATE_TIME = 1<<0, // time
mjSTATE_QPOS = 1<<1, // position
mjSTATE_QVEL = 1<<2, // velocity
mjSTATE_ACT = 1<<3, // actuator activation
mjSTATE_HISTORY = 1<<4, // history buffers (control, sensor)
mjSTATE_WARMSTART = 1<<5, // acceleration used for warmstart
mjSTATE_CTRL = 1<<6, // control
mjSTATE_QFRC_APPLIED = 1<<7, // applied generalized force
mjSTATE_XFRC_APPLIED = 1<<8, // applied Cartesian force/torque
mjSTATE_EQ_ACTIVE = 1<<9, // enable/disable constraints
mjSTATE_MOCAP_POS = 1<<10, // positions of mocap bodies
mjSTATE_MOCAP_QUAT = 1<<11, // orientations of mocap bodies
mjSTATE_USERDATA = 1<<12, // user data
mjSTATE_PLUGIN = 1<<13, // plugin state
mjNSTATE = 14, // number of state elements
// convenience values for commonly used state specifications
mjSTATE_PHYSICS = mjSTATE_QPOS | mjSTATE_QVEL | mjSTATE_ACT | mjSTATE_HISTORY,
mjSTATE_FULLPHYSICS = mjSTATE_TIME | mjSTATE_PHYSICS | mjSTATE_PLUGIN,
mjSTATE_USER = mjSTATE_CTRL | mjSTATE_QFRC_APPLIED | mjSTATE_XFRC_APPLIED |
mjSTATE_EQ_ACTIVE | mjSTATE_MOCAP_POS | mjSTATE_MOCAP_QUAT |
mjSTATE_USERDATA,
mjSTATE_INTEGRATION = mjSTATE_FULLPHYSICS | mjSTATE_USER | mjSTATE_WARMSTART
} mjtState;
typedef enum mjtConstraint_ { // type of constraint
mjCNSTR_EQUALITY = 0, // equality constraint
mjCNSTR_FRICTION_DOF, // dof friction
mjCNSTR_FRICTION_TENDON, // tendon friction
mjCNSTR_LIMIT_JOINT, // joint limit
mjCNSTR_LIMIT_TENDON, // tendon limit
mjCNSTR_CONTACT_FRICTIONLESS, // frictionless contact
mjCNSTR_CONTACT_PYRAMIDAL, // frictional contact, pyramidal friction cone
mjCNSTR_CONTACT_ELLIPTIC // frictional contact, elliptic friction cone
} mjtConstraint;
typedef enum mjtConstraintState_ { // constraint state
mjCNSTRSTATE_SATISFIED = 0, // constraint satisfied, zero cost (limit, contact)
mjCNSTRSTATE_QUADRATIC, // quadratic cost (equality, friction, limit, contact)
mjCNSTRSTATE_LINEARNEG, // linear cost, negative side (friction)
mjCNSTRSTATE_LINEARPOS, // linear cost, positive side (friction)
mjCNSTRSTATE_CONE // squared distance to cone cost (elliptic contact)
} mjtConstraintState;
typedef enum mjtWarning_ { // warning types
mjWARN_INERTIA = 0, // (near) singular inertia matrix
mjWARN_CONTACTFULL, // too many contacts in contact list
mjWARN_CNSTRFULL, // too many constraints
mjWARN_BADQPOS, // bad number in qpos
mjWARN_BADQVEL, // bad number in qvel
mjWARN_BADQACC, // bad number in qacc
mjWARN_BADCTRL, // bad number in ctrl
mjNWARNING // number of warnings
} mjtWarning;
typedef enum mjtTimer_ { // internal timers
// main api
mjTIMER_STEP = 0, // step
mjTIMER_FORWARD, // forward
mjTIMER_INVERSE, // inverse
// breakdown of step/forward
mjTIMER_POSITION, // fwdPosition
mjTIMER_VELOCITY, // fwdVelocity
mjTIMER_ACTUATION, // fwdActuation
mjTIMER_CONSTRAINT, // fwdConstraint
mjTIMER_ADVANCE, // mj_Euler, mj_implicit
// breakdown of fwdPosition
mjTIMER_POS_KINEMATICS, // kinematics, com, tendon, transmission
mjTIMER_POS_INERTIA, // inertia computations
mjTIMER_POS_COLLISION, // collision detection
mjTIMER_POS_MAKE, // make constraints
mjTIMER_POS_PROJECT, // project constraints
// breakdown of mj_collision
mjTIMER_COL_BROAD, // broadphase
mjTIMER_COL_NARROW, // narrowphase
mjNTIMER // number of timers
} mjtTimer;
typedef enum mjtSleepState_ { // sleep state of an object
mjS_STATIC = -1, // object is static
mjS_ASLEEP = 0, // object is asleep
mjS_AWAKE = 1 // object is awake
} mjtSleepState;
#endif // MUJOCO_INCLUDE_MJTYPE_H_
include/mujoco/mjvisualize.h
+1 -1
View File
@@ -17,7 +17,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#define mjNGROUP 6 // number of geom, site, joint, skin groups with visflags
include/mujoco/mjxmacro.h
+26 -38
View File
@@ -318,9 +318,9 @@
X ( int, jnt_bodyid, njnt, 1 ) \
X ( int, jnt_actuatorid, njnt, 1 ) \
X ( int, jnt_group, njnt, 1 ) \
X ( mjtByte, jnt_limited, njnt, 1 ) \
X ( mjtByte, jnt_actfrclimited, njnt, 1 ) \
X ( mjtByte, jnt_actgravcomp, njnt, 1 ) \
X ( mjtBool, jnt_limited, njnt, 1 ) \
X ( mjtBool, jnt_actfrclimited, njnt, 1 ) \
X ( mjtBool, jnt_actgravcomp, njnt, 1 ) \
X ( mjtNum, jnt_solref, njnt, mjNREF ) \
X ( mjtNum, jnt_solimp, njnt, mjNIMP ) \
X ( mjtNum, jnt_pos, njnt, 3 ) \
@@ -419,11 +419,11 @@
X ( int, light_targetbodyid, nlight, 1 ) \
X ( int, light_type, nlight, 1 ) \
X ( int, light_texid, nlight, 1 ) \
X ( mjtByte, light_castshadow, nlight, 1 ) \
X ( mjtBool, light_castshadow, nlight, 1 ) \
X ( float, light_bulbradius, nlight, 1 ) \
X ( float, light_intensity, nlight, 1 ) \
X ( float, light_range, nlight, 1 ) \
X ( mjtByte, light_active, nlight, 1 ) \
X ( mjtBool, light_active, nlight, 1 ) \
X ( mjtNum, light_pos, nlight, 3 ) \
X ( mjtNum, light_dir, nlight, 3 ) \
X ( mjtNum, light_poscom0, nlight, 3 ) \
@@ -447,7 +447,7 @@
X ( mjtNum, flex_friction, nflex, 3 ) \
X ( mjtNum, flex_margin, nflex, 1 ) \
X ( mjtNum, flex_gap, nflex, 1 ) \
X ( mjtByte, flex_internal, nflex, 1 ) \
X ( mjtBool, flex_internal, nflex, 1 ) \
X ( int, flex_selfcollide, nflex, 1 ) \
X ( int, flex_activelayers, nflex, 1 ) \
X ( int, flex_passive, nflex, 1 ) \
@@ -501,10 +501,10 @@
X ( mjtNum, flex_edgestiffness, nflex, 1 ) \
X ( mjtNum, flex_edgedamping, nflex, 1 ) \
X ( int, flex_edgeequality, nflex, 1 ) \
X ( mjtByte, flex_rigid, nflex, 1 ) \
X ( mjtByte, flexedge_rigid, nflexedge, 1 ) \
X ( mjtByte, flex_centered, nflex, 1 ) \
X ( mjtByte, flex_flatskin, nflex, 1 ) \
X ( mjtBool, flex_rigid, nflex, 1 ) \
X ( mjtBool, flexedge_rigid, nflexedge, 1 ) \
X ( mjtBool, flex_centered, nflex, 1 ) \
X ( mjtBool, flex_flatskin, nflex, 1 ) \
X ( int, flex_bvhadr, nflex, 1 ) \
X ( int, flex_bvhnum, nflex, 1 ) \
X ( int, flexedge_J_rownnz, nflexedge, 1 ) \
@@ -595,7 +595,7 @@
#define MJMODEL_POINTERS_MATERIAL \
X ( int, mat_texid, nmat, mjNTEXROLE ) \
X ( mjtByte, mat_texuniform, nmat, 1 ) \
X ( mjtBool, mat_texuniform, nmat, 1 ) \
X ( float, mat_texrepeat, nmat, 2 ) \
X ( float, mat_emission, nmat, 1 ) \
X ( float, mat_specular, nmat, 1 ) \
@@ -625,7 +625,7 @@
X ( int, eq_obj1id, neq, 1 ) \
X ( int, eq_obj2id, neq, 1 ) \
X ( int, eq_objtype, neq, 1 ) \
X ( mjtByte, eq_active0, neq, 1 ) \
X ( mjtBool, eq_active0, neq, 1 ) \
X ( mjtNum, eq_solref, neq, mjNREF ) \
X ( mjtNum, eq_solimp, neq, mjNIMP ) \
X ( mjtNum, eq_data, neq, mjNEQDATA )
@@ -641,8 +641,8 @@
X ( int, ten_J_rownnz, ntendon, 1 ) \
X ( int, ten_J_rowadr, ntendon, 1 ) \
X ( int, ten_J_colind, nJten, 1 ) \
X ( mjtByte, tendon_limited, ntendon, 1 ) \
X ( mjtByte, tendon_actfrclimited, ntendon, 1 ) \
X ( mjtBool, tendon_limited, ntendon, 1 ) \
X ( mjtBool, tendon_actfrclimited, ntendon, 1 ) \
X ( mjtNum, tendon_width, ntendon, 1 ) \
X ( mjtNum, tendon_solref_lim, ntendon, mjNREF ) \
X ( mjtNum, tendon_solimp_lim, ntendon, mjNIMP ) \
@@ -678,13 +678,13 @@
X ( int, actuator_history, nu, 2 ) \
X ( int, actuator_historyadr, nu, 1 ) \
X ( mjtNum, actuator_delay, nu, 1 ) \
X ( mjtByte, actuator_ctrllimited, nu, 1 ) \
X ( mjtByte, actuator_forcelimited, nu, 1 ) \
X ( mjtByte, actuator_actlimited, nu, 1 ) \
X ( mjtBool, actuator_ctrllimited, nu, 1 ) \
X ( mjtBool, actuator_forcelimited, nu, 1 ) \
X ( mjtBool, actuator_actlimited, nu, 1 ) \
X ( mjtNum, actuator_dynprm, nu, mjNDYN ) \
X ( mjtNum, actuator_gainprm, nu, mjNGAIN ) \
X ( mjtNum, actuator_biasprm, nu, mjNBIAS ) \
X ( mjtByte, actuator_actearly, nu, 1 ) \
X ( mjtBool, actuator_actearly, nu, 1 ) \
X ( mjtNum, actuator_ctrlrange, nu, 2 ) \
X ( mjtNum, actuator_forcerange, nu, 2 ) \
X ( mjtNum, actuator_actrange, nu, 2 ) \
@@ -828,7 +828,7 @@
X ( mjtNum, ctrl, nu, 1 ) \
X ( mjtNum, qfrc_applied, nv, 1 ) \
X ( mjtNum, xfrc_applied, nbody, 6 ) \
X ( mjtByte, eq_active, neq, 1 ) \
X ( mjtBool, eq_active, neq, 1 ) \
X ( mjtNum, mocap_pos, nmocap, 3 ) \
X ( mjtNum, mocap_quat, nmocap, 4 ) \
X ( mjtNum, qacc, nv, 1 ) \
@@ -879,7 +879,7 @@
XNV ( mjtNum, M, nC, 1 ) \
XNV ( mjtNum, qLD, nC, 1 ) \
X ( mjtNum, qLDiagInv, nv, 1 ) \
X ( mjtByte, bvh_active, nbvh, 1 ) \
X ( mjtBool, bvh_active, nbvh, 1 ) \
X ( int, tree_awake, ntree, 1 ) \
X ( int, body_awake, nbody, 1 ) \
X ( int, body_awake_ind, nbody, 1 ) \
@@ -969,12 +969,6 @@
X ( int, map_idof2dof, MJ_M(nv), 1 ) \
X ( mjtNum, ifrc_smooth, MJ_D(nidof), 1 ) \
X ( mjtNum, iacc_smooth, MJ_D(nidof), 1 ) \
XNV( int, iM_rownnz, MJ_D(nidof), 1 ) \
XNV( int, iM_rowadr, MJ_D(nidof), 1 ) \
XNV( int, iM_colind, MJ_M(nC), 1 ) \
XNV( mjtNum, iM, MJ_M(nC), 1 ) \
XNV( mjtNum, iLD, MJ_M(nC), 1 ) \
X ( mjtNum, iLDiagInv, MJ_D(nidof), 1 ) \
X ( mjtNum, iacc, MJ_D(nidof), 1 ) \
X ( int, efc_island, MJ_D(nefc), 1 ) \
X ( int, island_ne, MJ_D(nisland), 1 ) \
@@ -985,11 +979,6 @@
X ( int, map_iefc2efc, MJ_D(nefc), 1 ) \
X ( int, iefc_type, MJ_D(nefc), 1 ) \
X ( int, iefc_id, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_rownnz, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_rowadr, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_rowsuper, MJ_D(nefc), 1 ) \
XNV( int, iefc_J_colind, MJ_D(nJ), 1 ) \
XNV( mjtNum, iefc_J, MJ_D(nJ), 1 ) \
X ( mjtNum, iefc_frictionloss, MJ_D(nefc), 1 ) \
X ( mjtNum, iefc_D, MJ_D(nefc), 1 ) \
X ( mjtNum, iefc_R, MJ_D(nefc), 1 ) \
@@ -1014,6 +1003,7 @@
X( size_t, pstack ) \
X( size_t, pbase ) \
X( size_t, parena ) \
X( uintptr_t, threadpool ) \
X( size_t, maxuse_stack ) \
X( size_t, maxuse_arena ) \
X( int, maxuse_con ) \
@@ -1032,17 +1022,15 @@
X( int, nbody_awake ) \
X( int, nparent_awake ) \
X( int, nv_awake ) \
X( mjtByte, flg_energypos ) \
X( mjtByte, flg_energyvel ) \
X( mjtByte, flg_subtreevel ) \
X( mjtByte, flg_rnepost ) \
X( mjtNum, time ) \
X( uintptr_t, threadpool )
X( mjtBool, flg_energypos ) \
X( mjtBool, flg_energyvel ) \
X( mjtBool, flg_subtreevel ) \
X( mjtBool, flg_rnepost ) \
X( mjtNum, time )
// vector fields of mjData
#define MJDATA_VECTOR \
X( size_t, maxuse_threadstack, mjMAXTHREAD, 1 ) \
X( mjSolverStat, solver, mjNISLAND, mjNSOLVER ) \
X( int, solver_niter, mjNISLAND, 1 ) \
X( int, solver_nnz, mjNISLAND, 1 ) \
include/mujoco/mujoco.h
+23 -29
View File
@@ -16,7 +16,7 @@
#define MUJOCO_MUJOCO_H_
// header version; should match the library version as returned by mj_version()
#define mjVERSION_HEADER 3009000
#define mjVERSION_HEADER 3010000
// needed to define size_t, fabs and log10
#include <stdlib.h>
@@ -31,10 +31,10 @@
#include <mujoco/mjrender.h>
#include <mujoco/mjsan.h>
#include <mujoco/mjspec.h>
#include <mujoco/mjthread.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjui.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mjassert.h>
// this is a C-API
#ifdef __cplusplus
@@ -679,7 +679,7 @@ MJAPI const char* mj_versionString(void);
// geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.
// Nullable: geomgroup, geomid, normal
MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], const mjtNum vec[3],
const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
int geomid[1], mjtNum normal[3]);
// Intersect multiple rays emanating from a single point, compute normals if given.
@@ -687,7 +687,7 @@ MJAPI mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum pnt[3], cons
// Geoms further than cutoff are ignored.
// Nullable: geomgroup, geomid, normal
MJAPI void mj_multiRay(const mjModel* m, mjData* d, const mjtNum pnt[3], const mjtNum* vec,
const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
const mjtByte* geomgroup, mjtBool flg_static, int bodyexclude,
int* geomid, mjtNum* dist, mjtNum* normal, int nray, mjtNum cutoff);
// Intersect ray with hfield; return nearest distance or -1 if no intersection.
@@ -710,8 +710,8 @@ MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum
// and also output nearest vertex id and surface normal.
// Nullable: vertid, normal
MJAPI mjtNum mj_rayFlex(const mjModel* m, const mjData* d, int flex_layer,
mjtByte flg_vert, mjtByte flg_edge, mjtByte flg_face,
mjtByte flg_skin, int flexid, const mjtNum pnt[3],
mjtBool flg_vert, mjtBool flg_edge, mjtBool flg_face,
mjtBool flg_skin, int flexid, const mjtNum pnt[3],
const mjtNum vec[3], int vertid[1], mjtNum normal[3]);
// Intersect ray with skin; return nearest distance or -1 if no intersection,
@@ -1303,14 +1303,14 @@ MJAPI void mju_cholSolveBand(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
// Convert banded matrix to dense matrix, fill upper triangle if flg_sym>0.
MJAPI void mju_band2Dense(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense,
mjtByte flg_sym);
mjtBool flg_sym);
// Convert dense matrix to banded matrix.
MJAPI void mju_dense2Band(mjtNum* res, const mjtNum* mat, int ntotal, int nband, int ndense);
// Multiply band-diagonal matrix with nvec vectors, include upper triangle if flg_sym>0.
MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
int ntotal, int nband, int ndense, int nvec, mjtByte flg_sym);
int ntotal, int nband, int ndense, int nvec, mjtBool flg_sym);
// Address of diagonal element i in band-dense matrix representation.
MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
@@ -1460,7 +1460,7 @@ MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNu
// D: (nsensordata x 2*nv+na)
// C: (nsensordata x nu)
// Nullable: A, B, C, D
MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_centered,
mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
// Finite differenced Jacobians of (force, sensors) = mj_inverse(state, acceleration)
@@ -1479,7 +1479,7 @@ MJAPI void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg
// optionally computes mass matrix Jacobian DmDq
// flg_actuation specifies whether to subtract qfrc_actuator from qfrc_inverse
// Nullable: DfDq, DfDv, DfDa, DsDq, DsDv, DsDa, DmDq
MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
MJAPI void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtBool flg_actuation,
mjtNum *DfDq, mjtNum *DfDv, mjtNum *DfDa,
mjtNum *DsDq, mjtNum *DsDv, mjtNum *DsDa,
mjtNum *DmDq);
@@ -1592,23 +1592,8 @@ MJAPI mjSpec* mju_decodeResource(mjResource* resource, const char* content_type,
//---------------------------------- Threads -------------------------------------------------------
// Create a thread pool with the specified number of threads running.
MJAPI mjThreadPool* mju_threadPoolCreate(size_t number_of_threads);
// Adds a thread pool to mjData and configures it for multi-threaded use.
MJAPI void mju_bindThreadPool(mjData* d, void* thread_pool);
// Enqueue a task in a thread pool.
MJAPI void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
// Destroy a thread pool.
MJAPI void mju_threadPoolDestroy(mjThreadPool* thread_pool);
// Initialize an mjTask.
MJAPI void mju_defaultTask(mjTask* task);
// Wait for a task to complete.
MJAPI void mju_taskJoin(mjTask* task);
// Create a thread pool with nthread worker threads.
MJAPI void mju_threadpool(mjData* d, int nthread);
//---------------------------------- Attachment ----------------------------------------------------
@@ -1666,6 +1651,15 @@ MJAPI mjsSensor* mjs_addSensor(mjSpec* s);
// Add flex.
MJAPI mjsFlex* mjs_addFlex(mjSpec* s);
// Add flexcomp: create flex with auto-generated bodies/joints, return flex spec.
// Nullable: type, dof, count, cellcount, spacing, scale, pos, quat, origin, file, vfs
MJAPI mjsFlex* mjs_makeFlex(mjsBody* body, const char* name, const char* type, int dim,
const char* dof, const int count[3], const int cellcount[3],
const double spacing[3], const double scale[3], double radius,
double mass, double inertiabox, int equality, int rigid, int flatskin,
int elastic2d, const double pos[3], const double quat[4],
const double origin[3], const char* file, const mjVFS* vfs);
// Add contact pair.
// Nullable: def
MJAPI mjsPair* mjs_addPair(mjSpec* s, const mjsDefault* def);
@@ -1859,7 +1853,7 @@ MJAPI void mjs_setString(mjString* dest, const char* text);
MJAPI void mjs_setStringVec(mjStringVec* dest, const char* text);
// Set entry in string vector.
MJAPI mjtByte mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
MJAPI mjtBool mjs_setInStringVec(mjStringVec* dest, int i, const char* text);
// Append text entry to string vector.
MJAPI void mjs_appendString(mjStringVec* dest, const char* text);
mjx/cuda_requirements.txt
+5 -5
View File
@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
--hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
--hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
warp-lang==1.12.1 \
--hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
--hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
--hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
--hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
warp-lang==1.13.0 \
--hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
--hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
--hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
--hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8
mjx/mujoco/mjx/_src/io.py
+13 -4
View File
@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
wp_field.shape = wp_field.shape[1:]
# warp scalars
wp_dtype = type(wp_field)
if wp_dtype in wp.types.warp_type_to_np_dtype:
return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
if wp_dtype in wp._src.types.warp_type_to_np_dtype:
return wp.dtype_to_numpy(wp_dtype)(wp_field)
# warp arrays
if isinstance(wp_field, wp.array):
@@ -272,6 +272,10 @@ def _put_option(
if impl == types.Impl.WARP:
if not mjxw.mjwp_io.ENABLE_ISLANDS:
fields['disableflags'] = types.DisableBit(
fields['disableflags'] | mjwp_types.DisableBit.ISLAND
)
impl_fields = {k: _wp_to_np_type(v) for k, v in impl_fields.items()}
return types.Option(**fields, _impl=mjxw.types.OptionWarp(**impl_fields))
@@ -1272,8 +1276,11 @@ def _get_data_into_warp(
else:
value = getattr(d_i, field.name)
if field.name in ('ne', 'nl', 'nf'):
pass
if field.name in ('ne', 'nl', 'nf', 'nisland', 'nidof'):
if isinstance(value, np.ndarray) and value.size == 0:
value = 0
else:
value = int(value)
elif field.name in ('nefc', 'ncon'):
value = {'nefc': nefc, 'ncon': ncon}[field.name]
elif field.name.endswith('xmat') or field.name == 'ximat':
@@ -1287,9 +1294,11 @@ def _get_data_into_warp(
'contact',
'qM',
'qLD',
'qLU',
'qLDiagInv',
'ten_J',
'flexedge_J',
'M',
):
continue
if field.name.startswith('efc_'):
mjx/mujoco/mjx/_src/io_test.py
+1 -1
View File
@@ -554,7 +554,7 @@ class DataIOTest(parameterized.TestCase):
elif impl == 'warp':
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
np.testing.assert_allclose(dx._impl.qM, qm)
np.testing.assert_allclose(dx._impl.M[:m.nv, :m.nv], qm)
# TODO(taylorhowell): test efc__J
np.testing.assert_allclose(dx._impl.efc__aref[:3], d.efc_aref[:3])
mjx/mujoco/mjx/codegen/generate_warp_shim.py
+4 -3
View File
@@ -215,7 +215,7 @@ def _warp_function(
info = mjwarp_field_info[f]
expected_type = _clean_type(info.expected_type)
fn_args_model.append((f'{f}: {expected_type},', info.param_order))
fn_assignments.append(f' _m.{f.replace('__', '.')} = {f}')
fn_assignments.append(f" _m.{f.replace('__', '.')} = {f}")
fn_args_model = sorted(fn_args_model, key=lambda x: x[1])
fn_args_model = ['# Model'] + [f[0] for f in fn_args_model]
@@ -233,7 +233,7 @@ def _warp_function(
param_order = info.param_order
expected_type = _clean_type(info.expected_type)
fn_args_data.append((f'{f}: {expected_type},', (is_out, param_order)))
fn_assignments.append(f' _d.{f.replace('__', '.')} = {f}')
fn_assignments.append(f" _d.{f.replace('__', '.')} = {f}")
fn_args_data = sorted(fn_args_data, key=lambda x: x[1])
fn_args_data = ['# Data'] + [f[0] for f in fn_args_data]
@@ -418,6 +418,7 @@ def create_jax_warp_shim(
fn_args_raw_str = '\n'.join([' ' + arg for arg in fn_args_raw])
warp_fn_args = [arg.split(':')[0] for arg in fn_args_raw if '#' not in arg] # pytype: disable=attribute-error
fn_assignments_str = '\n'.join(fn_assignments)
src += f"""
@ffi.format_args_for_warp
def _{fn_name}_shim(
@@ -428,7 +429,7 @@ def _{fn_name}_shim(
_m.callback = _cb
_d.efc = _e
_d.contact = _c
{'\n'.join(fn_assignments)}
{fn_assignments_str}
{fn_call}
"""
src += '\n\n'
mjx/mujoco/mjx/codegen/generate_warp_types.py
+6 -3
View File
@@ -67,7 +67,8 @@ def _to_py_string(value, indent=0):
next_indent_str = ' ' * (indent + 1)
if isinstance(value, tuple):
items = [_to_py_string(item, indent) for item in value]
return f'({', '.join(items)})'
joined_items = ', '.join(items)
return f'({joined_items})'
if isinstance(value, type):
if value.__module__ == 'builtins':
@@ -79,12 +80,14 @@ def _to_py_string(value, indent=0):
f'\n{next_indent_str}{repr(k)}: {_to_py_string(v, indent + 1)}'
for k, v in sorted(value.items(), key=lambda x: x[0])
]
return f'{{{','.join(items)}\n{indent_str}}}'
joined_items = ','.join(items)
return f'{{{joined_items}\n{indent_str}}}'
if isinstance(value, set):
items = sorted([_to_py_string(item, indent) for item in value])
items = [f'\n{next_indent_str}{item}' for item in items]
return f'{{{",".join(items)}\n{indent_str}}}'
joined_items = ','.join(items)
return f'{{{joined_items}\n{indent_str}}}'
return repr(value)
mjx/mujoco/mjx/third_party/mujoco_warp/__init__.py
Vendored
+4
View File
@@ -47,6 +47,10 @@ from mujoco.mjx.third_party.mujoco_warp._src.forward import implicit as implicit
from mujoco.mjx.third_party.mujoco_warp._src.forward import rungekutta4 as rungekutta4
from mujoco.mjx.third_party.mujoco_warp._src.forward import step1 as step1
from mujoco.mjx.third_party.mujoco_warp._src.forward import step2 as step2
from mujoco.mjx.third_party.mujoco_warp._src.history import init_ctrl_history as init_ctrl_history
from mujoco.mjx.third_party.mujoco_warp._src.history import init_sensor_history as init_sensor_history
from mujoco.mjx.third_party.mujoco_warp._src.history import read_ctrl as read_ctrl
from mujoco.mjx.third_party.mujoco_warp._src.history import read_sensor as read_sensor
from mujoco.mjx.third_party.mujoco_warp._src.inverse import inverse as inverse
from mujoco.mjx.third_party.mujoco_warp._src.io import create_render_context as create_render_context
from mujoco.mjx.third_party.mujoco_warp._src.io import get_data_into as get_data_into
mjx/mujoco/mjx/third_party/mujoco_warp/_src/block_cholesky.py
Vendored
+73 -36
View File
@@ -19,48 +19,77 @@ import warp as wp
@lru_cache(maxsize=None)
def create_blocked_cholesky_func(block_size: int):
def create_blocked_cholesky_factorize_solve_func(block_size: int, matrix_size_static: int):
@wp.func
def blocked_cholesky_func(
def blocked_cholesky_factorize_solve_func(
# In:
A: wp.array2d[float],
b: wp.array2d[float],
matrix_size: int,
# Out:
L: wp.array2d[float],
U: wp.array2d[float],
x: wp.array2d[float],
):
"""Computes the Cholesky factorization of a symmetric positive definite matrix A in blocks.
"""Block Cholesky factorization and solve while keeping the forward RHS live."""
rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
It returns a lower-triangular matrix L such that A = L L^T.
"""
# Process the matrix in blocks along its leading dimension.
for k in range(0, matrix_size, block_size):
end = k + block_size
rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(k, 0))
# Load current diagonal block A[k:end, k:end]
# and update with contributions from previously computed blocks.
A_kk_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(k, k), storage="shared")
A_kk_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, k), storage="shared", bounds_check=False, aligned=True
)
for j in range(0, k, block_size):
L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
wp.tile_matmul(L_block, wp.tile_transpose(L_block), A_kk_tile, alpha=-1.0)
U_block = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_block), U_block, A_kk_tile, alpha=-1.0)
# Compute the Cholesky factorization for the block
wp.tile_cholesky_inplace(A_kk_tile)
wp.tile_store(L, A_kk_tile, offset=(k, k))
y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
wp.tile_cholesky_inplace(A_kk_tile, fill_mode="upper")
wp.tile_store(U, A_kk_tile, offset=(k, k), bounds_check=False, aligned=True)
wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), rhs_view)
# Process the blocks below the current block
for i in range(end, matrix_size, block_size):
A_ik_tile = wp.tile_load(A, shape=(block_size, block_size), offset=(i, k), storage="shared")
A_ki_tile = wp.tile_load(
A, shape=(block_size, block_size), offset=(k, i), storage="shared", bounds_check=False, aligned=True
)
for j in range(0, k, block_size):
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
L_2_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(k, j), storage="shared")
wp.tile_matmul(L_tile, wp.tile_transpose(L_2_tile), A_ik_tile, alpha=-1.0)
U_jk_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, k), storage="shared", bounds_check=False, aligned=True
)
U_ji_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_matmul(wp.tile_transpose(U_jk_tile), U_ji_tile, A_ki_tile, alpha=-1.0)
wp.tile_lower_solve_inplace(A_kk_tile, wp.tile_transpose(A_ik_tile))
wp.tile_store(L, A_ik_tile, offset=(i, k))
wp.tile_lower_solve_inplace(wp.tile_transpose(A_kk_tile), A_ki_tile)
wp.tile_store(U, A_ki_tile, offset=(k, i), bounds_check=False, aligned=True)
return blocked_cholesky_func
for i in range(matrix_size - block_size, -1, -block_size):
i_end = i + block_size
tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(i_end, matrix_size, block_size):
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
)
x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
return blocked_cholesky_factorize_solve_func
@lru_cache(maxsize=None)
@@ -68,41 +97,49 @@ def create_blocked_cholesky_solve_func(block_size: int, matrix_size_static: int)
@wp.func
def blocked_cholesky_solve_func(
# In:
L: wp.array2d[float],
U: wp.array2d[float],
b: wp.array2d[float],
matrix_size: int,
# Out:
x: wp.array2d[float],
):
"""Block Cholesky factorization and solve.
"""Block Cholesky solve.
Solves A x = b given the Cholesky factor L (A = L L^T) using blocked forward and backward
Solves A x = b given the Cholesky factor U (A = U^T U) using blocked forward and backward
substitution.
"""
rhs_tile = wp.tile_load(b, shape=(matrix_size_static, 1), offset=(0, 0), storage="shared", bounds_check=False)
# Forward substitution: solve L y = b
# Forward substitution: solve U^T y = b
for i in range(0, matrix_size, block_size):
rhs_view = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(0, i, block_size):
L_block = wp.tile_load(L, shape=(block_size, block_size), offset=(i, j), storage="shared")
U_block = wp.tile_load(
U, shape=(block_size, block_size), offset=(j, i), storage="shared", bounds_check=False, aligned=True
)
y_block = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(L_block, y_block, rhs_view, alpha=-1.0)
wp.tile_matmul(wp.tile_transpose(U_block), y_block, rhs_view, alpha=-1.0)
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
wp.tile_lower_solve_inplace(L_tile, rhs_view)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_lower_solve_inplace(wp.tile_transpose(U_tile), rhs_view)
# Backward substitution: solve L^T x = y
# Backward substitution: solve U x = y
for i in range(matrix_size - block_size, -1, -block_size):
i_end = i + block_size
tmp_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(i, 0))
for j in range(i_end, matrix_size, block_size):
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(j, i), storage="shared")
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, j), storage="shared", bounds_check=False, aligned=True
)
x_tile = wp.tile_view(rhs_tile, shape=(block_size, 1), offset=(j, 0))
wp.tile_matmul(wp.tile_transpose(L_tile), x_tile, tmp_tile, alpha=-1.0)
L_tile = wp.tile_load(L, shape=(block_size, block_size), offset=(i, i), storage="shared")
wp.tile_matmul(U_tile, x_tile, tmp_tile, alpha=-1.0)
U_tile = wp.tile_load(
U, shape=(block_size, block_size), offset=(i, i), storage="shared", bounds_check=False, aligned=True
)
wp.tile_upper_solve_inplace(wp.tile_transpose(L_tile), tmp_tile)
wp.tile_upper_solve_inplace(U_tile, tmp_tile)
wp.tile_store(x, rhs_tile, offset=(0, 0), bounds_check=False)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/cli.py
Vendored
+23 -6
View File
@@ -25,6 +25,7 @@ from absl import flags
from etils import epath
import mujoco.mjx.third_party.mujoco_warp as mjw
from mujoco.mjx.third_party.mujoco_warp._src import io
from mujoco.mjx.third_party.mujoco_warp._src import warp_util
from mujoco.mjx.third_party.mujoco_warp._src.io import load_trajectory
from mujoco.mjx.third_party.mujoco_warp._src.io import override_model
@@ -42,6 +43,12 @@ KEYFRAME = flags.DEFINE_integer("keyframe", 0, "keyframe to initialize simulatio
EVENT_TRACE = flags.DEFINE_bool("event_trace", False, "print an event trace report")
NOISE_STD = flags.DEFINE_float("noise_std", 0.01, "add noise to ctrl signal (standard deviation)")
NOISE_RATE = flags.DEFINE_float("noise_rate", 0.1, "add noise to ctrl signal (noise rate)")
ENABLE_ISLANDS = flags.DEFINE_bool(
"enable_islands",
False,
"Enable constraint islands solver",
)
DEVICE = flags.DEFINE_string("device", None, "override the default Warp device")
REPLAY = flags.DEFINE_string("replay", None, "NPZ file with ctrl sequence to replay")
@@ -52,6 +59,12 @@ RENDER_RGB = flags.DEFINE_bool("render_rgb", True, "render RGB image")
RENDER_DEPTH = flags.DEFINE_bool("render_depth", True, "render depth image")
RENDER_TEXTURES = flags.DEFINE_bool("render_textures", True, "use textures")
RENDER_SHADOWS = flags.DEFINE_bool("render_shadows", False, "use shadows")
RENDER_BACKFACE_CULLING = flags.DEFINE_bool(
"render_backface_culling",
True,
"enable renderer backface culling (RenderContext.enable_backface_culling)",
)
RENDER_SKYBOX = flags.DEFINE_bool("render_skybox", True, "render skybox texture if available")
def load_model(path: epath.Path) -> mujoco.MjModel:
@@ -128,6 +141,8 @@ def init_structs(
fn: Callable[..., None], mjm: mujoco.MjModel
) -> Tuple[mjw.Model, mjw.Data, mjw.RenderContext | None, list[np.ndarray] | None]:
"""Initialize device structs."""
io.ENABLE_ISLANDS = ENABLE_ISLANDS.value
mjd = mujoco.MjData(mjm)
ctrls = None
if REPLAY.value:
@@ -152,12 +167,14 @@ def init_structs(
rc = mjw.create_render_context(
mjm,
NWORLD.value,
(RENDER_WIDTH.value, RENDER_HEIGHT.value),
RENDER_RGB.value,
RENDER_DEPTH.value,
RENDER_TEXTURES.value,
RENDER_SHADOWS.value,
nworld=NWORLD.value,
cam_res=(RENDER_WIDTH.value, RENDER_HEIGHT.value),
render_rgb=RENDER_RGB.value,
render_depth=RENDER_DEPTH.value,
use_textures=RENDER_TEXTURES.value,
use_shadows=RENDER_SHADOWS.value,
enable_backface_culling=RENDER_BACKFACE_CULLING.value,
render_skybox=RENDER_SKYBOX.value,
)
return m, d, rc, ctrls
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_convex.py
Vendored
+34 -18
View File
@@ -23,6 +23,8 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import contact_param
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collision_pair
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import ccd
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import epa_phase
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import gjk_phase
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import multicontact
from mujoco.mjx.third_party.mujoco_warp._src.collision_gjk import support
from mujoco.mjx.third_party.mujoco_warp._src.collision_primitive import Geom
@@ -35,7 +37,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAFACES
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAX_EPAHORIZON
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
@@ -715,6 +717,7 @@ def ccd_kernel_builder(
opt_ccd_tolerance: wp.array[float],
# Data in:
naconmax_in: int,
naccdmax_in: int,
# In:
epa_vert_in: wp.array2d[wp.vec3],
epa_vert_index_in: wp.array2d[int],
@@ -737,7 +740,7 @@ def ccd_kernel_builder(
geom2: Geom,
geoms: wp.vec2i,
worldid: int,
ccdid: int,
nccd_in: wp.array[int],
margin: float,
gap: float,
condim: int,
@@ -774,30 +777,47 @@ def ccd_kernel_builder(
if is_collision_sensor:
cutoff = 1.0e32
else:
if wp.static(_NEW_GAP_SEMANTICS):
if wp.static(NEW_GAP_SEMANTICS):
cutoff = gap
else:
cutoff = 0.0
dist, ncollision, w1, w2, multiccd_idx = ccd(
needs_epa, dist, ncollision, w1, w2, gjk_result, geom1, geom2 = gjk_phase(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
cutoff,
gjk_iterations,
epa_iterations,
geom1,
geom2,
geomtype1,
geomtype2,
x1,
x2,
epa_vert_in[ccdid],
epa_vert_index_in[ccdid],
epa_face_in[ccdid],
epa_pr_in[ccdid],
epa_norm2_in[ccdid],
epa_horizon_in[ccdid],
)
if wp.static(_NEW_GAP_SEMANTICS):
ccdid = int(-1)
multiccd_idx = int(-1)
if needs_epa:
ccdid = wp.atomic_add(nccd_in, geomgeomid, 1)
if ccdid >= naccdmax_in:
wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
return 0
dist, ncollision, w1, w2, multiccd_idx = epa_phase(
opt_ccd_tolerance[worldid % opt_ccd_tolerance.shape[0]],
epa_iterations,
gjk_result,
geom1,
geom2,
geomtype1,
geomtype2,
epa_vert_in[ccdid],
epa_vert_index_in[ccdid],
epa_face_in[ccdid],
epa_pr_in[ccdid],
epa_norm2_in[ccdid],
epa_horizon_in[ccdid],
)
if wp.static(NEW_GAP_SEMANTICS):
if dist >= gap and pairid[1] == -1:
return 0
else:
@@ -990,11 +1010,6 @@ def ccd_kernel_builder(
if geom_type[g1] != geomtype1 or geom_type[g2] != geomtype2:
return
ccdid = wp.atomic_add(nccd_in, wp.static(geomgeomid), 1)
if ccdid >= naccdmax_in:
wp.printf("CCD overflow - please increase naccdmax to %u\n", ccdid)
return
worldid = collision_worldid_in[collisionid]
_, margin, gap, condim, friction, solref, solreffriction, solimp = contact_params(
@@ -1046,6 +1061,7 @@ def ccd_kernel_builder(
eval_ccd_write_contact(
opt_ccd_tolerance,
naconmax_in,
naccdmax_in,
epa_vert_in,
epa_vert_index_in,
epa_face_in,
@@ -1067,7 +1083,7 @@ def ccd_kernel_builder(
geom2,
geoms,
worldid,
ccdid,
nccd_in,
margin,
gap,
condim,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_core.py
Vendored
+2 -2
View File
@@ -22,8 +22,8 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.math import safe_div
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINMU
from mujoco.mjx.third_party.mujoco_warp._src.types import _NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import NEW_GAP_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import mat63
@@ -219,7 +219,7 @@ def write_contact(
contact_frame_out[cid] = frame_in
contact_geom_out[cid] = geoms_in
contact_worldid_out[cid] = worldid_in
if wp.static(_NEW_GAP_SEMANTICS):
if wp.static(NEW_GAP_SEMANTICS):
includemargin = margin_in
else:
includemargin = margin_in - gap_in
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_driver.py
Vendored
+8 -2
View File
@@ -375,6 +375,7 @@ def _binary_search(values: wp.array[Any], value: Any, lower: int, upper: int) ->
return upper
@cache_kernel
def _sap_project(opt_broadphase: int):
@wp.kernel(module="unique", enable_backward=False)
def sap_project(
@@ -531,6 +532,7 @@ def _sap_broadphase(opt_broadphase_filter: int, ngeom_aabb: int, ngeom_rbound: i
return kernel
@cache_kernel
def _segmented_sort(tile_size: int):
@wp.kernel(module="unique")
def segmented_sort(
@@ -629,7 +631,9 @@ def sap_broadphase(m: Model, d: Data, ctx: CollisionContext):
# assumes each geom has 5 other geoms (batched over all worlds)
nsweep = 5 * nworldgeom
wp.launch(
kernel=_sap_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
kernel=_sap_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
),
dim=nsweep,
inputs=[
m.ngeom,
@@ -718,7 +722,9 @@ def nxn_broadphase(m: Model, d: Data, ctx: CollisionContext):
`contype`/`conaffinity`, parent-child relationships, and explicit `<exclude>` tags.
"""
wp.launch(
_nxn_broadphase(m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]),
_nxn_broadphase(
m.opt.broadphase_filter, m.geom_aabb.shape[0], m.geom_rbound.shape[0], m.geom_margin.shape[0], m.geom_gap.shape[0]
),
dim=(d.nworld, m.nxn_geom_pair_filtered.shape[0]),
inputs=[
m.geom_type,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_gjk.py
Vendored
+80 -17
View File
@@ -2135,8 +2135,12 @@ def multicontact(
# face1 is an edge; clip face1 against face2
if is_edge_contact_geom1:
approx_dir = wp.norm_l2(dir) * n2[j]
return _polygon_clip(plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped)
approx_dir = -wp.norm_l2(dir) * n2[j]
nclipped, clipped1, clipped2 = _polygon_clip(
plane_normal, plane_dist, face2, nface2, face1, nface1, n2[j], approx_dir, polygon, clipped
)
# the faces were flipped in calling _polygon_clip so we need to flip them back
return nclipped, clipped2, clipped1
# face2 is an edge; clip face2 against face1
if is_edge_contact_geom2:
@@ -2197,34 +2201,29 @@ def _inflate(
@wp.func
def ccd(
def gjk_phase(
# In:
tolerance: float,
cutoff: float,
gjk_iterations: int,
epa_iterations: int,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
x_1: wp.vec3,
x_2: wp.vec3,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA."""
) -> Tuple[bool, float, int, wp.vec3, wp.vec3, GJKResult, Geom, Geom]:
"""Run GJK phase of CCD."""
full_margin1 = 0.0
full_margin2 = 0.0
size1 = 0.0
size2 = 0.0
empty = GJKResult()
# determine if the geoms being tested are discrete
is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
# special handling for sphere and capsule (shrink to point and line respectively)
if geomtype1 == GeomType.SPHERE or geomtype1 == GeomType.CAPSULE:
size1 = geom1.size[0]
full_margin1 = size1 + 0.5 * geom1.margin
@@ -2237,7 +2236,6 @@ def ccd(
geom2.margin = 0.0
geom2.size = wp.vec3(0.0, geom2.size[1], geom2.size[2])
# special handling for sphere and capsule (shrink to point and line respectively)
if size1 + size2 > 0.0:
cutoff += full_margin1 + full_margin2
result = gjk(tolerance, gjk_iterations, geom1, geom2, x_1, x_2, geomtype1, geomtype2, cutoff, is_discrete)
@@ -2245,11 +2243,11 @@ def ccd(
# shallow penetration, inflate contact
if result.dist > tolerance:
if result.dist == FLOAT_MAX:
return result.dist, 1, result.x1, result.x2, -1
return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
dist, x1, x2 = _inflate(result, geom1, geom2, geomtype1, geomtype2, full_margin1, full_margin2)
return dist, 1, x1, x2, -1
return False, dist, 1, x1, x2, empty, geom1, geom2
# deep penetration, reset initial conditions and rerun GJK + EPA
# deep penetration: reset initial conditions and rerun GJK + EPA
geom1.margin = full_margin1 - size1
geom1.size = wp.vec3(size1, geom1.size[1], geom1.size[2])
geom2.margin = full_margin2 - size2
@@ -2260,8 +2258,29 @@ def ccd(
# no penetration depth to recover
if result.dist > tolerance or result.dim < 2:
return result.dist, 1, result.x1, result.x2, -1
return False, result.dist, 1, result.x1, result.x2, empty, geom1, geom2
return True, result.dist, 1, result.x1, result.x2, result, geom1, geom2
@wp.func
def epa_phase(
# In:
tolerance: float,
epa_iterations: int,
result: GJKResult,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""Run EPA given GJK result. Returns (dist, ncontact, x1, x2, multiccd_idx)."""
pt = Polytope()
pt.nface = 0
pt.nvert = 0
@@ -2330,6 +2349,7 @@ def ccd(
if pt.status:
return result.dist, 1, result.x1, result.x2, -1
is_discrete = _discrete_geoms(geomtype1, geomtype2) and (geom1.margin == 0.0 and geom2.margin == 0.0)
dist, x1, x2, idx = _epa(tolerance, epa_iterations, pt, geom1, geom2, geomtype1, geomtype2, is_discrete)
if idx == -1:
return FLOAT_MAX, 0, wp.vec3(), wp.vec3(), -1
@@ -2343,3 +2363,46 @@ def ccd(
idx = -1
return dist, 1, x1, x2, idx
@wp.func
def ccd(
# In:
tolerance: float,
cutoff: float,
gjk_iterations: int,
epa_iterations: int,
geom1: Geom,
geom2: Geom,
geomtype1: int,
geomtype2: int,
x_1: wp.vec3,
x_2: wp.vec3,
vert: wp.array[wp.vec3],
vert_index: wp.array[int],
face: wp.array[int],
face_pr: wp.array[wp.vec3],
face_norm2: wp.array[float],
horizon: wp.array[int],
) -> Tuple[float, int, wp.vec3, wp.vec3, int]:
"""General convex collision detection via GJK/EPA."""
needs_epa, dist, ncontact, x1, x2, result, geom1, geom2 = gjk_phase(
tolerance, cutoff, gjk_iterations, geom1, geom2, geomtype1, geomtype2, x_1, x_2
)
if not needs_epa:
return dist, ncontact, x1, x2, -1
return epa_phase(
tolerance,
epa_iterations,
result,
geom1,
geom2,
geomtype1,
geomtype2,
vert,
vert_index,
face,
face_pr,
face_norm2,
horizon,
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/collision_sdf.py
Vendored
+83 -9
View File
@@ -23,6 +23,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_core import geom_collisio
from mujoco.mjx.third_party.mujoco_warp._src.collision_core import write_contact
from mujoco.mjx.third_party.mujoco_warp._src.math import make_frame
from mujoco.mjx.third_party.mujoco_warp._src.ray import ray_mesh
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import GeomType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
@@ -183,6 +184,24 @@ def ellipsoid(p: wp.vec3, size: wp.vec3) -> float:
return k0 * (k0 - 1.0) / denom
@wp.func
def capsule(p: wp.vec3, size: wp.vec3) -> float:
r = size[0]
h = size[1]
pz_clamped = wp.clamp(p[2], -h, h)
diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
return wp.length(diff) - r
@wp.func
def cylinder(p: wp.vec3, size: wp.vec3) -> float:
r = size[0]
h = size[1]
dx = wp.length(wp.vec2(p[0], p[1])) - r
dy = wp.abs(p[2]) - h
return wp.min(wp.max(dx, dy), 0.0) + wp.length(wp.vec2(wp.max(dx, 0.0), wp.max(dy, 0.0)))
@wp.func
def grad_sphere(p: wp.vec3) -> wp.vec3:
c = wp.length(p)
@@ -230,6 +249,52 @@ def grad_ellipsoid(p: wp.vec3, size: wp.vec3) -> wp.vec3:
return raw_grad / wp.length(raw_grad)
@wp.func
def grad_capsule(p: wp.vec3, size: wp.vec3) -> wp.vec3:
h = size[1]
pz_clamped = wp.clamp(p[2], -h, h)
diff = wp.vec3(p[0], p[1], p[2] - pz_clamped)
c = wp.length(diff)
if c > MJ_MINVAL:
return diff / c
else:
return wp.vec3(0.0)
@wp.func
def grad_cylinder(p: wp.vec3, size: wp.vec3) -> wp.vec3:
r = size[0]
h = size[1]
radial_dist = wp.length(wp.vec2(p[0], p[1]))
if radial_dist > MJ_MINVAL:
u = wp.vec3(p[0] / radial_dist, p[1] / radial_dist, 0.0)
else:
u = wp.vec3(0.0)
w = wp.vec3(0.0, 0.0, wp.sign(p[2]))
dx = radial_dist - r
dy = wp.abs(p[2]) - h
if dx > 0.0 and dy > 0.0:
v = wp.vec2(dx, dy)
len_v = wp.length(v)
if len_v > MJ_MINVAL:
return u * (dx / len_v) + w * (dy / len_v)
else:
return wp.vec3(0.0)
elif dx > 0.0:
return u
elif dy > 0.0:
return w
else:
if dx > dy:
return u
else:
return w
@wp.func
def user_sdf(p: wp.vec3, attr: vec_pluginattr, sdf_type: int) -> float:
"""User-defined SDF function.
@@ -287,16 +352,17 @@ def find_oct(
# check if the node is a leaf
# child indices are relative to root (mesh_octadr offset)
child0 = oct_child[node][0]
# Evaluate this hot leaf predicate eagerly to avoid branch-heavy codegen.
if (
child0 == -1
and oct_child[node][1] == -1
and oct_child[node][2] == -1
and oct_child[node][3] == -1
and oct_child[node][4] == -1
and oct_child[node][5] == -1
and oct_child[node][6] == -1
and oct_child[node][7] == -1
):
int(child0 == -1)
& int(oct_child[node][1] == -1)
& int(oct_child[node][2] == -1)
& int(oct_child[node][3] == -1)
& int(oct_child[node][4] == -1)
& int(oct_child[node][5] == -1)
& int(oct_child[node][6] == -1)
& int(oct_child[node][7] == -1)
) != 0:
for j in range(8):
if not grad:
rx[j] = (
@@ -394,6 +460,10 @@ def sdf(type: int, p: wp.vec3, attr: vec_pluginattr, sdf_type: int, volume_data:
return p[2]
elif type == GeomType.SPHERE:
return sphere(p, attr_vec3)
elif type == GeomType.CAPSULE:
return capsule(p, attr_vec3)
elif type == GeomType.CYLINDER:
return cylinder(p, attr_vec3)
elif type == GeomType.BOX:
return box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
@@ -452,6 +522,10 @@ def sdf_grad(
return grad
elif type == GeomType.SPHERE:
return grad_sphere(p)
elif type == GeomType.CAPSULE:
return grad_capsule(p, attr_vec3)
elif type == GeomType.CYLINDER:
return grad_cylinder(p, attr_vec3)
elif type == GeomType.BOX:
return grad_box(p, attr_vec3)
elif type == GeomType.ELLIPSOID:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/constraint.py
Vendored
+243 -2
View File
@@ -134,7 +134,10 @@ def _equality_connect(
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
body_invweight0: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jnt_dofadr: wp.array[int],
dof_bodyid: wp.array[int],
dof_jntid: wp.array[int],
dof_parentid: wp.array[int],
site_bodyid: wp.array[int],
eq_obj1id: wp.array[int],
@@ -154,6 +157,9 @@ def _equality_connect(
site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d[wp.spatial_vector],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
subtree_linvel_in: wp.array2d[wp.vec3],
njmax_in: int,
njmax_nnz_in: int,
# Data out:
@@ -214,6 +220,7 @@ def _equality_connect(
# compute Jacobian difference (opposite of contact: 0 - 1)
Jqvel = wp.vec3f(0.0, 0.0, 0.0)
Jdotv = wp.vec3f(0.0, 0.0, 0.0)
if is_sparse:
# TODO(team): pre-compute number of non-zeros
@@ -282,6 +289,42 @@ def _equality_connect(
)
j1mj2 = jacp1 - jacp2
jacp1_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
da,
worldid,
)
jacp2_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
da,
worldid,
)
j1mj2_dot = jacp1_dot - jacp2_dot
sparseid0 = rowadr + nnz
sparseid1 = rowadr + rownnz + nnz
sparseid2 = rowadr + 2 * rownnz + nnz
@@ -294,7 +337,9 @@ def _equality_connect(
efc_J_out[worldid, 0, sparseid1] = j1mj2[1]
efc_J_out[worldid, 0, sparseid2] = j1mj2[2]
Jqvel += j1mj2 * qvel_in[worldid, da]
qvel = qvel_in[worldid, da]
Jqvel += j1mj2 * qvel
Jdotv += j1mj2_dot * qvel
nnz += 1
else:
@@ -326,11 +371,49 @@ def _equality_connect(
)
j1mj2 = jacp1 - jacp2
jacp1_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
dofid,
worldid,
)
jacp2_dot, _ = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
dofid,
worldid,
)
j1mj2_dot = jacp1_dot - jacp2_dot
efc_J_out[worldid, efcid0, dofid] = j1mj2[0]
efc_J_out[worldid, efcid1, dofid] = j1mj2[1]
efc_J_out[worldid, efcid2, dofid] = j1mj2[2]
Jqvel += j1mj2 * qvel_in[worldid, dofid]
qvel = qvel_in[worldid, dofid]
Jqvel += j1mj2 * qvel
Jdotv += j1mj2_dot * qvel
body_invweight0_id = worldid % body_invweight0.shape[0]
invweight = body_invweight0[body_invweight0_id, body1][0] + body_invweight0[body_invweight0_id, body2][0]
@@ -368,6 +451,8 @@ def _equality_connect(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcidi] -= Jdotv[i]
@wp.kernel
def _equality_joint(
@@ -812,7 +897,10 @@ def _equality_weld(
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
body_invweight0: wp.array2d[wp.vec2],
jnt_type: wp.array[int],
jnt_dofadr: wp.array[int],
dof_bodyid: wp.array[int],
dof_jntid: wp.array[int],
dof_parentid: wp.array[int],
site_bodyid: wp.array[int],
site_quat: wp.array2d[wp.quat],
@@ -834,6 +922,9 @@ def _equality_weld(
site_xpos_in: wp.array2d[wp.vec3],
subtree_com_in: wp.array2d[wp.vec3],
cdof_in: wp.array2d[wp.spatial_vector],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
subtree_linvel_in: wp.array2d[wp.vec3],
njmax_in: int,
njmax_nnz_in: int,
# Data out:
@@ -903,9 +994,43 @@ def _equality_weld(
quat = math.mul_quat(xquat_in[worldid, body1], relpose)
quat1 = math.quat_inv(xquat_in[worldid, body2])
# quat1 = quat_inv(xquat_in[worldid, body2])
q2 = xquat_in[worldid, body2]
quat1 = wp.quat(q2[0], -q2[1], -q2[2], -q2[3])
# compute rotational Jdotv helper quaternions
omega1 = wp.spatial_top(cvel_in[worldid, body1])
omega2 = wp.spatial_top(cvel_in[worldid, body2])
domega = omega1 - omega2
omega1_q = wp.quat(0.0, omega1[0], omega1[1], omega1[2])
omega2_q = wp.quat(0.0, omega2[0], omega2[1], omega2[2])
domega_q = wp.quat(0.0, domega[0], domega[1], domega[2])
if is_site:
qdot0r = math.mul_quat(omega1_q, quat) * 0.5
qfull1 = math.mul_quat(xquat_in[worldid, body2], site_quat[site_quat_id, obj2id])
qdot1 = math.mul_quat(omega2_q, qfull1) * 0.5
negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
negq1 = wp.quat(qfull1[0], -qfull1[1], -qfull1[2], -qfull1[3])
else:
# qdot0 = mul_quat(xquat_in[worldid, body1], omega1_q) * 0.5
u7 = xquat_in[worldid, body1]
qdot0 = math.mul_quat(omega1_q, xquat_in[worldid, body1]) * 0.5
qdot0r = math.mul_quat(qdot0, relpose)
q1_non_site = xquat_in[worldid, body2]
qdot1 = math.mul_quat(omega2_q, q1_non_site) * 0.5
negqdot1 = wp.quat(-qdot1[0], -qdot1[1], -qdot1[2], -qdot1[3])
negq1 = wp.quat(q1_non_site[0], -q1_non_site[1], -q1_non_site[2], -q1_non_site[3])
# compute Jacobian difference (opposite of contact: 0 - 1)
Jqvelp = wp.vec3f(0.0, 0.0, 0.0)
Jqvelr = wp.vec3f(0.0, 0.0, 0.0)
Jdotv_p = wp.vec3f(0.0, 0.0, 0.0)
Jdotv_r0 = wp.vec3f(0.0, 0.0, 0.0)
if is_sparse:
# TODO(team): pre-compute number of non-zeros
@@ -985,6 +1110,44 @@ def _equality_weld(
jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
jacp1_dot, jacr1_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
da,
worldid,
)
jacp2_dot, jacr2_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
da,
worldid,
)
jacdifp_dot = jacp1_dot - jacp2_dot
jacdifr_dot = jacr1_dot - jacr2_dot
sparseid0 = rowadr + nnz
sparseid1 = rowadr + rownnz + nnz
sparseid2 = rowadr + 2 * rownnz + nnz
@@ -1008,6 +1171,8 @@ def _equality_weld(
Jqvelp += jacdifp * qvel_in[worldid, da]
Jqvelr += jacdifr * qvel_in[worldid, da]
Jdotv_p += jacdifp_dot * qvel_in[worldid, da]
Jdotv_r0 += jacdifr_dot * qvel_in[worldid, da]
nnz += 1
else:
@@ -1047,12 +1212,52 @@ def _equality_weld(
jacdifrq = math.mul_quat(math.quat_mul_axis(quat1, jacdifr), quat)
jacdifr = 0.5 * wp.vec3(jacdifrq[1], jacdifrq[2], jacdifrq[3])
jacp1_dot, jacr1_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos1,
body1,
dofid,
worldid,
)
jacp2_dot, jacr2_dot = support.jac_dot_dof(
body_parentid,
body_rootid,
jnt_type,
jnt_dofadr,
dof_bodyid,
dof_jntid,
body_isdofancestor,
subtree_com_in,
cdof_in,
cvel_in,
cdof_dot_in,
pos2,
body2,
dofid,
worldid,
)
jacdifp_dot = jacp1_dot - jacp2_dot
jacdifr_dot = jacr1_dot - jacr2_dot
efc_J_out[worldid, efcid3, dofid] = jacdifr[0]
efc_J_out[worldid, efcid4, dofid] = jacdifr[1]
efc_J_out[worldid, efcid5, dofid] = jacdifr[2]
Jqvelp += jacdifp * qvel_in[worldid, dofid]
Jqvelr += jacdifr * qvel_in[worldid, dofid]
Jdotv_p += jacdifp_dot * qvel_in[worldid, dofid]
Jdotv_r0 += jacdifr_dot * qvel_in[worldid, dofid]
# error is difference in global position and orientation
cpos = pos1 - pos2
@@ -1070,6 +1275,22 @@ def _equality_weld(
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
djrdv_q = wp.quat(0.0, Jdotv_r0[0], Jdotv_r0[1], Jdotv_r0[2])
# Term 1: negqdot1 * domega * q0r
t1a = math.mul_quat(negqdot1, domega_q)
t1 = math.mul_quat(t1a, quat)
# Term 2: negq1 * djrdv * q0r
t2a = math.mul_quat(negq1, djrdv_q)
t2 = math.mul_quat(t2a, quat)
# Term 3: negq1 * domega * qdot0r
t3a = math.mul_quat(negq1, domega_q)
t3 = math.mul_quat(t3a, qdot0r)
Jdotv_r = wp.vec3(t1[1] + t2[1] + t3[1], t1[2] + t2[2] + t3[2], t1[3] + t2[3] + t3[3]) * 0.5 * torquescale
for i in range(3):
_efc_row(
opt_disableflags,
@@ -1096,6 +1317,8 @@ def _equality_weld(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcid + i] -= Jdotv_p[i]
invweight_r = body_invweight0[body_invweight0_id, body1][1] + body_invweight0[body_invweight0_id, body2][1]
for i in range(3):
@@ -1124,6 +1347,8 @@ def _equality_weld(
efc_frictionloss_out,
)
efc_aref_out[worldid, efcid + 3 + i] -= Jdotv_r[i]
@wp.kernel
def _friction_dof(
@@ -1680,6 +1905,7 @@ def _limit_tendon(
)
@cache_kernel
def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
IS_SPARSE = is_sparse
@@ -1803,6 +2029,7 @@ def _efc_contact_init(cone_type: types.ConeType, is_sparse: bool):
return kernel
@cache_kernel
def _efc_contact_jac_sparse(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -1973,6 +2200,7 @@ def _efc_contact_jac_sparse(cone_type: types.ConeType):
return kernel
@cache_kernel
def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
TILE_SIZE = tile_size
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2117,6 +2345,7 @@ def _efc_contact_jac_dense(tile_size: int, cone_type: types.ConeType):
return kernel
@cache_kernel
def _efc_contact_update(cone_type: types.ConeType):
IS_ELLIPTIC = cone_type == types.ConeType.ELLIPTIC
@@ -2297,7 +2526,10 @@ def make_constraint(m: types.Model, d: types.Data):
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.jnt_type,
m.jnt_dofadr,
m.dof_bodyid,
m.dof_jntid,
m.dof_parentid,
m.site_bodyid,
m.eq_obj1id,
@@ -2316,6 +2548,9 @@ def make_constraint(m: types.Model, d: types.Data):
d.site_xpos,
d.subtree_com,
d.cdof,
d.cvel,
d.cdof_dot,
d.subtree_linvel,
d.njmax,
d.njmax_nnz,
],
@@ -2351,7 +2586,10 @@ def make_constraint(m: types.Model, d: types.Data):
m.body_dofnum,
m.body_dofadr,
m.body_invweight0,
m.jnt_type,
m.jnt_dofadr,
m.dof_bodyid,
m.dof_jntid,
m.dof_parentid,
m.site_bodyid,
m.site_quat,
@@ -2372,6 +2610,9 @@ def make_constraint(m: types.Model, d: types.Data):
d.site_xpos,
d.subtree_com,
d.cdof,
d.cvel,
d.cdof_dot,
d.subtree_linvel,
d.njmax,
d.njmax_nnz,
],
mjx/mujoco/mjx/third_party/mujoco_warp/_src/derivative.py
Vendored
+363 -60
View File
@@ -15,6 +15,7 @@
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
@@ -24,6 +25,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import DisableBit
from mujoco.mjx.third_party.mujoco_warp._src.types import DynType
from mujoco.mjx.third_party.mujoco_warp._src.types import GainType
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10f
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -60,13 +62,47 @@ def _qderiv_actuator_passive_vel(
actuator_gainprm_id = worldid % actuator_gainprm.shape[0]
actuator_biasprm_id = worldid % actuator_biasprm.shape[0]
bias = float(0.0)
if actuator_gaintype[actid] == GainType.AFFINE:
gain = actuator_gainprm[actuator_gainprm_id, actid][2]
elif actuator_gaintype[actid] == GainType.DCMOTOR:
gain = 0.0
dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
gainprm = actuator_gainprm[actuator_gainprm_id, actid]
te = dynprm[0]
# controller velocity derivative: dV/dω
input_mode = int(gainprm[8])
dVdw = 0.0
if input_mode == 1:
dVdw = -gainprm[6] # position: -kd
elif input_mode == 2:
dVdw = -gainprm[4] # velocity: -kp
if te > 0.0:
# stateful current with actearly: d(K*next_act)/dω
# includes both back-EMF (-K) and controller (dVdw) through act_dot
R = wp.max(MJ_MINVAL, gainprm[0])
K = gainprm[1]
s = 1.0 - wp.exp(-opt_timestep[worldid % opt_timestep.shape[0]] / te)
bias += K * (dVdw - K) * s / R
elif dVdw != 0.0:
# stateless: controller terms only (back-EMF handled in bias block)
R = wp.max(MJ_MINVAL, gainprm[0])
K = gainprm[1]
bias += K * dVdw / R
# LuGre: force includes -sigma1*z_dot, z_dot = a*z + v
# d(sigma1*z_dot)/dv = sigma1*(da/dv*z + 1), ignoring higher-order da/dv*z
sigma1 = dynprm[6]
if sigma1 > 0.0:
bias -= sigma1
else:
gain = 0.0
if actuator_biastype[actid] == BiasType.AFFINE:
bias = actuator_biasprm[actuator_biasprm_id, actid][2]
bias += actuator_biasprm[actuator_biasprm_id, actid][2]
elif actuator_biastype[actid] == BiasType.DCMOTOR:
dynprm = actuator_dynprm[worldid % actuator_dynprm.shape[0], actid]
te = dynprm[0]
@@ -86,11 +122,7 @@ def _qderiv_actuator_passive_vel(
Ta = dynprm[4]
R *= 1.0 + alpha * (T + Ta - T0)
bias = -K * K / wp.max(MJ_MINVAL, R)
else:
bias = 0.0
else:
bias = 0.0
bias += -K * K / wp.max(MJ_MINVAL, R)
if bias == 0.0 and gain == 0.0:
vel_out[worldid, actid] = 0.0
@@ -151,15 +183,15 @@ def _qderiv_actuator_passive_actuation_dense(
actuator_moment_in: wp.array2d[float],
# In:
vel_in: wp.array2d[float],
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
qderiv_contrib = float(0.0)
for actid in range(nu):
vel = vel_in[worldid, actid]
@@ -195,8 +227,7 @@ def _qderiv_actuator_passive_actuation_dense(
@wp.kernel
def _qderiv_actuator_passive_actuation_sparse(
# Model:
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
M_elemid: wp.array2d[int],
# Data in:
moment_rownnz_in: wp.array2d[int],
moment_rowadr_in: wp.array2d[int],
@@ -204,7 +235,6 @@ def _qderiv_actuator_passive_actuation_sparse(
actuator_moment_in: wp.array2d[float],
# In:
vel_in: wp.array2d[float],
qMj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
@@ -231,19 +261,10 @@ def _qderiv_actuator_passive_actuation_sparse(
continue
dofj = moment_colind_in[worldid, rowadrj]
contrib = moment_i * moment_j * vel
# Search the corresponding elemid
# TODO: This could be precalculated for improved performance
row = dofi
col = dofj
row_startk = M_rowadr[row] - 1
row_nnz = M_rownnz[row]
for k in range(row_nnz):
row_startk += 1
if qMj[row_startk] == col:
wp.atomic_add(qDeriv_out[worldid, 0], row_startk, contrib)
break
elemid = M_elemid[dofi, dofj]
if elemid >= 0:
contrib = moment_i * moment_j * vel
wp.atomic_add(qDeriv_out[worldid, 0], elemid, contrib)
@wp.kernel
@@ -254,23 +275,29 @@ def _qderiv_actuator_passive(
dof_damping: wp.array2d[float],
dof_dampingpoly: wp.array2d[wp.vec2],
is_sparse: bool,
M_elemid: wp.array2d[int],
# Data in:
qvel_in: wp.array2d[float],
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
qDeriv_in: wp.array3d[float],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
madr = M_elemid[dofiid, dofjid]
if is_sparse:
qderiv = qDeriv_in[worldid, 0, elemid]
if madr >= 0:
qderiv = qDeriv_in[worldid, 0, madr]
else:
qderiv = 0.0
else:
qderiv = qDeriv_in[worldid, dofiid, dofjid]
@@ -283,12 +310,13 @@ def _qderiv_actuator_passive(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
if is_sparse:
qDeriv_out[worldid, 0, elemid] = qM_in[worldid, 0, elemid] - qderiv
if madr >= 0:
qDeriv_out[worldid, 0, madr] = M_in[worldid, 0, madr] - qderiv
else:
qM = qM_in[worldid, dofiid, dofjid] - qderiv
qDeriv_out[worldid, dofiid, dofjid] = qM
M = M_in[worldid, dofiid, dofjid] - qderiv
qDeriv_out[worldid, dofiid, dofjid] = M
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] = qM
qDeriv_out[worldid, dofjid, dofiid] = M
# TODO(team): improve performance with tile operations?
@@ -303,18 +331,19 @@ def _qderiv_tendon_damping(
tendon_damping: wp.array2d[float],
tendon_dampingpoly: wp.array2d[wp.vec2],
is_sparse: bool,
M_elemid: wp.array2d[int],
# Data in:
ten_J_in: wp.array2d[float],
ten_velocity_in: wp.array2d[float],
# In:
qMi: wp.array[int],
qMj: wp.array[int],
Mi: wp.array[int],
Mj: wp.array[int],
# Out:
qDeriv_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
dofiid = qMi[elemid]
dofjid = qMj[elemid]
dofiid = Mi[elemid]
dofjid = Mj[elemid]
qderiv = float(0.0)
tendon_damping_id = worldid % tendon_damping.shape[0]
@@ -343,14 +372,283 @@ def _qderiv_tendon_damping(
qderiv *= opt_timestep[worldid % opt_timestep.shape[0]]
madr = M_elemid[dofiid, dofjid]
if is_sparse:
qDeriv_out[worldid, 0, elemid] -= qderiv
if madr >= 0:
qDeriv_out[worldid, 0, madr] -= qderiv
else:
qDeriv_out[worldid, dofiid, dofjid] -= qderiv
if dofiid != dofjid:
qDeriv_out[worldid, dofjid, dofiid] -= qderiv
@wp.kernel
def deriv_rne_cvel_cdof_dot(
# Model:
body_parentid: wp.array[int],
body_jntnum: wp.array[int],
body_jntadr: wp.array[int],
body_dofadr: wp.array[int],
jnt_type: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
# In:
body_tree_: wp.array[int],
# Out:
Dcvel_out: wp.array3d[wp.spatial_vector],
Dcdof_dot_out: wp.array3d[wp.spatial_vector],
):
"""Forward pass: compute d(cvel)/d(qvel_k) and d(cdof_dot)/d(qvel_k).
Mirrors the accumulation order of comvel for each joint type.
Dcdof_dot for rotation DOFs of free joints (dofid+0..2) is zero because the
forward pass sets cdof_dot[dofid+0..2] = 0. The Dcdof_dot array is
zero-initialized so no explicit write is needed.
"""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
dofadr = body_dofadr[bodyid]
jntid = body_jntadr[bodyid]
jntnum = body_jntnum[bodyid]
pid = body_parentid[bodyid]
cdof = cdof_in[worldid]
# Initialize from parent
cvel_k = Dcvel_out[worldid, pid, dofid]
if jntnum == 0:
Dcvel_out[worldid, bodyid, dofid] = cvel_k
return
dof_i = dofadr
for j in range(jntid, jntid + jntnum):
jnttype = jnt_type[j]
if jnttype == 0: # FREE
# rotation DOFs (dof_i+0..2) contribute to cvel
if dofid >= dof_i and dofid < dof_i + 3:
cvel_k += cdof[dofid]
# cdof_dot for rotation DOFs is zero (set in forward kinematics),
# so Dcdof_dot for rotation DOFs is zero (from wp.zeros init)
# derivative of cdof_dot for translation DOFs 3,4,5
Dcdof_dot_out[worldid, dof_i + 3, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 3])
Dcdof_dot_out[worldid, dof_i + 4, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 4])
Dcdof_dot_out[worldid, dof_i + 5, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 5])
# translation DOFs (dof_i+3..5) contribute to cvel
if dofid >= dof_i + 3 and dofid < dof_i + 6:
cvel_k += cdof[dofid]
dof_i += 6
elif jnttype == 1: # BALL
Dcdof_dot_out[worldid, dof_i + 0, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 0])
Dcdof_dot_out[worldid, dof_i + 1, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 1])
Dcdof_dot_out[worldid, dof_i + 2, dofid] = math.motion_cross(cvel_k, cdof[dof_i + 2])
if dofid >= dof_i and dofid < dof_i + 3:
cvel_k += cdof[dofid]
dof_i += 3
else: # HINGE or SLIDE
Dcdof_dot_out[worldid, dof_i, dofid] = math.motion_cross(cvel_k, cdof[dof_i])
if dofid == dof_i:
cvel_k += cdof[dof_i]
dof_i += 1
Dcvel_out[worldid, bodyid, dofid] = cvel_k
@wp.kernel
def deriv_rne_cacc_cfrcbody_forward(
# Model:
body_parentid: wp.array[int],
body_dofnum: wp.array[int],
body_dofadr: wp.array[int],
# Data in:
qvel_in: wp.array2d[float],
cinert_in: wp.array2d[vec10],
cvel_in: wp.array2d[wp.spatial_vector],
cdof_dot_in: wp.array2d[wp.spatial_vector],
# In:
body_tree_: wp.array[int],
Dcvel_in: wp.array3d[wp.spatial_vector],
Dcdof_dot_in: wp.array3d[wp.spatial_vector],
# Out:
Dcacc_out: wp.array3d[wp.spatial_vector],
Dcfrcbody_out: wp.array3d[wp.spatial_vector],
):
"""Forward pass: compute d(cacc)/d(qvel_k) and d(cfrc_body)/d(qvel_k)."""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
dofadr = body_dofadr[bodyid]
dofnum = body_dofnum[bodyid]
pid = body_parentid[bodyid]
qvel = qvel_in[worldid]
dcacc = Dcacc_out[worldid, pid, dofid]
for j in range(dofadr, dofadr + dofnum):
# Term 1: d(cdof_dot * qvel)/d(qvel_k) when j == dofid
if j == dofid:
dcacc += cdof_dot_in[worldid, j]
# Term 2: cdof_dot depends on cvel which depends on qvel_k
dcdofdot = Dcdof_dot_in[worldid, j, dofid]
dcacc += dcdofdot * qvel[j]
Dcacc_out[worldid, bodyid, dofid] = dcacc
# d(cfrc_body)/d(qvel_k)
cinert = cinert_in[worldid, bodyid]
cvel = cvel_in[worldid, bodyid]
dcvel = Dcvel_in[worldid, bodyid, dofid]
# term1 = cinert * d(cacc)/d(qvel_k)
term1 = math.inert_vec(cinert, dcacc)
# term2 = d(cvel x* (cinert * cvel))/d(qvel_k)
cinert_cvel = math.inert_vec(cinert, cvel)
cinert_dcvel = math.inert_vec(cinert, dcvel)
term2 = math.motion_cross_force(dcvel, cinert_cvel) + math.motion_cross_force(cvel, cinert_dcvel)
Dcfrcbody_out[worldid, bodyid, dofid] = term1 + term2
@wp.kernel
def deriv_rne_cfrcbody_backward(
# Model:
body_parentid: wp.array[int],
# In:
body_tree_: wp.array[int],
# Out:
Dcfrcbody_out: wp.array3d[wp.spatial_vector],
):
"""Backward pass: accumulate d(cfrc_body) from children to parents."""
worldid, nodeid, dofid = wp.tid()
bodyid = body_tree_[nodeid]
pid = body_parentid[bodyid]
# body_tree never contains bodyid=0 (worldbody), so pid >= 0 is always valid.
# Siblings at the same level may share a parent; atomic_add handles this.
val = Dcfrcbody_out[worldid, bodyid, dofid]
wp.atomic_add(Dcfrcbody_out[worldid, pid], dofid, val)
@wp.kernel
def deriv_rne_body2jnt_sparse(
# Model:
dof_bodyid: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
# In:
timestep: wp.array[float],
Di: wp.array[int],
Dj: wp.array[int],
Dcfrcbody_in: wp.array3d[wp.spatial_vector],
flg_subtract: bool,
# Out:
qDeriv_out: wp.array3d[float],
):
"""Project body-space RNE derivatives into joint-space qDeriv (sparse)."""
worldid, elemid = wp.tid()
dt = timestep[worldid % timestep.shape[0]]
i = Di[elemid]
j = Dj[elemid]
body_i = dof_bodyid[i]
dcfrc = Dcfrcbody_in[worldid, body_i, j]
term = wp.dot(cdof_in[worldid, i], dcfrc)
if flg_subtract:
wp.atomic_sub(qDeriv_out[worldid, 0], elemid, dt * term)
else:
wp.atomic_add(qDeriv_out[worldid, 0], elemid, dt * term)
def deriv_rne_vel(m: Model, d: Data, out: wp.array3d[float], flg_subtract: bool = False):
"""Compute RNE velocity derivatives and add/subtract from the output.
Implements the analytical derivative of inverse-dynamics Coriolis/centrifugal
forces with respect to joint velocities.
Args:
m: The model (device).
d: The data (device).
out: D-structure output array (nworld, 1, nD) to accumulate RNE terms into.
flg_subtract: If True, subtract the RNE derivatives from output instead of adding them.
"""
# TODO(team): consider caching these allocations
Dcvel = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
Dcdof_dot = wp.zeros((d.nworld, m.nv, m.nv), dtype=wp.spatial_vector)
Dcacc = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
Dcfrcbody = wp.zeros((d.nworld, m.nbody, m.nv), dtype=wp.spatial_vector)
# Forward pass 1: compute Dcvel and Dcdof_dot
for body_tree in m.body_tree:
wp.launch(
deriv_rne_cvel_cdof_dot,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[
m.body_parentid,
m.body_jntnum,
m.body_jntadr,
m.body_dofadr,
m.jnt_type,
d.cdof,
body_tree,
],
outputs=[Dcvel, Dcdof_dot],
)
# Forward pass 2: compute Dcacc and Dcfrcbody
for body_tree in m.body_tree:
wp.launch(
deriv_rne_cacc_cfrcbody_forward,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[
m.body_parentid,
m.body_dofnum,
m.body_dofadr,
d.qvel,
d.cinert,
d.cvel,
d.cdof_dot,
body_tree,
Dcvel,
Dcdof_dot,
],
outputs=[Dcacc, Dcfrcbody],
)
# Backward pass: accumulate Dcfrcbody from children to parents
for body_tree in reversed(m.body_tree):
wp.launch(
deriv_rne_cfrcbody_backward,
dim=(d.nworld, body_tree.size, m.nv),
inputs=[m.body_parentid, body_tree],
outputs=[Dcfrcbody],
)
# Project body-space derivatives into joint-space qDeriv (always sparse D-structure)
wp.launch(
deriv_rne_body2jnt_sparse,
dim=(d.nworld, m.qD_fullm_i.size),
inputs=[m.dof_bodyid, d.cdof, m.opt.timestep, m.qD_fullm_i, m.qD_fullm_j, Dcfrcbody, flg_subtract],
outputs=[out],
)
@event_scope
def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
"""Analytical derivative of smooth forces w.r.t. velocities.
@@ -358,12 +656,10 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
Args:
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state and output arrays (device).
out: qM - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
out: M - dt * qDeriv (derivatives of smooth forces w.r.t velocities).
"""
qMi = m.qM_fullm_i
qMj = m.qM_fullm_j
# TODO(team): implicit requires different sparsity structure
Mi = m.M_fullm_i
Mj = m.M_fullm_j
if ~(m.opt.disableflags & (DisableBit.ACTUATION | DisableBit.DAMPER)):
# TODO(team): only clear elements not set by _qderiv_actuator_passive
@@ -399,41 +695,49 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
wp.launch(
_qderiv_actuator_passive_actuation_sparse,
dim=(d.nworld, m.nu),
inputs=[m.M_rownnz, m.M_rowadr, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMj],
inputs=[
m.M_elemid,
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind,
d.actuator_moment,
vel,
],
outputs=[out],
)
else:
wp.launch(
_qderiv_actuator_passive_actuation_dense,
dim=(d.nworld, qMi.size),
inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, qMi, qMj],
dim=(d.nworld, Mi.size),
inputs=[m.nu, d.moment_rownnz, d.moment_rowadr, d.moment_colind, d.actuator_moment, vel, Mi, Mj],
outputs=[out],
)
wp.launch(
_qderiv_actuator_passive,
dim=(d.nworld, qMi.size),
dim=(d.nworld, Mi.size),
inputs=[
m.opt.timestep,
m.opt.disableflags,
m.dof_damping,
m.dof_dampingpoly,
m.is_sparse,
m.M_elemid,
d.qvel,
d.qM,
qMi,
qMj,
d.M,
Mi,
Mj,
out,
],
outputs=[out],
)
else:
# TODO(team): directly utilize qM for these settings
wp.copy(out, d.qM)
# TODO(team): directly utilize M for these settings
wp.copy(out, d.M)
if not (m.opt.disableflags & DisableBit.DAMPER):
wp.launch(
_qderiv_tendon_damping,
dim=(d.nworld, qMi.size),
dim=(d.nworld, Mi.size),
inputs=[
m.ntendon,
m.opt.timestep,
@@ -443,12 +747,11 @@ def deriv_smooth_vel(m: Model, d: Data, out: wp.array2d[float]):
m.tendon_damping,
m.tendon_dampingpoly,
m.is_sparse,
m.M_elemid,
d.ten_J,
d.ten_velocity,
qMi,
qMj,
Mi,
Mj,
],
outputs=[out],
)
# TODO(team): rne derivative
mjx/mujoco/mjx/third_party/mujoco_warp/_src/forward.py
Vendored
+82 -25
View File
@@ -20,12 +20,14 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import collision_driver
from mujoco.mjx.third_party.mujoco_warp._src import constraint
from mujoco.mjx.third_party.mujoco_warp._src import derivative
from mujoco.mjx.third_party.mujoco_warp._src import history
from mujoco.mjx.third_party.mujoco_warp._src import island
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import passive
from mujoco.mjx.third_party.mujoco_warp._src import sensor
from mujoco.mjx.third_party.mujoco_warp._src import smooth
from mujoco.mjx.third_party.mujoco_warp._src import solver
from mujoco.mjx.third_party.mujoco_warp._src import types
from mujoco.mjx.third_party.mujoco_warp._src import util_misc
from mujoco.mjx.third_party.mujoco_warp._src.support import next_act
from mujoco.mjx.third_party.mujoco_warp._src.support import xfrc_accumulate
@@ -302,6 +304,9 @@ def _advance(m: Model, d: Data, qacc: wp.array, qvel: Optional[wp.array] = None)
outputs=[d.qpos],
)
# advance history buffers before time advance
history.insert_ctrl_history(m, d)
wp.launch(
_next_time,
dim=d.nworld,
@@ -346,17 +351,18 @@ def _compute_damping_deriv(
def _euler_damp_qfrc_sparse(
# Model:
opt_timestep: wp.array[float],
dof_Madr: wp.array[int],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# In:
damp_deriv: wp.array2d[float],
# Out:
qM_integration_out: wp.array3d[float],
M_integration_out: wp.array3d[float],
):
worldid, tid = wp.tid()
timestep = opt_timestep[worldid % opt_timestep.shape[0]]
adr = dof_Madr[tid]
qM_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
adr = M_rowadr[tid] + M_rownnz[tid] - 1
M_integration_out[worldid, 0, adr] += timestep * damp_deriv[worldid, tid]
@cache_kernel
@@ -366,7 +372,7 @@ def _tile_euler_dense(tile: TileSet):
# Model:
opt_timestep: wp.array[float],
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
efc_Ma_in: wp.array2d[float],
# In:
damp_deriv: wp.array2d[float],
@@ -379,14 +385,14 @@ def _tile_euler_dense(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr_in[nodeid]
M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
damping_tile = wp.tile_load(damp_deriv[worldid], shape=(TILE_SIZE,), offset=(dofid,))
damping_scaled = damping_tile * timestep
qm_integration_tile = wp.tile_diag_add(M_tile, damping_scaled)
Ma_tile = wp.tile_load(efc_Ma_in[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_cholesky(qm_integration_tile)
qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile)
L_tile = wp.tile_cholesky(qm_integration_tile, fill_mode="upper")
qacc_tile = wp.tile_cholesky_solve(L_tile, Ma_tile, fill_mode="upper")
wp.tile_store(qacc_out[worldid], qacc_tile, offset=(dofid))
return euler_dense
@@ -409,22 +415,22 @@ def euler(m: Model, d: Data):
)
if m.is_sparse:
qM = wp.clone(d.qM)
M = wp.clone(d.M)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
wp.launch(
_euler_damp_qfrc_sparse,
dim=(d.nworld, m.nv),
inputs=[m.opt.timestep, m.dof_Madr, damp_deriv],
outputs=[qM],
inputs=[m.opt.timestep, m.M_rownnz, m.M_rowadr, damp_deriv],
outputs=[M],
)
smooth.factor_solve_i(m, d, qM, qLD, qLDiagInv, qacc, d.efc.Ma)
smooth.factor_solve_i(m, d, M, qLD, qLDiagInv, qacc, d.efc.Ma)
else:
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_euler_dense(tile),
dim=(d.nworld, tile.adr.size),
inputs=[m.opt.timestep, d.qM, d.efc.Ma, damp_deriv, tile.adr],
inputs=[m.opt.timestep, d.M, d.efc.Ma, damp_deriv, tile.adr],
outputs=[qacc],
block_dim=m.block_dim.euler_dense,
)
@@ -573,16 +579,62 @@ def rungekutta4(m: Model, d: Data):
_advance(m, d, qacc_rk, qvel_rk)
@wp.kernel
def _map_m2d(
# Model:
mapM2D: wp.array[int],
is_sparse: bool,
# In:
qDi: wp.array[int],
qDj: wp.array[int],
qH_M: wp.array3d[float],
# Data out:
qLU_out: wp.array3d[float],
):
worldid, elemid = wp.tid()
if is_sparse:
m_idx = mapM2D[elemid]
if m_idx >= 0:
qLU_out[worldid, 0, elemid] = qH_M[worldid, 0, m_idx]
else:
qLU_out[worldid, 0, elemid] = 0.0
else:
i = qDi[elemid]
j = qDj[elemid]
qLU_out[worldid, 0, elemid] = qH_M[worldid, i, j]
@event_scope
def implicit(m: Model, d: Data):
"""Integrates fully implicit in velocity."""
if ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
if m.opt.integrator == IntegratorType.IMPLICIT:
qH_M = wp.empty(d.M.shape, dtype=float)
# 1. Compute M - dt * qDeriv_smooth in M-structure
derivative.deriv_smooth_vel(m, d, qH_M)
# 2. Map M-structure to D-structure
wp.launch(
_map_m2d,
dim=(d.nworld, m.nD),
inputs=[m.mapM2D, m.is_sparse, m.qD_fullm_i, m.qD_fullm_j, qH_M],
outputs=[d.qLU],
)
# 3. Compute RNE derivatives, scale by timestep, and subtract in-place from qLU
derivative.deriv_rne_vel(m, d, d.qLU, flg_subtract=True)
# 4. Factorize and solve: qacc = qLU \ Ma
qacc = wp.empty((d.nworld, m.nv), dtype=float)
smooth.factor_solve_lu(m, d, d.qLU, qacc, d.efc.Ma)
_advance(m, d, qacc)
elif ~(m.opt.disableflags | ~(DisableBit.ACTUATION | DisableBit.SPRING | DisableBit.DAMPER)):
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
qLD = wp.empty((d.nworld, 1, m.nC), dtype=float)
else:
qDeriv = wp.empty(d.qM.shape, dtype=float)
qLD = wp.empty(d.qM.shape, dtype=float)
qDeriv = wp.empty(d.M.shape, dtype=float)
qLD = wp.empty(d.M.shape, dtype=float)
qLDiagInv = wp.empty((d.nworld, m.nv), dtype=float)
derivative.deriv_smooth_vel(m, d, qDeriv)
qacc = wp.empty((d.nworld, m.nv), dtype=float)
@@ -613,8 +665,7 @@ def fwd_position(m: Model, d: Data, factorize: bool = True):
if m.opt.run_collision_detection:
collision_driver.collision(m, d)
constraint.make_constraint(m, d)
# TODO(team): remove False after island features are more complete
if False and not (m.opt.disableflags & DisableBit.ISLAND):
if m.ntree > 1 and not (m.opt.disableflags & types.DisableBit.ISLAND):
island.island(m, d)
smooth.transmission(m, d)
@@ -1098,6 +1149,13 @@ def fwd_actuation(m: Model, d: Data):
d.actuator_force.zero_()
return
# read delayed ctrl (or direct copy if no delay)
if m.nhistory > 0:
ctrl = wp.empty((d.nworld, m.nu), dtype=float)
history.read_ctrl_delayed(m, d, ctrl)
else:
ctrl = d.ctrl
wp.launch(
_actuator_force,
dim=(d.nworld, m.nu),
@@ -1122,7 +1180,7 @@ def fwd_actuation(m: Model, d: Data):
m.actuator_acc0,
m.actuator_lengthrange,
d.act,
d.ctrl,
ctrl,
d.actuator_length,
d.actuator_velocity,
m.opt.disableflags & DisableBit.CLAMPCTRL,
@@ -1221,7 +1279,7 @@ def fwd_acceleration(m: Model, d: Data, factorize: bool = False):
xfrc_accumulate(m, d, d.qfrc_smooth)
if factorize:
smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, d.qacc_smooth, d.qfrc_smooth)
else:
smooth.solve_m(m, d, d.qacc_smooth, d.qfrc_smooth)
@@ -1266,7 +1324,7 @@ def step(m: Model, d: Data):
euler(m, d)
elif m.opt.integrator == IntegratorType.RK4:
rungekutta4(m, d)
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
elif m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
implicit(m, d)
else:
raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
@@ -1307,8 +1365,7 @@ def step2(m: Model, d: Data):
sensor.sensor_acc(m, d)
# integrate with Euler or implicitfast
# TODO(team): implicit
if m.opt.integrator == IntegratorType.IMPLICITFAST:
if m.opt.integrator in (IntegratorType.IMPLICITFAST, IntegratorType.IMPLICIT):
implicit(m, d)
else:
# note: RK4 defaults to Euler
mjx/mujoco/mjx/third_party/mujoco_warp/_src/history.py
Vendored
+925
View File
@@ -0,0 +1,925 @@
# Copyright 2026 The Newton Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import Data
from mujoco.mjx.third_party.mujoco_warp._src.types import Model
wp.set_module_options({"enable_backward": False})
@wp.func
def _history_physical_index(cursor: int, n: int, logical: int) -> int:
"""Convert logical index (0=oldest, n-1=newest) to physical index."""
return (cursor + 1 + logical) % n
@wp.func
def _history_find_index(
# In:
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
cursor: int,
t: float,
) -> int:
"""Find logical index i such that times[i-1] < t <= times[i].
Returns 0 if t <= times[oldest], n if t > times[newest].
Uses circular binary search matching MuJoCo C historyFindIndex.
"""
times_offset = buf_offset + 2
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# before or at first element
if t <= t_oldest:
return 0
# after last element
if t > t_newest:
return n
# circular binary search: find smallest logical i such that times[phys(i)] >= t
lo = int(0)
hi = int(n - 1)
while hi - lo > 1:
mid = int((lo + hi) >> 1)
mid_phys = _history_physical_index(cursor, n, mid)
if buf[worldid, times_offset + mid_phys] < t:
lo = mid
else:
hi = mid
return hi
@wp.func
def _history_read_scalar(
# In:
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
t: float,
interp: int,
) -> float:
"""Read a scalar value from history buffer at time t.
interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
"""
cursor = int(buf[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# extrapolate before oldest
if t <= t_oldest + 1e-6:
return buf[worldid, values_offset + oldest_phys]
# extrapolate after newest
if t >= t_newest - 1e-6:
return buf[worldid, values_offset + newest_phys]
# find bracketing index
i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
phys_i = _history_physical_index(cursor, n, i)
# exact match
if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
return buf[worldid, values_offset + phys_i]
phys_lo = _history_physical_index(cursor, n, i - 1)
phys_hi = phys_i
# zero-order hold
if interp == 0:
return buf[worldid, values_offset + phys_lo]
dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
v_lo = buf[worldid, values_offset + phys_lo]
v_hi = buf[worldid, values_offset + phys_hi]
# linear interpolation
if interp == 1:
return v_lo + alpha * (v_hi - v_lo)
# cubic spline interpolation (Catmull-Rom)
alpha2 = alpha * alpha
alpha3 = alpha2 * alpha
h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
h10 = alpha3 - 2.0 * alpha2 + alpha
h01 = -2.0 * alpha3 + 3.0 * alpha2
h11 = alpha3 - alpha2
# finite-differenced Catmull-Rom slopes, 0 at endpoints
m_lo = 0.0
if i > 1:
phys_lo_prev = _history_physical_index(cursor, n, i - 2)
dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev]) / dt_lo
m_hi = 0.0
if i < n - 1:
phys_hi_next = _history_physical_index(cursor, n, i + 1)
dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
m_hi = (buf[worldid, values_offset + phys_hi_next] - v_lo) / dt_hi
return h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
@wp.func
def _history_read_vector(
# In:
adr: int,
buf: wp.array2d[float],
worldid: int,
buf_offset: int,
n: int,
dim: int,
t: float,
interp: int,
# Data out:
sensordata_out: wp.array2d[float],
) -> int:
"""Read a vector value from history buffer at time t into sensordata.
Returns 1 on success (value written to sensordata).
interp: 0=zero-order-hold, 1=linear, 2=cubic (Catmull-Rom spline)
"""
cursor = int(buf[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
oldest_phys = _history_physical_index(cursor, n, 0)
newest_phys = _history_physical_index(cursor, n, n - 1)
t_oldest = buf[worldid, times_offset + oldest_phys]
t_newest = buf[worldid, times_offset + newest_phys]
# extrapolate before oldest: copy oldest
if t <= t_oldest + 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + oldest_phys * dim + d]
return 1
# extrapolate after newest: copy newest
if t >= t_newest - 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + newest_phys * dim + d]
return 1
# find bracketing index
i = _history_find_index(buf, worldid, buf_offset, n, cursor, t)
phys_i = _history_physical_index(cursor, n, i)
# exact match
if wp.abs(t - buf[worldid, times_offset + phys_i]) < 1e-6:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_i * dim + d]
return 1
phys_lo = _history_physical_index(cursor, n, i - 1)
phys_hi = phys_i
# zero-order hold
if interp == 0:
for d in range(dim):
sensordata_out[worldid, adr + d] = buf[worldid, values_offset + phys_lo * dim + d]
return 1
dt = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo]
alpha = (t - buf[worldid, times_offset + phys_lo]) / dt
# linear interpolation
if interp == 1:
for d in range(dim):
v_lo = buf[worldid, values_offset + phys_lo * dim + d]
v_hi = buf[worldid, values_offset + phys_hi * dim + d]
sensordata_out[worldid, adr + d] = v_lo + alpha * (v_hi - v_lo)
return 1
# cubic spline interpolation (Catmull-Rom)
alpha2 = alpha * alpha
alpha3 = alpha2 * alpha
h00 = 2.0 * alpha3 - 3.0 * alpha2 + 1.0
h10 = alpha3 - 2.0 * alpha2 + alpha
h01 = -2.0 * alpha3 + 3.0 * alpha2
h11 = alpha3 - alpha2
for d in range(dim):
v_lo = buf[worldid, values_offset + phys_lo * dim + d]
v_hi = buf[worldid, values_offset + phys_hi * dim + d]
# finite-differenced Catmull-Rom slopes, 0 at endpoints
m_lo = 0.0
if i > 1:
phys_lo_prev = _history_physical_index(cursor, n, i - 2)
dt_lo = buf[worldid, times_offset + phys_hi] - buf[worldid, times_offset + phys_lo_prev]
m_lo = (v_hi - buf[worldid, values_offset + phys_lo_prev * dim + d]) / dt_lo
m_hi = 0.0
if i < n - 1:
phys_hi_next = _history_physical_index(cursor, n, i + 1)
dt_hi = buf[worldid, times_offset + phys_hi_next] - buf[worldid, times_offset + phys_lo]
m_hi = (buf[worldid, values_offset + phys_hi_next * dim + d] - v_lo) / dt_hi
sensordata_out[worldid, adr + d] = h00 * v_lo + h10 * dt * m_lo + h01 * v_hi + h11 * dt * m_hi
return 1
@wp.func
def _history_insert_scalar(
# In:
worldid: int,
buf_offset: int,
n: int,
t: float,
value: float,
# Out:
buf_out: wp.array2d[float],
):
"""Insert a scalar value into history buffer at time t."""
cursor = int(buf_out[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
# exact match
if i < n:
phys_i = _history_physical_index(cursor, n, i)
if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
buf_out[worldid, values_offset + phys_i] = value
return
# older than oldest: replace oldest
if i == 0:
oldest_phys = _history_physical_index(cursor, n, 0)
buf_out[worldid, times_offset + oldest_phys] = t
buf_out[worldid, values_offset + oldest_phys] = value
return
# newer than newest: advance cursor
if i == n:
cursor = (cursor + 1) % n
buf_out[worldid, buf_offset + 1] = float(cursor)
buf_out[worldid, times_offset + cursor] = t
buf_out[worldid, values_offset + cursor] = value
return
# out-of-order: shift [1, i-1] left, insert at i-1
for j in range(i - 1):
src_phys = _history_physical_index(cursor, n, j + 1)
dst_phys = _history_physical_index(cursor, n, j)
buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
buf_out[worldid, values_offset + dst_phys] = buf_out[worldid, values_offset + src_phys]
insert_phys = _history_physical_index(cursor, n, i - 1)
buf_out[worldid, times_offset + insert_phys] = t
buf_out[worldid, values_offset + insert_phys] = value
@wp.func
def _history_insert_vector(
# In:
worldid: int,
buf_offset: int,
n: int,
dim: int,
t: float,
src: wp.array2d[float],
src_adr: int,
# Out:
buf_out: wp.array2d[float],
):
"""Insert a vector value from src[worldid, src_adr:src_adr+dim] into history buffer at time t."""
cursor = int(buf_out[worldid, buf_offset + 1])
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + n
i = _history_find_index(buf_out, worldid, buf_offset, n, cursor, t)
slot_phys = -1
# exact match
if i < n:
phys_i = _history_physical_index(cursor, n, i)
if wp.abs(t - buf_out[worldid, times_offset + phys_i]) < 1e-6:
slot_phys = phys_i
if slot_phys < 0:
if i == 0:
# older than oldest: replace oldest
slot_phys = _history_physical_index(cursor, n, 0)
buf_out[worldid, times_offset + slot_phys] = t
elif i == n:
# newer than newest: advance cursor
cursor = (cursor + 1) % n
buf_out[worldid, buf_offset + 1] = float(cursor)
slot_phys = cursor
buf_out[worldid, times_offset + slot_phys] = t
else:
# out-of-order: shift [1, i-1] left, insert at i-1
for j in range(i - 1):
src_phys = _history_physical_index(cursor, n, j + 1)
dst_phys = _history_physical_index(cursor, n, j)
buf_out[worldid, times_offset + dst_phys] = buf_out[worldid, times_offset + src_phys]
for d in range(dim):
buf_out[worldid, values_offset + dst_phys * dim + d] = buf_out[worldid, values_offset + src_phys * dim + d]
slot_phys = _history_physical_index(cursor, n, i - 1)
buf_out[worldid, times_offset + slot_phys] = t
# copy values
for d in range(dim):
buf_out[worldid, values_offset + slot_phys * dim + d] = src[worldid, src_adr + d]
@wp.kernel
def _read_ctrl_delayed_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
actuator_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
# Data out:
ctrl_out: wp.array2d[float],
):
"""Read delayed ctrl for each actuator."""
worldid, uid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
delay = actuator_delay[uid]
if nsample == 0 or delay == 0.0:
# no delay: direct copy
ctrl_out[worldid, uid] = ctrl_in[worldid, uid]
else:
interp = hist[1]
buf_offset = actuator_historyadr[uid]
t = time_in[worldid] - delay
ctrl_out[worldid, uid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp)
@wp.kernel
def _insert_ctrl_history_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
# Data in:
time_in: wp.array[float],
ctrl_in: wp.array2d[float],
# Data out:
history_out: wp.array2d[float],
):
"""Insert current ctrl into history buffers."""
worldid, uid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
if nsample == 0:
return
buf_offset = actuator_historyadr[uid]
t = time_in[worldid]
value = ctrl_in[worldid, uid]
_history_insert_scalar(worldid, buf_offset, nsample, t, value, history_out)
@wp.kernel
def _insert_sensor_history_stage(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
sensor_interval: wp.array[wp.vec2],
# Data in:
time_in: wp.array[float],
sensordata_in: wp.array2d[float],
# In:
sensor_ids: wp.array[int],
# Data out:
history_out: wp.array2d[float],
):
"""Insert current sensor values into history buffers for specific sensor IDs."""
worldid, idx = wp.tid()
sid = sensor_ids[idx]
hist = sensor_history[sid]
nsample = hist[0]
if nsample == 0:
return
buf_offset = sensor_historyadr[sid]
dim = sensor_dim[sid]
interval_val = sensor_interval[sid]
period = interval_val[0]
t = time_in[worldid]
if period > 0.0:
# interval mode: check if condition is satisfied
time_prev = history_out[worldid, buf_offset] # user slot stores time_prev
if time_prev + period <= t:
# advance time_prev by exact period
history_out[worldid, buf_offset] = time_prev + period
# insert sensor value
_history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
else:
_history_insert_vector(worldid, buf_offset, nsample, dim, t, sensordata_in, sensor_adr[sid], history_out)
@wp.kernel
def _apply_sensor_delay_kernel(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
sensor_interval: wp.array[wp.vec2],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
# In:
sensor_ids: wp.array[int],
# Data out:
sensordata_out: wp.array2d[float],
):
"""Apply delay/interval logic for sensors after computation.
TODO(team): Revisit always-compute decision for computationally expensive sensors
with interval/period (e.g., raytracers)
"""
worldid, idx = wp.tid()
sid = sensor_ids[idx]
hist = sensor_history[sid]
nsample = hist[0]
if nsample <= 0:
return
delay = sensor_delay[sid]
dim = sensor_dim[sid]
interp = hist[1]
buf_offset = sensor_historyadr[sid]
t = time_in[worldid]
if delay > 0.0:
# delay > 0: read delayed value from buffer
_history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t - delay, interp, sensordata_out)
else:
# interval-only (delay == 0, interval > 0): check interval condition
interval_val = sensor_interval[sid]
period = interval_val[0]
if period > 0.0:
time_prev = history_in[worldid, buf_offset] # user slot
if time_prev + period > t:
# interval condition not satisfied: read from buffer
_history_read_vector(sensor_adr[sid], history_in, worldid, buf_offset, nsample, dim, t, interp, sensordata_out)
# else: interval condition satisfied, keep computed value
def read_ctrl_delayed(m: Model, d: Data, ctrl: wp.array2d[float]):
"""Read delayed ctrl values for all actuators."""
if m.nhistory == 0:
wp.copy(ctrl, d.ctrl)
return
wp.launch(
_read_ctrl_delayed_kernel,
dim=(d.nworld, m.nu),
inputs=[
m.actuator_history,
m.actuator_historyadr,
m.actuator_delay,
d.time,
d.history,
d.ctrl,
],
outputs=[ctrl],
)
def insert_ctrl_history(m: Model, d: Data):
"""Insert current ctrl values into history buffers."""
if m.nhistory == 0 or m.nu == 0:
return
wp.launch(
_insert_ctrl_history_kernel,
dim=(d.nworld, m.nu),
inputs=[
m.actuator_history,
m.actuator_historyadr,
d.time,
d.ctrl,
],
outputs=[d.history],
)
def apply_sensor_delay(m: Model, d: Data, sensorid: wp.array[int]):
"""Apply delay/interval logic for given sensors after computation.
Matches MuJoCo C architecture where the delayed read (mj_sensorPos) occurs
before the fresh value insert (mj_advance). We save fresh sensordata,
overwrite with delayed values, then insert the saved fresh values.
"""
if m.nhistory == 0 or sensorid.shape[0] == 0:
return
# Save fresh sensordata before delay overwrite
fresh_sensordata = wp.empty_like(d.sensordata)
wp.copy(fresh_sensordata, d.sensordata)
# Read delayed values from buffer → overwrite sensordata
wp.launch(
_apply_sensor_delay_kernel,
dim=(d.nworld, sensorid.shape[0]),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
m.sensor_interval,
d.time,
d.history,
sensorid,
],
outputs=[d.sensordata],
)
# Insert saved fresh sensor values into history buffers
wp.launch(
_insert_sensor_history_stage,
dim=(d.nworld, sensorid.shape[0]),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
m.sensor_interval,
d.time,
fresh_sensordata,
sensorid,
],
outputs=[d.history],
)
@wp.kernel
def _read_ctrl_kernel(
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
actuator_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
# In:
uid: int,
interp: int,
# Out:
result_out: wp.array[float],
):
"""Read delayed ctrl for 1 actuator across all worlds."""
worldid = wp.tid()
hist = actuator_history[uid]
nsample = hist[0]
if nsample == 0:
result_out[worldid] = ctrl_in[worldid, uid]
else:
interp_val = interp
if interp_val < 0:
interp_val = hist[1]
delay = actuator_delay[uid]
buf_offset = actuator_historyadr[uid]
t = time_in[worldid] - delay
result_out[worldid] = _history_read_scalar(history_in, worldid, buf_offset, nsample, t, interp_val)
def read_ctrl(
m: Model,
d: Data,
ctrlid: int,
time: wp.array[float],
interp: int,
result: wp.array2d[float],
):
"""Read delayed ctrl for 1 actuator across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
ctrlid: actuator index.
time: query time per world (nworld,).
interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
result: output buffer (nworld,).
"""
wp.launch(
_read_ctrl_kernel,
dim=(d.nworld,),
inputs=[
m.actuator_history,
m.actuator_historyadr,
m.actuator_delay,
time,
d.history,
d.ctrl,
ctrlid,
interp,
],
outputs=[result],
)
@wp.kernel
def _read_sensor_kernel(
# Model:
sensor_dim: wp.array[int],
sensor_adr: wp.array[int],
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_delay: wp.array[float],
# Data in:
time_in: wp.array[float],
history_in: wp.array2d[float],
sensordata_in: wp.array2d[float],
# In:
sid: int,
interp: int,
# Out:
result_out: wp.array2d[float],
):
"""Read delayed sensor for 1 sensor across all worlds."""
worldid = wp.tid()
hist = sensor_history[sid]
nsample = hist[0]
dim = sensor_dim[sid]
adr = sensor_adr[sid]
if nsample == 0:
for i in range(dim):
result_out[worldid, i] = sensordata_in[worldid, adr + i]
else:
interp_val = interp
if interp_val < 0:
interp_val = hist[1]
delay = sensor_delay[sid]
buf_offset = sensor_historyadr[sid]
t = time_in[worldid] - delay
_history_read_vector(
0, # write to result_out starting at index 0 (not global sensor adr)
history_in,
worldid,
buf_offset,
nsample,
dim,
t,
interp_val,
result_out,
)
def read_sensor(
m: Model,
d: Data,
sensorid: int,
time: wp.array[float],
interp: int,
result: wp.array2d[float],
):
"""Read delayed sensor for 1 sensor across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
sensorid: sensor index.
time: query time per world (nworld,).
interp: interpolation mode (-1=model default, 0=ZOH, 1=linear, 2=cubic).
result: output buffer (nworld, dim).
"""
wp.launch(
_read_sensor_kernel,
dim=(d.nworld,),
inputs=[
m.sensor_dim,
m.sensor_adr,
m.sensor_history,
m.sensor_historyadr,
m.sensor_delay,
time,
d.history,
d.sensordata,
sensorid,
interp,
],
outputs=[result],
)
@wp.kernel
def _init_ctrl_history_kernel(
# kernel_analyzer: off
# Model:
actuator_history: wp.array[wp.vec2i],
actuator_historyadr: wp.array[int],
# In:
ctrlid: int,
times: wp.array[float],
values: wp.array2d[float],
has_times: int,
# Data out:
history_out: wp.array2d[float],
# kernel_analyzer: on
):
"""Initialize history buffer for 1 actuator across all worlds."""
worldid = wp.tid()
nsample = actuator_history[ctrlid][0]
buf_offset = actuator_historyadr[ctrlid]
# preserve user slot
user = history_out[worldid, buf_offset]
# cursor = 0 (samples in order, newest at index nsample-1)
history_out[worldid, buf_offset + 1] = float(nsample - 1)
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + nsample
for i in range(nsample):
if has_times != 0:
history_out[worldid, times_offset + i] = times[i]
else:
history_out[worldid, times_offset + i] = -MJ_MAXVAL
history_out[worldid, values_offset + i] = values[worldid, i]
# restore user slot
history_out[worldid, buf_offset] = user
def init_ctrl_history(
m: Model,
d: Data,
ctrlid: int,
times: wp.array[float],
values: wp.array2d[float],
):
"""Initialize history buffer for 1 actuator across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
ctrlid: actuator index.
times: timestamps or None (nsample,).
values: ctrl values (nworld, nsample).
Raises:
ValueError: If times are not strictly increasing.
"""
has_times = 0 if times is None else 1
if times is not None:
t_np = times.numpy()
for i in range(len(t_np) - 1):
if t_np[i + 1] - t_np[i] < MJ_MINVAL:
raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
if times is None:
times = wp.empty(0, dtype=float)
wp.launch(
_init_ctrl_history_kernel,
dim=(d.nworld,),
inputs=[
m.actuator_history,
m.actuator_historyadr,
ctrlid,
times,
values,
has_times,
],
outputs=[d.history],
)
# kernel_analyzer: off
@wp.kernel
def _init_sensor_history_kernel(
# Model:
sensor_history: wp.array[wp.vec2i],
sensor_historyadr: wp.array[int],
sensor_dim_arr: wp.array[int],
# In:
sensorid: int,
times: wp.array[float],
values: wp.array2d[float],
phase: wp.array[float],
has_times: int,
# Data out:
history_out: wp.array2d[float],
):
# kernel_analyzer: on
"""Initialize history buffer for 1 sensor across all worlds."""
worldid = wp.tid()
nsample = sensor_history[sensorid][0]
dim = sensor_dim_arr[sensorid]
buf_offset = sensor_historyadr[sensorid]
# set user slot (phase = last computation time for interval sensors)
history_out[worldid, buf_offset] = phase[worldid]
# cursor = 0 (samples in order, newest at index nsample-1)
history_out[worldid, buf_offset + 1] = float(nsample - 1)
times_offset = buf_offset + 2
values_offset = buf_offset + 2 + nsample
for i in range(nsample):
if has_times != 0:
history_out[worldid, times_offset + i] = times[i]
else:
history_out[worldid, times_offset + i] = -MJ_MAXVAL
for j in range(dim):
history_out[worldid, values_offset + i * dim + j] = values[worldid, i * dim + j]
def init_sensor_history(
m: Model,
d: Data,
sensorid: int,
times: wp.array[float],
values: wp.array2d[float],
phase: wp.array[float],
):
"""Initialize history buffer for 1 sensor across all worlds.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
sensorid: sensor index.
times: timestamps or None (nsample,).
values: sensor values (nworld, nsample * dim).
phase: user slot value per world (nworld,).
Raises:
ValueError: If times are not strictly increasing.
"""
has_times = 0 if times is None else 1
if times is not None:
t_np = times.numpy()
for i in range(len(t_np) - 1):
if t_np[i + 1] - t_np[i] < MJ_MINVAL:
raise ValueError(f"times must be strictly increasing, got times[{i}]={t_np[i]} >= times[{i + 1}]={t_np[i + 1]}")
if times is None:
times = wp.empty(0, dtype=float)
wp.launch(
_init_sensor_history_kernel,
dim=(d.nworld,),
inputs=[
m.sensor_history,
m.sensor_historyadr,
m.sensor_dim,
sensorid,
times,
values,
phase,
has_times,
],
outputs=[d.history],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/inverse.py
Vendored
+5 -5
View File
@@ -95,9 +95,9 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
# TODO(team): qacc = d.qacc if (m.dof_damping == 0.0).all()
# set qfrc = (d.qM + m.opt.timestep * diag(m.dof_damping)) * d.qacc
# set qfrc = (d.M + m.opt.timestep * diag(m.dof_damping)) * d.qacc
# d.qM @ d.qacc
# d.M @ d.qacc
support.mul_m(m, d, qfrc, d.qacc)
# qfrc += m.opt.timestep * damp_deriv * d.qacc
@@ -109,16 +109,16 @@ def discrete_acc(m: Model, d: Data, qacc: wp.array2d[float]):
)
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
if m.is_sparse:
qDeriv = wp.empty((d.nworld, 1, m.nM), dtype=float)
qDeriv = wp.empty((d.nworld, 1, m.nC), dtype=float)
else:
qDeriv = wp.empty((d.nworld, m.nv, m.nv), dtype=float)
derivative.deriv_smooth_vel(m, d, qDeriv)
mul_m(m, d, qfrc, d.qacc, M=qDeriv)
smooth.factor_solve_i(m, d, d.qM, d.qLD, d.qLDiagInv, qacc, qfrc)
smooth.factor_solve_i(m, d, d.M, d.qLD, d.qLDiagInv, qacc, qfrc)
else:
raise NotImplementedError(f"integrator {m.opt.integrator} not implemented.")
# solve for qacc: qfrc = d.qM @ d.qacc
# solve for qacc: qfrc = d.M @ d.qacc
smooth.solve_m(m, d, qacc, qfrc)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py
Vendored
+330 -102
View File
@@ -33,6 +33,9 @@ from mujoco.mjx.third_party.mujoco_warp._src.types import TrnType
from mujoco.mjx.third_party.mujoco_warp._src.types import vec10
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
# TODO(team): remove after improving island solver performance
ENABLE_ISLANDS = False
def _is_array_spec(typ) -> bool:
"""Check if a type annotation is an array spec (wp.array instance or bracket annotation)."""
@@ -64,6 +67,61 @@ def _create_array(data: Any, spec, sizes: dict[str, int]) -> wp.array | None:
return array
def _create_constraint(
mjm,
nworld: int,
njmax: int,
njmax_nnz: int,
sizes: dict,
island_enabled: bool,
mjd=None,
) -> types.Constraint:
"""Construct a types.Constraint with standard and island local fields allocated properly."""
efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
sparse = is_sparse(mjm)
for f in dataclasses.fields(types.Constraint):
if f.name == "itype":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
elif f.name == "iid":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
elif f.name == "iJ_rownnz":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
elif f.name == "iJ_rowadr":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0) if sparse else (nworld, 0), dtype=int)
elif f.name == "iJ_colind":
efc_kwargs[f.name] = wp.empty((nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, 0, 0), dtype=int)
elif f.name == "iJ":
efc_kwargs[f.name] = wp.empty(
(nworld, 1, njmax_nnz if island_enabled else 0) if sparse else (nworld, njmax if island_enabled else 0, mjm.nv),
dtype=float,
)
elif f.name == "iD":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "iaref":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "ifrictionloss":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "iforce":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=float)
elif f.name == "istate":
efc_kwargs[f.name] = wp.empty((nworld, njmax if island_enabled else 0), dtype=int)
else:
if f.name in efc_kwargs:
continue
if mjd is not None:
shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
val = np.zeros(shape, dtype=f.type.dtype)
if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
else:
efc_kwargs[f.name] = _create_array(None, f.type, sizes)
return types.Constraint(**efc_kwargs)
def is_sparse(mjm: mujoco.MjModel) -> bool:
if mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_AUTO:
if mjm.nv > 32:
@@ -183,6 +241,12 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
opt_kwargs["impratio_invsqrt"] = 1.0 / np.sqrt(np.maximum(mjm.opt.impratio, mujoco.mjMINVAL))
opt = types.Option(**opt_kwargs)
# islands are disabled by default while performance is being improved
# override by setting io.ENABLE_ISLANDS = True
# TODO(team): remove after improving island solver performance
if not ENABLE_ISLANDS:
opt.disableflags |= types.DisableBit.ISLAND
# C MuJoCo tolerance was chosen for float64 architecture, but we default to float32 on GPU
# adjust the tolerance for lower precision, to avoid the solver spending iterations needlessly
# bouncing around the optimal solution
@@ -271,7 +335,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
jnt_limited_slide_hinge = mjm.jnt_limited & np.isin(mjm.jnt_type, (mujoco.mjtJoint.mjJNT_SLIDE, mujoco.mjtJoint.mjJNT_HINGE))
m.jnt_limited_slide_hinge_adr = np.nonzero(jnt_limited_slide_hinge)[0]
m.jnt_limited_ball_adr = np.nonzero(mjm.jnt_limited & (mjm.jnt_type == mujoco.mjtJoint.mjJNT_BALL))[0]
m.dof_tri_row, m.dof_tri_col = np.tril_indices(mjm.nv)
m.dof_tri_row, m.dof_tri_col = np.triu_indices(mjm.nv)
# precompute body_isdofancestor: which DOFs affect each body
# TODO: Investigate alternative approach such as bitmap
@@ -585,14 +649,14 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
for j in range(mjm.mesh_vertnum[mjm.sensor_objid[i]])
]
# qM_tiles records the block diagonal structure of qM
# M_tiles records the block diagonal structure of M
tile_corners = [i for i in range(mjm.nv) if mjm.dof_parentid[i] == -1]
tiles = {}
for i in range(len(tile_corners)):
tile_beg = tile_corners[i]
tile_end = mjm.nv if i == len(tile_corners) - 1 else tile_corners[i + 1]
tiles.setdefault(tile_end - tile_beg, []).append(tile_beg)
m.qM_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
m.M_tiles = tuple(types.TileSet(adr=wp.array(tiles[sz], dtype=int), size=sz) for sz in sorted(tiles.keys()))
# qLD_updates has dof tree ordering of qLD updates for sparse factor m
qLD_updates, dof_depth = {}, np.zeros(mjm.nv, dtype=int) - 1
@@ -620,41 +684,69 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
m.qLD_all_updates = all_updates_flat if all_updates_flat else [(0, 0, 0)]
m.qLD_level_offsets = level_offsets
# indices for sparse qM_fullm (used in solver)
m.qM_fullm_i, m.qM_fullm_j = [], []
# Indices for sparse M_fullm (used in solver). M_fullm_i/j are built by
# walking dof_parentid for each dof, so for joint types whose internal block
# MuJoCo stores diagonal-only in the compact (M_rownnz, M_rowadr) layout
# (e.g. free joints), the chain-aware layout here has more entries per row
# than the compact layout.
m.M_fullm_i, m.M_fullm_j = [], []
for i in range(mjm.nv):
j = i
while j > -1:
m.qM_fullm_i.append(i)
m.qM_fullm_j.append(j)
m.M_fullm_i.append(i)
m.M_fullm_j.append(j)
j = mjm.dof_parentid[j]
# M_elemid maps (row, col) -> madr index in the native CSR M layout
M_elemid = np.full((mjm.nv, mjm.nv), -1, dtype=np.int32)
for i in range(mjm.nv):
rowadr = mjm.M_rowadr[i]
rownnz = mjm.M_rownnz[i]
for k in range(rownnz):
madr = rowadr + k
col = int(mjm.M_colind[madr])
M_elemid[i, col] = madr
m.M_elemid = M_elemid
upper_j, upper_i = np.triu_indices(mjm.nv)
upper_elemid = M_elemid[upper_i, upper_j]
valid_mask = upper_elemid != -1
m.M_fullm_upper_i = upper_j[valid_mask].tolist()
m.M_fullm_upper_j = upper_i[valid_mask].tolist()
m.M_fullm_upper_elemid = upper_elemid[valid_mask].tolist()
# indices for sparse qD_fullm (used in RNE derivatives)
# D-structure is the full square sparsity pattern (both upper and lower triangle)
m.qD_fullm_i, m.qD_fullm_j = [], []
for i in range(mjm.nv):
rowadr = mjm.D_rowadr[i]
rownnz = mjm.D_rownnz[i]
for k in range(rownnz):
m.qD_fullm_i.append(i)
m.qD_fullm_j.append(int(mjm.D_colind[rowadr + k]))
m.nD = mjm.nD
# Gather-based sparse mul_m: for each row, all (col, madr) including diagonal
row_elements = [[] for _ in range(mjm.nv)]
# Add diagonal
for i in range(mjm.nv):
row_elements[i].append((i, mjm.dof_Madr[i]))
# Add off-diagonals: ancestors (lower) and descendants (upper)
for i in range(mjm.nv):
madr_ij, j = mjm.dof_Madr[i], i
while True:
madr_ij, j = madr_ij + 1, mjm.dof_parentid[j]
if j == -1:
break
row_elements[i].append((j, madr_ij)) # row i gathers M[i,j] * vec[j]
row_elements[j].append((i, madr_ij)) # row j gathers M[j,i] * vec[i]
rowadr = mjm.M_rowadr[i]
rownnz = mjm.M_rownnz[i]
for k in range(rownnz):
madr = rowadr + k
col = int(mjm.M_colind[madr])
row_elements[i].append((col, madr)) # row i gathers M[i,col] * vec[col]
if i != col:
row_elements[col].append((i, madr)) # row col gathers M[i,col] * vec[i]
# Flatten into CSR-like arrays
m.qM_mulm_rowadr = [0]
m.qM_mulm_col = []
m.qM_mulm_madr = []
m.M_mulm_rowadr = [0]
m.M_mulm_col = []
m.M_mulm_madr = []
for i in range(mjm.nv):
for col, madr in row_elements[i]:
m.qM_mulm_col.append(col)
m.qM_mulm_madr.append(madr)
m.qM_mulm_rowadr.append(len(m.qM_mulm_col))
m.M_mulm_col.append(col)
m.M_mulm_madr.append(madr)
m.M_mulm_rowadr.append(len(m.M_mulm_col))
m.flexedge_J_rownnz = mjm.flexedge_J_rownnz
m.flexedge_J_rowadr = mjm.flexedge_J_rowadr
@@ -886,6 +978,42 @@ def _resolve_batch_size(na: int | None, n: int | None, nworld: int, default: int
return default
def _allocate_island_arrays(
mjm: mujoco.MjModel,
d: types.Data,
nworld: int,
njmax: int,
island_enabled: bool,
mjd: mujoco.MjData,
):
ntree_size = mjm.ntree if island_enabled else 0
nv_size = mjm.nv if island_enabled else 0
njmax_size = njmax if island_enabled else 0
d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1 if island_enabled else 0)), dtype=int)
d.dof_island = wp.array(np.tile(mjd.dof_island, (nworld, 1 if island_enabled else 0)), dtype=int)
d.island_dofadr = wp.empty((nworld, ntree_size), dtype=int)
d.island_nv = wp.empty((nworld, ntree_size), dtype=int)
d.island_nefc = wp.empty((nworld, ntree_size), dtype=int)
d.island_ne = wp.empty((nworld, ntree_size), dtype=int)
d.island_nf = wp.empty((nworld, ntree_size), dtype=int)
d.island_efcadr = wp.empty((nworld, ntree_size), dtype=int)
d.nidof = wp.empty((nworld if island_enabled else 0,), dtype=int)
d.map_dof2idof = wp.empty((nworld, nv_size), dtype=int)
d.map_idof2dof = wp.empty((nworld, nv_size), dtype=int)
d.map_efc2iefc = wp.empty((nworld, njmax_size), dtype=int)
d.map_iefc2efc = wp.empty((nworld, njmax_size), dtype=int)
d.dof_islandid = wp.empty((nworld, nv_size), dtype=int)
d.efc_islandid = wp.empty((nworld, njmax_size), dtype=int)
d.iqacc = wp.empty((nworld, nv_size), dtype=float)
d.iqacc_smooth = wp.empty((nworld, nv_size), dtype=float)
d.iqfrc_smooth = wp.empty((nworld, nv_size), dtype=float)
d.iqfrc_constraint = wp.empty((nworld, nv_size), dtype=float)
def make_data(
mjm: mujoco.MjModel,
nworld: int = 1,
@@ -967,7 +1095,8 @@ def make_data(
contact = types.Contact(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Contact)})
contact.efc_address = wp.array(np.full((naconmax, sizes["nmaxpyramid"]), -1, dtype=int), dtype=int)
efc = types.Constraint(**{f.name: _create_array(None, f.type, sizes) for f in dataclasses.fields(types.Constraint)})
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS)
if is_sparse(mjm):
efc.J_rownnz = wp.zeros((nworld, njmax), dtype=int)
@@ -1005,7 +1134,7 @@ def make_data(
"njmax": njmax,
"njmax_pad": sizes["njmax_pad"],
"njmax_nnz": njmax_nnz,
"qM": None,
"M": None,
"qLD": None,
# world body
"xquat": wp.array(np.tile(mjd.xquat, (nworld, 1)), shape=(nworld, mjm.nbody), dtype=wp.quat),
@@ -1024,6 +1153,24 @@ def make_data(
# island arrays
"nisland": None,
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
"island_nf": None,
"island_efcadr": None,
"nidof": None,
"map_dof2idof": None,
"map_idof2dof": None,
"map_efc2iefc": None,
"map_iefc2efc": None,
"dof_islandid": None,
"efc_islandid": None,
"iqacc": None,
"iqacc_smooth": None,
"iqfrc_smooth": None,
"iqfrc_constraint": None,
}
for f in dataclasses.fields(types.Data):
if f.name in d_kwargs:
@@ -1033,15 +1180,13 @@ def make_data(
d = types.Data(**d_kwargs)
if is_sparse(mjm):
d.qM = wp.zeros((nworld, 1, mjm.nM), dtype=float)
d.M = wp.zeros((nworld, 1, mjm.nC), dtype=float)
d.qLD = wp.zeros((nworld, 1, mjm.nC), dtype=float)
else:
d.qM = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
d.M = wp.zeros((nworld, sizes["nv_pad"], sizes["nv_pad"]), dtype=float)
d.qLD = wp.zeros((nworld, mjm.nv, mjm.nv), dtype=float)
# island discovery arrays
d.nisland = wp.zeros((nworld,), dtype=int)
d.tree_island = wp.zeros((nworld, mjm.ntree), dtype=int)
_allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
return d
@@ -1176,16 +1321,7 @@ def put_data(
# create efc
efc_kwargs = {"J_rownnz": None, "J_rowadr": None, "J_colind": None, "J": None}
for f in dataclasses.fields(types.Constraint):
if f.name in efc_kwargs:
continue
shape = tuple(sizes[dim] if isinstance(dim, str) else dim for dim in f.type.shape)
val = np.zeros(shape, dtype=f.type.dtype)
if f.name in ("type", "id", "pos", "margin", "D", "vel", "aref", "frictionloss", "force"):
val[:, : mjd.nefc] = np.tile(getattr(mjd, "efc_" + f.name), (nworld, 1))
efc_kwargs[f.name] = wp.array(val, dtype=f.type.dtype)
efc = types.Constraint(**efc_kwargs)
efc = _create_constraint(mjm, nworld, njmax, njmax_nnz, sizes, ENABLE_ISLANDS, mjd)
if is_sparse(mjm):
J_rownnz = np.zeros(njmax, dtype=np.int32)
@@ -1236,12 +1372,30 @@ def put_data(
"njmax_nnz": njmax_nnz,
# fields set after initialization:
"solver_niter": None,
"qM": None,
"M": None,
"qLD": None,
"nacon": None,
# island arrays
"nisland": None,
"tree_island": None,
"dof_island": None,
"island_dofadr": None,
"island_nv": None,
"island_nefc": None,
"island_ne": None,
"island_nf": None,
"island_efcadr": None,
"nidof": None,
"map_dof2idof": None,
"map_idof2dof": None,
"map_efc2iefc": None,
"map_iefc2efc": None,
"dof_islandid": None,
"efc_islandid": None,
"iqacc": None,
"iqacc_smooth": None,
"iqfrc_smooth": None,
"iqfrc_constraint": None,
}
for f in dataclasses.fields(types.Data):
if f.name in d_kwargs:
@@ -1256,20 +1410,25 @@ def put_data(
d.solver_niter = wp.full((nworld,), mjd.solver_niter[0], dtype=int)
if is_sparse(mjm):
d.qM = wp.array(np.full((nworld, 1, mjm.nM), mjd.qM), dtype=float)
if check_version("mujoco>=3.8.1.dev910242375"):
d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.M), dtype=float)
else:
d.M = wp.array(np.full((nworld, 1, mjm.nC), mjd.qM[mjm.mapM2M]), dtype=float)
d.qLD = wp.array(np.full((nworld, 1, mjm.nC), mjd.qLD), dtype=float)
else:
qM = np.zeros((mjm.nv, mjm.nv))
mujoco.mj_fullM(mjm, qM, mjd.qM)
qLD = np.linalg.cholesky(qM) if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
M = np.zeros((mjm.nv, mjm.nv))
if check_version("mujoco>=3.8.1.dev910242375"):
mujoco.mju_sym2dense(M, mjd.M, mjm.M_rownnz, mjm.M_rowadr, mjm.M_colind)
qLD = np.linalg.cholesky(M).T if (mjd.M != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
else:
mujoco.mj_fullM(mjm, M, mjd.qM)
qLD = np.linalg.cholesky(M).T if (mjd.qM != 0.0).any() and (mjd.qLD != 0.0).any() else np.zeros((mjm.nv, mjm.nv))
padding = sizes["nv_pad"] - mjm.nv
qM_padded = np.pad(qM, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
d.qM = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), qM_padded), dtype=float)
M_padded = np.pad(M, ((0, padding), (0, padding)), mode="constant", constant_values=0.0)
d.M = wp.array(np.full((nworld, sizes["nv_pad"], sizes["nv_pad"]), M_padded), dtype=float)
d.qLD = wp.array(np.full((nworld, mjm.nv, mjm.nv), qLD), dtype=float)
# island arrays
d.nisland = wp.array(np.full(nworld, mjd.nisland), dtype=int)
d.tree_island = wp.array(np.tile(mjd.tree_island, (nworld, 1)), dtype=int)
_allocate_island_arrays(mjm, d, nworld, njmax, ENABLE_ISLANDS, mjd)
d.nacon = wp.array([mjd.ncon * nworld], dtype=int)
@@ -1403,6 +1562,9 @@ def get_data_into(
result.qfrc_constraint[:] = d.qfrc_constraint.numpy()[world_id]
result.qfrc_inverse[:] = d.qfrc_inverse.numpy()[world_id]
if mjm.nhistory > 0:
result.history[:] = d.history.numpy()[world_id]
# contact
result.contact.dist[:ncon] = d.contact.dist.numpy()[ncon_filter]
result.contact.pos[:ncon] = d.contact.pos.numpy()[ncon_filter]
@@ -1417,17 +1579,27 @@ def get_data_into(
result.contact.efc_address[:ncon] = contact_efc_address_ordered[:ncon]
if is_sparse(mjm):
result.qM[:] = d.qM.numpy()[world_id, 0]
if check_version("mujoco>=3.8.1.dev910242375"):
result.M[:] = d.M.numpy()[world_id, 0]
else:
result.qM[mjm.mapM2M] = d.M.numpy()[world_id, 0]
result.qLD[:] = d.qLD.numpy()[world_id, 0]
else:
qM = d.qM.numpy()[world_id]
adr = 0
for i in range(mjm.nv):
j = i
while j >= 0:
result.qM[adr] = qM[i, j]
j = mjm.dof_parentid[j]
adr += 1
M = d.M.numpy()[world_id]
if check_version("mujoco>=3.8.1.dev910242375"):
for i in range(mjm.nv):
adr = mjm.M_rowadr[i]
for k in range(mjm.M_rownnz[i]):
col = mjm.M_colind[adr + k]
result.M[adr + k] = M[i, col]
else:
adr = 0
for i in range(mjm.nv):
j = i
while j >= 0:
result.qM[adr] = M[i, j]
j = mjm.dof_parentid[j]
adr += 1
mujoco.mj_factorM(mjm, result)
if nefc > 0:
@@ -1466,6 +1638,7 @@ def get_data_into(
result.efc_frictionloss[:] = d.efc.frictionloss.numpy()[world_id, efc_idx]
result.efc_state[:] = d.efc.state.numpy()[world_id, efc_idx]
result.efc_force[:] = d.efc.force.numpy()[world_id, efc_idx]
result.efc_island[:] = d.efc.island.numpy()[world_id, efc_idx]
# rne_postconstraint
result.cacc[:] = d.cacc.numpy()[world_id]
@@ -1487,8 +1660,51 @@ def get_data_into(
# islands
nisland = d.nisland.numpy()[world_id]
result.nisland = nisland
if nisland:
if d.tree_island.shape[1] > 0 and nisland:
result.tree_island[:] = d.tree_island.numpy()[world_id]
result.dof_island[:] = d.dof_island.numpy()[world_id]
result.island_dofadr[:nisland] = d.island_dofadr.numpy()[world_id, :nisland]
result.island_nv[:nisland] = d.island_nv.numpy()[world_id, :nisland]
result.island_nefc[:nisland] = d.island_nefc.numpy()[world_id, :nisland]
result.island_ne[:nisland] = d.island_ne.numpy()[world_id, :nisland]
result.island_nf[:nisland] = d.island_nf.numpy()[world_id, :nisland]
result.island_iefcadr[:nisland] = d.island_efcadr.numpy()[world_id, :nisland]
nv = mjm.nv
result.map_dof2idof[:nv] = d.map_dof2idof.numpy()[world_id, :nv]
result.map_idof2dof[:nv] = d.map_idof2dof.numpy()[world_id, :nv]
result.map_efc2iefc[:nefc] = d.map_efc2iefc.numpy()[world_id, :nefc]
result.map_iefc2efc[:nefc] = d.map_iefc2efc.numpy()[world_id, :nefc]
result.iefc_type[:nefc] = d.efc.itype.numpy()[world_id, :nefc]
result.iefc_id[:nefc] = d.efc.iid.numpy()[world_id, :nefc]
result.iefc_D[:nefc] = d.efc.iD.numpy()[world_id, :nefc]
result.iefc_aref[:nefc] = d.efc.iaref.numpy()[world_id, :nefc]
result.iefc_frictionloss[:nefc] = d.efc.ifrictionloss.numpy()[world_id, :nefc]
result.iefc_state[:nefc] = d.efc.istate.numpy()[world_id, :nefc]
result.iefc_force[:nefc] = d.efc.iforce.numpy()[world_id, :nefc]
if is_sparse(mjm):
iefc_J = np.zeros((nefc, mjm.nv))
mujoco.mju_sparse2dense(
iefc_J,
d.efc.iJ.numpy()[world_id, 0],
d.efc.iJ_rownnz.numpy()[world_id, :nefc],
d.efc.iJ_rowadr.numpy()[world_id, :nefc],
d.efc.iJ_colind.numpy()[world_id, 0],
)
else:
iefc_J = d.efc.iJ.numpy()[world_id, :nefc, : mjm.nv]
if mujoco.mj_isSparse(mjm):
mujoco.mju_dense2sparse(
result.iefc_J,
iefc_J,
result.iefc_J_rownnz,
result.iefc_J_rowadr,
result.iefc_J_colind,
)
else:
result.iefc_J[: nefc * mjm.nv] = iefc_J.flatten()
def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
@@ -1511,14 +1727,14 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
xfrc_applied_out[worldid, bodyid][elemid] = 0.0
@wp.kernel(module="unique", enable_backward=False)
def reset_qM(reset_in: wp.array[bool], qM_out: wp.array3d[float]):
def reset_M(reset_in: wp.array[bool], M_out: wp.array3d[float]):
worldid, elemid1, elemid2 = wp.tid()
if wp.static(reset is not None):
if not reset_in[worldid]:
return
qM_out[worldid, elemid1, elemid2] = 0.0
M_out[worldid, elemid1, elemid2] = 0.0
@wp.kernel(module="unique", enable_backward=False)
def reset_nworld(
@@ -1670,10 +1886,10 @@ def reset_data(m: types.Model, d: types.Data, reset: Optional[wp.array] = None):
wp.launch(reset_xfrc_applied, dim=(d.nworld, m.nbody, 6), inputs=[reset_input], outputs=[d.xfrc_applied])
wp.launch(
reset_qM,
dim=(d.nworld, d.qM.shape[1], d.qM.shape[2]),
reset_M,
dim=(d.nworld, d.M.shape[1], d.M.shape[2]),
inputs=[reset_input],
outputs=[d.qM],
outputs=[d.M],
)
# set mocap_pos/quat = body_pos/quat for mocap bodies
@@ -1786,11 +2002,12 @@ def _copy_tendon_length0(
def _compute_meaninertia(
nv: int,
is_sparse: bool,
dof_Madr_in: wp.array[int],
qM_in: wp.array3d[float],
M_rownnz_in: wp.array[int],
M_rowadr_in: wp.array[int],
M_in: wp.array3d[float],
meaninertia_out: wp.array[float],
):
"""Compute mean diagonal inertia from qM at qpos0."""
"""Compute mean diagonal inertia from M at qpos0."""
worldid = wp.tid()
if nv == 0:
@@ -1800,12 +2017,12 @@ def _compute_meaninertia(
total = float(0.0)
for i in range(nv):
if is_sparse:
# Sparse: qM is flattened lower triangular, diagonal at dof_Madr[i]
madr = dof_Madr_in[i]
total += qM_in[worldid, 0, madr]
# Sparse: M is in CSR format, diagonal at M_rowadr_in[i] + M_rownnz_in[i] - 1
madr = M_rowadr_in[i] + M_rownnz_in[i] - 1
total += M_in[worldid, 0, madr]
else:
# Dense: qM is 2D matrix, diagonal at [i,i]
total += qM_in[worldid, i, i]
# Dense: M is 2D matrix, diagonal at [i,i]
total += M_in[worldid, i, i]
meaninertia_out[worldid % meaninertia_out.shape[0]] = total / float(nv)
@@ -2288,11 +2505,11 @@ def set_const_0(m: types.Model, d: types.Data):
smooth.factor_m(m, d)
smooth.transmission(m, d)
# Compute meaninertia from qM diagonal at qpos0
# Compute meaninertia from M diagonal at qpos0
wp.launch(
_compute_meaninertia,
dim=d.nworld,
inputs=[m.nv, m.is_sparse, m.dof_Madr, d.qM],
inputs=[m.nv, m.is_sparse, m.M_rownnz, m.M_rowadr, d.M],
outputs=[m.stat.meaninertia],
)
@@ -2452,27 +2669,31 @@ def set_const(m: types.Model, d: types.Data):
Model fields that can be modified safely with set_const:
Field | Notes
---------------------------------|----------------------------------------------
qpos0, qpos_spring |
body_mass, body_inertia, | Mass and inertia are usually scaled together
body_ipos, body_iquat | since inertia is sum(m * r^2).
body_pos, body_quat | Unsafe for static bodies (invalidates BVH).
body_gravcomp | If changing from 0 to >0 bodies, required.
dof_armature |
eq_data | For connect/weld, offsets computed if not set.
hfield_size |
tendon_stiffness, tendon_damping | Only if changing from/to zero.
actuator_gainprm, actuator_biasprm | For position actuators with dampratio.
================================== ==============================================
Field Notes
================================== ==============================================
qpos0, qpos_spring
body_mass, body_inertia, Mass and inertia are usually scaled together
body_ipos, body_iquat since inertia is sum(m * r^2).
body_pos, body_quat Unsafe for static bodies (invalidates BVH).
body_gravcomp If changing from 0 to >0 bodies, required.
dof_armature
eq_data For connect/weld, offsets computed if not set.
hfield_size
tendon_stiffness, tendon_damping Only if changing from/to zero.
actuator_gainprm, actuator_biasprm For position actuators with dampratio.
================================== ==============================================
For selective updates, use the sub-functions directly based on what changed:
Modified Field | Call
----------------|------------------
body_mass | set_const
body_gravcomp | set_const_fixed
body_inertia | set_const_0
qpos0 | set_const_0
============== ===============
Modified Field Call
============== ===============
body_mass set_const
body_gravcomp set_const_fixed
body_inertia set_const_0
qpos0 set_const_0
============== ===============
Computes:
- Fixed quantities (via set_const_fixed):
@@ -2621,6 +2842,9 @@ def override_model(model: types.Model | mujoco.MjModel, overrides: dict[str, Any
else:
val = typ(val)
if attr == "disableflags" and isinstance(obj, types.Option) and not ENABLE_ISLANDS:
val = int(val) | types.DisableBit.ISLAND
setattr(obj, attr, val)
@@ -2724,11 +2948,13 @@ def create_render_context(
render_seg: list[bool] | bool | None = None,
use_textures: bool = True,
use_shadows: bool = False,
use_ambient_lighting: bool = True,
enabled_geom_groups: list[int] = [0, 1, 2],
cam_active: list[bool] | None = None,
flex_render_smooth: bool = True,
use_precomputed_rays: bool = True,
render_skybox: bool = False,
enable_backface_culling: bool = True,
) -> types.RenderContext:
"""Creates a render context on device.
@@ -2743,6 +2969,8 @@ def create_render_context(
If None, uses the MuJoCo model values.
use_textures: Whether to use textures.
use_shadows: Whether to use shadows.
use_ambient_lighting: Whether to add the renderer's hemispheric ambient
lighting term before applying model lights.
enabled_geom_groups: The geom groups to render.
cam_active: List of booleans indicating which cameras to include in rendering.
If None, all cameras are included.
@@ -2751,6 +2979,10 @@ def create_render_context(
When using domain randomization for camera intrinsics, set to False.
render_skybox: Whether to shade missed rays with the MuJoCo skybox texture.
Requires the model to contain a texture with type `mjTEXTURE_SKYBOX`.
enable_backface_culling: Drop primitive-ray hits whose normal faces away from
the ray (ray origin inside the geom). Matches MuJoCo's
mesh-ray rule. Default True. Disable for a small
performance gain when no camera is ever inside a geom.
Returns:
The render context containing rendering fields and output arrays on device.
@@ -2758,13 +2990,7 @@ def create_render_context(
mjd = mujoco.MjData(mjm)
mujoco.mj_forward(mjm, mjd)
constructor = "sah"
if check_version("warp>=1.13.0.dev20260325"):
# TODO: The cubql constructor and is_cubql_available exist only in
# recent Warp 1.13+ builds, modify this after warp is updated to 1.13+.
_cubql_avail = getattr(wp, "is_cubql_available", None)
if callable(_cubql_avail) and _cubql_avail():
constructor = "cubql"
constructor = "cubql"
# Mesh BVHs – build for all meshes so per-world variants are available
nmesh = mjm.nmesh
@@ -2937,6 +3163,7 @@ def create_render_context(
cam_id_map=wp.array(active_cam_indices, dtype=int),
use_textures=use_textures,
use_shadows=use_shadows,
use_ambient_lighting=use_ambient_lighting,
background_color=render_util.pack_rgba_to_uint32(0.1 * 255.0, 0.1 * 255.0, 0.2 * 255.0, 1.0 * 255.0),
use_precomputed_rays=use_precomputed_rays,
render_skybox=render_skybox,
@@ -2982,6 +3209,7 @@ def create_render_context(
render_seg=wp.array(render_seg, dtype=bool),
znear=znear,
total_rays=int(total),
enable_backface_culling=enable_backface_culling,
)
bvh.build_scene_bvh(mjm, mjd, rc, nworld)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/island.py
Vendored
+821 -20
View File
@@ -18,6 +18,7 @@ import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import types
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
from mujoco.mjx.third_party.mujoco_warp._src.types import EqType
from mujoco.mjx.third_party.mujoco_warp._src.types import IslandSolverContext
from mujoco.mjx.third_party.mujoco_warp._src.types import ObjType
from mujoco.mjx.third_party.mujoco_warp._src.warp_util import event_scope
@@ -35,11 +36,15 @@ def _tree_edges(
eq_obj1id: wp.array[int],
eq_obj2id: wp.array[int],
eq_objtype: wp.array[int],
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
contact_geom_in: wp.array[wp.vec2i],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
efc_J_in: wp.array3d[float],
njmax_in: int,
# Out:
@@ -131,26 +136,39 @@ def _tree_edges(
first_tree = int(-1)
has_cross_edge = int(0)
for dof in range(nv):
# TODO(team): sparse efc_J
# TODO(team): tree dof skip
J_val = efc_J_in[worldid, efcid, dof]
if J_val != 0.0:
tree = dof_treeid[dof]
if tree < 0:
count = nv
rowadr = 0
if is_sparse:
count = efc_J_rownnz_in[worldid, efcid]
rowadr = efc_J_rowadr_in[worldid, efcid]
for i in range(count):
dof = i
if is_sparse:
sparseid = rowadr + i
dof = efc_J_colind_in[worldid, 0, sparseid]
else:
J_val = efc_J_in[worldid, efcid, dof]
if J_val == 0.0:
continue
if first_tree == -1:
first_tree = tree
elif tree != first_tree:
t1 = wp.min(first_tree, tree)
t2 = wp.max(first_tree, tree)
wp.atomic_max(tree_tree, worldid, t1, t2, 1)
has_cross_edge = 1
tree = dof_treeid[dof]
if tree < 0:
continue
if first_tree == -1:
first_tree = tree
elif tree != first_tree:
t1 = wp.min(first_tree, tree)
t2 = wp.max(first_tree, tree)
wp.atomic_max(tree_tree, worldid, t1, t2, 1)
wp.atomic_max(tree_tree, worldid, t2, t1, 1)
has_cross_edge = 1
if first_tree >= 0 and has_cross_edge == 0:
wp.atomic_max(tree_tree, worldid, first_tree, first_tree, 1)
@event_scope
def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
"""Compute tree-tree adjacency matrix."""
tree_tree.zero_()
@@ -168,10 +186,14 @@ def tree_edges(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
m.eq_obj1id,
m.eq_obj2id,
m.eq_objtype,
m.is_sparse,
d.nefc,
d.contact.geom,
d.efc.type,
d.efc.id,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.efc.J,
d.njmax,
],
@@ -242,6 +264,19 @@ def _flood_fill(
nisland_out[worldid] = nisland
@event_scope
def flood_fill(m: types.Model, d: types.Data, tree_tree: wp.array3d[int]):
d.tree_island.fill_(-1)
stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
wp.launch(
_flood_fill,
dim=d.nworld,
inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
outputs=[d.nisland, d.tree_island, stack_scratch],
)
@event_scope
def island(m: types.Model, d: types.Data):
"""Discover constraint islands."""
@@ -254,12 +289,778 @@ def island(m: types.Model, d: types.Data):
tree_edges(m, d, tree_tree)
# Step 2: DFS flood fill
d.tree_island.fill_(-1)
stack_scratch = wp.empty((d.nworld, m.ntree * m.ntree), dtype=int)
flood_fill(m, d, tree_tree)
@wp.kernel
def _island_count_dofs(
dof_treeid: wp.array[int],
tree_island_in: wp.array2d[int],
dof_island_out: wp.array2d[int],
island_nv_out: wp.array2d[int],
):
worldid, dofid = wp.tid()
island_id = tree_island_in[worldid, dof_treeid[dofid]]
dof_island_out[worldid, dofid] = island_id
if island_id >= 0:
wp.atomic_add(island_nv_out, worldid, island_id, 1)
@wp.kernel
def _island_scan_sizes(
nisland_in: wp.array[int],
island_idofadr_out: wp.array2d[int],
island_nv_inout: wp.array2d[int],
island_nefc_inout: wp.array2d[int],
island_iefcadr_out: wp.array2d[int],
nidof_out: wp.array[int],
):
worldid = wp.tid()
nisland = nisland_in[worldid]
if nisland == 0:
nidof_out[worldid] = 0
return
# Scan DOFs and Constraints
island_idofadr_out[worldid, 0] = 0
island_iefcadr_out[worldid, 0] = 0
for i in range(1, nisland):
island_idofadr_out[worldid, i] = island_idofadr_out[worldid, i - 1] + island_nv_inout[worldid, i - 1]
island_iefcadr_out[worldid, i] = island_iefcadr_out[worldid, i - 1] + island_nefc_inout[worldid, i - 1]
nidof = island_idofadr_out[worldid, nisland - 1] + island_nv_inout[worldid, nisland - 1]
nidof_out[worldid] = nidof
# Reset for recount
for i in range(nisland):
island_nv_inout[worldid, i] = 0
island_nefc_inout[worldid, i] = 0
@wp.kernel
def _island_map_dofs(
nv: int,
dof_island_in: wp.array2d[int],
island_idofadr_in: wp.array2d[int],
nidof_in: wp.array[int],
island_nv_inout: wp.array2d[int],
map_dof2idof_out: wp.array2d[int],
map_idof2dof_out: wp.array2d[int],
idof_islandid_out: wp.array2d[int],
unconstrained_cnt_inout: wp.array2d[int],
):
worldid, dofid = wp.tid()
nidof = nidof_in[worldid]
island_id = dof_island_in[worldid, dofid]
if island_id >= 0:
local_idx = wp.atomic_add(island_nv_inout, worldid, island_id, 1)
idof = island_idofadr_in[worldid, island_id] + local_idx
idof_islandid_out[worldid, idof] = island_id
else:
cnt = wp.atomic_add(unconstrained_cnt_inout, worldid, 0, 1)
idof = nidof + cnt
map_dof2idof_out[worldid, dofid] = idof
map_idof2dof_out[worldid, idof] = dofid
@wp.kernel
def _island_count_constraints(
nefc_in: wp.array[int],
njmax_in: int,
efc_tree_in: wp.array2d[int],
tree_island_in: wp.array2d[int],
efc_type_in: wp.array2d[int],
efc_island_out: wp.array2d[int],
island_nefc_out: wp.array2d[int],
island_ne_out: wp.array2d[int],
island_nf_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
efc_tree = efc_tree_in[worldid, efcid]
if efc_tree < 0:
efc_island_out[worldid, efcid] = -1
return
island_id = tree_island_in[worldid, efc_tree]
efc_island_out[worldid, efcid] = island_id
if island_id >= 0:
wp.atomic_add(island_nefc_out, worldid, island_id, 1)
efc_type = efc_type_in[worldid, efcid]
if efc_type == ConstraintType.EQUALITY:
wp.atomic_add(island_ne_out, worldid, island_id, 1)
elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
wp.atomic_add(island_nf_out, worldid, island_id, 1)
@wp.kernel
def _island_map_constraints(
nefc_in: wp.array[int],
njmax_in: int,
efc_island_in: wp.array2d[int],
island_iefcadr_in: wp.array2d[int],
island_ne_in: wp.array2d[int],
island_nf_in: wp.array2d[int],
efc_type_in: wp.array2d[int],
# Counters (inout):
island_ne_mapped_inout: wp.array2d[int],
island_nf_mapped_inout: wp.array2d[int],
island_nother_mapped_inout: wp.array2d[int],
island_nefc_inout: wp.array2d[int],
# Out:
map_efc2iefc_out: wp.array2d[int],
map_iefc2efc_out: wp.array2d[int],
iefc_islandid_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
island_id = efc_island_in[worldid, efcid]
if island_id >= 0:
efc_type = efc_type_in[worldid, efcid]
# 1. Determine local index and absolute index ic based on category
if efc_type == ConstraintType.EQUALITY:
local_idx = wp.atomic_add(island_ne_mapped_inout, worldid, island_id, 1)
ic = island_iefcadr_in[worldid, island_id] + local_idx
elif efc_type == ConstraintType.FRICTION_DOF or efc_type == ConstraintType.FRICTION_TENDON:
local_idx = wp.atomic_add(island_nf_mapped_inout, worldid, island_id, 1)
ic = island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + local_idx
else:
local_idx = wp.atomic_add(island_nother_mapped_inout, worldid, island_id, 1)
ic = (
island_iefcadr_in[worldid, island_id] + island_ne_in[worldid, island_id] + island_nf_in[worldid, island_id] + local_idx
)
# 2. Increment overall island_nefc counter to reconstruct d.island_nefc
wp.atomic_add(island_nefc_inout, worldid, island_id, 1)
# 3. Store mappings
map_efc2iefc_out[worldid, efcid] = ic
map_iefc2efc_out[worldid, ic] = efcid
iefc_islandid_out[worldid, ic] = island_id
@wp.kernel
def _island_scan_sparse_rows(
nisland_in: wp.array[int],
efc_J_rownnz_in: wp.array2d[int],
island_nefc_in: wp.array2d[int],
island_iefcadr_in: wp.array2d[int],
map_iefc2efc_in: wp.array2d[int],
iefc_J_rownnz_out: wp.array2d[int],
iefc_J_rowadr_out: wp.array2d[int],
):
worldid = wp.tid()
nisland = nisland_in[worldid]
if nisland == 0:
return
total_gathered_efc = island_iefcadr_in[worldid, nisland - 1] + island_nefc_in[worldid, nisland - 1]
running_rowadr = int(0)
for ic in range(total_gathered_efc):
c = map_iefc2efc_in[worldid, ic]
rownnz = efc_J_rownnz_in[worldid, c]
iefc_J_rowadr_out[worldid, ic] = running_rowadr
iefc_J_rownnz_out[worldid, ic] = rownnz
running_rowadr += rownnz
@wp.kernel
def _compute_efc_tree(
# Model:
nv: int,
body_treeid: wp.array[int],
jnt_dofadr: wp.array[int],
dof_treeid: wp.array[int],
geom_bodyid: wp.array[int],
site_bodyid: wp.array[int],
eq_type: wp.array[int],
eq_obj1id: wp.array[int],
eq_obj2id: wp.array[int],
eq_objtype: wp.array[int],
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
contact_geom_in: wp.array[wp.vec2i],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_J_in: wp.array3d[float],
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
njmax_in: int,
# Out:
efc_tree_out: wp.array2d[int],
):
"""Compute the first non-negative tree for each constraint."""
worldid, efcid = wp.tid()
if efcid >= wp.min(njmax_in, nefc_in[worldid]):
return
efc_type = efc_type_in[worldid, efcid]
efc_id = efc_id_in[worldid, efcid]
tree = int(-1)
use_generic = int(0)
# equality (connect/weld)
if efc_type == ConstraintType.EQUALITY:
eq_t = eq_type[efc_id]
if eq_t == EqType.CONNECT or eq_t == EqType.WELD:
b1 = eq_obj1id[efc_id]
b2 = eq_obj2id[efc_id]
if eq_objtype[efc_id] == ObjType.SITE:
b1 = site_bodyid[b1]
b2 = site_bodyid[b2]
t1 = body_treeid[b1]
t2 = body_treeid[b2]
if t1 >= 0:
tree = t1
else:
tree = t2
else:
# JOINT, TENDON, FLEX: generic scan
use_generic = 1
# joint friction
elif efc_type == ConstraintType.FRICTION_DOF:
tree = dof_treeid[efc_id]
# joint limit
elif efc_type == ConstraintType.LIMIT_JOINT:
tree = dof_treeid[jnt_dofadr[efc_id]]
# contact
elif (
efc_type == ConstraintType.CONTACT_FRICTIONLESS
or efc_type == ConstraintType.CONTACT_PYRAMIDAL
or efc_type == ConstraintType.CONTACT_ELLIPTIC
):
geom_pair = contact_geom_in[efc_id]
g1 = geom_pair[0]
g2 = geom_pair[1]
if g1 >= 0 and g2 >= 0:
t1 = body_treeid[geom_bodyid[g1]]
t2 = body_treeid[geom_bodyid[g2]]
if t1 >= 0:
tree = t1
else:
tree = t2
else:
# flex contacts: generic scan
use_generic = 1
else:
# generic: scan Jacobian row
use_generic = 1
if use_generic:
count = nv
rowadr = 0
if is_sparse:
count = efc_J_rownnz_in[worldid, efcid]
rowadr = efc_J_rowadr_in[worldid, efcid]
for i in range(count):
dof = i
if is_sparse:
sparseid = rowadr + i
dof = efc_J_colind_in[worldid, 0, sparseid]
else:
J_val = efc_J_in[worldid, efcid, dof]
if J_val == 0.0:
continue
t = dof_treeid[dof]
if t >= 0:
tree = t
break
efc_tree_out[worldid, efcid] = tree
@wp.kernel
def _gather_efc_and_jacobian(
# Model:
is_sparse: bool,
# Data in:
nefc_in: wp.array[int],
efc_D_in: wp.array2d[float],
efc_type_in: wp.array2d[int],
efc_id_in: wp.array2d[int],
efc_frictionloss_in: wp.array2d[float],
efc_aref_in: wp.array2d[float],
efc_J_in: wp.array3d[float],
# Sparse Jacobian arrays:
efc_J_rownnz_in: wp.array2d[int],
efc_J_rowadr_in: wp.array2d[int],
efc_J_colind_in: wp.array3d[int],
iefc_J_rowadr_in: wp.array2d[int],
# In:
njmax_in: int,
map_iefc2efc_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
map_dof2idof_in: wp.array2d[int],
nidof_in: wp.array[int],
# Out:
iefc_D_out: wp.array2d[float],
iefc_type_out: wp.array2d[int],
iefc_id_out: wp.array2d[int],
iefc_frictionloss_out: wp.array2d[float],
iefc_aref_out: wp.array2d[float],
iefc_J_out: wp.array3d[float],
iefc_J_colind_out: wp.array3d[int],
):
"""Gather constraint arrays and Jacobian into island-local dense order."""
worldid, iefcid = wp.tid()
if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
return
c = map_iefc2efc_in[worldid, iefcid]
# Scalar constraint fields
iefc_D_out[worldid, iefcid] = efc_D_in[worldid, c]
iefc_type_out[worldid, iefcid] = efc_type_in[worldid, c]
iefc_id_out[worldid, iefcid] = efc_id_in[worldid, c]
iefc_frictionloss_out[worldid, iefcid] = efc_frictionloss_in[worldid, c]
iefc_aref_out[worldid, iefcid] = efc_aref_in[worldid, c]
# Jacobian: convert to island-local dense format
nid = nidof_in[worldid]
if is_sparse:
rownnz = efc_J_rownnz_in[worldid, c]
rowadr_in = efc_J_rowadr_in[worldid, c]
rowadr_out = iefc_J_rowadr_in[worldid, iefcid]
for i in range(rownnz):
sparseid_in = rowadr_in + i
sparseid_out = rowadr_out + i
dof = efc_J_colind_in[worldid, 0, sparseid_in]
idof = map_dof2idof_in[worldid, dof]
# Store in sparse iefc_J_out and iefc_J_colind_out
iefc_J_out[worldid, 0, sparseid_out] = efc_J_in[worldid, 0, sparseid_in]
iefc_J_colind_out[worldid, 0, sparseid_out] = idof
else:
# Dense path: reorder rows by iefc, columns by idof
for idof in range(nid):
dof = map_idof2dof_in[worldid, idof]
iefc_J_out[worldid, iefcid, idof] = efc_J_in[worldid, c, dof]
@wp.kernel
def _gather_dof_arrays(
# Data in:
qacc_in: wp.array2d[float],
qacc_smooth_in: wp.array2d[float],
qfrc_smooth_in: wp.array2d[float],
# In:
nidof_in: wp.array[int],
map_idof2dof_in: wp.array2d[int],
# Out:
iacc_out: wp.array2d[float],
iacc_smooth_out: wp.array2d[float],
ifrc_smooth_out: wp.array2d[float],
):
"""Gather DOF arrays into island-local order."""
worldid, idofid = wp.tid()
if idofid >= nidof_in[worldid]:
return
dof = map_idof2dof_in[worldid, idofid]
iacc_out[worldid, idofid] = qacc_in[worldid, dof]
iacc_smooth_out[worldid, idofid] = qacc_smooth_in[worldid, dof]
ifrc_smooth_out[worldid, idofid] = qfrc_smooth_in[worldid, dof]
@wp.kernel
def _scatter_dof_arrays(
# In:
qacc_smooth_in: wp.array2d[float],
qfrc_smooth_in: wp.array2d[float],
dof_island_in: wp.array2d[int],
iacc_in: wp.array2d[float],
ifrc_constraint_in: wp.array2d[float],
iMa_in: wp.array2d[float],
map_dof2idof_in: wp.array2d[int],
scatter_Ma: bool,
# Out:
qacc_out: wp.array2d[float],
qfrc_constraint_out: wp.array2d[float],
Ma_out: wp.array2d[float],
):
"""Scatter island results to global arrays, copy qacc_smooth for non-island DOFs."""
worldid, dofid = wp.tid()
if dof_island_in[worldid, dofid] < 0:
qacc_out[worldid, dofid] = qacc_smooth_in[worldid, dofid]
qfrc_constraint_out[worldid, dofid] = 0.0
if scatter_Ma:
Ma_out[worldid, dofid] = qfrc_smooth_in[worldid, dofid]
else:
idof = map_dof2idof_in[worldid, dofid]
qacc_out[worldid, dofid] = iacc_in[worldid, idof]
qfrc_constraint_out[worldid, dofid] = ifrc_constraint_in[worldid, idof]
if scatter_Ma:
Ma_out[worldid, dofid] = iMa_in[worldid, idof]
@wp.kernel
def _scatter_efc_arrays(
# In:
nefc_in: wp.array[int],
njmax_in: int,
map_iefc2efc_in: wp.array2d[int],
iefc_force_in: wp.array2d[float],
iefc_state_in: wp.array2d[int],
# Out:
efc_force_out: wp.array2d[float],
efc_state_out: wp.array2d[int],
):
"""Scatter island-local constraint results back to global order."""
worldid, iefcid = wp.tid()
if iefcid >= wp.min(njmax_in, nefc_in[worldid]):
return
c = map_iefc2efc_in[worldid, iefcid]
efc_force_out[worldid, c] = iefc_force_in[worldid, iefcid]
efc_state_out[worldid, c] = iefc_state_in[worldid, iefcid]
@wp.kernel
def _init_island_arrays(
island_idofadr_out: wp.array2d[int],
island_nv_out: wp.array2d[int],
island_nefc_out: wp.array2d[int],
island_ne_out: wp.array2d[int],
island_nf_out: wp.array2d[int],
island_iefcadr_out: wp.array2d[int],
nidof_out: wp.array[int],
):
worldid, islandid = wp.tid()
island_nv_out[worldid, islandid] = 0
island_nefc_out[worldid, islandid] = 0
island_ne_out[worldid, islandid] = 0
island_nf_out[worldid, islandid] = 0
island_idofadr_out[worldid, islandid] = 0
island_iefcadr_out[worldid, islandid] = 0
if islandid == 0:
nidof_out[worldid] = 0
@wp.kernel
def _init_dof_arrays(
dof_island_out: wp.array2d[int],
map_dof2idof_out: wp.array2d[int],
map_idof2dof_out: wp.array2d[int],
idof_islandid_out: wp.array2d[int],
):
worldid, dofid = wp.tid()
dof_island_out[worldid, dofid] = -1
map_dof2idof_out[worldid, dofid] = 0
map_idof2dof_out[worldid, dofid] = 0
idof_islandid_out[worldid, dofid] = -1
@wp.kernel
def _init_efc_arrays(
efc_island_out: wp.array2d[int],
map_efc2iefc_out: wp.array2d[int],
map_iefc2efc_out: wp.array2d[int],
iefc_islandid_out: wp.array2d[int],
efc_tree_out: wp.array2d[int],
):
worldid, efcid = wp.tid()
efc_island_out[worldid, efcid] = -1
map_efc2iefc_out[worldid, efcid] = 0
map_iefc2efc_out[worldid, efcid] = 0
iefc_islandid_out[worldid, efcid] = -1
efc_tree_out[worldid, efcid] = -1
@event_scope
def compute_island_mapping(m: types.Model, d: types.Data, ctx: IslandSolverContext):
"""Compute DOF/constraint island mappings after island discovery.
Populates island solver context arrays via ctx: nv, nefc, ne, nf,
iefcadr, nidof, map_dof2idof, map_idof2dof, dof_islandid, map_efc2iefc,
map_iefc2efc, efc_islandid. Also populates d.dof_island, d.efc.island,
and d.island_dofadr.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext.
"""
# Ensure dof_islandid / efc_islandid are allocated at the right shape
if d.dof_islandid.shape[1] != m.nv:
d.dof_islandid = wp.empty((d.nworld, m.nv), dtype=int)
if d.efc_islandid.shape[1] != d.njmax:
d.efc_islandid = wp.empty((d.nworld, d.njmax), dtype=int)
# Ensure island-local DOF arrays are allocated at the right shape
if d.iqacc.shape[1] != m.nv:
nw = d.nworld
d.iqacc = wp.empty((nw, m.nv), dtype=float)
d.iqacc_smooth = wp.empty((nw, m.nv), dtype=float)
d.iqfrc_smooth = wp.empty((nw, m.nv), dtype=float)
d.iqfrc_constraint = wp.empty((nw, m.nv), dtype=float)
wp.launch(
_init_island_arrays,
dim=(d.nworld, m.ntree),
inputs=[],
outputs=[
d.island_dofadr,
d.island_nv,
d.island_nefc,
d.island_ne,
d.island_nf,
d.island_efcadr,
d.nidof,
],
)
wp.launch(
_init_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[],
outputs=[d.dof_island, d.map_dof2idof, d.map_idof2dof, d.dof_islandid],
)
efc_tree = wp.empty((d.nworld, d.njmax), dtype=int)
wp.launch(
_init_efc_arrays,
dim=(d.nworld, d.njmax),
inputs=[],
outputs=[d.efc.island, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid, efc_tree],
)
wp.launch(
_flood_fill,
dim=d.nworld,
inputs=[m.ntree, tree_tree, d.tree_island, stack_scratch],
outputs=[d.nisland, d.tree_island, stack_scratch],
_compute_efc_tree,
dim=(d.nworld, d.njmax),
inputs=[
m.nv,
m.body_treeid,
m.jnt_dofadr,
m.dof_treeid,
m.geom_bodyid,
m.site_bodyid,
m.eq_type,
m.eq_obj1id,
m.eq_obj2id,
m.eq_objtype,
m.is_sparse,
d.nefc,
d.contact.geom,
d.efc.type,
d.efc.id,
d.efc.J,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.njmax,
],
outputs=[efc_tree],
)
# 1. Count DOFs per island
wp.launch(
_island_count_dofs,
dim=(d.nworld, m.nv),
inputs=[m.dof_treeid, d.tree_island],
outputs=[d.dof_island, d.island_nv],
)
# 2. Count Constraints per island
wp.launch(
_island_count_constraints,
dim=(d.nworld, d.njmax),
inputs=[d.nefc, d.njmax, efc_tree, d.tree_island, d.efc.type],
outputs=[d.efc.island, d.island_nefc, d.island_ne, d.island_nf],
)
# 3. Scan sizes and reset counters for mapping
wp.launch(
_island_scan_sizes,
dim=d.nworld,
inputs=[d.nisland],
outputs=[d.island_dofadr, d.island_nv, d.island_nefc, d.island_efcadr, d.nidof],
)
# 4. Map DOFs
unconstrained_cnt = wp.zeros((d.nworld, 1), dtype=int)
wp.launch(
_island_map_dofs,
dim=(d.nworld, m.nv),
inputs=[m.nv, d.dof_island, d.island_dofadr, d.nidof],
outputs=[d.island_nv, d.map_dof2idof, d.map_idof2dof, d.dof_islandid, unconstrained_cnt],
)
# 5. Map Constraints
ne_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
nf_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
nother_mapped = wp.zeros((d.nworld, m.ntree), dtype=int)
wp.launch(
_island_map_constraints,
dim=(d.nworld, d.njmax),
inputs=[
d.nefc,
d.njmax,
d.efc.island,
d.island_efcadr,
d.island_ne,
d.island_nf,
d.efc.type,
ne_mapped,
nf_mapped,
nother_mapped,
],
outputs=[d.island_nefc, d.map_efc2iefc, d.map_iefc2efc, d.efc_islandid],
)
# 6. Scan Sparse Rows (if sparse)
if m.is_sparse:
wp.launch(
_island_scan_sparse_rows,
dim=d.nworld,
inputs=[d.nisland, d.efc.J_rownnz, d.island_nefc, d.island_efcadr, d.map_iefc2efc],
outputs=[d.efc.iJ_rownnz, d.efc.iJ_rowadr],
)
@event_scope
def gather_island_inputs(m: types.Model, d: types.Data, ctx: IslandSolverContext):
"""Gather constraint and DOF arrays into island-local order.
Populates d.iefc (D, type, id, frictionloss, aref, J, J_colind) and
d.iqacc, d.iqacc_smooth, d.iqfrc_smooth.
Must be called after compute_island_mapping() and before per-island solving.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext whose arrays are populated.
"""
# Gather constraint arrays and dense Jacobian (fused)
wp.launch(
_gather_efc_and_jacobian,
dim=(d.nworld, d.njmax),
inputs=[
m.is_sparse,
d.nefc,
d.efc.D,
d.efc.type,
d.efc.id,
d.efc.frictionloss,
d.efc.aref,
d.efc.J,
d.efc.J_rownnz,
d.efc.J_rowadr,
d.efc.J_colind,
d.efc.iJ_rowadr,
d.njmax,
d.map_iefc2efc,
d.map_idof2dof,
d.map_dof2idof,
d.nidof,
],
outputs=[
d.efc.iD,
d.efc.itype,
d.efc.iid,
d.efc.ifrictionloss,
d.efc.iaref,
d.efc.iJ,
d.efc.iJ_colind,
],
)
# Gather DOF arrays
wp.launch(
_gather_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[
d.qacc,
d.qacc_smooth,
d.qfrc_smooth,
d.nidof,
d.map_idof2dof,
],
outputs=[
d.iqacc,
d.iqacc_smooth,
d.iqfrc_smooth,
],
)
@event_scope
def scatter_island_results(m: types.Model, d: types.Data, ctx: IslandSolverContext, scatter_Ma: bool):
"""Scatter island-local solver results back to global arrays.
Reads ctx qacc, qfrc_constraint, Ma, d.efc iforce, istate and
writes them back to d.qacc, d.qfrc_constraint, d.efc.Ma, d.efc.force,
d.efc.state. Unconstrained DOFs receive qacc_smooth and zero qfrc.
Args:
m: Model.
d: Data.
ctx: IslandSolverContext which contains results.
scatter_Ma: Whether to scatter Ma for Euler/implicit integrators.
"""
# Scatter DOF results (and optionally Ma for Euler/implicit integrators)
wp.launch(
_scatter_dof_arrays,
dim=(d.nworld, m.nv),
inputs=[
d.qacc_smooth,
d.qfrc_smooth,
d.dof_island,
d.iqacc,
d.iqfrc_constraint,
ctx.Ma,
d.map_dof2idof,
scatter_Ma,
],
outputs=[
d.qacc,
d.qfrc_constraint,
d.efc.Ma,
],
)
# Scatter constraint results (force and state from island-local arrays)
wp.launch(
_scatter_efc_arrays,
dim=(d.nworld, d.njmax),
inputs=[
d.nefc,
d.njmax,
d.map_iefc2efc,
d.efc.iforce,
d.efc.istate,
],
outputs=[
d.efc.force,
d.efc.state,
],
)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/ray.py
Vendored
+13 -1
View File
@@ -707,10 +707,14 @@ def ray_mesh_with_bvh(
pnt: wp.vec3,
vec: wp.vec3,
max_t: float,
cull_backfaces: bool,
) -> Tuple[float, wp.vec3, float, float, int, int]:
"""Returns intersection information for ray mesh intersections.
Requires wp.Mesh be constructed and their ids to be passed.
When ``cull_backfaces`` is True, the function rejects exit-face hits in the
local-space frame. This matches MuJoCo OpenGL rendering's backface culling rule.
"""
t = float(-1.0)
u = float(0.0)
@@ -722,7 +726,12 @@ def ray_mesh_with_bvh(
lpnt, lvec = _ray_map(pos, mat, pnt, vec)
hit = wp.mesh_query_ray(mesh_bvh_id[mesh_geom_id], lpnt, lvec, max_t, t, u, v, sign, n, f)
if hit and wp.dot(lvec, n) < 0.0: # Backface culling in local space
if not hit:
return -1.0, wp.vec3(0.0, 0.0, 0.0), 0.0, 0.0, -1, -1
# Front-face hit, or back-face hit when cull is disabled: rotate the
# local-space normal into world space and return the hit.
if (not cull_backfaces) or wp.dot(lvec, n) < 0.0:
normal = mat @ n
normal = wp.normalize(normal)
return t, normal, u, v, f, mesh_geom_id
@@ -1055,6 +1064,8 @@ def _ray_geom_mesh_bvh(
if gtype == GeomType.MESH or gtype == GeomType.HFIELD:
bvh_ids = mesh_bvh_id if gtype == GeomType.MESH else hfield_bvh_id
# Public ray API (mjw.ray / mjw.rays) preserves MuJoCo's mj_ray cull rule:
# rangefinder sensors and user-facing ray casts always cull back-faces.
t, n, u, v, f, geom_mesh_id = ray_mesh_with_bvh(
bvh_ids,
geom_dataid[worldid % geom_dataid.shape[0], geomid],
@@ -1063,6 +1074,7 @@ def _ray_geom_mesh_bvh(
pnt,
vec,
min_dist,
True,
)
if t >= 0.0 and t < min_dist:
return t, n
mjx/mujoco/mjx/third_party/mujoco_warp/_src/render.py
Vendored
+27 -6
View File
@@ -180,6 +180,7 @@ def cast_ray(
flex_group_root: wp.array2d[int],
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
cull_backfaces: bool,
) -> Tuple[int, float, wp.vec3, float, float, int, int]:
dist = float(MJ_MAXVAL)
normal = wp.vec3(0.0, 0.0, 0.0)
@@ -232,6 +233,7 @@ def cast_ray(
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
@@ -281,6 +283,7 @@ def cast_ray(
ray_origin_world,
ray_dir_world,
dist,
cull_backfaces,
)
if gtype == GeomType.FLEX:
hit_geom_id = -2
@@ -308,6 +311,11 @@ def cast_ray(
if d >= 0.0:
hit_mesh_id = flexid
# Backface cull: drop exit-face hits when the ray origin is inside the geom,
# matching ray_mesh_with_bvh's `dot(lvec, n) < 0` rule.
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
if d >= 0.0 and d < dist:
dist = d
normal = n
@@ -349,6 +357,7 @@ def cast_ray_first_hit(
ray_origin_world: wp.vec3,
ray_dir_world: wp.vec3,
max_dist: float,
cull_backfaces: bool,
) -> bool:
"""A simpler version of casting rays that only checks for the first hit."""
query = wp.bvh_query_ray(bvh_id, ray_origin_world, ray_dir_world, group_root)
@@ -387,6 +396,7 @@ def cast_ray_first_hit(
ray_origin_world,
ray_dir_world,
max_dist,
cull_backfaces,
)
if gtype == GeomType.SPHERE:
d, n = ray_sphere(
@@ -467,6 +477,11 @@ def cast_ray_first_hit(
)
d = 0.0 if hit else -1.0
# Backface cull: see cast_ray for rationale. Strict `> 0` keeps tangent
# hits and skips branches with a zero-vector normal (mesh/flex anyhit).
if cull_backfaces and d >= 0.0 and wp.dot(ray_dir_world, n) > 0.0:
d = -1.0
if d >= 0.0 and d < max_dist:
return True
@@ -507,6 +522,7 @@ def compute_lighting(
lightdir: wp.vec3,
normal: wp.vec3,
hitpoint: wp.vec3,
cull_backfaces: bool,
) -> float:
light_contribution = float(0.0)
@@ -571,6 +587,7 @@ def compute_lighting(
shadow_origin,
L,
max_t,
cull_backfaces,
)
if shadow_hit:
@@ -725,6 +742,7 @@ def render(m: Model, d: Data, rc: RenderContext):
flex_group_root,
ray_origin_world,
ray_dir_world,
wp.static(rc.enable_backface_culling),
)
if render_seg[cam_idx] and geom_id != -1:
@@ -798,12 +816,14 @@ def render(m: Model, d: Data, rc: RenderContext):
)
base_color = wp.cw_mul(base_color, tex_color)
len_n = wp.length(normal)
n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
n = wp.normalize(n)
hemispheric = 0.5 * (n[2] + 1.0)
ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
result = 0.5 * wp.cw_mul(base_color, ambient_color)
result = wp.vec3(0.0, 0.0, 0.0)
if wp.static(rc.use_ambient_lighting):
len_n = wp.length(normal)
n = normal if len_n > 0.0 else wp.vec3(0.0, 0.0, 1.0)
n = wp.normalize(n)
hemispheric = 0.5 * (n[2] + 1.0)
ambient_color = wp.vec3(0.4, 0.4, 0.45) * hemispheric + wp.vec3(0.1, 0.1, 0.12) * (1.0 - hemispheric)
result = 0.5 * wp.cw_mul(base_color, ambient_color)
# Apply lighting and shadows
for l in range(wp.static(m.nlight)):
@@ -837,6 +857,7 @@ def render(m: Model, d: Data, rc: RenderContext):
light_xdir_in[worldid, l],
normal,
hit_point,
wp.static(rc.enable_backface_culling),
)
result = result + base_color * light_contribution
mjx/mujoco/mjx/third_party/mujoco_warp/_src/sensor.py
Vendored
+26 -6
View File
@@ -17,6 +17,7 @@ from typing import Any, Tuple
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src import history
from mujoco.mjx.third_party.mujoco_warp._src import math
from mujoco.mjx.third_party.mujoco_warp._src import ray
from mujoco.mjx.third_party.mujoco_warp._src import smooth
@@ -26,6 +27,7 @@ from mujoco.mjx.third_party.mujoco_warp._src.collision_sdf import sdf
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXCONPAIR
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MAXVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import MJ_MINVAL
from mujoco.mjx.third_party.mujoco_warp._src.types import TACTILE_DEPTH_SEMANTICS
from mujoco.mjx.third_party.mujoco_warp._src.types import ConeType
from mujoco.mjx.third_party.mujoco_warp._src.types import ConstraintType
from mujoco.mjx.third_party.mujoco_warp._src.types import ContactType
@@ -917,6 +919,10 @@ def sensor_pos(m: Model, d: Data):
],
)
# apply sensor delay/interval for position sensors
history.apply_sensor_delay(m, d, m.sensor_pos_adr)
history.apply_sensor_delay(m, d, m.sensor_limitpos_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.POS)
@@ -1459,6 +1465,10 @@ def sensor_vel(m: Model, d: Data):
],
)
# apply sensor delay/interval for velocity sensors
history.apply_sensor_delay(m, d, m.sensor_vel_adr)
history.apply_sensor_delay(m, d, m.sensor_limitvel_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.VEL)
@@ -2252,18 +2262,24 @@ def _sensor_tactile(
)
vel_rel = vel_sensor - vel_other
kMaxDepth = 0.05
pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
force = wp.mul(normal, pressure)
forceT = wp.vec3(0.0, 0.0, 0.0)
forceT[0] = wp.dot(force, normal)
if wp.static(TACTILE_DEPTH_SEMANTICS):
forceT[0] = -depth
else:
kMaxDepth = 0.05
pressure = depth / wp.max(kMaxDepth - depth, MJ_MINVAL)
force = wp.mul(normal, pressure)
forceT[0] = wp.dot(force, normal)
if has_frame:
forceT[1] = wp.abs(wp.dot(vel_rel, tang1))
forceT[2] = wp.abs(wp.dot(vel_rel, tang2))
dim = sensor_dim[sensor_id] // 3
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
if wp.static(TACTILE_DEPTH_SEMANTICS):
wp.atomic_max(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
else:
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 0 * dim + vertid, forceT[0])
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 1 * dim + vertid, forceT[1])
wp.atomic_add(sensordata_out, worldid, sensor_adr[sensor_id] + 2 * dim + vertid, forceT[2])
@@ -2709,6 +2725,10 @@ def sensor_acc(m: Model, d: Data):
],
)
# apply sensor delay/interval for acceleration sensors
history.apply_sensor_delay(m, d, m.sensor_acc_adr)
history.apply_sensor_delay(m, d, m.sensor_limitfrc_adr)
if m.callback.sensor:
m.callback.sensor(m, d, Stage.ACC)
mjx/mujoco/mjx/third_party/mujoco_warp/_src/smooth.py
Vendored
+179 -62
View File
@@ -823,37 +823,38 @@ def _crb_accumulate(
@wp.kernel
def _qM_sparse(
def _M_sparse(
# Model:
dof_bodyid: wp.array[int],
dof_parentid: wp.array[int],
dof_Madr: wp.array[int],
dof_armature: wp.array2d[float],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
# Data in:
cdof_in: wp.array2d[wp.spatial_vector],
crb_in: wp.array2d[vec10],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, dofid = wp.tid()
madr_ij = dof_Madr[dofid] # dof_Madr is not batched
bodyid = dof_bodyid[dofid]
madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
# init M(i,i) with armature inertia
qM_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
M_out[worldid, 0, madr_ij] = dof_armature[worldid % dof_armature.shape[0], dofid]
# precompute buf = crb_body_i * cdof_i
buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
# sparse backward pass over ancestors
while dofid >= 0:
qM_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
madr_ij += 1
M_out[worldid, 0, madr_ij] += wp.dot(cdof_in[worldid, dofid], buf)
madr_ij -= 1
dofid = dof_parentid[dofid]
@wp.kernel
def _qM_dense(
def _M_dense(
# Model:
dof_bodyid: wp.array[int],
dof_parentid: wp.array[int],
@@ -862,7 +863,7 @@ def _qM_dense(
cdof_in: wp.array2d[wp.spatial_vector],
crb_in: wp.array2d[vec10],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, dofid = wp.tid()
bodyid = dof_bodyid[dofid]
@@ -873,15 +874,15 @@ def _qM_dense(
buf = math.inert_vec(crb_in[worldid, bodyid], cdof_in[worldid, dofid])
M += wp.dot(cdof_in[worldid, dofid], buf)
qM_out[worldid, dofid, dofid] = M
M_out[worldid, dofid, dofid] = M
# sparse backward pass over ancestors
dofidi = dofid
dofid = dof_parentid[dofid]
while dofid >= 0:
qMij = wp.dot(cdof_in[worldid, dofid], buf)
qM_out[worldid, dofidi, dofid] += qMij
qM_out[worldid, dofid, dofidi] += qMij
Mij = wp.dot(cdof_in[worldid, dofid], buf)
M_out[worldid, dofidi, dofid] += Mij
M_out[worldid, dofid, dofidi] += Mij
dofid = dof_parentid[dofid]
@@ -898,17 +899,17 @@ def crb(m: Model, d: Data):
body_tree = m.body_tree[i]
wp.launch(_crb_accumulate, dim=(d.nworld, body_tree.size), inputs=[m.body_parentid, d.crb, body_tree], outputs=[d.crb])
d.qM.zero_()
d.M.zero_()
if m.is_sparse:
wp.launch(
_qM_sparse,
_M_sparse,
dim=(d.nworld, m.nv),
inputs=[m.dof_bodyid, m.dof_parentid, m.dof_Madr, m.dof_armature, d.cdof, d.crb],
outputs=[d.qM],
inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, m.M_rownnz, m.M_rowadr, d.cdof, d.crb],
outputs=[d.M],
)
else:
wp.launch(
_qM_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.qM]
_M_dense, dim=(d.nworld, m.nv), inputs=[m.dof_bodyid, m.dof_parentid, m.dof_armature, d.cdof, d.crb], outputs=[d.M]
)
@@ -916,16 +917,17 @@ def crb(m: Model, d: Data):
def _tendon_armature(
# Model:
dof_parentid: wp.array[int],
dof_Madr: wp.array[int],
ten_J_rownnz: wp.array[int],
ten_J_rowadr: wp.array[int],
ten_J_colind: wp.array[int],
tendon_armature: wp.array2d[float],
M_rownnz: wp.array[int],
M_rowadr: wp.array[int],
is_sparse: bool,
# Data in:
ten_J_in: wp.array2d[float],
# Data out:
qM_out: wp.array3d[float],
M_out: wp.array3d[float],
):
worldid, tenid, dofid = wp.tid()
@@ -947,7 +949,7 @@ def _tendon_armature(
return
if is_sparse:
madr_ij = dof_Madr[dofid]
madr_ij = M_rowadr[dofid] + M_rownnz[dofid] - 1
# sparse backward pass over ancestors
dofidi = dofid
@@ -967,52 +969,40 @@ def _tendon_armature(
else:
ten_Jj = float(0.0)
qMij = armature * ten_Jj * ten_Ji
Mij = armature * ten_Jj * ten_Ji
if is_sparse:
wp.atomic_add(qM_out[worldid, 0], madr_ij, qMij)
madr_ij += 1
wp.atomic_add(M_out[worldid, 0], madr_ij, Mij)
madr_ij -= 1
else:
wp.atomic_add(qM_out[worldid, dofidi], dofid, qMij)
wp.atomic_add(M_out[worldid, dofidi], dofid, Mij)
if dofidi != dofid:
wp.atomic_add(qM_out[worldid, dofid], dofidi, qMij)
wp.atomic_add(M_out[worldid, dofid], dofidi, Mij)
dofid = dof_parentid[dofid]
@event_scope
def tendon_armature(m: Model, d: Data):
"""Add tendon armature to qM."""
"""Add tendon armature to M."""
wp.launch(
_tendon_armature,
dim=(d.nworld, m.ntendon, m.max_ten_J_rownnz),
inputs=[
m.dof_parentid,
m.dof_Madr,
m.ten_J_rownnz,
m.ten_J_rowadr,
m.ten_J_colind,
m.tendon_armature,
m.M_rownnz,
m.M_rowadr,
m.is_sparse,
d.ten_J,
],
outputs=[d.qM],
outputs=[d.M],
)
@wp.kernel
def _copy_CSR(
# Model:
mapM2M: wp.array[int],
# In:
M_in: wp.array3d[float],
# Out:
L_out: wp.array3d[float],
):
worldid, ind = wp.tid()
L_out[worldid, 0, ind] = M_in[worldid, 0, mapM2M[ind]]
@wp.kernel
def _qLD_acc(
# Model:
@@ -1055,7 +1045,7 @@ def _qLDiag_div(
def _factor_i_sparse(m: Model, d: Data, M: wp.array3d[float], L: wp.array3d[float], D: wp.array2d[float]):
"""Sparse L'*D*L factorization of inertia-like matrix M, assumed spd."""
wp.launch(_copy_CSR, dim=(d.nworld, m.nC), inputs=[m.mapM2M, M], outputs=[L])
wp.copy(L, M)
for i in reversed(range(len(m.qLD_updates))):
qLD_updates = m.qLD_updates[i]
@@ -1071,7 +1061,7 @@ def _tile_cholesky_factorize(tile: TileSet):
@wp.kernel(module="unique", enable_backward=False)
def cholesky_factorize(
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
adr: wp.array[int],
# Out:
@@ -1081,8 +1071,8 @@ def _tile_cholesky_factorize(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
M_tile = wp.tile_load(qM_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
L_tile = wp.tile_cholesky(M_tile)
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
wp.tile_store(L_out[worldid], L_tile, offset=(dofid, dofid))
return cholesky_factorize
@@ -1090,7 +1080,7 @@ def _tile_cholesky_factorize(tile: TileSet):
def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
"""Dense Cholesky factorization of inertia-like matrix M, assumed spd."""
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize(tile),
dim=(d.nworld, tile.adr.size),
@@ -1104,9 +1094,9 @@ def _factor_i_dense(m: Model, d: Data, M: wp.array, L: wp.array):
def factor_m(m: Model, d: Data):
"""Factorization of inertia-like matrix M, assumed spd."""
if m.is_sparse:
_factor_i_sparse(m, d, d.qM, d.qLD, d.qLDiagInv)
_factor_i_sparse(m, d, d.M, d.qLD, d.qLDiagInv)
else:
_factor_i_dense(m, d, d.qM, d.qLD)
_factor_i_dense(m, d, d.M, d.qLD)
@wp.kernel
@@ -2842,15 +2832,15 @@ def _tile_cholesky_solve(tile: TileSet):
dofid = adr[nodeid]
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
wp.tile_store(x[worldid], x_slice, offset=(dofid,))
return cholesky_solve
def _solve_LD_dense(m: Model, d: Data, L: wp.array3d[float], x: wp.array2d[float], y: wp.array2d[float]):
"""Computes dense backsubstitution: x = inv(L'*L)*y."""
for tile in m.qM_tiles:
"""Computes dense backsubstitution: x = inv(U.T @ U) * y."""
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_solve(tile),
dim=(d.nworld, tile.adr.size),
@@ -2868,9 +2858,9 @@ def solve_LD(
x: wp.array2d[float],
y: wp.array2d[float],
):
"""Computes backsubstitution to solve a linear system of the form x = inv(L'*D*L) * y.
"""Computes backsubstitution for the inertia factorization.
L and D are the factors from the Cholesky factorization of the inertia matrix.
Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
This function dispatches to either a sparse or dense solver depending on Model options.
@@ -2907,8 +2897,9 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
@wp.kernel(module="unique", enable_backward=False)
def cholesky_factorize_solve(
# Data in:
M_in: wp.array3d[float],
# In:
M: wp.array3d[float],
y: wp.array2d[float],
adr: wp.array[int],
# Out:
@@ -2919,12 +2910,12 @@ def _tile_cholesky_factorize_solve(tile: TileSet):
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
M_tile = wp.tile_load(M[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
M_tile = wp.tile_load(M_in[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(dofid,))
L_tile = wp.tile_cholesky(M_tile)
L_tile = wp.tile_cholesky(M_tile, fill_mode="upper")
wp.tile_store(L[worldid], L_tile, offset=(dofid, dofid))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
x_slice = wp.tile_cholesky_solve(L_tile, y_slice, fill_mode="upper")
wp.tile_store(x[worldid], x_slice, offset=(dofid,))
return cholesky_factorize_solve
@@ -2938,7 +2929,7 @@ def _factor_solve_i_dense(
y: wp.array2d[float],
L: wp.array3d[float],
):
for tile in m.qM_tiles:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_factorize_solve(tile),
dim=(d.nworld, tile.adr.size),
@@ -2949,9 +2940,9 @@ def _factor_solve_i_dense(
def factor_solve_i(m, d, M, L, D, x, y):
"""Factorizes and solves the linear system: x = inv(L'*D*L) * y or x = inv(L'*L) * y.
"""Factorizes and solves the inertia-like linear system.
M is an inertia-like matrix and L, D are its Cholesky-like factors.
Sparse models use MuJoCo's L'*D*L factors; dense models use an upper Cholesky factor U.
This function first factorizes the matrix M (sparse or dense depending on model options),
then solves the system for x given right-hand side y.
@@ -2960,7 +2951,7 @@ def factor_solve_i(m, d, M, L, D, x, y):
m: The model containing factorization and sparsity information.
d: The data object containing workspace and factorization results.
M: The inertia-like matrix to factorize.
L: Output lower-triangular factor from the factorization (sparse or dense).
L: Output sparse factor or dense upper Cholesky factor.
D: Output diagonal factor from the factorization (only used for sparse).
x: Output array for the solution.
y: Input right-hand side array.
@@ -2972,6 +2963,132 @@ def factor_solve_i(m, d, M, L, D, x, y):
_factor_solve_i_dense(m, d, M, x, y, L)
@cache_kernel
def _factor_solve_lu_sparse_fused(nv: int):
"""Fused sparse LU factorization and solve in a single kernel."""
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Model:
D_rownnz: wp.array[int],
D_rowadr: wp.array[int],
D_diag: wp.array[int],
D_colind: wp.array[int],
# In:
qfrc: wp.array2d[float],
# Data out:
qacc_out: wp.array2d[float],
qLU_out: wp.array3d[float],
):
worldid = wp.tid()
NV = wp.static(nv)
# Phase 1: LU factorization (in-place on qLU_out)
for i in range(NV):
qacc_out[worldid, i] = float(D_rownnz[i])
# process diagonal elements from n-1 down to 0
for r_rev in range(NV):
i = NV - 1 - r_rev
rem_i = int(qacc_out[worldid, i])
rowadr_i = D_rowadr[i]
ii = rowadr_i + rem_i - 1
qacc_out[worldid, i] = float(rem_i - 1)
# cache diagonal element for row i
LUii = qLU_out[worldid, 0, ii]
# rows j above i (j < i), processed from i-1 down to 0
for c in range(i):
j = i - 1 - c
# get address of last remaining element of row j
rem_j = int(qacc_out[worldid, j])
rowadr_j = D_rowadr[j]
ji = rowadr_j + rem_j - 1
# process row j if (j,i) is non-zero
if D_colind[ji] == i:
# adjust remaining counter
rem_j = rem_j - 1
qacc_out[worldid, j] = float(rem_j)
# (j,i) = (j,i) / (i,i)
LUji = qLU_out[worldid, 0, ji] / LUii
qLU_out[worldid, 0, ji] = LUji
# (j,k) = (j,k) - (i,k) * (j,i) for k < i
icnt = rowadr_i
jcnt = rowadr_j
jend = rowadr_j + rem_j
while jcnt < jend:
col_i = D_colind[icnt]
col_j = D_colind[jcnt]
if col_i == col_j:
qLU_out[worldid, 0, jcnt] = qLU_out[worldid, 0, jcnt] - qLU_out[worldid, 0, icnt] * LUji
icnt = icnt + 1
jcnt = jcnt + 1
elif col_i > col_j:
jcnt = jcnt + 1
else:
icnt = icnt + 1
# Phase 2: LU solve (backward + forward substitution)
# Backward substitution: solve (U+I)*qacc = qfrc
for k_rev in range(NV):
i = NV - 1 - k_rev
diag_i = D_diag[i]
rowadr_i = D_rowadr[i]
d1 = diag_i + 1
nnz_upper = D_rownnz[i] - d1
acc = qfrc[worldid, i]
for j in range(nnz_upper):
adr_j = rowadr_i + d1 + j
col = D_colind[adr_j]
acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
qacc_out[worldid, i] = acc
# Forward substitution: solve L*qacc = qacc
for i in range(NV):
diag_i = D_diag[i]
rowadr_i = D_rowadr[i]
acc = qacc_out[worldid, i]
for j in range(diag_i):
adr_j = rowadr_i + j
col = D_colind[adr_j]
acc = acc - qLU_out[worldid, 0, adr_j] * qacc_out[worldid, col]
qacc_out[worldid, i] = acc / qLU_out[worldid, 0, rowadr_i + diag_i]
return kernel
@event_scope
def factor_solve_lu(m: Model, d: Data, qLU: wp.array3d[float], qacc: wp.array2d[float], qfrc: wp.array2d[float]):
r"""Factorize and solve non-symmetric implicit system: qacc = A \\ qfrc.
qLU is overwritten in-place with the LU factors, then used to solve for qacc.
Args:
m: The model containing D-structure sparsity information.
d: The data object.
qLU: array containing the system matrix, overwritten with LU factors.
qacc: output array for the solution.
qfrc: input right-hand side.
"""
wp.launch(
_factor_solve_lu_sparse_fused(m.nv),
dim=(d.nworld,),
inputs=[m.D_rownnz, m.D_rowadr, m.D_diag, m.D_colind, qfrc],
outputs=[qacc, qLU],
)
@wp.kernel
def _subtree_vel_forward(
# Model:
mjx/mujoco/mjx/third_party/mujoco_warp/_src/solver.py
Vendored
+2227 -110
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/third_party/mujoco_warp/_src/support.py
Vendored
+365 -15
View File
@@ -69,11 +69,11 @@ def mul_m_sparse(check_skip: bool):
@wp.kernel(module="unique")
def _mul_m_sparse(
# Model:
qM_mulm_rowadr: wp.array[int],
qM_mulm_col: wp.array[int],
qM_mulm_madr: wp.array[int],
M_mulm_rowadr: wp.array[int],
M_mulm_col: wp.array[int],
M_mulm_madr: wp.array[int],
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
vec: wp.array2d[float],
skip: wp.array[bool],
@@ -89,12 +89,12 @@ def mul_m_sparse(check_skip: bool):
# Gather all contributions (diagonal + off-diagonal)
acc = float(0.0)
start = qM_mulm_rowadr[dofid]
end = qM_mulm_rowadr[dofid + 1]
start = M_mulm_rowadr[dofid]
end = M_mulm_rowadr[dofid + 1]
for k in range(start, end):
col = qM_mulm_col[k]
madr = qM_mulm_madr[k]
acc += qM_in[worldid, 0, madr] * vec[worldid, col]
col = M_mulm_col[k]
madr = M_mulm_madr[k]
acc += M_in[worldid, 0, madr] * vec[worldid, col]
res[worldid, dofid] = acc
@@ -108,7 +108,7 @@ def mul_m_dense(nv: int, check_skip: bool):
@wp.kernel(module="unique")
def _mul_m_dense(
# Data in:
qM_in: wp.array3d[float],
M_in: wp.array3d[float],
# In:
vec: wp.array2d[float],
skip: wp.array[bool],
@@ -123,7 +123,7 @@ def mul_m_dense(nv: int, check_skip: bool):
acc = float(0.0)
for j in range(wp.static(nv)):
acc += qM_in[worldid, i, j] * vec[worldid, j]
acc += M_in[worldid, i, j] * vec[worldid, j]
res[worldid, i] = acc
return _mul_m_dense
@@ -143,8 +143,8 @@ def mul_m(
Args:
m: The model containing kinematic and dynamic information (device).
d: The data object containing the current state and output arrays (device).
res: Result: qM @ vec.
vec: Input vector to multiply by qM.
res: Result: M @ vec.
vec: Input vector to multiply by M.
skip: Per-world bitmask to skip computing output.
M: Input matrix: M @ vec.
"""
@@ -152,13 +152,13 @@ def mul_m(
skip = skip or wp.empty(0, dtype=bool)
if M is None:
M = d.qM
M = d.M
if m.is_sparse:
wp.launch(
mul_m_sparse(check_skip),
dim=(d.nworld, m.nv),
inputs=[m.qM_mulm_rowadr, m.qM_mulm_col, m.qM_mulm_madr, M, vec, skip],
inputs=[m.M_mulm_rowadr, m.M_mulm_col, m.M_mulm_madr, M, vec, skip],
outputs=[res],
)
@@ -171,6 +171,340 @@ def mul_m(
)
@wp.kernel
def _mul_m_island_sparse(
# Model:
M_mulm_rowadr: wp.array[int],
M_mulm_col: wp.array[int],
M_mulm_madr: wp.array[int],
# Data in:
nidof_in: wp.array[int],
M_in: wp.array3d[float],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
idof_islandid_in: wp.array2d[int],
vec: wp.array2d[float],
island_done_in: wp.array2d[bool],
check_skip: int,
# Out:
res: wp.array2d[float],
):
"""Sparse island mul_m for ALL islands in parallel."""
worldid, idofid = wp.tid()
nidof = nidof_in[worldid]
if idofid >= nidof:
return
islandid = idof_islandid_in[worldid, idofid]
if islandid < 0:
return
if check_skip:
if island_done_in[worldid, islandid]:
return
dof = map_idof2dof_in[worldid, idofid]
acc = float(0.0)
start = M_mulm_rowadr[dof]
end = M_mulm_rowadr[dof + 1]
for k in range(start, end):
col = M_mulm_col[k]
madr = M_mulm_madr[k]
col_idof = map_dof2idof_in[worldid, col]
# skip unconstrained DOFs
if col_idof < nidof:
acc += M_in[worldid, 0, madr] * vec[worldid, col_idof]
res[worldid, idofid] = acc
@wp.kernel
def _mul_m_island_dense(
# Model:
nv: int,
# Data in:
nidof_in: wp.array[int],
M_in: wp.array3d[float],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
idof_islandid_in: wp.array2d[int],
vec: wp.array2d[float],
island_done_in: wp.array2d[bool],
check_skip: int,
# Out:
res: wp.array2d[float],
):
"""Dense island mul_m for ALL islands in parallel."""
worldid, idofid = wp.tid()
nidof = nidof_in[worldid]
if idofid >= nidof:
return
islandid = idof_islandid_in[worldid, idofid]
if islandid < 0:
return
if check_skip:
if island_done_in[worldid, islandid]:
return
dof = map_idof2dof_in[worldid, idofid]
acc = float(0.0)
for j in range(nv):
col_idof = map_dof2idof_in[worldid, j]
# skip unconstrained DOFs
if col_idof < nidof:
acc += M_in[worldid, dof, j] * vec[worldid, col_idof]
res[worldid, idofid] = acc
@event_scope
def mul_m_island(
m: Model,
d: Data,
res: wp.array2d[float],
vec: wp.array2d[float],
nidof: wp.array[int],
map_idof2dof: wp.array2d[int],
map_dof2idof: wp.array2d[int],
idof_islandid: wp.array2d[int],
island_done: Optional[wp.array] = None,
M: Optional[wp.array] = None,
):
"""Multiply island-local vectors by inertia matrix for all islands in parallel.
Args:
m: The model containing kinematic and dynamic information.
d: The data object containing the current state and output arrays.
res: Result: qM @ vec (island-local DOF order).
vec: Input vector (island-local DOF order).
nidof: Number of island DOFs per world.
map_idof2dof: Island-local DOF → global DOF map.
map_dof2idof: Global DOF → island-local DOF map.
idof_islandid: Island ID per island-local DOF.
island_done: Per-island done flags (nworld, ntree).
M: Optional mass matrix override.
"""
check_skip = int(island_done is not None)
island_done = island_done or wp.empty((0, 0), dtype=bool)
if M is None:
M = d.M
if m.is_sparse:
wp.launch(
_mul_m_island_sparse,
dim=(d.nworld, m.nv),
inputs=[
m.M_mulm_rowadr,
m.M_mulm_col,
m.M_mulm_madr,
nidof,
M,
map_dof2idof,
map_idof2dof,
idof_islandid,
vec,
island_done,
check_skip,
],
outputs=[res],
)
else:
wp.launch(
_mul_m_island_dense,
dim=(d.nworld, m.nv),
inputs=[
m.nv,
nidof,
M,
map_dof2idof,
map_idof2dof,
idof_islandid,
vec,
island_done,
check_skip,
],
outputs=[res],
)
@cache_kernel
def _solve_LD_sparse_island(nv: int, nlevels: int):
"""Sparse backsubstitution with island-local index remapping.
Same algorithm as _solve_LD_sparse_fused, but reads/writes x/y in
island-local DOF order. The L/D factorization stays in global DOF order.
"""
@wp.func_native(snippet="WP_TILE_SYNC();")
def _syncthreads():
pass
@wp.kernel(module="unique", enable_backward=False)
def kernel(
# Data in:
nidof_in: wp.array[int],
map_dof2idof_in: wp.array2d[int],
map_idof2dof_in: wp.array2d[int],
# In:
L: wp.array3d[float],
D: wp.array2d[float],
all_updates: wp.array[wp.vec3i],
level_offsets: wp.array[int],
y: wp.array2d[float],
# Out:
x_out: wp.array2d[float],
):
worldid, tid = wp.tid()
NLEVELS = wp.static(nlevels)
BLOCK_DIM = wp.block_dim()
nid = nidof_in[worldid]
# Copy y to x_out (island-local, only iterate up to nid)
for idof in range(tid, nid, BLOCK_DIM):
x_out[worldid, idof] = y[worldid, idof]
_syncthreads()
# Forward substitution
for level in range(NLEVELS):
level_idx = NLEVELS - 1 - level
level_offset = level_offsets[level_idx]
level_size = level_offsets[level_idx + 1] - level_offset
for u in range(tid, level_size, BLOCK_DIM):
update = all_updates[level_offset + u]
i, k, Madr_ki = update[0], update[1], update[2]
idof_i = map_dof2idof_in[worldid, i]
if idof_i < nid:
idof_k = map_dof2idof_in[worldid, k]
wp.atomic_sub(x_out[worldid], idof_i, L[worldid, 0, Madr_ki] * x_out[worldid, idof_k])
_syncthreads()
# Diagonal multiply (only iterate up to nid)
for idof in range(tid, nid, BLOCK_DIM):
dofid = map_idof2dof_in[worldid, idof]
x_out[worldid, idof] *= D[worldid, dofid]
_syncthreads()
# Backward substitution
for level in range(NLEVELS):
level_idx = level
level_offset = level_offsets[level_idx]
level_size = level_offsets[level_idx + 1] - level_offset
for u in range(tid, level_size, BLOCK_DIM):
update = all_updates[level_offset + u]
i, k, Madr_ki = update[0], update[1], update[2]
idof_k = map_dof2idof_in[worldid, k]
if idof_k < nid:
idof_i = map_dof2idof_in[worldid, i]
wp.atomic_sub(x_out[worldid], idof_k, L[worldid, 0, Madr_ki] * x_out[worldid, idof_i])
_syncthreads()
return kernel
@cache_kernel
def _tile_cholesky_solve_island(tile):
"""Dense Cholesky backsubstitution with island-local index remapping.
L is loaded from global DOF offsets (factorization unchanged).
y/x are loaded/stored at island-local DOF offsets via map_dof2idof.
"""
@wp.kernel(module="unique", enable_backward=False)
def cholesky_solve(
# Data in:
nidof_in: wp.array[int],
map_dof2idof_in: wp.array2d[int],
# In:
L: wp.array3d[float],
y: wp.array2d[float],
adr: wp.array[int],
# Out:
x: wp.array2d[float],
):
worldid, nodeid = wp.tid()
TILE_SIZE = wp.static(tile.size)
dofid = adr[nodeid]
idofid = map_dof2idof_in[worldid, dofid]
# Skip unconstrained trees (uniform branch — all threads in block agree)
if idofid >= nidof_in[worldid]:
return
# L stays in global order
L_tile = wp.tile_load(L[worldid], shape=(TILE_SIZE, TILE_SIZE), offset=(dofid, dofid))
# y and x use island-local offsets
y_slice = wp.tile_load(y[worldid], shape=(TILE_SIZE,), offset=(idofid,))
x_slice = wp.tile_cholesky_solve(L_tile, y_slice)
wp.tile_store(x[worldid], x_slice, offset=(idofid,))
return cholesky_solve
@event_scope
def solve_m_island(
m: Model,
d: Data,
res: wp.array2d[float],
vec: wp.array2d[float],
nidof: wp.array[int],
map_idof2dof: wp.array2d[int],
):
"""Compute res = M^{-1} @ vec for island-local DOFs.
Args:
m: Model.
d: Data.
res: Output in island-local DOF order.
vec: Input in island-local DOF order.
nidof: Number of island DOFs per world.
map_idof2dof: Island-local DOF -> global DOF map.
"""
if m.is_sparse:
nlevels = len(m.qLD_updates)
if wp.get_device().is_cuda:
dim_block = m.block_dim.solve_LD_sparse_fused
else:
dim_block = 1
wp.launch(
_solve_LD_sparse_island(m.nv, nlevels),
dim=(d.nworld, dim_block),
inputs=[
d.nidof,
d.map_dof2idof,
map_idof2dof,
d.qLD,
d.qLDiagInv,
m.qLD_all_updates,
m.qLD_level_offsets,
vec,
],
outputs=[res],
block_dim=dim_block,
)
else:
for tile in m.M_tiles:
wp.launch_tiled(
_tile_cholesky_solve_island(tile),
dim=(d.nworld, tile.adr.size),
inputs=[d.nidof, d.map_dof2idof, d.qLD, vec, tile.adr],
outputs=[res],
block_dim=m.block_dim.cholesky_solve,
)
@wp.kernel
def _apply_ft(
# Model:
@@ -604,11 +938,13 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
nbody: int,
neq: int,
nmocap: int,
nhistory: int,
# Data in:
time_in: wp.array[float],
qpos_in: wp.array2d[float],
qvel_in: wp.array2d[float],
act_in: wp.array2d[float],
history_in: wp.array2d[float],
qacc_warmstart_in: wp.array2d[float],
ctrl_in: wp.array2d[float],
qfrc_applied_in: wp.array2d[float],
@@ -647,6 +983,10 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
for j in range(na):
state_out[worldid, adr + j] = act_in[worldid, j]
adr += na
elif element == State.HISTORY:
for j in range(nhistory):
state_out[worldid, adr + j] = history_in[worldid, j]
adr += nhistory
elif element == State.WARMSTART:
for j in range(nv):
state_out[worldid, adr + j] = qacc_warmstart_in[worldid, j]
@@ -700,10 +1040,12 @@ def get_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
m.nbody,
m.neq,
m.nmocap,
m.nhistory,
d.time,
d.qpos,
d.qvel,
d.act,
d.history,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
@@ -743,6 +1085,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
nbody: int,
neq: int,
nmocap: int,
nhistory: int,
# In:
sig_in: int,
active_in: wp.array[bool],
@@ -752,6 +1095,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
qpos_out: wp.array2d[float],
qvel_out: wp.array2d[float],
act_out: wp.array2d[float],
history_out: wp.array2d[float],
qacc_warmstart_out: wp.array2d[float],
ctrl_out: wp.array2d[float],
qfrc_applied_out: wp.array2d[float],
@@ -785,6 +1129,10 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
for j in range(na):
act_out[worldid, j] = state_in[worldid, adr + j]
adr += na
elif element == State.HISTORY:
for j in range(nhistory):
history_out[worldid, j] = state_in[worldid, adr + j]
adr += nhistory
elif element == State.WARMSTART:
for j in range(nv):
qacc_warmstart_out[worldid, j] = state_in[worldid, adr + j]
@@ -844,6 +1192,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
m.nbody,
m.neq,
m.nmocap,
m.nhistory,
int(sig),
active or wp.ones(d.nworld, dtype=bool),
state,
@@ -853,6 +1202,7 @@ def set_state(m: Model, d: Data, state: wp.array2d[float], sig: int, active: Opt
d.qpos,
d.qvel,
d.act,
d.history,
d.qacc_warmstart,
d.ctrl,
d.qfrc_applied,
mjx/mujoco/mjx/third_party/mujoco_warp/_src/types.py
Vendored
+258 -43
View File
@@ -17,18 +17,19 @@ import enum
from typing import Callable
import mujoco
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
import numpy as np
import warp as wp
from mujoco.mjx.third_party.mujoco_warp._src.util_pkg import check_version
MJ_MINVAL = mujoco.mjMINVAL
MJ_MAXVAL = mujoco.mjMAXVAL
MJ_MINIMP = mujoco.mjMINIMP # minimum constraint impedance
MJ_MAXIMP = mujoco.mjMAXIMP # maximum constraint impedance
MJ_MAXCONPAIR = mujoco.mjMAXCONPAIR
MJ_MINMU = mujoco.mjMINMU # minimum friction
# True if MuJoCo >= 3.9.0 (new margin/gap semantics: includemargin = margin)
_NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0")
NEW_GAP_SEMANTICS = check_version("mujoco>=3.9.0.dev914519929")
TACTILE_DEPTH_SEMANTICS = check_version("mujoco>=3.9.0.dev921980899")
# maximum size (by number of edges) of an horizon in EPA algorithm
MJ_MAX_EPAHORIZON = 24
# maximum average number of trianglarfaces EPA can insert at each iteration
@@ -47,6 +48,25 @@ class BlockDim:
"""Block dimension 'block_dim' settings for wp.launch_tiled.
TODO(team): experimental and may be removed
Attributes:
segmented_sort: segmented sort block dimension (collision_driver)
euler_dense: Euler dense block dimension (forward)
actuator_velocity: actuator velocity block dimension (forward)
ray: ray block dimension (ray)
contact_sort: contact sort block dimension (sensor)
energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
cholesky_factorize: Cholesky factorize block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
"""
# collision_driver
@@ -131,13 +151,8 @@ class ProjectionType(enum.IntEnum):
ORTHOGRAPHIC: orthographic projection
"""
# TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
if hasattr(mujoco, "mjtProjection"):
PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
else:
PERSPECTIVE = 0
ORTHOGRAPHIC = 1
PERSPECTIVE = mujoco.mjtProjection.mjPROJ_PERSPECTIVE
ORTHOGRAPHIC = mujoco.mjtProjection.mjPROJ_ORTHOGRAPHIC
class Stage(enum.IntEnum):
@@ -188,6 +203,7 @@ class DisableBit(enum.IntFlag):
EULERDAMP: implicit damping for Euler integration
NATIVECCD: native convex collision detection (ignored in MJWarp)
ISLAND: constraint islands
MULTICCD: disable multiple contacts with CCD
"""
CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT
@@ -217,7 +233,6 @@ class EnableBit(enum.IntFlag):
Attributes:
ENERGY: energy computation
INVDISCRETE: discrete-time inverse dynamics
MULTICCD: multiple contacts with CCD
"""
ENERGY = mujoco.mjtEnableBit.mjENBL_ENERGY
@@ -338,12 +353,13 @@ class IntegratorType(enum.IntEnum):
EULER: semi-implicit Euler
RK4: 4th-order Runge Kutta
IMPLICITFAST: implicit in velocity, no rne derivative
IMPLICIT: implicit in velocity, with rne derivative
"""
EULER = mujoco.mjtIntegrator.mjINT_EULER
RK4 = mujoco.mjtIntegrator.mjINT_RK4
IMPLICITFAST = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
# unsupported: IMPLICIT
IMPLICIT = mujoco.mjtIntegrator.mjINT_IMPLICIT
class GeomType(enum.IntEnum):
@@ -617,6 +633,7 @@ class State(enum.IntEnum):
QPOS: position
QVEL: velocity
ACT: actuator activation
HISTORY: delay/interval history buffers
WARMSTART: acceleration used for warmstart
CTRL: control
QFRC_APPLIED: applied generalized force
@@ -625,7 +642,7 @@ class State(enum.IntEnum):
MOCAP_POS: positions of mocap bodies
MOCAP_QUAT: orientations of mocap bodies
NSTATE: number of state elements
PHYSICS: QPOS | QVEL | ACT
PHYSICS: TIME | QPOS | QVEL | ACT | HISTORY
FULLPHYSICS: TIME | PHYSICS | PLUGIN
USER: CTRL | QFRC_APPLIED | XFRC_APPLIED | EQ_ACTIVE | MOCAP_POS | MOCAP_QUAT | USERDATA
INTEGRATION: FULLPHYSICS | USER | WARMSTART
@@ -635,6 +652,7 @@ class State(enum.IntEnum):
QPOS = mujoco.mjtState.mjSTATE_QPOS
QVEL = mujoco.mjtState.mjSTATE_QVEL
ACT = mujoco.mjtState.mjSTATE_ACT
HISTORY = mujoco.mjtState.mjSTATE_HISTORY
WARMSTART = mujoco.mjtState.mjSTATE_WARMSTART
CTRL = mujoco.mjtState.mjSTATE_CTRL
QFRC_APPLIED = mujoco.mjtState.mjSTATE_QFRC_APPLIED
@@ -643,7 +661,7 @@ class State(enum.IntEnum):
MOCAP_POS = mujoco.mjtState.mjSTATE_MOCAP_POS
MOCAP_QUAT = mujoco.mjtState.mjSTATE_MOCAP_QUAT
NSTATE = mujoco.mjtState.mjNSTATE
PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS
PHYSICS = mujoco.mjtState.mjSTATE_PHYSICS # includes HISTORY
FULLPHYSICS = mujoco.mjtState.mjSTATE_FULLPHYSICS
USER = mujoco.mjtState.mjSTATE_USER
INTEGRATION = mujoco.mjtState.mjSTATE_INTEGRATION
@@ -861,6 +879,7 @@ class Model:
ntree: number of kinematic trees
nM: number of non-zeros in sparse inertia matrix
nC: number of non-zeros in sparse body-dof matrix
nD: number of non-zeros in sparse derivative matrix
ngeom: number of geoms
nsite: number of sites
ncam: number of cameras
@@ -898,6 +917,7 @@ class Model:
nJmom: number of non-zeros in actuator_moment
ngravcomp: number of bodies with nonzero gravcomp
nsensordata: number of elements in sensor data vector
nhistory: number of history buffer entries
opt: physics options
stat: model statistics
qpos0: qpos values at default pose (*, nq)
@@ -1129,6 +1149,9 @@ class Model:
actuator_trnid: transmission id: joint, tendon, site (nu, 2)
actuator_actadr: first activation address; -1: stateless (nu,)
actuator_actnum: number of activation variables (nu,)
actuator_history: history buffer sizes (nu, 2)
actuator_historyadr: history buffer address (nu,)
actuator_delay: delay in seconds (nu,)
actuator_ctrllimited: is control limited (nu,)
actuator_forcelimited: is force limited (nu,)
actuator_actlimited: is activation limited (nu,)
@@ -1153,12 +1176,22 @@ class Model:
sensor_dim: number of scalar outputs (nsensor,)
sensor_adr: address in sensor array (nsensor,)
sensor_cutoff: cutoff for real and positive; 0: ignore (nsensor,)
sensor_history: history buffer sizes (nsensor, 2)
sensor_historyadr: history buffer address (nsensor,)
sensor_delay: delay in seconds (nsensor,)
sensor_interval: sensor interval and phase (nsensor, 2)
plugin: globally registered plugin slot number (nplugin,)
plugin_attr: config attributes of geom plugin (nplugin, _NPLUGINATTR)
M_rownnz: number of non-zeros in each row of qM (nv,)
M_rowadr: index of each row in qM (nv,)
M_colind: column indices of non-zeros in qM (nM,)
M_rownnz: number of non-zeros in each row of M (nv,)
M_rowadr: index of each row in M (nv,)
M_colind: column indices of non-zeros in M (nC,)
mapM2M: index mapping from M (legacy) to M (CSR) (nC)
D_rownnz: non-zeros per row in D-structure (nv,)
D_rowadr: row start addresses in D-structure (nv,)
D_diag: diagonal element index within each row (nv,)
D_colind: column indices in D-structure (nD,)
mapM2D: index mapping from M to D (nD,)
mapD2M: index mapping from D to M (nC,)
flex_vertflexid: flex id for each flex vertex (nflexvert,)
warp only fields:
@@ -1187,8 +1220,8 @@ class Model:
jnt_limited_slide_hinge_adr: limited/slide/hinge jntadr
jnt_limited_ball_adr: limited/ball jntadr
body_isdofancestor: precomputed mask of which DOFs affect each body
dof_tri_row: dof lower triangle row (used in solver)
dof_tri_col: dof lower triangle col (used in solver)
dof_tri_row: dof upper triangle row (used in solver)
dof_tri_col: dof upper triangle col (used in solver)
nxn_geom_pair: collision pair geom ids [-2, ngeom-1]
nxn_geom_pair_filtered: valid collision pair geom ids
[-1, ngeom - 1]
@@ -1243,15 +1276,21 @@ class Model:
sensor_rangefinder_bodyid: bodyid for rangefinder (nrangefinder,)
taxel_vertadr: tactile sensor vertex address (nsensortaxel,)
taxel_sensorid: address for tactile sensors
qM_tiles: tiling configuration
M_tiles: tiling configuration
qLD_updates: tuple of index triples for sparse factorization
qLD_all_updates: tuple of all levels concatenated
qLD_level_offsets: tuple of start offsets for each level
qM_fullm_i: sparse mass matrix addressing
qM_fullm_j: sparse mass matrix addressing
qM_mulm_rowadr: sparse matmul row pointers
qM_mulm_col: sparse matmul column indices
qM_mulm_madr: sparse matmul matrix addresses
M_fullm_i: sparse mass matrix addressing
M_fullm_j: sparse mass matrix addressing
M_elemid: (row, col) -> CSR madr addresses; -1 if not a chain ancestor
M_fullm_upper_i: upper-triangle row indices for solver h seeding
M_fullm_upper_j: upper-triangle column indices for solver h seeding
M_fullm_upper_elemid: source elemid into M_fullm_i/M_fullm_j
qD_fullm_i: D-structure row indices for RNE derivatives
qD_fullm_j: D-structure column indices for RNE derivatives
M_mulm_rowadr: sparse matmul row pointers
M_mulm_col: sparse matmul column indices
M_mulm_madr: sparse matmul matrix addresses
"""
nq: int
@@ -1264,6 +1303,7 @@ class Model:
ntree: int
nM: int
nC: int
nD: int
ngeom: int
nsite: int
ncam: int
@@ -1301,6 +1341,7 @@ class Model:
nJmom: int
ngravcomp: int
nsensordata: int
nhistory: int
opt: Option
stat: Statistic
qpos0: array("*", "nq", float)
@@ -1532,6 +1573,9 @@ class Model:
actuator_trnid: array("nu", wp.vec2i)
actuator_actadr: array("nu", int)
actuator_actnum: array("nu", int)
actuator_history: array("nu", wp.vec2i)
actuator_historyadr: array("nu", int)
actuator_delay: array("nu", float)
actuator_ctrllimited: array("nu", bool)
actuator_forcelimited: array("nu", bool)
actuator_actlimited: array("nu", bool)
@@ -1556,12 +1600,22 @@ class Model:
sensor_dim: array("nsensor", int)
sensor_adr: array("nsensor", int)
sensor_cutoff: array("nsensor", float)
sensor_history: array("nsensor", wp.vec2i)
sensor_historyadr: array("nsensor", int)
sensor_delay: array("nsensor", float)
sensor_interval: array("nsensor", wp.vec2)
plugin: array("nplugin", int)
plugin_attr: array("nplugin", vec_pluginattr)
M_rownnz: array("nv", int)
M_rowadr: array("nv", int)
M_colind: array("nC", int)
mapM2M: array("nC", int)
D_rownnz: array("nv", int)
D_rowadr: array("nv", int)
D_diag: array("nv", int)
D_colind: array("nD", int)
mapM2D: array("nD", int)
mapD2M: array("nC", int)
flex_vertflexid: array("nflexvert", int)
# warp only fields:
callback: Callback
@@ -1635,23 +1689,32 @@ class Model:
sensor_rangefinder_bodyid: array("nrangefinder", int)
taxel_vertadr: array("nsensortaxel", int)
taxel_sensorid: wp.array[int]
qM_tiles: tuple[TileSet, ...]
M_tiles: tuple[TileSet, ...]
qLD_updates: tuple[wp.array[wp.vec3i], ...]
qLD_all_updates: wp.array[wp.vec3i]
qLD_level_offsets: wp.array[int]
qM_fullm_i: wp.array[int]
qM_fullm_j: wp.array[int]
# TODO(team): Remove M_fullm_i/j and M_elemid by iterating the M CSR layout
# directly in the solver/derivative kernels
M_fullm_i: wp.array[int]
M_fullm_j: wp.array[int]
M_elemid: wp.array2d[int] # (row, col) -> CSR madr address; -1 if col is not a chain ancestor of row
M_fullm_upper_i: wp.array[int]
M_fullm_upper_j: wp.array[int]
M_fullm_upper_elemid: wp.array[int]
qD_fullm_i: wp.array[int] # D-structure (full square) row indices for RNE derivatives
qD_fullm_j: wp.array[int] # D-structure (full square) column indices for RNE derivatives
# Gather-based sparse mul_m indices (thread per DOF, no atomics)
qM_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
qM_mulm_col: wp.array[int] # column index to gather from
qM_mulm_madr: wp.array[int] # matrix address to read
M_mulm_rowadr: wp.array[int] # start address for each row [nv+1]
M_mulm_col: wp.array[int] # column index to gather from
M_mulm_madr: wp.array[int] # matrix address to read
class ContactType(enum.IntFlag):
"""Type of contact.
CONSTRAINT: contact for constraint solver.
SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO).
Attributes:
CONSTRAINT: contact for constraint solver
SENSOR: contact for collision sensor (GEOMDIST, GEOMNORMAL, GEOMFROMTO)
"""
CONSTRAINT = 1
@@ -1673,6 +1736,8 @@ class Contact:
solimp: constraint solver impedance (naconmax, 5)
dim: contact space dimensionality: 1, 3, 4 or 6 (naconmax,)
geom: geom ids; -1 for flex (naconmax, 2)
flex: flex ids; -1 for geom (naconmax, 2)
vert: vertex ids for flex/mesh contact (naconmax, 2)
efc_address: address in efc; -1: not included (naconmax, nmaxpyramid)
worldid: world id (naconmax,)
type: ContactType (naconmax,)
@@ -1711,9 +1776,9 @@ class Constraint:
J_rowadr: row start address in colind array (nworld, 0) dense
(nworld, njmax) sparse
J_colind: column indices in J (nworld, 0, 0) dense
(nworld, 1, njmax * nv) sparse
(nworld, 1, njmax_nnz) sparse
J: constraint Jacobian (nworld, njmax_pad, nv_pad) dense
(nworld, 1, njmax * nv) sparse
(nworld, 1, njmax_nnz) sparse
pos: constraint position (equality, contact) (nworld, njmax)
margin: inclusion margin (contact) (nworld, njmax)
D: constraint mass (nworld, njmax_pad)
@@ -1722,6 +1787,20 @@ class Constraint:
frictionloss: frictionloss (friction) (nworld, njmax)
force: constraint force in constraint space (nworld, njmax)
state: constraint state (nworld, njmax_pad)
island: island ID per constraint (nworld, njmax)
itype: island constraint type (nworld, njmax)
iid: island constraint id (nworld, njmax)
iJ_rownnz: island J_rownnz (nworld, njmax)
iJ_rowadr: island J_rowadr (nworld, njmax)
iJ_colind: island J_colind (nworld, 0, 0) dense
(nworld, 1, njmax_nnz) sparse
iJ: island J (nworld, njmax, nv) dense
(nworld, 1, njmax_nnz) sparse
iD: island constraint mass (nworld, njmax_pad)
iaref: island aref (nworld, njmax)
ifrictionloss: island frictionloss (nworld, njmax)
iforce: island force (nworld, njmax)
istate: island state (nworld, njmax_pad)
warp only fields:
Ma: M*qacc (nworld, nv)
Jqvel: J*qvel (nworld, njmax)
@@ -1729,8 +1808,8 @@ class Constraint:
type: array("nworld", "njmax", int)
id: array("nworld", "njmax", int)
J_rownnz: wp.array2d[int]
J_rowadr: wp.array2d[int]
J_rownnz: array("nworld", "njmax", int)
J_rowadr: array("nworld", "njmax", int)
J_colind: wp.array3d[int]
J: wp.array3d[float]
pos: array("nworld", "njmax", float)
@@ -1741,9 +1820,22 @@ class Constraint:
frictionloss: array("nworld", "njmax", float)
force: array("nworld", "njmax", float)
state: array("nworld", "njmax_pad", int)
island: array("nworld", "njmax", int)
Ma: array("nworld", "nv", float)
Jqvel: array("nworld", "njmax", float)
itype: array("nworld", "njmax", int)
iid: array("nworld", "njmax", int)
iJ_rownnz: array("nworld", "njmax", int)
iJ_rowadr: array("nworld", "njmax", int)
iJ_colind: wp.array3d[int]
iJ: wp.array3d[float]
iD: array("nworld", "njmax_pad", float)
iaref: array("nworld", "njmax", float)
ifrictionloss: array("nworld", "njmax", float)
iforce: array("nworld", "njmax", float)
istate: array("nworld", "njmax_pad", int)
@dataclasses.dataclass
class Data:
@@ -1756,11 +1848,13 @@ class Data:
nl: number of limit constraints (nworld,)
nefc: number of constraints (nworld,)
nisland: number of constraint islands (nworld,)
nidof: total DOFs in islands (nworld,)
time: simulation time (nworld,)
energy: potential, kinetic energy (nworld, 2)
qpos: position (nworld, nq)
qvel: velocity (nworld, nv)
act: actuator activation (nworld, na)
history: history buffer for delays (nworld, nhistory)
qacc_warmstart: acceleration used for warmstart (nworld, nv)
ctrl: control (nworld, nu)
qfrc_applied: applied generalized force (nworld, nv)
@@ -1791,7 +1885,7 @@ class Data:
cinert: com-based body inertia and mass (nworld, nbody, 10)
flexvert_xpos: cartesian flex vertex positions (nworld, nflexvert, 3)
flexedge_J: edge length Jacobian (nworld, nJfe)
flexedge_length: flex edge lengths (nworld, nflexedge, 1)
flexedge_length: flex edge lengths (nworld, nflexedge)
ten_wrapadr: start address of tendon's path (nworld, ntendon)
ten_wrapnum: number of wrap points in path (nworld, ntendon)
ten_J: tendon Jacobian (nworld, nJten)
@@ -1804,10 +1898,10 @@ class Data:
moment_colind: column indices in sparse actuator_moment (nworld, nJmom)
actuator_moment: actuator moments (nworld, nJmom)
crb: com-based composite inertia and mass (nworld, nbody, 10)
qM: total inertia (nworld, nv, nv) if dense
(nworld, 1, nM) if sparse
qLD: L'*D*L factorization of M (nworld, nv, nv) if dense
M: total inertia (nworld, nv_pad, nv_pad) if dense
(nworld, 1, nC) if sparse
qLD: upper Cholesky factorization (nworld, nv, nv) if dense
L'*D*L factorization of M (nworld, 1, nC) if sparse
qLDiagInv: 1/diag(D) (nworld, nv)
flexedge_velocity: flex edge velocities (nworld, nflexedge)
ten_velocity: tendon velocities (nworld, ntendon)
@@ -1822,6 +1916,7 @@ class Data:
qfrc_passive: total passive force (nworld, nv)
subtree_linvel: linear velocity of subtree com (nworld, nbody, 3)
subtree_angmom: angular momentum about subtree com (nworld, nbody, 3)
qLU: sparse LU factorization of (M - dt*qDeriv) (nworld, 1, nD)
actuator_force: actuator force in actuation space (nworld, nu)
qfrc_actuator: actuator force (nworld, nv)
qfrc_smooth: net unconstrained force (nworld, nv)
@@ -1836,6 +1931,23 @@ class Data:
contact: contact data
efc: constraint data
tree_island: island ID per tree (-1 if unconstrained) (nworld, ntree)
dof_island: island ID per DOF (-1 if unconstrained) (nworld, nv)
island_dofadr: island start address in dof vector (nworld, ntree)
island_nv: DOFs per island (nworld, ntree)
island_nefc: constraints per island (nworld, ntree)
island_ne: equality constraints per island (nworld, ntree)
island_nf: friction constraints per island (nworld, ntree)
island_efcadr: island start address in efc vector (nworld, ntree)
map_dof2idof: global DOF -> island-local DOF (nworld, nv)
map_idof2dof: island-local DOF -> global DOF (nworld, nv)
map_efc2iefc: global EFC -> island-local EFC (nworld, njmax)
map_iefc2efc: island-local EFC -> global EFC (nworld, njmax)
dof_islandid: island ID per island-DOF (nworld, nv)
efc_islandid: island ID per island-EFC (nworld, njmax)
iqacc: island-local qacc (nworld, nv)
iqacc_smooth: island-local qacc_smooth (nworld, nv)
iqfrc_smooth: island-local qfrc_smooth (nworld, nv)
iqfrc_constraint: island-local qfrc_constraint (nworld, nv)
warp only fields:
nworld: number of worlds
@@ -1854,11 +1966,13 @@ class Data:
nl: array("nworld", int)
nefc: array("nworld", int)
nisland: array("nworld", int)
nidof: array("nworld", int)
time: array("nworld", float)
energy: array("nworld", wp.vec2)
qpos: array("nworld", "nq", float)
qvel: array("nworld", "nv", float)
act: array("nworld", "na", float)
history: array("nworld", "nhistory", float)
qacc_warmstart: array("nworld", "nv", float)
ctrl: array("nworld", "nu", float)
qfrc_applied: array("nworld", "nv", float)
@@ -1902,7 +2016,7 @@ class Data:
moment_colind: array("nworld", "nJmom", int)
actuator_moment: array("nworld", "nJmom", float)
crb: array("nworld", "nbody", vec10)
qM: wp.array3d[float]
M: wp.array3d[float]
qLD: wp.array3d[float]
qLDiagInv: array("nworld", "nv", float)
flexedge_velocity: array("nworld", "nflexedge", float)
@@ -1918,6 +2032,7 @@ class Data:
qfrc_passive: array("nworld", "nv", float)
subtree_linvel: array("nworld", "nbody", wp.vec3)
subtree_angmom: array("nworld", "nbody", wp.vec3)
qLU: array("nworld", 1, "nD", float)
actuator_force: array("nworld", "nu", float)
qfrc_actuator: array("nworld", "nv", float)
qfrc_smooth: array("nworld", "nv", float)
@@ -1930,6 +2045,23 @@ class Data:
contact: Contact
efc: Constraint
tree_island: array("nworld", "ntree", int)
dof_island: array("nworld", "nv", int)
island_dofadr: array("nworld", "ntree", int)
island_nv: array("nworld", "ntree", int)
island_nefc: array("nworld", "ntree", int)
island_ne: array("nworld", "ntree", int)
island_nf: array("nworld", "ntree", int)
island_efcadr: array("nworld", "ntree", int)
map_dof2idof: array("nworld", "nv", int)
map_idof2dof: array("nworld", "nv", int)
map_efc2iefc: array("nworld", "njmax", int)
map_iefc2efc: array("nworld", "njmax", int)
dof_islandid: array("nworld", "nv", int)
efc_islandid: array("nworld", "njmax", int)
iqacc: wp.array2d[float]
iqacc_smooth: wp.array2d[float]
iqfrc_smooth: wp.array2d[float]
iqfrc_constraint: wp.array2d[float]
# warp only fields:
nworld: int
@@ -1942,6 +2074,77 @@ class Data:
ncollision: array(1, int)
@dataclasses.dataclass
class InverseContext:
"""Workspace arrays for inverse dynamics."""
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
@dataclasses.dataclass
class IslandSolverContext:
"""Workspace arrays for island constraint solver."""
# Re-ordered workspace arrays (sized per-dof / per-constraint)
Jaref: wp.array2d[float]
jv: wp.array2d[float]
search: wp.array2d[float]
mv: wp.array2d[float]
grad: wp.array2d[float]
Mgrad: wp.array2d[float]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
h: wp.array3d[float]
# Per-island solver scalars (nworld, ntree)
cost: wp.array2d[float]
prev_cost: wp.array2d[float]
gauss: wp.array2d[float]
search_dot: wp.array2d[float]
grad_dot: wp.array2d[float]
done: wp.array2d[bool] # per-island convergence
solver_niter: wp.array2d[int] # iterations per island
beta: wp.array2d[float]
alpha: wp.array2d[float]
Ma: wp.array2d[float] # island-local Ma (nworld, nv)
@dataclasses.dataclass
class SolverContext:
"""Workspace arrays for constraint solver."""
Jaref: wp.array2d[float]
search_dot: wp.array[float]
gauss: wp.array[float]
cost: wp.array[float]
prev_cost: wp.array[float]
done: wp.array[bool]
grad: wp.array2d[float]
grad_dot: wp.array[float]
Mgrad: wp.array2d[float]
search: wp.array2d[float]
mv: wp.array2d[float]
jv: wp.array2d[float]
quad: wp.array2d[wp.vec3]
quad_gauss: wp.array[wp.vec3]
alpha: wp.array[float]
prev_grad: wp.array2d[float]
prev_Mgrad: wp.array2d[float]
beta: wp.array[float]
h: wp.array3d[float]
hfactor: wp.array3d[float]
# Incremental Hessian update (Newton only)
changed_efc_ids: wp.array2d[int]
changed_efc_count: wp.array[int]
@dataclasses.dataclass
class RenderContext:
"""Context for rendering.
@@ -1952,6 +2155,8 @@ class RenderContext:
cam_id_map: camera id map
use_textures: whether to use textures
use_shadows: whether to use shadows
use_ambient_lighting: whether to use ambient lighting
background_color: background color
use_precomputed_rays: whether to use precomputed rays
bvh_ngeom: number of geometries in the BVH
enabled_geom_ids: enabled geometry ids
@@ -1971,7 +2176,10 @@ class RenderContext:
flex_bvh_id: per-flex BVH ids
flex_group_root: per-flex group roots (nworld x n_flex_bvh)
flex_render_smooth: whether to render flex meshes smoothly
flex_dim: flex dimension per flex (1D/2D/3D)
bvh_nflexgeom: number of flex geometries in the BVH
flex_dim_np: flex dimension per flex (1D/2D/3D)
flex_geom_flexid: map from flex geom ID to flex ID
flex_geom_edgeid: map from flex geom ID to flex edge ID
bvh: scene BVH
bvh_id: scene BVH id
lower: lower bounds
@@ -1993,6 +2201,11 @@ class RenderContext:
render_skybox: whether to shade missed rays with the MuJoCo skybox texture
skybox_tex_id: index into textures of the skybox (MuJoCo tex_type == SKYBOX), -1 if none
skybox_face_width: pixel width of one skybox cube face (0 if no skybox)
enable_backface_culling: drop primitive ray hits whose normal faces away
from the ray (i.e. the ray origin is inside the geom). Matches MuJoCo's
mesh-ray rule. When False, the renderer reports inner-surface hits, which
is faster but causes a camera placed inside a geom to render that geom's
back wall.
"""
nrender: int
@@ -2000,6 +2213,7 @@ class RenderContext:
cam_id_map: array("ncam", int)
use_textures: bool
use_shadows: bool
use_ambient_lighting: bool
background_color: wp.uint32
use_precomputed_rays: bool
render_skybox: bool
@@ -2045,3 +2259,4 @@ class RenderContext:
render_seg: array("ncam", bool)
znear: float
total_rays: int
enable_backface_culling: bool
mjx/mujoco/mjx/third_party/mujoco_warp/_src/util_pkg.py
Vendored
+3 -3
View File
@@ -23,7 +23,7 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
"""Parse a version string into comparable components.
Dot-separated components form the version. Hyphen-separated suffixes (e.g.,
"-foo3") are treated as local build identifiers and stripped before
"-newton") are treated as local build identifiers and stripped before
parsing. Each component is wrapped in a tuple: (0, int) for numeric parts,
(-1, str) for non-numeric. A (0, 0) sentinel is appended so that stable
releases sort above pre-release suffixes during Python tuple comparison
@@ -32,13 +32,13 @@ def _parse_version(version_str: str) -> tuple[tuple[int, int | str], ...]:
Args:
version_str: Version string like "3.5.0", "3.5.0.dev869102767", or
"3.9.0-foo3".
"3.9.0-newton".
Returns:
Tuple of (type_order, value) pairs for comparison, where type_order is 0
for integers and -1 for strings, followed by a (0, 0) sentinel.
"""
# Strip local build identifier (e.g., "3.9.0-foo3" -> "3.9.0")
# Strip local build identifier (e.g., "3.9.0-newton" -> "3.9.0")
version_str = version_str.split("-", 1)[0]
parts = version_str.split(".")
return tuple([(0, int(p)) if p.isdigit() else (-1, p) for p in parts] + [(0, 0)])
mjx/mujoco/mjx/third_party/mujoco_warp/pyproject.toml
Vendored
+2 -2
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name="mujoco-warp"
version = "3.8.0.1"
version = "3.9.0.1"
# TODO(team): create a distribution list
authors = [
{name = "Newton Developers", email = "mujoco@deepmind.com"},
@@ -31,7 +31,7 @@ dependencies = [
"etils[epath]",
"mujoco>=3.8.0",
"numpy",
"warp-lang>=1.12",
"warp-lang>=1.13",
]
[[tool.uv.index]]
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py
Vendored
+7 -5
View File
@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
# Launch the kernel.
wp._src.context.runtime.core.wp_cuda_launch_kernel(
device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
)
def _create_jax_warp_primitive():
import jax # noqa: PLC0415
from jax._src.interpreters import batching # noqa: PLC0415
from jax.interpreters import batching # noqa: PLC0415
from jax.interpreters import mlir # noqa: PLC0415
from jax.interpreters.mlir import ir # noqa: PLC0415
from jaxlib.hlo_helpers import custom_call # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
# Figure out the number of outputs.
wp_kernel = _registered_kernels[params["kernel"]]
output_count = len(wp_kernel.adj.args) - len(args)
shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
size = shape[dim]
args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
# Create the batched primitive.
return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
def base_type_to_jax_ir(warp_dtype):
warp_to_jax_dict = {
wp.float16: ir.F16Type.get(),
wp.bfloat16: ir.BF16Type.get(),
wp.float32: ir.F32Type.get(),
wp.float64: ir.F64Type.get(),
wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
def base_type_is_compatible(warp_type, jax_ir_type):
jax_ir_to_warp = {
"f16": wp.float16,
"bf16": wp.bfloat16,
"f32": wp.float32,
"f64": wp.float64,
"i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
# Figure out the types and shapes of the input arrays.
arg_strings = []
operand_layouts = []
for actual, warg in zip(args, wp_kernel.adj.args):
for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
wtype = warg.type
rtt = ir.RankedTensorType(actual.type)
mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py
Vendored
+14 -13
View File
@@ -8,8 +8,8 @@ import ctypes
import inspect
import threading
import traceback
from collections.abc import Callable
from enum import IntEnum
from typing import Callable
import jax
@@ -25,6 +25,7 @@ from warp._src.types import (
type_size_in_bytes,
type_to_warp,
)
from warp._src.utils import warn
from .xla_ffi import *
@@ -249,19 +250,18 @@ class FfiKernel:
# check dtype
if input_value.dtype != input_arg.jax_scalar_type:
raise TypeError(
f"Invalid data type for array argument '{input_arg.name}',"
f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
)
# check ndim
if input_value.ndim != input_arg.jax_ndim:
raise TypeError(
f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
)
# check inner dims
for d in range(input_arg.dtype_ndim):
if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
raise TypeError(
f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
)
else:
# make sure scalar is not a traced variable, should be static
@@ -455,6 +455,7 @@ class FfiKernel:
hooks.forward_smem_bytes,
kernel_params,
stream,
None, # apic_info
)
else:
wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
except Exception as e:
print(traceback.format_exc())
return create_ffi_error(
call_frame.contents.api,
XLA_FFI_Error_Code.UNKNOWN,
f"FFI callback error: {type(e).__name__}: {e}",
call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
)
@@ -617,7 +616,9 @@ class FfiCallable:
self.input_output_aliases = input_output_aliases
# register the callback
FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
FFI_CCALLFUNC = ctypes.CFUNCTYPE(
ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
)
self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1071,7 +1072,7 @@ class FfiCallable:
call_desc.capture = capture
else:
# not capturing or on CPU
# not capturing
self.func(*arg_list)
except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
try:
gi.zero_()
except Exception as e:
wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
warn(f"Failed to zero gradient array: {e}", stacklevel=2)
raise e
# NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
non_static_inputs, output_vals_tuple = residuals
input_vals = list(non_static_inputs)
for i, v in zip(static_args, nondiff_vals):
for i, v in zip(static_args, nondiff_vals, strict=True):
input_vals.insert(i, v)
# Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(
out_dims_map = {}
param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
for name, val in zip(differentiable_input_names, non_static_inputs):
for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
ann = param_ann.get(name)
if ann is None:
continue
mjx/mujoco/mjx/warp/__init__.py
+6
View File
@@ -23,6 +23,7 @@ if not typing.TYPE_CHECKING:
warp: Any = None
mujoco_warp: Any = None
mjwp_types: Any = None
mjwp_io: Any = None
WARP_INSTALLED: bool = False
# pylint: disable=g-import-not-at-top
@@ -35,6 +36,7 @@ if not typing.TYPE_CHECKING:
try:
from mujoco.mjx.third_party import mujoco_warp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.third_party.mujoco_warp._src import io as mjwp_io
except (ImportError, RuntimeError) as e:
print('Failed to import mujoco_warp:', e)
pass
@@ -67,7 +69,11 @@ else:
def BlockDim(self, *args, **kwargs): # pylint: disable=invalid-name
pass
class _MjwpIoStub:
ENABLE_ISLANDS = True
WARP_INSTALLED: bool = True
warp: Any = _WpStub()
mujoco_warp: Any = _MjwpStub()
mjwp_types: Any = _MjwpTypesStub()
mjwp_io: Any = _MjwpIoStub()
mjx/mujoco/mjx/warp/bvh.py
+7 -4
View File
@@ -14,17 +14,19 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
import warp as wp
from mujoco.mjx._src import types
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
from mujoco.mjx.warp.render_context import _MJX_RENDER_CONTEXT_BUFFERS
from mujoco.mjx.warp.render_context import RenderContextPytree
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -48,6 +50,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _refit_bvh_shim(
# Model
mjx/mujoco/mjx/warp/collision_driver.py
+7 -4
View File
@@ -14,15 +14,17 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
from mujoco.mjx._src import types
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _collision_shim(
# Model
mjx/mujoco/mjx/warp/forward.py
+2587 -2122
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/warp/forward_test.py
+2 -2
View File
@@ -176,8 +176,8 @@ class ForwardTest(parameterized.TestCase):
qm = np.zeros((m.nv, m.nv))
mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
# mjwarp adds padding to qM
tu.assert_eq(qm, dx._impl.qM[: m.nv, : m.nv], 'qM')
# mjwarp adds padding to M
tu.assert_eq(qm, dx._impl.M[: m.nv, : m.nv], 'M')
# qLD is fused in a cholesky factorize and solve, and not written to.
tu.assert_contact_eq(d, dx, worldid=i)
mjx/mujoco/mjx/warp/smooth.py
+7 -4
View File
@@ -14,15 +14,17 @@
# ==============================================================================
"""DO NOT EDIT. This file is auto-generated."""
import dataclasses
import functools
import jax
from mujoco.mjx._src import types
from mujoco.mjx.warp import ffi
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
import warp as wp
from mujoco.mjx._src import types
import mujoco.mjx.third_party.mujoco_warp as mjwarp
from mujoco.mjx.third_party.mujoco_warp._src import types as mjwp_types
from mujoco.mjx.warp import ffi
_m = mjwarp.Model(
**{f.name: None for f in dataclasses.fields(mjwarp.Model) if f.init}
@@ -46,6 +48,7 @@ _cb = mjwp_types.Callback(
**{f.name: None for f in dataclasses.fields(mjwp_types.Callback) if f.init}
)
@ffi.format_args_for_warp
def _kinematics_shim(
# Model
mjx/mujoco/mjx/warp/smooth_test.py
+2 -2
View File
@@ -112,7 +112,7 @@ class SmoothTest(parameterized.TestCase):
d = mujoco.MjData(m)
mx = mjx.put_model(m, impl='warp')
mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
rng = jax.random.PRNGKey(0)
dx = mjx.make_data(m, impl='warp')
@@ -188,7 +188,7 @@ class SmoothTest(parameterized.TestCase):
batch_size = 7
d = mujoco.MjData(m)
mx = mjx.put_model(m, impl='warp')
mx = mx.replace(_impl=mx._impl.replace(qM_tiles=()))
mx = mx.replace(_impl=mx._impl.replace(M_tiles=()))
worldids = jp.arange(batch_size)
dx_batch = jax.vmap(functools.partial(tu.make_data, m))(worldids)
mjx/mujoco/mjx/warp/types.py
+197 -21
View File
@@ -78,6 +78,25 @@ class BlockDim:
"""Block dimension 'block_dim' settings for wp.launch_tiled.
TODO(team): experimental and may be removed
Attributes:
segmented_sort: segmented sort block dimension (collision_driver)
euler_dense: Euler dense block dimension (forward)
actuator_velocity: actuator velocity block dimension (forward)
ray: ray block dimension (ray)
contact_sort: contact sort block dimension (sensor)
energy_vel_kinetic: energy velocity kinetic block dimension (sensor)
cholesky_factorize: Cholesky factorize block dimension (smooth)
cholesky_solve: Cholesky solve block dimension (smooth)
cholesky_factorize_solve: Cholesky factorize and solve block dimension (smooth)
solve_LD_sparse_fused: solve LD sparse fused block dimension (smooth)
update_gradient_cholesky: update gradient Cholesky block dimension (solver)
update_gradient_cholesky_blocked: update gradient Cholesky blocked block dimension (solver)
update_gradient_JTDAJ_sparse: update gradient JTDAJ sparse block dimension (solver)
update_gradient_JTDAJ_dense: update gradient JTDAJ dense block dimension (solver)
linesearch_iterative: linesearch iterative block dimension (solver)
contact_jac_tiled: contact Jacobian tiled block dimension (solver)
qderiv_actuator_dense: qderiv actuator dense block dimension (derivative)
"""
actuator_velocity: int
@@ -135,9 +154,26 @@ class OptionWarp(PyTreeNode):
class ModelWarp(PyTreeNode):
"""Derived fields from Model."""
D_colind: np.ndarray
D_diag: np.ndarray
D_rowadr: np.ndarray
D_rownnz: np.ndarray
M_colind: np.ndarray
M_elemid: np.ndarray
M_fullm_i: np.ndarray
M_fullm_j: np.ndarray
M_fullm_upper_elemid: np.ndarray
M_fullm_upper_i: np.ndarray
M_fullm_upper_j: np.ndarray
M_mulm_col: np.ndarray
M_mulm_madr: np.ndarray
M_mulm_rowadr: np.ndarray
M_rowadr: np.ndarray
M_rownnz: np.ndarray
M_tiles: Tuple[TileSet, ...]
actuator_delay: np.ndarray
actuator_history: np.ndarray
actuator_historyadr: np.ndarray
actuator_trntype_body_adr: np.ndarray
block_dim: BlockDim
body_branch_start: np.ndarray
@@ -203,6 +239,8 @@ class ModelWarp(PyTreeNode):
jnt_limited_slide_hinge_adr: np.ndarray
light_bodyid: np.ndarray
light_targetbodyid: np.ndarray
mapD2M: np.ndarray
mapM2D: np.ndarray
mapM2M: np.ndarray
mat_texrepeat: jax.Array
max_ten_J_rownnz: int
@@ -248,22 +286,22 @@ class ModelWarp(PyTreeNode):
oct_coeff: np.ndarray
plugin: np.ndarray
plugin_attr: np.ndarray
qD_fullm_i: np.ndarray
qD_fullm_j: np.ndarray
qLD_all_updates: np.ndarray
qLD_level_offsets: np.ndarray
qLD_updates: Tuple[np.ndarray, ...]
qM_fullm_i: np.ndarray
qM_fullm_j: np.ndarray
qM_mulm_col: np.ndarray
qM_mulm_madr: np.ndarray
qM_mulm_rowadr: np.ndarray
qM_tiles: Tuple[TileSet, ...]
rangefinder_sensor_adr: np.ndarray
sensor_acc_adr: np.ndarray
sensor_adr_to_contact_adr: np.ndarray
sensor_collision_start_adr: np.ndarray
sensor_contact_adr: np.ndarray
sensor_delay: np.ndarray
sensor_e_kinetic: bool
sensor_e_potential: bool
sensor_history: np.ndarray
sensor_historyadr: np.ndarray
sensor_interval: np.ndarray
sensor_limitfrc_adr: np.ndarray
sensor_limitpos_adr: np.ndarray
sensor_limitvel_adr: np.ndarray
@@ -299,6 +337,7 @@ class ModelWarp(PyTreeNode):
class DataWarp(PyTreeNode):
"""Derived fields from Data."""
M: jax.Array
actuator_moment: jax.Array
actuator_velocity: jax.Array
cacc: jax.Array
@@ -322,6 +361,8 @@ class DataWarp(PyTreeNode):
contact__vert: jax.Array
contact__worldid: jax.Array
crb: jax.Array
dof_island: jax.Array
dof_islandid: jax.Array
efc__D: jax.Array
efc__J: jax.Array
efc__J_colind: jax.Array
@@ -332,19 +373,46 @@ class DataWarp(PyTreeNode):
efc__aref: jax.Array
efc__force: jax.Array
efc__frictionloss: jax.Array
efc__iD: jax.Array
efc__iJ: jax.Array
efc__iJ_colind: jax.Array
efc__iJ_rowadr: jax.Array
efc__iJ_rownnz: jax.Array
efc__iaref: jax.Array
efc__id: jax.Array
efc__iforce: jax.Array
efc__ifrictionloss: jax.Array
efc__iid: jax.Array
efc__island: jax.Array
efc__istate: jax.Array
efc__itype: jax.Array
efc__margin: jax.Array
efc__pos: jax.Array
efc__state: jax.Array
efc__type: jax.Array
efc__vel: jax.Array
efc_islandid: jax.Array
energy: jax.Array
flexedge_J: jax.Array
flexedge_length: jax.Array
flexedge_velocity: jax.Array
flexvert_xpos: jax.Array
iqacc: jax.Array
iqacc_smooth: jax.Array
iqfrc_constraint: jax.Array
iqfrc_smooth: jax.Array
island_dofadr: jax.Array
island_efcadr: jax.Array
island_ne: jax.Array
island_nefc: jax.Array
island_nf: jax.Array
island_nv: jax.Array
light_xdir: jax.Array
light_xpos: jax.Array
map_dof2idof: jax.Array
map_efc2iefc: jax.Array
map_idof2dof: jax.Array
map_iefc2efc: jax.Array
moment_colind: jax.Array
moment_rowadr: jax.Array
moment_rownnz: jax.Array
@@ -355,6 +423,7 @@ class DataWarp(PyTreeNode):
ne: jax.Array
nefc: jax.Array
nf: jax.Array
nidof: jax.Array
nisland: jax.Array
njmax: int
njmax_nnz: int
@@ -363,7 +432,7 @@ class DataWarp(PyTreeNode):
nworld: int
qLD: jax.Array
qLDiagInv: jax.Array
qM: jax.Array
qLU: jax.Array
qfrc_damper: jax.Array
qfrc_spring: jax.Array
solver_niter: jax.Array
@@ -400,6 +469,7 @@ DATA_NON_VMAP = {
'nacon',
'naconmax',
'ncollision',
'nidof',
'njmax',
'njmax_nnz',
'njmax_pad',
@@ -433,6 +503,7 @@ batching.register_vmappable(DataWarp, int, int, _to_elt, _from_elt, None)
_NDIM = {
'Data': {
'M': 3,
'act': 2,
'act_dot': 2,
'actuator_force': 2,
@@ -466,6 +537,8 @@ _NDIM = {
'crb': 3,
'ctrl': 2,
'cvel': 3,
'dof_island': 2,
'dof_islandid': 2,
'efc__D': 2,
'efc__J': 3,
'efc__J_colind': 3,
@@ -476,12 +549,25 @@ _NDIM = {
'efc__aref': 2,
'efc__force': 2,
'efc__frictionloss': 2,
'efc__iD': 2,
'efc__iJ': 3,
'efc__iJ_colind': 3,
'efc__iJ_rowadr': 2,
'efc__iJ_rownnz': 2,
'efc__iaref': 2,
'efc__id': 2,
'efc__iforce': 2,
'efc__ifrictionloss': 2,
'efc__iid': 2,
'efc__island': 2,
'efc__istate': 2,
'efc__itype': 2,
'efc__margin': 2,
'efc__pos': 2,
'efc__state': 2,
'efc__type': 2,
'efc__vel': 2,
'efc_islandid': 2,
'energy': 2,
'eq_active': 2,
'flexedge_J': 2,
@@ -490,8 +576,23 @@ _NDIM = {
'flexvert_xpos': 3,
'geom_xmat': 4,
'geom_xpos': 3,
'history': 2,
'iqacc': 2,
'iqacc_smooth': 2,
'iqfrc_constraint': 2,
'iqfrc_smooth': 2,
'island_dofadr': 2,
'island_efcadr': 2,
'island_ne': 2,
'island_nefc': 2,
'island_nf': 2,
'island_nv': 2,
'light_xdir': 3,
'light_xpos': 3,
'map_dof2idof': 2,
'map_efc2iefc': 2,
'map_idof2dof': 2,
'map_iefc2efc': 2,
'mocap_pos': 3,
'mocap_quat': 3,
'moment_colind': 2,
@@ -504,6 +605,7 @@ _NDIM = {
'ne': 1,
'nefc': 1,
'nf': 1,
'nidof': 1,
'nisland': 1,
'njmax': 0,
'njmax_nnz': 0,
@@ -512,7 +614,7 @@ _NDIM = {
'nworld': 0,
'qLD': 3,
'qLDiagInv': 2,
'qM': 3,
'qLU': 3,
'qacc': 2,
'qacc_smooth': 2,
'qacc_warmstart': 2,
@@ -555,9 +657,23 @@ _NDIM = {
'xquat': 3,
},
'Model': {
'D_colind': 1,
'D_diag': 1,
'D_rowadr': 1,
'D_rownnz': 1,
'M_colind': 1,
'M_elemid': 2,
'M_fullm_i': 1,
'M_fullm_j': 1,
'M_fullm_upper_elemid': 1,
'M_fullm_upper_i': 1,
'M_fullm_upper_j': 1,
'M_mulm_col': 1,
'M_mulm_madr': 1,
'M_mulm_rowadr': 1,
'M_rowadr': 1,
'M_rownnz': 1,
'M_tiles': -1,
'actuator_acc0': 2,
'actuator_actadr': 1,
'actuator_actearly': 1,
@@ -569,6 +685,7 @@ _NDIM = {
'actuator_cranklength': 2,
'actuator_ctrllimited': 1,
'actuator_ctrlrange': 3,
'actuator_delay': 1,
'actuator_dynprm': 3,
'actuator_dyntype': 1,
'actuator_forcelimited': 1,
@@ -576,6 +693,8 @@ _NDIM = {
'actuator_gainprm': 3,
'actuator_gaintype': 1,
'actuator_gear': 3,
'actuator_history': 2,
'actuator_historyadr': 1,
'actuator_lengthrange': 3,
'actuator_trnid': 2,
'actuator_trntype': 1,
@@ -768,6 +887,8 @@ _NDIM = {
'light_poscom0': 3,
'light_targetbodyid': 1,
'light_type': 2,
'mapD2M': 1,
'mapM2D': 1,
'mapM2M': 1,
'mat_rgba': 3,
'mat_texid': 3,
@@ -796,6 +917,7 @@ _NDIM = {
'mesh_vertnum': 1,
'mocap_bodyid': 1,
'nC': 0,
'nD': 0,
'nJfe': 0,
'nJmom': 0,
'nJten': 0,
@@ -820,6 +942,7 @@ _NDIM = {
'ngravcomp': 0,
'nhfield': 0,
'nhfielddata': 0,
'nhistory': 0,
'njnt': 0,
'nlight': 0,
'nmat': 0,
@@ -896,15 +1019,11 @@ _NDIM = {
'pair_solreffriction': 3,
'plugin': 1,
'plugin_attr': 2,
'qD_fullm_i': 1,
'qD_fullm_j': 1,
'qLD_all_updates': 2,
'qLD_level_offsets': 1,
'qLD_updates': -1,
'qM_fullm_i': 1,
'qM_fullm_j': 1,
'qM_mulm_col': 1,
'qM_mulm_madr': 1,
'qM_mulm_rowadr': 1,
'qM_tiles': -1,
'qpos0': 2,
'qpos_spring': 2,
'rangefinder_sensor_adr': 1,
@@ -915,9 +1034,13 @@ _NDIM = {
'sensor_contact_adr': 1,
'sensor_cutoff': 1,
'sensor_datatype': 1,
'sensor_delay': 1,
'sensor_dim': 1,
'sensor_e_kinetic': 0,
'sensor_e_potential': 0,
'sensor_history': 2,
'sensor_historyadr': 1,
'sensor_interval': 2,
'sensor_intprm': 2,
'sensor_limitfrc_adr': 1,
'sensor_limitpos_adr': 1,
@@ -1015,6 +1138,7 @@ _NDIM = {
}
_BATCH_DIM = {
'Data': {
'M': True,
'act': True,
'act_dot': True,
'actuator_force': True,
@@ -1048,6 +1172,8 @@ _BATCH_DIM = {
'crb': True,
'ctrl': True,
'cvel': True,
'dof_island': True,
'dof_islandid': True,
'efc__D': True,
'efc__J': True,
'efc__J_colind': True,
@@ -1058,12 +1184,25 @@ _BATCH_DIM = {
'efc__aref': True,
'efc__force': True,
'efc__frictionloss': True,
'efc__iD': True,
'efc__iJ': True,
'efc__iJ_colind': True,
'efc__iJ_rowadr': True,
'efc__iJ_rownnz': True,
'efc__iaref': True,
'efc__id': True,
'efc__iforce': True,
'efc__ifrictionloss': True,
'efc__iid': True,
'efc__island': True,
'efc__istate': True,
'efc__itype': True,
'efc__margin': True,
'efc__pos': True,
'efc__state': True,
'efc__type': True,
'efc__vel': True,
'efc_islandid': True,
'energy': True,
'eq_active': True,
'flexedge_J': True,
@@ -1072,8 +1211,23 @@ _BATCH_DIM = {
'flexvert_xpos': True,
'geom_xmat': True,
'geom_xpos': True,
'history': True,
'iqacc': True,
'iqacc_smooth': True,
'iqfrc_constraint': True,
'iqfrc_smooth': True,
'island_dofadr': True,
'island_efcadr': True,
'island_ne': True,
'island_nefc': True,
'island_nf': True,
'island_nv': True,
'light_xdir': True,
'light_xpos': True,
'map_dof2idof': True,
'map_efc2iefc': True,
'map_idof2dof': True,
'map_iefc2efc': True,
'mocap_pos': True,
'mocap_quat': True,
'moment_colind': True,
@@ -1086,6 +1240,7 @@ _BATCH_DIM = {
'ne': True,
'nefc': True,
'nf': True,
'nidof': False,
'nisland': True,
'njmax': False,
'njmax_nnz': False,
@@ -1094,7 +1249,7 @@ _BATCH_DIM = {
'nworld': False,
'qLD': True,
'qLDiagInv': True,
'qM': True,
'qLU': True,
'qacc': True,
'qacc_smooth': True,
'qacc_warmstart': True,
@@ -1137,9 +1292,23 @@ _BATCH_DIM = {
'xquat': True,
},
'Model': {
'D_colind': False,
'D_diag': False,
'D_rowadr': False,
'D_rownnz': False,
'M_colind': False,
'M_elemid': False,
'M_fullm_i': False,
'M_fullm_j': False,
'M_fullm_upper_elemid': False,
'M_fullm_upper_i': False,
'M_fullm_upper_j': False,
'M_mulm_col': False,
'M_mulm_madr': False,
'M_mulm_rowadr': False,
'M_rowadr': False,
'M_rownnz': False,
'M_tiles': False,
'actuator_acc0': True,
'actuator_actadr': False,
'actuator_actearly': False,
@@ -1151,6 +1320,7 @@ _BATCH_DIM = {
'actuator_cranklength': True,
'actuator_ctrllimited': False,
'actuator_ctrlrange': True,
'actuator_delay': False,
'actuator_dynprm': True,
'actuator_dyntype': False,
'actuator_forcelimited': False,
@@ -1158,6 +1328,8 @@ _BATCH_DIM = {
'actuator_gainprm': True,
'actuator_gaintype': False,
'actuator_gear': True,
'actuator_history': False,
'actuator_historyadr': False,
'actuator_lengthrange': True,
'actuator_trnid': False,
'actuator_trntype': False,
@@ -1350,6 +1522,8 @@ _BATCH_DIM = {
'light_poscom0': True,
'light_targetbodyid': False,
'light_type': True,
'mapD2M': False,
'mapM2D': False,
'mapM2M': False,
'mat_rgba': True,
'mat_texid': True,
@@ -1378,6 +1552,7 @@ _BATCH_DIM = {
'mesh_vertnum': False,
'mocap_bodyid': False,
'nC': False,
'nD': False,
'nJfe': False,
'nJmom': False,
'nJten': False,
@@ -1402,6 +1577,7 @@ _BATCH_DIM = {
'ngravcomp': False,
'nhfield': False,
'nhfielddata': False,
'nhistory': False,
'njnt': False,
'nlight': False,
'nmat': False,
@@ -1478,15 +1654,11 @@ _BATCH_DIM = {
'pair_solreffriction': True,
'plugin': False,
'plugin_attr': False,
'qD_fullm_i': False,
'qD_fullm_j': False,
'qLD_all_updates': False,
'qLD_level_offsets': False,
'qLD_updates': False,
'qM_fullm_i': False,
'qM_fullm_j': False,
'qM_mulm_col': False,
'qM_mulm_madr': False,
'qM_mulm_rowadr': False,
'qM_tiles': False,
'qpos0': True,
'qpos_spring': True,
'rangefinder_sensor_adr': False,
@@ -1497,9 +1669,13 @@ _BATCH_DIM = {
'sensor_contact_adr': False,
'sensor_cutoff': False,
'sensor_datatype': False,
'sensor_delay': False,
'sensor_dim': False,
'sensor_e_kinetic': False,
'sensor_e_potential': False,
'sensor_history': False,
'sensor_historyadr': False,
'sensor_interval': False,
'sensor_intprm': False,
'sensor_limitfrc_adr': False,
'sensor_limitpos_adr': False,
mjx/pyproject.toml
+5 -5
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name="mujoco-mjx"
version = "3.9.0"
version = "3.10.0"
authors = [
{name = "Google DeepMind", email = "mujoco@deepmind.com"},
]
@@ -30,14 +30,14 @@ dependencies = [
"etils[epath]",
"jax",
"jaxlib",
"mujoco>=3.9.0.dev0",
"mujoco>=3.10.0.dev0",
"scipy",
"trimesh",
]
[project.optional-dependencies]
warp = [
"warp-lang==1.12.1",
"warp-lang==1.13.0",
]
dev = [
"isort",
@@ -50,9 +50,9 @@ 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.9.0"
Documentation = "https://mujoco.readthedocs.io/en/3.10.0"
Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
Changelog = "https://mujoco.readthedocs.io/en/3.9.0/changelog.html"
Changelog = "https://mujoco.readthedocs.io/en/3.10.0/changelog.html"
[tool.isort]
force_single_line = true
mjx/requirements.txt
+8 -5
View File
@@ -64,11 +64,14 @@ jaxlib==0.4.38; python_version >= '3.10' and python_version <= '3.13' and sys_pl
--hash=sha256:248cca3771ebf24b070f49701364ceada33e6139445b06c782cca5ac5ad92bf4 \
--hash=sha256:2ce77ba8cda9259a4bca97afc1c722e4291a6c463a63f8d372c6edc85117d625 \
--hash=sha256:4103db0b3a38a5dc132741237453c24d8547290a22079ba1b577d6c88c95300a
pip==25.2 \
--hash=sha256:578283f006390f85bb6282dffb876454593d637f5d1be494b5202ce4877e71f2 \
--hash=sha256:6d67a2b4e7f14d8b31b8b52648866fa717f45a1eb70e83002f4331d07e953717
pytest==8.3.3 \
--hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2
pip==26.1 \
--hash=sha256:4e8486d821d814b77319acb7b9e8bf5a4ee7590a643e7cb21029f209be8573c1 \
--hash=sha256:81e13ebcca3ffa8cc85e4deff5c27e1ee26dea0aa7fc2f294a073ac208806ff3
pytest==9.0.3 \
--hash=sha256:2c5efc453d45394fdd706ade797c0a81091eccd1d6e4bccfcd476e2b8e0ab5d9 \
--hash=sha256:b86ada508af81d19edeb213c681b1d48246c1a91d304c6c81a427674c17eb91c
pygments==2.20.0 \
--hash=sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176
pytest-xdist==3.6.1 \
--hash=sha256:9ed4adfb68a016610848639bb7e02c9352d5d9f03d04809919e2dafc3be4cca7
scipy==1.17.0; python_version >= '3.13' \
plugin/actuator/pid.h
+1 -1
View File
@@ -21,7 +21,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
namespace mujoco::plugin::actuator {
plugin/elasticity/cable.cc
+1 -1
View File
@@ -18,7 +18,7 @@
#include <optional>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mujoco.h>
#include "cable.h"
plugin/elasticity/cable.h
+1 -1
View File
@@ -20,7 +20,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
plugin/sdf/bolt.cc
+1 -1
View File
@@ -19,7 +19,7 @@
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "bolt.h"
plugin/sdf/bolt.h
+1 -1
View File
@@ -19,7 +19,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
#include "sdf.h"
plugin/sdf/bowl.cc
+1 -1
View File
@@ -17,7 +17,7 @@
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "bowl.h"
plugin/sdf/bowl.h
+1 -1
View File
@@ -19,7 +19,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
#include "sdf.h"
plugin/sdf/gear.cc
+1 -1
View File
@@ -18,7 +18,7 @@
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
plugin/sdf/gear.h
+1 -1
View File
@@ -19,7 +19,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
#include "sdf.h"
plugin/sdf/nut.cc
+1 -1
View File
@@ -19,7 +19,7 @@
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "nut.h"
plugin/sdf/nut.h
+1 -1
View File
@@ -19,7 +19,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjtype.h>
#include <mujoco/mjvisualize.h>
#include "sdf.h"

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