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

Merge main branch from mujoco

Browse Source
This commit is contained in:
Vidur
2025-06-03 13:00:27 -05:00
parent bf3843210f 3eb31f56cd
commit 35fdc279fe
289 changed files with 29274 additions and 11525 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build.yml
+65 -41
View File
@@ -24,8 +24,24 @@ jobs:
fail-fast: false
matrix:
include:
- os: ubuntu-24.04
label: "ubuntu-24.04-gcc-14"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-14
-DCMAKE_CXX_COMPILER:STRING=g++-14
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
tmpdir: "/tmp"
- os: ubuntu-24.04
label: "ubuntu-24.04-gcc-13"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-13
-DCMAKE_CXX_COMPILER:STRING=g++-13
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-gcc-12"
label: "ubuntu-22.04-gcc-12"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-12
@@ -33,7 +49,7 @@ jobs:
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-gcc-11"
label: "ubuntu-22.04-gcc-11"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-11
@@ -41,23 +57,47 @@ jobs:
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-gcc-10"
label: "ubuntu-22.04-gcc-10"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-10
-DCMAKE_CXX_COMPILER:STRING=g++-10
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-gcc-9"
- os: ubuntu-24.04
label: "ubuntu-24.04-clang-18"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=gcc-9
-DCMAKE_CXX_COMPILER:STRING=g++-9
-DCMAKE_EXE_LINKER_FLAGS:STRING=-Wl,--no-as-needed
-DCMAKE_C_COMPILER:STRING=clang-18
-DCMAKE_CXX_COMPILER:STRING=clang++-18
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-24.04
label: "ubuntu-24.04-clang-17"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-17
-DCMAKE_CXX_COMPILER:STRING=clang++-17
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-24.04
label: "ubuntu-24.04-clang-16"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-16
-DCMAKE_CXX_COMPILER:STRING=clang++-16
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-clang-14"
label: "ubuntu-22.04-clang-15"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-15
-DCMAKE_CXX_COMPILER:STRING=clang++-15
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-22.04
label: "ubuntu-22.04-clang-14"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-14
@@ -65,49 +105,33 @@ jobs:
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-22.04
additional_label: "-clang-13"
label: "ubuntu-22.04-clang-13"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-13
-DCMAKE_CXX_COMPILER:STRING=clang++-13
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-20.04
additional_label: "-clang-12"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-12
-DCMAKE_CXX_COMPILER:STRING=clang++-12
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-20.04
additional_label: "-clang-11"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-11
-DCMAKE_CXX_COMPILER:STRING=clang++-11
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: ubuntu-20.04
additional_label: "-clang-10"
cmake_args: >-
-G Ninja
-DCMAKE_C_COMPILER:STRING=clang-10
-DCMAKE_CXX_COMPILER:STRING=clang++-10
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: macos-13
- os: macos-15
label: "macos-15-arm64"
cmake_args: >-
-G Ninja
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: windows-2022
- os: macos-15-large
label: "macos-15-x86_64"
cmake_args: >-
-DCMAKE_SYSTEM_VERSION="10.0.22621.0"
-G Ninja
-DMUJOCO_HARDEN:BOOL=ON
tmpdir: "/tmp"
- os: windows-2025
label: "windows-2025"
cmake_args: >-
-DCMAKE_SYSTEM_VERSION="10.0.26100.0"
cmake_build_args: "-- -m"
tmpdir: "C:/Temp"
name: "${{ matrix.os }}${{ matrix.additional_label }}"
name: "${{ matrix.label }}"
runs-on: ${{ matrix.os }}
steps:
@@ -179,7 +203,7 @@ jobs:
cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
cp lib/libsdf_plugin.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
- name: Copy plugins (Windows)
if: ${{ runner.os == 'Windows' }}
working-directory: build
@@ -273,7 +297,7 @@ jobs:
CHATMSG_AUTHOR_NAME: ${{ github.event.head_commit.author.name }}
CHATMSG_AUTHOR_EMAIL: ${{ github.event.head_commit.author.email }}
CHATMSG_COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
CHATMSG_JOB_ID: ${{ matrix.os }}${{ matrix.additional_label }}
CHATMSG_JOB_ID: ${{ matrix.label }}
if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API_URL != '' }}
run: |
CHATMSG="$(cat <<-'EOF' | python3
@@ -288,7 +312,7 @@ jobs:
email=env('CHATMSG_AUTHOR_EMAIL'),
msg=env('CHATMSG_COMMIT_MESSAGE').replace('```', '')
)
text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name}<{email}>\n\n{msg}```'.format(**data)
text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name} <{email}>\n\n{msg}```'.format(**data)
print(json.dumps({'text' : text}))
EOF
)" &&
CMakeLists.txt
+13 -1
View File
@@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
project(
mujoco
VERSION 3.3.1
VERSION 3.3.3
DESCRIPTION "MuJoCo Physics Simulator"
HOMEPAGE_URL "https://mujoco.org"
)
@@ -42,6 +42,14 @@ option(MUJOCO_BUILD_EXAMPLES "Build samples for MuJoCo" ON)
option(MUJOCO_BUILD_SIMULATE "Build simulate library for MuJoCo" ON)
option(MUJOCO_BUILD_TESTS "Build tests for MuJoCo" ON)
option(MUJOCO_TEST_PYTHON_UTIL "Build and test utility libraries for Python bindings" ON)
option(MUJOCO_BUILD_USD_PLUGINS "Build OpenUSD plugins" OFF)
# USD libs to compile against.
set(MUJOCO_USD_ALLOWED_TARGET_VALUES "USD" "Houdini")
set(MUJOCO_USD_TARGET "USD" CACHE STRING "Select the USD target for the project.")
set_property(CACHE MUJOCO_USD_TARGET
PROPERTY STRINGS ${MUJOCO_USD_ALLOWED_TARGET_VALUES}
)
if(APPLE AND (MUJOCO_BUILD_EXAMPLES OR MUJOCO_BUILD_SIMULATE))
enable_language(OBJC)
@@ -181,6 +189,10 @@ if(MUJOCO_BUILD_EXAMPLES)
add_subdirectory(sample)
endif()
if(MUJOCO_BUILD_USD_PLUGINS)
add_subdirectory(src/experimental/usd/plugins)
endif()
if(BUILD_TESTING AND MUJOCO_BUILD_TESTS)
enable_testing()
add_subdirectory(test)
README.md
+2
View File
@@ -53,6 +53,8 @@ running on Google Colab:
- The **introductory** tutorial teaches MuJoCo basics:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/tutorial.ipynb)
- The **Model Editing** tutorial shows how to create and edit models procedurally:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb)
- The **rollout** tutorial shows how to use the multithreaded `rollout` module:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/rollout.ipynb)
- The **LQR** tutorial synthesizes a linear-quadratic controller, balancing a
cmake/MujocoDependencies.cmake
+2 -2
View File
@@ -39,12 +39,12 @@ set(MUJOCO_DEP_VERSION_qhull
CACHE STRING "Version of `qhull` to be fetched."
)
set(MUJOCO_DEP_VERSION_Eigen3
66f7f51b7e069d0a03a21157fa60b24aece69aeb
464c1d097891a1462ab28bf8bb763c1683883892
CACHE STRING "Version of `Eigen3` to be fetched."
)
set(MUJOCO_DEP_VERSION_abseil
9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0
d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1
CACHE STRING "Version of `abseil` to be fetched."
)
dist/mujoco.rc
Vendored
+4 -4
View File
@@ -1,6 +1,6 @@
1 VERSIONINFO
FILEVERSION 3,3,1,0
PRODUCTVERSION 3,3,1,0
FILEVERSION 3,3,3,0
PRODUCTVERSION 3,3,3,0
FILEOS 0x4
FILETYPE 0x1
{
@@ -9,9 +9,9 @@ FILETYPE 0x1
BLOCK "040904b0"
{
VALUE "ProductName", "MuJoCo"
VALUE "ProductVersion", "3.3.1"
VALUE "ProductVersion", "3.3.3"
VALUE "FileDescription", "MuJoCo"
VALUE "FileVersion", "3.3.1"
VALUE "FileVersion", "3.3.3"
VALUE "InternalName", "mujoco.dll"
VALUE "OriginalFilename", "mujoco.dll"
VALUE "CompanyName", "Google DeepMind"
dist/simulate.rc
Vendored
+4 -4
View File
@@ -1,8 +1,8 @@
MUJOCO ICON "mujoco.ico"
1 VERSIONINFO
FILEVERSION 3,3,1,0
PRODUCTVERSION 3,3,1,0
FILEVERSION 3,3,3,0
PRODUCTVERSION 3,3,3,0
FILEOS 0x4
FILETYPE 0x1
{
@@ -11,9 +11,9 @@ FILETYPE 0x1
BLOCK "040904b0"
{
VALUE "ProductName", "MuJoCo"
VALUE "ProductVersion", "3.3.1"
VALUE "ProductVersion", "3.3.3"
VALUE "FileDescription", "MuJoCo"
VALUE "FileVersion", "3.3.1"
VALUE "FileVersion", "3.3.3"
VALUE "InternalName", "simulate.exe"
VALUE "OriginalFilename", "simulate.exe"
VALUE "CompanyName", "Google DeepMind"
doc/APIreference/APIglobals.rst
+2 -2
View File
@@ -448,7 +448,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
- 200
- The number of iterations where solver statistics can be stored in ``mjData.solver``. This array is used
to store diagnostic information about each iteration of the constraint solver.
The actual number of iterations is given by ``mjData.solver_iter``.
The actual number of iterations is given by ``mjData.solver_niter``.
* - ``mjNISLAND``
- 20
- The number of islands for which solver statistics can be stored in ``mjData.solver``. This array is
@@ -517,7 +517,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
- Maximum number of UI rectangles.
Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
* - ``mjVERSION_HEADER``
- 331
- 333
- The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software
version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a
number with the same meaning but for the compiled library.
doc/APIreference/APItypes.rst
+32 -12
View File
@@ -165,6 +165,17 @@ values are used in ``m->cam_mode`` and ``m->light_mode``.
.. mujoco-include:: mjtCamLight
.. _mjtLightType:
mjtLightType
~~~~~~~~~~~~
The type of a light source describing how its position, orientation and other properties will interact with the
objects in the scene. These values are used in ``m->light_type``.
.. mujoco-include:: mjtLightType
.. _mjtTexture:
mjtTexture
@@ -357,6 +368,26 @@ last argument to :ref:`mj_local2global`.
.. mujoco-include:: mjtSameFrame
.. _mjtFlexSelf:
mjtFlexSelf
~~~~~~~~~~~~
Types of flex self-collisions midphase.
.. mujoco-include:: mjtFlexSelf
.. _mjtSDFType:
mjtSDFType
~~~~~~~~~~~
Formulas used to combine SDFs when calling mjc_distance and mjc_gradient.
.. mujoco-include:: mjtSDFType
.. _tyDataEnums:
Data
@@ -883,7 +914,7 @@ mjSolverStat
This is the data structure holding information about one solver iteration. ``mjData.solver`` is a preallocated array
of mjSolverStat data structures, one for each iteration of the solver, up to a maximum of mjNSOLVER. The actual number
of solver iterations is given by ``mjData.solver_iter``.
of solver iterations is given by ``mjData.solver_niter``.
.. mujoco-include:: mjSolverStat
@@ -967,17 +998,6 @@ This structure contains everything needed to render the 3D scene in OpenGL.
.. mujoco-include:: mjvScene
.. _mjvSceneState:
mjvSceneState
~~~~~~~~~~~~~
This structure contains the portions of :ref:`mjModel` and :ref:`mjData` that are required for
various ``mjv_*`` functions.
.. mujoco-include:: mjvSceneState
.. _mjvFigure:
mjvFigure
doc/APIreference/functions.rst
+58 -99
View File
@@ -51,6 +51,15 @@ Compile :ref:`mjSpec` to :ref:`mjModel`. A spec can be edited and compiled multi
:ref:`mjModel` instance that takes the edits into account.
If compilation fails, :ref:`mj_compile` returns ``NULL``; the error can be read with :ref:`mjs_getError`.
.. _mj_copyBack:
`mj_copyBack <#mj_copyBack>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mj_copyBack
Copy real-valued arrays from model to spec, returns 1 on success.
.. _mj_recompile:
`mj_recompile <#mj_recompile>`__
@@ -459,7 +468,7 @@ Multiply vector by (inertia matrix)^(1/2).
.. mujoco-include:: mj_addM
Add inertia matrix to destination matrix.
Destination can be sparse uncompressed, or dense when all int* are NULL
Destination can be sparse or dense when all int* are NULL.
.. _mj_applyFT:
@@ -748,8 +757,7 @@ Compare forward and inverse dynamics, save results in fwdinv.
Sub components
^^^^^^^^^^^^^^
These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
that the user will need to call them.
These are sub-components of the simulation pipeline, called internally from the components above.
.. _mj_sensorPos:
@@ -886,6 +894,18 @@ Compute actuator transmission lengths and moments.
Run composite rigid body inertia algorithm (CRB).
.. _mj_makeM:
`mj_makeM <#mj_makeM>`__
~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mj_makeM
Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due
to :ref:`tendon armature<tendon-spatial-armature>`. The joint-space inertia matrix is stored in both ``mjData.qM`` and
``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row,
respectively).
.. _mj_factorM:
`mj_factorM <#mj_factorM>`__
@@ -1364,6 +1384,15 @@ If the model buffer is unallocated the initial configuration will not be set.
Copy mjData.
m is only required to contain the size fields from MJMODEL_INTS.
.. _mjv_copyData:
`mjv_copyData <#mjv_copyData>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_copyData
Copy mjData, skip large arrays not required for visualization.
.. _mj_resetData:
`mj_resetData <#mj_resetData>`__
@@ -1995,15 +2024,6 @@ Rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_
Move camera with mouse; action is mjtMouse.
.. _mjv_moveCameraFromState:
`mjv_moveCameraFromState <#mjv_moveCameraFromState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_moveCameraFromState
Move camera with mouse given a scene state; action is mjtMouse.
.. _mjv_movePerturb:
`mjv_movePerturb <#mjv_movePerturb>`__
@@ -2013,15 +2033,6 @@ Move camera with mouse given a scene state; action is mjtMouse.
Move perturb object with mouse; action is mjtMouse.
.. _mjv_movePerturbFromState:
`mjv_movePerturbFromState <#mjv_movePerturbFromState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_movePerturbFromState
Move perturb object with mouse given a scene state; action is mjtMouse.
.. _mjv_moveModel:
`mjv_moveModel <#mjv_moveModel>`__
@@ -2164,15 +2175,6 @@ Free abstract scene.
Update entire scene given model state.
.. _mjv_updateSceneFromState:
`mjv_updateSceneFromState <#mjv_updateSceneFromState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_updateSceneFromState
Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings.
.. _mjv_copyModel:
`mjv_copyModel <#mjv_copyModel>`__
@@ -2182,42 +2184,6 @@ Update entire scene from a scene state, return the number of new mjWARN_VGEOMFUL
Copy mjModel, skip large arrays not required for abstract visualization.
.. _mjv_defaultSceneState:
`mjv_defaultSceneState <#mjv_defaultSceneState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_defaultSceneState
Set default scene state.
.. _mjv_makeSceneState:
`mjv_makeSceneState <#mjv_makeSceneState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_makeSceneState
Allocate resources and initialize a scene state object.
.. _mjv_freeSceneState:
`mjv_freeSceneState <#mjv_freeSceneState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_freeSceneState
Free scene state.
.. _mjv_updateSceneState:
`mjv_updateSceneState <#mjv_updateSceneState>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjv_updateSceneState
Update a scene state from model and data.
.. _mjv_addGeoms:
`mjv_addGeoms <#mjv_addGeoms>`__
@@ -3808,41 +3774,14 @@ Free all pointers with ``mju_free()``.
Attachment
^^^^^^^^^^
.. _mjs_attachBody:
.. _mjs_attach:
`mjs_attachBody <#mjs_attachBody>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
`mjs_attach <#mjs_attach>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_attachBody
.. mujoco-include:: mjs_attach
Attach child body to a parent frame, return the attached body if success or NULL otherwise.
.. _mjs_attachFrame:
`mjs_attachFrame <#mjs_attachFrame>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_attachFrame
Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
.. _mjs_attachToSite:
`mjs_attachToSite <#mjs_attachToSite>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_attachToSite
Attach child body to a parent site, return the attached body if success or NULL otherwise.
.. _mjs_attachFrameToSite:
`mjs_attachFrameToSite <#mjs_attachFrameToSite>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_attachFrameToSite
Attach child frame to a parent site, return the attached frame if success or NULL otherwise.
Attach child to a parent, return the attached element if success or NULL otherwise.
.. _mjs_detachBody:
@@ -4428,6 +4367,15 @@ Get string contents.
Get double array contents and optionally its size.
.. _mjs_getPluginAttributes:
`mjs_getPluginAttributes <#mjs_getPluginAttributes>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_getPluginAttributes
Get plugin attributes.
.. _SpecUtilities:
Spec utilities
@@ -4448,7 +4396,7 @@ Set element's default.
.. mujoco-include:: mjs_setFrame
Set element's enclosing frame.
Set element's enclosing frame, return 0 on success.
.. _mjs_resolveOrientation:
@@ -4477,6 +4425,17 @@ Transform body into a frame.
Set user payload, overriding the existing value for the specified key if present.
.. _mjs_setUserValueWithCleanup:
`mjs_setUserValueWithCleanup <#mjs_setUserValueWithCleanup>`__
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. mujoco-include:: mjs_setUserValueWithCleanup
Set user payload, overriding the existing value for the specified key if
present. This version differs from mjs_setUserValue in that it takes a
cleanup function that will be called when the user payload is deleted.
.. _mjs_getUserValue:
`mjs_getUserValue <#mjs_getUserValue>`__
doc/APIreference/functions_override.rst
+8 -2
View File
@@ -113,8 +113,14 @@ Integrates the simulation state using an implicit-in-velocity integrator (either
.. _Subcomponents:
These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
that the user will need to call them.
These are sub-components of the simulation pipeline, called internally from the components above.
.. _mj_makeM:
Compute the composite rigid body inertia with :ref:`mj_crb`, add terms due
to :ref:`tendon armature<tendon-spatial-armature>`. The joint-space inertia matrix is stored in both ``mjData.qM`` and
``mjData.M``. These arrays represent the same quantity using different layouts (parent-based and compressed sparse row,
respectively).
.. _mj_factorM:
doc/XMLreference.rst
+225 -67
View File
@@ -175,11 +175,10 @@ replicating 200 times, suffixes will be ``000, 001, ...`` etc). All referencing
and namespaced appropriately. Detailed examples of models using replicate can be found in the
`model/replicate/ <https://github.com/google-deepmind/mujoco/tree/main/model/replicate>`__ directory.
There is a caveat concerning :ref:`keyframes<keyframe>` when using replicate. Since :ref:`mjs_attachFrame` is used to
There is a caveat concerning :ref:`keyframes<keyframe>` when using replicate. Since :ref:`mjs_attach` is used to
self-attach multiple times the enclosed kinematic tree, if this tree contains further :ref:`attach<body-attach>`
elements, keyframes will not be replicated nor namespaced by :ref:`replicate<replicate>`, but they will be attached and
namespaced once by the innermost call of :ref:`mjs_attachFrame` or :ref:`mjs_attachBody`. See the limitations discussed
in :ref:`attach<body-attach>`.
namespaced once by the innermost call of :ref:`mjs_attach`. See the limitations discussed in :ref:`attach<body-attach>`.
.. _replicate-count:
@@ -833,6 +832,10 @@ has any effect. The settings here are global and apply to the entire model.
necessary to adjust this attribute and the geom-specific groups so as to exclude world geoms from the inertial
computation.
.. _compiler-saveinertial:
:at:`saveinertial`: :at-val:`[false, true], "false"`
If set to "true", the compiler will save explicit :ref:`inertial <body-inertial>` clauses for all bodies.
.. _compiler-lengthrange:
@@ -1461,16 +1464,19 @@ still be specified here but this functionality is now deprecated and will be rem
:el-prefix:`asset/` |-| **texture** (*)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
| This element creates a texture asset, which is then referenced from a :ref:`material <asset-material>` asset, which is
finally referenced from a model element that needs to be textured. MuJoCo provides access to the texture mapping
mechanism in OpenGL. Texture coordinates are generated automatically in GL_OBJECT_PLANE mode, using either 2D or cube
mapping. MIP maps are always enabled in GL_LINEAR_MIPMAP_LINEAR mode. The texture color is combined with the object
color in GL_MODULATE mode. The texture data can be loaded from PNG files, with provisions for loading cube and skybox
textures. Alternatively the data can be generated by the compiler as a procedural texture. Because different texture
types require different parameters, only a subset of the attributes below are used for any given texture.
| A second file format is supported for loading textures, in addition to PNG. If the file name extension is
different from .png or .PNG, or if the ``content_type`` attribute is set to ``image/vnd.mujoco.texture``, then MuJoCo
assumes that the texture is in this format. This is a custom binary file format, containing the following data:
This element creates a texture asset, which is then referenced from a :ref:`material <asset-material>` asset, which
is finally referenced from a model element that needs to be textured.
The texture data can be loaded from files or can be generated by the compiler as a procedural texture. Because
different texture types require different parameters, only a subset of the attributes below are used for any given
texture. Provisions are provided for loading cube and skybox textures from individual image files.
Currently, three file formats are supported for loading textures: PNG, KTX, and a custom MuJoCo texture format. The
loader will use the extension of the file name to determine which format to use, defaulting to the custom format if
the extension is not recognized. Alternatively, the content_type attribute can be used to specify the format
explicitly. Only ``image/png``, ``image/ktx``, or ``image/vnd.mujoco.texture`` are supported.
The custom MuJoCo format is assumed to be a binary file containing the following data:
.. code:: Text
@@ -1523,13 +1529,19 @@ still be specified here but this functionality is now deprecated and will be rem
with the texrepeat attribute of :ref:`material <asset-material>`. The data can be loaded from a single file or
created procedurally.
.. _asset-texture-colorspace:
:at:`colorspace`: :at-val:`[auto, linear, sRGB], "auto"`
This attribute determines the color space of the texture. The default value ``auto`` means that the color space will
be determined from the image file itself. If no color space is defined in the file, then ``linear`` is assumed.
.. _asset-texture-content_type:
:at:`content_type`: :at-val:`string, optional`
If the file attribute is specified, then this sets the
`Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`_ (formerly known as MIME types) of the
file to be loaded. Any filename extensions will be ignored. Currently ``image/png`` and ``image/vnd.mujoco.texture``
are supported.
file to be loaded. Any filename extensions will be ignored. Currently ``image/png``, ``image/ktx``, and
``image/vnd.mujoco.texture`` are supported.
.. _asset-texture-file:
@@ -1777,11 +1789,11 @@ properties are grouped together.
.. _asset-material-rgba:
:at:`rgba`: :at-val:`real(4), "1 1 1 1"`
Color and transparency of the material. All components should be in the range [0 1]. Note that textures are applied
in GL_MODULATE mode, meaning that the texture color and the color specified here are multiplied component-wise. Thus
the default value of "1 1 1 1" has the effect of leaving the texture unchanged. When the material is applied to a
model element which defines its own local rgba attribute, the local definition has precedence. Note that this "local"
definition could in fact come from a defaults class. The remaining material properties always apply.
Color and transparency of the material. All components should be in the range [0 1]. Note that the texture color (if
assigned) and the color specified here are multiplied component-wise. Thus the default value of "1 1 1 1" has the
effect of leaving the texture unchanged. When the material is applied to a model element which defines its own local
rgba attribute, the local definition has precedence. Note that this "local" definition could in fact come from a
defaults class. The remaining material properties always apply.
.. _material-layer:
@@ -2151,14 +2163,28 @@ rotations as unit quaternions.
corresponding to mjModel.qpos_spring is also used to compute the spring reference lengths of all tendons, stored in
mjModel.tendon_lengthspring. This is because :ref:`tendons <tendon>` can also have springs.
.. image:: images/XMLreference/armature.gif
:width: 40%
:align: right
:class: only-light
:target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml
.. image:: images/XMLreference/armature_dark.gif
:width: 40%
:align: right
:class: only-dark
:target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/armature_equivalence.xml
.. _body-joint-armature:
:at:`armature`: :at-val:`real, "0"`
Additional inertia associated with movement of the joint that is not due to body mass. This added inertia is usually
due to a rotor (a.k.a `armature <https://en.wikipedia.org/wiki/Armature_(electrical)>`__) spinning faster than the
joint itself due to a geared transmission; in this case the added inertia is known as "reflected inertia" and its
value is the rotational inertia of the spinning element multiplied by the square of the gear ratio. The value applies
to all degrees of freedom created by this joint.
joint itself due to a geared transmission. In the illustration, we compare (*left*) a 2-dof system with an armature
body (purple box), coupled with a gear ratio of :math:`3` to the pendulum using a :ref:`joint
equality<equality-joint>` constraint, and (*right*) a simple 1-dof pendulum with an equivalent :at:`armature`.
Because the gear ratio appears twice, multiplying both forces and lengths, the effect is known as "reflected
inertia" and the equivalent value is the inertia of the spinning body multiplied by the *square of the gear ratio*,
in this case :math:`9=3^2`. The value applies to all degrees of freedom created by this joint.
Besides increasing the realism of joints with geared transmission, positive :at:`armature` significantly improves
simulation stability, even for small values, and is a recommended possible fix when encountering stability issues.
@@ -2849,11 +2875,15 @@ and the +Y axis points up. Thus the frame position and orientation are the key a
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This element creates a light, which moves with the body where it is defined. To create a fixed light, define it in the
world body. The lights created here are in addition to the default headlight which is always defined and is adjusted via
the :ref:`visual <visual>` element. MuJoCo relies on the standard lighting model in OpenGL (fixed functionality)
augmented with shadow mapping. The effects of lights are additive, thus adding a light always makes the scene brighter.
The maximum number of lights that can be active simultaneously is 8, counting the headlight. The light is shining along
the direction specified by the dir attribute. It does not have a full spatial frame with three orthogonal axes.
world body. The lights created here are in addition to the headlight which is always defined and is configured via the
:ref:`visual <visual>` element. Lights shine along the direction specified by the dir attribute. They do not have
a full spatial frame with three orthogonal axes.
By default, MuJoCo uses the standard OpenGL (fixed functional) Phong lighting model for its rendering, with augmented
with shadow mapping. (See the OpenGL documentation for more information, including details about various attributes.)
MJCF also supports alternative lighting models (e.g. physically-based rendering) by providing additional attributes.
Attributes may be applied or ignored depending on the lighting model being used.
.. _body-light-name:
@@ -2879,10 +2909,17 @@ the direction specified by the dir attribute. It does not have a full spatial fr
This is identical to the target attribute of :ref:`camera <body-camera>` above. It specifies which body should be
targeted in "targetbody" and "targetbodycom" modes.
.. _body-light-type:
:at:`type`: :at-val:`[spot, directional, point, image], "spot"`
Determines the type of light. Note that some light types may not be supported by some renderers (e.g. only spot and
directional lights are supported by the default native renderer).
.. _body-light-directional:
:at:`directional`: :at-val:`[false, true], "false"`
The light is directional if this attribute is "true", otherwise it is a spotlight.
This is a deprecated legacy attribute. Please use :ref:`light <body-light-type>` type instead. If set to "true", and
no type is specified, this will change the light type to be directional.
.. _body-light-castshadow:
@@ -2902,12 +2939,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr
these clipping planes bound the cone or box shadow volume in the light direction. As a result, some shadows
(especially those very close to the light) may be clipped.
.. _body-light-bulbradius:
:at:`radius`: :at-val:`real, "0.02"`
Radius of the light, affects shadow softness. This attribute has no effect in MuJoCo's native renderer, but it can be
useful when rendering scenes with an external renderer.
.. _body-light-active:
:at:`active`: :at-val:`[false, true], "true"`
@@ -2924,11 +2955,50 @@ the direction specified by the dir attribute. It does not have a full spatial fr
:at:`dir`: :at-val:`real(3), "0 0 -1"`
Direction of the light.
.. _body-light-diffuse:
:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"`
The color of the light. For the Phong (default) lighting model, this defines the diffuse color of
the light.
.. _body-light-texture:
:at:`texture`: :at-val:`string, optional`
The texture to use for image-based lighting. This is unused by the default Phong lighting model.
.. _body-light-intensity:
:at:`intensity`: :at-val:`real, "0.0"`
The intensity of the light source, measured in candela, used for physically-based lighting models.
This is unused by the default Phong lighting model.
.. _body-light-ambient:
:at:`ambient`: :at-val:`real(3), "0 0 0"`
The ambient color of the light, used by the default Phong lighting model.
.. _body-light-specular:
:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"`
The specular color of the light, used by the default Phong lighting model.
.. _body-light-range:
:at:`range`: :at-val:`real, "10.0"`
The effective range of the light. Objects further than this distance from the light position
will not be illuminated by this light. This only applies to spotlights.
.. _body-light-bulbradius:
:at:`bulbradius`: :at-val:`real, "0.02"`
The radius of the light source which can affect shadow softness depending on the
renderer. This only applies to spotlights.
.. _body-light-attenuation:
:at:`attenuation`: :at-val:`real(3), "1 0 0"`
These are the constant, linear and quadratic attenuation coefficients in OpenGL. The default corresponds to no
attenuation. See the OpenGL documentation for more information on this and all other OpenGL-related properties.
These are the constant, linear and quadratic attenuation coefficients for Phong lighting.
The default corresponds to no attenuation.
.. _body-light-cutoff:
@@ -2940,22 +3010,6 @@ the direction specified by the dir attribute. It does not have a full spatial fr
:at:`exponent`: :at-val:`real, "10"`
Exponent for spotlights. This setting controls the softness of the spotlight cutoff.
.. _body-light-ambient:
:at:`ambient`: :at-val:`real(3), "0 0 0"`
The ambient color of the light.
.. _body-light-diffuse:
:at:`diffuse`: :at-val:`real(3), "0.7 0.7 0.7"`
The diffuse color of the light.
.. _body-light-specular:
:at:`specular`: :at-val:`real(3), "0.3 0.3 0.3"`
The specular color of the light.
.. _body-composite:
:el-prefix:`body/` |-| **composite** (*)
@@ -3003,8 +3057,13 @@ cable, which produces an inextensible chain of bodies connected with ball joints
.. _body-composite-offset:
:at:`offset`: :at-val:`real(3), "0 0 0"`
It specifies a 3D offset from the center of the parent body to the center of the grid of elements. The offset is
expressed in the local coordinate frame of the parent body.
It specifies a 3D offset from the center of the parent body to the center of the first body of the cable. The offset
is expressed in the local coordinate frame of the parent body.
.. _body-composite-quat:
:at:`quat`: :at-val:`real(4), "1 0 0 0"`
It specifies a quaternion that rotates the first body frame. The quaternion is expressed in the parent body frame.
.. _body-composite-vertex:
@@ -3532,6 +3591,7 @@ saving the XML:
.. _flexcomp-contact-internal:
.. _flexcomp-contact-selfcollide:
.. _flexcomp-contact-vertcollide:
.. _flexcomp-contact-activelayers:
.. _flexcomp-contact-contype:
.. _flexcomp-contact-conaffinity:
@@ -3545,8 +3605,8 @@ saving the XML:
.. _flexcomp-contact-gap:
.. |body/flexcomp/contact attrib list| replace::
:at:`internal`, :at:`selfcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`, :at:`condim`,
:at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap`
:at:`internal`, :at:`selfcollide`, :at:`vertcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`,
:at:`condim`, :at:`priority`, :at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap`
|body/flexcomp/contact attrib list|
Same as in :ref:`flex/contact<flex-contact>`. All attributes are passed through to the automatically-generated flex.
@@ -3588,9 +3648,10 @@ element is used to adjust the properties of all edges in the flex.
.. _flexcomp-elasticity-poisson:
.. _flexcomp-elasticity-damping:
.. _flexcomp-elasticity-thickness:
.. _flexcomp-elasticity-elastic2d:
.. |body/flexcomp/elasticity attrib list| replace::
:at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness`
:at:`young`, :at:`poisson`, :at:`damping`, :at:`thickness`, :at:`elastic2d`
|body/flexcomp/elasticity attrib list|
Same as in :ref:`flex/elasticity<flex-elasticity>`.
@@ -3698,8 +3759,9 @@ all attachments will appear in the saved XML file.
.. _body-attach-body:
:at:`body`: :at-val:`string, required`
Name of the body in the sub-model to attach here. The body and its subtree will be attached.
:at:`body`: :at-val:`string, optional`
Name of the body in the sub-model to attach here. The body and its subtree will be attached. If this attribute is not
specified, the contents of the world body will be attached in a new :ref:`frame<body-frame>`.
.. _body-attach-prefix:
@@ -4030,6 +4092,11 @@ stress-strain relationship.. See also :ref:`deformable <CDeformable>` objects.
This thickness can be set equal to 2 times the :ref:`radius <deformable-flex-radius>` in order to match the geometry,
but is exposed separately since the radius might be constrained by considerations related to collision detection.
.. _flex-elasticity-elastic2d:
:at:`elastic2d`: :at-val:`[none, bend, stretch, both], "none"`
Elastic contribution to passive forces of 2D flexes. "none": none, "bend": bending only, "stretch": stretching only,
"both": bending and stretching.
.. _flex-contact:
@@ -4063,6 +4130,13 @@ extensions specific to flexes.
**sap** in 1D and 2D, and **bvh** in 3D. Which strategy performs better depends on the specifics of the model. The
automatic setting is just a simple rule which we have found to perform well in general.
.. _flex-contact-vertcollide:
:at:`vertcollide`: :at-val:`[true, false], "false"`
Enables or disables vertex collisions. if **true**, spherical geoms are added at the vertices of flex, with radius
equal to the radius of the flex. These geoms can collide with other geoms and are not visible by default. If
**false**, no additional geoms are added.
.. _flex-contact-activelayers:
:at:`activelayers`: :at-val:`int(1), "1"`
@@ -4607,12 +4681,28 @@ length X, as in the clip on the right of `this example model
solver. If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler <compiler>`, length limits
will be enabled if range is defined.
.. _tendon-spatial-actuatorfrclimited:
:at:`actuatorfrclimited`: :at-val:`[false, true, auto], "auto"`
This attribute specifies whether actuator forces acting on the tendon should be clamped. See :ref:`CForceRange` for
details. This attribute interacts with the :ref:`actuatorfrcrange<tendon-spatial-actuatorfrcrange>` attribute. If
this attribute is "false", actuator force clamping is disabled. If it is "true", actuator force clamping is enabled.
If this attribute is "auto", and :at:`autolimits` is set in :ref:`compiler <compiler>`, actuator force clamping will
be enabled if :at:`actuatorfrcrange` is defined.
.. _tendon-spatial-range:
:at:`range`: :at-val:`real(2), "0 0"`
Range of allowed tendon lengths. Setting this attribute without specifying :at:`limited` is an error, unless
:at:`autolimits` is set in :ref:`compiler <compiler>`.
.. _tendon-spatial-actuatorfrcrange:
:at:`actuatorfrcrange`: :at-val:`real(2), "0 0"`
Range for clamping total actuator forces acting on this tendon. See :ref:`CForceRange` for details. The compiler
expects the lower bound to be nonpositive and the upper bound to be nonnegative. |br| Setting this attribute without
specifying :at:`actuatorfrclimited` is an error if :at:`compiler-autolimits` is "false".
.. _tendon-spatial-solreflimit:
.. _tendon-spatial-solimplimit:
@@ -4683,6 +4773,31 @@ length X, as in the clip on the right of `this example model
joint damping which is integrated implicitly by the Euler method, tendon damping is not integrated implicitly, thus
joint damping should be used if possible.
.. image:: images/XMLreference/tendon_armature.gif
:width: 30%
:align: right
:class: only-light
:target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml
.. image:: images/XMLreference/tendon_armature_dark.gif
:width: 30%
:align: right
:class: only-dark
:target: https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/core_smooth/ten_armature_1_compare.xml
.. _tendon-spatial-armature:
:at:`armature`: :at-val:`real, "0"`
Inertia associated with changes in tendon length. Setting this attribute to a positive value :math:`m` adds a kinetic
energy term :math:`\frac{1}{2}mv^2`, where :math:`v` is the tendon velocity. Tendon inertia is most valuable
when modeling the :ref:`armature<body-joint-armature>` inertia in a linear actuator which contains a spinning element
or the inertial motion of a fluid in a linear hydraulic actuator. In the illustration, we compare (*left*) a 3-dof
system with a "tendon" implemented with a rotational joint and a slider joint with
:ref:`armature<body-joint-armature>`, attached to the world with a :ref:`connect<equality-connect>` constraint and
(*right*) an equivalent 1-dof model with an armature-bearing tendon. Like joint :ref:`armature<body-joint-armature>`,
this added inertia is only associated with changes in tendon length, and would not affect the dynamics of a moving
fixed-length tendon. Because the tendon Jacobian :math:`J` is position-dependent, tendon armature leads to an
additional bias-force term :math:`c = m J \dot{J}^T \dot{q}`.
.. _tendon-spatial-user:
:at:`user`: :at-val:`real(nuser_tendon), "0 0 ..."`
@@ -4772,8 +4887,12 @@ as above.
.. _tendon-fixed-limited:
.. _tendon-fixed-actuatorfrclimited:
.. _tendon-fixed-range:
.. _tendon-fixed-actuatorfrcrange:
.. _tendon-fixed-solreflimit:
.. _tendon-fixed-solimplimit:
@@ -4792,6 +4911,8 @@ as above.
.. _tendon-fixed-damping:
.. _tendon-fixed-armature:
.. _tendon-fixed-user:
.. |tendon/fixed attrib list| replace::
@@ -6335,6 +6456,33 @@ joint or when a single actuator acts on multiple joints. See :ref:`CForceRange`
The joint where actuator forces will be sensed. The sensor output is copied from ``mjData.qfrc_actuator``.
.. _sensor-tendonactuatorfrc:
:el-prefix:`sensor/` |-| **tendonactuatorfrc** (*)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This element creates an actuator force sensor, measured at a tendon. The quantity being sensed is the total force
contributed by all actuators to a single tendon. This type of sensor is important when multiple actuators act on a
single tendon. See :ref:`CForceRange` for details.
.. _sensor-tendonactuatorfrc-name:
.. _sensor-tendonactuatorfrc-noise:
.. _sensor-tendonactuatorfrc-cutoff:
.. _sensor-tendonactuatorfrc-user:
:at:`name`, :at:`noise`, :at:`cutoff`, :at:`user`
See :ref:`CSensor`.
.. _sensor-tendonactuatorfrc-tendon:
:at:`tendon`: :at-val:`string, required`
The tendon where actuator forces will be sensed.
.. _sensor-ballquat:
:el-prefix:`sensor/` |-| **ballquat** (*)
@@ -8152,7 +8300,7 @@ if omitted.
.. _default-light-dir:
.. _default-light-bulbradius:
.. _default-light-type:
.. _default-light-directional:
@@ -8160,18 +8308,24 @@ if omitted.
.. _default-light-active:
.. _default-light-diffuse:
.. _default-light-intensity:
.. _default-light-ambient:
.. _default-light-specular:
.. _default-light-bulbradius:
.. _default-light-range:
.. _default-light-attenuation:
.. _default-light-cutoff:
.. _default-light-exponent:
.. _default-light-ambient:
.. _default-light-diffuse:
.. _default-light-specular:
.. _default-light-mode:
:el-prefix:`default/` |-| **light** (?)
@@ -8227,8 +8381,12 @@ if omitted.
.. _default-tendon-limited:
.. _default-tendon-actuatorfrclimited:
.. _default-tendon-range:
.. _default-tendon-actuatorfrcrange:
.. _default-tendon-solreflimit:
.. _default-tendon-solimplimit:
doc/XMLschema.rst
+45 -24
View File
@@ -54,7 +54,7 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`discardvisual<compiler-discardvisual>` | :ref:`usethread<compiler-usethread>` | :ref:`fusestatic<compiler-fusestatic>` | :ref:`inertiafromgeom<compiler-inertiafromgeom>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`inertiagrouprange<compiler-inertiagrouprange>` | :ref:`assetdir<compiler-assetdir>` | :ref:`alignfree<compiler-alignfree>` | | |
| | | | :ref:`inertiagrouprange<compiler-inertiagrouprange>` | :ref:`saveinertial<compiler-saveinertial>` | :ref:`assetdir<compiler-assetdir>` | :ref:`alignfree<compiler-alignfree>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| compiler |br| |_| |L| | | .. table:: |
@@ -153,17 +153,17 @@
| :ref:`texture | \* | :class: mjcf-attributes |
| <asset-texture>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<asset-texture-name>` | :ref:`type<asset-texture-type>` | :ref:`content_type<asset-texture-content_type>` | :ref:`file<asset-texture-file>` | |
| | | | :ref:`name<asset-texture-name>` | :ref:`type<asset-texture-type>` | :ref:`colorspace<asset-texture-colorspace>` | :ref:`content_type<asset-texture-content_type>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`gridsize<asset-texture-gridsize>` | :ref:`gridlayout<asset-texture-gridlayout>` | :ref:`fileright<asset-texture-fileright>` | :ref:`fileleft<asset-texture-fileleft>` | |
| | | | :ref:`file<asset-texture-file>` | :ref:`gridsize<asset-texture-gridsize>` | :ref:`gridlayout<asset-texture-gridlayout>` | :ref:`fileright<asset-texture-fileright>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`fileup<asset-texture-fileup>` | :ref:`filedown<asset-texture-filedown>` | :ref:`filefront<asset-texture-filefront>` | :ref:`fileback<asset-texture-fileback>` | |
| | | | :ref:`fileleft<asset-texture-fileleft>` | :ref:`fileup<asset-texture-fileup>` | :ref:`filedown<asset-texture-filedown>` | :ref:`filefront<asset-texture-filefront>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`builtin<asset-texture-builtin>` | :ref:`rgb1<asset-texture-rgb1>` | :ref:`rgb2<asset-texture-rgb2>` | :ref:`mark<asset-texture-mark>` | |
| | | | :ref:`fileback<asset-texture-fileback>` | :ref:`builtin<asset-texture-builtin>` | :ref:`rgb1<asset-texture-rgb1>` | :ref:`rgb2<asset-texture-rgb2>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`markrgb<asset-texture-markrgb>` | :ref:`random<asset-texture-random>` | :ref:`width<asset-texture-width>` | :ref:`height<asset-texture-height>` | |
| | | | :ref:`mark<asset-texture-mark>` | :ref:`markrgb<asset-texture-markrgb>` | :ref:`random<asset-texture-random>` | :ref:`width<asset-texture-width>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`hflip<asset-texture-hflip>` | :ref:`vflip<asset-texture-vflip>` | :ref:`nchannel<asset-texture-nchannel>` | | |
| | | | :ref:`height<asset-texture-height>` | :ref:`hflip<asset-texture-hflip>` | :ref:`vflip<asset-texture-vflip>` | :ref:`nchannel<asset-texture-nchannel>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| asset |br| |_| |L| | | .. table:: |
@@ -313,13 +313,15 @@
| :ref:`light | \* | :class: mjcf-attributes |
| <body-light>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<body-light-name>` | :ref:`class<body-light-class>` | :ref:`directional<body-light-directional>` | :ref:`castshadow<body-light-castshadow>` | |
| | | | :ref:`name<body-light-name>` | :ref:`class<body-light-class>` | :ref:`directional<body-light-directional>` | :ref:`type<body-light-type>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`active<body-light-active>` | :ref:`pos<body-light-pos>` | :ref:`dir<body-light-dir>` | :ref:`bulbradius<body-light-bulbradius>` | |
| | | | :ref:`castshadow<body-light-castshadow>` | :ref:`active<body-light-active>` | :ref:`pos<body-light-pos>` | :ref:`dir<body-light-dir>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`attenuation<body-light-attenuation>` | :ref:`cutoff<body-light-cutoff>` | :ref:`exponent<body-light-exponent>` | :ref:`ambient<body-light-ambient>` | |
| | | | :ref:`bulbradius<body-light-bulbradius>` | :ref:`intensity<body-light-intensity>` | :ref:`range<body-light-range>` | :ref:`attenuation<body-light-attenuation>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`diffuse<body-light-diffuse>` | :ref:`specular<body-light-specular>` | :ref:`mode<body-light-mode>` | :ref:`target<body-light-target>` | |
| | | | :ref:`cutoff<body-light-cutoff>` | :ref:`exponent<body-light-exponent>` | :ref:`ambient<body-light-ambient>` | :ref:`diffuse<body-light-diffuse>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`specular<body-light-specular>` | :ref:`mode<body-light-mode>` | :ref:`target<body-light-target>` | :ref:`texture<body-light-texture>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| body |br| |_| |L| | | .. table:: |
@@ -344,6 +346,8 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`vertex<body-composite-vertex>` | :ref:`initial<body-composite-initial>` | :ref:`curve<body-composite-curve>` | :ref:`size<body-composite-size>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`quat<body-composite-quat>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| composite |br| |_2| |L| | | .. table:: |
| :ref:`joint | \* | :class: mjcf-attributes |
@@ -439,6 +443,8 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`young<flexcomp-elasticity-young>` | :ref:`poisson<flexcomp-elasticity-poisson>` | :ref:`damping<flexcomp-elasticity-damping>` | :ref:`thickness<flexcomp-elasticity-thickness>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`elastic2d<flexcomp-elasticity-elastic2d>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
| :ref:`contact | ? | :class: mjcf-attributes |
@@ -450,7 +456,7 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`margin<flexcomp-contact-margin>` | :ref:`gap<flexcomp-contact-gap>` | :ref:`internal<flexcomp-contact-internal>` | :ref:`selfcollide<flexcomp-contact-selfcollide>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`activelayers<flexcomp-contact-activelayers>` | | | | |
| | | | :ref:`activelayers<flexcomp-contact-activelayers>` | :ref:`vertcollide<flexcomp-contact-vertcollide>` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
@@ -500,7 +506,7 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`margin<flex-contact-margin>` | :ref:`gap<flex-contact-gap>` | :ref:`internal<flex-contact-internal>` | :ref:`selfcollide<flex-contact-selfcollide>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`activelayers<flex-contact-activelayers>` | | | | |
| | | | :ref:`activelayers<flex-contact-activelayers>` | :ref:`vertcollide<flex-contact-vertcollide>` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flex |br| |_2| |L| | | .. table:: |
@@ -516,6 +522,8 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`young<flex-elasticity-young>` | :ref:`poisson<flex-elasticity-poisson>` | :ref:`damping<flex-elasticity-damping>` | :ref:`thickness<flex-elasticity-thickness>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`elastic2d<flex-elasticity-elastic2d>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| deformable |br| |_| |L| | | .. table:: |
| :ref:`skin | \* | :class: mjcf-attributes |
@@ -619,13 +627,15 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<tendon-spatial-name>` | :ref:`class<tendon-spatial-class>` | :ref:`group<tendon-spatial-group>` | :ref:`limited<tendon-spatial-limited>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`range<tendon-spatial-range>` | :ref:`solreflimit<tendon-spatial-solreflimit>` | :ref:`solimplimit<tendon-spatial-solimplimit>` | :ref:`solreffriction<tendon-spatial-solreffriction>` | |
| | | | :ref:`actuatorfrclimited<tendon-spatial-actuatorfrclimited>` | :ref:`range<tendon-spatial-range>` | :ref:`actuatorfrcrange<tendon-spatial-actuatorfrcrange>` | :ref:`solreflimit<tendon-spatial-solreflimit>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`solimpfriction<tendon-spatial-solimpfriction>` | :ref:`frictionloss<tendon-spatial-frictionloss>` | :ref:`springlength<tendon-spatial-springlength>` | :ref:`width<tendon-spatial-width>` | |
| | | | :ref:`solimplimit<tendon-spatial-solimplimit>` | :ref:`solreffriction<tendon-spatial-solreffriction>` | :ref:`solimpfriction<tendon-spatial-solimpfriction>` | :ref:`frictionloss<tendon-spatial-frictionloss>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`material<tendon-spatial-material>` | :ref:`margin<tendon-spatial-margin>` | :ref:`stiffness<tendon-spatial-stiffness>` | :ref:`damping<tendon-spatial-damping>` | |
| | | | :ref:`springlength<tendon-spatial-springlength>` | :ref:`width<tendon-spatial-width>` | :ref:`material<tendon-spatial-material>` | :ref:`margin<tendon-spatial-margin>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`rgba<tendon-spatial-rgba>` | :ref:`user<tendon-spatial-user>` | | | |
| | | | :ref:`stiffness<tendon-spatial-stiffness>` | :ref:`damping<tendon-spatial-damping>` | :ref:`armature<tendon-spatial-armature>` | :ref:`rgba<tendon-spatial-rgba>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`user<tendon-spatial-user>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| spatial |br| |_2| |L| | | .. table:: |
@@ -655,11 +665,13 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<tendon-fixed-name>` | :ref:`class<tendon-fixed-class>` | :ref:`group<tendon-fixed-group>` | :ref:`limited<tendon-fixed-limited>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`range<tendon-fixed-range>` | :ref:`solreflimit<tendon-fixed-solreflimit>` | :ref:`solimplimit<tendon-fixed-solimplimit>` | :ref:`solreffriction<tendon-fixed-solreffriction>` | |
| | | | :ref:`actuatorfrclimited<tendon-fixed-actuatorfrclimited>` | :ref:`range<tendon-fixed-range>` | :ref:`actuatorfrcrange<tendon-fixed-actuatorfrcrange>` | :ref:`solreflimit<tendon-fixed-solreflimit>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`solimpfriction<tendon-fixed-solimpfriction>` | :ref:`frictionloss<tendon-fixed-frictionloss>` | :ref:`springlength<tendon-fixed-springlength>` | :ref:`margin<tendon-fixed-margin>` | |
| | | | :ref:`solimplimit<tendon-fixed-solimplimit>` | :ref:`solreffriction<tendon-fixed-solreffriction>` | :ref:`solimpfriction<tendon-fixed-solimpfriction>` | :ref:`frictionloss<tendon-fixed-frictionloss>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`stiffness<tendon-fixed-stiffness>` | :ref:`damping<tendon-fixed-damping>` | :ref:`user<tendon-fixed-user>` | | |
| | | | :ref:`springlength<tendon-fixed-springlength>` | :ref:`margin<tendon-fixed-margin>` | :ref:`stiffness<tendon-fixed-stiffness>` | :ref:`damping<tendon-fixed-damping>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`armature<tendon-fixed-armature>` | :ref:`user<tendon-fixed-user>` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| fixed |br| |_2| |L| | | .. table:: |
@@ -1002,6 +1014,15 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| sensor |br| |_| |L| | | .. table:: |
| :ref:`tendonactuatorfrc | \* | :class: mjcf-attributes |
| <sensor-tendonactuatorfrc>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<sensor-tendonactuatorfrc-name>` | :ref:`tendon<sensor-tendonactuatorfrc-tendon>` | :ref:`cutoff<sensor-tendonactuatorfrc-cutoff>` | :ref:`noise<sensor-tendonactuatorfrc-noise>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`user<sensor-tendonactuatorfrc-user>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| sensor |br| |_| |L| | | .. table:: |
| :ref:`ballquat | \* | :class: mjcf-attributes |
| <sensor-ballquat>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
@@ -1442,13 +1463,13 @@
| :ref:`light | ? | :class: mjcf-attributes |
| <default-light>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`pos<default-light-pos>` | :ref:`dir<default-light-dir>` | :ref:`bulbradius<default-light-bulbradius>` | :ref:`directional<default-light-directional>` | |
| | | | :ref:`pos<default-light-pos>` | :ref:`dir<default-light-dir>` | :ref:`bulbradius<default-light-bulbradius>` | :ref:`intensity<default-light-intensity>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`castshadow<default-light-castshadow>` | :ref:`active<default-light-active>` | :ref:`attenuation<default-light-attenuation>` | :ref:`cutoff<default-light-cutoff>` | |
| | | | :ref:`range<default-light-range>` | :ref:`directional<default-light-directional>` | :ref:`type<default-light-type>` | :ref:`castshadow<default-light-castshadow>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`exponent<default-light-exponent>` | :ref:`ambient<default-light-ambient>` | :ref:`diffuse<default-light-diffuse>` | :ref:`specular<default-light-specular>` | |
| | | | :ref:`active<default-light-active>` | :ref:`attenuation<default-light-attenuation>` | :ref:`cutoff<default-light-cutoff>` | :ref:`exponent<default-light-exponent>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`mode<default-light-mode>` | | | | |
| | | | :ref:`ambient<default-light-ambient>` | :ref:`diffuse<default-light-diffuse>` | :ref:`specular<default-light-specular>` | :ref:`mode<default-light-mode>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| default |br| |_| |L| | | .. table:: |
doc/changelog.rst
+83 -9
View File
@@ -2,26 +2,100 @@
Changelog
=========
Upcoming version (not yet released)
-----------------------------------
Upcoming version (not yet release)
----------------------------------
General
^^^^^^^
- Refactored island implementation so that island data is memory-contiguous. This speeds up island processing in the
solver and clears the way for the addition of the Newton and PGS solvers (currently only CG is supported).
- Removed the :at:`shell` plugin. This is now supported by :ref:`flexcomp<body-flexcomp>` and is active depending on
the :ref:`elastic2d<flexcomp-elasticity-elastic2d>` attribute (off by default).
- Replaced the :ref:`directional<body-light-directional>` (boolean) field for lights with a
:ref:`type<body-light-type>` field (of type :ref:`mjtLightType<mjtLightType>`) to allow for additional lighting
types.
- Added new sub-component :ref:`mj_makeM` which combines the :ref:`mj_crb` call with additional logic to support the
introduction in 3.3.1 of :ref:`tendon armature<tendon-spatial-armature>`. In addition to the traditional
``mjData.qM``, :ref:`mj_makeM` also computes ``mjData.M``, a CSR representation of the same matrix.
- Added a new function :ref:`mj_copyBack` to copy real-valued arrays in an mjModel to a compatible mjSpec.
Simulate
^^^^^^^^
- The struct ``mjv_sceneState`` has been removed. This struct was used for partial synchronization of ``mjModel`` and
``mjData`` when the Python viewer is used in passive mode. This functionality is now provided by :ref:`mjv_copyModel`
and :ref:`mjv_copyData`, which don't copy arrays which are not required for visualization.
.. image:: images/changelog/procedural_terrain_generation.png
:width: 25%
:align: right
Python bindings
^^^^^^^^^^^^^^^
- Added examples of procedural terrain generation to the Model Editing tutorial: |mjspec_colab|
Version 3.3.2 (April 28, 2025)
------------------------------
MJX
^^^
1. Added inverse dynamics.
2. Added tendon actuator force sensor.
3. Fix :github:issue:`2606` such that ``make_data`` copies over ``mocap_pos`` and ``mocap_quat``
from ``body_pos`` and ``body_quat``.
Version 3.3.1 (Apr 9, 2025)
----------------------------
.. admonition:: Breaking API changes
:class: attention
- The default value of the flag for toggling :ref:`internal flex contacts<flex-contact-internal>` was changed from
"true" to "false". This feature has proven to be counterintuitive for users.
1. The default value of the flag for toggling :ref:`internal flex contacts<flex-contact-internal>` was changed from
"true" to "false". This feature has proven to be counterintuitive for users.
2. All of the attach functions (``mjs_attachBody``, ``mjs_attachFrame``, ``mjs_attachToSite``,
``mjs_attachFrameToSite``) have been removed and replaced by a single function :ref:`mjs_attach`.
General
^^^^^^^
3. Added :ref:`tendon armature<tendon-spatial-armature>`: inertia associated with changes in tendon length.
4. Added the :ref:`compiler/saveinertial<compiler-saveinertial>` flag, writing explicit inertial clauses for all
bodies when saving to XML.
5. Added :ref:`orientation<body-composite-quat>` attribute to :ref:`composite<body-composite>`. Moreover, allow the
composite to be the direct child of a frame.
6. Added :ref:`tendon actuator force limits<tendon-spatial-actuatorfrclimited>` and
:ref:`tendon actuator force sensor<sensor-tendonactuatorfrc>`.
MJX
^^^
7. Added tendon actuator force limits.
Bug fixes
^^^^^^^^^
- :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
which the Jacobian is computed, now fixed.
- Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
mjSpec to a frame or a site.
8. :ref:`mj_jacDot` was missing a term that accounts for the motion of the point with respect to
which the Jacobian is computed, now fixed.
9. Fixed a bug that caused the parent frame of elements in the child worldbody to be incorrectly set when attaching an
mjSpec to a frame or a site.
10. Fixed a bug that caused shadow rendering to flicker on platforms (e.g., MacOS) that do not support
ARB_clip_control. Fixed in collaboration with :github:user:`aftersomemath`.
Python bindings
^^^^^^^^^^^^^^^
.. youtube:: LbANnKMDOHg
:aspect: 16:7
:align: right
:width: 240px
11. Added examples of procedural model creation to the Model Editing tutorial: |mjspec_colab|
12. Added support for nameless :ref:`mjSpec` objects in the ``bind`` method, see the corresponding
:ref:`section<PyMJCF>` in the documentation.
.. |mjspec_colab| image:: https://colab.research.google.com/assets/colab-badge.svg
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/mjspec.ipynb
Version 3.3.0 (Feb 26, 2025)
----------------------------
Feature promotion
^^^^^^^^^^^^^^^^^
.. youtube:: qJFbx-FR7Bc
doc/computation/index.rst
+1 -1
View File
@@ -1692,7 +1692,7 @@ The stages below compute quantities that depend on the generalized positions ``m
4. Compute quantities related to :ref:`flex<deformable-flex>` objects: :ref:`mj_flex`
5. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
tendons: :ref:`mj_tendon`
6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb`
6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_makeM`
7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM`
8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection:
:ref:`mj_collision`
doc/images/XMLreference/armature.gif
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 134 KiB

doc/images/XMLreference/armature_dark.gif
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 146 KiB

doc/images/XMLreference/tendon_armature.gif
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 123 KiB

doc/images/XMLreference/tendon_armature_dark.gif
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 133 KiB

doc/images/changelog/procedural_terrain_generation.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 83 KiB

doc/includes/references.h
+145 -344
View File
@@ -154,7 +154,6 @@ struct mjData_ {
// solver statistics
mjSolverStat solver[mjNISLAND*mjNSOLVER]; // solver statistics per island, per iteration
int solver_nisland; // number of islands processed by solver
int solver_niter[mjNISLAND]; // number of solver iterations, per island
int solver_nnz[mjNISLAND]; // number of nonzeros in Hessian or efc_AR, per island
mjtNum solver_fwdinv[2]; // forward-inverse comparison: qfrc, efc
@@ -172,6 +171,7 @@ struct mjData_ {
int nJ; // number of non-zeros in constraint Jacobian
int nA; // number of non-zeros in constraint inverse inertia matrix
int nisland; // number of detected constraint islands
int nidof; // number of dofs in all islands
// global properties
mjtNum time; // simulation time
@@ -267,12 +267,13 @@ struct mjData_ {
int* moment_colind; // column indices in sparse Jacobian (nJmom x 1)
mjtNum* actuator_moment; // actuator moments (nJmom x 1)
// computed by mj_fwdPosition/mj_crb
// computed by mj_fwdPosition/mj_makeM
mjtNum* crb; // com-based composite inertia and mass (nbody x 10)
mjtNum* qM; // total inertia (sparse) (nM x 1)
mjtNum* qM; // inertia (sparse) (nM x 1)
mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1)
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1)
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
// computed by mj_collisionTree
@@ -305,27 +306,23 @@ struct mjData_ {
mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3)
// computed by mj_Euler or mj_implicit
mjtNum* qH; // L'*D*L factorization of modified M (nM x 1)
mjtNum* qH; // L'*D*L factorization of modified M (nC x 1)
mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1)
// computed by mj_resetData
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
int* M_rownnz; // inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // inertia: address of each row in M_colind (nv x 1)
int* M_colind; // inertia: column indices of non-zeros (nM x 1)
int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1)
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
int* mapM2C; // index mapping from M to C (nC x 1)
int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1)
int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1)
int* D_diag; // dof-dof: index of diagonal element (nv x 1)
int* D_colind; // dof-dof: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1)
int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1)
int* mapM2M; // index mapping from qM to M (nC x 1)
int* D_rownnz; // full inertia: non-zeros in each row (nv x 1)
int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
// computed by mj_implicit/mj_derivative
mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1)
@@ -347,8 +344,8 @@ struct mjData_ {
mjtNum* qfrc_constraint; // constraint force (nv x 1)
// computed by mj_inverse
mjtNum* qfrc_inverse; // net external force; should equal: (nv x 1)
// qfrc_applied + J'*xfrc_applied + qfrc_actuator
mjtNum* qfrc_inverse; // net external force; should equal:
// qfrc_applied + J'*xfrc_applied + qfrc_actuator (nv x 1)
// computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format
mjtNum* cacc; // com-based acceleration (nbody x 6)
@@ -382,16 +379,50 @@ struct mjData_ {
mjtNum* efc_R; // inverse constraint mass (nefc x 1)
int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1)
// computed by mj_island
// computed by mj_island (island dof structure)
int* dof_island; // island id of this dof; -1: none (nv x 1)
int* island_dofnum; // number of dofs in island (nisland x 1)
int* island_dofadr; // start address in island_dofind (nisland x 1)
int* island_dofind; // island dof indices; -1: none (nv x 1)
int* dof_islandind; // dof island indices; -1: none (nv x 1)
int* island_nv; // number of dofs in this island (nisland x 1)
int* island_idofadr; // island start address in idof vector (nisland x 1)
int* island_dofadr; // island start address in dof vector (nisland x 1)
int* map_dof2idof; // map from dof to idof (nv x 1)
int* map_idof2dof; // map from idof to dof; >= nidof: unconstrained (nv x 1)
// 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)
int* efc_island; // island id of this constraint (nefc x 1)
int* island_efcnum; // number of constraints in island (nisland x 1)
int* island_efcadr; // start address in island_efcind (nisland x 1)
int* island_efcind; // island constraint indices (nefc x 1)
int* island_ne; // number of equality constraints in island (nisland x 1)
int* island_nf; // number of friction constraints in island (nisland x 1)
int* island_nefc; // number of constraints in island (nisland x 1)
int* island_iefcadr; // start address in iefc vector (nisland x 1)
int* map_efc2iefc; // map from efc to iefc (nefc x 1)
int* map_iefc2efc; // map from iefc to efc (nefc x 1)
// 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)
int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1)
int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1)
int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1)
int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* iefc_J; // constraint Jacobian (nJ x 1)
mjtNum* iefc_JT; // constraint Jacobian transposed (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)
// computed by mj_projectConstraint (PGS solver)
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
@@ -409,11 +440,18 @@ struct mjData_ {
// computed by mj_fwdConstraint/mj_inverse
mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1)
mjtNum* efc_force; // constraint force in constraint space (nefc x 1)
mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1)
int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1)
mjtNum* iefc_force; // constraint force in constraint space (nefc x 1)
int* efc_state; // constraint state (mjtConstraintState) (nefc x 1)
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
};
typedef struct mjData_ mjData;
typedef enum mjtDisableBit_ { // disable default feature bitflags
@@ -488,6 +526,12 @@ typedef enum mjtCamLight_ { // tracking mode for camera and light
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
@@ -506,6 +550,11 @@ typedef enum mjtTextureRole_ { // role of texture map in rendering
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
@@ -604,7 +653,8 @@ typedef enum mjtObj_ { // type of MujoCo object
// meta elements, do not appear in mjModel
mjOBJ_FRAME = 100, // frame
mjOBJ_DEFAULT // default
mjOBJ_DEFAULT, // default
mjOBJ_MODEL // entire model
} mjtObj;
typedef enum mjtConstraint_ { // type of constraint
@@ -645,6 +695,7 @@ typedef enum mjtSensor_ { // type of sensor
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
@@ -722,6 +773,12 @@ typedef enum mjtFlexSelf_ { // mode for flex selfcollide
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;
struct mjLROpt_ { // options for mj_setLengthRange()
// flags
int mode; // which actuators to process (mjtLRMode)
@@ -907,7 +964,7 @@ struct mjModel_ {
int ncam; // number of cameras
int nlight; // number of lights
int nflex; // number of flexes
int nflexnode; // number of dofs in all flexes
int nflexnode; // number of dofs in all flexes
int nflexvert; // number of vertices in all flexes
int nflexedge; // number of edges in all flexes
int nflexelem; // number of elements in all flexes
@@ -1128,9 +1185,12 @@ struct mjModel_ {
int* light_mode; // light tracking mode (mjtCamLight) (nlight x 1)
int* light_bodyid; // id of light's body (nlight x 1)
int* light_targetbodyid; // id of targeted body; -1: none (nlight x 1)
mjtByte* light_directional; // directional light (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)
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)
mjtNum* light_pos; // position rel. to body frame (nlight x 3)
mjtNum* light_dir; // direction rel. to body frame (nlight x 3)
@@ -1182,6 +1242,7 @@ struct mjModel_ {
int* flex_nodebodyid; // node body ids (nflexnode x 1)
int* flex_vertbodyid; // vertex body ids (nflexvert x 1)
int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2)
int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2)
int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1)
int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1)
int* flex_elemedge; // element edge ids (nflexelemedge x 1)
@@ -1196,6 +1257,7 @@ struct mjModel_ {
mjtNum* flexedge_invweight0; // edge inv. weight in qpos0 (nflexedge x 1)
mjtNum* flex_radius; // radius around primitive element (nflex x 1)
mjtNum* flex_stiffness; // finite element stiffness matrix (nflexelem x 21)
mjtNum* flex_bending; // bending stiffness (nflexedge x 16)
mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1)
mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1)
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
@@ -1276,6 +1338,7 @@ struct mjModel_ {
// textures
int* tex_type; // texture type (mjtTexture) (ntex x 1)
int* tex_colorspace; // texture colorspace (mjtColorSpace) (ntex x 1)
int* tex_height; // number of rows in texture image (ntex x 1)
int* tex_width; // number of columns in texture image (ntex x 1)
int* tex_nchannel; // number of channels in texture image (ntex x 1)
@@ -1326,15 +1389,18 @@ struct mjModel_ {
int* tendon_matid; // material id for rendering (ntendon x 1)
int* tendon_group; // group for visibility (ntendon 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)
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)
mjtNum* tendon_solref_fri; // constraint solver reference: friction (ntendon x mjNREF)
mjtNum* tendon_solimp_fri; // constraint solver impedance: friction (ntendon x mjNIMP)
mjtNum* tendon_range; // tendon length limits (ntendon x 2)
mjtNum* tendon_actfrcrange; // range of total actuator force (ntendon x 2)
mjtNum* tendon_margin; // min distance for limit detection (ntendon x 1)
mjtNum* tendon_stiffness; // stiffness coefficient (ntendon x 1)
mjtNum* tendon_damping; // damping coefficient (ntendon x 1)
mjtNum* tendon_armature; // inertia associated with tendon velocity (ntendon x 1)
mjtNum* tendon_frictionloss; // loss due to friction (ntendon x 1)
mjtNum* tendon_lengthspring; // spring resting length range (ntendon x 2)
mjtNum* tendon_length0; // tendon length in qpos0 (ntendon x 1)
@@ -1451,6 +1517,9 @@ struct mjModel_ {
// paths
char* paths; // paths to assets, 0-terminated (npaths x 1)
// compilation signature
uint64_t signature; // also held by the mjSpec that compiled this model
};
typedef struct mjModel_ mjModel;
struct mjResource_ {
@@ -1536,6 +1605,15 @@ struct mjpPlugin_ {
void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
};
typedef struct mjpPlugin_ mjpPlugin;
struct mjSDF_ {
const mjpPlugin** plugin;
int* id;
mjtSDFType type;
mjtNum* relpos;
mjtNum* relmat;
mjtGeom* geomtype;
};
typedef struct mjSDF_ mjSDF;
typedef enum mjtGridPos_ { // grid position for overlay
mjGRID_TOPLEFT = 0, // top left
mjGRID_TOPRIGHT, // top right
@@ -1713,6 +1791,7 @@ typedef enum mjtOrientation_ { // type of orientation specifier
} mjtOrientation;
typedef struct mjsElement_ { // element type, do not modify
mjtObj elemtype; // element type
uint64_t signature; // compilation signature
} mjsElement;
typedef struct mjsCompiler_ { // compiler options
mjtByte autolimits; // infer "limited" attribute based on range
@@ -1728,6 +1807,7 @@ typedef struct mjsCompiler_ { // compiler options
mjtByte fusestatic; // fuse static bodies with parent
int inertiafromgeom; // use geom inertias (mjtInertiaFromGeom)
int inertiagrouprange[2]; // range of geom groups used to compute inertia
mjtByte saveinertial; // save explicit inertial clause for all bodies to XML
int alignfree; // align free joints with inertial frame
mjLROpt LRopt; // options for lengthrange computation
} mjsCompiler;
@@ -1963,9 +2043,12 @@ typedef struct mjsLight_ { // light specification
// intrinsics
mjtByte active; // is light active
mjtByte directional; // is light directional or spot
mjtLightType type; // type of light
mjString* texture; // texture name for image lights
mjtByte castshadow; // does light cast shadows
double bulbradius; // bulb radius, for soft shadows
float bulbradius; // bulb radius, for soft shadows
float intensity; // intensity, in candelas
float range; // range of effectiveness
float attenuation[3]; // OpenGL attenuation (quadratic model)
float cutoff; // OpenGL cutoff
float exponent; // OpenGL exponent
@@ -1997,7 +2080,8 @@ typedef struct mjsFlex_ { // flex specification
double radius; // radius around primitive element
mjtByte internal; // enable internal collisions
mjtByte flatskin; // render flex skin with flat shading
int selfcollide; // mode for flex self colllision
int selfcollide; // mode for flex self collision
int vertcollide; // mode for vertex collision
int activelayers; // number of active element layers in 3D
int group; // group for visualizatioh
double edgestiffness; // edge stiffness
@@ -2008,6 +2092,7 @@ typedef struct mjsFlex_ { // flex specification
double poisson; // Poisson's ratio
double damping; // Rayleigh's damping
double thickness; // thickness (2D only)
int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
// mesh properties
mjStringVec* nodebody; // node body names
@@ -2079,6 +2164,7 @@ typedef struct mjsTexture_ { // texture specification
mjsElement* element; // element type
mjString* name; // name
mjtTexture type; // texture type
mjtColorSpace colorspace; // colorspace
// method 1: builtin
int builtin; // builtin type (mjtBuiltin)
@@ -2165,17 +2251,20 @@ typedef struct mjsTendon_ { // tendon specification
mjsElement* element; // element type
mjString* name; // name
// stiffness, damping, friction
// stiffness, damping, friction, armature
double stiffness; // stiffness coefficient
double springlength[2]; // spring resting length; {-1, -1}: use qpos_spring
double damping; // damping coefficient
double frictionloss; // friction loss
mjtNum solref_friction[mjNREF]; // solver reference: tendon friction
mjtNum solimp_friction[mjNIMP]; // solver impedance: tendon friction
double armature; // inertia associated with tendon velocity
// length range
int limited; // does tendon have limits (mjtLimited)
int actfrclimited; // does tendon have actuator force limits
double range[2]; // length limits
double actfrcrange[2]; // actuator force limits
double margin; // margin value for tendon limit detection
mjtNum solref_limit[mjNREF]; // solver reference: tendon limits
mjtNum solimp_limit[mjNIMP]; // solver impedance: tendon limits
@@ -2728,7 +2817,7 @@ typedef struct mjvGLCamera_ mjvGLCamera;
struct mjvGeom_ { // abstract geom
// type info
int type; // geom type (mjtGeom)
int dataid; // mesh, hfield or plane id; -1: none
int dataid; // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull)
int objtype; // mujoco object type; mjOBJ_UNKNOWN for decor
int objid; // mujoco object id; -1 for decor
int category; // visual category
@@ -2759,6 +2848,8 @@ typedef struct mjvGeom_ mjvGeom;
struct mjvLight_ { // OpenGL light
float pos[3]; // position rel. to body frame
float dir[3]; // direction rel. to body frame
int type; // type (mjtLightType)
int texid; // texture id for image lights
float attenuation[3]; // OpenGL attenuation (quadratic model)
float cutoff; // OpenGL cutoff
float exponent; // OpenGL exponent
@@ -2766,9 +2857,10 @@ struct mjvLight_ { // OpenGL light
float diffuse[3]; // diffuse rgb (alpha=1)
float specular[3]; // specular rgb (alpha=1)
mjtByte headlight; // headlight
mjtByte directional; // directional light
mjtByte castshadow; // does light cast shadows
float bulbradius; // bulb radius for soft shadows
float intensity; // intensity, in candelas
float range; // range of effectiveness
};
typedef struct mjvLight_ mjvLight;
struct mjvOption_ { // abstract visualization options
@@ -2889,287 +2981,6 @@ struct mjvFigure_ { // abstract 2D figure passed to OpenGL rendere
float yaxisdata[2]; // range of y-axis in data units
};
typedef struct mjvFigure_ mjvFigure;
struct mjvSceneState_ {
int nbuffer; // size of the buffer in bytes
void* buffer; // heap-allocated memory for all arrays in this struct
int maxgeom; // maximum number of mjvGeom supported by this state object
mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins
// fields in mjModel that are necessary to re-render a scene
struct {
int nv;
int nu;
int na;
int nbody;
int nbvh;
int nbvhstatic;
int njnt;
int ngeom;
int nsite;
int ncam;
int nlight;
int nmesh;
int nskin;
int nflex;
int nflexvert;
int nflextexcoord;
int nskinvert;
int nskinface;
int nskinbone;
int nskinbonevert;
int nmat;
int neq;
int ntendon;
int ntree;
int nwrap;
int nsensor;
int nnames;
int npaths;
int nsensordata;
int narena;
mjOption opt;
mjVisual vis;
mjStatistic stat;
int* body_parentid;
int* body_rootid;
int* body_weldid;
int* body_mocapid;
int* body_jntnum;
int* body_jntadr;
int* body_dofnum;
int* body_dofadr;
int* body_geomnum;
int* body_geomadr;
mjtNum* body_iquat;
mjtNum* body_mass;
mjtNum* body_inertia;
int* body_bvhadr;
int* body_bvhnum;
int* bvh_depth;
int* bvh_child;
int* bvh_nodeid;
mjtNum* bvh_aabb;
int* jnt_type;
int* jnt_bodyid;
int* jnt_group;
int* geom_type;
int* geom_bodyid;
int* geom_contype;
int* geom_conaffinity;
int* geom_dataid;
int* geom_matid;
int* geom_group;
mjtNum* geom_size;
mjtNum* geom_aabb;
mjtNum* geom_rbound;
float* geom_rgba;
int* site_type;
int* site_bodyid;
int* site_matid;
int* site_group;
mjtNum* site_size;
float* site_rgba;
int* cam_orthographic;
mjtNum* cam_fovy;
mjtNum* cam_ipd;
int* cam_resolution;
float* cam_sensorsize;
float* cam_intrinsic;
mjtByte* light_directional;
mjtByte* light_castshadow;
float* light_bulbradius;
mjtByte* light_active;
float* light_attenuation;
float* light_cutoff;
float* light_exponent;
float* light_ambient;
float* light_diffuse;
float* light_specular;
mjtByte* flex_flatskin;
int* flex_dim;
int* flex_matid;
int* flex_group;
int* flex_interp;
int* flex_nodeadr;
int* flex_nodenum;
int* flex_nodebodyid;
int* flex_vertadr;
int* flex_vertnum;
int* flex_elem;
int* flex_elemtexcoord;
int* flex_elemlayer;
int* flex_elemadr;
int* flex_elemnum;
int* flex_elemdataadr;
int* flex_shell;
int* flex_shellnum;
int* flex_shelldataadr;
int* flex_texcoordadr;
int* flex_bvhadr;
int* flex_bvhnum;
mjtByte* flex_centered;
mjtNum* flex_node;
mjtNum* flex_radius;
float* flex_rgba;
float* flex_texcoord;
int* hfield_pathadr;
int* mesh_bvhadr;
int* mesh_bvhnum;
int* mesh_texcoordadr;
int* mesh_graphadr;
int* mesh_pathadr;
int* skin_matid;
int* skin_group;
float* skin_rgba;
float* skin_inflate;
int* skin_vertadr;
int* skin_vertnum;
int* skin_texcoordadr;
int* skin_faceadr;
int* skin_facenum;
int* skin_boneadr;
int* skin_bonenum;
float* skin_vert;
int* skin_face;
int* skin_bonevertadr;
int* skin_bonevertnum;
float* skin_bonebindpos;
float* skin_bonebindquat;
int* skin_bonebodyid;
int* skin_bonevertid;
float* skin_bonevertweight;
int* skin_pathadr;
int* tex_pathadr;
int* mat_texid;
mjtByte* mat_texuniform;
float* mat_texrepeat;
float* mat_emission;
float* mat_specular;
float* mat_shininess;
float* mat_reflectance;
float* mat_metallic;
float* mat_roughness;
float* mat_rgba;
int* eq_type;
int* eq_obj1id;
int* eq_obj2id;
int* eq_objtype;
mjtNum* eq_data;
int* tendon_num;
int* tendon_matid;
int* tendon_group;
mjtByte* tendon_limited;
mjtNum* tendon_width;
mjtNum* tendon_range;
mjtNum* tendon_stiffness;
mjtNum* tendon_damping;
mjtNum* tendon_frictionloss;
mjtNum* tendon_lengthspring;
float* tendon_rgba;
int* actuator_trntype;
int* actuator_dyntype;
int* actuator_trnid;
int* actuator_actadr;
int* actuator_actnum;
int* actuator_group;
mjtByte* actuator_ctrllimited;
mjtByte* actuator_actlimited;
mjtNum* actuator_ctrlrange;
mjtNum* actuator_actrange;
mjtNum* actuator_cranklength;
int* sensor_type;
int* sensor_objid;
int* sensor_adr;
int* name_bodyadr;
int* name_jntadr;
int* name_geomadr;
int* name_siteadr;
int* name_camadr;
int* name_lightadr;
int* name_eqadr;
int* name_tendonadr;
int* name_actuatoradr;
char* names;
char* paths;
} model;
// fields in mjData that are necessary to re-render a scene
struct {
mjWarningStat warning[mjNWARNING];
int nefc;
int ncon;
int nisland;
mjtNum time;
mjtNum* act;
mjtNum* ctrl;
mjtNum* xfrc_applied;
mjtByte* eq_active;
mjtNum* sensordata;
mjtNum* xpos;
mjtNum* xquat;
mjtNum* xmat;
mjtNum* xipos;
mjtNum* ximat;
mjtNum* xanchor;
mjtNum* xaxis;
mjtNum* geom_xpos;
mjtNum* geom_xmat;
mjtNum* site_xpos;
mjtNum* site_xmat;
mjtNum* cam_xpos;
mjtNum* cam_xmat;
mjtNum* light_xpos;
mjtNum* light_xdir;
mjtNum* subtree_com;
int* ten_wrapadr;
int* ten_wrapnum;
int* wrap_obj;
mjtNum* ten_length;
mjtNum* wrap_xpos;
mjtNum* bvh_aabb_dyn;
mjtByte* bvh_active;
int* island_dofadr;
int* island_dofind;
int* dof_island;
int* efc_island;
int* tendon_efcadr;
mjtNum* flexvert_xpos;
mjContact* contact;
mjtNum* efc_force;
void* arena;
} data;
};
typedef struct mjvSceneState_ mjvSceneState;
//----------------------------- MJAPI FUNCTIONS --------------------------------
void mj_defaultVFS(mjVFS* vfs);
@@ -3181,6 +2992,7 @@ mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int err
mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz);
mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
int mj_copyBack(mjSpec* s, const mjModel* m);
int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
int mj_saveLastXML(const char* filename, const mjModel* m, char* error, int error_sz);
void mj_freeLastXML(void);
@@ -3204,6 +3016,7 @@ void mj_deleteModel(mjModel* m);
int mj_sizeModel(const mjModel* m);
mjData* mj_makeData(const mjModel* m);
mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src);
void mj_resetData(const mjModel* m, mjData* d);
void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
@@ -3259,6 +3072,7 @@ void mj_flex(const mjModel* m, mjData* d);
void mj_tendon(const mjModel* m, mjData* d);
void mj_transmission(const mjModel* m, mjData* d);
void mj_crb(const mjModel* m, mjData* d);
void mj_makeM(const mjModel* m, mjData* d);
void mj_factorM(const mjModel* m, mjData* d);
void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y,
@@ -3357,14 +3171,8 @@ mjtNum mjv_frustumHeight(const mjvScene* scn);
void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum forward[3]);
void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvCamera* cam);
void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action,
mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvCamera* cam);
void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action,
mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvPerturb* pert);
void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
const mjtNum roomup[3], mjvScene* scn);
void mjv_initPerturb(const mjModel* m, mjData* d, const mjvScene* scn, mjvPerturb* pert);
@@ -3387,16 +3195,7 @@ void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom);
void mjv_freeScene(mjvScene* scn);
void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt,
const mjvPerturb* pert, mjvCamera* cam, int catmask,
mjvScene* scn);
void mjv_copyModel(mjModel* dest, const mjModel* src);
void mjv_defaultSceneState(mjvSceneState* scnstate);
void mjv_makeSceneState(const mjModel* m, const mjData* d,
mjvSceneState* scnstate, int maxgeom);
void mjv_freeSceneState(mjvSceneState* scnstate);
void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
mjvSceneState* scnstate);
void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
const mjvPerturb* pert, int catmask, mjvScene* scn);
void mjv_makeLights(const mjModel* m, const mjData* d, mjvScene* scn);
@@ -3580,6 +3379,10 @@ void mju_insertionSortInt(int* list, int n);
mjtNum mju_Halton(int index, int base);
char* mju_strncpy(char *dst, const char *src, int n);
mjtNum mju_sigmoid(mjtNum x);
const mjpPlugin* mjc_getSDF(const mjModel* m, int id);
mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]);
void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3],
const mjtNum x[3]);
void mjd_transitionFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_centered,
mjtNum* A, mjtNum* B, mjtNum* C, mjtNum* D);
void mjd_inverseFD(const mjModel* m, mjData* d, mjtNum eps, mjtByte flg_actuation,
@@ -3605,14 +3408,8 @@ void mju_threadPoolEnqueue(mjThreadPool* thread_pool, mjTask* task);
void mju_threadPoolDestroy(mjThreadPool* thread_pool);
void mju_defaultTask(mjTask* task);
void mju_taskJoin(mjTask* task);
mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
const char* prefix, const char* suffix);
mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix);
mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
const char* prefix, const char* suffix);
int mjs_detachBody(mjSpec* s, mjsBody* b);
int mjs_detachDefault(mjSpec* s, mjsDefault* d);
mjsBody* mjs_addBody(mjsBody* body, const mjsDefault* def);
@@ -3675,12 +3472,16 @@ void mjs_setDouble(mjDoubleVec* dest, const double* array, int size);
void mjs_setPluginAttributes(mjsPlugin* plugin, void* attributes);
const char* mjs_getString(const mjString* source);
const double* mjs_getDouble(const mjDoubleVec* source, int* size);
const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
void mjs_setDefault(mjsElement* element, const mjsDefault* def);
void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
int mjs_setFrame(mjsElement* dest, mjsFrame* frame);
const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
const mjsOrientation* orientation);
mjsFrame* mjs_bodyToFrame(mjsBody** body);
void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
void mjs_setUserValueWithCleanup(mjsElement* element, const char* key,
const void* data,
void (*cleanup)(const void*));
const void* mjs_getUserValue(mjsElement* element, const char* key);
void mjs_deleteUserValue(mjsElement* element, const char* key);
void mjs_defaultSpec(mjSpec* spec);
doc/mjx.rst
+3 -3
View File
@@ -235,6 +235,8 @@ The following features are **fully supported** in MJX:
- 1, 3, 4, 6 (1 is not supported with ``ELLIPTIC``)
* - :ref:`Solver <mjtSolver>`
- ``CG``, ``NEWTON``
* - Dynamics
- :ref:`Inverse <mj_inverse>`
* - Fluid Model
- :ref:`flInertia`
* - :ref:`Tendons <tendon>`
@@ -244,7 +246,7 @@ The following features are **fully supported** in MJX:
``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``,
``VELOCIMETER``, ``GYRO``, ``JOINTVEL``, ``TENDONVEL``, ``ACTUATORVEL``, ``BALLANGVEL``, ``FRAMELINVEL``,
``FRAMEANGVEL``, ``SUBTREELINVEL``, ``SUBTREEANGMOM``, ``TOUCH``, ``ACCELEROMETER``, ``FORCE``, ``TORQUE``,
``ACTUATORFRC``, ``JOINTACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``
``ACTUATORFRC``, ``JOINTACTFRC``, ``TENDONACTFRC``, ``FRAMELINACC``, ``FRAMEANGACC``
(``ACCELEROMETER``, ``FORCE``, ``TORQUE`` not supported with connect or weld equality constraints)
The following features are **in development** and coming soon:
@@ -262,8 +264,6 @@ The following features are **in development** and coming soon:
(``BOX``, ``MESH``, ``HFIELD``) and ``ELLIPSOID``.
* - :ref:`Integrator <mjtIntegrator>`
- ``IMPLICIT``
* - Dynamics
- :ref:`Inverse <mj_inverse>`
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Sensors <mjtSensor>`
doc/modeling.rst
+4 -1
View File
@@ -724,7 +724,10 @@ Force clamping at joint input with :ref:`joint/actuatorfrcrange<body-joint-actua
:ref:`jointactuatorfrc<sensor-jointactuatorfrc>` sensor to report the total actuator force acting on a joint.
The standard :ref:`actuatorfrc<sensor-actuatorfrc>` sensor will continue to report the pre-clamped actuator force.
The three clamping options above are non-exclusive and can be combined as required.
Force clamping at tendon input with :ref:`tendon/actuatorfrcrange<tendon-spatial-actuatorfrcrange>`:
This tendon attribute clamps input forces from all actuators acting on the tendon.
The clamping options above are non-exclusive and can be combined as required.
.. _CLengthRange:
doc/programming/modeledit.rst
+15 -14
View File
@@ -111,9 +111,10 @@ This framework introduces a powerful new feature: attaching and detaching model
to power the :ref:`attach<body-attach>` an :ref:`replicate<replicate>` meta-elements in MJCF. Attachment allows the user
to move or copy a subtree from one model into another, while also copying or moving related referenced assets and
referencing elements from outside the kinematic tree (e.g., actuators and sensors). Similarly, detaching a subtree will
remove all associated elements from the model. The default behavior is to move during attach. The user can select to
instead copy by passing the corresponding flag to ``mjs_setDeepCopy``. This flag is temporary set to true while parsing
XMLs. It is possible to :ref:`attach a body to a frame<mjs_attachBody>`:
remove all associated elements from the model. The default behavior is to move the child into the parent while
attaching, so subsequent changes to the child will also change the parent. Alternatively, the user can choose to make an
entirely new copy during attach using :ref:`mjs_setDeepCopy`. This flag is temporarily set to true while parsing XMLs.
It is possible to :ref:`attach a body to a frame<mjs_attach>`:
.. code-block:: C
@@ -121,29 +122,29 @@ XMLs. It is possible to :ref:`attach a body to a frame<mjs_attachBody>`:
mjSpec* child = mj_makeSpec();
parent->compiler.degree = 0;
child->compiler.degree = 1;
mjsFrame* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL);
mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL);
mjsBody* attached_body_1 = mjs_attachBody(frame, body, "attached-", "-1");
mjsElement* frame = mjs_addFrame(mjs_findBody(parent, "world"), NULL)->element;
mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element;
mjsBody* attached_body_1 = mjs_asBody(mjs_attach(frame, body, "attached-", "-1"));
or :ref:`attach a body to a site<mjs_attachToSite>`:
or :ref:`attach a body to a site<mjs_attach>`:
.. code-block:: C
mjSpec* parent = mj_makeSpec();
mjSpec* child = mj_makeSpec();
mjsSite* site = mjs_addSite(mjs_findBody(parent, "world"), NULL);
mjsBody* body = mjs_addBody(mjs_findBody(child, "world"), NULL);
mjsBody* attached_body_2 = mjs_attachToSite(site, body, "attached-", "-2");
mjsElement* site = mjs_addSite(mjs_findBody(parent, "world"), NULL)->element;
mjsElement* body = mjs_addBody(mjs_findBody(child, "world"), NULL)->element;
mjsBody* attached_body_2 = mjs_asBody(mjs_attach(site, body, "attached-", "-2"));
or :ref:`attach a frame to a body<mjs_attachFrame>`:
or :ref:`attach a frame to a body<mjs_attach>`:
.. code-block:: C
mjSpec* parent = mj_makeSpec();
mjSpec* child = mj_makeSpec();
mjsBody* body = mjs_addBody(mjs_findBody(parent, "world"), NULL);
mjsFrame* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL);
mjsFrame* attached_frame = mjs_attachFrame(body, frame, "attached-", "-1");
mjsElement* body = mjs_addBody(mjs_findBody(parent, "world"), NULL)->element;
mjsElement* frame = mjs_addFrame(mjs_findBody(child, "world"), NULL)->element;
mjsFrame* attached_frame = mjs_asFrame(mjs_attach(body, frame, "attached-", "-1"));
Note that in the above examples, the parent and child models have different values for ``compiler.degree``,
corresponding to the :ref:`compiler/angle<compiler-angle>` attribute, specifying the units in which angles are
doc/programming/simulation.rst
+2 -2
View File
@@ -807,7 +807,7 @@ to implement high-resolution timers in C without bringing in additional dependen
does not need timing, and in that case there is no reason to call timing functions.
One part of the simulation pipeline that needs to be monitored closely is the iterative constraint solver. The
simplest diagnostic here is ``mjData.solver_iter`` which shows how many iterations the solver took on the last call to
simplest diagnostic here is ``mjData.solver_niter`` which shows how many iterations the solver took on the last call to
mj_step or ``mj_forward``. Note that the solver has tolerance parameters for early termination, so this number is
usually smaller than the maximum number of iterations allowed. The array ``mjData.solver`` contains one
:ref:`mjSolverStat` data structure per iteration of the constraint solver, with information about the constraint state
@@ -913,7 +913,7 @@ elliptic, depending on which solver is selected in ``mjModel.opt``. The function
can be used to determine which friction cone type is used. For pyramidal cones, the interpretation of the contact force
(whose address we computed above) is non-trivial, because the components are forces along redundant non-orthogonal axes
corresponding to the edges of the pyramid. The function :ref:`mj_contactForce` can be
used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D toque.
used to convert the force generated by a given contact into a more intuitive format: a 3D force followed by a 3D torque.
The torque component will be zero when :at:`condim` is 1 or 3, and non-zero otherwise. This force and torque are
expressed in the contact frame given by mjContact.frame. Unlike all other matrices in mjData, this matrix is stored in
transposed form. Normally a 3-by-3 matrix corresponding to a coordinate frame would have the frame axes along the
doc/python.rst
+23 -5
View File
@@ -601,6 +601,14 @@ Lists of all elements in a spec can be accessed using named properties, using th
``equalities``, ``tendons``, ``actuators``, ``skins``, ``textures``, ``texts``, ``tuples``, ``flexes``, ``hfields``,
``keys``, ``numerics``, ``excludes``, ``sensors``, ``plugins``.
Element removal
^^^^^^^^^^^^^^^
For elements that can have children (bodies and defaults), the methods ``spec.detach_body(body)`` and
``spec.detach_default(def)`` remove, respectively, ``body`` and ``def`` from the spec, together with all of their
children. When detaching body subtrees, all elements which reference elements in the subtree, will also be removed. For
all other elements, the method ``delete()`` removes the corresponding element from the spec, e.g.
``spec.geom('my_geom').delete()`` will remove the geom named "my_geom" and all of the elements that reference it.
Tree traversal
^^^^^^^^^^^^^^
Traversal of the kinematic tree is aided by the following methods which return tree-related lists of elements:
@@ -620,8 +628,16 @@ Parent:
The parent body of a given element -- including bodies and frames -- can be accessed via the ``parent`` property.
For example, the parent of a site can be accessed via ``site.parent``.
Relationship to ``PyMJCF``
--------------------------
Serialization
^^^^^^^^^^^^^
The ``MjSpec`` object can be serialized with all of its assets using the function ``spec.to_zip(file)``, where ``file``
can be either a path to a file or a file object. In order to load the spec from a zip file, use ``spec =
MjSpec.from_zip(file)``, where ``file`` is a path to a zip file or a zip file object.
.. _PyMJCF:
Relationship to ``PyMJCF`` and ``bind``
---------------------------------------
`dm_control <https://github.com/google-deepmind/dm_control/tree/main>`__'s
`PyMJCF <https://github.com/google-deepmind/dm_control/blob/main/dm_control/mjcf/README.md>`__ module provides similar
@@ -637,9 +653,8 @@ includes a reimplementation of the ``PyMJCF`` example in the ``dm_control``
``PyMJCF`` provides a notion of "binding", giving access to :ref:`mjModel` and :ref:`mjData` values via a helper class.
In the native API, the helper class is not needed, so it is possible to directly bind an ``mjs`` object to
:ref:`mjModel` and :ref:`mjData`. This requires the objects to have a non-empty name. For example, say we have multiple
geoms containing the string "torso" in their name. We want to get their Cartesian positions in the XY plane from
``mjData``. This can be done as follows:
:ref:`mjModel` and :ref:`mjData`. For example, say we have multiple geoms containing the string "torso" in their name.
We want to get their Cartesian positions in the XY plane from ``mjData``. This can be done as follows:
.. code-block:: python
@@ -647,6 +662,9 @@ geoms containing the string "torso" in their name. We want to get their Cartesia
pos_x = [torso.xpos[0] for torso in torsos]
pos_y = [torso.xpos[1] for torso in torsos]
Using the ``bind`` method requires the :ref:`mjModel` and :ref:`mjData` to be compiled from the :`ref:`mjSpec`. If
objects are added or removed from the :ref:`mjSpec` since the last compilation, an error is raised.
Notes
-----
doc/requirements.txt
+2 -2
View File
@@ -1,6 +1,6 @@
Sphinx==5.3.0
furo==2022.9.29
sphinxcontrib-bibtex==2.6.1
sphinxcontrib-bibtex==2.6.3
sphinxcontrib-katex==0.9.4
sphinxcontrib-youtube==1.2.0
sphinx-copybutton==0.5.2
@@ -10,7 +10,7 @@ sphinx-toolbox==3.8.2
nbsphinx==0.9.1
pandoc==1.1.0
pygments==2.15.0
jq==1.4.1
jq==1.8.0
Jinja2~=3.0
wheel
# see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe
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.3.1.dylib`` (it can be
``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.3.3.dylib`` (it can be
found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``.
Linux
_____
Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from
``~/.mujoco/mujoco-3.3.1/lib/libmujoco.so.3.3.1`` and rename it as ``libmujoco.so``.
``~/.mujoco/mujoco-3.3.3/lib/libmujoco.so.3.3.3`` and rename it as ``libmujoco.so``.
Windows
_______
include/mujoco/mjdata.h
+68 -30
View File
@@ -182,7 +182,6 @@ struct mjData_ {
// solver statistics
mjSolverStat solver[mjNISLAND*mjNSOLVER]; // solver statistics per island, per iteration
int solver_nisland; // number of islands processed by solver
int solver_niter[mjNISLAND]; // number of solver iterations, per island
int solver_nnz[mjNISLAND]; // number of nonzeros in Hessian or efc_AR, per island
mjtNum solver_fwdinv[2]; // forward-inverse comparison: qfrc, efc
@@ -200,6 +199,7 @@ struct mjData_ {
int nJ; // number of non-zeros in constraint Jacobian
int nA; // number of non-zeros in constraint inverse inertia matrix
int nisland; // number of detected constraint islands
int nidof; // number of dofs in all islands
// global properties
mjtNum time; // simulation time
@@ -295,12 +295,13 @@ struct mjData_ {
int* moment_colind; // column indices in sparse Jacobian (nJmom x 1)
mjtNum* actuator_moment; // actuator moments (nJmom x 1)
// computed by mj_fwdPosition/mj_crb
// computed by mj_fwdPosition/mj_makeM
mjtNum* crb; // com-based composite inertia and mass (nbody x 10)
mjtNum* qM; // total inertia (sparse) (nM x 1)
mjtNum* qM; // inertia (sparse) (nM x 1)
mjtNum* M; // reduced inertia (compressed sparse row) (nC x 1)
// computed by mj_fwdPosition/mj_factorM
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nM x 1)
mjtNum* qLD; // L'*D*L factorization of M (sparse) (nC x 1)
mjtNum* qLDiagInv; // 1/diag(D) (nv x 1)
// computed by mj_collisionTree
@@ -333,27 +334,23 @@ struct mjData_ {
mjtNum* subtree_angmom; // angular momentum about subtree com (nbody x 3)
// computed by mj_Euler or mj_implicit
mjtNum* qH; // L'*D*L factorization of modified M (nM x 1)
mjtNum* qH; // L'*D*L factorization of modified M (nC x 1)
mjtNum* qHDiagInv; // 1/diag(D) of modified M (nv x 1)
// computed by mj_resetData
int* B_rownnz; // body-dof: non-zeros in each row (nbody x 1)
int* B_rowadr; // body-dof: address of each row in B_colind (nbody x 1)
int* B_colind; // body-dof: column indices of non-zeros (nB x 1)
int* M_rownnz; // inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // inertia: address of each row in M_colind (nv x 1)
int* M_colind; // inertia: column indices of non-zeros (nM x 1)
int* mapM2M; // index mapping from M (legacy) to M (CSR) (nM x 1)
int* C_rownnz; // reduced dof-dof: non-zeros in each row (nv x 1)
int* C_rowadr; // reduced dof-dof: address of each row in C_colind (nv x 1)
int* C_colind; // reduced dof-dof: column indices of non-zeros (nC x 1)
int* mapM2C; // index mapping from M to C (nC x 1)
int* D_rownnz; // dof-dof: non-zeros in each row (nv x 1)
int* D_rowadr; // dof-dof: address of each row in D_colind (nv x 1)
int* D_diag; // dof-dof: index of diagonal element (nv x 1)
int* D_colind; // dof-dof: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from M to D (nD x 1)
int* mapD2M; // index mapping from D to M (nM x 1)
int* M_rownnz; // reduced inertia: non-zeros in each row (nv x 1)
int* M_rowadr; // reduced inertia: address of each row in M_colind (nv x 1)
int* M_colind; // reduced inertia: column indices of non-zeros (nC x 1)
int* mapM2M; // index mapping from qM to M (nC x 1)
int* D_rownnz; // full inertia: non-zeros in each row (nv x 1)
int* D_rowadr; // full inertia: address of each row in D_colind (nv x 1)
int* D_diag; // full inertia: index of diagonal element (nv x 1)
int* D_colind; // full inertia: column indices of non-zeros (nD x 1)
int* mapM2D; // index mapping from qM to D (nD x 1)
int* mapD2M; // index mapping from D to qM (nM x 1)
// computed by mj_implicit/mj_derivative
mjtNum* qDeriv; // d (passive + actuator - bias) / d qvel (nD x 1)
@@ -375,8 +372,8 @@ struct mjData_ {
mjtNum* qfrc_constraint; // constraint force (nv x 1)
// computed by mj_inverse
mjtNum* qfrc_inverse; // net external force; should equal: (nv x 1)
// qfrc_applied + J'*xfrc_applied + qfrc_actuator
mjtNum* qfrc_inverse; // net external force; should equal:
// qfrc_applied + J'*xfrc_applied + qfrc_actuator (nv x 1)
// computed by mj_sensorAcc/mj_rnePostConstraint if needed; rotation:translation format
mjtNum* cacc; // com-based acceleration (nbody x 6)
@@ -410,16 +407,50 @@ struct mjData_ {
mjtNum* efc_R; // inverse constraint mass (nefc x 1)
int* tendon_efcadr; // first efc address involving tendon; -1: none (ntendon x 1)
// computed by mj_island
// computed by mj_island (island dof structure)
int* dof_island; // island id of this dof; -1: none (nv x 1)
int* island_dofnum; // number of dofs in island (nisland x 1)
int* island_dofadr; // start address in island_dofind (nisland x 1)
int* island_dofind; // island dof indices; -1: none (nv x 1)
int* dof_islandind; // dof island indices; -1: none (nv x 1)
int* island_nv; // number of dofs in this island (nisland x 1)
int* island_idofadr; // island start address in idof vector (nisland x 1)
int* island_dofadr; // island start address in dof vector (nisland x 1)
int* map_dof2idof; // map from dof to idof (nv x 1)
int* map_idof2dof; // map from idof to dof; >= nidof: unconstrained (nv x 1)
// 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)
int* efc_island; // island id of this constraint (nefc x 1)
int* island_efcnum; // number of constraints in island (nisland x 1)
int* island_efcadr; // start address in island_efcind (nisland x 1)
int* island_efcind; // island constraint indices (nefc x 1)
int* island_ne; // number of equality constraints in island (nisland x 1)
int* island_nf; // number of friction constraints in island (nisland x 1)
int* island_nefc; // number of constraints in island (nisland x 1)
int* island_iefcadr; // start address in iefc vector (nisland x 1)
int* map_efc2iefc; // map from efc to iefc (nefc x 1)
int* map_iefc2efc; // map from iefc to efc (nefc x 1)
// 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)
int* iefc_JT_rownnz; // number of non-zeros in constraint Jacobian row T (nidof x 1)
int* iefc_JT_rowadr; // row start address in colind array T (nidof x 1)
int* iefc_JT_rowsuper; // number of subsequent rows in supernode T (nidof x 1)
int* iefc_JT_colind; // column indices in constraint Jacobian T (nJ x 1)
mjtNum* iefc_J; // constraint Jacobian (nJ x 1)
mjtNum* iefc_JT; // constraint Jacobian transposed (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)
// computed by mj_projectConstraint (PGS solver)
int* efc_AR_rownnz; // number of non-zeros in AR (nefc x 1)
@@ -437,11 +468,18 @@ struct mjData_ {
// computed by mj_fwdConstraint/mj_inverse
mjtNum* efc_b; // linear cost term: J*qacc_smooth - aref (nefc x 1)
mjtNum* efc_force; // constraint force in constraint space (nefc x 1)
mjtNum* iefc_aref; // reference pseudo-acceleration (nefc x 1)
int* iefc_state; // constraint state (mjtConstraintState) (nefc x 1)
mjtNum* iefc_force; // constraint force in constraint space (nefc x 1)
int* efc_state; // constraint state (mjtConstraintState) (nefc x 1)
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
};
typedef struct mjData_ mjData;
include/mujoco/mjmodel.h
+40 -3
View File
@@ -129,6 +129,14 @@ typedef enum mjtCamLight_ { // tracking mode for camera and light
} 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
@@ -151,6 +159,13 @@ typedef enum mjtTextureRole_ { // role of texture map in rendering
} 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
@@ -269,7 +284,8 @@ typedef enum mjtObj_ { // type of MujoCo object
// meta elements, do not appear in mjModel
mjOBJ_FRAME = 100, // frame
mjOBJ_DEFAULT // default
mjOBJ_DEFAULT, // default
mjOBJ_MODEL // entire model
} mjtObj;
@@ -316,6 +332,7 @@ typedef enum mjtSensor_ { // type of sensor
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
@@ -405,6 +422,14 @@ typedef enum mjtFlexSelf_ { // mode for flex selfcollide
} 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 -------------------------------------------------------
struct mjLROpt_ { // options for mj_setLengthRange()
@@ -611,7 +636,7 @@ struct mjModel_ {
int ncam; // number of cameras
int nlight; // number of lights
int nflex; // number of flexes
int nflexnode; // number of dofs in all flexes
int nflexnode; // number of dofs in all flexes
int nflexvert; // number of vertices in all flexes
int nflexedge; // number of edges in all flexes
int nflexelem; // number of elements in all flexes
@@ -832,9 +857,12 @@ struct mjModel_ {
int* light_mode; // light tracking mode (mjtCamLight) (nlight x 1)
int* light_bodyid; // id of light's body (nlight x 1)
int* light_targetbodyid; // id of targeted body; -1: none (nlight x 1)
mjtByte* light_directional; // directional light (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)
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)
mjtNum* light_pos; // position rel. to body frame (nlight x 3)
mjtNum* light_dir; // direction rel. to body frame (nlight x 3)
@@ -886,6 +914,7 @@ struct mjModel_ {
int* flex_nodebodyid; // node body ids (nflexnode x 1)
int* flex_vertbodyid; // vertex body ids (nflexvert x 1)
int* flex_edge; // edge vertex ids (2 per edge) (nflexedge x 2)
int* flex_edgeflap; // adjacent vertex ids (dim=2 only) (nflexedge x 2)
int* flex_elem; // element vertex ids (dim+1 per elem) (nflexelemdata x 1)
int* flex_elemtexcoord; // element texture coordinates (dim+1) (nflexelemdata x 1)
int* flex_elemedge; // element edge ids (nflexelemedge x 1)
@@ -900,6 +929,7 @@ struct mjModel_ {
mjtNum* flexedge_invweight0; // edge inv. weight in qpos0 (nflexedge x 1)
mjtNum* flex_radius; // radius around primitive element (nflex x 1)
mjtNum* flex_stiffness; // finite element stiffness matrix (nflexelem x 21)
mjtNum* flex_bending; // bending stiffness (nflexedge x 16)
mjtNum* flex_damping; // Rayleigh's damping coefficient (nflex x 1)
mjtNum* flex_edgestiffness; // edge stiffness (nflex x 1)
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
@@ -980,6 +1010,7 @@ struct mjModel_ {
// textures
int* tex_type; // texture type (mjtTexture) (ntex x 1)
int* tex_colorspace; // texture colorspace (mjtColorSpace) (ntex x 1)
int* tex_height; // number of rows in texture image (ntex x 1)
int* tex_width; // number of columns in texture image (ntex x 1)
int* tex_nchannel; // number of channels in texture image (ntex x 1)
@@ -1030,15 +1061,18 @@ struct mjModel_ {
int* tendon_matid; // material id for rendering (ntendon x 1)
int* tendon_group; // group for visibility (ntendon 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)
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)
mjtNum* tendon_solref_fri; // constraint solver reference: friction (ntendon x mjNREF)
mjtNum* tendon_solimp_fri; // constraint solver impedance: friction (ntendon x mjNIMP)
mjtNum* tendon_range; // tendon length limits (ntendon x 2)
mjtNum* tendon_actfrcrange; // range of total actuator force (ntendon x 2)
mjtNum* tendon_margin; // min distance for limit detection (ntendon x 1)
mjtNum* tendon_stiffness; // stiffness coefficient (ntendon x 1)
mjtNum* tendon_damping; // damping coefficient (ntendon x 1)
mjtNum* tendon_armature; // inertia associated with tendon velocity (ntendon x 1)
mjtNum* tendon_frictionloss; // loss due to friction (ntendon x 1)
mjtNum* tendon_lengthspring; // spring resting length range (ntendon x 2)
mjtNum* tendon_length0; // tendon length in qpos0 (ntendon x 1)
@@ -1155,6 +1189,9 @@ struct mjModel_ {
// paths
char* paths; // paths to assets, 0-terminated (npaths x 1)
// compilation signature
uint64_t signature; // also held by the mjSpec that compiled this model
};
typedef struct mjModel_ mjModel;
include/mujoco/mjplugin.h
+10
View File
@@ -135,6 +135,16 @@ struct mjpPlugin_ {
};
typedef struct mjpPlugin_ mjpPlugin;
struct mjSDF_ {
const mjpPlugin** plugin;
int* id;
mjtSDFType type;
mjtNum* relpos;
mjtNum* relmat;
mjtGeom* geomtype;
};
typedef struct mjSDF_ mjSDF;
#if defined(__has_attribute)
#if __has_attribute(constructor)
include/mujoco/mjspec.h
+16 -4
View File
@@ -23,6 +23,7 @@
// this is a C-API
#ifdef __cplusplus
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
@@ -119,6 +120,7 @@ typedef enum mjtOrientation_ { // type of orientation specifier
typedef struct mjsElement_ { // element type, do not modify
mjtObj elemtype; // element type
uint64_t signature; // compilation signature
} mjsElement;
@@ -136,6 +138,7 @@ typedef struct mjsCompiler_ { // compiler options
mjtByte fusestatic; // fuse static bodies with parent
int inertiafromgeom; // use geom inertias (mjtInertiaFromGeom)
int inertiagrouprange[2]; // range of geom groups used to compute inertia
mjtByte saveinertial; // save explicit inertial clause for all bodies to XML
int alignfree; // align free joints with inertial frame
mjLROpt LRopt; // options for lengthrange computation
} mjsCompiler;
@@ -391,9 +394,12 @@ typedef struct mjsLight_ { // light specification
// intrinsics
mjtByte active; // is light active
mjtByte directional; // is light directional or spot
mjtLightType type; // type of light
mjString* texture; // texture name for image lights
mjtByte castshadow; // does light cast shadows
double bulbradius; // bulb radius, for soft shadows
float bulbradius; // bulb radius, for soft shadows
float intensity; // intensity, in candelas
float range; // range of effectiveness
float attenuation[3]; // OpenGL attenuation (quadratic model)
float cutoff; // OpenGL cutoff
float exponent; // OpenGL exponent
@@ -427,7 +433,8 @@ typedef struct mjsFlex_ { // flex specification
double radius; // radius around primitive element
mjtByte internal; // enable internal collisions
mjtByte flatskin; // render flex skin with flat shading
int selfcollide; // mode for flex self colllision
int selfcollide; // mode for flex self collision
int vertcollide; // mode for vertex collision
int activelayers; // number of active element layers in 3D
int group; // group for visualizatioh
double edgestiffness; // edge stiffness
@@ -438,6 +445,7 @@ typedef struct mjsFlex_ { // flex specification
double poisson; // Poisson's ratio
double damping; // Rayleigh's damping
double thickness; // thickness (2D only)
int elastic2d; // 2D passive forces; 0: none, 1: bending, 2: stretching, 3: both
// mesh properties
mjStringVec* nodebody; // node body names
@@ -518,6 +526,7 @@ typedef struct mjsTexture_ { // texture specification
mjsElement* element; // element type
mjString* name; // name
mjtTexture type; // texture type
mjtColorSpace colorspace; // colorspace
// method 1: builtin
int builtin; // builtin type (mjtBuiltin)
@@ -614,17 +623,20 @@ typedef struct mjsTendon_ { // tendon specification
mjsElement* element; // element type
mjString* name; // name
// stiffness, damping, friction
// stiffness, damping, friction, armature
double stiffness; // stiffness coefficient
double springlength[2]; // spring resting length; {-1, -1}: use qpos_spring
double damping; // damping coefficient
double frictionloss; // friction loss
mjtNum solref_friction[mjNREF]; // solver reference: tendon friction
mjtNum solimp_friction[mjNIMP]; // solver impedance: tendon friction
double armature; // inertia associated with tendon velocity
// length range
int limited; // does tendon have limits (mjtLimited)
int actfrclimited; // does tendon have actuator force limits
double range[2]; // length limits
double actfrcrange[2]; // actuator force limits
double margin; // margin value for tendon limit detection
mjtNum solref_limit[mjNREF]; // solver reference: tendon limits
mjtNum solimp_limit[mjNIMP]; // solver impedance: tendon limits
include/mujoco/mjvisualize.h
+5 -287
View File
@@ -226,7 +226,7 @@ typedef struct mjvGLCamera_ mjvGLCamera;
struct mjvGeom_ { // abstract geom
// type info
int type; // geom type (mjtGeom)
int dataid; // mesh, hfield or plane id; -1: none
int dataid; // mesh, hfield or plane id; -1: none; mesh: 2*id or 2*id+1 (hull)
int objtype; // mujoco object type; mjOBJ_UNKNOWN for decor
int objid; // mujoco object id; -1 for decor
int category; // visual category
@@ -261,6 +261,8 @@ typedef struct mjvGeom_ mjvGeom;
struct mjvLight_ { // OpenGL light
float pos[3]; // position rel. to body frame
float dir[3]; // direction rel. to body frame
int type; // type (mjtLightType)
int texid; // texture id for image lights
float attenuation[3]; // OpenGL attenuation (quadratic model)
float cutoff; // OpenGL cutoff
float exponent; // OpenGL exponent
@@ -268,9 +270,10 @@ struct mjvLight_ { // OpenGL light
float diffuse[3]; // diffuse rgb (alpha=1)
float specular[3]; // specular rgb (alpha=1)
mjtByte headlight; // headlight
mjtByte directional; // directional light
mjtByte castshadow; // does light cast shadows
float bulbradius; // bulb radius for soft shadows
float intensity; // intensity, in candelas
float range; // range of effectiveness
};
typedef struct mjvLight_ mjvLight;
@@ -404,289 +407,4 @@ struct mjvFigure_ { // abstract 2D figure passed to OpenGL rendere
};
typedef struct mjvFigure_ mjvFigure;
//---------------------------------- mjvSceneState -------------------------------------------------
struct mjvSceneState_ {
int nbuffer; // size of the buffer in bytes
void* buffer; // heap-allocated memory for all arrays in this struct
int maxgeom; // maximum number of mjvGeom supported by this state object
mjvScene scratch; // scratch space for vis geoms inserted by the user and plugins
// fields in mjModel that are necessary to re-render a scene
struct {
int nv;
int nu;
int na;
int nbody;
int nbvh;
int nbvhstatic;
int njnt;
int ngeom;
int nsite;
int ncam;
int nlight;
int nmesh;
int nskin;
int nflex;
int nflexvert;
int nflextexcoord;
int nskinvert;
int nskinface;
int nskinbone;
int nskinbonevert;
int nmat;
int neq;
int ntendon;
int ntree;
int nwrap;
int nsensor;
int nnames;
int npaths;
int nsensordata;
int narena;
mjOption opt;
mjVisual vis;
mjStatistic stat;
int* body_parentid;
int* body_rootid;
int* body_weldid;
int* body_mocapid;
int* body_jntnum;
int* body_jntadr;
int* body_dofnum;
int* body_dofadr;
int* body_geomnum;
int* body_geomadr;
mjtNum* body_iquat;
mjtNum* body_mass;
mjtNum* body_inertia;
int* body_bvhadr;
int* body_bvhnum;
int* bvh_depth;
int* bvh_child;
int* bvh_nodeid;
mjtNum* bvh_aabb;
int* jnt_type;
int* jnt_bodyid;
int* jnt_group;
int* geom_type;
int* geom_bodyid;
int* geom_contype;
int* geom_conaffinity;
int* geom_dataid;
int* geom_matid;
int* geom_group;
mjtNum* geom_size;
mjtNum* geom_aabb;
mjtNum* geom_rbound;
float* geom_rgba;
int* site_type;
int* site_bodyid;
int* site_matid;
int* site_group;
mjtNum* site_size;
float* site_rgba;
int* cam_orthographic;
mjtNum* cam_fovy;
mjtNum* cam_ipd;
int* cam_resolution;
float* cam_sensorsize;
float* cam_intrinsic;
mjtByte* light_directional;
mjtByte* light_castshadow;
float* light_bulbradius;
mjtByte* light_active;
float* light_attenuation;
float* light_cutoff;
float* light_exponent;
float* light_ambient;
float* light_diffuse;
float* light_specular;
mjtByte* flex_flatskin;
int* flex_dim;
int* flex_matid;
int* flex_group;
int* flex_interp;
int* flex_nodeadr;
int* flex_nodenum;
int* flex_nodebodyid;
int* flex_vertadr;
int* flex_vertnum;
int* flex_elem;
int* flex_elemtexcoord;
int* flex_elemlayer;
int* flex_elemadr;
int* flex_elemnum;
int* flex_elemdataadr;
int* flex_shell;
int* flex_shellnum;
int* flex_shelldataadr;
int* flex_texcoordadr;
int* flex_bvhadr;
int* flex_bvhnum;
mjtByte* flex_centered;
mjtNum* flex_node;
mjtNum* flex_radius;
float* flex_rgba;
float* flex_texcoord;
int* hfield_pathadr;
int* mesh_bvhadr;
int* mesh_bvhnum;
int* mesh_texcoordadr;
int* mesh_graphadr;
int* mesh_pathadr;
int* skin_matid;
int* skin_group;
float* skin_rgba;
float* skin_inflate;
int* skin_vertadr;
int* skin_vertnum;
int* skin_texcoordadr;
int* skin_faceadr;
int* skin_facenum;
int* skin_boneadr;
int* skin_bonenum;
float* skin_vert;
int* skin_face;
int* skin_bonevertadr;
int* skin_bonevertnum;
float* skin_bonebindpos;
float* skin_bonebindquat;
int* skin_bonebodyid;
int* skin_bonevertid;
float* skin_bonevertweight;
int* skin_pathadr;
int* tex_pathadr;
int* mat_texid;
mjtByte* mat_texuniform;
float* mat_texrepeat;
float* mat_emission;
float* mat_specular;
float* mat_shininess;
float* mat_reflectance;
float* mat_metallic;
float* mat_roughness;
float* mat_rgba;
int* eq_type;
int* eq_obj1id;
int* eq_obj2id;
int* eq_objtype;
mjtNum* eq_data;
int* tendon_num;
int* tendon_matid;
int* tendon_group;
mjtByte* tendon_limited;
mjtNum* tendon_width;
mjtNum* tendon_range;
mjtNum* tendon_stiffness;
mjtNum* tendon_damping;
mjtNum* tendon_frictionloss;
mjtNum* tendon_lengthspring;
float* tendon_rgba;
int* actuator_trntype;
int* actuator_dyntype;
int* actuator_trnid;
int* actuator_actadr;
int* actuator_actnum;
int* actuator_group;
mjtByte* actuator_ctrllimited;
mjtByte* actuator_actlimited;
mjtNum* actuator_ctrlrange;
mjtNum* actuator_actrange;
mjtNum* actuator_cranklength;
int* sensor_type;
int* sensor_objid;
int* sensor_adr;
int* name_bodyadr;
int* name_jntadr;
int* name_geomadr;
int* name_siteadr;
int* name_camadr;
int* name_lightadr;
int* name_eqadr;
int* name_tendonadr;
int* name_actuatoradr;
char* names;
char* paths;
} model;
// fields in mjData that are necessary to re-render a scene
struct {
mjWarningStat warning[mjNWARNING];
int nefc;
int ncon;
int nisland;
mjtNum time;
mjtNum* act;
mjtNum* ctrl;
mjtNum* xfrc_applied;
mjtByte* eq_active;
mjtNum* sensordata;
mjtNum* xpos;
mjtNum* xquat;
mjtNum* xmat;
mjtNum* xipos;
mjtNum* ximat;
mjtNum* xanchor;
mjtNum* xaxis;
mjtNum* geom_xpos;
mjtNum* geom_xmat;
mjtNum* site_xpos;
mjtNum* site_xmat;
mjtNum* cam_xpos;
mjtNum* cam_xmat;
mjtNum* light_xpos;
mjtNum* light_xdir;
mjtNum* subtree_com;
int* ten_wrapadr;
int* ten_wrapnum;
int* wrap_obj;
mjtNum* ten_length;
mjtNum* wrap_xpos;
mjtNum* bvh_aabb_dyn;
mjtByte* bvh_active;
int* island_dofadr;
int* island_dofind;
int* dof_island;
int* efc_island;
int* tendon_efcadr;
mjtNum* flexvert_xpos;
mjContact* contact;
mjtNum* efc_force;
void* arena;
} data;
};
typedef struct mjvSceneState_ mjvSceneState;
#endif // MUJOCO_MJVISUALIZE_H_
include/mujoco/mjxmacro.h
+393 -361
View File
File diff suppressed because it is too large Load Diff
include/mujoco/mujoco.h
+39 -48
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 331
#define mjVERSION_HEADER 333
// needed to define size_t, fabs and log10
#include <stdlib.h>
@@ -107,6 +107,9 @@ MJAPI mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error,
// Compile spec to model.
MJAPI mjModel* mj_compile(mjSpec* s, const mjVFS* vfs);
// Copy real-valued arrays from model to spec, returns 1 on success.
MJAPI int mj_copyBack(mjSpec* s, const mjModel* m);
// Recompile spec to model, preserving the state, return 0 on success.
MJAPI int mj_recompile(mjSpec* s, const mjVFS* vfs, mjModel* m, mjData* d);
@@ -187,6 +190,9 @@ MJAPI mjData* mj_makeData(const mjModel* m);
// m is only required to contain the size fields from MJMODEL_INTS.
MJAPI mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
// Copy mjData, skip large arrays not required for visualization.
MJAPI mjData* mjv_copyData(mjData* dest, const mjModel* m, const mjData* src);
// Reset data to defaults.
MJAPI void mj_resetData(const mjModel* m, mjData* d);
@@ -361,6 +367,9 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
// Run composite rigid body inertia algorithm (CRB).
MJAPI void mj_crb(const mjModel* m, mjData* d);
// Make inertia matrix.
MJAPI void mj_makeM(const mjModel* m, mjData* d);
// Compute sparse L'*D*L factorizaton of inertia matrix.
MJAPI void mj_factorM(const mjModel* m, mjData* d);
@@ -485,7 +494,7 @@ MJAPI void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
MJAPI void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
// Add inertia matrix to destination matrix.
// Destination can be sparse uncompressed, or dense when all int* are NULL
// Destination can be sparse or dense when all int* are NULL.
MJAPI void mj_addM(const mjModel* m, mjData* d, mjtNum* dst, int* rownnz, int* rowadr, int* colind);
// Apply Cartesian force and torque (outside xfrc_applied mechanism).
@@ -622,20 +631,10 @@ MJAPI void mjv_alignToCamera(mjtNum res[3], const mjtNum vec[3], const mjtNum fo
MJAPI void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvCamera* cam);
// Move camera with mouse given a scene state; action is mjtMouse.
MJAPI void mjv_moveCameraFromState(const mjvSceneState* scnstate, int action,
mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvCamera* cam);
// Move perturb object with mouse; action is mjtMouse.
MJAPI void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
// Move perturb object with mouse given a scene state; action is mjtMouse.
MJAPI void mjv_movePerturbFromState(const mjvSceneState* scnstate, int action,
mjtNum reldx, mjtNum reldy,
const mjvScene* scn, mjvPerturb* pert);
// Move model with mouse; action is mjtMouse.
MJAPI void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
const mjtNum roomup[3], mjvScene* scn);
@@ -692,28 +691,9 @@ MJAPI void mjv_freeScene(mjvScene* scn);
MJAPI void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
// Update entire scene from a scene state, return the number of new mjWARN_VGEOMFULL warnings.
MJAPI int mjv_updateSceneFromState(const mjvSceneState* scnstate, const mjvOption* opt,
const mjvPerturb* pert, mjvCamera* cam, int catmask,
mjvScene* scn);
// Copy mjModel, skip large arrays not required for abstract visualization.
MJAPI void mjv_copyModel(mjModel* dest, const mjModel* src);
// Set default scene state.
MJAPI void mjv_defaultSceneState(mjvSceneState* scnstate);
// Allocate resources and initialize a scene state object.
MJAPI void mjv_makeSceneState(const mjModel* m, const mjData* d,
mjvSceneState* scnstate, int maxgeom);
// Free scene state.
MJAPI void mjv_freeSceneState(mjvSceneState* scnstate);
// Update a scene state from model and data.
MJAPI void mjv_updateSceneState(const mjModel* m, mjData* d, const mjvOption* opt,
mjvSceneState* scnstate);
// Add geoms from selected categories.
MJAPI void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
const mjvPerturb* pert, int catmask, mjvScene* scn);
@@ -1308,6 +1288,19 @@ MJAPI char* mju_strncpy(char *dst, const char *src, int n);
MJAPI mjtNum mju_sigmoid(mjtNum x);
//---------------------------------- Signed Distance Function --------------------------------------
// get sdf from geom id
MJAPI const mjpPlugin* mjc_getSDF(const mjModel* m, int id);
// signed distance function
MJAPI mjtNum mjc_distance(const mjModel* m, const mjData* d, const mjSDF* s, const mjtNum x[3]);
// gradient of sdf
MJAPI void mjc_gradient(const mjModel* m, const mjData* d, const mjSDF* s, mjtNum gradient[3],
const mjtNum x[3]);
//---------------------------------- Derivatives ---------------------------------------------------
// Finite differenced transition matrices (control theory notation)
@@ -1414,21 +1407,9 @@ MJAPI void mju_taskJoin(mjTask* task);
//---------------------------------- Attachment ----------------------------------------------------
// Attach child body to a parent frame, return the attached body if success or NULL otherwise.
MJAPI mjsBody* mjs_attachBody(mjsFrame* parent, const mjsBody* child,
const char* prefix, const char* suffix);
// Attach child frame to a parent body, return the attached frame if success or NULL otherwise.
MJAPI mjsFrame* mjs_attachFrame(mjsBody* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
// Attach child body to a parent site, return the attached body if success or NULL otherwise.
MJAPI mjsBody* mjs_attachToSite(mjsSite* parent, const mjsBody* child,
const char* prefix, const char* suffix);
// Attach child frame to a parent site, return the attached frame if success or NULL otherwise.
MJAPI mjsFrame* mjs_attachFrameToSite(mjsSite* parent, const mjsFrame* child,
const char* prefix, const char* suffix);
// Attach child to a parent, return the attached element if success or NULL otherwise.
MJAPI mjsElement* mjs_attach(mjsElement* parent, const mjsElement* child,
const char* prefix, const char* suffix);
// Delete body and descendants from mjSpec, remove all references, return 0 on success.
MJAPI int mjs_detachBody(mjSpec* s, mjsBody* b);
@@ -1634,14 +1615,17 @@ MJAPI const char* mjs_getString(const mjString* source);
// Get double array contents and optionally its size.
MJAPI const double* mjs_getDouble(const mjDoubleVec* source, int* size);
// Get plugin attributes.
MJAPI const void* mjs_getPluginAttributes(const mjsPlugin* plugin);
//---------------------------------- Spec utilities ------------------------------------------------
// Set element's default.
MJAPI void mjs_setDefault(mjsElement* element, const mjsDefault* def);
// Set element's enclosing frame.
MJAPI void mjs_setFrame(mjsElement* dest, mjsFrame* frame);
// Set element's enclosing frame, return 0 on success.
MJAPI int mjs_setFrame(mjsElement* dest, mjsFrame* frame);
// Resolve alternative orientations to quat, return error if any.
MJAPI const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* sequence,
@@ -1653,6 +1637,13 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body);
// Set user payload, overriding the existing value for the specified key if present.
MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data);
// Set user payload, overriding the existing value for the specified key if
// present. This version differs from mjs_setUserValue in that it takes a
// cleanup function that will be called when the user payload is deleted.
MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key,
const void* data,
void (*cleanup)(const void*));
// Return user payload or NULL if none found.
MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key);
mjx/mujoco/mjx/__init__.py
+2
View File
@@ -17,6 +17,7 @@
# pylint:disable=g-importing-member
from mujoco.mjx._src.collision_driver import collision
from mujoco.mjx._src.constraint import make_constraint
from mujoco.mjx._src.derivative import deriv_smooth_vel
from mujoco.mjx._src.forward import euler
from mujoco.mjx._src.forward import forward
from mujoco.mjx._src.forward import fwd_acceleration
@@ -26,6 +27,7 @@ from mujoco.mjx._src.forward import fwd_velocity
from mujoco.mjx._src.forward import implicit
from mujoco.mjx._src.forward import rungekutta4
from mujoco.mjx._src.forward import step
from mujoco.mjx._src.inverse import inverse
from mujoco.mjx._src.io import get_data
from mujoco.mjx._src.io import get_data_into
from mujoco.mjx._src.io import make_data
mjx/mujoco/mjx/_src/collision_convex.py
+6 -1
View File
@@ -28,8 +28,10 @@ from mujoco.mjx._src.collision_types import FunctionKey
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.collision_types import HFieldInfo
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
# pylint: enable=g-importing-member
_GeomInfo = Union[GeomInfo, ConvexInfo]
@@ -42,6 +44,9 @@ def collider(ncon: int):
def collide(
m: Model, d: Data, key: FunctionKey, geom: jax.Array
) -> Collision:
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('collider requires JAX backend implementation.')
g1, g2 = geom.T
infos = [
GeomInfo(d.geom_xpos[g1], d.geom_xmat[g1], m.geom_size[g1]),
@@ -56,7 +61,7 @@ def collider(ncon: int):
pos=0, mat=0, size=0, face=0, vert=0
)
elif key.types[i] == GeomType.MESH:
c, cm = infos[i], m.mesh_convex[key.data_ids[i]]
c, cm = infos[i], m._impl.mesh_convex[key.data_ids[i]]
infos[i] = ConvexInfo(**vars(c), **vars(cm))
in_axes[i] = jax.tree_util.tree_map(lambda x: None, infos[i]).replace(
pos=0, mat=0, size=0
mjx/mujoco/mjx/_src/collision_driver.py
+16 -7
View File
@@ -71,9 +71,11 @@ from mujoco.mjx._src.collision_sdf import sphere_ellipsoid
from mujoco.mjx._src.collision_types import FunctionKey
from mujoco.mjx._src.types import Contact
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
# pylint: enable=g-importing-member
import numpy as np
@@ -227,7 +229,7 @@ def _geom_groups(
if types[0] == mujoco.mjtGeom.mjGEOM_HFIELD:
# add static grid bounds to the grouping key for hfield collisions
geom_rbound_hfield = (
m.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound
m._impl.geom_rbound_hfield if isinstance(m, Model) else m.geom_rbound # pytype: disable=attribute-error
)
nrow, ncol = m.hfield_nrow[data_ids[0]], m.hfield_ncol[data_ids[0]]
xsize, ysize = m.hfield_size[data_ids[0]][:2]
@@ -323,11 +325,11 @@ def _contact_groups(m: Model, d: Data) -> Dict[FunctionKey, Contact]:
solref=solref,
solreffriction=solreffriction,
solimp=solimp,
dim=d.contact.dim,
dim=d._impl.contact.dim, # pytype: disable=attribute-error
geom1=jp.array(geom[:, 0]),
geom2=jp.array(geom[:, 1]),
geom=jp.array(geom[:, :2]),
efc_address=d.contact.efc_address,
efc_address=d._impl.contact.efc_address, # pytype: disable=attribute-error
)
return groups
@@ -361,8 +363,12 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
condim_counts = {}
for k, v in group_counts.items():
func = _COLLISION_FUNC[k.types]
num_contacts = condim_counts.get(k.condim, 0) + func.ncon * v # pytype: disable=attribute-error
if k.types[1] == mujoco.mjtGeom.mjGEOM_SDF:
ncon = m.opt.sdf_initpoints
else:
func = _COLLISION_FUNC[k.types]
ncon = func.ncon # pytype: disable=attribute-error
num_contacts = condim_counts.get(k.condim, 0) + ncon * v
if max_contact_points > -1:
num_contacts = min(max_contact_points, num_contacts)
condim_counts[k.condim] = num_contacts
@@ -374,7 +380,10 @@ def make_condim(m: Union[Model, mujoco.MjModel]) -> np.ndarray:
def collision(m: Model, d: Data) -> Data:
"""Collides geometries."""
if d.ncon == 0:
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('collision requires JAX backend implementation.')
if d._impl.ncon == 0: # pytype: disable=attribute-error
return d
max_geom_pairs = _numeric(m, 'max_geom_pairs')
@@ -424,4 +433,4 @@ def collision(m: Model, d: Data) -> Data:
contacts = sum([condim_groups[k] for k in sorted(condim_groups)], [])
contact = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *contacts)
return d.replace(contact=contact)
return d.tree_replace({'_impl.contact': contact})
mjx/mujoco/mjx/_src/collision_driver_test.py
+99 -87
View File
@@ -119,7 +119,7 @@ class SphereCollisionTest(parameterized.TestCase):
def test_sphere(self, name, mjcf):
d, dx = _collide(mjcf)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5)
_assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-5)
_SPHERE_CONVEX = """
<mujoco>
@@ -143,7 +143,7 @@ class SphereCollisionTest(parameterized.TestCase):
)
d, dx = _collide(xml)
self.assertEmpty(d.contact.dist)
self.assertGreater(dx.contact.dist, 0)
self.assertGreater(dx._impl.contact.dist, 0)
# face contact
xml = self._SPHERE_CONVEX.replace(
@@ -151,38 +151,38 @@ class SphereCollisionTest(parameterized.TestCase):
)
d, dx = _collide(xml)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'face', 1e-4)
_assert_attr_eq(dx._impl.contact, d.contact, field.name, 'face', 1e-4)
# deep face contact
xml = self._SPHERE_CONVEX.replace(
'<body pos="0.52 0 0.52">', '<body pos="0.48 0 0.47">'
)
d, dx = _collide(xml)
self.assertTrue((dx.contact.dist < 0).all())
self.assertTrue((dx._impl.contact.dist < 0).all())
self.assertTrue((d.contact.dist < 0).all())
np.testing.assert_allclose(dx.contact.dist, [-0.07], atol=1e-5)
np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
np.testing.assert_allclose(dx._impl.contact.dist, [-0.07], atol=1e-5)
np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
np.testing.assert_array_almost_equal(
dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
)
def test_sphere_convex_edge(self):
# edge contact
d, dx = _collide(self._SPHERE_CONVEX)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'edge', 1e-4)
_assert_attr_eq(dx._impl.contact, d.contact, field.name, 'edge', 1e-4)
# deep edge penetration
xml = self._SPHERE_CONVEX.replace(
'<body pos="0.52 0 0.52">', '<body pos="0.49 0 0.49">'
)
d, dx = _collide(xml)
self.assertTrue((dx.contact.dist < 0).all())
self.assertTrue((dx._impl.contact.dist < 0).all())
self.assertTrue((d.contact.dist < 0).all())
np.testing.assert_allclose(dx.contact.dist, [-0.06], atol=1e-5)
np.testing.assert_array_almost_equal(dx.contact.pos, d.contact.pos)
np.testing.assert_allclose(dx._impl.contact.dist, [-0.06], atol=1e-5)
np.testing.assert_array_almost_equal(dx._impl.contact.pos, d.contact.pos)
np.testing.assert_array_almost_equal(
dx.contact.frame, d.contact.frame.reshape((-1, 3, 3))
dx._impl.contact.frame, d.contact.frame.reshape((-1, 3, 3))
)
# vertex contact
@@ -191,7 +191,7 @@ class SphereCollisionTest(parameterized.TestCase):
)
d, dx = _collide(xml)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'vertex', 1e-4)
_assert_attr_eq(dx._impl.contact, d.contact, field.name, 'vertex', 1e-4)
# sphere center on vertex
xml = self._SPHERE_CONVEX.replace(
@@ -199,7 +199,9 @@ class SphereCollisionTest(parameterized.TestCase):
)
d, dx = _collide(xml)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'vertex_center', 1e-4)
_assert_attr_eq(
dx._impl.contact, d.contact, field.name, 'vertex_center', 1e-4
)
class EllipsoidCollisionTest(parameterized.TestCase):
@@ -219,10 +221,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
def test_plane_ellipsoid(self):
"""Tests ellipsoid plane contact."""
d, dx = _collide(self._ELLIPSOID_PLANE)
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
dx._impl.contact, d.contact, field.name, 'ellipsoid-plane', 1e-5
)
_ELLIPSOID_ELLIPSOID = """
@@ -242,10 +244,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
def test_ellipsoid_ellipsoid(self):
"""Tests ellipsoid ellipsoid contact."""
d, dx = _collide(self._ELLIPSOID_ELLIPSOID)
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
dx._impl.contact, d.contact, field.name, 'ellipsoid-ellipsoid', 1e-2
)
_ELLIPSOID_SPHERE = """
@@ -266,10 +268,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
"""Tests ellipsoid capsule contact."""
d, dx = _collide(self._ELLIPSOID_SPHERE)
d.contact.pos[0][2] = 0.03 # MJX finds the point on the surface
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
dx._impl.contact, d.contact, field.name, 'ellipsoid-sphere', 1e-4
)
_ELLIPSOID_CAPSULE = """
@@ -289,10 +291,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
def test_capsule_ellipsoid(self):
"""Tests ellipsoid capsule contact."""
d, dx = _collide(self._ELLIPSOID_CAPSULE)
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
dx._impl.contact, d.contact, field.name, 'ellipsoid-capsule', 1e-5
)
_ELLIPSOID_CYLINDER = """
@@ -313,10 +315,10 @@ class EllipsoidCollisionTest(parameterized.TestCase):
"""Tests ellipsoid cylinder contact."""
d, dx = _collide(self._ELLIPSOID_CYLINDER)
d.contact.pos[0][2] = 0.04 # MJX finds the deepest point on the surface
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
dx._impl.contact, d.contact, field.name, 'ellipsoid-cylinder', 1e-4
)
@@ -357,7 +359,7 @@ class CapsuleCollisionTest(parameterized.TestCase):
def test_capsule(self, name, mjcf):
d, dx = _collide(mjcf)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4)
_assert_attr_eq(dx._impl.contact, d.contact, field.name, name, 1e-4)
_PARALLEL_CAP = """
<mujoco>
@@ -378,14 +380,14 @@ class CapsuleCollisionTest(parameterized.TestCase):
"""Tests that two parallel capsules are colliding at the midpoint."""
_, dx = _collide(self._PARALLEL_CAP)
np.testing.assert_allclose(dx.contact.dist, -0.05)
np.testing.assert_allclose(dx._impl.contact.dist, -0.05)
np.testing.assert_allclose(
dx.contact.pos[0],
dx._impl.contact.pos[0],
np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]),
atol=1e-5,
)
np.testing.assert_allclose(
dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
dx._impl.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
)
_CAP_BOX = """
@@ -408,30 +410,32 @@ class CapsuleCollisionTest(parameterized.TestCase):
d, dx = _collide(self._CAP_BOX)
# sort positions for comparison
idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
d.contact.pos[:] = d.contact.pos[idx]
d.contact.frame[:] = d.contact.frame[idx]
d.contact.dist[:] = d.contact.dist[idx]
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4)
_assert_attr_eq(
dx._impl.contact, d.contact, field.name, 'capsule_convex', 1e-4
)
def test_capsule_convex_face_deep(self):
"""Tests deep face penetration."""
xml = self._CAP_BOX.replace('<body pos="0 0 0.54">', '<body pos="0 0 0.4">')
_, dx = _collide(xml)
self.assertTrue((dx.contact.dist < 0).all())
self.assertTrue((dx._impl.contact.dist < 0).all())
np.testing.assert_array_almost_equal(
dx.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
dx._impl.contact.pos, np.array([[0.5, 0, 0.425], [-0.4, 0, 0.425]])
)
np.testing.assert_array_almost_equal(
dx.contact.dist, np.array([-0.15, -0.15])
dx._impl.contact.dist, np.array([-0.15, -0.15])
)
np.testing.assert_array_almost_equal(
dx.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
dx._impl.contact.frame[:, 0], np.array([[0, 0, -1]] * 2)
)
_CAP_EDGE_BOX = """
@@ -453,11 +457,11 @@ class CapsuleCollisionTest(parameterized.TestCase):
"""Tests edge contact."""
d, dx = _collide(self._CAP_EDGE_BOX)
c = dx.contact
c = dx._impl.contact
self.assertEqual(c.pos.shape[0], 2)
self.assertGreater(c.dist[1], 0)
# extract the contact point with penetration
c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
for field in dataclasses.fields(Contact):
_assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4)
@@ -469,13 +473,17 @@ class CapsuleCollisionTest(parameterized.TestCase):
)
_, dx = _collide(xml)
np.testing.assert_array_equal(dx.contact.dist < 0, np.array([True, False]))
np.testing.assert_array_almost_equal(dx.contact.dist[0], np.array([-0.13]))
np.testing.assert_array_almost_equal(
dx.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
np.testing.assert_array_equal(
dx._impl.contact.dist < 0, np.array([True, False])
)
np.testing.assert_array_almost_equal(
dx.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
dx._impl.contact.dist[0], np.array([-0.13])
)
np.testing.assert_array_almost_equal(
dx._impl.contact.pos[0], np.array([0.5, 0, 0.435]), decimal=3
)
np.testing.assert_array_almost_equal(
dx._impl.contact.frame[0, 0], np.array([0, 0, -1]), decimal=3
)
def test_capsule_convex_edge_shallow_tip(self):
@@ -489,16 +497,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.58 0 0.55"')
d, dx = _collide(xml)
c = dx.contact
c = dx._impl.contact
self.assertEqual(c.pos.shape[0], 2)
self.assertGreater(c.dist[1], 0)
# extract the contact point with penetration
c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
for field in dataclasses.fields(Contact):
_assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip1', 1e-4)
np.testing.assert_array_almost_equal(
dx.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
dx._impl.contact.frame[0][0, :3], np.array([-0.43952, 0.0, -0.898233])
)
# the capsule sphere is outside the edge voronoi region, so there is a
@@ -510,16 +518,16 @@ class CapsuleCollisionTest(parameterized.TestCase):
xml = xml.replace('<body pos="0.5 0 0.55"', '<body pos="0.5 0 0.52"')
d, dx = _collide(xml)
c = dx.contact
c = dx._impl.contact
self.assertEqual(c.pos.shape[0], 2)
self.assertGreater(c.dist[1], 0)
# extract the contact point with penetration
c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
for field in dataclasses.fields(Contact):
_assert_attr_eq(c, d.contact, field.name, 'edge_shallow_tip2', 1e-4)
np.testing.assert_array_almost_equal(
dx.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
dx._impl.contact.frame[0][0, :3], np.array([0.0, 0.0, -1.0])
)
@@ -542,17 +550,17 @@ class CylinderTest(absltest.TestCase):
d, dx = _collide(self._CYLINDER_PLANE)
# cylinder is lying flat
np.testing.assert_array_less(dx.contact.dist[:2], 0)
np.testing.assert_array_less(-dx.contact.dist[2:], 0)
np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
# sort position for comparison
idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
d.contact.pos[:] = d.contact.pos[idx]
# extract the contact points with penetration
c = jax.tree_util.tree_map(lambda x: x[:2], dx.contact)
c = jax.tree_util.tree_map(lambda x: x[:2], dx._impl.contact)
c = c.replace(dim=c.dim[:2], efc_address=c.efc_address[:2])
for field in dataclasses.fields(Contact):
_assert_attr_eq(c, d.contact, field.name, 'cylinder_plane', 1e-5)
@@ -564,9 +572,11 @@ class CylinderTest(absltest.TestCase):
xml = xml.replace('pos="0 0 0.04"', 'pos="0 0 0.095"')
d, dx = _collide(xml)
np.testing.assert_array_less(dx.contact.dist, 0)
np.testing.assert_array_less(dx._impl.contact.dist, 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'cylinder_plane', 1e-5)
_assert_attr_eq(
dx._impl.contact, d.contact, field.name, 'cylinder_plane', 1e-5
)
_SPHERE_CYLINDER = """
<mujoco>
@@ -586,10 +596,10 @@ class CylinderTest(absltest.TestCase):
"""Tests sphere cylinder contact."""
d, dx = _collide(self._SPHERE_CYLINDER)
d.contact.pos[0][2] = 0.05 # MJX finds the deepest point on the surface
self.assertLess(dx.contact.dist[0], 0)
self.assertLess(dx._impl.contact.dist[0], 0)
for field in dataclasses.fields(Contact):
_assert_attr_eq(
dx.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
dx._impl.contact, d.contact, field.name, 'sphere-cylinder', 1e-4
)
@@ -612,11 +622,11 @@ class ConvexTest(absltest.TestCase):
"""Tests box collision with a plane."""
d, dx = _collide(self._BOX_PLANE)
np.testing.assert_array_less(dx.contact.dist[:2], 0)
np.testing.assert_array_less(-dx.contact.dist[2:], 0)
np.testing.assert_array_less(dx._impl.contact.dist[:2], 0)
np.testing.assert_array_less(-dx._impl.contact.dist[2:], 0)
# extract the contact points with penetration
c = jax.tree_util.tree_map(
lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact
lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx._impl.contact
)
c = c.replace(dim=c.dim[[0, 1]], efc_address=c.efc_address[[0, 1]])
for field in dataclasses.fields(Contact):
@@ -638,16 +648,18 @@ class ConvexTest(absltest.TestCase):
"""Tests box collision with a plane."""
d, dx = _collide(self._FLAT_BOX_PLANE)
np.testing.assert_array_less(dx.contact.dist, 0)
np.testing.assert_array_less(dx._impl.contact.dist, 0)
# sort positions for comparison
idx = np.lexsort((dx.contact.pos[:, 0], dx.contact.pos[:, 1]))
dx = dx.tree_replace({'contact.pos': dx.contact.pos[idx]})
idx = np.lexsort((dx._impl.contact.pos[:, 0], dx._impl.contact.pos[:, 1]))
dx = dx.tree_replace({'_impl.contact.pos': dx._impl.contact.pos[idx]})
idx = np.lexsort((d.contact.pos[:, 0], d.contact.pos[:, 1]))
d.contact.pos[:] = d.contact.pos[idx]
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'flat_box_plane', 1e-5)
_assert_attr_eq(
dx._impl.contact, d.contact, field.name, 'flat_box_plane', 1e-5
)
_BOX_BOX = """
<mujoco>
@@ -669,7 +681,7 @@ class ConvexTest(absltest.TestCase):
def test_box_box(self):
"""Tests a face contact for a box-box collision."""
d, dx = _collide(self._BOX_BOX, keyframe=0)
c = dx.contact
c = dx._impl.contact
self.assertEqual(c.pos.shape[0], 4)
np.testing.assert_array_less(c.dist, 0)
@@ -701,10 +713,10 @@ class ConvexTest(absltest.TestCase):
d, dx = _collide(self._BOX_BOX_EDGE)
# Only one contact point.
np.testing.assert_array_less(dx.contact.dist[:1], 0)
np.testing.assert_array_less(-dx.contact.dist[1:], 0)
np.testing.assert_array_less(dx._impl.contact.dist[:1], 0)
np.testing.assert_array_less(-dx._impl.contact.dist[1:], 0)
# extract the contact point with penetration
c = jax.tree_util.tree_map(lambda x: x[:1], dx.contact)
c = jax.tree_util.tree_map(lambda x: x[:1], dx._impl.contact)
c = c.replace(dim=c.dim[:1], efc_address=c.efc_address[:1])
for field in dataclasses.fields(Contact):
_assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
@@ -736,7 +748,7 @@ class ConvexTest(absltest.TestCase):
).read_bytes(),
}
_, dx = _collide(self._CONVEX_CONVEX, assets=assets)
c = dx.contact
c = dx._impl.contact
# Only one contact point for an edge contact.
self.assertLess(c.dist[0], 0)
@@ -766,7 +778,7 @@ class ConvexTest(absltest.TestCase):
def test_convex_convex_edge(self):
"""Tests convex-convex collisions with edge contact via _sat_gaussmap."""
_, dx = _collide(self._CONVEX_CONVEX_THIN)
c = dx.contact
c = dx._impl.contact
# Only one contact point for an edge contact.
self.assertLess(c.dist[0], 0)
@@ -781,7 +793,7 @@ class ConvexTest(absltest.TestCase):
'pos="0.0 2.0 0.35"', 'pos="0.0 2.0 0"'
)
)
c = dx.contact
c = dx._impl.contact
self.assertTrue((c.dist > 0).all())
@@ -837,8 +849,8 @@ class HFieldTest(absltest.TestCase):
# check that all geoms are colliding with the hfield
for geom_id in [1, 2, 3]:
mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
self.assertTrue((c.dist < 0).any())
self.assertTrue((c.dist > -1e-3).any())
# all contact normals are roughly pointing in the right direction
@@ -850,7 +862,7 @@ class HFieldTest(absltest.TestCase):
for p in positions:
xml = self._HFIELD.replace('<body pos="0 0', f'<body pos="{p}')
_, dx = _collide(xml)
self.assertTrue((dx.contact.dist >= 0).all())
self.assertTrue((dx._impl.contact.dist >= 0).all())
def test_hfield_deep(self):
"""Tests that objects with deep penetration do not get stuck."""
@@ -868,8 +880,8 @@ class HFieldTest(absltest.TestCase):
# check that all geoms are colliding with the hfield
for geom_id in [1, 2, 3]:
mask = (dx.contact.geom == np.array([0, geom_id])).all(axis=1)
c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx.contact)
mask = (dx._impl.contact.geom == np.array([0, geom_id])).all(axis=1)
c = jax.tree_util.tree_map(lambda x, m=mask: x[m], dx._impl.contact)
# all contact normals are in the top half-face of the hfield
self.assertTrue((c.frame[:, 0].dot(np.array([0, 0, 1])) > 0.7).all())
@@ -892,8 +904,8 @@ class BodyPairFilterTest(absltest.TestCase):
def test_filter_self_collision(self):
"""Tests that self collisions get filtered."""
d, dx = _collide(self._SELF_COLLISION)
self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx.contact.pos.shape[0], 0)
self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx._impl.contact.pos.shape[0], 0)
_PARENT_CHILD = """
<mujoco>
@@ -923,8 +935,8 @@ class BodyPairFilterTest(absltest.TestCase):
dx = kinematics_jit_fn(mx, dx)
dx = collision_jit_fn(mx, dx)
self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx.contact.pos.shape[0], 0)
self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx._impl.contact.pos.shape[0], 0)
def test_disable_filter_parent_child(self):
"""Tests that filterparent flag disables parent-child filtering."""
@@ -941,8 +953,8 @@ class BodyPairFilterTest(absltest.TestCase):
dx = collision_jit_fn(mx, dx)
# one collision between parent-child spheres
self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx.contact.pos.shape[0], 1)
self.assertEqual(dx._impl.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx._impl.contact.pos.shape[0], 1)
class DimTest(parameterized.TestCase):
@@ -1027,8 +1039,8 @@ class TopKContactTest(absltest.TestCase):
dx_all = collision_jit_fn(mx_all, dx)
dx_top_k = collision_jit_fn(mx_top_k, dx)
self.assertEqual(dx_all.contact.dist.shape, (3,))
self.assertEqual(dx_top_k.contact.dist.shape, (2,))
self.assertEqual(dx_all._impl.contact.dist.shape, (3,))
self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
_CAPSULES_MAX_PAIR = """
<mujoco>
@@ -1073,9 +1085,9 @@ class TopKContactTest(absltest.TestCase):
dx_all = collision_jit_fn(mx_all, dx)
dx_top_k = collision_jit_fn(mx_top_k, dx)
self.assertEqual(dx_all.contact.dist.shape, (6,))
self.assertEqual(dx_top_k.contact.dist.shape, (2,))
self.assertTrue((dx_top_k.contact.dist < 0).all())
self.assertEqual(dx_all._impl.contact.dist.shape, (6,))
self.assertEqual(dx_top_k._impl.contact.dist.shape, (2,))
self.assertTrue((dx_top_k._impl.contact.dist < 0).all())
if __name__ == '__main__':
mjx/mujoco/mjx/_src/collision_sdf.py
+2
View File
@@ -32,7 +32,9 @@ from mujoco.mjx._src.collision_types import Collision
from mujoco.mjx._src.collision_types import GeomInfo
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
# pylint: enable=g-importing-member
# the SDF function takes position in, and returns a distance or objective
mjx/mujoco/mjx/_src/constraint.py
+58 -26
View File
@@ -28,10 +28,12 @@ from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import ConstraintType
from mujoco.mjx._src.types import Contact
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import EqType
from mujoco.mjx._src.types import JointType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import ObjType
# pylint: enable=g-importing-member
import numpy as np
@@ -281,6 +283,10 @@ def _efc_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for tendon equality constraints."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError(
'_efc_equality_tendon requires JAX backend implementation.'
)
eq_id = np.nonzero(m.eq_type == EqType.TENDON)[0]
@@ -314,9 +320,9 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
return jax.tree_util.tree_map(lambda x: x * active, efc)
inv1, inv2 = m.tendon_invweight0[obj1id], m.tendon_invweight0[obj2id]
jac1, jac2 = d.ten_J[obj1id], d.ten_J[obj2id]
pos1 = d.ten_length[obj1id] - m.tendon_length0[obj1id]
pos2 = d.ten_length[obj2id] - m.tendon_length0[obj2id]
jac1, jac2 = d._impl.ten_J[obj1id], d._impl.ten_J[obj2id]
pos1 = d._impl.ten_length[obj1id] - m.tendon_length0[obj1id]
pos2 = d._impl.ten_length[obj2id] - m.tendon_length0[obj2id]
invweight = inv1 + inv2 * (obj2id > -1)
return rows(
@@ -326,8 +332,11 @@ def _efc_equality_tendon(m: Model, d: Data) -> Optional[_Efc]:
def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for dof frictionloss."""
dof_id = np.nonzero(m.dof_hasfrictionloss)[0]
tendon_id = np.nonzero(m.tendon_hasfrictionloss)[0]
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('_efc_friction requires JAX backend implementation.')
dof_id = np.nonzero(m._impl.dof_hasfrictionloss)[0]
tendon_id = np.nonzero(m._impl.tendon_hasfrictionloss)[0]
size = dof_id.size + tendon_id.size
if (m.opt.disableflags & DisableBit.FRICTIONLOSS) or (size == 0):
@@ -337,7 +346,7 @@ def _efc_friction(m: Model, d: Data) -> Optional[_Efc]:
args_dof += (m.dof_solimp,)
args_dof = jax.tree_util.tree_map(lambda x: x[dof_id], args_dof)
args_ten = (d.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
args_ten = (d._impl.ten_J, m.tendon_frictionloss, m.tendon_invweight0)
args_ten += (m.tendon_solref_fri, m.tendon_solimp_fri)
args_ten = jax.tree_util.tree_map(lambda x: x[tendon_id], args_ten)
@@ -415,6 +424,9 @@ def _efc_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for tendon limits."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('_efc_limit_tendon requires JAX backend implementation.')
tendon_id = np.nonzero(m.tendon_limited)[0]
if (m.opt.disableflags & DisableBit.LIMIT) or tendon_id.size == 0:
@@ -423,8 +435,8 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
length, j, range_, margin, invweight, solref, solimp = jax.tree_util.tree_map(
lambda x: x[tendon_id],
(
d.ten_length,
d.ten_J,
d._impl.ten_length,
d._impl.ten_J,
m.tendon_range,
m.tendon_margin,
m.tendon_invweight0,
@@ -446,8 +458,12 @@ def _efc_limit_tendon(m: Model, d: Data) -> Optional[_Efc]:
def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for frictionless contacts."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError(
'_efc_contact_frictionless requires JAX backend implementation.'
)
con_id = np.nonzero(d.contact.dim == 1)[0]
con_id = np.nonzero(d._impl.contact.dim == 1)[0]
if con_id.size == 0:
return None
@@ -473,15 +489,19 @@ def _efc_contact_frictionless(m: Model, d: Data) -> Optional[_Efc]:
jp.zeros_like(pos),
)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
return rows(contact)
def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
"""Calculates constraint rows for frictional pyramidal contacts."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError(
'_efc_contact_pyramidal requires JAX backend implementation.'
)
con_id = np.nonzero(d.contact.dim == condim)[0]
con_id = np.nonzero(d._impl.contact.dim == condim)[0]
if con_id.size == 0:
return None
@@ -518,15 +538,19 @@ def _efc_contact_pyramidal(m: Model, d: Data, condim: int) -> Optional[_Efc]:
jp.zeros_like(pos),
)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
# concatenate to drop row grouping
return jax.tree_util.tree_map(jp.concatenate, rows(contact))
def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
"""Calculates constraint rows for frictional elliptic contacts."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError(
'_efc_contact_elliptic requires JAX backend implementation.'
)
con_id = np.nonzero(d.contact.dim == condim)[0]
con_id = np.nonzero(d._impl.contact.dim == condim)[0]
if con_id.size == 0:
return None
@@ -563,7 +587,7 @@ def _efc_contact_elliptic(m: Model, d: Data, condim: int) -> Optional[_Efc]:
jp.zeros_like(pos),
)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d.contact)
contact = jax.tree_util.tree_map(lambda x: x[con_id], d._impl.contact)
# concatenate to drop row grouping
return jax.tree_util.tree_map(jp.concatenate, rows(contact))
@@ -602,14 +626,14 @@ def make_efc_type(
if not m.opt.disableflags & DisableBit.FRICTIONLOSS:
nf_dof = (
m.dof_hasfrictionloss.sum()
if isinstance(m, Model)
m._impl.dof_hasfrictionloss.sum()
if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
else (m.dof_frictionloss > 0).sum()
)
efc_types += [ConstraintType.FRICTION_DOF] * nf_dof
nf_tendon = (
m.tendon_hasfrictionloss.sum()
if isinstance(m, Model)
m._impl.tendon_hasfrictionloss.sum()
if isinstance(m, Model) and isinstance(m._impl, ModelJAX)
else (m.tendon_frictionloss > 0).sum()
)
efc_types += [ConstraintType.FRICTION_TENDON] * nf_tendon
@@ -683,10 +707,14 @@ def make_constraint(m: Model, d: Data) -> Data:
if not efcs:
z = jp.empty(0)
d = d.replace(efc_J=jp.empty((0, m.nv)))
d = d.replace(
efc_D=z, efc_aref=z, efc_frictionloss=z, efc_pos=z, efc_margin=z
)
d = d.tree_replace({'_impl.efc_J': jp.empty((0, m.nv))})
d = d.tree_replace({
'_impl.efc_D': z,
'_impl.efc_aref': z,
'_impl.efc_frictionloss': z,
'_impl.efc_pos': z,
'_impl.efc_margin': z,
})
return d
efc = jax.tree_util.tree_map(lambda *x: jp.concatenate(x), *efcs)
@@ -699,9 +727,13 @@ def make_constraint(m: Model, d: Data) -> Data:
return aref, r, efc.pos_aref + efc.margin, efc.margin, efc.frictionloss
aref, r, pos, margin, frictionloss = fn(efc)
d = d.replace(
efc_J=efc.J, efc_D=1 / r, efc_aref=aref, efc_pos=pos, efc_margin=margin
)
d = d.replace(efc_frictionloss=frictionloss)
d = d.tree_replace({
'_impl.efc_J': efc.J,
'_impl.efc_D': 1 / r,
'_impl.efc_aref': aref,
'_impl.efc_pos': pos,
'_impl.efc_margin': margin,
})
d = d.tree_replace({'_impl.efc_frictionloss': frictionloss})
return d
mjx/mujoco/mjx/_src/constraint_test.py
+18 -14
View File
@@ -69,16 +69,16 @@ class ConstraintTest(parameterized.TestCase):
order = test_util.efc_order(m, d, dx)
d_efc_j = d.efc_J.reshape((-1, m.nv))
_assert_eq(d_efc_j, dx.efc_J[order][: d.nefc], 'efc_J')
_assert_eq(0, dx.efc_J[order][d.nefc :], 'efc_J')
_assert_eq(d.efc_aref, dx.efc_aref[order][: d.nefc], 'efc_aref')
_assert_eq(0, dx.efc_aref[order][d.nefc :], 'efc_aref')
_assert_eq(d.efc_D, dx.efc_D[order][: d.nefc], 'efc_D')
_assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
_assert_eq(dx.efc_pos[order][d.nefc :], 0, 'efc_pos')
_assert_eq(d_efc_j, dx._impl.efc_J[order][: d.nefc], 'efc_J')
_assert_eq(0, dx._impl.efc_J[order][d.nefc :], 'efc_J')
_assert_eq(d.efc_aref, dx._impl.efc_aref[order][: d.nefc], 'efc_aref')
_assert_eq(0, dx._impl.efc_aref[order][d.nefc :], 'efc_aref')
_assert_eq(d.efc_D, dx._impl.efc_D[order][: d.nefc], 'efc_D')
_assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
_assert_eq(dx._impl.efc_pos[order][d.nefc :], 0, 'efc_pos')
_assert_eq(
d.efc_frictionloss,
dx.efc_frictionloss[order][: d.nefc],
dx._impl.efc_frictionloss[order][: d.nefc],
'efc_frictionloss',
)
@@ -104,7 +104,7 @@ class ConstraintTest(parameterized.TestCase):
self.assertEqual(nl, 0)
self.assertEqual(nc, 0)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 0)
self.assertEqual(dx._impl.efc_J.shape[0], 0)
def test_disable_equality(self):
m = test_util.load_test_file('constraints.xml')
@@ -115,7 +115,9 @@ class ConstraintTest(parameterized.TestCase):
self.assertEqual(nl, 5)
self.assertEqual(nc, 180)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 187) # only joint/tendon limit, contact
self.assertEqual(
dx._impl.efc_J.shape[0], 187
) # only joint/tendon limit, contact
def test_disable_contact(self):
m = test_util.load_test_file('constraints.xml')
@@ -126,7 +128,9 @@ class ConstraintTest(parameterized.TestCase):
self.assertEqual(nl, 5)
self.assertEqual(nc, 0)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 27) # only equality, joint/tendon limit
self.assertEqual(
dx._impl.efc_J.shape[0], 27
) # only equality, joint/tendon limit
def test_disable_frictionloss(self):
m = test_util.load_test_file('constraints.xml')
@@ -137,7 +141,7 @@ class ConstraintTest(parameterized.TestCase):
self.assertEqual(nl, 5)
self.assertEqual(nc, 180)
dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
self.assertEqual(dx.efc_J.shape[0], 205)
self.assertEqual(dx._impl.efc_J.shape[0], 205)
def test_margin(self):
"""Test margin."""
@@ -163,8 +167,8 @@ class ConstraintTest(parameterized.TestCase):
dx = mjx.make_constraint(mx, dx)
order = test_util.efc_order(m, d, dx)
_assert_eq(d.efc_pos, dx.efc_pos[order][: d.nefc], 'efc_pos')
_assert_eq(d.efc_margin, dx.efc_margin[order][: d.nefc], 'efc_margin')
_assert_eq(d.efc_pos, dx._impl.efc_pos[order][: d.nefc], 'efc_pos')
_assert_eq(d.efc_margin, dx._impl.efc_margin[order][: d.nefc], 'efc_margin')
if __name__ == '__main__':
mjx/mujoco/mjx/_src/dataclasses.py
+1
View File
@@ -18,6 +18,7 @@ import copy
import dataclasses
import typing
from typing import Dict, Optional, Sequence, Tuple, TypeVar, Union
import warnings
import jax
import numpy as np
mjx/mujoco/mjx/_src/derivative.py
+61
View File
@@ -0,0 +1,61 @@
# Copyright 2025 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.
# ==============================================================================
"""Derivative functions."""
from typing import Optional
import jax
from jax import numpy as jp
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import BiasType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import DynType
from mujoco.mjx._src.types import GainType
from mujoco.mjx._src.types import Model
def deriv_smooth_vel(m: Model, d: Data) -> Optional[jax.Array]:
"""Analytical derivative of smooth forces w.r.t velocities."""
qderiv = None
# qDeriv += d qfrc_actuator / d qvel
if not m.opt.disableflags & DisableBit.ACTUATION:
affine_bias = m.actuator_biastype == BiasType.AFFINE
bias_vel = m.actuator_biasprm[:, 2] * affine_bias
affine_gain = m.actuator_gaintype == GainType.AFFINE
gain_vel = m.actuator_gainprm[:, 2] * affine_gain
ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
vel = bias_vel + gain_vel * ctrl
qderiv = d._impl.actuator_moment.T @ jax.vmap(jp.multiply)(
d._impl.actuator_moment, vel
)
# qDeriv += d qfrc_passive / d qvel
if not m.opt.disableflags & DisableBit.PASSIVE:
if qderiv is None:
qderiv = -jp.diag(m.dof_damping)
else:
qderiv -= jp.diag(m.dof_damping)
if m.ntendon:
qderiv -= d._impl.ten_J.T @ jp.diag(m.tendon_damping) @ d._impl.ten_J
# TODO(robotics-simulation): fluid drag model
if m.opt.has_fluid_params: # pytype: disable=attribute-error
raise NotImplementedError('fluid drag not supported for implicitfast')
# TODO(team): rne derivative
return qderiv
mjx/mujoco/mjx/_src/forward.py
+63 -40
View File
@@ -22,6 +22,7 @@ from jax import numpy as jp
import mujoco
from mujoco.mjx._src import collision_driver
from mujoco.mjx._src import constraint
from mujoco.mjx._src import derivative
from mujoco.mjx._src import math
from mujoco.mjx._src import passive
from mujoco.mjx._src import scan
@@ -32,12 +33,15 @@ from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import BiasType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import DynType
from mujoco.mjx._src.types import GainType
from mujoco.mjx._src.types import IntegratorType
from mujoco.mjx._src.types import JointType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import TrnType
# pylint: enable=g-importing-member
import numpy as np
@@ -79,10 +83,13 @@ def fwd_position(m: Model, d: Data) -> Data:
@named_scope
def fwd_velocity(m: Model, d: Data) -> Data:
"""Velocity-dependent computations."""
d = d.replace(
actuator_velocity=d.actuator_moment @ d.qvel,
ten_velocity=d.ten_J @ d.qvel,
)
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('fwd_velocity requires JAX backend implementation.')
d = d.tree_replace({
'_impl.actuator_velocity': d._impl.actuator_moment @ d.qvel,
'_impl.ten_velocity': d._impl.ten_J @ d.qvel,
})
d = smooth.com_vel(m, d)
d = passive.passive(m, d)
d = smooth.rne(m, d)
@@ -171,13 +178,41 @@ def fwd_actuation(m: Model, d: Data) -> Data:
m.actuator_gainprm,
m.actuator_biastype,
m.actuator_biasprm,
d.actuator_length,
d.actuator_velocity,
d._impl.actuator_length,
d._impl.actuator_velocity,
ctrl_act,
jp.array(m.actuator_lengthrange),
jp.array(m.actuator_acc0),
group_by='u',
)
# tendon total force clamping
if np.any(m.tendon_actfrclimited):
(tendon_actfrclimited_id,) = np.nonzero(m.tendon_actfrclimited)
actuator_tendon = m.actuator_trntype == TrnType.TENDON
force_mask = [
actuator_tendon & (m.actuator_trnid[:, 0] == tendon_id)
for tendon_id in tendon_actfrclimited_id
]
force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask])
force_mat = np.array(force_mask)[:, force_ids]
tendon_total_force = force_mat @ force[force_ids]
force_scaling = jp.where(
tendon_total_force < m.tendon_actfrcrange[tendon_actfrclimited_id, 0],
m.tendon_actfrcrange[tendon_actfrclimited_id, 0] / tendon_total_force,
1,
)
force_scaling = jp.where(
tendon_total_force > m.tendon_actfrcrange[tendon_actfrclimited_id, 1],
m.tendon_actfrcrange[tendon_actfrclimited_id, 1] / tendon_total_force,
force_scaling,
)
tendon_forces = force[force_ids] * (force_mat.T @ force_scaling)
force = force.at[force_ids].set(tendon_forces)
forcerange = jp.where(
m.actuator_forcelimited[:, None],
m.actuator_forcerange,
@@ -185,7 +220,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
)
force = jp.clip(force, forcerange[:, 0], forcerange[:, 1])
qfrc_actuator = d.actuator_moment.T @ force
qfrc_actuator = d._impl.actuator_moment.T @ force
if m.ngravcomp:
# actuator-level gravity compensation, skip if added as passive force
@@ -301,13 +336,17 @@ def _advance(
@named_scope
def euler(m: Model, d: Data) -> Data:
"""Euler integrator, semi-implicit in velocity."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('euler requires JAX backend implementation.')
# integrate damping implicitly
qacc = d.qacc
if not m.opt.disableflags & DisableBit.EULERDAMP:
if support.is_sparse(m):
dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
else:
dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping))
qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
dh = d.tree_replace({'_impl.qM': qM})
dh = smooth.factor_m(m, dh)
qfrc = d.qfrc_smooth + d.qfrc_constraint
qacc = smooth.solve_m(m, dh, qfrc)
@@ -317,7 +356,7 @@ def euler(m: Model, d: Data) -> Data:
@named_scope
def rungekutta4(m: Model, d: Data) -> Data:
"""Runge-Kutta explicit order 4 integrator."""
d_t0 = d
d0 = d
# pylint: disable=invalid-name
A, B = _RK4_A, _RK4_B
C = jp.tril(A).sum(axis=0) # C(i) = sum_j A(i,j)
@@ -338,9 +377,9 @@ def rungekutta4(m: Model, d: Data) -> Data:
lambda k: a * k, (kqvel, d.qacc, d.act_dot)
)
# get intermediate RK solutions
kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel)
kact = d_t0.act + dact_dot * m.opt.timestep
kqvel = d_t0.qvel + dqacc * m.opt.timestep
kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d0.qpos, dqvel)
kact = d0.act + dact_dot * m.opt.timestep
kqvel = d0.qvel + dqacc * m.opt.timestep
d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t)
d = forward(m, d)
@@ -352,44 +391,25 @@ def rungekutta4(m: Model, d: Data) -> Data:
abt = jp.vstack([jp.diag(A), B[1:4], T]).T
out, _ = jax.lax.scan(f, (qvel, qacc, act_dot, kqvel, d), abt, unroll=3)
qvel, qacc, act_dot, *_ = out
qvel, qacc, act_dot, _, d1 = out
d = _advance(m, d_t0, act_dot, qacc, qvel)
d = d1.replace(qpos=d0.qpos, qvel=d0.qvel, act=d0.act, time=d0.time)
d = _advance(m, d, act_dot, qacc, qvel)
return d
@named_scope
def implicit(m: Model, d: Data) -> Data:
"""Integrates fully implicit in velocity."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('implicit requires JAX backend implementation.')
qderiv = None
# qDeriv += d qfrc_actuator / d qvel
if not m.opt.disableflags & DisableBit.ACTUATION:
affine_bias = m.actuator_biastype == BiasType.AFFINE
bias_vel = m.actuator_biasprm[:, 2] * affine_bias
affine_gain = m.actuator_gaintype == GainType.AFFINE
gain_vel = m.actuator_gainprm[:, 2] * affine_gain
ctrl = d.ctrl.at[m.actuator_dyntype != DynType.NONE].set(d.act)
vel = bias_vel + gain_vel * ctrl
qderiv = d.actuator_moment.T @ jp.diag(vel) @ d.actuator_moment
# qDeriv += d qfrc_passive / d qvel
if not m.opt.disableflags & DisableBit.PASSIVE:
if qderiv is None:
qderiv = -jp.diag(m.dof_damping)
else:
qderiv -= jp.diag(m.dof_damping)
if m.ntendon:
qderiv -= d.ten_J.T @ jp.diag(m.tendon_damping) @ d.ten_J
# TODO(robotics-simulation): fluid drag model
if m.opt.has_fluid_params:
raise NotImplementedError('fluid drag not supported for implicitfast')
qderiv = derivative.deriv_smooth_vel(m, d)
qacc = d.qacc
if qderiv is not None:
# TODO(robotics-simulation): use smooth.factor_m / solve_m here:
qm = support.full_m(m, d) if support.is_sparse(m) else d.qM
qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
qm -= m.opt.timestep * qderiv
qh, _ = jax.scipy.linalg.cho_factor(qm)
qfrc = d.qfrc_smooth + d.qfrc_constraint
@@ -401,6 +421,9 @@ def implicit(m: Model, d: Data) -> Data:
@named_scope
def forward(m: Model, d: Data) -> Data:
"""Forward dynamics."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('forward requires JAX backend implementation.')
d = fwd_position(m, d)
d = sensor.sensor_pos(m, d)
d = fwd_velocity(m, d)
@@ -409,7 +432,7 @@ def forward(m: Model, d: Data) -> Data:
d = fwd_acceleration(m, d)
d = sensor.sensor_acc(m, d)
if d.efc_J.size == 0:
if d._impl.efc_J.size == 0:
d = d.replace(qacc=d.qacc_smooth)
return d
mjx/mujoco/mjx/_src/forward_test.py
+23
View File
@@ -17,6 +17,7 @@
from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
@@ -130,6 +131,7 @@ class ForwardTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'qpos')
_assert_attr_eq(d, dx, 'act')
_assert_attr_eq(d, dx, 'time')
_assert_attr_eq(d, dx, 'xpos')
def test_eulerdamp(self):
m = test_util.load_test_file('pendula.xml')
@@ -195,6 +197,27 @@ class ActuatorTest(parameterized.TestCase):
dx = jax.jit(mjx.euler)(mx, dx)
_assert_attr_eq(d, dx, 'act')
def test_tendon_force_clamp(self):
m = test_util.load_test_file('actuator/tendon_force_clamp.xml')
d = mujoco.MjData(m)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
dx = dx.replace(ctrl=jp.array([1.0, 1.0, 1.0, -4.0, 1.0, -20.0, 5.0, -5.0]))
dx = mjx.forward(mx, dx)
_assert_eq(
dx.actuator_force,
jp.array([1.0, 1.0, 1.0, -4.0 / 3.0, 1.0 / 3.0, -10.0, 5.0, -5.0]),
'actuator_force',
)
_assert_eq(
dx.sensordata,
jp.array([3.0, -1.0, -10.0, 0.0]),
'sensordata',
)
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/inverse.py
+106
View File
@@ -0,0 +1,106 @@
# Copyright 2025 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.
# ==============================================================================
"""Inverse dynamics functions."""
from jax import numpy as jp
from mujoco.mjx._src import derivative
from mujoco.mjx._src import forward
from mujoco.mjx._src import sensor
from mujoco.mjx._src import smooth
from mujoco.mjx._src import solver
from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import EnableBit
from mujoco.mjx._src.types import IntegratorType
from mujoco.mjx._src.types import Model
def discrete_acc(m: Model, d: Data) -> Data:
"""Convert discrete-time qacc to continuous-time qacc."""
if m.opt.integrator == IntegratorType.RK4:
raise RuntimeError(
'discrete inverse dynamics is not supported by RK4 integrator'
)
elif m.opt.integrator == IntegratorType.EULER:
dsbl_eulerdamp = m.opt.disableflags & DisableBit.EULERDAMP
no_dof_damping = (m.dof_damping == 0).all()
if dsbl_eulerdamp or no_dof_damping:
return d
# set qfrc = (M + h*diag(B)) * qacc
qfrc = support.mul_m(m, d, d.qacc)
qfrc += m.opt.timestep * m.dof_damping * d.qacc
elif m.opt.integrator == IntegratorType.IMPLICITFAST:
qm = support.full_m(m, d)
# compute analytical derivative qDeriv; skip rne derivative
qderiv = derivative.deriv_smooth_vel(m, d)
if qderiv is not None:
# M = M - dt*qDeriv
qm -= m.opt.timestep * qderiv
# set qfrc = (M - dt*qDeriv) * qacc
qfrc = qm @ d.qacc
else:
raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.')
# solve for qacc: qfrc = M * qacc
qacc = smooth.solve_m(m, d, qfrc)
return d.replace(qacc=qacc)
def inv_constraint(m: Model, d: Data) -> Data:
"""Inverse constraint solver."""
# no constraints
if d._impl.efc_J.size == 0: # pytype: disable=attribute-error
return d.replace(qfrc_constraint=jp.zeros(m.nv))
# update
ctx = solver.Context.create(m, d, grad=False)
return d.tree_replace({
'qfrc_constraint': ctx.qfrc_constraint,
'_impl.efc_force': ctx.efc_force,
})
def inverse(m: Model, d: Data) -> Data:
"""Inverse dynamics."""
d = forward.fwd_position(m, d)
d = sensor.sensor_pos(m, d)
d = forward.fwd_velocity(m, d)
d = sensor.sensor_vel(m, d)
qacc = d.qacc
if m.opt.enableflags & EnableBit.INVDISCRETE:
d = discrete_acc(m, d)
d = inv_constraint(m, d)
d = smooth.rne(m, d, flg_acc=True)
d = sensor.sensor_acc(m, d)
qfrc_inverse = (
d.qfrc_bias + m.dof_armature * d.qacc - d.qfrc_passive - d.qfrc_constraint
)
if m.opt.enableflags & EnableBit.INVDISCRETE:
return d.replace(qfrc_inverse=qfrc_inverse, qacc=qacc)
else:
return d.replace(qfrc_inverse=qfrc_inverse)
mjx/mujoco/mjx/_src/inverse_test.py
+114
View File
@@ -0,0 +1,114 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for inverse dynamics functions."""
from absl.testing import absltest
from absl.testing import parameterized
from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import support
import numpy as np
# tolerance for difference between MuJoCo and MJX calculations - mostly
# due to float precision
_TOLERANCE = 1e-5
def _assert_eq(a, b, name, tol=_TOLERANCE):
tol = tol * 10 # avoid test noise
err_msg = f'mismatch: {name}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
class InverseTest(parameterized.TestCase):
@parameterized.parameters(
(mujoco.mjtIntegrator.mjINT_EULER, False, False),
(mujoco.mjtIntegrator.mjINT_EULER, False, True),
(mujoco.mjtIntegrator.mjINT_EULER, True, False),
(mujoco.mjtIntegrator.mjINT_EULER, True, True),
(mujoco.mjtIntegrator.mjINT_IMPLICITFAST, False, False),
(mujoco.mjtIntegrator.mjINT_IMPLICITFAST, True, False),
)
def test_forward_inverse_match(self, integrator, invdiscrete, eulerdamp):
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<option timestep=".005" gravity="-1 -1 -10"/>
<worldbody>
<geom type="plane" size="10 10 .001"/>
<body pos="0 0 1">
<geom type="sphere" size=".1" pos=".1 .2 .3"/>
<joint name="jnt1" type="hinge" axis="0 1 0" stiffness=".25" damping=".125"/>
<body pos="0 0 1">
<geom type="sphere" size=".1" pos=".1 .2 .3"/>
<joint name="jnt2" type="hinge" axis="0 1 0" stiffness=".6" damping=".3"/>
</body>
</body>
</worldbody>
<actuator>
<motor joint="jnt1"/>
</actuator>
<equality>
<joint joint1="jnt1" joint2="jnt2"/>
</equality>
</mujoco>
""")
m.opt.integrator = integrator
if invdiscrete:
m.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_INVDISCRETE
if not eulerdamp:
m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_EULERDAMP
d = mujoco.MjData(m)
d.qvel = np.random.uniform(low=-0.01, high=0.01, size=d.qvel.shape)
d.ctrl = np.random.uniform(low=-0.01, high=0.01, size=d.ctrl.shape)
d.qfrc_applied = np.random.uniform(
low=-0.01, high=0.01, size=d.qfrc_applied.shape
)
d.xfrc_applied = np.random.uniform(
low=-0.01, high=0.01, size=d.xfrc_applied.shape
)
mujoco.mj_step(m, d, 100)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d)
dx_next = mjx.step(mx, dx)
qacc_fd = (dx_next.qvel - dx.qvel) / mx.opt.timestep
dx = mjx.forward(mx, dx)
if invdiscrete:
dx = dx.replace(qacc=qacc_fd)
dxinv = mjx.inverse(mx, dx)
fwdinv0 = jp.linalg.norm(
dxinv.qfrc_constraint - dx.qfrc_constraint, ord=np.inf
)
fwdinv1 = jp.linalg.norm(
dxinv.qfrc_inverse
- (
dx.qfrc_applied + dx.qfrc_actuator + support.xfrc_accumulate(mx, dx)
),
ord=np.inf,
)
self.assertLess(fwdinv0, 1.0e-3)
self.assertLess(fwdinv1, 1.0e-3)
_assert_eq(dxinv.qacc, dx.qacc, 'qacc')
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/io.py
+970 -427
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/_src/io_test.py
+518 -100
View File
@@ -14,14 +14,19 @@
# ==============================================================================
"""Tests for io functions."""
import os
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import io as mjx_io
from mujoco.mjx._src import test_util
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import BackendImpl
from mujoco.mjx._src.types import ConeType
# pylint: enable=g-importing-member
import numpy as np
@@ -91,14 +96,29 @@ _MULTIPLE_CONSTRAINTS = """
</mujoco>
"""
_SIMPLE_BODY = """
<mujoco>
<worldbody>
<!-- nC < nM because this body has an inertia-aligned free joint -->
<body name="simplebody">
<freejoint/>
<geom type="sphere" size="0.01"/>
</body>
</worldbody>
</mujoco>
"""
class ModelIOTest(parameterized.TestCase):
"""IO tests for mjx.Model."""
@parameterized.parameters(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS)
def test_put_model(self, xml):
@parameterized.product(
xml=(_MULTIPLE_CONVEX_OBJECTS, _MULTIPLE_CONSTRAINTS),
backend_impl=('jax', 'c'),
)
def test_put_model(self, xml, backend_impl):
m = mujoco.MjModel.from_xml_string(xml)
mx = mjx.put_model(m)
mx = mjx.put_model(m, backend_impl=backend_impl)
def assert_not_weak_type(x):
if isinstance(x, jax.Array):
@@ -120,8 +140,14 @@ class ModelIOTest(parameterized.TestCase):
self.assertEqual(mx.nM, m.nM)
self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
# fields restricted to MuJoCo should not be populated
self.assertEqual(mx.bvh_aabb.shape, (0,))
if backend_impl == 'jax':
# fields restricted to MuJoCo should not be populated
self.assertFalse(hasattr(mx, 'bvh_aabb'))
elif backend_impl == 'c':
# Options specific to C are populated.
self.assertEqual(mx.opt.apirate, m.opt.apirate)
# Fields private to C backend impl are populated.
self.assertTrue(hasattr(mx._impl, 'bvh_aabb'))
np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
np.testing.assert_allclose(mx.geom_type, m.geom_type)
@@ -150,7 +176,8 @@ class ModelIOTest(parameterized.TestCase):
m = mjx.put_model(
mujoco.MjModel.from_xml_string(
'<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
)
),
backend_impl='jax',
)
self.assertTrue(m.opt.has_fluid_params)
@@ -174,7 +201,8 @@ class ModelIOTest(parameterized.TestCase):
def test_margin_gap_mesh_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
mjx.put_model(
mujoco.MjModel.from_xml_string("""
<mujoco>
<asset>
<mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 1"/>
@@ -189,27 +217,32 @@ class ModelIOTest(parameterized.TestCase):
<geom size="0.05"/>
</body>
</worldbody>
</mujoco>"""))
</mujoco>"""),
backend_impl='jax',
)
def test_implicitfast_fluid_not_implemented(self):
with self.assertRaises(NotImplementedError):
mjx.put_model(mujoco.MjModel.from_xml_string("""
mjx.put_model(
mujoco.MjModel.from_xml_string("""
<mujoco>
<option viscosity="3.0" integrator="implicitfast"/>
<worldbody/>
</mujoco>"""))
</mujoco>"""),
backend_impl='jax',
)
def test_wrap_inside(self):
m = test_util.load_test_file('tendon/wrap_sidesite.xml')
mx0 = mjx.put_model(m)
mx0 = mjx.put_model(m, backend_impl='jax')
np.testing.assert_equal(
mx0.is_wrap_inside,
mx0._impl.is_wrap_inside,
np.array([1, 0, 1, 0, 1, 1, 0]),
)
m.site_pos[2] = m.site_pos[1]
mx1 = mjx.put_model(m)
mx1 = mjx.put_model(m, backend_impl='jax')
np.testing.assert_equal(
mx1.is_wrap_inside,
mx1._impl.is_wrap_inside,
np.array([0, 0, 1, 0, 1, 0, 0]),
)
@@ -217,11 +250,11 @@ class ModelIOTest(parameterized.TestCase):
class DataIOTest(parameterized.TestCase):
"""IO tests for mjx.Data."""
def test_make_data(self):
@parameterized.parameters('jax', 'c')
def test_make_data(self, backend_impl: str):
"""Test that make_data returns the correct shapes."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
d = mjx.make_data(m)
d = mjx.make_data(m, backend_impl=backend_impl)
nq = 22
nbody = 5
@@ -230,7 +263,7 @@ class DataIOTest(parameterized.TestCase):
nv = 19
nefc = 185
self.assertEqual(d.nefc, nefc)
self.assertEqual(d._impl.nefc, nefc)
self.assertEqual(d.qpos.shape, (nq,))
self.assertEqual(d.qvel.shape, (nv,))
self.assertEqual(d.act.shape, (0,))
@@ -251,57 +284,68 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d.geom_xpos.shape, (6, 3))
self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
self.assertEqual(d.subtree_com.shape, (nbody, 3))
self.assertEqual(d.cdof.shape, (nv, 6))
self.assertEqual(d.cinert.shape, (nbody, 10))
self.assertEqual(d.crb.shape, (nbody, 10))
self.assertEqual(d.actuator_length.shape, (1,))
self.assertEqual(d.actuator_moment.shape, (1, nv))
self.assertEqual(d.qM.shape, (nv, nv))
self.assertEqual(d.qLD.shape, (nv, nv))
self.assertEqual(d.qLDiagInv.shape, (0,))
self.assertEqual(d.contact.dist.shape, (ncon,))
self.assertEqual(d.contact.pos.shape, (ncon, 3))
self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
self.assertEqual(d.contact.solref.shape, (ncon, 2))
self.assertEqual(d.contact.solimp.shape, (ncon, 5))
self.assertEqual(d.contact.geom1.shape, (ncon,))
self.assertEqual(d.contact.geom2.shape, (ncon,))
self.assertEqual(d.efc_J.shape, (nefc, nv))
self.assertEqual(d.efc_frictionloss.shape, (nefc,))
self.assertEqual(d.efc_D.shape, (nefc,))
self.assertEqual(d.actuator_velocity.shape, (1,))
self.assertEqual(d._impl.cdof.shape, (nv, 6))
self.assertEqual(d._impl.cinert.shape, (nbody, 10))
self.assertEqual(d._impl.crb.shape, (nbody, 10))
self.assertEqual(d._impl.actuator_length.shape, (1,))
self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
self.assertEqual(d._impl.contact.dist.shape, (ncon,))
self.assertEqual(d._impl.contact.pos.shape, (ncon, 3))
self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3))
self.assertEqual(d._impl.contact.solref.shape, (ncon, 2))
self.assertEqual(d._impl.contact.solimp.shape, (ncon, 5))
self.assertEqual(d._impl.contact.geom1.shape, (ncon,))
self.assertEqual(d._impl.contact.geom2.shape, (ncon,))
self.assertEqual(d._impl.efc_J.shape, (nefc, nv))
self.assertEqual(d._impl.efc_frictionloss.shape, (nefc,))
self.assertEqual(d._impl.efc_D.shape, (nefc,))
self.assertEqual(d._impl.actuator_velocity.shape, (1,))
self.assertEqual(d.cvel.shape, (nbody, 6))
self.assertEqual(d.cdof_dot.shape, (nv, 6))
self.assertEqual(d._impl.cdof_dot.shape, (nv, 6))
self.assertEqual(d.qfrc_bias.shape, (nv,))
self.assertEqual(d.qfrc_passive.shape, (nv,))
self.assertEqual(d.efc_aref.shape, (nefc,))
self.assertEqual(d._impl.efc_aref.shape, (nefc,))
self.assertEqual(d.qfrc_actuator.shape, (nv,))
self.assertEqual(d.qfrc_smooth.shape, (nv,))
self.assertEqual(d.qacc_smooth.shape, (nv,))
self.assertEqual(d.qfrc_constraint.shape, (nv,))
self.assertEqual(d.qfrc_inverse.shape, (nv,))
self.assertEqual(d.efc_force.shape, (nefc,))
self.assertEqual(d._impl.efc_force.shape, (nefc,))
if backend_impl == 'jax':
self.assertEqual(d._impl.qM.shape, (nv, nv))
self.assertEqual(d._impl.qLD.shape, (nv, nv))
self.assertEqual(d._impl.qLDiagInv.shape, (0,))
elif backend_impl == 'c':
self.assertEqual(d._impl.qM.shape, (nm,))
self.assertEqual(d._impl.qLD.shape, (nm,))
self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
# test sparse
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mjx.make_data(m)
self.assertEqual(d.qM.shape, (nm,))
self.assertEqual(d.qLD.shape, (nm,))
self.assertEqual(d.qLDiagInv.shape, (nv,))
d = mjx.make_data(m, backend_impl=backend_impl)
self.assertEqual(d._impl.qM.shape, (nm,))
self.assertEqual(d._impl.qLD.shape, (nm,))
self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
def test_put_data(self):
if backend_impl == 'c':
# check C specific fields
self.assertEqual(d._impl.light_xpos.shape, (m.nlight, 3))
self.assertEqual(d._impl.bvh_active.shape, (m.nbvh,))
@parameterized.parameters('jax', 'c')
def test_put_data(self, backend_impl: str):
"""Test that put_data puts the correct data for dense and sparse."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
# check a few fields
np.testing.assert_allclose(dx.qpos, d.qpos)
np.testing.assert_allclose(dx.xpos, d.xpos)
np.testing.assert_allclose(dx.cvel, d.cvel)
np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot)
# check that there are no weak types
self.assertFalse(
@@ -312,21 +356,26 @@ class DataIOTest(parameterized.TestCase):
)
)
# check that qM is transformed properly
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
if backend_impl == 'jax':
# check that qM is transformed properly
qm = np.zeros((m.nv, m.nv), dtype=np.float64)
mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
elif backend_impl == 'c':
np.testing.assert_allclose(dx._impl.qM, d.qM)
np.testing.assert_allclose(dx._impl.qLD, d.qLD)
np.testing.assert_allclose(dx._impl.qLDiagInv, d.qLDiagInv)
# 4 contacts, 2 for each capsule against the plane
self.assertEqual(dx.contact.dist.shape, (4,))
self.assertEqual(dx._impl.contact.dist.shape, (4,))
self.assertEqual(d.ncon, 1) # however only 1 contact in this step
np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
self.assertTrue((dx.contact.dist[1:] > 0).all())
self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
np.testing.assert_allclose(dx._impl.contact.dist[0], d.contact.dist[0])
self.assertTrue((dx._impl.contact.dist[1:] > 0).all())
self.assertEqual(dx._impl.contact.frame.shape, (4, 3, 3))
np.testing.assert_allclose(
dx.contact.frame[0].reshape(9), d.contact.frame[0]
dx._impl.contact.frame[0].reshape(9), d.contact.frame[0]
)
np.testing.assert_allclose(dx.contact.frame[1:], 0)
np.testing.assert_allclose(dx._impl.contact.frame[1:], 0)
# xmat, ximat, geom_xmat are all shape transformed
self.assertEqual(dx.xmat.shape, (3, 3, 3))
@@ -339,57 +388,68 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
# tendon data is correct
np.testing.assert_allclose(dx.ten_length, d.ten_length)
np.testing.assert_equal(dx.ten_wrapadr, np.zeros((1,)))
np.testing.assert_equal(dx.ten_wrapnum, np.zeros((1,)))
np.testing.assert_equal(dx.wrap_obj, np.zeros((2, 2)))
np.testing.assert_equal(dx.wrap_xpos, np.zeros((2, 6)))
np.testing.assert_allclose(dx._impl.ten_length, d.ten_length)
np.testing.assert_equal(dx._impl.ten_wrapadr, np.zeros((1,)))
np.testing.assert_equal(dx._impl.ten_wrapnum, np.zeros((1,)))
np.testing.assert_equal(dx._impl.wrap_obj, np.zeros((2, 2)))
np.testing.assert_equal(dx._impl.wrap_xpos, np.zeros((2, 6)))
# efc_ are also shape transformed and padded
self.assertEqual(dx.efc_J.shape, (45, 8)) # nefc, nv
self.assertEqual(dx._impl.efc_J.shape, (45, 8)) # nefc, nv
d_efc_j = d.efc_J.reshape((-1, 8))
np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3]) # connect eq
np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3]) # one active limit
np.testing.assert_allclose(dx.efc_J[4], 0) # one inactive limit
np.testing.assert_allclose(dx.efc_J[5:15], d_efc_j[4:14]) # contact
np.testing.assert_allclose(dx.efc_J[15:], 0) # no contact
np.testing.assert_allclose(dx._impl.efc_J[:3], d_efc_j[:3]) # connect eq
np.testing.assert_allclose(
dx._impl.efc_J[3], d_efc_j[3]
) # one active limit
np.testing.assert_allclose(dx._impl.efc_J[4], 0) # one inactive limit
np.testing.assert_allclose(dx._impl.efc_J[5:15], d_efc_j[4:14]) # contact
np.testing.assert_allclose(dx._impl.efc_J[15:], 0) # no contact
# check another efc_ too
self.assertEqual(dx.efc_aref.shape, (45,)) # nefc
np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
np.testing.assert_allclose(dx.efc_aref[4], 0)
np.testing.assert_allclose(dx.efc_aref[5:15], d.efc_aref[4:14])
np.testing.assert_allclose(dx.efc_aref[15:], 0)
self.assertEqual(dx._impl.efc_aref.shape, (45,)) # nefc
np.testing.assert_allclose(dx._impl.efc_aref[:3], d.efc_aref[:3])
np.testing.assert_allclose(dx._impl.efc_aref[3], d.efc_aref[3])
np.testing.assert_allclose(dx._impl.efc_aref[4], 0)
np.testing.assert_allclose(dx._impl.efc_aref[5:15], d.efc_aref[4:14])
np.testing.assert_allclose(dx._impl.efc_aref[15:], 0)
# check sparse transform is correct
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx_sparse = mjx.put_data(m, d)
np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
dx_sparse = mjx.put_data(m, d, backend_impl=backend_impl)
np.testing.assert_allclose(dx_sparse._impl.efc_J, dx._impl.efc_J, atol=1e-8)
# check sparse mass matrices are correct
np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8)
np.testing.assert_allclose(
dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8
)
# check dense mass matrices are correct
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx_from_dense = mjx.put_data(m, d)
qm = np.zeros((m.nv, m.nv))
mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
dx_from_dense = mjx.put_data(m, d, backend_impl=backend_impl)
if backend_impl == 'jax':
qm = np.zeros((m.nv, m.nv))
mujoco.mj_fullM(m, qm, d.qM)
np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
elif backend_impl == 'c':
np.testing.assert_allclose(dx_from_dense._impl.qM, d.qM, atol=1e-8)
def test_get_data(self):
@parameterized.parameters(
('jax', False), ('jax', True), ('c', False), ('c', True)
)
def test_get_data(self, backend_impl: str, sparse: bool):
"""Test that get_data makes correct MjData."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
if sparse:
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
d_2: mujoco.MjData = mjx.get_data(m, dx)
# check a few fields
@@ -398,6 +458,8 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(d_2.cvel, d.cvel)
np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
np.testing.assert_allclose(d_2.qM, d.qM)
np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6)
np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6)
# only 1 contact active
self.assertEqual(d_2.contact.dist.shape, (1,))
@@ -418,12 +480,45 @@ class DataIOTest(parameterized.TestCase):
# efc_* are also shape transformed and filtered
self.assertEqual(d_2.nefc, 14)
self.assertEqual(d_2.efc_J.shape, (112,)) # nefc * nv
np.testing.assert_allclose(d_2.efc_J, d.efc_J)
if sparse:
efc_j = np.zeros((d.nefc, m.nv))
mujoco.mju_sparse2dense(
efc_j,
d.efc_J,
d.efc_J_rownnz,
d.efc_J_rowadr,
d.efc_J_colind,
)
efc_j2 = np.zeros((d_2.nefc, m.nv))
mujoco.mju_sparse2dense(
efc_j2,
d_2.efc_J,
d_2.efc_J_rownnz,
d_2.efc_J_rowadr,
d_2.efc_J_colind,
)
np.testing.assert_allclose(efc_j, efc_j2)
else:
self.assertEqual(d_2.efc_J.shape, (112,)) # nefc * nv
np.testing.assert_allclose(d_2.efc_J, d.efc_J)
self.assertEqual(d_2.efc_aref.shape, (14,)) # nefc
np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
if backend_impl == 'c':
# check fields specific to the C implementation
np.testing.assert_allclose(d_2.bvh_active, d.bvh_active)
def test_get_data_simplebody(self):
"""Test that get_data works with simple bodies where nC < nM."""
m = mujoco.MjModel.from_xml_string(_SIMPLE_BODY)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
d_2: mujoco.MjData = mjx.get_data(m, dx)
np.testing.assert_allclose(d_2.qLD, d.qLD, atol=1e-6)
np.testing.assert_allclose(d_2.qLDiagInv, d.qLDiagInv, atol=1e-6)
def test_get_data_runs(self):
xml = """
<mujoco>
@@ -444,13 +539,14 @@ class DataIOTest(parameterized.TestCase):
dx = mjx.put_data(m, d)
mjx.get_data(m, dx)
def test_get_data_batched(self):
@parameterized.parameters('jax', 'c')
def test_get_data_batched(self, backend_impl):
"""Test that get_data makes correct List[MjData] for batched Data."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
# second data in batch has contact dist > 0, disables contact
dx_b = jax.tree_util.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
ds = mjx.get_data(m, dx_b)
@@ -460,13 +556,14 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(ds[0].ncon, 1)
self.assertEqual(ds[1].ncon, 0)
def test_get_data_into(self):
@parameterized.parameters('jax', 'c')
def test_get_data_into(self, backend_impl):
"""Test that get_data_into correctly populates an MjData."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
d_2 = mujoco.MjData(m)
mjx.get_data_into(d_2, m, dx)
@@ -482,32 +579,33 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d_2.contact.frame.shape, (1, 9))
np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
def test_get_data_into_wrong_shape(self):
@parameterized.parameters('jax', 'c')
def test_get_data_into_wrong_shape(self, backend_impl):
"""Tests that get_data_into throwsif input and output shapes don't match."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
m_2 = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
d_2 = mujoco.MjData(m_2)
with self.assertRaisesRegex(ValueError, r'Input field.*has shape.*'):
mjx.get_data_into(d_2, m, dx)
def test_make_matches_put(self):
@parameterized.parameters('jax', 'c')
def test_make_matches_put(self, backend_impl):
"""Test that make_data produces a pytree that matches put_data."""
m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
d = mujoco.MjData(m)
mujoco.mj_step(m, d, 2)
dx = mjx.put_data(m, d)
dx = mjx.put_data(m, d, backend_impl=backend_impl)
step_fn = lambda d: d.replace(time=d.time + 1)
step_fn_jit = jax.jit(step_fn).lower(dx).compile()
# placing an MjData onto device should yield the same treedef mjx.Data as
# calling make_data. they should be interchangeable for jax functions:
step_fn_jit(mjx.make_data(m))
step_fn_jit(mjx.make_data(m, backend_impl=backend_impl))
def test_contact_elliptic_condim1(self):
"""Test that condim=1 with ConeType.ELLIPTIC is not implemented."""
@@ -556,7 +654,327 @@ class DataIOTest(parameterized.TestCase):
</mujoco>
""")
with self.assertRaises(NotImplementedError):
mjx.put_model(m)
mjx.put_model(m, backend_impl='jax')
class FullCompatTest(parameterized.TestCase):
"""Tests for the _full_compat flag."""
def test_full_compat_deprecated(self):
"""Tests that _full_compat is deprecated."""
xml = """
<mujoco>
<worldbody>
<body name="box">
<joint name="slide1" type="slide" axis="1 0 0" />
<geom type="box" size=".05 .05 .05" mass="1"/>
</body>
</worldbody>
<actuator>
<motor joint="slide1"/>
</actuator>
</mujoco>
"""
m = mujoco.MjModel.from_xml_string(xml)
with self.assertWarns(DeprecationWarning):
out = mjx_io.put_model(m, _full_compat=True)
self.assertEqual(out.backend_impl, BackendImpl.C)
with self.assertWarns(DeprecationWarning):
out = mjx_io.make_data(m, _full_compat=True)
self.assertEqual(out.backend_impl, BackendImpl.C)
# Test cases for `_resolve_backend_impl_and_device` where the device is
# specified by the user and the device is available.
_DEVICE_TEST_CASES = [
# Arguments use the following format:
# (device_type_str, backend_impl_str,
# (expected_device, expected_backend_impl)))
# No backend specified.
('cpu', None, ('cpu', BackendImpl.C)),
('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
('tpu', None, ('tpu', BackendImpl.JAX)),
# JAX backend specified.
('cpu', 'jax', ('cpu', BackendImpl.JAX)),
('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
('tpu', 'jax', ('tpu', BackendImpl.JAX)),
# WARP backend specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
# C backend specified.
('cpu', 'c', ('cpu', BackendImpl.C)),
('gpu-notnvidia', 'c', ('cpu', 'error')),
('gpu-nvidia', 'c', ('cpu', 'error')),
('tpu', 'c', ('tpu', 'error')),
]
# Test cases for `_resolve_backend_impl_and_device` where the user does NOT
# specify a device. We mock the JAX default device.
_DEFAULT_DEVICE_TEST_CASES = [
# Arguments use the following format:
# (jax.default_device, backend_impl_str,
# (expected_device, expected_backend_impl))
# No backend impl specified.
('cpu', None, ('cpu', BackendImpl.C)),
('gpu-notnvidia', None, ('gpu', BackendImpl.JAX)),
('gpu-nvidia', None, ('gpu', BackendImpl.WARP)),
('tpu', None, ('tpu', BackendImpl.JAX)),
# JAX backend impl specified.
('cpu', 'jax', ('cpu', BackendImpl.JAX)),
('gpu-notnvidia', 'jax', ('gpu', BackendImpl.JAX)),
('gpu-nvidia', 'jax', ('gpu', BackendImpl.JAX)),
('tpu', 'jax', ('tpu', BackendImpl.JAX)),
# WARP backend impl specified.
('cpu', 'warp', ('cpu', 'error')),
('gpu-notnvidia', 'warp', ('cpu', 'error')),
('gpu-nvidia', 'warp', ('gpu', BackendImpl.WARP)),
('tpu', 'warp', ('tpu', 'error')),
# C backend impl specified, CPU should always be available.
('cpu', 'c', ('cpu', BackendImpl.C)),
('gpu-notnvidia', 'c', ('cpu', BackendImpl.C)),
('gpu-nvidia', 'c', ('cpu', BackendImpl.C)),
('tpu', 'c', ('cpu', BackendImpl.C)),
]
class ResolveBackendImplAndDeviceTest(parameterized.TestCase):
"""Tests for the _resolve_backend_impl_and_device function."""
def setUp(self):
super().setUp()
# Create mock devices
self.mock_cpu = mock.Mock(spec=jax.Device)
self.mock_cpu.platform = 'cpu'
self.mock_cpu.device_kind = 'Mock CPU'
self.mock_cpu.id = 0
self.mock_nvidia_gpu = mock.Mock(spec=jax.Device)
self.mock_nvidia_gpu.platform = 'gpu'
self.mock_nvidia_gpu.device_kind = 'NVIDIA Mocked GPU'
self.mock_nvidia_gpu.id = 0
self.mock_other_gpu = mock.Mock(spec=jax.Device)
self.mock_other_gpu.platform = 'gpu'
self.mock_other_gpu.device_kind = 'Other Mocked GPU'
self.mock_other_gpu.id = 1
self.mock_tpu = mock.Mock(spec=jax.Device)
self.mock_tpu.platform = 'tpu'
self.mock_tpu.device_kind = 'Mock TPU'
self.mock_tpu.id = 0
# Patch jax.devices for the entire test class using enter_context
self.mock_jax_devices = self.enter_context(mock.patch('jax.devices'))
self.mock_default_backend = self.enter_context(
mock.patch('jax.default_backend')
)
@parameterized.named_parameters(
(f'{str(args[0])}_{str(args[1])}', *args) for args in _DEVICE_TEST_CASES
)
@mock.patch.dict(
os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
)
def test_resolve_with_device(
self,
device_type_str,
backend_impl_str,
expected,
):
"""Tests various combinations of device and backend impls."""
input_device = {
'cpu': self.mock_cpu,
'gpu-nvidia': self.mock_nvidia_gpu,
'gpu-notnvidia': self.mock_other_gpu,
'tpu': self.mock_tpu,
}[device_type_str]
def devices_side_effect(backend=None):
# assume the user-specified device is always available
if backend == 'cpu':
return [self.mock_cpu]
elif backend == 'gpu':
if 'nvidia' in device_type_str:
return [self.mock_nvidia_gpu]
return [self.mock_other_gpu]
elif backend == 'tpu':
return [self.mock_tpu]
elif backend == 'cuda':
return [self.mock_nvidia_gpu]
raise AssertionError('Should not be called.')
self.mock_jax_devices.side_effect = devices_side_effect
expected_device, expected_backend_impl = expected
if expected_backend_impl == 'error':
with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device(
backend_impl=backend_impl_str, device=input_device
)
return
actual_backend_impl, actual_device = (
mjx_io._resolve_backend_impl_and_device(
backend_impl=backend_impl_str, device=input_device
)
)
self.assertEqual(actual_backend_impl, expected_backend_impl)
self.assertIsNotNone(actual_device)
self.assertEqual(actual_device.platform, expected_device)
@parameterized.named_parameters(
(f'{str(args[0])}_{str(args[1])}', *args)
for args in _DEFAULT_DEVICE_TEST_CASES
)
@mock.patch.dict(
os.environ, {'MJX_WARP_ENABLED': 'true', 'MJX_C_DEFAULT_ENABLED': 'true'}
)
def test_resolve_without_device(
self,
default_device_str,
backend_impl_str,
expected,
):
"""Tests various combinations of jax.default_device and backend impls."""
default_devices = {
'cpu': [self.mock_cpu],
'gpu-nvidia': [self.mock_nvidia_gpu, self.mock_cpu],
'gpu-notnvidia': [self.mock_other_gpu, self.mock_cpu],
'tpu': [self.mock_tpu, self.mock_cpu],
}[default_device_str]
def devices_side_effect(backend=None):
if backend == 'cpu':
return [self.mock_cpu] # CPU is always available
if backend == 'gpu' and default_device_str == 'gpu-notnvidia':
return [self.mock_other_gpu]
if backend == 'gpu' and default_device_str == 'gpu-nvidia':
return [self.mock_nvidia_gpu]
if backend == 'cuda' and default_device_str == 'gpu-nvidia':
return [self.mock_nvidia_gpu]
if backend == 'tpu' and default_device_str == 'tpu':
return [self.mock_tpu]
if backend is None:
return default_devices
if backend == 'cuda':
raise RuntimeError('cuda backend not supported')
raise AssertionError('jax.devices error')
self.mock_jax_devices.side_effect = devices_side_effect
default_device_side_effect_str = {
'cpu': 'cpu',
'gpu-nvidia': 'gpu',
'gpu-notnvidia': 'gpu',
'tpu': 'tpu',
}[default_device_str]
self.mock_default_backend.side_effect = (
lambda: default_device_side_effect_str
)
expected_device, expected_backend_impl = expected
if (
expected_backend_impl == 'error'
and default_device_str != 'gpu-nvidia'
and backend_impl_str == 'warp'
):
with self.assertRaisesRegex(RuntimeError, 'cuda backend not supported'):
mjx_io._resolve_backend_impl_and_device(
backend_impl=backend_impl_str, device=None
)
return
if expected_backend_impl == 'error':
with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device(
backend_impl=backend_impl_str, device=None
)
return
actual_backend_impl, actual_device = (
mjx_io._resolve_backend_impl_and_device(
backend_impl=backend_impl_str, device=None
)
)
self.assertEqual(actual_backend_impl, expected_backend_impl)
self.assertIsNotNone(actual_device)
self.assertEqual(actual_device.platform, expected_device)
@mock.patch.dict(os.environ, {'MJX_WARP_ENABLED': 'false'})
def test_resolve_warp_disabled(self):
"""Tests behavior when MJX_WARP_ENABLED is false."""
self.mock_jax_devices.side_effect = lambda backend=None: (
[self.mock_nvidia_gpu, self.mock_cpu]
if backend is None
else ([self.mock_nvidia_gpu] if backend == 'gpu' else [self.mock_cpu])
)
self.mock_default_backend.side_effect = lambda: 'gpu'
# Default to JAX instead of WARP on NVIDIA GPU.
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl=None, device=None
)
self.assertEqual(backend_impl, BackendImpl.JAX)
self.assertEqual(device.platform, 'gpu')
# Specifying an NVIDIA GPU should still choose JAX.
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl=None, device=self.mock_nvidia_gpu
)
self.assertEqual(backend_impl, BackendImpl.JAX)
self.assertEqual(device.platform, 'gpu')
# Requesting warp explicitly should fail since it is disabled.
with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device(
backend_impl='warp', device=self.mock_nvidia_gpu
)
with self.assertRaises(AssertionError):
mjx_io._resolve_backend_impl_and_device(backend_impl='warp', device=None)
@mock.patch.dict(os.environ, {'MJX_C_DEFAULT_ENABLED': 'false'})
def test_resolve_c_disabled(self):
"""Tests behavior when MJX_C_DEFAULT_ENABLED is false."""
# Users expect that CPU defaults to the JAX impl. But in the future, it will
# default to the C backend implementation. This test checks that
# MJX_C_DEFAULT_ENABLED=false defaults to the old behavior, until the
# migration to MJEP-15 is complete.
self.mock_jax_devices.side_effect = lambda backend=None: ([self.mock_cpu])
self.mock_default_backend.side_effect = lambda: 'cpu'
# Default to JAX instead of C on CPU.
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl=None, device=None
)
self.assertEqual(backend_impl, BackendImpl.JAX)
self.assertEqual(device.platform, 'cpu')
# Specifing CPU should still choose JAX.
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl=None, device=self.mock_cpu
)
self.assertEqual(backend_impl, BackendImpl.JAX)
self.assertEqual(device.platform, 'cpu')
# Specifying C should choose C!
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl='c', device=None
)
self.assertEqual(backend_impl, BackendImpl.C)
self.assertEqual(device.platform, 'cpu')
backend_impl, device = mjx_io._resolve_backend_impl_and_device(
backend_impl='c', device=self.mock_cpu
)
self.assertEqual(backend_impl, BackendImpl.C)
self.assertEqual(device.platform, 'cpu')
if __name__ == '__main__':
mjx/mujoco/mjx/_src/passive.py
+10 -4
View File
@@ -23,9 +23,12 @@ from mujoco.mjx._src import scan
from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import JointType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import OptionJAX
# pylint: enable=g-importing-member
@@ -72,11 +75,11 @@ def _spring_damper(m: Model, d: Data) -> jax.Array:
qfrc -= m.dof_damping * d.qvel
# tendon-level spring-dampers
below, above = m.tendon_lengthspring.T - d.ten_length
below, above = m.tendon_lengthspring.T - d._impl.ten_length
frc_spring = jp.where(below > 0, m.tendon_stiffness * below, 0)
frc_spring = jp.where(above < 0, m.tendon_stiffness * above, frc_spring)
frc_damper = -m.tendon_damping * d.ten_velocity
qfrc += d.ten_J.T @ (frc_spring + frc_damper)
frc_damper = -m.tendon_damping * d._impl.ten_velocity
qfrc += d._impl.ten_J.T @ (frc_spring + frc_damper)
return qfrc
@@ -113,6 +116,9 @@ def _fluid(m: Model, d: Data) -> jax.Array:
def passive(m: Model, d: Data) -> Data:
"""Adds all passive forces."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('passive requires JAX backend implementation.')
if m.opt.disableflags & DisableBit.PASSIVE:
return d.replace(qfrc_passive=jp.zeros(m.nv), qfrc_gravcomp=jp.zeros(m.nv))
@@ -124,7 +130,7 @@ def passive(m: Model, d: Data) -> Data:
# add gravcomp unless added via actuators
qfrc_passive += qfrc_gravcomp * (1 - m.jnt_actgravcomp[m.dof_jntid])
if m.opt.has_fluid_params:
if m.opt.has_fluid_params: # pytype: disable=attribute-error
qfrc_passive += _fluid(m, d)
d = d.replace(qfrc_passive=qfrc_passive, qfrc_gravcomp=qfrc_gravcomp)
mjx/mujoco/mjx/_src/scan.py
+2
View File
@@ -132,6 +132,8 @@ def _nvmap(f: Callable[..., Y], *args) -> Y:
def _check_input(m: Model, args: Any, in_types: str) -> None:
"""Checks that scan input has the right shape."""
if m.nv == 0:
raise ValueError('Scan across Model with zero DoFs unsupported.')
size = {
'b': m.nbody,
'j': m.njnt,
mjx/mujoco/mjx/_src/scan_test.py
+1 -4
View File
@@ -61,10 +61,7 @@ class ScanTest(absltest.TestCase):
return body_id + 1
b_in = jp.array([1])
b_expect = jp.array([2])
b_out = scan.flat(m, fn, 'b', 'b', b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
self.assertRaises(ValueError, scan.flat, m, fn, 'b', 'b', b_in)
def test_flat_joints(self):
"""Tests scanning over bodies with joints of different types."""
mjx/mujoco/mjx/_src/sensor.py
+47 -24
View File
@@ -22,16 +22,20 @@ from mujoco.mjx._src import math
from mujoco.mjx._src import ray
from mujoco.mjx._src import smooth
from mujoco.mjx._src import support
from mujoco.mjx._src.types import BackendImpl
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import ObjType
from mujoco.mjx._src.types import SensorType
from mujoco.mjx._src.types import TrnType
# pylint: enable=g-importing-member
import numpy as np
def apply_cutoff(
def _apply_cutoff(
sensor: jax.Array, cutoff: jax.Array, data_type: int
) -> jax.Array:
"""Clip sensor to cutoff value."""
@@ -50,6 +54,8 @@ def apply_cutoff(
def sensor_pos(m: Model, d: Data) -> Data:
"""Compute position-dependent sensors values."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('sensor_pos requires JAX backend implementation.')
if m.opt.disableflags & DisableBit.SENSOR:
return d
@@ -162,15 +168,15 @@ def sensor_pos(m: Model, d: Data) -> Data:
sensor, _ = jax.vmap(
ray.ray, in_axes=(None, None, 0, 0, None, None, None)
)(m, d, site_xpos, site_mat, (), True, sid)
sensors.append(apply_cutoff(sensor, cutoffs, data_type[0]))
sensors.append(_apply_cutoff(sensor, cutoffs, data_type[0]))
adrs.append(adr[idxs])
continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.JOINTPOS:
sensor = d.qpos[m.jnt_qposadr[objid]]
elif sensor_type == SensorType.TENDONPOS:
sensor = d.ten_length[objid]
sensor = d._impl.ten_length[objid]
elif sensor_type == SensorType.ACTUATORPOS:
sensor = d.actuator_length[objid]
sensor = d._impl.actuator_length[objid]
elif sensor_type == SensorType.BALLQUAT:
jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
quat = d.qpos[jnt_qposadr]
@@ -193,7 +199,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
cutofft = cutoff[idxt]
sensor = jax.vmap(_framepos)(xpos, xpos_ref, xmat_ref, refidt)
adrt = adr[idxt, None] + np.arange(3)[None]
sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time
elif sensor_type in frame_axis:
@@ -213,7 +219,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
cutofft = cutoff[idxt]
sensor = jax.vmap(_frameaxis)(xmat, xmat_ref, refidt)
adrt = adr[idxt, None] + np.arange(3)[None]
sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.FRAMEQUAT:
@@ -254,7 +260,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
)
)(quat, refquat, refidt)
adrt = adr[idxt, None] + np.arange(4)[None]
sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.SUBTREECOM:
@@ -266,7 +272,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
# TODO(taylorhowell): raise error after adding sensor check to io.py
continue # unsupported sensor type
sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
adrs.append(adr)
if not adrs:
@@ -281,6 +287,8 @@ def sensor_pos(m: Model, d: Data) -> Data:
def sensor_vel(m: Model, d: Data) -> Data:
"""Compute velocity-dependent sensors values."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('sensor_vel requires JAX backend implementation.')
if m.opt.disableflags & DisableBit.SENSOR:
return d
@@ -332,9 +340,9 @@ def sensor_vel(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.JOINTVEL:
sensor = d.qvel[m.jnt_dofadr[objid]]
elif sensor_type == SensorType.TENDONVEL:
sensor = d.ten_velocity[objid]
sensor = d._impl.ten_velocity[objid]
elif sensor_type == SensorType.ACTUATORVEL:
sensor = d.actuator_velocity[objid]
sensor = d._impl.actuator_velocity[objid]
elif sensor_type == SensorType.BALLANGVEL:
jnt_dotadr = m.jnt_dofadr[objid, None] + np.arange(3)[None]
sensor = d.qvel[jnt_dotadr]
@@ -392,20 +400,20 @@ def sensor_vel(m: Model, d: Data) -> Data:
adrt = adr[idxt, None] + np.arange(3)[None]
sensors.append(apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutofft, data_type[0]).reshape(-1))
adrs.append(adrt.reshape(-1))
continue # avoid adding to sensors/adrs list a second time
elif sensor_type == SensorType.SUBTREELINVEL:
sensor = d.subtree_linvel[objid]
sensor = d._impl.subtree_linvel[objid]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.SUBTREEANGMOM:
sensor = d.subtree_angmom[objid]
sensor = d._impl.subtree_angmom[objid]
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
else:
# TODO(taylorhowell): raise error after adding sensor check to io.py
continue # unsupported sensor type
sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
adrs.append(adr)
if not adrs:
@@ -420,6 +428,8 @@ def sensor_vel(m: Model, d: Data) -> Data:
def sensor_acc(m: Model, d: Data) -> Data:
"""Compute acceleration/force-dependent sensors values."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('sensor_acc requires JAX backend implementation.')
if m.opt.disableflags & DisableBit.SENSOR:
return d
@@ -441,6 +451,8 @@ def sensor_acc(m: Model, d: Data) -> Data:
SensorType.ACCELEROMETER,
SensorType.FORCE,
SensorType.TORQUE,
SensorType.FRAMELINACC,
SensorType.FRAMEANGACC,
}:
d = smooth.rne_postconstraint(m, d)
@@ -457,14 +469,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
# compute contact forces
forces = []
condim_ids = []
for dim in set(d.contact.dim):
for dim in set(d._impl.contact.dim):
force, condim_id = support.contact_force_dim(m, d, dim)
forces.append(force)
condim_ids.append(condim_id)
forces = jp.concatenate(forces)[np.argsort(np.concatenate(condim_ids))]
# get bodies of contact geoms
conbody = jp.array(m.geom_bodyid)[d.contact.geom]
conbody = jp.array(m.geom_bodyid)[d._impl.contact.geom]
# get site information
site_bodyid = m.site_bodyid[objid]
@@ -474,12 +486,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
site_type = m.site_type[objid]
conbody0 = site_bodyid[:, None] == conbody[:, 0]
conbody1 = site_bodyid[:, None] == conbody[:, 1]
contacts = (d.contact.efc_address >= 0)[None] & (conbody0 | conbody1)
contacts = (d._impl.contact.efc_address >= 0)[None] & (
conbody0 | conbody1
)
# compute conray, flip if second body
conray = jax.vmap(
lambda frame, force: math.normalize(frame[0] * force[0])
)(d.contact.frame, forces)
)(d._impl.contact.frame, forces)
conray = jp.where(conbody1[..., None], -conray, conray)
# compute distance, mapping over sites and contacts
@@ -501,7 +515,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
site_xpos[dist_id_site],
site_xmat[dist_id_site],
st,
d.contact.pos,
d._impl.contact.pos,
conray[dist_id_site],
)
dist.append(jp.where(jp.isinf(dist_site), 0, dist_site))
@@ -523,14 +537,14 @@ def sensor_acc(m: Model, d: Data) -> Data:
bodyid = m.site_bodyid[objid]
rot = d.site_xmat[objid]
cvel = d.cvel[bodyid]
cacc = d.cacc[bodyid]
cacc = d._impl.cacc[bodyid]
dif = d.site_xpos[objid] - d.subtree_com[m.body_rootid[bodyid]]
sensor = _accelerometer(cvel, cacc, dif, rot)
adr = (adr[:, None] + np.arange(3)[None]).reshape(-1)
elif sensor_type == SensorType.FORCE:
bodyid = m.site_bodyid[objid]
cfrc_int = d.cfrc_int[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid]
site_xmat = d.site_xmat[objid]
sensor = jax.vmap(lambda mat, vec: mat.T @ vec)(
site_xmat, cfrc_int[:, 3:]
@@ -539,7 +553,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.TORQUE:
bodyid = m.site_bodyid[objid]
rootid = m.body_rootid[bodyid]
cfrc_int = d.cfrc_int[bodyid]
cfrc_int = d._impl.cfrc_int[bodyid]
site_xmat = d.site_xmat[objid]
dif = d.site_xpos[objid] - d.subtree_com[rootid]
sensor = jax.vmap(
@@ -550,6 +564,15 @@ def sensor_acc(m: Model, d: Data) -> Data:
sensor = d.actuator_force[objid]
elif sensor_type == SensorType.JOINTACTFRC:
sensor = d.qfrc_actuator[m.jnt_dofadr[objid]]
elif sensor_type == SensorType.TENDONACTFRC:
force_mask = [
(m.actuator_trntype == TrnType.TENDON)
& (m.actuator_trnid[:, 0] == tendon_id)
for tendon_id in objid
]
force_ids = np.concatenate([np.nonzero(mask)[0] for mask in force_mask])
force_mat = np.array(force_mask)[:, force_ids]
sensor = force_mat @ d.actuator_force[force_ids]
elif sensor_type in (SensorType.FRAMELINACC, SensorType.FRAMEANGACC):
objtype = m.sensor_objtype[idx]
@@ -559,7 +582,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
pos, bodyid = objtype_data[ot]
pos = pos[objidt]
bodyid = bodyid[objidt]
cacc = d.cacc[bodyid]
cacc = d._impl.cacc[bodyid]
if sensor_type == SensorType.FRAMELINACC:
@@ -589,7 +612,7 @@ def sensor_acc(m: Model, d: Data) -> Data:
# TODO(taylorhowell): raise error after adding sensor check to io.py
continue # unsupported sensor type
sensors.append(apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
sensors.append(_apply_cutoff(sensor, cutoff, data_type[0]).reshape(-1))
adrs.append(adr)
if not adrs:
mjx/mujoco/mjx/_src/sensor_test.py
+8 -8
View File
@@ -64,14 +64,14 @@ class SensorTest(parameterized.TestCase):
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = mjx.put_data(m, d).replace(
sensordata=jp.zeros_like(d.sensordata),
subtree_linvel=jp.zeros_like(d.subtree_linvel),
subtree_angmom=jp.zeros_like(d.subtree_angmom),
cacc=jp.zeros_like(d.cacc),
cfrc_int=jp.zeros_like(d.cfrc_int),
cfrc_ext=jp.zeros_like(d.cfrc_ext),
)
dx = mjx.put_data(m, d).tree_replace({
'sensordata': jp.zeros_like(d.sensordata),
'_impl.subtree_linvel': jp.zeros_like(d.subtree_linvel),
'_impl.subtree_angmom': jp.zeros_like(d.subtree_angmom),
'_impl.cacc': jp.zeros_like(d.cacc),
'_impl.cfrc_int': jp.zeros_like(d.cfrc_int),
'_impl.cfrc_ext': jp.zeros_like(d.cfrc_ext),
})
dx = jax.jit(mjx.sensor_pos)(mx, dx)
dx = jax.jit(mjx.sensor_vel)(mx, dx)
dx = jax.jit(mjx.sensor_acc)(mx, dx)
mjx/mujoco/mjx/_src/smooth.py
+101 -63
View File
@@ -23,10 +23,12 @@ from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import CamLightType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import EqType
from mujoco.mjx._src.types import JointType
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import TrnType
from mujoco.mjx._src.types import WrapType
# pylint: enable=g-importing-member
@@ -35,7 +37,6 @@ import numpy as np
def kinematics(m: Model, d: Data) -> Data:
"""Converts position/velocity from generalized coordinates to maximal."""
def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
# calculate joint anchors, axes, body pos and quat in global frame
# also normalize qpos while we're at it
@@ -131,6 +132,8 @@ def kinematics(m: Model, d: Data) -> Data:
def com_pos(m: Model, d: Data) -> Data:
"""Maps inertias and motion dofs to global frame centered at subtree-CoM."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('com_pos requires JAX backend implementation.')
# calculate center of mass of each subtree
def subtree_sum(carry, xipos, body_mass):
@@ -162,7 +165,7 @@ def com_pos(m: Model, d: Data) -> Data:
root_com = subtree_com[m.body_rootid]
offset = d.xipos - root_com
cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
d = d.replace(cinert=cinert)
d = d.tree_replace({'_impl.cinert': cinert})
# map motion dofs to global frame centered at subtree_com
def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis):
@@ -201,13 +204,16 @@ def com_pos(m: Model, d: Data) -> Data:
d.xanchor,
d.xaxis,
)
d = d.replace(cdof=cdof)
d = d.tree_replace({'_impl.cdof': cdof})
return d
def camlight(m: Model, d: Data) -> Data:
"""Computes camera and light positions and orientations."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('camlight requires JAX backend implementation.')
if m.ncam == 0:
return d.replace(cam_xpos=jp.zeros((0, 3)), cam_xmat=jp.zeros((0, 3, 3)))
@@ -278,32 +284,33 @@ def camlight(m: Model, d: Data) -> Data:
def crb(m: Model, d: Data) -> Data:
"""Runs composite rigid body inertia algorithm."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('crb requires JAX backend implementation.')
def crb_fn(crb_child, crb_body):
if crb_child is not None:
crb_body += crb_child
return crb_body
crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d.cinert, reverse=True)
crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d._impl.cinert, reverse=True)
crb_body = crb_body.at[0].set(0.0)
d = d.replace(crb=crb_body)
d = d.tree_replace({'_impl.crb': crb_body})
crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.replace(qM=qm)
if support.is_sparse(m) and d._qM_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qM_sparse=qm)
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof)
qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature)
d = d.tree_replace({'_impl.qM': qm})
return d
def factor_m(m: Model, d: Data) -> Data:
"""Gets factorizaton of inertia-like matrix M, assumed spd."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('factor_m requires JAX backend implementation.')
if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d.qM)
d = d.replace(qLD=qh)
qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM)
d = d.tree_replace({'_impl.qLD': qh})
return d
# build up indices for where we will do backwards updates over qLD
@@ -325,7 +332,7 @@ def factor_m(m: Model, d: Data) -> Data:
(out_beg, out_end, madr_d, madr_ij)
)
qld = d.qM
qld = d._impl.qM
for _, updates in sorted(updates.items(), reverse=True):
# combine the updates into one update batch (per depth level)
@@ -353,20 +360,17 @@ def factor_m(m: Model, d: Data) -> Data:
qld_diag = qld[m.dof_Madr]
qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag)
if d._qLD_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qLD_sparse=d.qLD)
if d._qLDiagInv_sparse.size > 0: # pylint: disable=protected-access
d = d.replace(_qLDiagInv_sparse=d.qLDiagInv)
d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
return d
def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
"""Computes sparse backsubstitution: x = inv(L'*D*L)*y ."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('solve_m requires JAX backend implementation.')
if not support.is_sparse(m):
return jax.scipy.linalg.cho_solve((d.qLD, False), x)
return jax.scipy.linalg.cho_solve((d._impl.qLD, False), x)
depth = []
for i in range(m.nv):
@@ -385,21 +389,23 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
# x <- inv(L') * x
for _, vals in sorted(updates_j.items(), reverse=True):
j, madr_ij, i = np.array(vals).T
x = x.at[j].add(-d.qLD[madr_ij] * x[i])
x = x.at[j].add(-d._impl.qLD[madr_ij] * x[i])
# x <- inv(D) * x
x = x * d.qLDiagInv
x = x * d._impl.qLDiagInv
# x <- inv(L) * x
for _, vals in sorted(updates_i.items()):
i, madr_ij, j = np.array(vals).T
x = x.at[i].add(-d.qLD[madr_ij] * x[j])
x = x.at[i].add(-d._impl.qLD[madr_ij] * x[j])
return x
def com_vel(m: Model, d: Data) -> Data:
"""Computes cvel, cdof_dot."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('com_vel requires JAX backend implementation.')
# forward scan down tree: accumulate link center of mass velocity
def fn(parent, jnt_typs, cdof, qvel):
@@ -431,17 +437,19 @@ def com_vel(m: Model, d: Data) -> Data:
'jvv',
'bv',
m.jnt_type,
d.cdof,
d._impl.cdof,
d.qvel,
)
d = d.replace(cvel=cvel, cdof_dot=cdof_dot)
d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot})
return d
def subtree_vel(m: Model, d: Data) -> Data:
"""Subtree linear velocity and angular momentum."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('subtree_vel requires JAX backend implementation.')
# bodywise quantities
def _forward(cvel, xipos, ximat, subtree_com_root, mass, inertia):
@@ -529,14 +537,22 @@ def subtree_vel(m: Model, d: Data) -> Data:
reverse=True,
)
return d.replace(subtree_linvel=subtree_linvel, subtree_angmom=subtree_angmom)
return d.tree_replace({
'_impl.subtree_linvel': subtree_linvel,
'_impl.subtree_angmom': subtree_angmom,
})
def rne(m: Model, d: Data) -> Data:
"""Computes inverse dynamics using the recursive Newton-Euler algorithm."""
def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
"""Computes inverse dynamics using the recursive Newton-Euler algorithm.
flg_acc=False removes inertial term.
"""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('rne requires JAX backend implementation.')
# forward scan over tree: accumulate link center of mass acceleration
def cacc_fn(cacc, cdof_dot, qvel):
def cacc_fn(cacc, cdof_dot, qvel, cdof, qacc):
if cacc is None:
if m.opt.disableflags & DisableBit.GRAVITY:
cacc = jp.zeros((6,))
@@ -545,9 +561,15 @@ def rne(m: Model, d: Data) -> Data:
cacc += jp.sum(jax.vmap(jp.multiply)(cdof_dot, qvel), axis=0)
# cacc += cdof * qacc
if flg_acc:
cacc += jp.sum(jax.vmap(jp.multiply)(cdof, qacc), axis=0)
return cacc
cacc = scan.body_tree(m, cacc_fn, 'vv', 'b', d.cdof_dot, d.qvel)
cacc = scan.body_tree(
m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc
)
def frc(cinert, cacc, cvel):
frc = math.inert_mul(cinert, cacc)
@@ -555,7 +577,7 @@ def rne(m: Model, d: Data) -> Data:
return frc
loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel)
loc_cfrc = jax.vmap(frc)(d._impl.cinert, cacc, d.cvel)
# backward scan up tree: accumulate body forces
def cfrc_fn(cfrc_child, cfrc):
@@ -564,7 +586,7 @@ def rne(m: Model, d: Data) -> Data:
return cfrc
cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True)
qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)])
qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)])
d = d.replace(qfrc_bias=qfrc_bias)
@@ -573,6 +595,8 @@ def rne(m: Model, d: Data) -> Data:
def rne_postconstraint(m: Model, d: Data) -> Data:
"""RNE with complete data: compute cacc, cfrc_ext, cfrc_int."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('rne_postconstraint requires JAX backend implementation.')
def _transform_force(frc, offset):
force, torque = jp.split(frc, 2)
@@ -593,7 +617,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
# compute contact forces for each condim
forces = []
condim_idx = []
for dim in set(d.contact.dim):
for dim in set(d._impl.contact.dim):
force, idx = support.contact_force_dim(m, d, dim)
forces.append(force)
condim_idx.append(idx)
@@ -620,10 +644,10 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
)
condim_idx = jp.concatenate(condim_idx)
frame = d.contact.frame[condim_idx]
pos = d.contact.pos[condim_idx]
id1 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 0]]
id2 = jp.array(m.geom_bodyid)[d.contact.geom[condim_idx, 1]]
frame = d._impl.contact.frame[condim_idx]
pos = d._impl.contact.pos[condim_idx]
id1 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 0]]
id2 = jp.array(m.geom_bodyid)[d._impl.contact.geom[condim_idx, 1]]
com1 = d.subtree_com[jp.array(m.body_rootid)][id1]
com2 = d.subtree_com[jp.array(m.body_rootid)][id2]
@@ -659,8 +683,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
)
# cacc = cacc_parent + cdofdot * qvel + cdof * qacc
cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
cacc_acc = d.cdof.T @ (mask * d.qacc)
cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel)
cacc_acc = d._impl.cdof.T @ (mask * d.qacc)
cacc = cacc_parent + cacc_vel + cacc_acc
# cfrc_body = cinert * cacc + cvel x (cinert * cvel)
@@ -678,7 +702,7 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
'bbbbb',
'bb',
cfrc_ext,
d.cinert,
d._impl.cinert,
d.cvel,
jp.array(m.body_dofadr),
jp.array(m.body_dofnum),
@@ -695,11 +719,18 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
)
# update data
return d.replace(cacc=cacc, cfrc_int=cfrc_int, cfrc_ext=cfrc_ext)
return d.tree_replace({
'_impl.cacc': cacc,
'_impl.cfrc_int': cfrc_int,
'_impl.cfrc_ext': cfrc_ext,
})
def tendon(m: Model, d: Data) -> Data:
"""Computes tendon lengths and moments."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('tendon requires JAX backend implementation.')
if not m.ntendon:
return d
@@ -827,8 +858,8 @@ def tendon(m: Model, d: Data) -> Data:
# wrap inside
# TODO(taylorhowell): check that is_wrap_inside is consistent with
# site and geom relative positions
(wrap_inside_id,) = np.nonzero(m.is_wrap_inside)
(wrap_outside_id,) = np.nonzero(~m.is_wrap_inside)
(wrap_inside_id,) = np.nonzero(m._impl.is_wrap_inside)
(wrap_outside_id,) = np.nonzero(~m._impl.is_wrap_inside)
# compute geom wrap length and connect points (if wrap occurs)
v_wrap = jax.vmap(
@@ -844,9 +875,9 @@ def tendon(m: Model, d: Data) -> Data:
has_sidesite[wrap_inside_id],
is_sphere[wrap_inside_id],
True,
m.wrap_inside_maxiter,
m.wrap_inside_tolerance,
m.wrap_inside_z_init,
m._impl.wrap_inside_maxiter,
m._impl.wrap_inside_tolerance,
m._impl.wrap_inside_z_init,
)
lengths_outside, pnt0_outside, pnt1_outside = v_wrap(
@@ -859,9 +890,9 @@ def tendon(m: Model, d: Data) -> Data:
has_sidesite[wrap_outside_id],
is_sphere[wrap_outside_id],
False,
m.wrap_inside_maxiter,
m.wrap_inside_tolerance,
m.wrap_inside_z_init,
m._impl.wrap_inside_maxiter,
m._impl.wrap_inside_tolerance,
m._impl.wrap_inside_z_init,
)
wrap_id = np.argsort(np.concatenate([wrap_inside_id, wrap_outside_id]))
@@ -943,12 +974,14 @@ def tendon(m: Model, d: Data) -> Data:
)
# assemble length and moment
ten_length = jp.zeros_like(d.ten_length).at[tendon_id_jnt].set(length_jnt)
ten_length = (
jp.zeros_like(d._impl.ten_length).at[tendon_id_jnt].set(length_jnt)
)
ten_length = ten_length.at[tendon_id_site].add(length_site)
ten_length = ten_length.at[tendon_id_geom].add(length_geom)
ten_moment = (
jp.zeros_like(d.ten_J)
jp.zeros_like(d._impl.ten_J)
.at[adr_moment_jnt, dofadr_moment_jnt]
.set(moment_jnt)
)
@@ -1011,14 +1044,14 @@ def tendon(m: Model, d: Data) -> Data:
[wrap_obj[sort], jp.zeros(2 * m.nwrap - count, dtype=int)]
).reshape((m.nwrap, 2))
return d.replace(
ten_length=ten_length,
ten_J=ten_moment,
ten_wrapadr=jp.array(ten_wrapadr, dtype=int),
ten_wrapnum=jp.array(ten_wrapnum, dtype=int),
wrap_xpos=wrap_xpos,
wrap_obj=jp.array(wrap_obj, dtype=int),
)
return d.tree_replace({
'_impl.ten_length': ten_length,
'_impl.ten_J': ten_moment,
'_impl.ten_wrapadr': jp.array(ten_wrapadr, dtype=int),
'_impl.ten_wrapnum': jp.array(ten_wrapnum, dtype=int),
'_impl.wrap_xpos': wrap_xpos,
'_impl.wrap_obj': jp.array(wrap_obj, dtype=int),
})
def _site_dof_mask(m: Model) -> np.ndarray:
@@ -1052,6 +1085,9 @@ def _site_dof_mask(m: Model) -> np.ndarray:
def transmission(m: Model, d: Data) -> Data:
"""Computes actuator/transmission lengths and moments."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('transmission requires JAX backend implementation.')
if not m.nu:
return d
@@ -1111,8 +1147,8 @@ def transmission(m: Model, d: Data) -> Data:
wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
moment = jac @ wrench
elif trntype == TrnType.TENDON:
length = d.ten_length[trnid[0]] * gear[:1]
moment = d.ten_J[trnid[0]] * gear[0]
length = d._impl.ten_length[trnid[0]] * gear[:1]
moment = d._impl.ten_J[trnid[0]] * gear[0]
else:
raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
@@ -1144,5 +1180,7 @@ def transmission(m: Model, d: Data) -> Data:
length = length.reshape((m.nu,))
moment = moment.reshape((m.nu, m.nv))
d = d.replace(actuator_length=length, actuator_moment=moment)
d = d.tree_replace(
{'_impl.actuator_length': length, '_impl.actuator_moment': moment}
)
return d
mjx/mujoco/mjx/_src/smooth_test.py
+33 -29
View File
@@ -20,7 +20,7 @@ import jax
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import ConeType # pylint: disable=g-importing-member
import numpy as np
# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -78,33 +78,37 @@ class SmoothTest(absltest.TestCase):
# com_pos
dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'subtree_com')
_assert_attr_eq(d, dx, 'cinert')
_assert_attr_eq(d, dx, 'cdof')
_assert_attr_eq(d, dx._impl, 'cinert')
_assert_attr_eq(d, dx._impl, 'cdof')
# camlight
dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'cam_xpos')
_assert_eq(d.cam_xmat.reshape((-1, 3, 3)), dx.cam_xmat, 'cam_xmat')
# crb
dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'crb')
_assert_attr_eq(d, dx, 'qM')
_assert_eq(dx._qM_sparse, np.zeros(0), '_qM_sparse')
_assert_attr_eq(d, dx._impl, 'crb')
_assert_attr_eq(d, dx._impl, 'qM')
# factor_m
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
for i in range(m.nM):
for i in range(m.nC):
qLDLegacy[d.mapM2M[i]] = d.qLD[i]
_assert_eq(qLDLegacy, dx.qLD, 'qLD')
_assert_attr_eq(d, dx, 'qLDiagInv')
_assert_eq(dx._qLD_sparse, np.zeros(0), '_qLD_sparse')
_assert_eq(dx._qLDiagInv_sparse, np.zeros(0), '_qLDiagInv_sparse')
_assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
_assert_attr_eq(d, dx._impl, 'qLDiagInv')
# com_vel
dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'cvel')
_assert_attr_eq(d, dx, 'cdof_dot')
_assert_attr_eq(d, dx._impl, 'cdof_dot')
# rne
dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_bias')
# rne (flg_acc=True)
qfrc_bias = np.zeros(m.nv)
mujoco.mj_rne(m, d, 1, qfrc_bias)
dx = jax.jit(mjx.rne, static_argnums=(2,))(
mx, mjx.put_data(m, d), flg_acc=True
)
_assert_eq(dx.qfrc_bias, qfrc_bias, 'qfrc_bias')
# set dense jacobian for tendon:
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
@@ -115,11 +119,11 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_forward(m, d)
# tendon
dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'ten_J')
_assert_attr_eq(d, dx, 'ten_length')
_assert_attr_eq(d, dx._impl, 'ten_J')
_assert_attr_eq(d, dx._impl, 'ten_length')
# transmission
dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx, 'actuator_length')
_assert_attr_eq(d, dx._impl, 'actuator_length')
# convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv))
@@ -130,7 +134,7 @@ class SmoothTest(absltest.TestCase):
d.moment_rowadr,
d.moment_colind,
)
_assert_eq(moment, dx.actuator_moment, 'actuator_moment')
_assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
def test_disable_gravity(self):
m = mujoco.MjModel.from_xml_string("""
@@ -190,7 +194,7 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_transmission(m, d)
dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx, 'actuator_length')
_assert_attr_eq(d, dx._impl, 'actuator_length')
# convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv))
@@ -201,7 +205,7 @@ class SmoothTest(absltest.TestCase):
d.moment_rowadr,
d.moment_colind,
)
_assert_eq(moment, dx.actuator_moment, 'actuator_moment')
_assert_eq(moment, dx._impl.actuator_moment, 'actuator_moment')
def test_subtree_vel(self):
"""Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel."""
@@ -219,8 +223,8 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_subtreeVel(m, d)
dx = jax.jit(mjx.subtree_vel)(mx, dx)
_assert_attr_eq(d, dx, 'subtree_linvel')
_assert_attr_eq(d, dx, 'subtree_angmom')
_assert_attr_eq(d, dx._impl, 'subtree_linvel')
_assert_attr_eq(d, dx._impl, 'subtree_angmom')
class RnePostConstraintTest(parameterized.TestCase):
@@ -271,9 +275,9 @@ class RnePostConstraintTest(parameterized.TestCase):
mujoco.mj_rnePostConstraint(m, d)
dx = jax.jit(mjx.rne_postconstraint)(mx, dx)
_assert_eq(d.cacc, dx.cacc, 'cacc')
_assert_eq(d.cfrc_ext, dx.cfrc_ext, 'cfrc_ext')
_assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int')
_assert_eq(d.cacc, dx._impl.cacc, 'cacc')
_assert_eq(d.cfrc_ext, dx._impl.cfrc_ext, 'cfrc_ext')
_assert_eq(d.cfrc_int, dx._impl.cfrc_int, 'cfrc_int')
class TendonTest(parameterized.TestCase):
@@ -305,12 +309,12 @@ class TendonTest(parameterized.TestCase):
mujoco.mj_forward(m, d)
dx = jax.jit(mjx.forward)(mx, dx)
_assert_eq(d.ten_length, dx.ten_length, 'ten_length')
_assert_eq(d.ten_J, dx.ten_J, 'ten_J')
_assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum')
_assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr')
_assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj')
_assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos')
_assert_eq(d.ten_length, dx._impl.ten_length, 'ten_length')
_assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
_assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
_assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
_assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
_assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos')
if __name__ == '__main__':
mjx/mujoco/mjx/_src/solver.py
+72 -53
View File
@@ -24,13 +24,15 @@ from mujoco.mjx._src import support
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DataJAX
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import ModelJAX
from mujoco.mjx._src.types import SolverType
# pylint: enable=g-importing-member
class _Context(PyTreeNode):
class Context(PyTreeNode):
"""Data updated during each solver iteration.
Attributes:
@@ -72,25 +74,30 @@ class _Context(PyTreeNode):
h: jax.Array
@classmethod
def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context':
jaref = d.efc_J @ d.qacc - d.efc_aref
def create(cls, m: Model, d: Data, grad: bool = True) -> 'Context':
if not isinstance(d._impl, DataJAX):
raise ValueError(
'Constraint context requires JAX backend implementation.'
)
jaref = d._impl.efc_J @ d.qacc - d._impl.efc_aref
# TODO(robotics-team): determine nv at which sparse mul is faster
ma = support.mul_m(m, d, d.qacc)
nv_0 = jp.zeros(m.nv)
fri = 0.0
if m.opt.cone == ConeType.ELLIPTIC:
friction = d.contact.friction[d.contact.dim > 1]
dim = d.contact.dim[d.contact.dim > 1]
friction = d._impl.contact.friction[d._impl.contact.dim > 1]
dim = d._impl.contact.dim[d._impl.contact.dim > 1]
mu = friction[:, 0] / jp.sqrt(m.opt.impratio)
fri = jp.concatenate((mu[:, None], friction), axis=1)
for condim in (3, 4, 6):
fri = fri.at[dim == condim, condim:].set(0)
ctx = _Context(
ctx = Context(
qacc=d.qacc,
qfrc_constraint=d.qfrc_constraint,
Jaref=jaref,
efc_force=d.efc_force,
efc_force=d._impl.efc_force,
Ma=ma,
grad=nv_0,
Mgrad=nv_0,
@@ -133,7 +140,7 @@ class _LSPoint(PyTreeNode):
cls,
m: Model,
d: Data,
ctx: _Context,
ctx: Context,
alpha: jax.Array,
jv: jax.Array,
quad: jax.Array,
@@ -145,18 +152,20 @@ class _LSPoint(PyTreeNode):
) -> '_LSPoint':
"""Creates a linesearch point with first and second derivatives."""
# roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('LSPoint requires JAX backend implementation.')
cost, deriv_0, deriv_1 = 0.0, 0.0, 0.0
quad_total = quad_gauss
x = ctx.Jaref + alpha * jv
active = (x < 0).at[: d.ne + d.nf].set(True)
active = (x < 0).at[: d._impl.ne + d._impl.nf].set(True)
dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
if (dof_fl.any() or ten_fl.any()) and not (
m.opt.disableflags & DisableBit.FRICTIONLOSS
):
f = d.efc_frictionloss
r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
f = d._impl.efc_frictionloss
r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
rf, z = r * f, jp.zeros_like(f)
linear_neg = (x <= -rf)[:, None]
linear_pos = (x >= rf)[:, None]
@@ -174,13 +183,13 @@ class _LSPoint(PyTreeNode):
middle_zone = (tsqr > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
# quadratic cost for equality, friction, limits, frictionless contacts
dim1 = d.contact.efc_address[d.contact.dim == 1]
nefl = d.ne + d.nf + d.nl
dim1 = d._impl.contact.efc_address[d._impl.contact.dim == 1]
nefl = d._impl.ne + d._impl.nf + d._impl.nl
active = active.at[nefl:].set(False).at[dim1].set(active[dim1])
quad_efld = jax.vmap(jp.multiply)(quad, active)
quad_total += jp.sum(quad_efld, axis=0)
# elliptic bottom zone: quadratic cost
efc_elliptic = d.contact.efc_address[d.contact.dim > 1]
efc_elliptic = d._impl.contact.efc_address[d._impl.contact.dim > 1]
quad_c = jax.vmap(jp.multiply)(quad[efc_elliptic], bottom_zone)
quad_total += jp.sum(quad_c, axis=0)
# elliptic middle zone
@@ -241,7 +250,7 @@ def _while_loop_scan(cond_fun, body_fun, init_val, max_iter):
return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]
def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
def _update_constraint(m: Model, d: Data, ctx: Context) -> Context:
"""Updates constraint force and resulting cost given last solver iteration.
Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
@@ -254,17 +263,20 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
Returns:
context with new constraint force and costs
"""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('_update_constraint requires JAX backend implementation.')
# ne constraints are always active, nf are conditionally active, others are
# non-negative constraints.
active = (ctx.Jaref < 0).at[: d.ne + d.nf].set(True)
active = (ctx.Jaref < 0).at[: d._impl.ne + d._impl.nf].set(True)
floss_force, floss_cost = jp.zeros(d.nefc), 0.0
dof_fl, ten_fl = m.dof_hasfrictionloss, m.tendon_hasfrictionloss
floss_force, floss_cost = jp.zeros(d._impl.nefc), 0.0
dof_fl, ten_fl = m._impl.dof_hasfrictionloss, m._impl.tendon_hasfrictionloss
if (dof_fl.any() or ten_fl.any()) and not (
m.opt.disableflags & DisableBit.FRICTIONLOSS
):
f = d.efc_frictionloss
r = 1.0 / (d.efc_D + (d.efc_D == 0.0) * mujoco.mjMINVAL)
f = d._impl.efc_frictionloss
r = 1.0 / (d._impl.efc_D + (d._impl.efc_D == 0.0) * mujoco.mjMINVAL)
linear_neg = (ctx.Jaref <= -r * f) * (f > 0)
linear_pos = (ctx.Jaref >= r * f) * (f > 0)
active = active & ~linear_neg & ~linear_pos
@@ -274,13 +286,13 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
floss_cost = floss_cost.sum()
if m.opt.cone == ConeType.PYRAMIDAL:
efc_force = d.efc_D * -ctx.Jaref * active + floss_force
cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
dm, u, h = 0.0, 0.0, 0.0
elif m.opt.cone == ConeType.ELLIPTIC:
friction = d.contact.friction[d.contact.dim > 1]
efc_address = d.contact.efc_address[d.contact.dim > 1]
dim = d.contact.dim[d.contact.dim > 1]
friction = d._impl.contact.friction[d._impl.contact.dim > 1]
efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
dim = d._impl.contact.dim[d._impl.contact.dim > 1]
# to prevent out of range append zeros to ctx.Jaref
slice_fn = jax.vmap(
lambda x: jax.lax.dynamic_slice(
@@ -297,12 +309,12 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
adr_i.extend(range(addr, addr + condim))
adr_j.extend([i] * condim)
active = active.at[jp.array(adr_i)].set(bottom_zone[jp.array(adr_j)])
efc_force = d.efc_D * -ctx.Jaref * active + floss_force
cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active)
efc_force = d._impl.efc_D * -ctx.Jaref * active + floss_force
cost = 0.5 * jp.sum(d._impl.efc_D * ctx.Jaref * ctx.Jaref * active)
# middle zone: cone
middle_zone = (t > 0) & (n < (mu * t)) & ((mu * n + t) > 0)
dm = d.efc_D[efc_address] / jp.maximum(
dm = d._impl.efc_D[efc_address] / jp.maximum(
mu * mu * (1 + mu * mu), mujoco.mjMINVAL
)
nmt = n - mu * t
@@ -339,7 +351,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
else:
raise NotImplementedError(f'unsupported cone type: {m.opt.cone}')
qfrc_constraint = d.efc_J.T @ efc_force
qfrc_constraint = d._impl.efc_J.T @ efc_force
gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
ctx = ctx.replace(
qfrc_constraint=qfrc_constraint,
@@ -356,7 +368,7 @@ def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
return ctx
def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
def _update_gradient(m: Model, d: Data, ctx: Context) -> Context:
"""Updates grad and M / grad given latest solver iteration.
Corresponds to CGupdateGradient in mujoco/src/engine/engine_solver.c
@@ -371,6 +383,8 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
Raises:
NotImplementedError: for unsupported solver type
"""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('_update_gradient requires JAX backend implementation.')
grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint
@@ -378,16 +392,16 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
mgrad = smooth.solve_m(m, d, grad)
elif m.opt.solver == SolverType.NEWTON:
if m.opt.cone == ConeType.ELLIPTIC:
cm = jp.diag(d.efc_D * ctx.active)
efc_address = d.contact.efc_address[d.contact.dim > 1]
dim = d.contact.dim[d.contact.dim > 1]
cm = jp.diag(d._impl.efc_D * ctx.active)
efc_address = d._impl.contact.efc_address[d._impl.contact.dim > 1]
dim = d._impl.contact.dim[d._impl.contact.dim > 1]
# set efc of cone H along diagonal
for i, (condim, addr) in enumerate(zip(dim, efc_address)):
h_cone = ctx.h[i, :condim, :condim]
cm = cm.at[addr : addr + condim, addr : addr + condim].add(h_cone)
h = d.efc_J.T @ cm @ d.efc_J
h = d._impl.efc_J.T @ cm @ d._impl.efc_J
else:
h = (d.efc_J.T * d.efc_D * ctx.active) @ d.efc_J
h = (d._impl.efc_J.T * d._impl.efc_D * ctx.active) @ d._impl.efc_J
h = support.full_m(m, d) + h
h_ = jax.scipy.linalg.cho_factor(h)
mgrad = jax.scipy.linalg.cho_solve(h_, grad)
@@ -403,7 +417,7 @@ def _rescale(m: Model, value: jax.Array) -> jax.Array:
return value / (m.stat.meaninertia * max(1, m.nv))
def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
def _linesearch(m: Model, d: Data, ctx: Context) -> Context:
"""Performs a zoom linesearch to find optimal search step size.
Args:
@@ -414,12 +428,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
Returns:
updated context with new qacc, Ma, Jaref
"""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('_lineasearch requires JAX backend implementation.')
smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv)
gtol = m.opt.tolerance * m.opt.ls_tolerance * smag
# compute Mv, Jv
mv = support.mul_m(m, d, ctx.search)
jv = d.efc_J @ ctx.search
jv = d._impl.efc_J @ ctx.search
# prepare quadratics
quad_gauss = jp.stack((
@@ -428,13 +445,15 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
0.5 * jp.dot(ctx.search, mv),
))
quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
quad = (quad * d.efc_D).T
quad = (quad * d._impl.efc_D).T
uu, v0, uv, vv = 0.0, 0.0, 0.0, 0.0
if m.opt.cone == ConeType.ELLIPTIC:
mask = d.contact.dim > 1
mask = d._impl.contact.dim > 1
# complete vector quadratic (for bottom zone)
efc_con, efc_fri = [], []
for condim, addr in zip(d.contact.dim[mask], d.contact.efc_address[mask]):
for condim, addr in zip(
d._impl.contact.dim[mask], d._impl.contact.efc_address[mask]
):
efc_con.extend([addr] * (condim - 1))
efc_fri.extend(range(addr + 1, addr + condim))
quad = quad.at[jp.array(efc_con)].add(quad[jp.array(efc_fri)])
@@ -446,7 +465,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
jp.concatenate((jv, jp.zeros(3))), (x,), (6,)
)
)
efc_elliptic = d.contact.efc_address[mask]
efc_elliptic = d._impl.contact.efc_address[mask]
v = jv_fn(efc_elliptic) * ctx.fri
uu = jp.sum(ctx.u[:, 1:] * ctx.u[:, 1:], axis=1)
v0 = v[:, 0]
@@ -529,7 +548,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
def solve(m: Model, d: Data) -> Data:
"""Finds forces that satisfy constraints using conjugate gradient descent."""
def cond(ctx: _Context) -> jax.Array:
def cond(ctx: Context) -> jax.Array:
improvement = _rescale(m, ctx.prev_cost - ctx.cost)
gradient = _rescale(m, math.norm(ctx.grad))
@@ -539,7 +558,7 @@ def solve(m: Model, d: Data) -> Data:
return ~done
def body(ctx: _Context) -> _Context:
def body(ctx: Context) -> Context:
ctx = _linesearch(m, d, ctx)
prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad # pylint: disable=invalid-name
ctx = _update_constraint(m, d, ctx)
@@ -560,22 +579,22 @@ def solve(m: Model, d: Data) -> Data:
# warmstart:
qacc = d.qacc_smooth
if not m.opt.disableflags & DisableBit.WARMSTART:
warm = _Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
smth = _Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False)
warm = Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
smth = Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False)
qacc = jp.where(warm.cost < smth.cost, d.qacc_warmstart, d.qacc_smooth)
d = d.replace(qacc=qacc)
ctx = _Context.create(m, d)
ctx = Context.create(m, d)
if m.opt.iterations == 1:
ctx = body(ctx)
else:
ctx = jax.lax.while_loop(cond, body, ctx)
d = d.replace(
qacc_warmstart=ctx.qacc,
qacc=ctx.qacc,
qfrc_constraint=ctx.qfrc_constraint,
efc_force=ctx.efc_force,
)
d = d.tree_replace({
'qacc_warmstart': ctx.qacc,
'qfrc_constraint': ctx.qfrc_constraint,
'qacc': ctx.qacc,
'_impl.efc_force': ctx.efc_force,
})
return d
mjx/mujoco/mjx/_src/solver_test.py
+9 -9
View File
@@ -74,7 +74,7 @@ class SolverTest(parameterized.TestCase):
# compare costs
mj_cost = cost(d.qacc)
ctx = solver._Context.create(mjx.put_model(m), mjx.put_data(m, d))
ctx = solver.Context.create(mjx.put_model(m), mjx.put_data(m, d))
mjx_cost = ctx.cost - ctx.gauss
_assert_eq(mj_cost, mjx_cost, 'cost')
@@ -87,8 +87,8 @@ class SolverTest(parameterized.TestCase):
# MJX finds very similar solutions with the newton solver
if solver_ == mujoco.mjtSolver.mjSOL_NEWTON:
nnz = dx.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
nnz = dx._impl.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
_assert_attr_eq(d, dx, 'qfrc_constraint')
_assert_attr_eq(d, dx, 'qacc')
@@ -108,9 +108,9 @@ class SolverTest(parameterized.TestCase):
mujoco.mj_forward(m, d)
mx = mjx.put_model(m)
dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
nnz = dx.efc_J.any(axis=1)
nnz = dx._impl.efc_J.any(axis=1)
# even without warmstart, newton converges quickly
_assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-4)
_assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force', tol=2e-4)
def test_sparse(self):
"""Test solver works with sparse mass matrices."""
@@ -130,8 +130,8 @@ class SolverTest(parameterized.TestCase):
_assert_attr_eq(d, dx, 'qacc')
_assert_attr_eq(d, dx, 'qfrc_constraint')
nnz = dx.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
nnz = dx._impl.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
def test_quad_frictionloss(self):
"""Test a case with quadratic frictionloss constraints."""
@@ -144,8 +144,8 @@ class SolverTest(parameterized.TestCase):
_assert_attr_eq(d, dx, 'qacc')
_assert_attr_eq(d, dx, 'qfrc_constraint')
nnz = dx.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
nnz = dx._impl.efc_J.any(axis=1)
_assert_eq(d.efc_force, dx._impl.efc_force[nnz], 'efc_force')
# TODO(taylorhowell): condim=1 with ConeType.ELLIPTIC
@parameterized.product(condim=(3, 4, 6), cone=tuple(ConeType))
mjx/mujoco/mjx/_src/support.py
+56 -65
View File
@@ -13,6 +13,7 @@
# limitations under the License.
# ==============================================================================
"""Engine support functions."""
from collections.abc import Iterable, Sequence
from typing import Optional, Tuple, Union
@@ -92,7 +93,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
"""Reconstitute dense mass matrix from qM."""
if not is_sparse(m):
return d.qM
return d._impl.qM # pytype: disable=attribute-error
ij = []
for i in range(m.nv):
@@ -103,7 +104,7 @@ def full_m(m: Model, d: Data) -> jax.Array:
i, j = (jp.array(x) for x in zip(*ij))
mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM) # pytype: disable=attribute-error
# also set upper triangular
mat = mat + jp.tril(mat, -1).T
@@ -115,9 +116,9 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
"""Multiply vector by inertia matrix."""
if not is_sparse(m):
return d.qM @ vec
return d._impl.qM @ vec # pytype: disable=attribute-error
diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec # pytype: disable=attribute-error
is_, js, madr_ijs = [], [], []
for i in range(m.nv):
@@ -131,8 +132,8 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j])
out = out.at[j].add(d.qM[madr_ij] * vec[i])
out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j]) # pytype: disable=attribute-error
out = out.at[j].add(d._impl.qM[madr_ij] * vec[i]) # pytype: disable=attribute-error
return out
@@ -147,9 +148,9 @@ def jac(
mask = mask[jp.array(m.dof_bodyid)] > 0
offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof)
jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof) # pytype: disable=attribute-error
jacp = jax.vmap(jp.multiply)(jacp, mask)
jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask)
jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask) # pytype: disable=attribute-error
return jacp, jacr
@@ -295,73 +296,56 @@ class BindModel(object):
self.prefix = ''
ids = []
for spec in specs:
if not spec.name:
raise KeyError(f'cannot bind spec with empty name')
if model.signature != spec.signature:
raise ValueError(
'mjSpec signature does not match mjx.Model signature:'
f' {spec.signature} != {model.signature}'
)
elif spec.id < 0:
raise KeyError(f'invalid id: {spec.id}')
elif isinstance(spec, mujoco.MjsBody):
self.prefix = 'body_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
elif isinstance(spec, mujoco.MjsJoint):
self.prefix = 'jnt_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
elif isinstance(spec, mujoco.MjsGeom):
self.prefix = 'geom_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
elif isinstance(spec, mujoco.MjsSite):
self.prefix = 'site_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
elif isinstance(spec, mujoco.MjsLight):
self.prefix = 'light_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
elif isinstance(spec, mujoco.MjsCamera):
self.prefix = 'cam_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
elif isinstance(spec, mujoco.MjsMesh):
self.prefix = 'mesh_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_MESH, spec.name)
elif isinstance(spec, mujoco.MjsHField):
self.prefix = 'hfield_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_HFIELD, spec.name)
elif isinstance(spec, mujoco.MjsPair):
self.prefix = 'pair_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_PAIR, spec.name)
elif isinstance(spec, mujoco.MjsTendon):
self.prefix = 'tendon_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
elif isinstance(spec, mujoco.MjsActuator):
self.prefix = 'actuator_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
elif isinstance(spec, mujoco.MjsSensor):
self.prefix = 'sensor_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
elif isinstance(spec, mujoco.MjsNumeric):
self.prefix = 'numeric_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_NUMERIC, spec.name)
elif isinstance(spec, mujoco.MjsText):
self.prefix = 'text_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_TEXT, spec.name)
elif isinstance(spec, mujoco.MjsTuple):
self.prefix = 'tuple_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_TUPLE, spec.name)
elif isinstance(spec, mujoco.MjsKey):
self.prefix = 'key_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_KEY, spec.name)
elif isinstance(spec, mujoco.MjsEquality):
self.prefix = 'eq_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
elif isinstance(spec, mujoco.MjsExclude):
self.prefix = 'exclude_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_EXCLUDE, spec.name)
elif isinstance(spec, mujoco.MjsSkin):
self.prefix = 'skin_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_SKIN, spec.name)
elif isinstance(spec, mujoco.MjsMaterial):
self.prefix = 'material_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_MATERIAL, spec.name)
else:
raise ValueError('invalid spec type')
if idx < 0:
raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error
ids.append(idx)
ids.append(spec.id)
if len(ids) == 1:
self.id = ids[0]
else:
@@ -402,42 +386,36 @@ class BindData(object):
self.prefix = ''
ids = []
for spec in specs:
if not spec.name:
raise KeyError(f'cannot bind spec with empty name')
if model.signature != spec.signature:
raise ValueError(
'mjSpec signature does not match mjx.Model signature:'
f' {spec.signature} != {model.signature}'
)
if spec.id < 0:
raise KeyError(f'invalid id: {spec.id}')
elif isinstance(spec, mujoco.MjsBody):
idx = name2id(model, mujoco.mjtObj.mjOBJ_BODY, spec.name)
pass
elif isinstance(spec, mujoco.MjsJoint):
self.prefix = 'jnt_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_JOINT, spec.name)
elif isinstance(spec, mujoco.MjsGeom):
self.prefix = 'geom_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_GEOM, spec.name)
elif isinstance(spec, mujoco.MjsSite):
self.prefix = 'site_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_SITE, spec.name)
elif isinstance(spec, mujoco.MjsLight):
self.prefix = 'light_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_LIGHT, spec.name)
elif isinstance(spec, mujoco.MjsCamera):
self.prefix = 'cam_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, spec.name)
elif isinstance(spec, mujoco.MjsTendon):
self.prefix = 'ten_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_TENDON, spec.name)
elif isinstance(spec, mujoco.MjsActuator):
self.prefix = 'actuator_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, spec.name)
elif isinstance(spec, mujoco.MjsSensor):
self.prefix = 'sensor_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, spec.name)
elif isinstance(spec, mujoco.MjsEquality):
self.prefix = 'eq_'
idx = name2id(model, mujoco.mjtObj.mjOBJ_EQUALITY, spec.name)
else:
raise ValueError('invalid spec type')
if idx < 0:
raise KeyError(f'invalid name: {spec.name}') # pytype: disable=attribute-error
ids.append(idx)
ids.append(spec.id)
if len(ids) == 1:
self.id = ids[0]
else:
@@ -455,11 +433,13 @@ class BindData(object):
return name
else:
raise AttributeError('ctrl is not available for this type')
if name == 'qpos' or name == 'qvel' or name == 'qacc':
if name == 'qpos' or name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
if self.prefix == 'jnt_':
return name
else:
raise AttributeError('qpos, qvel, qacc are not available for this type')
raise AttributeError(
'qpos, qvel, qacc, qfrc are not available for this type'
)
else:
return self.prefix + name
@@ -473,7 +453,9 @@ class BindData(object):
return var[..., idx, :]
def __getattr__(self, name: str):
if name in ('sensordata', 'qpos', 'qvel', 'qacc'):
if name in ('sensordata', 'qpos', 'qvel', 'qacc') or (
name.startswith('qfrc_')
):
adr = num = 0
if name == 'sensordata':
adr = self.model.sensor_adr[self.id]
@@ -482,7 +464,7 @@ class BindData(object):
adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
elif name == 'qvel' or name == 'qacc':
elif name == 'qvel' or name == 'qacc' or name.startswith('qfrc_'):
adr = self.model.jnt_dofadr[self.id]
typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.dof_width() for jt in JointType)
@@ -495,6 +477,8 @@ class BindData(object):
return self._slice(self.__getname(name), slice(adr, adr + num))
else:
return self._slice(self.__getname(name), adr)
elif name in ('mocap_pos', 'mocap_quat'):
return self._slice(self.__getname(name), self.model.body_mocapid[self.id])
return self._slice(self.__getname(name), self.id)
def set(self, name: str, value: jax.Array) -> Data:
@@ -507,7 +491,7 @@ class BindData(object):
iter(value)
except TypeError:
value = [value]
if name in ('qpos', 'qvel', 'qacc'):
if name in ('qpos', 'qvel', 'qacc', 'mocap_pos', 'mocap_quat'):
adr = num = 0
if name == 'qpos':
adr = self.model.jnt_qposadr[self.id]
@@ -517,16 +501,23 @@ class BindData(object):
adr = self.model.jnt_dofadr[self.id]
typ = self.model.jnt_type[self.id]
num = sum((typ == jt) * jt.dof_width() for jt in JointType)
elif name == 'mocap_pos':
adr = self.model.body_mocapid[self.id] * 3
num = np.ones_like(self.id, dtype=int) * 3
elif name == 'mocap_quat':
adr = self.model.body_mocapid[self.id] * 4
num = np.ones_like(self.id, dtype=int) * 4
if not isinstance(self.id, list):
adr = [adr]
num = [num]
elif isinstance(self.id, list):
adr = self.id * dim
adr = (np.array(self.id) * dim).tolist()
num = [dim for _ in range(len(self.id))]
else:
adr = [self.id * dim]
num = [dim]
i = 0
value = jax.numpy.array(value).flatten()
for a, n in zip(adr, num):
shape = array.shape
array = array.flatten().at[a : a + n].set(value[i : i + n]).reshape(shape)
@@ -571,20 +562,20 @@ def contact_force(
m: Model, d: Data, contact_id: int, to_world_frame: bool = False
) -> jax.Array:
"""Extract 6D force:torque for one contact, in contact frame by default."""
efc_address = d.contact.efc_address[contact_id]
condim = d.contact.dim[contact_id]
efc_address = d._impl.contact.efc_address[contact_id] # pytype: disable=attribute-error
condim = d._impl.contact.dim[contact_id] # pytype: disable=attribute-error
if m.opt.cone == ConeType.PYRAMIDAL:
force = _decode_pyramid(
d.efc_force[efc_address:], d.contact.friction[contact_id], condim
d._impl.efc_force[efc_address:], d._impl.contact.friction[contact_id], condim # pytype: disable=attribute-error
)
elif m.opt.cone == ConeType.ELLIPTIC:
force = d.efc_force[efc_address : efc_address + condim]
force = d._impl.efc_force[efc_address : efc_address + condim] # pytype: disable=attribute-error
force = jp.concatenate([force, jp.zeros((6 - condim))])
else:
raise ValueError(f'Unknown cone type: {m.opt.cone}')
if to_world_frame:
force = force.reshape((-1, 3)) @ d.contact.frame[contact_id]
force = force.reshape((-1, 3)) @ d._impl.contact.frame[contact_id] # pytype: disable=attribute-error
force = force.reshape(-1)
return force * (efc_address >= 0)
@@ -595,21 +586,21 @@ def contact_force_dim(
) -> Tuple[jax.Array, np.ndarray]:
"""Extract 6D force:torque for contacts with dimension dim."""
# valid contact and condim indices
idx_dim = (d.contact.efc_address >= 0) & (d.contact.dim == dim)
idx_dim = (d._impl.contact.efc_address >= 0) & (d._impl.contact.dim == dim) # pytype: disable=attribute-error
# contact force from efc
if m.opt.cone == ConeType.PYRAMIDAL:
efc_address = (
d.contact.efc_address[idx_dim, None]
d._impl.contact.efc_address[idx_dim, None] # pytype: disable=attribute-error
+ np.arange(np.where(dim == 1, 1, 2 * (dim - 1)))[None]
)
efc_force = d.efc_force[efc_address]
efc_force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error
force = jax.vmap(_decode_pyramid, in_axes=(0, 0, None))(
efc_force, d.contact.friction[idx_dim], dim
efc_force, d._impl.contact.friction[idx_dim], dim # pytype: disable=attribute-error
)
elif m.opt.cone == ConeType.ELLIPTIC:
efc_address = d.contact.efc_address[idx_dim, None] + np.arange(dim)[None]
force = d.efc_force[efc_address]
efc_address = d._impl.contact.efc_address[idx_dim, None] + np.arange(dim)[None] # pytype: disable=attribute-error
force = d._impl.efc_force[efc_address] # pytype: disable=attribute-error
force = jp.hstack([force, jp.zeros((force.shape[0], 6 - dim))])
else:
raise ValueError(f'Unknown cone type: {m.opt.cone}.')
mjx/mujoco/mjx/_src/support_test.py
+59 -30
View File
@@ -173,18 +173,20 @@ class SupportTest(parameterized.TestCase):
<joint axis="0 0 1" type="slide" name="joint3"/>
<geom size="7 8 9" type="box" name="geom3"/>
</body>
<body pos="100 110 120" name="body4" mocap="true"/>
<body pos="130 140 150" name="body5" mocap="true"/>
</worldbody>
<actuator>
<motor name="actuator1" joint="joint1"/>
<motor name="actuator2" joint="joint2"/>
<motor name="actuator3" joint="joint3"/>
<motor joint="joint1"/>
<motor joint="joint2"/>
<motor joint="joint3"/>
</actuator>
<sensor>
<framepos name="sensor1" objtype="body" objname="body1"/>
<framepos name="sensor2" objtype="body" objname="body2"/>
<framepos name="sensor3" objtype="body" objname="body3"/>
<framepos objtype="body" objname="body1"/>
<framepos objtype="body" objname="body2"/>
<framepos objtype="body" objname="body3"/>
</sensor>
</mujoco>
"""
@@ -247,6 +249,11 @@ class SupportTest(parameterized.TestCase):
dx.bind(mx, s.joints[i]).qacc,
d.qacc[m.jnt_dofadr[i]:m.jnt_dofadr[i] + dofnum[i]], decimal=6
)
np.testing.assert_array_almost_equal(
dx.bind(mx, s.joints[i]).qfrc_actuator,
d.qfrc_actuator[m.jnt_dofadr[i] : m.jnt_dofadr[i] + dofnum[i]],
decimal=6,
)
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
for i in range(m.nu):
@@ -308,30 +315,49 @@ class SupportTest(parameterized.TestCase):
np.testing.assert_array_equal(
dx7.bind(mx, body).xfrc_applied, [0, 0, 0, 0, 0, 0]
)
dx10 = dx.bind(mx, s.bodies[0:2]).set(
'xfrc_applied',
np.array([np.array([1, 1, 1, 1, 1, 1]), np.array([2, 2, 2, 2, 2, 2])]),
)
np.testing.assert_array_equal(
dx10.bind(mx, s.bodies[0]).xfrc_applied, [1, 1, 1, 1, 1, 1]
)
np.testing.assert_array_equal(
dx10.bind(mx, s.bodies[1]).xfrc_applied, [2, 2, 2, 2, 2, 2]
)
dx11 = dx.bind(mx, s.bodies[-1]).set(
'mocap_pos',
[1, 2, 3],
)
np.testing.assert_array_equal(
dx11.bind(mx, s.bodies[-2]).mocap_pos, [100, 110, 120]
)
np.testing.assert_array_equal(
dx11.bind(mx, s.bodies[-1]).mocap_pos, [1, 2, 3]
)
# test invalid name
with self.assertRaises(
AttributeError, msg='ctrl is not available for this type'
# test attribute and type mismatches
with self.assertRaisesRegex(
AttributeError, 'ctrl is not available for this type'
):
print(dx.bind(mx, s.geoms).ctrl)
with self.assertRaises(
KeyError, msg='actuator_actuator_ctrl'
):
with self.assertRaises(KeyError):
print(dx.bind(mx, s.actuators).actuator_ctrl)
with self.assertRaises(
AttributeError, msg='actuator_actuator_ctrl'
with self.assertRaisesRegex(
AttributeError,
"'Data' object has no attribute 'actuator_actuator_ctrl'",
):
print(dx.bind(mx, s.actuators).set('actuator_ctrl', [1, 2, 3]))
with self.assertRaises(
AttributeError, msg='qpos, qvel, qacc are not available for this type'
with self.assertRaisesRegex(
AttributeError, 'qpos, qvel, qacc, qfrc are not available for this type'
):
print(dx.bind(mx, s.geoms).qpos)
with self.assertRaises(KeyError, msg='invalid name: invalid_actuator_name'):
s.actuators[0].name = 'invalid_actuator_name'
print(dx.bind(mx, s.actuators).set('ctrl', [1, 2, 3]))
with self.assertRaises(KeyError, msg='invalid name: invalid_geom_name'):
s.geoms[0].name = 'invalid_geom_name'
print(mx.bind(s.geoms).pos)
# test that modified names do not raise an error
s.actuators[0].name = 'modified_actuator_name'
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
s.geoms[0].name = 'modified_geom_name'
np.testing.assert_array_equal(mx.bind(s.geoms[0]).pos, m.geom_pos[0, :])
# test batched data
batch_size = 16
@@ -343,12 +369,15 @@ class SupportTest(parameterized.TestCase):
vdx.bind(mx, s.bodies[i]).xpos, [d.xpos[i, :]] * batch_size
)
# test emtpy name
# test that adding a body requires recompilation
s.worldbody.add_body()
m = s.compile()
mx = mjx.put_model(m)
with self.assertRaises(KeyError, msg='cannot bind spec with empty name'):
with self.assertRaises(ValueError) as e:
mx.bind(s.bodies)
self.assertEqual(
str(e.exception),
'mjSpec signature does not match mjx.Model signature:'
' 15297169659434471387 != 2785811613804955188',
)
_CONTACTS = """
<mujoco>
@@ -383,7 +412,7 @@ class SupportTest(parameterized.TestCase):
# map MJX contacts to MJ ones
def _find(g):
val = (g == dx.contact.geom).sum(axis=1)
val = (g == dx._impl.contact.geom).sum(axis=1)
return np.where(val == 2)[0][0]
contact_id_map = {i: _find(d.contact.geom[i]) for i in range(d.ncon)}
@@ -397,7 +426,7 @@ class SupportTest(parameterized.TestCase):
np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
# check for zeros after first condim elements
condim = dx.contact.dim[j]
condim = dx._impl.contact.dim[j]
if condim < 6:
np.testing.assert_allclose(force[condim:], 0, rtol=1e-5, atol=1e-5)
@@ -410,8 +439,8 @@ class SupportTest(parameterized.TestCase):
),
)(mx, dx, j, True)
# back to contact frame
force = force.at[:3].set(dx.contact.frame[j] @ force[:3])
force = force.at[3:].set(dx.contact.frame[j] @ force[3:])
force = force.at[:3].set(dx._impl.contact.frame[j] @ force[:3])
force = force.at[3:].set(dx._impl.contact.frame[j] @ force[3:])
np.testing.assert_allclose(result, force, rtol=1e-5, atol=2)
def test_wrap_inside(self):
mjx/mujoco/mjx/_src/test_util.py
+10 -10
View File
@@ -106,20 +106,20 @@ def benchmark(
def efc_order(m: mujoco.MjModel, d: mujoco.MjData, dx: Data) -> np.ndarray:
"""Returns a sort order such that dx.efc_*[order][:d.nefc] == d.efc_*."""
"""Returns a sort order such that dx.efc_*[order][:d._impl.nefc] == d.efc_*.""" # pytype: disable=attribute-error
# reorder efc rows to skip inactive constraints and match contact order
efl = dx.ne + dx.nf + dx.nl
efl = dx._impl.ne + dx._impl.nf + dx._impl.nl # pytype: disable=attribute-error
order = np.arange(efl)
order[(dx.efc_J[:efl] == 0).all(axis=1)] = 2**16 # move empty rows to end
for i in range(dx.ncon):
num_rows = dx.contact.dim[i]
if dx.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL:
num_rows = (dx.contact.dim[i] - 1) * 2
if dx.contact.dist[i] > 0: # move empty contacts to end
order[(dx._impl.efc_J[:efl] == 0).all(axis=1)] = 2**16 # move empty rows to end # pytype: disable=attribute-error
for i in range(dx._impl.ncon): # pytype: disable=attribute-error
num_rows = dx._impl.contact.dim[i] # pytype: disable=attribute-error
if dx._impl.contact.dim[i] > 1 and m.opt.cone == mujoco.mjtCone.mjCONE_PYRAMIDAL: # pytype: disable=attribute-error
num_rows = (dx._impl.contact.dim[i] - 1) * 2 # pytype: disable=attribute-error
if dx._impl.contact.dist[i] > 0: # move empty contacts to end # pytype: disable=attribute-error
order = np.append(order, np.repeat(2**16, num_rows))
continue
contact_match = (d.contact.geom == dx.contact.geom[i]).all(axis=-1)
contact_match &= (d.contact.pos == dx.contact.pos[i]).all(axis=-1)
contact_match = (d.contact.geom == dx._impl.contact.geom[i]).all(axis=-1) # pytype: disable=attribute-error
contact_match &= (d.contact.pos == dx._impl.contact.pos[i]).all(axis=-1) # pytype: disable=attribute-error
assert contact_match.any(), f'contact {i} not found'
contact_id = np.nonzero(contact_match)[0][0]
order = np.append(order, np.repeat(efl + contact_id, num_rows))
mjx/mujoco/mjx/_src/types.py
+352 -699
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/test_data/actuator/tendon_force_clamp.xml
+55
View File
@@ -0,0 +1,55 @@
<mujoco model="fixed_site">
<worldbody>
<body>
<joint name="joint0" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site0" pos="0.25 0 0.1" size="0.025"/>
<body pos="0.5 0 0">
<joint name="joint1" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site1" pos="0.25 0 0.1" size="0.025"/>
<body pos="0.5 0 0">
<joint name="joint2" type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
<site name="site2" pos="0.25 0 0.1" size="0.025"/>
</body>
</body>
</body>
</worldbody>
<tendon>
<spatial name="spatial0" width="0.0125">
<site site="site0"/>
<site site="site1"/>
<site site="site2"/>
</spatial>
<spatial name="spatial1" width="0.0125" actuatorfrclimited="true" actuatorfrcrange="-1 1">
<site site="site1"/>
<site site="site2"/>
</spatial>
<fixed name="fixed0" actuatorfrclimited="true" actuatorfrcrange="-10 10">
<joint joint="joint0" coef=".1"/>
<joint joint="joint1" coef=".2"/>
<joint joint="joint2" coef=".3"/>
</fixed>
<fixed name="fixed1" actuatorfrclimited="true" actuatorfrcrange="0 1">
<joint joint="joint0" coef=".1"/>
<joint joint="joint2" coef=".3"/>
</fixed>
</tendon>
<actuator>
<motor tendon="spatial0"/>
<motor tendon="spatial0"/>
<motor tendon="spatial0"/>
<motor tendon="spatial1"/>
<motor tendon="spatial1"/>
<motor tendon="fixed0"/>
<motor tendon="fixed1"/>
<motor tendon="fixed1"/>
</actuator>
<sensor>
<tendonactuatorfrc tendon="spatial0"/>
<tendonactuatorfrc tendon="spatial1"/>
<tendonactuatorfrc tendon="fixed0"/>
<tendonactuatorfrc tendon="fixed1"/>
</sensor>
</mujoco>
mjx/pyproject.toml
+4 -4
View File
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
[project]
name="mujoco-mjx"
version = "3.3.1"
version = "3.3.3"
authors = [
{name = "Google DeepMind", email = "mujoco@deepmind.com"},
]
@@ -30,7 +30,7 @@ dependencies = [
"etils[epath]",
"jax",
"jaxlib",
"mujoco>=3.3.1.dev0",
"mujoco>=3.3.3.dev0",
"scipy",
"trimesh",
]
@@ -41,9 +41,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.3.1"
Documentation = "https://mujoco.readthedocs.io/en/3.3.3"
Repository = "https://github.com/google-deepmind/mujoco/tree/main/mjx"
Changelog = "https://mujoco.readthedocs.io/en/3.3.1/changelog.html"
Changelog = "https://mujoco.readthedocs.io/en/3.3.3/changelog.html"
[tool.isort]
force_single_line = true
mjx/requirements.txt
+3 -2
View File
@@ -76,8 +76,9 @@ scipy==1.13.1; python_version == '3.9' \
--hash=sha256:d533654b7d221a6a97304ab63c41c96473ff04459e404b83275b60aa8f4b7004 \
--hash=sha256:8335549ebbca860c52bf3d02f80784e91a004b71b059e3eea9678ba994796a24 \
--hash=sha256:436bbb42a94a8aeef855d755ce5a465479c721e9d684de76bf61a62e7c2b81d5
setuptools==75.5.0 \
--hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
setuptools==78.1.1 \
--hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
--hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
trimesh==4.5.2 \
--hash=sha256:2e50f3a7fd135c3045da887a1b9f91230528f3ce11d2ec1ba44750d82d6b4f73
wheel==0.45.0 \
model/flex/bunny.xml
+1 -1
View File
@@ -31,7 +31,7 @@
</body>
<flexcomp type="mesh" file="bunny.obj" pos="0 0 .1" dim="2" euler="90 0 0"
radius=".001" rgba="0 .7 .7 1" mass=".05" name="softbody" dof="trilinear">
<elasticity young="1e3" poisson="0.1" damping="0.001"/>
<elasticity young="1e3" poisson="0.1" damping="0.001" elastic2d="stretch"/>
<contact selfcollide="none" internal="false"/>
</flexcomp>
</worldbody>
model/flex/bunny_with_uv.xml
+1 -1
View File
@@ -37,7 +37,7 @@
</body>
<flexcomp type="mesh" file="bunny_with_uv.obj" pos="0 0 .1" dim="2" euler="90 0 0"
radius=".001" material="matsponge" mass=".05" name="softbody" dof="trilinear">
<elasticity young="1e3" poisson="0.1" damping="0.001"/>
<elasticity young="1e3" poisson="0.1" damping="0.001" elastic2d="stretch"/>
<contact selfcollide="none" internal="false"/>
</flexcomp>
</worldbody>
model/plugin/elasticity/flag_flex.xml → model/flex/flag.xml
+1 -10
View File
@@ -26,10 +26,6 @@
<option wind="5 5 0" density="10" solver="CG" tolerance="1e-6">
</option>
<extension>
<plugin plugin="mujoco.elasticity.shell"/>
</extension>
<worldbody>
<geom name="floor" type="plane" size="0 0 .1"/>
<light diffuse=".6 .6 .6" specular="0.2 0.2 0.2" pos="0 0 4" dir="0 0 -1"/>
@@ -37,12 +33,7 @@
<flexcomp type="grid" count="9 19 1" spacing=".05 .05 .05" mass="10"
name="flag" radius="0.001">
<edge equality="true" damping="0.001"/>
<plugin plugin="mujoco.elasticity.shell">
<config key="poisson" value="0"/>
<config key="thickness" value="1e-2"/>
<!--Units are in Pa (SI)-->
<config key="young" value="3e6"/>
</plugin>
<elasticity poisson="0" thickness="1e-2" young="3e6" elastic2d="none"/>
</flexcomp>
</body>
</worldbody>
model/plugin/elasticity/mannequin.xml → model/flex/mannequin.xml
View File
model/plugin/elasticity/pancake_flex.xml → model/flex/pancake.xml
+1 -10
View File
@@ -16,10 +16,6 @@
<mujoco model="Pancake">
<include file="scene.xml"/>
<extension>
<plugin plugin="mujoco.elasticity.shell"/>
</extension>
<option timestep="0.001" solver="CG" tolerance="1e-6" integrator="implicitfast"/>
<size memory="10M"/>
@@ -39,12 +35,7 @@
radius=".01" mass=".5" name="plate" dim="2">
<contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
<edge equality="true" damping="10"/>
<plugin plugin="mujoco.elasticity.shell">
<config key="poisson" value="0"/>
<config key="thickness" value="8e-3"/>
<!--Units are in Pa (SI)-->
<config key="young" value="3e5"/>
</plugin>
<elasticity poisson="0" thickness="8e-3" young="3e5" elastic2d="bend"/>
</flexcomp>
</worldbody>
</mujoco>
model/plugin/elasticity/plate_flex.xml → model/flex/plate.xml
+1 -10
View File
@@ -16,10 +16,6 @@
<mujoco model="Plate">
<include file="scene.xml"/>
<extension>
<plugin plugin="mujoco.elasticity.shell"/>
</extension>
<option solver="CG" tolerance="1e-6" integrator="implicitfast"/>
<size memory="10M"/>
@@ -42,12 +38,7 @@
radius=".001" mass="100" name="plate">
<contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
<edge equality="true" damping="0.001"/>
<plugin plugin="mujoco.elasticity.shell">
<config key="poisson" value="0"/>
<config key="thickness" value="1e-2"/>
<!--Units are in Pa (SI)-->
<config key="young" value="3e7"/>
</plugin>
<elasticity poisson="0" thickness="1e-2" young="3e7" elastic2d="bend"/>
</flexcomp>
</worldbody>
</mujoco>
model/plugin/elasticity/poncho_flex.xml → model/flex/poncho.xml
+1 -10
View File
@@ -19,10 +19,6 @@
<flag energy="enable"/>
</option>
<extension>
<plugin plugin="mujoco.elasticity.shell"/>
</extension>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
@@ -1418,13 +1414,8 @@
398 399 418
398 376 378">
<edge equality="true" damping="0.1"/>
<elasticity young="3e5" poisson="0" thickness="8e-3" elastic2d="bend"/>
<contact solref="0.003"/>
<plugin plugin="mujoco.elasticity.shell">
<config key="poisson" value="0"/>
<config key="thickness" value="8e-3"/>
<!--Units are in Pa (SI)-->
<config key="young" value="3e5"/>
</plugin>
</flexcomp>
</worldbody>
model/flex/poncho_vertcollide.xml
+1423
View File
File diff suppressed because it is too large Load Diff
model/flex/trampoline.xml
+1 -1
View File
@@ -39,7 +39,7 @@
radius=".001" mass="10" name="plate" dim="2">
<contact condim="3" solref="0.01 1" solimp=".95 .99 .0001"/>
<edge equality="false" damping="10"/>
<elasticity young="3e5" poisson="0" thickness="1e-2"/>
<elasticity young="3e5" poisson="0" thickness="1e-2" elastic2d="stretch"/>
<pin id="0 15 240 255"/>
</flexcomp>
</worldbody>
model/plugin/sensor/touch_stress.xml
+64
View File
@@ -0,0 +1,64 @@
<mujoco model="touchtest">
<compiler autolimits="true"/>
<extension>
<plugin plugin="mujoco.sdf.gear">
<instance name="gear">
<config key="alpha" value="0"/>
</instance>
</plugin>
<plugin plugin="mujoco.sensor.touch_stress">
<instance name="touch_stress">
<config key="size" value="37 37"/>
<config key="fov" value="45 45"/>
<config key="gamma" value="0"/>
<config key="nchannel" value="3"/>
</instance>
</plugin>
</extension>
<asset>
<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4"
width="300" height="300" mark="edge" markrgb=".2 .3 .4"/>
<material name="grid" texture="grid" texrepeat="3 1" texuniform="true"/>
<mesh name="gear">
<plugin instance="gear"/>
</mesh>
</asset>
<visual>
<headlight ambient=".7 .7 .7" diffuse=".2 .2 .2" specular="0.1 0.1 0.1"/>
<map znear="0.01"/>
<scale contactwidth=".02" contactheight=".5"/>
</visual>
<default>
<geom friction="0.4" solimp="0 0.95 0.02"/>
</default>
<statistic center="0 0 1" extent="1" meansize=".1"/>
<worldbody>
<light pos="1 0 .3" dir="-1 0 -.3"/>
<light pos="-1 0 .3" dir="1 0 -.3"/>
<geom name="floor" pos="0 0 -0.01" type="plane" size="3 3 .01"/>
<geom type="sdf" name="gear" mesh="gear" rgba="0.4 0.4 0.4 1" pos="-1.3 0 0">
<plugin instance="gear"/>
</geom>
<body name="ball" pos="0 0 1">
<joint name="x" type="slide" axis="1 0 0" damping="1"/>
<joint name="y" type="slide" axis="0 1 0" damping="1"/>
<joint name="z" type="slide" axis="0 0 1"/>
<joint name="rx" axis="1 0 0" springdamper="0.2 1"/>
<joint name="ry" axis="0 1 0" springdamper="0.2 1"/>
<geom type="sphere" size=".3" mass="0.1" rgba=".5 .5 .5 .3"/>
<site name="touch"/>
</body>
</worldbody>
<sensor>
<plugin instance="touch_stress" objtype="site" objname="touch"/>
</sensor>
</mujoco>
plugin/elasticity/CMakeLists.txt
-2
View File
@@ -22,8 +22,6 @@ set(MUJOCO_ELASTICITY_SRCS
elasticity.cc
elasticity.h
register.cc
shell.cc
shell.h
)
add_library(elasticity SHARED)
plugin/elasticity/README.md
-12
View File
@@ -18,15 +18,3 @@ Parameters:
- `bend` [Pa]: bending stiffness.
- `flat` [bool]: if true, the stress-equilibrium configuration is that of a straight cable; if false or unspecified, it is the configuration defined in the XML.
- `vmax` [N/m^2]: If greater than zero, the cable is colored using mechanical stresses; the value represent the maximum stress in the color scale.
### Shell
Implemented in [shell.cc](shell.cc).
The shell plugin discretizes an inextensible 2D continuum. It is intended to simulate the bending of plates where the stretching is negligible compared to other deformation modes.
Parameters:
- `young` [Pa]: Young's modulus.
- `poisson` [Pa]: Poisson's ratio; if 0, then the material only opposed shear deformations; if near 0.5, then the material is nearly incompressible (rubber-like).
- `thickness` [m]: shell thickness, used to scale the bending stiffness.
plugin/elasticity/register.cc
-2
View File
@@ -14,13 +14,11 @@
#include <mujoco/mjplugin.h>
#include "cable.h"
#include "shell.h"
namespace mujoco::plugin::elasticity {
mjPLUGIN_LIB_INIT {
Cable::RegisterPlugin();
Shell::RegisterPlugin();
}
} // namespace mujoco::plugin::elasticity
plugin/elasticity/shell.cc
-271
View File
@@ -1,271 +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.
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mujoco.h>
#include "elasticity.h"
#include "shell.h"
namespace mujoco::plugin::elasticity {
namespace {
// local tetrahedron numbering
constexpr int kNumEdges = Stencil2D::kNumEdges;
constexpr int kNumVerts = Stencil2D::kNumVerts;
constexpr int edge[kNumEdges][2] = {{1, 2}, {2, 0}, {0, 1}};
// cotangent between two edges
mjtNum cot(mjtNum* x, int v0, int v1, int v2) {
mjtNum normal[3];
mjtNum edge1[3];
mjtNum edge2[3];
mju_sub3(edge1, x+3*v1, x+3*v0);
mju_sub3(edge2, x+3*v2, x+3*v0);
mju_cross(normal, edge1, edge2);
return mju_dot3(edge1, edge2) / mju_norm3(normal);
}
// area of a triangle
mjtNum ComputeVolume(const mjtNum* x, const int v[kNumVerts]) {
mjtNum normal[3];
mjtNum edge1[3];
mjtNum edge2[3];
mju_sub3(edge1, x+3*v[1], x+3*v[0]);
mju_sub3(edge2, x+3*v[2], x+3*v[0]);
mju_cross(normal, edge1, edge2);
return mju_norm3(normal) / 2;
}
} // namespace
// factory function
std::optional<Shell> Shell::Create(const mjModel* m, mjData* d, int instance) {
if (CheckAttr("face", m, instance) &&
CheckAttr("edge", m, instance) &&
CheckAttr("poisson", m, instance) &&
CheckAttr("young", m, instance) &&
CheckAttr("thickness", m, instance)) {
mjtNum nu = strtod(mj_getPluginConfig(m, instance, "poisson"), nullptr);
mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
mjtNum thick =
strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
std::vector<int> face, edge;
String2Vector(mj_getPluginConfig(m, instance, "face"), face);
String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
return Shell(m, d, instance, nu, E, thick, face, edge);
} else {
mju_warning("Invalid parameter specification in shell plugin");
return std::nullopt;
}
}
// create map from triangles to vertices and edges and from edges to vertices
void Shell::CreateStencils(const std::vector<int>& simplex,
const std::vector<int>& edgeidx) {
// populate stencil
nt = simplex.size() / kNumVerts;
elements.resize(nt);
for (int t = 0; t < nt; t++) {
for (int v = 0; v < kNumVerts; v++) {
elements[t].vertices[v] = simplex[kNumVerts*t+v];
}
}
// map from edge vertices to their index in `edges` vector
std::unordered_map<std::pair<int, int>, int, PairHash> edge_indices;
// loop over all triangles
for (int t = 0; t < nt; t++) {
int* v = elements[t].vertices;
// compute edges to vertices map for fast computations
for (int e = 0; e < kNumEdges; e++) {
auto pair = std::pair(
std::min(v[edge[e][0]], v[edge[e][1]]),
std::max(v[edge[e][0]], v[edge[e][1]])
);
// if edge is already present in the vector only store its index
auto [it, inserted] = edge_indices.insert({pair, ne});
if (inserted) {
StencilFlap flap;
flap.vertices[0] = v[edge[e][0]];
flap.vertices[1] = v[edge[e][1]];
flap.vertices[2] = v[(edge[e][1]+1) % 3];
flap.vertices[3] = -1;
flaps.push_back(flap);
elements[t].edges[e] = ne++;
} else {
elements[t].edges[e] = it->second;
flaps[it->second].vertices[3] = v[(edge[e][1]+1) % 3];
}
if (!edgeidx.empty()) {
assert(elements[t].edges[e] == edgeidx[kNumEdges*t+e]);
}
}
}
}
// plugin constructor
Shell::Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum thick, const std::vector<int>& face,
const std::vector<int>& edgeidx)
: thickness(thick) {
// count plugin bodies
nv = ne = 0;
for (int i = 1; i < m->nbody; i++) {
if (m->body_plugin[i] == instance) {
if (!nv++) {
i0 = i;
}
}
}
// generate triangles from the vertices
CreateStencils(face, edgeidx);
// material parameters
mjtNum mu = E / (2*(1+nu));
// loop over all triangles
for (int t = 0; t < nt; t++) {
int* v = elements[t].vertices;
for (int i = 0; i < kNumVerts; i++) {
if (m->body_plugin[i0+v[i]] != instance) {
mju_error("This body does not have the requested plugin instance");
}
}
}
// allocate array
position.assign(nv*3, 0);
bending.assign(ne*16, 0);
// store previous positions
mju_copy(position.data(), m->body_pos+3*i0, 3*nv);
// assemble bending Hessian
for (int e = 0; e < ne; e++) {
int* v = flaps[e].vertices;
int vadj[3] = {v[1], v[0], v[3]};
if (v[3]== -1) {
// skip boundary edges
continue;
}
// cotangent operator from Wardetzky at al., "Discrete Quadratic Curvature
// Energies", https://cims.nyu.edu/gcl/papers/wardetzky2007dqb.pdf
mjtNum a01 = cot(m->body_pos+3*i0, v[0], v[1], v[2]);
mjtNum a02 = cot(m->body_pos+3*i0, v[0], v[3], v[1]);
mjtNum a03 = cot(m->body_pos+3*i0, v[1], v[2], v[0]);
mjtNum a04 = cot(m->body_pos+3*i0, v[1], v[0], v[3]);
mjtNum c[4] = {a03 + a04, a01 + a02, -(a01 + a03), -(a02 + a04)};
mjtNum volume = ComputeVolume(m->body_pos+3*i0, v) +
ComputeVolume(m->body_pos+3*i0, vadj);
for (int v1 = 0; v1 < StencilFlap::kNumVerts; v1++) {
for (int v2 = 0; v2 < StencilFlap::kNumVerts; v2++) {
bending[16 * e + 4 * v1 + v2] +=
1.5 * c[v1] * c[v2] / volume * mu * pow(thickness, 3) / 12;
}
}
}
}
void Shell::Compute(const mjModel* m, mjData* d, int instance) {
for (int e = 0; e < ne; e++) {
int* v = flaps[e].vertices;
mjtNum force[12] = {0};
if (v[3] == -1) {
// skip boundary edges
continue;
}
for (int i = 0; i < StencilFlap::kNumVerts; i++) {
for (int j = 0; j < StencilFlap::kNumVerts; j++) {
for (int x = 0; x < 3; x++) {
force[3*i+x] += bending[16*e+4*i+j] * d->xpos[3*(i0+v[j])+x];
}
}
}
// update stored positions
mju_copy(position.data(), d->xpos+3*i0, 3*nv);
// insert into global force
for (int i = 0; i < StencilFlap::kNumVerts; i++) {
for (int x = 0; x < 3; x++) {
d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
}
}
}
}
void Shell::RegisterPlugin() {
mjpPlugin plugin;
mjp_defaultPlugin(&plugin);
plugin.name = "mujoco.elasticity.shell";
plugin.capabilityflags |= mjPLUGIN_PASSIVE;
const char* attributes[] = {"face", "edge", "young",
"poisson", "thickness", "damping"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
auto elasticity_or_null = Shell::Create(m, d, instance);
if (!elasticity_or_null.has_value()) {
return -1;
}
d->plugin_data[instance] = reinterpret_cast<uintptr_t>(
new Shell(std::move(*elasticity_or_null)));
return 0;
};
plugin.destroy = +[](mjData* d, int instance) {
delete reinterpret_cast<Shell*>(d->plugin_data[instance]);
d->plugin_data[instance] = 0;
};
plugin.compute = +[](const mjModel* m, mjData* d, int instance, int type) {
auto* elasticity = reinterpret_cast<Shell*>(d->plugin_data[instance]);
elasticity->Compute(m, d, instance);
};
mjp_registerPlugin(&plugin);
}
} // namespace mujoco::plugin::elasticity
plugin/elasticity/shell.h
-74
View File
@@ -1,74 +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_PLUGIN_ELASTICITY_SHELL_H_
#define MUJOCO_PLUGIN_ELASTICITY_SHELL_H_
#include <optional>
#include <vector>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include "elasticity.h"
namespace mujoco::plugin::elasticity {
struct StencilFlap {
static constexpr int kNumVerts = 4;
int vertices[kNumVerts];
};
class Shell {
public:
// Returns a new Shell instance or nullopt on failure.
static std::optional<Shell> Create(const mjModel* m, mjData* d,
int instance);
Shell(Shell&&) = default;
Shell& operator=(Shell&& other) = default;
void Compute(const mjModel* m, mjData* d, int instance);
static void RegisterPlugin();
int i0; // index of first body
int nc; // number of quads in the grid
int nv; // number of vertices (bodies) in the Shell
int nt; // number of area elements (triangles)
int ne; // number of edges in the Shell
// connectivity info for mapping tetrahedra to edges and vertices
std::vector<Stencil2D> elements; // triangles (nt x 6)
std::vector<StencilFlap> flaps; // adjacent triangles (ne x 4)
// precomputed quantities
std::vector<mjtNum> position; // previous-step positions (nv x 3)
std::vector<mjtNum> bending; // bending Hessian (ne x 16)
mjtNum thickness;
private:
Shell(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum thick, const std::vector<int>& face,
const std::vector<int>& edgeidx);
void CreateStencils(const std::vector<int>& simplex,
const std::vector<int>& edgeidx);
};
} // namespace mujoco::plugin::elasticity
#endif // MUJOCO_PLUGIN_ELASTICITY_SHELL_H_
plugin/sdf/CMakeLists.txt
+6 -6
View File
@@ -34,19 +34,19 @@ set(MUJOCO_SDF_SRCS
torus.h
)
add_library(sdf SHARED)
target_sources(sdf PRIVATE ${MUJOCO_SDF_SRCS})
target_include_directories(sdf PRIVATE ${MUJOCO_SDF_INCLUDE})
target_link_libraries(sdf PRIVATE mujoco SdfLib)
add_library(sdf_plugin SHARED)
target_sources(sdf_plugin PRIVATE ${MUJOCO_SDF_SRCS})
target_include_directories(sdf_plugin PRIVATE ${MUJOCO_SDF_INCLUDE})
target_link_libraries(sdf_plugin PRIVATE mujoco SdfLib)
target_compile_options(
sdf
sdf_plugin
PRIVATE ${AVX_COMPILE_OPTIONS}
${MUJOCO_MACOS_COMPILE_OPTIONS}
${EXTRA_COMPILE_OPTIONS}
${MUJOCO_CXX_FLAGS}
)
target_link_options(
sdf
sdf_plugin
PRIVATE
${MUJOCO_MACOS_LINK_OPTIONS}
${EXTRA_LINK_OPTIONS}
plugin/sensor/CMakeLists.txt
+3 -1
View File
@@ -18,9 +18,11 @@ set(MUJOCO_SENSOR_INCLUDE
)
set(MUJOCO_SENSOR_SRCS
sensor.cc
register.cc
touch_grid.h
touch_grid.cc
touch_stress.h
touch_stress.cc
)
add_library(sensor SHARED)
plugin/sensor/README.md
+65 -1
View File
@@ -9,6 +9,9 @@ plugins](https://mujoco.readthedocs.io/en/latest/programming/extension.html#engi
- [Illustration of foveal deformation](#illustration-of-foveal-deformation)
- [Illustration combining resolution, fields-of-view and foveal deformation](#illustration-combining-resolution-fields-of-view-and-foveal-deformation)
- [Touch Stress](#touch-stress)
- [Example model with analytical SDF](#example-model-with-analytical-sdf)
## [Touch Grid](touch_grid.h)
This sensor aggregates contact forces into "taxels": a rectangular array of pixel-like elements.
@@ -23,7 +26,7 @@ The output of the sensor is a stack of 1 to 6 "touch images" corresponding to fo
and torques in the frame of the sensor. Forces and torques are in the in [z, x,
y] order, corresponding to the ordering in contact frames: [normal, tangent,
tangent] and [torsional, rolling, rolling]. Each "taxel" corresponds to an angular bin
in spherical coordinates, and aggregates all the forces occuring inside this bin, which occur
in spherical coordinates, and aggregates all the forces occurring inside this bin, which occur
between the body in which the sensor's site is defined and any other body.
The sensor is parametrized by 6 numbers:
@@ -84,3 +87,64 @@ See [touch_grid.xml](../../model/plugin/sensor/touch_grid.xml) to play with the
### Illustration combining resolution, fields-of-view and foveal deformation
[![touch grid illustration](https://img.youtube.com/vi/YScjmR8LwQI/0.jpg)](https://www.youtube.com/watch?v=YScjmR8LwQI)
## [Touch Stress](touch_stress.h)
This sensor is based on similar concepts and parametrization as the `touch_grid`,
while overcoming some of its limitations. In particular, the `touch_grid` can
only provide sparse information, depending on the number of contact points
generated. The `touch_stress` sensor can instead generate a high-resolution
touch image. In order to do this, it requires a signed distance function (SDF)
of the object that is in contact with the sensor. This is handled internally for
primitives or it must be declared explicitly in the model using SDF plugins.
There is one important difference with respect to the `touch_grid`: in this case,
the force is computed in the local taxel frame and not in the frame of the sensor.
This allows for a more intuitive interpretation of normal and tangential stresses,
as shown in the images below.
Note that in this case, the absolute values of the stresses reported by the
sensor are unrelated to the contact forces. They are purely based on geometric
and kinematic considerations, i.e. the SDF for the normal stress and the sliding
velocity for the tangential contributions.
### Example model with analytical SDF
```xml
<extension>
<plugin plugin="mujoco.sdf.gear">
<instance name="gear">
<config key="alpha" value="0"/>
</instance>
</plugin>
<plugin plugin="mujoco.sensor.touch_stress">
<instance name="touch_stress">
<config key="size" value="37 37"/>
<config key="fov" value="45 45"/>
<config key="gamma" value="0"/>
<config key="nchannel" value="3"/>
</instance>
</plugin>
...
<asset>
<mesh name="gear">
<plugin instance="gear"/>
</mesh>
</asset>
...
<worldbody>
<geom type="sdf" name="gear" mesh="gear">
<plugin instance="gear"/>
</geom>
...
<sensor>
<plugin instance="touch_stress" objtype="site" objname="touch_site"/>
</sensor>
</extension>
```
The images below show a static sphere over a gear described by an analytic SDF
and the same sphere dragged along the x and y axes.
<img src="images/normal.png" style="width: 300px;"/>
<img src="images/tangential1.png" style="width: 300px;"/>
<img src="images/tangential2.png" style="width: 300px;"/>
plugin/sensor/images/normal.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 129 KiB

plugin/sensor/images/tangential1.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 89 KiB

plugin/sensor/images/tangential2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 93 KiB

plugin/sensor/sensor.cc → plugin/sensor/register.cc
+5 -1
View File
@@ -14,9 +14,13 @@
#include <mujoco/mjplugin.h>
#include "touch_grid.h"
#include "touch_stress.h"
namespace mujoco::plugin::sensor {
mjPLUGIN_LIB_INIT { TouchGrid::RegisterPlugin(); }
mjPLUGIN_LIB_INIT {
TouchGrid::RegisterPlugin();
TouchStress::RegisterPlugin();
}
} // namespace mujoco::plugin::sensor
plugin/sensor/touch_stress.cc
+557
View File
@@ -0,0 +1,557 @@
// Copyright 2023 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "touch_stress.h"
#include <algorithm>
#include <cctype>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <sstream>
#include <string>
#include <unordered_set>
#include <vector>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mujoco.h>
namespace mujoco::plugin::sensor {
namespace {
// Checks that a plugin config attribute exists.
bool CheckAttr(const std::string& input) {
char* end;
std::string value = input;
value.erase(std::remove_if(value.begin(), value.end(), isspace), value.end());
strtod(value.c_str(), &end);
return end == value.data() + value.size();
}
// Converts a string into a numeric vector
template <typename T>
void ReadVector(std::vector<T>& output, const std::string& input) {
std::stringstream ss(input);
std::string item;
char delim = ' ';
while (getline(ss, item, delim)) {
CheckAttr(item);
output.push_back(strtod(item.c_str(), nullptr));
}
}
// Evenly spaced numbers over a specified interval.
void LinSpace(mjtNum lower, mjtNum upper, int n, mjtNum array[]) {
mjtNum increment = n > 1 ? (upper - lower) / (n - 1) : 0;
for (int i = 0; i < n; ++i) {
*array = lower;
++array;
lower += increment;
}
}
// Parametrized linear/quintic interpolated nonlinearity.
mjtNum Fovea(mjtNum x, mjtNum gamma) {
// Quick return.
if (!gamma) return x;
// Foveal deformation.
mjtNum g = mjMAX(0, mjMIN(1, gamma));
return g*mju_pow(x, 5) + (1 - g)*x;
}
// Make bin edges.
void BinEdges(mjtNum* x_edges, mjtNum* y_edges, int size[2], mjtNum fov[2],
mjtNum gamma) {
// Make unit bin edges.
LinSpace(-1, 1, size[0] + 1, x_edges);
LinSpace(-1, 1, size[1] + 1, y_edges);
// Apply foveal deformation.
for (int i = 0; i < size[0] + 1; i++) {
x_edges[i] = Fovea(x_edges[i], gamma);
}
for (int i = 0; i < size[1] + 1; i++) {
y_edges[i] = Fovea(y_edges[i], gamma);
}
// Scale by field-of-view.
mju_scl(x_edges, x_edges, fov[0]*mjPI / 180, size[0] + 1);
mju_scl(y_edges, y_edges, fov[1]*mjPI / 180, size[1] + 1);
}
// Permute 3-vector from 0,1,2 to 2,0,1.
static void xyz2zxy(mjtNum* x) {
mjtNum z = x[2];
x[2] = x[1];
x[1] = x[0];
x[0] = z;
}
// Transform spherical (azimuth, elevation, radius) to Cartesian (x,y,z).
void SphericalToCartesian(const mjtNum aer[3], mjtNum xyz[3]) {
mjtNum a = aer[0], e = aer[1], r = aer[2];
xyz[0] = r * mju_cos(e) * mju_sin(a);
xyz[1] = r * mju_sin(e);
xyz[2] = -r * mju_cos(e) * mju_cos(a);
}
// Tangent frame in Cartesian coordinates.
void TangentFrame(const mjtNum aer[3], mjtNum mat[9]) {
mjtNum a = aer[0], e = aer[1], r = aer[2];
mjtNum ta[3] = {r * mju_cos(e) * mju_cos(a), 0,
r * mju_cos(e) * mju_sin(a)};
mjtNum te[3] = {-r * mju_sin(e) * mju_sin(a), r * mju_cos(e),
r * mju_sin(e) * mju_cos(a)};
mju_normalize3(ta);
mju_normalize3(te);
mju_copy3(mat, ta);
mju_copy3(mat+3, te);
mju_cross(mat+6, te, ta);
}
} // namespace
// Creates a TouchStress instance if all config attributes are defined and
// within their allowed bounds.
TouchStress* TouchStress::Create(const mjModel* m, mjData* d,
int instance) {
if (CheckAttr(std::string(mj_getPluginConfig(m, instance, "gamma"))) &&
CheckAttr(std::string(mj_getPluginConfig(m, instance, "nchannel")))) {
// nchannel
int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr);
if (!nchannel) nchannel = 1;
if (nchannel < 1 || nchannel > 3) {
mju_error("nchannel must be between 1 and 3");
return nullptr;
}
// size
std::vector<int> size;
std::string size_str = std::string(mj_getPluginConfig(m, instance, "size"));
ReadVector(size, size_str.c_str());
if (size.size()!= 2) {
mju_error("Both horizontal and vertical resolutions must be specified");
return nullptr;
}
if (size[0] <= 0 || size[1] <= 0) {
mju_error("Horizontal and vertical resolutions must be positive");
return nullptr;
}
// field of view
std::vector<mjtNum> fov;
std::string fov_str = std::string(mj_getPluginConfig(m, instance, "fov"));
ReadVector(fov, fov_str.c_str());
if (fov.size()!= 2) {
mju_error(
"Both horizontal and vertical fields of view must be specified");
return nullptr;
}
if (fov[0] <= 0 || fov[0] > 180) {
mju_error("`fov[0]` must be a float between (0, 180] degrees");
return nullptr;
}
if (fov[1] <= 0 || fov[1] > 90) {
mju_error("`fov[1]` must be a float between (0, 90] degrees");
return nullptr;
}
// gamma
mjtNum gamma = strtod(mj_getPluginConfig(m, instance, "gamma"), nullptr);
if (gamma < 0 || gamma > 1) {
mju_error("`gamma` must be a nonnegative float between [0, 1]");
return nullptr;
}
return new TouchStress(m, d, instance, nchannel, size.data(), fov.data(),
gamma);
} else {
mju_error("Invalid or missing parameters in touch_grid sensor plugin");
return nullptr;
}
}
TouchStress::TouchStress(const mjModel* m, mjData* d, int instance,
int nchannel, int size[2], mjtNum fov[2], mjtNum gamma)
: nchannel_(nchannel),
size_{size[0], size[1]},
fov_{fov[0], fov[1]},
gamma_(gamma) {
// Make sure sensor is attached to a site.
for (int i = 0; i < m->nsensor; ++i) {
if (m->sensor_type[i] == mjSENS_PLUGIN && m->sensor_plugin[i] == instance) {
if (m->sensor_objtype[i] != mjOBJ_SITE) {
mju_error("Touch Grid sensor must be attached to a site");
}
}
}
// Get sensor id.
for (id_ = 0; id_ < m->nsensor; ++id_) {
if (m->sensor_type[id_] == mjSENS_PLUGIN &&
m->sensor_plugin[id_] == instance) {
break;
}
}
// Get parent weld id.
int site_id = m->sensor_objid[id_];
int parent_body = m->body_weldid[m->site_bodyid[site_id]];
parent_weld_ = m->body_weldid[parent_body];
// Get geom id.
if (m->body_geomnum[parent_body] != 1) {
mju_error("Touch sensor must be attached to a body with exactly one geom");
}
geom_id_ = m->body_geomadr[parent_body];
// Create bin edges.
x_edges_.assign(size[0] + 1, 0);
y_edges_.assign(size[1] + 1, 0);
BinEdges(x_edges_.data(), y_edges_.data(), size_, fov_, gamma_);
dist_.resize(size[0]*size[1], 0);
pos_.resize(3*size[0]*size[1], 0);
mat_.resize(9*size[0]*size[1], 0);
// Precompute spherical coordinates.
for (int i = 0; i < size[0]; i++) {
for (int j = 0; j < size[1]; j++) {
mjtNum aer[3];
aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]);
aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]);
aer[2] = m->geom_size[3*geom_id_];
SphericalToCartesian(aer, pos_.data() + 3 * (i * size[1] + j));
dist_[i*size[1]+j] = mju_abs(aer[2]);
TangentFrame(aer, mat_.data() + 9 * (i * size[1] + j));
}
}
}
void TouchStress::Reset(const mjModel* m, int instance) {}
void TouchStress::Compute(const mjModel* m, mjData* d, int instance) {
mj_markStack(d);
// Clear sensordata and distance matrix.
mjtNum* sensordata = d->sensordata + m->sensor_adr[id_];
mju_zero(sensordata, m->sensor_dim[id_]);
// Get site id.
int site_id = m->sensor_objid[id_];
// Count contacts and get contact geom ids.
std::unordered_set<int> contact_geom_ids;
for (int i = 0; i < d->ncon; i++) {
int body1 = m->body_weldid[m->geom_bodyid[d->contact[i].geom1]];
int body2 = m->body_weldid[m->geom_bodyid[d->contact[i].geom2]];
if (body1 == parent_weld_) {
contact_geom_ids.insert(d->contact[i].geom2);
}
if (body2 == parent_weld_) {
contact_geom_ids.insert(d->contact[i].geom1);
}
}
// No contacts, return.
if (contact_geom_ids.empty()) {
mj_freeStack(d);
return;
}
// All of the quadrature points are contact points.
int ncon = size_[0]*size_[1];
// Get site frame.
mjtNum* site_pos = d->site_xpos + 3*site_id;
mjtNum* site_mat = d->site_xmat + 9*site_id;
// Allocate contact forces and positions.
mjtNum* forces = mj_stackAllocNum(d, ncon*3);
mjtNum* forcesT = mj_stackAllocNum(d, ncon*3);
// Iterate over colliding geoms.
for (auto geom : contact_geom_ids) {
int body = m->geom_bodyid[geom];
// Get sdf plugin of the geoms.
int sdf_instance[2] = {-1, geom_id_};
mjtGeom geomtype[2] = {mjGEOM_SDF, mjGEOM_SPHERE};
const mjpPlugin* sdf_ptr[2] = {NULL, NULL};
if (m->geom_type[geom] == mjGEOM_SDF) {
sdf_instance[0] = m->geom_plugin[geom];
sdf_ptr[0] = mjc_getSDF(m, geom);
} else {
sdf_instance[0] = geom;
geomtype[0] = (mjtGeom)m->geom_type[geom];
}
// Set SDF parameters.
mjSDF geom_sdf;
geom_sdf.id = &sdf_instance[0];
geom_sdf.type = mjSDFTYPE_SINGLE;
geom_sdf.plugin = &sdf_ptr[0];
geom_sdf.geomtype = &geomtype[0];
mjSDF sensor_sdf;
sensor_sdf.id = &sdf_instance[1];
sensor_sdf.type = mjSDFTYPE_SINGLE;
sensor_sdf.plugin = &sdf_ptr[1];
sensor_sdf.geomtype = &geomtype[1];
// Get forces and positions in spherical coordinates.
int node = 0;
for (int j = 0; j < size_[1]; j++) {
for (int i = 0; i < size_[0]; i++) {
// Position in site frame.
mjtNum* pos = pos_.data() + 3*(i*size_[1] + j);
mjtNum* mat = mat_.data() + 9*(i*size_[1] + j);
// Position in global frame.
mjtNum xpos[3];
mju_mulMatVec3(xpos, site_mat, pos);
mju_addTo3(xpos, site_pos);
// Position in other geom frame.
mjtNum lpos[3], tmp[3];
mju_sub3(tmp, xpos, d->geom_xpos + 3*geom);
mju_mulMatTVec3(lpos, d->geom_xmat + 9*geom, tmp);
// Add mesh position if needed.
if (m->geom_type[geom] == mjGEOM_MESH ||
m->geom_type[geom] == mjGEOM_SDF) {
mjtNum mesh_mat[9];
mju_quat2Mat(mesh_mat, m->mesh_quat + 4 * m->geom_dataid[geom]);
mju_mulMatVec3(lpos, mesh_mat, lpos);
mju_addTo3(lpos, m->mesh_pos + 3 * m->geom_dataid[geom]);
}
// Compute distance.
mjtNum depth = mju_min(mjc_distance(m, d, &geom_sdf, lpos), 0);
if (depth == 0) {
mju_zero3(forces + 3*node);
node++;
continue;
}
// Get velocity in global frame.
mjtNum vel_sensor[6], vel_other[6], vel_rel[3];
mju_transformSpatial(
vel_sensor, d->cvel + 6 * parent_weld_, 0, xpos,
d->subtree_com + 3 * m->body_rootid[parent_weld_], NULL);
mju_transformSpatial(
vel_other, d->cvel + 6 * body, 0, d->geom_xpos + 3 * geom,
d->subtree_com + 3 * m->body_rootid[body], NULL);
mju_sub3(vel_rel, vel_sensor+3, vel_other+3);
// Get contact force/torque, rotate into node frame.
mjtNum tmp_force[3], normal[3];
mjtNum kMaxDepth = 0.05;
mjtNum pressure = 1 / (kMaxDepth - depth) - 1 / kMaxDepth;
mjc_gradient(m, d, &sensor_sdf, normal, pos);
mju_scl3(tmp_force, normal, pressure);
mju_mulMatTVec3(forces + 3*node, mat, tmp_force);
forces[3*node+0] = mju_abs(mju_dot3(vel_rel, mat + 0));
forces[3*node+1] = mju_abs(mju_dot3(vel_rel, mat + 3));
// Permute forces from x,y,z to z,x,y (normal, tangent, tangent)
xyz2zxy(forces + 3*node);
node++;
}
}
// Transpose forces.
mju_transpose(forcesT, forces, ncon, 3);
// Compute sensor output.
for (int c = 0; c < nchannel_; c++) {
if (!mju_isZero(forcesT + c*ncon, ncon)) {
mju_addTo(sensordata + c*ncon, forcesT + c*ncon, size_[0]*size_[1]);
}
}
}
mj_freeStack(d);
}
// Thickness of taxel-visualization boxes relative to contact distance.
static const mjtNum kRelativeThickness = 0.02;
void TouchStress::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
mjvScene* scn, int instance) {
mj_markStack(d);
// Get sensor data.
mjtNum* sensordata = d->sensordata + m->sensor_adr[id_];
// Get maximum absolute normal force.
mjtNum maxval = 0;
int frame = size_[0]*size_[1];
for (int j=0; j < frame; j++) {
maxval = mju_max(maxval, mju_abs(sensordata[j]));
}
// If no normal force readings, quick return.
if (!maxval) {
mj_freeStack(d);
return;
}
// Get site id and frame.
int site_id = m->sensor_objid[id_];
mjtNum* site_pos = d->site_xpos + 3*site_id;
mjtNum* site_mat = d->site_xmat + 9*site_id;
mjtNum site_quat[4];
mju_mat2Quat(site_quat, site_mat);
// Draw geoms.
for (int i=0; i < size_[0]; i++) {
for (int j=0; j < size_[1]; j++) {
mjtNum dist = dist_[i*size_[1]+j];
if (!dist) {
continue;
}
if (scn->ngeom >= scn->maxgeom) {
mj_warning(d, mjWARN_VGEOMFULL, scn->maxgeom);
mj_freeStack(d);
return;
} else {
// size
mjtNum size[3];
size[0] = dist*0.5*(x_edges_[i+1]-x_edges_[i]);
size[1] = dist*0.5*(y_edges_[j+1]-y_edges_[j]);
size[2] = dist*kRelativeThickness;
// position
mjtNum pos[3];
mjtNum aer[3];
aer[0] = 0.5*(x_edges_[i+1]+x_edges_[i]);
aer[1] = 0.5*(y_edges_[j+1]+y_edges_[j]);
aer[2] = dist*(1-kRelativeThickness);
SphericalToCartesian(aer, pos);
mju_mulMatVec3(pos, site_mat, pos);
mju_addTo3(pos, site_pos);
// orientation
mjtNum a_quat[4];
mjtNum site_y[3] = {-site_mat[1], -site_mat[4], -site_mat[7]};
mju_axisAngle2Quat(a_quat, site_y, aer[0]);
mjtNum e_quat[4];
mjtNum site_x[3] = {site_mat[0], site_mat[3], site_mat[6]};
mju_axisAngle2Quat(e_quat, site_x, aer[1]);
mjtNum quat[4];
mju_mulQuat(quat, e_quat, site_quat);
mju_mulQuat(quat, a_quat, quat);
mjtNum mat[9];
mju_quat2Mat(mat, quat);
// color
float rgba[4] = {1, 1, 1, 1.0};
for (int k=0; k < mjMIN(nchannel_, 3); k++) {
rgba[k] = mju_abs(sensordata[k*frame + j*size_[0] + i]) / maxval;
}
// draw box geom
mjvGeom* thisgeom = scn->geoms + scn->ngeom;
mjv_initGeom(thisgeom, mjGEOM_BOX, size, pos, mat, rgba);
thisgeom->objtype = mjOBJ_UNKNOWN;
thisgeom->objid = id_;
thisgeom->category = mjCAT_DECOR;
thisgeom->segid = scn->ngeom;
scn->ngeom++;
}
}
}
mj_freeStack(d);
}
void TouchStress::RegisterPlugin() {
mjpPlugin plugin;
mjp_defaultPlugin(&plugin);
plugin.name = "mujoco.sensor.touch_stress";
plugin.capabilityflags |= mjPLUGIN_SENSOR;
// Parameterized by 4 attributes.
const char* attributes[] = {"nchannel", "size", "fov", "gamma"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
// Stateless.
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
// Sensor dimension = nchannel * size[0] * size[1]
plugin.nsensordata = +[](const mjModel* m, int instance, int sensor_id) {
int nchannel = strtod(mj_getPluginConfig(m, instance, "nchannel"), nullptr);
if (!nchannel) nchannel = 1;
std::vector<int> size;
std::string size_str = std::string(mj_getPluginConfig(m, instance, "size"));
ReadVector(size, size_str.c_str());
return nchannel * size[0] * size[1];
};
// Can only run after forces have been computed.
plugin.needstage = mjSTAGE_ACC;
// Initialization callback.
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
auto* TouchStress = TouchStress::Create(m, d, instance);
if (!TouchStress) {
return -1;
}
d->plugin_data[instance] = reinterpret_cast<uintptr_t>(TouchStress);
return 0;
};
// Destruction callback.
plugin.destroy = +[](mjData* d, int instance) {
delete reinterpret_cast<TouchStress*>(d->plugin_data[instance]);
d->plugin_data[instance] = 0;
};
// Reset callback.
plugin.reset = +[](const mjModel* m, mjtNum* plugin_state, void* plugin_data,
int instance) {
auto* TouchStress = reinterpret_cast<class TouchStress*>(plugin_data);
TouchStress->Reset(m, instance);
};
// Compute callback.
plugin.compute =
+[](const mjModel* m, mjData* d, int instance, int capability_bit) {
auto* TouchStress =
reinterpret_cast<class TouchStress*>(d->plugin_data[instance]);
TouchStress->Compute(m, d, instance);
};
// Visualization callback.
plugin.visualize = +[](const mjModel* m, mjData* d, const mjvOption* opt,
mjvScene* scn, int instance) {
auto* TouchStress =
reinterpret_cast<class TouchStress*>(d->plugin_data[instance]);
TouchStress->Visualize(m, d, opt, scn, instance);
};
// Register the plugin.
mjp_registerPlugin(&plugin);
}
} // namespace mujoco::plugin::sensor
plugin/sensor/touch_stress.h
+80
View File
@@ -0,0 +1,80 @@
// Copyright 2025 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_
#define MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_
#include <vector>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjvisualize.h>
namespace mujoco::plugin::sensor {
// A touch grid sensor is associated with a site and senses contact stresses
// between the site's parent body and all other bodies. The site's
// frame determines the orientation of the sensor with the same convention used
// for cameras and lights: the sensor points in the frame's negative-Z
// direction, so the X and Y axes correspond to horizontal and vertical
// directions respectively.
//
// The output of the sensor is a stack of 3 "touch images" corresponding to
// forces in the local frame of the taxels. Forces are in the [z, x, y] order,
// corresponding to the ordering in contact frames: [normal, tangent, tangent].
//
// The sensor has 6 parameters:
// 1. (int) Number of channels [1-3]. Defaults to 1.
// 2. (int) Horizontal resolution.
// 3. (int) Vertical resolution.
// 4. (float) Horizontal field-of-view (fov_x), in degrees.
// 5. (float) Vertical field-of-view (fov_y), in degrees.
// 6. (float) Foveal deformation. Defaults to 0.
class TouchStress {
public:
static TouchStress* Create(const mjModel* m, mjData* d, int instance);
TouchStress(TouchStress&&) = default;
~TouchStress() = default;
void Reset(const mjModel* m, int instance);
void Compute(const mjModel* m, mjData* d, int instance);
void Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
mjvScene* scn, int instance);
static void RegisterPlugin();
int nchannel_; // number of channels (1-3)
int size_[2]; // horizontal and vertical resolution
mjtNum fov_[2]; // horizontal and vertical field of view, in degrees
mjtNum gamma_; // foveal deformation
private:
TouchStress(const mjModel* m, mjData* d, int instance, int nchannel,
int* size, mjtNum* fov_x, mjtNum gamma);
std::vector<mjtNum> x_edges_;
std::vector<mjtNum> y_edges_;
std::vector<mjtNum> dist_;
std::vector<mjtNum> pos_;
std::vector<mjtNum> mat_;
int id_;
int parent_weld_;
int geom_id_;
};
} // namespace mujoco::plugin::sensor
#endif // MUJOCO_PLUGIN_SENSOR_TOUCH_STRESS_H_
python/build_requirements.txt
+3 -2
View File
@@ -57,8 +57,9 @@ PyOpenGL==3.1.7 \
--hash=sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a
pytest==8.3.3 \
--hash=sha256:a6853c7375b2663155079443d2e45de913a911a11d669df02a50814944db57b2
setuptools==75.5.0 \
--hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
setuptools==78.1.1 \
--hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
--hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
wheel==0.45.0 \
--hash=sha256:52f0baa5e6522155090a09c6bd95718cc46956d1b51d537ea5454249edb671c7
python/make_sdist_requirements.txt
+3 -2
View File
@@ -4,8 +4,9 @@ build==1.2.2.post1 \
--hash=sha256:1d61c0887fa860c01971625baae8bdd338e517b836a2f70dd1f7aa3a6b2fc5b5
pip==24.3.1 \
--hash=sha256:3790624780082365f47549d032f3770eeb2b1e8bd1f7b2e02dace1afa361b4ed
setuptools==75.5.0 \
--hash=sha256:87cb777c3b96d638ca02031192d40390e0ad97737e27b6b4fa831bea86f2f829
setuptools==78.1.1 \
--hash=sha256:c3a9c4211ff4c309edb8b8c4f1cbfa7ae324c4ba9f91ff254e3d305b9fd54561 \
--hash=sha256:fcc17fd9cd898242f6b4adfaca46137a9edef687f43e6f78469692a5e70d851d
# Transitive dependencies of build
colorama==0.4.6; platform_system == 'Windows' \
python/mjspec.ipynb
+1534 -326
View File
File diff suppressed because it is too large Load Diff
python/mujoco/CMakeLists.txt
+26 -4
View File
@@ -84,7 +84,7 @@ if(NOT TARGET mujoco)
if(MUJOCO_FRAMEWORK)
message("MuJoCo framework is at ${MUJOCO_FRAMEWORK}/mujoco.framework")
set(MUJOCO_LIBRARY
${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.1.dylib
${MUJOCO_FRAMEWORK}/mujoco.framework/Versions/A/libmujoco.3.3.3.dylib
)
target_compile_options(mujoco INTERFACE -F${MUJOCO_FRAMEWORK})
endif()
@@ -92,7 +92,7 @@ if(NOT TARGET mujoco)
if(NOT MUJOCO_FRAMEWORK)
find_library(
MUJOCO_LIBRARY mujoco mujoco.3.3.1 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
MUJOCO_LIBRARY mujoco mujoco.3.3.3 HINTS ${MUJOCO_LIBRARY_DIR} REQUIRED
)
find_path(MUJOCO_INCLUDE mujoco/mujoco.h HINTS ${MUJOCO_INCLUDE_DIR} REQUIRED)
message("MuJoCo is at ${MUJOCO_LIBRARY}")
@@ -140,7 +140,7 @@ findorfetch(
GIT_REPO
https://github.com/abseil/abseil-cpp
GIT_TAG
9ac7062b1860d895fb5a8cbf58c3e9ef8f674b5f # LTS 20250127.0
d9e4955c65cd4367dd6bf46f4ccb8cd3d100540b # LTS 20250127.1
TARGETS
${MUJOCO_PYTHON_ABSL_TARGETS}
EXCLUDE_FROM_ALL
@@ -173,7 +173,7 @@ findorfetch(
GIT_REPO
https://gitlab.com/libeigen/eigen
GIT_TAG
66f7f51b7e069d0a03a21157fa60b24aece69aeb
464c1d097891a1462ab28bf8bb763c1683883892
TARGETS
Eigen3::Eigen
EXCLUDE_FROM_ALL
@@ -278,6 +278,25 @@ target_link_libraries(
raw
)
add_library(structs_wrappers STATIC
structs_wrappers.cc
serialization.h
indexers.cc
)
target_include_directories(structs_wrappers PRIVATE ${Python3_INCLUDE_DIRS})
target_link_libraries(
structs_wrappers
PRIVATE absl::flat_hash_map
absl::span
crossplatform
mujoco
raw
structs_header
pybind11::headers
Eigen3::Eigen
)
add_library(functions_header INTERFACE)
target_sources(functions_header INTERFACE functions.h)
set_target_properties(functions_header PROPERTIES PUBLIC_HEADER functions.h)
@@ -403,6 +422,7 @@ target_link_libraries(
func_wrap
function_traits
structs_header
structs_wrappers
)
if(NOT EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/specs.cc.inc)
@@ -417,6 +437,8 @@ endif()
mujoco_pybind11_module(
_specs
specs.cc
specs_wrapper.h
specs_wrapper.cc
specs.cc.inc
)
target_link_libraries(
python/mujoco/bindings_test.py
+17 -1
View File
@@ -623,10 +623,12 @@ class MuJoCoBindingsTest(parameterized.TestCase):
ncon = 13
nefc = 17
mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc)
nj = 21
mujoco._functions._realloc_con_efc(self.data, ncon=ncon, nefc=nefc, nJ=nj)
self.assertLen(self.data.contact, ncon)
self.assertEqual(self.data.efc_id.shape, (nefc,))
self.assertEqual(self.data.efc_J.shape, (nj,))
self.assertEqual(self.data.efc_KBIP.shape, (nefc, 4))
expected_error = 'insufficient arena memory available'
@@ -1083,6 +1085,20 @@ Euler integrator, semi-implicit in velocity.
):
mujoco.mj_forward(self.model, self.data)
def test_timer_installed_by_default(self):
timer_step = mujoco.mjtTimer.mjTIMER_STEP
self.assertEqual(self.data.timer[timer_step].number, 0)
self.assertEqual(self.data.timer[timer_step].duration, 0.0)
mujoco.mj_step(self.model, self.data)
self.assertEqual(self.data.timer[timer_step].number, 1)
duration_1 = self.data.timer[timer_step].duration
self.assertGreater(duration_1, 0.0)
mujoco.mj_step(self.model, self.data, 5)
self.assertEqual(self.data.timer[timer_step].number, 6)
self.assertGreater(self.data.timer[timer_step].duration, duration_1)
def test_mjcb_time(self):
class CallCounter:
python/mujoco/codegen/generate_spec_bindings.py
+34 -1
View File
@@ -84,6 +84,7 @@ def _value_binding_code(
fulltype = fulltype.replace('mjOption', 'raw::MjOption')
fulltype = fulltype.replace('mjVisual', 'raw::MjVisual')
fulltype = fulltype.replace('mjStatistic', 'raw::MjStatistic')
element = '.element' if fullvarname == 'plugin' else ''
def_property_args = (
f'"{varname}"',
@@ -91,7 +92,7 @@ def _value_binding_code(
return self.{fullvarname};
}}""",
f"""[]({rawclassname}& self, {fulltype} {varname}) {{
self.{fullvarname} = {varname};
self.{fullvarname}{element} = {varname}{element};
}}""",
)
@@ -613,12 +614,44 @@ def generate_find() -> None:
print(code)
def generate_signature() -> None:
"""Generate signature functions."""
for key, _, _, _, _ in SPECS:
elem = key.removeprefix('mjs')
titlecase = 'Mjs' + elem
code = f"""\n
{key}.def_property_readonly("signature",
[](raw::{titlecase}& self) -> uint64_t {{
return mjs_getSpec(self.element)->element->signature;
}});
"""
print(code)
def generate_id() -> None:
"""Generate id functions."""
for key, _, _, _, _ in SPECS:
if key == 'mjsPlugin':
continue
elem = key.removeprefix('mjs')
titlecase = 'Mjs' + elem
code = f"""\n
{key}.def_property_readonly("id",
[](raw::{titlecase}& self) -> int {{
return mjs_getId(self.element);
}});
"""
print(code)
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
generate()
generate_add()
generate_find()
generate_signature()
generate_id()
if __name__ == '__main__':
python/mujoco/functions.cc
+9 -6
View File
@@ -211,6 +211,7 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_tendon>(pymodule);
Def<traits::mj_transmission>(pymodule);
Def<traits::mj_crb>(pymodule);
Def<traits::mj_makeM>(pymodule);
Def<traits::mj_factorM>(pymodule);
DEF_WITH_OMITTED_PY_ARGS(traits::mj_solveM, "n")(
pymodule,
@@ -1467,10 +1468,10 @@ PYBIND11_MODULE(_functions, pymodule) {
pymodule.def(
"_realloc_con_efc",
[](MjDataWrapper& d, int ncon, int nefc) {
[](MjDataWrapper& d, int ncon, int nefc, int nJ) {
raw::MjData* data = d.get();
auto cleanup = [](raw::MjData* data) {
auto cleanup = [](raw::MjData* data, int nJ) {
#ifdef ADDRESS_SANITIZER
ASAN_POISON_MEMORY_REGION(
static_cast<char*>(data->arena),
@@ -1479,6 +1480,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->parena = 0;
data->ncon = 0;
data->nefc = 0;
if (nJ > -1) data->nJ = 0;
data->contact = static_cast<raw::MjContact*>(data->arena);
#define X(type, name, nr, nc) data->name = nullptr;
MJDATA_ARENA_POINTERS_SOLVER
@@ -1486,14 +1488,15 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef X
};
cleanup(data);
cleanup(data, nJ);
data->ncon = ncon;
data->nefc = nefc;
if (nJ > -1) data->nJ = nJ;
data->contact =
static_cast<raw::MjContact*>(InterceptMjErrors(::mj_arenaAllocByte)(
data, ncon * sizeof(raw::MjContact), alignof(raw::MjContact)));
if (!data->contact) {
cleanup(data);
cleanup(data, nJ);
throw FatalError("insufficient arena memory available");
}
@@ -1505,7 +1508,7 @@ PYBIND11_MODULE(_functions, pymodule) {
data->name = static_cast<type*>(InterceptMjErrors(::mj_arenaAllocByte)( \
data, sizeof(type) * (nr) * (nc), alignof(type))); \
if (!data->name) { \
cleanup(data); \
cleanup(data, nJ); \
throw FatalError("insufficient arena memory available"); \
}
@@ -1519,7 +1522,7 @@ PYBIND11_MODULE(_functions, pymodule) {
#undef MJ_M
#define MJ_M(x) x
},
py::arg("d"), py::arg("ncon"), py::arg("nefc"),
py::arg("d"), py::arg("ncon"), py::arg("nefc"), py::arg("nJ") = -1,
py::call_guard<py::gil_scoped_release>());
} // PYBIND11_MODULE NOLINT(readability/fn_size)
} // namespace
python/mujoco/indexer_xmacro.h
+1 -1
View File
@@ -152,7 +152,7 @@
X( int, light_, mode, nlight, 1 ) \
X( int, light_, bodyid, nlight, 1 ) \
X( int, light_, targetbodyid, nlight, 1 ) \
X( mjtByte, light_, directional, nlight, 1 ) \
X( int, light_, type, nlight, 1 ) \
X( mjtByte, light_, castshadow, nlight, 1 ) \
X( mjtByte, light_, active, nlight, 1 ) \
X( mjtNum, light_, pos, nlight, 3 ) \
python/mujoco/introspect/codegen/generate_functions.py
+7 -2
View File
@@ -99,8 +99,13 @@ class MjFunctionVisitor:
return ''.join(strings)
def visit(self, node: ClangJsonNode) -> None:
if (node.get('kind') == 'FunctionDecl' and
node.get('name', '').startswith('mj')):
# Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated
# objects and doesn't need a python wrapper.
if (
node.get('kind') == 'FunctionDecl'
and node.get('name', '').startswith('mj')
and node.get('name', '') != 'mjs_setUserValueWithCleanup'
):
func_decl = self._make_function(node)
self._functions[func_decl.name] = func_decl
python/mujoco/introspect/codegen/generate_structs.py
+8 -5
View File
@@ -47,7 +47,7 @@ _EXCLUDED = (
'mjResource_',
)
_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s\s+\((.+) x (.+)\)\Z')
_ARRAY_COMMENT_PATTERN = re.compile(r'(.+?)\s+\(([^\(\)]+) x ([^\(\)]+)\)\Z')
def traverse(node, visitor):
@@ -93,16 +93,19 @@ class MjStructVisitor:
# No valid normalization, just parse the declname.
return type_parsing.parse_type(declname)
def _make_comment(self, node: ClangJsonNode) -> str:
def _make_comment(self, node: ClangJsonNode, strip: bool = True) -> str:
"""Makes a comment string from a Clang AST FullComment node."""
kind = node.get('kind')
if kind == 'TextComment':
return node['text'].replace('\N{NO-BREAK SPACE}', '&nbsp;').strip()
retval = node['text'].replace('\N{NO-BREAK SPACE}', '&nbsp;')
else:
strings = []
for child in node['inner']:
strings.append(self._make_comment(child))
return ''.join(strings).strip()
strings.append(self._make_comment(child, strip=False))
retval = ''.join(strings)
if strip:
retval = retval.strip()
return retval
def _make_field(
self, node: ClangJsonNode
python/mujoco/introspect/enums.py
+62 -28
View File
@@ -111,6 +111,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjCAMLIGHT_TARGETBODYCOM', 4),
]),
)),
('mjtLightType',
EnumDecl(
name='mjtLightType',
declname='enum mjtLightType_',
values=dict([
('mjLIGHT_SPOT', 0),
('mjLIGHT_DIRECTIONAL', 1),
('mjLIGHT_POINT', 2),
('mjLIGHT_IMAGE', 3),
]),
)),
('mjtTexture',
EnumDecl(
name='mjtTexture',
@@ -139,6 +150,16 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjNTEXROLE', 10),
]),
)),
('mjtColorSpace',
EnumDecl(
name='mjtColorSpace',
declname='enum mjtColorSpace_',
values=dict([
('mjCOLORSPACE_AUTO', 0),
('mjCOLORSPACE_LINEAR', 1),
('mjCOLORSPACE_SRGB', 2),
]),
)),
('mjtIntegrator',
EnumDecl(
name='mjtIntegrator',
@@ -288,6 +309,7 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjNOBJECT', 26),
('mjOBJ_FRAME', 100),
('mjOBJ_DEFAULT', 101),
('mjOBJ_MODEL', 102),
]),
)),
('mjtConstraint',
@@ -339,34 +361,35 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjSENS_ACTUATORVEL', 14),
('mjSENS_ACTUATORFRC', 15),
('mjSENS_JOINTACTFRC', 16),
('mjSENS_BALLQUAT', 17),
('mjSENS_BALLANGVEL', 18),
('mjSENS_JOINTLIMITPOS', 19),
('mjSENS_JOINTLIMITVEL', 20),
('mjSENS_JOINTLIMITFRC', 21),
('mjSENS_TENDONLIMITPOS', 22),
('mjSENS_TENDONLIMITVEL', 23),
('mjSENS_TENDONLIMITFRC', 24),
('mjSENS_FRAMEPOS', 25),
('mjSENS_FRAMEQUAT', 26),
('mjSENS_FRAMEXAXIS', 27),
('mjSENS_FRAMEYAXIS', 28),
('mjSENS_FRAMEZAXIS', 29),
('mjSENS_FRAMELINVEL', 30),
('mjSENS_FRAMEANGVEL', 31),
('mjSENS_FRAMELINACC', 32),
('mjSENS_FRAMEANGACC', 33),
('mjSENS_SUBTREECOM', 34),
('mjSENS_SUBTREELINVEL', 35),
('mjSENS_SUBTREEANGMOM', 36),
('mjSENS_GEOMDIST', 37),
('mjSENS_GEOMNORMAL', 38),
('mjSENS_GEOMFROMTO', 39),
('mjSENS_E_POTENTIAL', 40),
('mjSENS_E_KINETIC', 41),
('mjSENS_CLOCK', 42),
('mjSENS_PLUGIN', 43),
('mjSENS_USER', 44),
('mjSENS_TENDONACTFRC', 17),
('mjSENS_BALLQUAT', 18),
('mjSENS_BALLANGVEL', 19),
('mjSENS_JOINTLIMITPOS', 20),
('mjSENS_JOINTLIMITVEL', 21),
('mjSENS_JOINTLIMITFRC', 22),
('mjSENS_TENDONLIMITPOS', 23),
('mjSENS_TENDONLIMITVEL', 24),
('mjSENS_TENDONLIMITFRC', 25),
('mjSENS_FRAMEPOS', 26),
('mjSENS_FRAMEQUAT', 27),
('mjSENS_FRAMEXAXIS', 28),
('mjSENS_FRAMEYAXIS', 29),
('mjSENS_FRAMEZAXIS', 30),
('mjSENS_FRAMELINVEL', 31),
('mjSENS_FRAMEANGVEL', 32),
('mjSENS_FRAMELINACC', 33),
('mjSENS_FRAMEANGACC', 34),
('mjSENS_SUBTREECOM', 35),
('mjSENS_SUBTREELINVEL', 36),
('mjSENS_SUBTREEANGMOM', 37),
('mjSENS_GEOMDIST', 38),
('mjSENS_GEOMNORMAL', 39),
('mjSENS_GEOMFROMTO', 40),
('mjSENS_E_POTENTIAL', 41),
('mjSENS_E_KINETIC', 42),
('mjSENS_CLOCK', 43),
('mjSENS_PLUGIN', 44),
('mjSENS_USER', 45),
]),
)),
('mjtStage',
@@ -426,6 +449,17 @@ ENUMS: Mapping[str, EnumDecl] = dict([
('mjFLEXSELF_AUTO', 4),
]),
)),
('mjtSDFType',
EnumDecl(
name='mjtSDFType',
declname='enum mjtSDFType_',
values=dict([
('mjSDFTYPE_SINGLE', 0),
('mjSDFTYPE_INTERSECTION', 1),
('mjSDFTYPE_MIDSURFACE', 2),
('mjSDFTYPE_COLLISION', 3),
]),
)),
('mjtTaskStatus',
EnumDecl(
name='mjtTaskStatus',

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