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

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

Browse Source
This commit is contained in:
Abhishek Joshi
2023-11-15 12:38:46 -06:00
committed by GitHub
parent c37212138e 8d5966eec9
commit 1491ac15d8
133 changed files with 4003 additions and 1936 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/build.yml
+51 -18
View File
@@ -219,7 +219,7 @@ jobs:
working-directory: python/dist
run: >
source ${{ matrix.tmpdir }}/venv/bin/activate &&
pip install --no-deps --no-index mujoco-*.whl
pip install --no-index mujoco-*.whl
- name: Test Python bindings
if: ${{ runner.os != 'Windows' }}
shell: bash
@@ -228,23 +228,56 @@ jobs:
run: >
source ${{ matrix.tmpdir }}/venv/bin/activate &&
pytest -v --pyargs mujoco
- name: Package MJX
if: ${{ runner.os != 'Windows' }}
shell: bash
working-directory: mjx
run:
source ${{ matrix.tmpdir }}/venv/bin/activate &&
python -m build .
- name: Install MJX
if: ${{ runner.os != 'Windows' }}
shell: bash
working-directory: mjx
run:
source ${{ matrix.tmpdir }}/venv/bin/activate &&
pip install --require-hashes -r requirements.txt &&
pip install --no-index dist/mujoco_mjx-*.whl
- name: Test MJX
if: ${{ runner.os != 'Windows' }}
shell: bash
run:
source ${{ matrix.tmpdir }}/venv/bin/activate &&
pytest -n auto -v --pyargs mujoco.mjx
- name: Notify team chat
shell: bash
env:
GCHAT_API: ${{ secrets.GCHAT_API }}
JOB_URL: "${{ github.server_url }}/${{ github.repository }}/actions/runs/\
${{ github.run_id }}"
if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API != '' }}
run: >
curl
"$GCHAT_API&threadKey=$GITHUB_SHA&messageReplyOption=REPLY_MESSAGE_FALLBACK_TO_NEW_THREAD"
-X POST
-H "Content-Type: application/json"
-d "{
'text': '<$JOB_URL|*FAILURE*>:
job \`${{ matrix.os }}${{ matrix.additional_label }}\`
commit \`$(echo $GITHUB_SHA | head -c6)\`\n
\`\`\`Author: ${{ github.event.head_commit.author.name }}
<${{ github.event.head_commit.author.email }}>
\n\n${{ github.event.head_commit.message }}\`\`\`
'}"
GCHAT_API_URL: ${{ secrets.GCHAT_API }}
JOB_URL: "${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
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 }}
if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API_URL != '' }}
run: |
CHATMSG="$(cat <<-'EOF' | python3
import json
import os
env = lambda x: os.getenv(x, '')
data = dict(
result=env('JOB_URL'),
job=env('CHATMSG_JOB_ID'),
commit=env('GITHUB_SHA')[:6],
name=env('CHATMSG_AUTHOR_NAME').replace('```', ''),
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)
print(json.dumps({'text' : text}))
EOF
)" &&
curl "$GCHAT_API_URL&threadKey=$GITHUB_SHA&messageReplyOption=REPLY_MESSAGE_FALLBACK_TO_NEW_THREAD" \
-X POST \
-H "Content-Type: application/json" \
--data-raw "${CHATMSG}"
README.md
+4
View File
@@ -55,6 +55,10 @@ running on Google Colab:
- For a more advanced example, see the LQR tutorial which creates an LQR
controller to balance a humanoid on one leg using MuJoCo's dynamics
derivatives: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb)
- The MJX tutorial provides usage examples of
[MuJoCo XLA](https://mujoco.readthedocs.io/en/stable/mjx.html), a branch of MuJoCo written in
JAX:
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb)
## Installation
doc/APIreference/APIglobals.rst
+1 -1
View File
@@ -498,7 +498,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
- Maximum number of UI sections.
Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
* - ``mjMAXUIITEM``
- 80
- 200
- Maximum number of items per UI section.
Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
* - ``mjMAXUITEXT``
doc/APIreference/APItypes.rst
+15 -3
View File
@@ -34,6 +34,8 @@ MuJoCo defines a large number of types:
- Structs used by :ref:`engine plugins<tyPluginStructure>`.
- Several :ref:`tyFunction` for user-defined callbacks.
- :ref:`tyNotes` regarding specific data structures that require detailed description.
.. _tyPrimitive:
@@ -909,6 +911,7 @@ This structure contains the custom OpenGL rendering context, with the ids of all
User Interface
^^^^^^^^^^^^^^
For a high-level description of the UI framework, see :ref:`UI`.
The names of these struct types are prefixed with ``mjui``, except for the main :ref:`mjUI` struct itself.
@@ -917,7 +920,9 @@ The names of these struct types are prefixed with ``mjui``, except for the main
mjuiState
~~~~~~~~~
This structure contains the keyboard and mouse state used by the UI framework.
This C struct represents the global state of the window, keyboard and mouse, input event descriptors, and all window
rectangles (including the visible UI rectangles). There is only one ``mjuiState`` per application, even if there are
multiple UIs. This struct would normally be defined as a global variable.
.. mujoco-include:: mjuiState
@@ -967,7 +972,9 @@ This structure defines one section of the UI.
mjuiDef
~~~~~~~
This structure defines one entry in the definition table used for simplified UI construction.
This structure defines one entry in the definition table used for simplified UI construction. It contains everything
needed to define one UI item. Some translation is performed by the helper functions, so that multiple mjuiDefs can be
defined as a static table.
.. mujoco-include:: mjuiDef
@@ -977,7 +984,12 @@ This structure defines one entry in the definition table used for simplified UI
mjUI
~~~~
This structure defines the entire UI.
This C struct represents an entire UI. The same application could have multiple UIs, for example on the left and the
right of the window. This would normally be defined as a global variable. As explained earlier, it contains static
allocation for a maximum number of supported UI sections (:ref:`mjuiSection<mjuiSection>`) each with a maximum number
of supported items (:ref:`mjuiItem<mjuiItem>`). It also contains the color and spacing themes, enable/disable
callback, virtual window descriptor, text edit state, mouse focus. Some of these fields are set only once when the UI
is initialized, others change at runtime.
.. mujoco-include:: mjUI
doc/APIreference/functions.rst
+48 -11
View File
@@ -392,7 +392,9 @@ mj_objectAcceleration
.. mujoco-include:: mj_objectAcceleration
Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation.
Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation. If acceleration or force
sensors are not present in the model, :ref:`mj_rnePostConstraint` must be manually called in order to calculate
:ref:`mjData`.cacc -- the total body acceleration, including contributions from the constraint solver.
.. _mj_contactForce:
@@ -803,7 +805,7 @@ mj_passive
.. mujoco-include:: mj_passive
Compute qfrc_passive from spring-dampers, viscosity and density.
Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.
.. _mj_subtreeVel:
@@ -812,7 +814,10 @@ mj_subtreeVel
.. mujoco-include:: mj_subtreeVel
subtree linear velocity and angular momentum
Sub-tree linear velocity and angular momentum: compute ``subtree_linvel``, ``subtree_angmom``.
This function is triggered automatically if the subtree :ref:`velocity<sensor-subtreelinvel>` or
:ref:`momentum<sensor-subtreeangmom>` sensors are present in the model.
It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "vel".
.. _mj_rne:
@@ -821,7 +826,8 @@ mj_rne
.. mujoco-include:: mj_rne
RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term.
Recursive Newton Euler: compute :math:`M(q) \ddot q + C(q,\dot q)`. ``flg_acc=0`` removes the inertial term (i.e.
assumes :math:`\ddot q = 0`).
.. _mj_rnePostConstraint:
@@ -830,7 +836,21 @@ mj_rnePostConstraint
.. mujoco-include:: mj_rnePostConstraint
RNE with complete data: compute cacc, cfrc_ext, cfrc_int.
Recursive Newton Euler with final computed forces and accelerations.
Computes three body-level ``nv x 6`` arrays, all defined in the subtreecom-based
:ref:`c-frame<tyNotesCom>` and arranged in ``[rotation(3), translation(3)]`` order.
- ``cacc``: Body acceleration, required for :ref:`mj_objectAcceleration`.
- ``cfrc_int``: Interaction force with the parent body.
- ``cfrc_ext``: External force acting on the body.
This function is triggered automatically if the following sensors are present in the model:
:ref:`accelerometer<sensor-accelerometer>`, :ref:`force<sensor-force>`, :ref:`torque<sensor-torque>`,
:ref:`framelinacc<sensor-framelinacc>`, :ref:`frameangacc<sensor-frameangacc>`.
It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "acc".
The computed force arrays ``cfrc_int`` and ``cfrc_ext`` currently suffer from a know bug, they do not take into account
the effect of spatial tendons, see :github:issue:`832`.
.. _mj_collision:
@@ -884,8 +904,8 @@ mj_constraintUpdate
.. mujoco-include:: mj_constraintUpdate
Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set \*cost = s(jar)
where jar = Jac*qacc-aref.
Compute ``efc_state``, ``efc_force``, ``qfrc_constraint``, and (optionally) cone Hessians.
If ``cost`` is not ``NULL``, set ``*cost = s(jar)`` where ``jar = Jac*qacc - aref``.
.. _Raycollisions:
@@ -1945,6 +1965,8 @@ Find first rectangle containing mouse, -1: not found.
UI framework
^^^^^^^^^^^^
For a high-level description of the UI framework, see :ref:`UI`.
.. _mjui_themeSpacing:
mjui_themeSpacing
@@ -1970,7 +1992,12 @@ mjui_add
.. mujoco-include:: mjui_add
Add definitions to UI.
This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs,
each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added
after the end of the last used section. There is also another version of this function
(:ref:`mjui_addToSection<mjui_addToSection>`) which adds items to a specified section instead of adding them at the end
of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by
:ref:`mjMAXUISECT<glNumeric>` and :ref:`mjMAXUIITEM<glNumeric>`. Exceeding these maxima results in low-level errors.
.. _mjui_addToSection:
@@ -1997,7 +2024,12 @@ mjui_update
.. mujoco-include:: mjui_update
Update specific section/item; -1: update all.
This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers)
changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user
(``UiModify`` in :ref:`simulate.cc <saSimulate>`) which also recomputes the layout of all rectangles and associated
auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user
specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated
(which is needed following major changes.)
.. _mjui_event:
@@ -2006,7 +2038,10 @@ mjui_event
.. mujoco-include:: mjui_event
Handle UI event, return pointer to changed item, NULL if no change.
This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item
that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler
implemented by the user (``UiEvent`` in :ref:`simulate.cc <saSimulate>`), and then some action is taken by user code
depending on which UI item was modified and what the state of that item is after the event is handled.
.. _mjui_render:
@@ -2015,7 +2050,9 @@ mjui_render
.. mujoco-include:: mjui_render
Copy UI image to current buffer.
This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If
there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the
time, while :ref:`mjui_update` is called only when changes in the UI take place.
.. _Errorandmemory:
doc/APIreference/functions_override.rst
+75 -6
View File
@@ -2,10 +2,6 @@
This file contains each section text along with function doc overrides. By default the docs use the function doc
pulled from the header files.
.. _Activation:
The functions in this section are maintained for backward compatibility with the now-removed activation mechanism.
.. _Virtualfilesystem:
Virtual file system (VFS) enables the user to load all necessary files in memory, including MJB binary model files, XML
@@ -90,10 +86,40 @@ Solve linear system :math:`M x = y` using factorization: :math:`x = (L^T D L)^{-
Half of linear solve: :math:`x = \sqrt{D^{-1}} (L^T)^{-1} y`
.. _mj_subtreeVel:
Sub-tree linear velocity and angular momentum: compute ``subtree_linvel``, ``subtree_angmom``.
This function is triggered automatically if the subtree :ref:`velocity<sensor-subtreelinvel>` or
:ref:`momentum<sensor-subtreeangmom>` sensors are present in the model.
It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "vel".
.. _mj_rne:
Recursive Newton Euler: compute :math:`M(q) \ddot q + C(q,\dot q)`. ``flg_acc=0`` removes the inertial term (i.e.
assumes :math:`\ddot q = 0`).
.. _mj_rnePostConstraint:
Recursive Newton Euler with final computed forces and accelerations.
Computes three body-level ``nv x 6`` arrays, all defined in the subtreecom-based
:ref:`c-frame<tyNotesCom>` and arranged in ``[rotation(3), translation(3)]`` order.
- ``cacc``: Body acceleration, required for :ref:`mj_objectAcceleration`.
- ``cfrc_int``: Interaction force with the parent body.
- ``cfrc_ext``: External force acting on the body.
This function is triggered automatically if the following sensors are present in the model:
:ref:`accelerometer<sensor-accelerometer>`, :ref:`force<sensor-force>`, :ref:`torque<sensor-torque>`,
:ref:`framelinacc<sensor-framelinacc>`, :ref:`frameangacc<sensor-frameangacc>`.
It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "acc".
The computed force arrays ``cfrc_int`` and ``cfrc_ext`` currently suffer from a know bug, they do not take into account
the effect of spatial tendons, see :github:issue:`832`.
.. _mj_constraintUpdate:
Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set \*cost = s(jar)
where jar = Jac*qacc-aref.
Compute ``efc_state``, ``efc_force``, ``qfrc_constraint``, and (optionally) cone Hessians.
If ``cost`` is not ``NULL``, set ``*cost = s(jar)`` where ``jar = Jac*qacc - aref``.
.. _Support:
@@ -154,6 +180,12 @@ This function can be used to apply a Cartesian force and torque to a point on a
mjData.qfrc_applied of all applied forces. Note that the function requires a pointer to this vector, because sometimes
we want to add the result to a different vector.
.. _mj_objectAcceleration:
Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation. If acceleration or force
sensors are not present in the model, :ref:`mj_rnePostConstraint` must be manually called in order to calculate
mjData.cacc -- the total body acceleration, including contributions from the constraint solver.
.. _mj_differentiatePos:
This function subtracts two vectors in the format of qpos (and divides the result by dt), while respecting the
@@ -218,6 +250,43 @@ of how to use these functions.
.. _UIframework:
For a high-level description of the UI framework, see :ref:`UI`.
.. _mjui_add:
This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs,
each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added
after the end of the last used section. There is also another version of this function
(:ref:`mjui_addToSection<mjui_addToSection>`) which adds items to a specified section instead of adding them at the end
of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by
:ref:`mjMAXUISECT<glNumeric>` and :ref:`mjMAXUIITEM<glNumeric>`. Exceeding these maxima results in low-level errors.
.. _mjui_update:
This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers)
changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user
(``UiModify`` in :ref:`simulate.cc <saSimulate>`) which also recomputes the layout of all rectangles and associated
auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user
specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated
(which is needed following major changes.)
.. _mjui_event:
This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item
that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler
implemented by the user (``UiEvent`` in :ref:`simulate.cc <saSimulate>`), and then some action is taken by user code
depending on which UI item was modified and what the state of that item is after the event is handled.
.. _mjui_render:
This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If
there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the
time, while :ref:`mjui_update` is called only when changes in the UI take place.
.. _Errorandmemory:
.. _Standardmath:
doc/XMLreference.rst
+94 -64
View File
@@ -920,6 +920,11 @@ documented below.
The radius of the capsules used to render slider-crank mechanisms. The second part of the mechanism is automatically
scaled relative to this setting.
.. _visual-scale-frustum:
:at:`frustum`: :at-val:`real, "10"`
The distance of the zfar plane from the camera pinhole for rendering the frustum.
.. _visual-rgba:
@@ -1046,6 +1051,11 @@ disables the rendering of the corresponding object.
Color used to render the crank of slide-crank mechanisms, in model configurations where the specified rod length
cannot be maintained, i.e., it is "broken".
.. _visual-rgba-frustum:
:at:`frustum`: :at-val:`real(4), "1 1 0 0.2"`
Color used to render the camera frustum.
.. _asset:
@@ -1072,9 +1082,9 @@ file.
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.
| MuJoCo 2.0 introduced a second file format for loading textures, in addition to PNG. If the file name extension is
| 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 the new format. This is a custom binary file format, containing the following data:
assumes that the texture is in this format. This is a custom binary file format, containing the following data:
.. code:: Text
@@ -1387,7 +1397,7 @@ attribute of the :ref:`compiler <compiler>` element which controls the automatic
appearance (including texture mapping) is controlled by the :at:`material` and :at:`rgba` attributes of the referencing
geom, similarly to height fields.
Starting with MuJoCo 2.0, meshes can have explicit texture coordinates instead of relying on the automated texture
Meshes can have explicit texture coordinates instead of relying on the automated texture
mapping mechanism. When provided, these explicit coordinates have priority. Note that texture coordinates can be
specified with OBJ files and MSH files, as well as explicitly in the XML with the :at:`texcoord` attribute, but not via
STL files. These mechanism cannot be mixed. So if you have an STL mesh, the only way to add texture coordinates to it is
@@ -1421,17 +1431,17 @@ specified as OBJ or XML and an error message is returned.
The size of the mesh is determined by the 3D coordinates of the vertex data in the mesh file, multiplied by the
components of the :at:`scale` attribute below. Scaling is applied separately for each coordinate axis. Note that
negative scaling values can be used to flip the mesh; this is a legitimate operation. The size parameters of the
referening geoms are ignored, similarly to height fields. As of MuJoCo 2.0 we also provide a mechanism to translate and
rotate the 3D coordinates, using the attributes refpos and refquat.
referening geoms are ignored, similarly to height fields. We also provide a mechanism to translate and
rotate the 3D coordinates, using the attributes :ref:`refpos<asset-mesh-refpos>` and :ref:`refquat<asset-mesh-refquat>`.
Another new feature in MuJoCo 2.0 is that a mesh can be defined without faces (a point cloud essentially). In that case
A mesh can also be defined without faces (a point cloud essentially). In that case
the convex hull is constructed automatically, even if the compiler attribute convexhull is false. This makes it easy to
construct simple shapes directly in the XML. For example, a pyramid can be created as:
.. code-block:: xml
<asset>
<mesh name="pyramid" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
<mesh name="tetrahedron" vertex="0 0 0 1 0 0 0 1 0 0 0 1"/>
</asset>
Positioning and orienting is complicated by the fact that vertex data are often designed relative to coordinate frames
@@ -1869,6 +1879,12 @@ adjust it properly through the XML.
:at:`sdf_initpoints`: :at-val:`int, "40"`
Number of starting points used for fining contacts with Signed Distance Field collisions.
.. _option-actuatorgroupdisable:
:at:`actuatorgroupdisable`: :at-val:`int(30), ""`
List of actuator groups to disable. Actuators whose :ref:`group<actuator-general-group>` is in this list will produce
no force. If they are stateful, their activation states will not be integrated. Internally this list is
implemented as an integer bitfield, so values must be in the range ``0 <= group <= 30``.
.. _option-flag:
@@ -2161,10 +2177,10 @@ This element creates a joint. As explained in :ref:`Kinematic tree <CTree>`, a j
between the body where it is defined and the body's parent. If multiple joints are defined in the same body, the
corresponding spatial transformations (of the body frame relative to the parent frame) are applied in order. If no
joints are defined, the body is welded to its parent. Joints cannot be defined in the world body. At runtime the
positions and orientations of all joints defined in the model are stored in the vector mjData.qpos, in the order in
which the appear in the kinematic tree. The linear and angular velocities are stored in the vector mjData.qvel. These
two vectors have different dimensionality when free or ball joints are used, because such joints represent rotations as
unit quaternions.
positions and orientations of all joints defined in the model are stored in the vector ``mjData.qpos``, in the order in
which the appear in the kinematic tree. The linear and angular velocities are stored in the vector ``mjData.qvel``.
These two vectors have different dimensionality when free or ball joints are used, because such joints represent
rotations as unit quaternions.
.. _body-joint-name:
@@ -2211,7 +2227,8 @@ unit quaternions.
.. _body-joint-pos:
:at:`pos`: :at-val:`real(3), "0 0 0"`
Position of the joint, specified in the frame of the parent body. For free joints this attribute is ignored.
Position of the joint, specified in the frame of the body where the joint is defined.
For free joints this attribute is ignored.
.. _body-joint-axis:
@@ -2654,7 +2671,7 @@ helps clarify the role of bodies and geoms in MuJoCo.
.. _body-geom-pos:
:at:`pos`: :at-val:`real(3), "0 0 0"`
Position of the geom, specified in the frame of the parent body.
Position of the geom, specified in the frame of the body where the geom is defined.
.. _body-geom-quat:
@@ -2698,8 +2715,8 @@ helps clarify the role of bodies and geoms in MuJoCo.
:at:`fluidshape`: :at-val:`[none, ellipsoid], "none"`
"ellipsoid" activates the geom-level fluid interaction model based on an ellipsoidal approximation of the geom
shape. When active, the model based on :ref:`body inertia sizes <flInertia>` is disabled for the parent body.
See section on :ref:`ellipsoid-based<flEllipsoid>` fluid interaction model for details.
shape. When active, the model based on :ref:`body inertia sizes <flInertia>` is disabled for the body in which the
geom is defined. See section on :ref:`ellipsoid-based<flEllipsoid>` fluid interaction model for details.
.. _body-geom-fluidcoef:
@@ -2879,15 +2896,15 @@ and the +Y axis points up. Thus the frame position and orientation are the key a
:at:`mode`: :at-val:`[fixed, track, trackcom, targetbody, targetbodycom], "fixed"`
This attribute specifies how the camera position and orientation in world coordinates are computed in forward
kinematics (which in turn determine what the camera sees). "fixed" means that the position and orientation specified
below are fixed relative to the parent (i.e., the body where the camera is defined). "track" means that the camera
position is at a constant offset from the parent in world coordinates, while the camera orientation is constant in
world coordinates. These constants are determined by applying forward kinematics in qpos0 and treating the camera as
fixed. Tracking can be used for example to position a camera above a body, point it down so it sees the body, and
have it always remain above the body no matter how the body translates and rotates. "trackcom" is similar to "track"
but the constant spatial offset is defined relative to the center of mass of the kinematic subtree starting at the
parent body. This can be used to keep an entire mechanism in view. Note that the subtree center of mass for the world
body is the center of mass of the entire model. So if a camera is defined in the world body in mode "trackcom", it
will track the entire model. "targetbody" means that the camera position is fixed in the parent body, while the
below are fixed relative to the the body where the camera is defined. "track" means that the camera position is at a
constant offset from the body in world coordinates, while the camera orientation is constant in world coordinates.
These constants are determined by applying forward kinematics in qpos0 and treating the camera as fixed. Tracking can
be used for example to position a camera above a body, point it down so it sees the body, and have it always remain
above the body no matter how the body translates and rotates. "trackcom" is similar to "track" but the constant
spatial offset is defined relative to the center of mass of the kinematic subtree starting at the body in which the
camera is defined. This can be used to keep an entire mechanism in view. Note that the subtree center of mass for the
world body is the center of mass of the entire model. So if a camera is defined in the world body in mode "trackcom",
it will track the entire model. "targetbody" means that the camera position is fixed in the body frame, while the
camera orientation is adjusted so that it always points towards the targeted body (which is specified with the target
attribute below). This can be used for example to model an eye that fixates a moving object; the object will be the
target, and the camera/eye will be defined in the body corresponding to the head. "targetbodycom" is the same as
@@ -2965,6 +2982,8 @@ and the +Y axis points up. Thus the frame position and orientation are the key a
:at:`quat`, :at:`axisangle`, :at:`xyaxes`, :at:`zaxis`, :at:`euler`
Orientation of the camera frame. See :ref:`COrientation`.
Note that specifically for cameras, the :at:`xyaxes` attribute is semantically convenient as the X and Y axes
correspond to the directions "right" and "up" in pixel space, respectively.
.. _body-camera-user:
@@ -3636,11 +3655,6 @@ saving the XML:
</mujoco>
.. _body-flexcomp-class:
:at:`class`: :at-val:`string, optional`
Defaults class for setting unspecified attributes.
.. _body-flexcomp-name:
:at:`name`: :at-val:`string, required`
@@ -3689,12 +3703,13 @@ saving the XML:
**gmsh** is similar to mesh, but it loads a `GMSH file <https://gmsh.info//doc/texinfo/gmsh.html#MSH-file-format>`__
in format 4.1 (ascii or binary). The file extension can be anything; the parser recognizes the format by examining
the file header. This is a very rich file format, allowing all kinds of elements with different dimensionality and
topology. MuJoCo only supports GMSH element types 1, 2, 4 which happen to correspond to our 1D, 2D and 3D flexes.
Only the Nodes and Elements sections of the GMHS file are processed, and used to populate the point and element data
of the flexcomp. The parser will generate an error if the GMSH file contains meshes that are not supported by MuJoCo.
:at:`dim` is automatically set to the dimensionality specified in the GMSH file. Presently this is the only mechanism
to load a large tetrahedral mesh in MuJoCo and generate a corresponding soft entity. If such a mesh is available in a
different file format, use the freely available `GMSH software <https://gmsh.info/>`__ to convert it to GMSH 4.1.
topology. MuJoCo only supports GMSH element types 1, 2, 4 which happen to correspond to our 1D, 2D and 3D flexes and
assumes that the nodes are specified in a single block. Only the Nodes and Elements sections of the GMHS file are
processed, and used to populate the point and element data of the flexcomp. The parser will generate an error if the
GMSH file contains meshes that are not supported by MuJoCo. :at:`dim` is automatically set to the dimensionality
specified in the GMSH file. Presently this is the only mechanism to load a large tetrahedral mesh in MuJoCo and
generate a corresponding soft entity. If such a mesh is available in a different file format, use the freely
available `GMSH software <https://gmsh.info/>`__ to convert it to GMSH 4.1.
**direct** allows the user to specify the point and element data of the flexcomp directly in the XML. Note that
flexcomp will still generate moving bodies automatically, as well as automate other settings; so it still provides
@@ -3715,21 +3730,18 @@ saving the XML:
.. _body-flexcomp-point:
:at:`point`: :at-val:`real(3*npoint), optional`
The 3D coordinates of the points. This attribute is only used with type **direct**. All other flexcomp types generate
their own points. The points are used to construct bodies and vertices as explained earlier.
.. _body-flexcomp-element:
:at:`element`: :at-val:`int((dim+1)*npoint), optional`
The zero-based point ids forming each flex elements. This attribute is only used with type **direct**. All other
flexcomp types generate their own elements. This data is passed through to the automatically-generated flex.
.. _body-flexcomp-texcoord:
:at:`texcoord`: :at-val:`real(2*npoint), optional`
Texture coordinates of each point, passed through to the automatically-generated flex. Note that flexcomp does not
generate texture coordinates automatically, except for 2D grids. For all other types, the user can specify explicit
texture coordinates here, even if the points themselves were generated automatically. This requires understanding of
@@ -3779,7 +3791,6 @@ saving the XML:
.. _body-flexcomp-euler:
:at:`axisangle`, :at:`xyaxes`, :at:`zaxis`, :at:`euler`
Alternative specification of rotation, that can be used instead of :at:`quat`.
.. _body-flexcomp-scale:
@@ -3795,9 +3806,8 @@ saving the XML:
.. _body-flexcomp-flatskin:
:at:`radius`, :at:`material`, :at:`rgba`, :at:`group`, :at:`flatskin`
These attributes are directly passed through to the automatically-generated :ref:`flex<deformable-flex>` object and have
the same meaning.
These attributes are directly passed through to the automatically-generated :ref:`flex<deformable-flex>` object and
have the same meaning.
.. _flexcomp-contact:
@@ -3818,10 +3828,12 @@ the same meaning.
.. _flexcomp-contact-margin:
.. _flexcomp-contact-gap:
:at:`internal`, :at:`selfcollide`, :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| 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`
Same as in :ref:`flex/contact<flex-contact>`. All attributes are passed through to the automatically-generated flex.
|body/flexcomp/contact attrib list|
Same as in :ref:`flex/contact<flex-contact>`. All attributes are passed through to the automatically-generated flex.
.. _flexcomp-edge:
@@ -4119,7 +4131,7 @@ cases, the user will specify a :el:`flexcomp` which will then automatically cons
.. _flex-edge:
:el-prefix:`flex/` |-| **edge** (?)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
'''''''''''''''''''''''''''''''''''
This element adjusts the passive or constraint properties of all edges of the flex. A flex edge can have a damping
passive force and an :ref:`equality constraint<equality-flex>` associated with it, resulting in edge constraint forces.
@@ -4141,7 +4153,7 @@ these mechanisms to be combined as desired.
.. _flex-contact:
:el-prefix:`flex/` |-| **contact** (?)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
''''''''''''''''''''''''''''''''''''''
This element adjusts the contact properties of the flex. It is mostly identical to geom contact properties, with some
extensions specific to flexes.
@@ -4225,7 +4237,7 @@ Finally, the skin can be inflated by applying an offset to each vertex position
Skins are one-sided for rendering purposes; this is because back-face culling is needed to avoid shading and aliasing
artifacts. When the skin is a closed 3D shape this does not matter because the back sides cannot be seen. But if the
skin is a 2D object, we have to specify both sides and offset them slightly to avoid artifacts. Note that the
composite objects introduced in MuJoCo 2.0 generate skins automatically. So one can save an XML model with a composite
composite objects generate skins automatically. So one can save an XML model with a composite
object, and obtain an elaborate example of how a skin is specified in the XML.
Similar to meshes, skins can be specified directly in the XML via attributes documented later, or loaded from a binary
@@ -4475,11 +4487,12 @@ of the other body, without any joint elements in the child body.
.. _equality-weld-torquescale:
:at:`torquescale`: :at-val:`real, "1"`
Relative torque-to-force ratio. This ratio is used by the weld to scale how much it "cares" about rotational
displacements vs. translational displacements. Setting this value to 0 makes the :el:`weld` behave like a
:el:`connect` constraint. Note that this value has units of length and can therefore be interpreted as follows.
Imagining that the weld is implemented by a patch of glue sticking the two bodies together, :at:`torquescale` can be
interpreted as the diameter of this glue patch.
A constant that scales the angular residual (angular constraint violation). Notionally in units of
:math:`\textrm{torque}/\textrm{force}=\textrm{length}`. Intuitively this coefficient defines how much the weld
"cares" about rotational displacements vs. translational displacements. Setting this value to 0 makes the :el:`weld`
behave like a :el:`connect` constraint. Note that this value has units of length and can therefore be understood as
follows. Imagining that the weld is implemented by a flat patch of glue sticking the two bodies together,
:at:`torquescale` can be interpreted as the diameter of this glue patch.
.. _equality-joint:
@@ -4623,7 +4636,7 @@ has multiple obstacle geoms they must be separated by sites - so as to avoid the
tendon level. This example illustrates a multi-branch tendon acting as a finger extensor, with a counter-weight
instead of an actuator: `tendon.xml <_static/tendon.xml>`__.
MuJoCo 2.0 introduced a second form of wrapping, where the tendon is constrained to pass through a geom rather than
A second form of wrapping is where the tendon is constrained to pass *through* a geom rather than
wrap around it. This is enabled automatically when a sidesite is specified and its position is inside the volume of
the obstacle geom.
@@ -4871,7 +4884,7 @@ multiplied by the corresponding coef value, and added up to obtain the tendon le
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This is a grouping element for actuator definitions. Recall the discussion of MuJoCo's :ref:`Actuation model
<geActuation>` in the Computation chapter, and the :ref:`Actuator shortcuts <CActuator>` discussed earlier in this
<geActuation>` in the Computation chapter, and the :ref:`Actuator shortcuts <CActShortcuts>` discussed earlier in this
chapter. The first 13 attributes of all actuator-related elements below are the same, so we document them only once,
under the :el:`general` actuator.
@@ -5144,7 +5157,7 @@ specify them independently.
:el-prefix:`actuator/` |-| **motor** (*)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This and the next three elements are the :ref:`Actuator shortcuts <CActuator>` discussed earlier. When a
This and the next three elements are the :ref:`Actuator shortcuts <CActShortcuts>` discussed earlier. When a
such shortcut is encountered, the parser creates a :el:`general` actuator and sets its dynprm, gainprm and biasprm
attributes to the internal defaults shown above, regardless of any default settings. It then adjusts dyntype, gaintype
and biastype depending on the shortcut, parses any custom attributes (beyond the common ones), and translates them
@@ -5805,9 +5818,9 @@ Associate this actuator with an :ref:`engine plugin<exPlugin>`. Either :at:`plug
.. _actuator-plugin-user:
.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`ctrllimited`, :at:`forcelimited`
:at:`ctrlrange`, :at:`forcerange`, :at:`lengthrange`, :at:`gear`, :at:`cranklength`, :at:`joint`, :at:`jointinparent`
:at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`
.. |actuator/plugin attrib list| replace:: :at:`name`, :at:`class`, :at:`group`, :at:`ctrllimited`,
:at:`forcelimited`, :at:`ctrlrange`, :at:`forcerange`, :at:`lengthrange`, :at:`gear`, :at:`cranklength`,
:at:`joint`, :at:`jointinparent`, :at:`site`, :at:`tendon`, :at:`cranksite`, :at:`slidersite`, :at:`user`
|actuator/plugin attrib list|
Same as in actuator/ :ref:`general <actuator-general>`.
@@ -5867,6 +5880,8 @@ This element creates a 3-axis accelerometer. The sensor is mounted at a site, an
as the site frame. This sensor outputs three numbers, which are the linear acceleration of the site (including gravity)
in local coordinates.
The presence of this sensor in a model triggers a call to :ref:`mj_rnePostConstraint` during sensor computation.
.. _sensor-accelerometer-name:
.. _sensor-accelerometer-noise:
@@ -5947,6 +5962,8 @@ the child body, and the force points from the child towards the parent. The comp
forces acting on the system, including contacts as well as external perturbations. Using this sensor often requires
creating a dummy body welded to its parent (i.e., having no joint elements).
The presence of this sensor in a model triggers a call to :ref:`mj_rnePostConstraint` during sensor computation.
.. _sensor-force-name:
.. _sensor-force-noise:
@@ -5976,6 +5993,8 @@ creating a dummy body welded to its parent (i.e., having no joint elements).
This element creates a 3-axis torque sensor. This is similar to the :ref:`force <sensor-force>` sensor above, but
measures torque rather than force.
The presence of this sensor in a model triggers a call to :ref:`mj_rnePostConstraint` during sensor computation.
.. _sensor-torque-name:
.. _sensor-torque-noise:
@@ -6779,6 +6798,8 @@ coordinates.
This element creates a sensor that returns the 3D linear acceleration of the spatial frame of the object, in global
coordinates.
The presence of this sensor in a model triggers a call to :ref:`mj_rnePostConstraint` during sensor computation.
.. _sensor-framelinacc-name:
.. _sensor-framelinacc-noise:
@@ -6809,6 +6830,8 @@ coordinates.
This element creates a sensor that returns the 3D angular acceleration of the spatial frame of the object, in global
coordinates.
The presence of this sensor in a model triggers a call to :ref:`mj_rnePostConstraint` during sensor computation.
.. _sensor-frameangacc-name:
.. _sensor-frameangacc-noise:
@@ -6864,6 +6887,8 @@ global coordinates.
This element creates sensor that returns the linear velocity of the center of mass of the kinematic subtree rooted at a
specified body, in global coordinates.
The presence of this sensor in a model triggers a call to :ref:`mj_subtreeVel` during sensor computation.
.. _sensor-subtreelinvel-name:
.. _sensor-subtreelinvel-noise:
@@ -6889,6 +6914,8 @@ specified body, in global coordinates.
This element creates sensor that returns the angular momentum around the center of mass of the kinematic subtree rooted
at a specified body, in global coordinates.
The presence of this sensor in a model triggers a call to :ref:`mj_subtreeVel` during sensor computation.
.. _sensor-subtreeangmom-name:
.. _sensor-subtreeangmom-noise:
@@ -6931,12 +6958,15 @@ This element creates sensor that returns the simulation time.
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This element creates a user sensor. MuJoCo does not know how to compute the output of this sensor. Instead the user
should install the callback :ref:`mjcb_sensor` which is expected to fill in the sensor data in mjData.sensordata. The
specification in the XML is used to allocate space for this sensor, and also determine which MuJoCo object it is
should install the callback :ref:`mjcb_sensor` which is expected to fill in the sensor data in ``mjData.sensordata``.
The specification in the XML is used to allocate space for this sensor, and also determine which MuJoCo object it is
attached to and what stage of computation it needs before the data can be computed. Note that the MuJoCo object
referenced here can be a tuple, which in turn can reference a custom collection of MuJoCo objects - for example several
referenced here can be a tuple, which in turn can reference a custom collection of MuJoCo objects -- for example several
bodies whose center of mass is of interest.
If a user sensor is of :ref:`stage<sensor-user-needstage>` "vel" or "acc", then :ref:`mj_subtreeVel` or
:ref:`mj_rnePostConstraint` will be triggered, respectively.
.. _sensor-user-name:
.. _sensor-user-noise:
@@ -7009,7 +7039,7 @@ Ascociate this sensor with an :ref:`engine plugin<exPlugin>`. Either :at:`plugin
.. _sensor-plugin-user:
.. |sensor/plugin attrib list| replace:: :at:`name`, :at:`cutoff`, :at:`objtype`, :at:`objname`, :at:`reftype`
.. |sensor/plugin attrib list| replace:: :at:`name`, :at:`cutoff`, :at:`objtype`, :at:`objname`, :at:`reftype`,
:at:`refname`, :at:`user`
|sensor/plugin attrib list|
@@ -7514,7 +7544,7 @@ if omitted.
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
This and the next three elements set the attributes of the :ref:`general <actuator-general>` element using
:ref:`Actuator shortcuts <CActuator>`. It does not make sense to use more than one such shortcut in the same defaults
:ref:`Actuator shortcuts <CActShortcuts>`. It does not make sense to use more than one such shortcut in the same defaults
class, because they set the same underlying attributes, replacing any previous settings. All
:ref:`motor <actuator-motor>` attributes are available here except: name, class, joint, jointinparent, site, tendon,
slidersite, cranksite.
doc/XMLschema.rst
+10 -10
View File
@@ -109,6 +109,8 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`framelength<visual-scale-framelength>` | :ref:`framewidth<visual-scale-framewidth>` | :ref:`constraint<visual-scale-constraint>` | :ref:`slidercrank<visual-scale-slidercrank>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`frustum<visual-scale-frustum>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| visual |br| |_| |L| | | .. table:: |
| :ref:`rgba | ? | :class: mjcf-attributes |
@@ -124,7 +126,7 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`contacttorque<visual-rgba-contacttorque>` | :ref:`contactgap<visual-rgba-contactgap>` | :ref:`rangefinder<visual-rgba-rangefinder>` | :ref:`constraint<visual-rgba-constraint>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`slidercrank<visual-rgba-slidercrank>` | :ref:`crankbroken<visual-rgba-crankbroken>` | | | |
| | | | :ref:`slidercrank<visual-rgba-slidercrank>` | :ref:`crankbroken<visual-rgba-crankbroken>` | :ref:`frustum<visual-rgba-frustum>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| mujoco |br| |L| | | *no attributes* |
@@ -228,7 +230,7 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`iterations<option-iterations>` | :ref:`ls_iterations<option-ls_iterations>` | :ref:`noslip_iterations<option-noslip_iterations>` | :ref:`mpr_iterations<option-mpr_iterations>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`sdf_iterations<option-sdf_iterations>` | :ref:`sdf_initpoints<option-sdf_initpoints>` | | | |
| | | | :ref:`sdf_iterations<option-sdf_iterations>` | :ref:`sdf_initpoints<option-sdf_initpoints>` | :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_| option |br| |_| |L| | | .. table:: |
@@ -485,19 +487,17 @@
| :ref:`flexcomp | \* | :class: mjcf-attributes |
| <body-flexcomp>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<body-flexcomp-name>` | :ref:`class<body-flexcomp-class>` | :ref:`type<body-flexcomp-type>` | :ref:`group<body-flexcomp-group>` | |
| | | | :ref:`name<body-flexcomp-name>` | :ref:`type<body-flexcomp-type>` | :ref:`group<body-flexcomp-group>` | :ref:`dim<body-flexcomp-dim>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`dim<body-flexcomp-dim>` | :ref:`count<body-flexcomp-count>` | :ref:`spacing<body-flexcomp-spacing>` | :ref:`radius<body-flexcomp-radius>` | |
| | | | :ref:`count<body-flexcomp-count>` | :ref:`spacing<body-flexcomp-spacing>` | :ref:`radius<body-flexcomp-radius>` | :ref:`rigid<body-flexcomp-rigid>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`rigid<body-flexcomp-rigid>` | :ref:`mass<body-flexcomp-mass>` | :ref:`inertiabox<body-flexcomp-inertiabox>` | :ref:`scale<body-flexcomp-scale>` | |
| | | | :ref:`mass<body-flexcomp-mass>` | :ref:`inertiabox<body-flexcomp-inertiabox>` | :ref:`scale<body-flexcomp-scale>` | :ref:`file<body-flexcomp-file>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`file<body-flexcomp-file>` | :ref:`point<body-flexcomp-point>` | :ref:`element<body-flexcomp-element>` | :ref:`texcoord<body-flexcomp-texcoord>` | |
| | | | :ref:`point<body-flexcomp-point>` | :ref:`element<body-flexcomp-element>` | :ref:`texcoord<body-flexcomp-texcoord>` | :ref:`material<body-flexcomp-material>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`material<body-flexcomp-material>` | :ref:`rgba<body-flexcomp-rgba>` | :ref:`flatskin<body-flexcomp-flatskin>` | :ref:`pos<body-flexcomp-pos>` | |
| | | | :ref:`rgba<body-flexcomp-rgba>` | :ref:`flatskin<body-flexcomp-flatskin>` | :ref:`pos<body-flexcomp-pos>` | :ref:`quat<body-flexcomp-quat>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`quat<body-flexcomp-quat>` | :ref:`axisangle<body-flexcomp-axisangle>` | :ref:`xyaxes<body-flexcomp-xyaxes>` | :ref:`zaxis<body-flexcomp-zaxis>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`euler<body-flexcomp-euler>` | | | | |
| | | | :ref:`axisangle<body-flexcomp-axisangle>` | :ref:`xyaxes<body-flexcomp-xyaxes>` | :ref:`zaxis<body-flexcomp-zaxis>` | :ref:`euler<body-flexcomp-euler>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
doc/_static/gyroscopic.xml
Vendored
+22
View File
@@ -0,0 +1,22 @@
<mujoco>
<asset>
<texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
<texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
</asset>
<statistic extent="1"/>
<!-- switch the integrator to "implicit" for better gyroscopic stabillity -->
<option timestep="0.01" integrator="implicitfast"/>
<worldbody>
<geom pos="0 0 -.5" type="plane" size="100 100 .1" zaxis=".5 0 1" material="grid"/>
<body pos="0 0 .3">
<freejoint/>
<geom type="ellipsoid" size="0.15 0.1 0.07" euler="20 40 60"/>
<light pos="0 0 1.5" mode="track"/>
<camera pos="0 -1.498 0.1" xyaxes="1 0 0 0 0 1" mode="track"/>
</body>
</worldbody>
</mujoco>
doc/_static/pendulum.xml
Vendored
+27
View File
@@ -0,0 +1,27 @@
<mujoco model="energy conserving pendulum">
<option integrator="RK4">
<flag constraint="disable" energy="enable"/>
</option>
<worldbody>
<light pos="0 0 2"/>
<geom pos="0 0 -.5" type="plane" size="1 1 .01"/>
<body pos="0 0 0">
<joint type="hinge" axis="0 1 0"/>
<geom type="cylinder" size="0.02" fromto="0 -.02 0 0 .02 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body pos="0.1 0 0">
<joint type="slide" axis="1 0 0" stiffness="200"/>
<geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
<body pos=".1 0 0">
<joint type="ball"/>
<geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
<body pos="0 .1 0">
<joint axis="1 0 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 0 .1 0"/>
</body>
</body>
</body>
</body>
</worldbody>
</mujoco>
doc/changelog.rst
+81 -4
View File
@@ -2,13 +2,90 @@
Changelog
=========
Upcoming version (not yet released)
-----------------------------------
Version 3.0.1 (November 15, 2023)
---------------------------------
General
^^^^^^^
1. Added sub-terms of total passive forces in ``mjData.qfrc_passive`` to :ref:`mjData`:
``qfrc_{spring, damper, gravcomp, fluid}``. The sum of these vectors equals ``qfrc_passive``.
.. youtube:: H9qG9Zf2W44
:align: right
:width: 240px
2. Added :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` attribute and associated
:ref:`mjOption.disableactuator<mjOption>` integer bitfield, which can be used to disable sets of actuators at runtime
according to their :ref:`group<actuator-general-group>`. Fixes :github:issue:`1092`. See :ref:`CActDisable`.
- The first 6 actuator groups are toggleable in the :ref:`simulate<saSimulate>` viewer. See `example model
<https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
and associated screen-capture on the right.
3. Increased ``mjMAXUIITEM`` (maximum number of UI elements per section in Simulate) to 200.
MJX
^^^
4. Added support for Newton solver (``mjSOL_NEWTON`` in :ref:`mjtSolver`). The Newton solver significantly speeds up
simulation on GPU:
.. list-table:: Steps-per-second, Conjugate Gradient vs. Newton on A100
:header-rows: 1
:align: left
* - Model
- CG
- Newton
- Speedup
* - `Humanoid <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/humanoid>`__
- 640,000
- 1,020,000
- **1.6 x**
* - `Barkour v0 <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/barkour_v0>`__
- 1,290,000
- 1,750,000
- **1.35 x**
* - `Shadow Hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
- 215,000
- 270,000
- **1.25 x**
Humanoid is the standard MuJoCo humanoid,
`Google Barkour <https://blog.research.google/2023/05/barkour-benchmarking-animal-level.html>`__ and the Shadow Hand
are both available in the :ref:`MuJoCo Menagerie<Menagerie>`.
5. Added support for joint equality constraints (``mjEQ_JOINT`` in :ref:`mjtEq`).
6. Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly.
7. Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`).
8. Removed empty EFC rows from ``MJX``, for joints with no limits (fixes :github:issue:`1117`).
9. Fixed bug in ``scan.body_tree`` that led to incorrect smooth dynamics for some kinematic tree layouts.
Python bindings
^^^^^^^^^^^^^^^
10. Fix the macOS ``mjpython`` launcher to work with the Python interpreter from Apple Command Line Tools.
11. Fixed a crash when copying instances of ``mujoco.MjData`` for models that use plugins. Introduced a ``model``
attribute to ``MjData`` which is reference to the model that was used to create that ``MjData`` instance.
Simulate
^^^^^^^^
12. :ref:`simulate<saSimulate>`: correct handling of "Pause update", "Fullscreen" and "VSync" buttons.
Documentation
^^^^^^^^^^^^^
.. youtube:: cE3s_IfO4g4
:align: right
:width: 240px
13. Added cell to the `tutorial colab <https://github.com/google-deepmind/mujoco#getting-started>`__ providing an
example of procedural camera control:
14. Added documentation for the :ref:`UI` framework.
15. Fixed typos and supported fields in docs (fixes :github:issue:`1105` and :github:issue:`1106`).
Bug fixes
^^^^^^^^^
1. Fix in simulate: correct handling of "Pause update", "Fullscreen" and "VSync" buttons.
16. Fixed bug relating to welds modified with :ref:`torquescale<equality-weld-torquescale>`.
Version 3.0.0 (October 18, 2023)
--------------------------------
doc/computation/index.rst
+213 -187
View File
@@ -143,47 +143,59 @@ Our notation is summarized in the table below. Additional notation specific to c
When available, we also show the fields of main data structures :ref:`mjModel` and :ref:`mjData` corresponding to the
mathematical notation.
+-----------------+----------------+----------------+----------------------+
| Symbol | Size | Description | MuJoCo field |
+=================+================+================+======================+
| :math:`n_Q` | | number of | ``mjModel.nq`` |
| | | position | |
| | | coordinates | |
+-----------------+----------------+----------------+----------------------+
| :math:`n_V` | | number of | ``mjModel.nv`` |
| | | degrees of | |
| | | freedom | |
+-----------------+----------------+----------------+----------------------+
| :math:`n_C` | | number of | ``mjData.nefc`` |
| | | active | |
| | | constraints | |
+-----------------+----------------+----------------+----------------------+
| :math:`q` | :math:`n_Q` | joint position | ``mjData.qpos`` |
+-----------------+----------------+----------------+----------------------+
| :math:`v` | :math:`n_V` | joint velocity | ``mjData.qvel`` |
+-----------------+----------------+----------------+----------------------+
| :math:`\tau` | :math:`n_V` | applied force: | |
| | | passive, | |
| | | actuation, | |
| | | external | |
+-----------------+----------------+----------------+----------------------+
| :math:`c(q, v)` | :math:`n_V` | bias force: | ``mjData.qfrc_bias`` |
| | | Coriolis, | |
| | | centrifugal, | |
| | | gravitational | |
+-----------------+----------------+----------------+----------------------+
| :math:`M(q)` | :math:`n_V | inertia in | ``mjData.qM`` |
| | \times n_V` | joint space | |
+-----------------+----------------+----------------+----------------------+
| :math:`J(q)` | :math:`n_C | constraint | ``mjData.efc_J`` |
| | \times n_V` | Jacobian | |
+-----------------+----------------+----------------+----------------------+
| :math:`r(q)` | :math:`n_C` | constraint | ``mjData.efc_pos`` |
| | | residual | |
+-----------------+----------------+----------------+----------------------+
| :math:`f(q, v, | :math:`n_C` | constraint | ``mjData.efc_force`` |
| \tau)` | | force | |
+-----------------+----------------+----------------+----------------------+
.. list-table::
:widths: 2 2 7 4
:header-rows: 1
* - Symbol
- Size
- Description
- MuJoCo field
* - :math:`n_Q`
-
- number of position coordinates
- ``mjModel.nq``
* - :math:`n_V`
-
- number of degrees of freedom
- ``mjModel.nv``
* - :math:`n_C`
-
- number of active constraints
- ``mjData.nefc``
* - :math:`q`
- :math:`n_Q`
- joint position
- ``mjData.qpos``
* - :math:`v`
- :math:`n_V`
- joint velocity
- ``mjData.qvel``
* - :math:`\tau`
- :math:`n_V`
- applied force: passive, actuation, external
- ``mjData.qfrc_passive`` + ``mjData.qfrc_actuator`` + ``mjData.qfrc_applied``
* - :math:`c(q, v)`
- :math:`n_V`
- bias force: Coriolis, centrifugal, gravitational
- ``mjData.qfrc_bias``
* - :math:`M(q)`
- :math:`n_V \times n_V`
- inertia in joint space
- ``mjData.qM``
* - :math:`J(q)`
- :math:`n_C \times n_V`
- constraint
Jacobian
- ``mjData.efc_J``
* - :math:`r(q)`
- :math:`n_C`
- constraint residual
- ``mjData.efc_pos``
* - :math:`f(q, v,\tau)`
- :math:`n_C`
- constraint force
- ``mjData.efc_force``
All model elements are enumerated at compile time and assembled into the above system-level vectors and matrices. In our
earlier arm model :ref:`example <Examples>` the model has :math:`n_V = 13` degrees of freedom: 3 for the ball joint, one
@@ -259,7 +271,7 @@ the force outputs are stored in ``mjData.actuator_force``, and the activation st
These three components of an actuator - transmission, activation dynamics, and force generation - determine how the
actuator works. The user can set them independently for maximum flexibility, or use :ref:`Actuator shortcuts
<CActuator>` which instantiate common actuator types.
<CActShortcuts>` which instantiate common actuator types.
.. _geTransmission:
@@ -304,8 +316,8 @@ is attached; the possible attachment object types are :at:`joint`, :at:`tendon`,
.. _geActivation:
Activation dynamics
^^^^^^^^^^^^^^^^^^^
Stateful actuators
^^^^^^^^^^^^^^^^^^
Some actuators such as pneumatic and hydraulic cylinders as well as biological muscles have an internal state called
"activation". This is a true dynamic state, beyond the joint positions :math:`q` and velocities :math:`v`. Including
@@ -321,27 +333,33 @@ independent of the other actuators. The activation types currently implemented a
.. math::
\begin{aligned}
\text{integrator}: & & \dot{w}_i &= u_i \\
\text{filter}: & & \dot{w}_i &= (u_i - w_i) / t \\
\text{filterexact}: & & \dot{w}_i &= (u_i - w_i) / t \\
\text{filter}: & & \dot{w}_i &= (u_i - w_i) / \texttt{t} \\
\text{filterexact}: & & \dot{w}_i &= (u_i - w_i) / \texttt{t} \\
\text{muscle}: & & \dot{w}_i &= \textrm{muscle}(u_i, w_i, l_i, \dot{l}_i)
\end{aligned}
where :math:`t` is an actuator-specific time constant stored in ``mjModel.actuator_dynprm``. In addition the type can
be "user", in which case :math:`w_i` is computed by the user-defined callback :ref:`mjcb_act_dyn`. The type can also
be "none" which corresponds to a regular actuator with no activation state. The dimensionality of :math:`w` equals
where :math:`\texttt{t}` is an actuator-specific time-constant stored in ``mjModel.actuator_dynprm``. In addition, the
type can be "user", in which case :math:`w_i` is computed by the user-defined callback :ref:`mjcb_act_dyn`. The type can
also be "none" which corresponds to a regular actuator with no activation state. The dimensionality of :math:`w` equals
the number of actuators whose activation type is different from "none".
For more information regarding muscle activation dynamics, see :ref:`CMuscle`.
For ``filterexact`` activation dynamics, Euler integration of :math:`\dot{w}` is replaced with the analytic integral:
.. math::
\begin{aligned}
\text{filter}: & & w_{i+1} &= w_i + h (u_i - w_i) / t \\
\text{filterexact}: & & w_{i+1} &= w_i + (u_i - w_i) (1 - e^{-h / t}) \\
\text{filter}: & & w_{i+1} &= w_i + h (u_i - w_i) / \texttt{t} \\
\text{filterexact}: & & w_{i+1} &= w_i + (u_i - w_i) (1 - e^{-h / \texttt{t}}) \\
\end{aligned}
The two expressions converge to the same value in the :math:`h \rightarrow 0` limit.
The two expressions converge to the same value in the :math:`h \rightarrow 0` limit. Note that Euler-integrated filters
diverge for :math:`\texttt{t} < h`, while exactly-integrated filters are stable for any positive :math:`\texttt{t}`.
Note that Euler-integrated filters diverge for :math:`t < h`, while exactly-integrated filters are stable for any
:math:`t > 0`.
:ref:`actearly<actuator-general-actearly>`:
If the :ref:`actearly<actuator-general-actearly>` attribute is set to "true", ``mjData.actuator_force`` is computed
based on :math:`w_{i+1}` (the next activation), reducing the delay between changes to :math:`u` and their effects on
the acceleration by one time step (so the total dynamics are second-order rather than third order).
.. _geActuatorForce:
@@ -379,10 +397,6 @@ This quantity is stored in ``mjData.qfrc_actuator``. It is added to the applied
with any user-defined forces in joint or Cartesian coordinates (which are stored in ``mjData.qfrc_applied`` and
``mjData.xfrc_applied`` respectively).
Optionally, the :ref:`actearly<actuator-general-actearly>` attribute on an actuator computes ``mjData.qfrc_actuator``
based on the value of :math:`w_{i+1}` after integration, reducing the delay between changes to :math:`u` and
:math:`t`.
.. _gePassive:
Passive forces
@@ -526,29 +540,48 @@ Fast implicit-in-velocity (``implicitfast``)
derivatives are also the main source of asymmetry of :math:`D`, by dropping them and symmetrizing, we can use the
faster Cholesky rather than LU decomposition.
.. tip::
The implicitfast integrator has similar computational cost to Euler, yet provides increased stability, and is
therefore a strict improvement. It is the recommended integrator and will become the default in a future version.
4th-order Runge-Kutta (``RK4``)
One advantage of our continuous-time formulation is that we can use higher order integrators such as Runge-Kutta or
multistep methods. The only such integrator currently implemented is the fixed-step `4th-order Runge-Kutta method
<https://en.wikipedia.org/wiki/Runge–Kutta_methods#Derivation_of_the_Runge–Kutta_fourth-order_method>`_, though users
can easily implement other integrators by calling :ref:`mj_forward` and integrating accelerations themselves. We have
observed that for energy-conserving systems (`example
<https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/derivative/energy_conserving_pendulum.xml>`_) RK4
is qualitatively better than the single-step methods, both in terms of stability and accuracy, even when the timestep
is decreased by a factor of 4 (so the computational effort is identical). In the presence of large velocity-
dependent forces, if the chosen single-step method integrates those forces implicitly, single-step methods can be
significantly more stable than RK4.
<https://en.wikipedia.org/wiki/Runge–Kutta_methods#Derivation_of_the_Runge–Kutta_fourth-order_method>`__, though
users can easily implement other integrators by calling :ref:`mj_forward` and integrating accelerations themselves.
We have observed that for energy-conserving systems (`example <../_static/pendulum.xml>`__), RK4 is qualitatively
better than the single-step methods, both in terms of stability and accuracy, even when the timestep is decreased by
a factor of 4 (so the computational effort is identical). In the presence of large velocity- dependent forces, if the
chosen single-step method integrates those forces implicitly, single-step methods can be significantly more stable
than RK4.
.. note::
The accuracy and stability of all integrators can be improved by reducing the time step :math:`h` which is stored in
``mjModel.opt.timestep``. Of course this also slows down the simulation. The time step is perhaps the most important
parameter that the user can adjust. If it is too large, the simulation will become unstable. If it is too small, CPU
time will be wasted without meaningful improvement in accuracy. There is always a comfortable range where the time
step is "just right", but that range is model-dependent.
.. admonition:: Choosing timestep and integrator
:class: tip
:ref:`timestep<option-timestep>`
The accuracy and stability of all integrators can be improved by reducing the time step :math:`h`.
Of course a smaller time step also slows down the simulation. The time step is perhaps the single most important
parameter that the user can adjust. If it is too large, the simulation will become unstable. If it is too small, CPU
time will be wasted without meaningful improvement in accuracy. There is always a comfortable range where the time
step is "just right", but that range is model-dependent.
:ref:`integrator<option-integrator>`
Summary: The recommended integrator is ``implicitfast`` which usually has the best tradeoff of stabillity and
performance.
**Euler**:
Use ``Euler`` for compatibillity with older models and :ref:`MJX<Mjx>`. Specifically for MJX,
setting the :ref:`eulerdamp<option-flag-eulerdamp>` disable flag can :ref:`improve performance<MjxPerformance>`.
**implicitfast**:
The ``implicitfast`` integrator has similar computational cost to ``Euler``, yet provides
increased stability, and is therefore a strict improvement. It is the recommended integrator for most models.
**implicit**:
The benefit over ``implicitfast`` is the implicit integration of Coriolis and centripetal forces, including
gyroscopic forces. The most common case where integrating such forces implicitly leads to noticable improvement is
when free objects with assymetric inertia are spinning quickly. `gyroscopic.xml <../_static/gyroscopic.xml>`__
shows an ellipsoid rolling on an inclined plane which quickly diverges with ``implicitfast`` but is stable with
``implicit``.
**RK4**:
This integrator is best for systems which are energy conserving, or almost energy-conserving. `pendulum.xml
<../_static/pendulum.xml>`__ shows a complicated pendulum mechanism which diverges quickly using ``Euler`` or
``implicitfast`` yet conserves energy well under ``RK4``. Note that under ``implicit``, this model doesn't diverge
but rather loses energy.
.. _geState:
@@ -850,33 +883,25 @@ spatial frame and normal distance are given by the collision detector.
In addition to the above quantities which are computed online, each contact has several parameters obtained from the
model definition.
+---------------------------+-----------------------------------+
| Parameter | Description |
+===========================+===================================+
| ``condim`` | Dimensionality of the contact |
| | force/torque in the contact |
| | frame. It can be 1, 3, 4 or 6. |
+---------------------------+-----------------------------------+
| ``friction`` | Vector of friction coefficients, |
| | with dimensionality ``condim-1``. |
+---------------------------+-----------------------------------+
| ``margin`` | The distance margin used to |
| | determine if the contact should |
| | be included in the global contact |
| | array ``mjData.contact``. |
+---------------------------+-----------------------------------+
| ``gap`` | For custom computations it is |
| | sometimes convenient to include |
| | contacts in ``mjData.contact`` |
| | but not generate contact forces. |
| | This is what ``gap`` does: |
| | contact forces are generated only |
| | when the normal distance is below |
| | margin-gap. |
+---------------------------+-----------------------------------+
| ``solref`` and ``solimp`` | :ref:`Solver <Solver>` |
| | parameters explained later. |
+---------------------------+-----------------------------------+
.. list-table::
:widths: 1 5
:header-rows: 1
* - Parameter
- Description
* - ``condim``
- Dimensionality of the contact force/torque in the contact frame. |br| It can be 1, 3, 4 or 6.
* - ``friction``
- Vector of friction coefficients with dimensionality ``condim-1``.
* - ``margin``
- The distance margin used to determine if the contact should be included in the global contact array
``mjData.contact``.
* - ``gap``
- For custom computations it is sometimes convenient to include contacts in ``mjData.contact`` but not generate
contact forces. This is what ``gap`` does: contact forces are generated only when the normal distance is below
(margin - gap).
* - ``solref`` and ``solimp``
- :ref:`Solver <Solver>` parameters, explained later.
The contact friction cone can be either elliptic or pyramidal. This is a global setting determined by the choice of
constraint solver: the elliptic solvers work with elliptic cones, while the pyramidal solvers work with pyramidal cones,
@@ -967,11 +992,11 @@ to be solved numerically. In inverse dynamics, the problem becomes diagonal and
The primal formulation is based on a generalization of the Gauss principle of least constraint. In its basic form, the
Gauss principle states that if we have unconstrained dynamics :math:`M \dot{v} = \tau` and impose acceleration
constraint :math:`J \dot{v} = a^*`, the resulting acceleration will be
constraint :math:`J \dot{v} = \ar`, the resulting acceleration will be
.. math::
\dot{v} = \arg \min_x \left\| x-M^{-1} \tau \right\|^2_M \\
\textrm{subject to} \; J x = a^*
\textrm{subject to} \; J x = \ar
where the weighted :math:`L_2` norm is the usual :math:`\|x\|^2_M = x^T M x`. Thus the constraint causes the smallest
possible deviation from the unconstrained acceleration :math:`M^{-1}\tau`, where the metric for measuring deviations in
@@ -981,58 +1006,55 @@ will be done by generalizing both the cost function and the constraints in the G
We will use the following notation beyond the notation introduced earlier:
+----------------------+----------------------+----------------------+
| Symbol | Size | Description |
+======================+======================+======================+
| :math:`z` | :math:`n_C` | constraint |
| | | deformations |
+----------------------+----------------------+----------------------+
| :math:`\omega` | :math:`n_C` | velocity of |
| | | constraint |
| | | deformations |
+----------------------+----------------------+----------------------+
| :math:`d` | :math:`n_C` | constraint impedance |
+----------------------+----------------------+----------------------+
| :math:`b` | :math:`n_C` | virtual constraint |
| | | damping |
+----------------------+----------------------+----------------------+
| :math:`k` | :math:`n_C` | virtual constraint |
| | | stiffness |
+----------------------+----------------------+----------------------+
| :math:`A(q)` | :math:`n_C \times | inverse inertia in |
| | n_C` | constraint space |
+----------------------+----------------------+----------------------+
| :math:`R(q)` | :math:`n_C \times | diagonal regularizer |
| | n_C` | in constraint space |
+----------------------+----------------------+----------------------+
| :math:`a^*(q,v)` | :math:`n_C` | reference |
| | | acceleration in |
| | | constraint space |
+----------------------+----------------------+----------------------+
| :math:`a^0(q, v, | :math:`n_C` | unconstrained |
| \tau)` | | acceleration in |
| | | constraint space |
+----------------------+----------------------+----------------------+
| :math:`a^1(q, v, | :math:`n_C` | constrained |
| \dot{v})` | | acceleration in |
| | | constraint space |
+----------------------+----------------------+----------------------+
| :math:`\mathcal{K} | | product of all |
| (q)` | | contact friction |
| | | cones |
+----------------------+----------------------+----------------------+
| :math:`\eta` | | upper bounds on |
| | | friction loss forces |
+----------------------+----------------------+----------------------+
| :math:`\Omega(q)` | | convex set of |
| | | admissible |
| | | constraint forces |
+----------------------+----------------------+----------------------+
| :math:`\mathcal{E}, | | index sets for |
| \mathcal{F}, | | Equality, Friction |
| \mathcal{C}` | | loss, Contact |
| | | constraints |
+----------------------+----------------------+----------------------+
.. list-table::
:widths: 1 1 4
:header-rows: 1
* - Symbol
- Size
- Description
* - :math:`z`
- :math:`n_C`
- constraint deformations
* - :math:`\omega`
- :math:`n_C`
- velocity of constraint deformations
* - :math:`k`
- :math:`n_C`
- virtual constraint stiffness
* - :math:`b`
- :math:`n_C`
- virtual constraint damping
* - :math:`d`
- :math:`n_C`
- constraint impedance
* - :math:`A(q)`
- :math:`n_C \times n_C`
- inverse inertia in constraint space
* - :math:`R(q)`
- :math:`n_C \times n_C`
- diagonal regularizer in constraint space
* - :math:`\ar`
- :math:`n_C`
- reference acceleration in constraint space
* - :math:`\au(q, v, \tau)`
- :math:`n_C`
- unconstrained acceleration in constraint space
* - :math:`\ac(q, v, \dot{v})`
- :math:`n_C`
- constrained acceleration in constraint space
* - :math:`\mathcal{K}(q)`
-
- product of all contact friction cones
* - :math:`\eta`
-
- upper bounds on friction loss forces
* - :math:`\Omega(q)`
-
- convex set of admissible constraint forces
* - :math:`\mathcal{E}, \mathcal{F}, \mathcal{C}`
-
- index sets for Equality, Friction loss, Contact constraints
The index sets will be used to refer to parts of vectors and matrices. For example, :math:`J_\mathcal{C}` is the
sub-matrix of all rows of the Jacobian that correspond to contact constraints.
@@ -1048,7 +1070,7 @@ explain what it means and why it makes sense. That problem is
.. math::
(\dot{v}, \dot{\omega}) = \arg \min_{(x, y)}
\left\|x-M^{-1}(\tau-c)\right\|^2_M +
\left\|y-a^*\right\|^{\text{Huber}(\eta)}_{R^{-1}} \\
\left\|y-\ar\right\|^{\text{Huber}(\eta)}_{R^{-1}} \\
\textrm{subject to} \;
J_\mathcal{E} x_\mathcal{E} - y_\mathcal{E} = 0, \;
J_\mathcal{F} x_\mathcal{F} - y_\mathcal{F} = 0, \;
@@ -1056,15 +1078,15 @@ explain what it means and why it makes sense. That problem is
:label: eq:primal
The new players here are the diagonal regularizer :math:`R > 0` which makes the constraints soft, and the reference
acceleration :math:`a^*` which stabilizes the constraints. The latter is similar in spirit to Baumgarte stabilization,
acceleration :math:`\ar` which stabilizes the constraints. The latter is similar in spirit to Baumgarte stabilization,
but instead of adding a constraint force directly, it modifies the optimization problem whose solution is the constraint
force. Since this problem is itself constrained, the relation between :math:`a^*` and :math:`f` is generally non-linear.
The quantities :math:`R` and :math:`a^*` are computed from the solver :ref:`parameters <soParameters>` as described
force. Since this problem is itself constrained, the relation between :math:`\ar` and :math:`f` is generally non-linear.
The quantities :math:`R` and :math:`\ar` are computed from the solver :ref:`parameters <soParameters>` as described
later. For now we assume they are given.
The optimization variable :math:`x` stands for acceleration as in the Gauss principle, while :math:`y` is a slack
variable in constraint space. It is needed to model soft constraints. If we forced the solution to reach
:math:`y = a^*`, which we could do by taking the limit :math:`R \to 0`, we would obtain a hard constraint model. This
:math:`y = \ar`, which we could do by taking the limit :math:`R \to 0`, we would obtain a hard constraint model. This
limit is not allowed in MuJoCo, but nevertheless one can construct models that are phenomenologically hard.
The symbol :math:`\mathcal{K}^*` denotes the dual to the friction cone. It is motivated by mathematical reverse
@@ -1095,7 +1117,7 @@ constraint space (which the constraint force aims to prevent), there is no posit
\tilde{q} &= {q \brack z}, &
\tilde{v} &= {v \brack \omega}, &
\tilde{c} &= {c \brack 0}, \\
\tilde{\tau} &= {\tau \brack {R^{-1} a^*}}, &
\tilde{\tau} &= {\tau \brack {R^{-1} \ar}}, &
\tilde{M} &= \left[\begin{array}{cc}
M & 0 \\
0 & R^{-1}
@@ -1114,20 +1136,20 @@ Unpacking all the tildes yields the explicit form of the original and the deform
.. math::
\begin{aligned}
M \dot{v} + c &= \tau +J^T f \\
\dot{\omega} &= a^* - R f \\
\dot{\omega} &= \ar - R f \\
\end{aligned}
Thus :math:`R` has the meaning of inverse deformation inertia, while :math:`a^*` has the meaning of unforced deformation
Thus :math:`R` has the meaning of inverse deformation inertia, while :math:`\ar` has the meaning of unforced deformation
acceleration.
Does MuJoCo keep these deformation variables as part of the system state and integrate their dynamics together with the
joint positions and velocities? No, although such an option may be worth providing in the future. Recall that we defined
the functional dependence of the regularizer and the reference acceleration as :math:`R(q)` and :math:`a^*(q, v)`. This
the functional dependence of the regularizer and the reference acceleration as :math:`R(q)` and :math:`\ar(q, v)`. This
makes problem :eq:`eq:primal` dependent only on :math:`(q, v, \tau)`, and so the original dynamics are not actually
affected by the deformation dynamics. Since the general constraint model we developed up to now makes no assumptions
about how :math:`R` and :math:`a^*` are computed, our choice is consistent and improves simulator efficiency.
about how :math:`R` and :math:`\ar` are computed, our choice is consistent and improves simulator efficiency.
Nevertheless, given that these quantities turned out to be related to the deformation dynamics, it may be more natural
to define them as :math:`R(z)` and :math:`a^* (z, \omega)` and simulate the entire augmented system. Below we clarify
to define them as :math:`R(z)` and :math:`\ar (z, \omega)` and simulate the entire augmented system. Below we clarify
some of the benefits of such a simulation.
When do the deformation dynamics "track" the original dynamics exactly? One can verify that this happens when the
@@ -1159,7 +1181,7 @@ we obtain the unconstrained problem
.. math::
\dot{v} = \arg \min_{x} \left\|x-M^{-1}(\tau-c)\right\|^2_M +
s \left( J x - a^* \right)
s \left( J x - \ar \right)
:label: eq:reduced
The function :math:`s(\cdot)` plays the role of a soft-constraint penalty. It can be shown to be convex and
@@ -1169,7 +1191,7 @@ Another appealing feature of the reduced formulation is that the inverse dynamic
above problem is unconstrained and convex, the unique global minimum makes the gradient vanish. This yields the identity
.. math::
M \dot{v} + c = \tau - J^T \nabla s \left( J \dot{v} - a^* \right)
M \dot{v} + c = \tau - J^T \nabla s \left( J \dot{v} - \ar \right)
which is the analytical inverse dynamics in the presence of soft constraints. Comparing to the equations of motion
:eq:`eq:motion`, we see that the constraint forces :math:`f` are given by the negative gradient of the function
@@ -1191,7 +1213,7 @@ Lagrange dual to the primal problem defined above is
.. math::
f = \arg\min_\lambda \frac{1}{2} \lambda^{T} \left( A+R \right) \lambda +
\lambda^T \left( a^0 - a^* \right) \\
\lambda^T \left( \au - \ar \right) \\
\text{subject to} \; \lambda \in \Omega
:label: eq:dual
@@ -1203,7 +1225,7 @@ where the inverse inertia in constraint space is
and the unconstrained acceleration in constraint space is
.. math::
a^0 = J M^{-1} (\tau-c) + \dot{J} v
\au = J M^{-1} (\tau-c) + \dot{J} v
The constraint set :math:`\Omega` is as follows. :math:`\lambda_\mathcal{E}` is unconstrained, because it is the
Lagrange multiplier for an equality constraint in the primal problem. For friction loss we have the box constraint
@@ -1218,28 +1240,28 @@ problem are described later.
As mentioned earlier, MuJoCo's constraint model has uniquely-defined inverse dynamics, and we already saw one way to
derive it in the reduced formulation above. Here we derive it again from the dual formulation. Recall that in inverse
dynamics we have access to :math:`(q, v, \dot{v})` instead of :math:`(q, v, \tau)`, so the unconstrained acceleration
:math:`a^0` is unknown. However we can compute the constrained acceleration
:math:`\au` is unknown. However we can compute the constrained acceleration
.. math::
a^1 = J \dot{v} + \dot{J} v
\ac = J \dot{v} + \dot{J} v
Inverse dynamics can now be computed by solving the optimization problem
.. math::
f = \arg \min_\lambda \frac{1}{2} \lambda^{T} R \lambda +
\lambda^T \left( a^1 - a^* \right) \\
\lambda^T \left( \ac - \ar \right) \\
\text{subject to} \; \lambda \in \Omega
By comparing the KKT conditions for these two convex optimization problems, one can verify that their solutions coincide
when
.. math::
a^1 = a^0 + Af
\ac = \au + Af
:label: eq:identity
This key identity is essentially Newton's second law projected in constraint space. It is derived by moving the term
:math:`c` in the equations of motion :eq:`eq:motion` to the right hand side, multiplying by :math:`J M^{-1}` from the
left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, a^0, a^1`. In terms of
left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, \au, \ac`. In terms of
implementation, we do not actually compute the acceleration term :math:`\dot{J} v`. This is because our optimization
problems depend on differences of constraint-space accelerations, and so this term would cancel out even if we were to
compute it.
@@ -1316,12 +1338,12 @@ representations of the constraint Jacobian and related matrices.
Parameters
~~~~~~~~~~
Here we explain how the quantities :math:`R, a^*` are computed from model parameters. For the chosen parameterization to
make sense, we first need to understand how these quantities affect the dynamics. We focus on the unconstrained
minimizer of :eq:`eq:dual`, namely
Here we explain how the quantities :math:`R, \ar` are computed from model parameters. For the chosen
parameterization to make sense, we first need to understand how these quantities affect the dynamics. We focus on the
unconstrained minimizer of :eq:`eq:dual`, namely
.. math::
f^+ = (A+R)^{-1} (a^* - a^0)
f^+ = (A+R)^{-1} (\ar - \au)
If it happens that :math:`f^+ \in \Omega`, then :math:`f^+ = f` is the actual constraint force generated by our model.
We focus on this case because it is common, in the sense that the subset of the constraints in :math:`\Omega` that are
@@ -1329,11 +1351,11 @@ active at any given time is usually small, and furthermore it is the only case t
Substituting :math:`f^+` in the constraint dynamics :eq:`eq:identity` and rearranging terms yields
.. math::
a^1 = A(A+R)^{-1} a^* + R (A+R)^{-1} a^0
\ac = A(A+R)^{-1} \ar + R (A+R)^{-1} \au
Thus the constrained acceleration interpolates between the unconstrained and the reference acceleration. In particular,
in the limit :math:`R \to 0` we have a hard constraint and :math:`a^1 = a^*`, while in the limit :math:`R \to \infty` we
have have an infinitely soft constraint (i.e., no constraint) and :math:`a^1 = a^0`. It is then natural to introduce a
in the limit :math:`R \to 0` we have a hard constraint and :math:`\ac = \ar`, while in the limit :math:`R \to \infty` we
have have an infinitely soft constraint (i.e., no constraint) and :math:`\ac = \au`. It is then natural to introduce a
model parameter which directly controls the interpolation. We call this parameter *impedance* and denote it :math:`d`.
It is a vector with dimensionality :math:`n_C` satisfying :math:`0<d<1` element-wise. Once it is specified, we compute
the diagonal elements of the regularizer as
@@ -1349,7 +1371,7 @@ approximation happened to be exact, and :math:`A` itself happened to be diagonal
constraint would satisfy
.. math::
a^1_i = d_i a^*_i + (1-d_i) a^0_i
\aci = d_i \ari + (1-d_i) \aui
and so we would achieve the desired interpolation effect. This of course does not hold exactly in general, but the goal
here is to construct a sensible and intuitive parameterization of the constraint model and get the scaling right.
@@ -1358,14 +1380,14 @@ Next we explain how the reference acceleration is computed. As already mentioned
parameterized by *damping* and *stiffness* coefficients element-wise:
.. math::
a^*_i = -b_i (J v)_i - k_i r_i
\ari = -b_i (J v)_i - k_i r_i
Recall that :math:`r` is the position residual (which is zero for friction loss and friction dimensions of elliptic
cones), while :math:`J v` is the joint velocity projected in constraint space; the indexing notation refers to one
component of the projected velocity vector.
To summarize, the user specifies the vectors of impedance coefficients :math:`0<d<1`, damping coefficients :math:`b > 0`
and stiffness coefficients :math:`k > 0`. The quantities :math:`R, a^*` are then computed by MuJoCo as shown above, and
and stiffness coefficients :math:`k > 0`. The quantities :math:`R, \ar` are then computed by MuJoCo as shown above, and
the selected optimization algorithm is applied to solve problem :eq:`eq:dual`. As explained in the :ref:`solver
parameters <CSolver>` section of the Modeling chapter, MuJoCo offers additional automation for setting :math:`d, b, k`
so as to achieve critical damping, or model a soft contact layer by varying :math:`d` with distance.
@@ -1534,6 +1556,7 @@ The top-level function :ref:`mj_step` invokes the sequence of computations below
the subtree centers of mass.
#. Compute all passive forces: spring-dampers in joints and tendons, and fluid dynamics forces.
#. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
If required by sensors, call :ref:`mj_subtreeVel`.
#. Compute the reference constraint acceleration.
#. Compute the vector of Coriolis, centrifugal and gravitational forces.
#. Compute the actuator forces and activation dynamics if defined.
@@ -1541,6 +1564,7 @@ The top-level function :ref:`mj_step` invokes the sequence of computations below
#. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics.
#. Compute sensor data that depends on force and acceleration if enabled.
If required by sensors, call :ref:`mj_rnePostConstraint`.
#. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
automatically reset and the corresponding warning is raised.
#. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
@@ -1567,6 +1591,7 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
#. Compute the tendon and actuator velocities.
#. Compute the body velocities and joint axes rates of change.
#. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
If required by sensors, call :ref:`mj_subtreeVel`.
#. Compute all passive forces.
#. Compute the reference constraint acceleration.
#. If the :ref:`invdiscrete<option-flag-invdiscrete>` flag is set and the :ref:`integrator<option-integrator>` is not
@@ -1574,6 +1599,7 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
#. Compute the constraint force. This is done analytically, without using a numerical solver.
#. Compute the inverse dynamics for the unconstrained system.
#. Compute sensor data that depends on force and acceleration if enabled.
If required by sensors, call :ref:`mj_rnePostConstraint`.
#. Compute the vector ``mjData.qfrc_inverse`` by combining all results. This is the main output of inverse dynamics. It
equals the sum of external and actuation forces.
doc/conf.py
+7
View File
@@ -191,8 +191,15 @@ favicons = [
# -- Options for katex ------------------------------------------------------
# See: https://sphinxcontrib-katex.readthedocs.io/en/0.4.1/macros.html
# {ar au, ac} are {reference, unconstrained, constrained} acceleration, resp.
latex_macros = r"""
\def \d #1{\operatorname{#1}}
\def \ar {a_{\rm ref}}
\def \au {a_0}
\def \ac {a_1}
\def \ari {a_{{\rm ref},i}}
\def \aui {a_{0,i}}
\def \aci {a_{1,i}}
"""
# Translate LaTeX macros to KaTeX and add to options for HTML builder
doc/css/theme_overrides.css
+6 -1
View File
@@ -19,9 +19,14 @@ body {
}
/* Make the search bar query text readable in dark theme. */
body:not([data-theme="light"]) {
body[data-theme="dark"] {
--color-sidebar-search-foreground: white;
}
@media (prefers-color-scheme: dark) {
body[data-theme="auto"] {
--color-sidebar-search-foreground: white;
}
}
.schema-small {
font-size: 93%;
doc/includes/references.h
+19 -5
View File
@@ -84,7 +84,7 @@ struct mjContact_ { // result of collision detection functions
// contact parameters set by near-phase collision function
mjtNum dist; // distance between nearest points; neg: penetration
mjtNum pos[3]; // position of contact point: midpoint between geoms
mjtNum frame[9]; // normal is in [0-2]
mjtNum frame[9]; // normal is in [0-2], points from geom[0] to geom[1]
// contact parameters set by mj_collideGeoms
mjtNum includemargin; // include if dist<includemargin=margin-gap
@@ -99,8 +99,8 @@ struct mjContact_ { // result of collision detection functions
// contact descriptors set by mj_collideXXX
int dim; // contact space dimensionality: 1, 3, 4 or 6
int geom1; // id of geom 1
int geom2; // id of geom 2
int geom1; // id of geom 1; deprecated, use geom[0]
int geom2; // id of geom 2; deprecated, use geom[1]
int geom[2]; // geom ids; -1 for flex
int flex[2]; // flex ids; -1 for geom
int elem[2]; // element ids; -1 for geom or flex vertex
@@ -292,7 +292,11 @@ struct mjData_ {
mjtNum* qfrc_bias; // C(qpos,qvel) (nv x 1)
// computed by mj_fwdVelocity/mj_passive
mjtNum* qfrc_passive; // passive force (nv x 1)
mjtNum* qfrc_spring; // passive spring force (nv x 1)
mjtNum* qfrc_damper; // passive damper force (nv x 1)
mjtNum* qfrc_gravcomp; // passive gravity compensation force (nv x 1)
mjtNum* qfrc_fluid; // passive fluid force (nv x 1)
mjtNum* qfrc_passive; // total passive force (nv x 1)
// computed by mj_sensorVel/mj_subtreeVel if needed
mjtNum* subtree_linvel; // linear velocity of subtree com (nbody x 3)
@@ -728,6 +732,7 @@ struct mjOption_ { // physics options
int mpr_iterations; // maximum number of MPR solver iterations
int disableflags; // bit flags for disabling standard features
int enableflags; // bit flags for enabling optional features
int disableactuator; // bit flags for disabling actuators by group id
// sdf collision settings
int sdf_initpoints; // number of starting points for gradient descent
@@ -796,6 +801,7 @@ struct mjVisual_ { // visualization options
float framewidth; // bodyframe axis width
float constraint; // constraint width
float slidercrank; // slidercrank width
float frustum; // frustum zfar plane
} scale;
struct { // color of decor elements
@@ -821,6 +827,7 @@ struct mjVisual_ { // visualization options
float constraint[4]; // constraint
float slidercrank[4]; // slidercrank
float crankbroken[4]; // used when crank must be stretched/broken
float frustum[4]; // camera frustum
} rgba;
};
typedef struct mjVisual_ mjVisual;
@@ -944,7 +951,7 @@ struct mjModel_ {
int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1)
int* body_geomnum; // number of geoms (nbody x 1)
int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1)
mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1)
mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1)
mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1)
mjtNum* body_pos; // position offset rel. to parent body (nbody x 3)
mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4)
@@ -1122,6 +1129,7 @@ struct mjModel_ {
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all verices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
@@ -1425,6 +1433,9 @@ struct mjpPlugin_ {
// called during compilation for marching cubes
mjtNum (*sdf_staticdistance)(const mjtNum point[3], const mjtNum* attributes);
// convert attributes and provide defaults if not present
void (*sdf_attribute)(mjtNum attribute[], const char* name[], const char* value[]);
// bounding box of implicit surface
void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
};
@@ -2130,6 +2141,7 @@ struct mjvSceneState_ {
int nskin;
int nflex;
int nflexvert;
int nflextexcoord;
int nskinvert;
int nskinface;
int nskinbone;
@@ -2194,6 +2206,7 @@ struct mjvSceneState_ {
mjtNum* cam_fovy;
mjtNum* cam_ipd;
float* cam_intrinsic;
float* cam_sensorsize;
mjtByte* light_directional;
@@ -2219,6 +2232,7 @@ struct mjvSceneState_ {
int* flex_shell;
int* flex_shellnum;
int* flex_shelldataadr;
int* flex_texcoordadr;
int* flex_bvhadr;
int* flex_bvhnum;
mjtNum* flex_radius;
doc/mjx.rst
+21 -12
View File
@@ -1,6 +1,8 @@
==========
MuJoCo XLA
==========
.. _Mjx:
================
MuJoCo XLA (MJX)
================
Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the
`mjx <https://github.com/google-deepmind/mujoco/tree/main/mjx>`__ directory. MJX allows MuJoCo to run on compute
@@ -94,7 +96,7 @@ Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An
mjx_model = mjx.device_put(model)
mjx_data = mjx.make_data(model)
Using ``mx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``.
Using ``mjx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``.
.. _MjxFunctions:
@@ -186,7 +188,9 @@ The following features are **fully supported** in MJX:
* - :ref:`Geom <mjtGeom>`
- ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
* - :ref:`Constraint <mjtConstraint>`
- ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
- ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
* - :ref:`Equality <mjtEq>`
- ``CONNECT``, ``WELD``, ``JOINT``
* - :ref:`Integrator <mjtIntegrator>`
- ``EULER``, ``RK4``
* - :ref:`Cone <mjtCone>`
@@ -194,7 +198,7 @@ The following features are **fully supported** in MJX:
* - :ref:`Condim <coContact>`
- 3
* - :ref:`Solver <mjtSolver>`
- ``CG``
- ``CG``, ``NEWTON``
* - Fluid Model
- :ref:`flInertia`
@@ -214,18 +218,22 @@ The following features are **in development** and coming soon:
- ``TRN_TENDON``
* - :ref:`Geom <mjtGeom>`
- ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF``
* - :ref:`Constraint <mjtConstraint>`
- ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
* - :ref:`Integrator <mjtIntegrator>`
- ``IMPLICIT``, ``IMPLICITFAST``
* - :ref:`Cone <mjtCone>`
- ``ELLIPTIC``
* - :ref:`Condim <coContact>`
- 1, 4, 6
* - :ref:`Solver <mjtSolver>`
- ``NEWTON``
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Tendons <tendon>`
- :ref:`Spatial <tendon-spatial>`, :ref:`Fixed <tendon-fixed>`
* - :ref:`Equality <mjtEq>`
- ``TENDON``
* - :ref:`Sensors <mjtSensor>`
- All except ``mjSENS_PLUGIN``, ``mjSENS_USER``
The following features are **unsupported**:
@@ -296,10 +304,11 @@ Performance tuning
For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values:
:ref:`option` element
For now, solver must be set to ``CG`` (but Newton is on its way!). The ``iterations`` and ``ls_iterations``
attributes---which control solver and linesearch iterations, respectively---should be brought down to just low enough
that the simulation remains stable. Accurate solver forces are not so important in reinforcement learning in which
domain randomization is often used to add noise to physics for sim2real.
The ``iterations`` and ``ls_iterations`` attributes---which control solver and linesearch iterations, respectively---
should be brought down to just low enough that the simulation remains stable. Accurate solver forces are not so
important in reinforcement learning in which domain randomization is often used to add noise to physics for sim-to-real.
The ``NEWTON`` :ref:`Solver <mjtSolver>` often delivers reasonable convergence with one solver iteration, and performs
well on GPU. ``CG`` is currently a better choice for TPU.
:ref:`contact-pair` element
Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider
doc/modeling.rst
+145 -90
View File
@@ -190,7 +190,7 @@ element, it must be undefined in the active defaults class.
A final twist here is actuators. They are different because some of the actuator-related elements are actually
shortcuts, and shortcuts interact with the defaults setting mechanism in a non-obvious way. This is explained in the
:ref:`Actuator shortcuts <CActuator>` section below.
:ref:`Actuator shortcuts <CActShortcuts>` section below.
.. _CFrame:
@@ -198,8 +198,8 @@ Coordinate frames
~~~~~~~~~~~~~~~~~
The positions and orientations of all elements defined in the kinematic tree are expressed in local coordinates,
relative to the parent body for bodies, and relative to the body that owns the element for geoms, joints, sites, cameras
and lights.
relative to the parent body for bodies, and relative to the body that contains the element for geoms, joints, sites,
cameras and lights.
A related attribute is :ref:`compiler/angle<compiler-angle>`. It specifies whether angles in the MJCF file are expressed
in degrees or radians (after compilation, angles are always expressed in radians).
@@ -259,14 +259,14 @@ available in :ref:`option <option>`; it can be used to change all contact-relate
experiment interactively with parameter settings or implement continuation methods for numerical optimization.
Here we focus on a single scalar constraint. Using slightly different notation from the Computation chapter, let
:math:`a_1` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in
:math:`\ac` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in
friction dimensions), :math:`k` and :math:`b` the stiffness and damping of the virtual spring used to define the
reference acceleration :math:`a_{\rm ref} = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`a_0` the
reference acceleration :math:`\ar = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`\au` the
acceleration in the absence of constraint force. Our earlier analysis revealed that the dynamics in constraint space are
approximately
.. math::
a_1 + d \cdot (b v + k r) = (1 - d)\cdot a_0
\ac + d \cdot (b v + k r) = (1 - d)\cdot \au
Again, the parameters that are under the user's control are :math:`d, b, k`. The remaining quantities are functions of
the system state and are computed automatically at each time step.
@@ -278,15 +278,16 @@ Impedance
We begin by explaining the constraint impedance :math:`d`.
.. admonition:: Intuitive description
.. admonition:: Intuitive description of the **impedance**
The *impedance* :math:`d \in (0, 1)` determines a constraint's **ability to generate force**.
The *impedance* :math:`d \in (0, 1)` corresponds to a constraint's **ability to generate force**.
Small values of :math:`d` correspond to weak constraints while large values of :math:`d`
correspond to strong constraints. Impedance is set using the :at:`solimp` attribute.
correspond to strong constraints. The impedance affects the constraint at all times, in particular when the system is
at rest. Impedance is set using the :at:`solimp` attribute.
Recall that :math:`d` must lie between 0 and 1; internally MuJoCo clamps it to the range [:ref:`mjMINIMP mjMAXIMP
<glNumeric>`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced
acceleration :math:`a_0` and reference acceleration :math:`a_{\rm ref}`. The user can set :math:`d` to a constant, or
acceleration :math:`\au` and reference acceleration :math:`\ar`. The user can set :math:`d` to a constant, or
take advantage of its interpolating property and make it position-dependent, i.e., a function of the constraint
violation :math:`r`. Position-dependent impedance can be used to model soft contact layers around objects, or define
equality constraints that become stronger with larger violation (so as to approximate backlash, for example). The shape
@@ -333,15 +334,26 @@ Reference
^^^^^^^^^
Next we explain the setting of the stiffness :math:`k` and damping :math:`b` which control the reference acceleration
:math:`a_{\rm ref}`.
:math:`\ar`.
.. admonition:: Intuitive description
.. admonition:: Intuitive description of the **reference acceleration**
The *reference acceleration* :math:`a_{\rm ref}` determines **what the constraint is trying to achieve** (as opposed
to how well it can achieve it). This acceleration is defined by two numbers, a stiffness :math:`k` and damping
:math:`b` which can be set directly or re-parameterized as the time-constant and damping ratio of a
mass-spring-damper system (a `harmonic oscillator <https://en.wikipedia.org/wiki/Harmonic_oscillator>`__).
The reference acceleration is controlled by the :at:`solref` attribute.
The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in
order to rectify violation. For example, consider a contact between a motionless free body pulled down by gravity
onto a static plane geom. Since there is no motion, the penetration will be entirely determined by the impedance
while the reference has no effect. Now imagine that the body is dropped onto the plane. Upon impact the constraint
will generate a normal force which attempts to rectify the penetration using a particular motion; this motion is
the reference acceleration.
Another way of understanding the reference acceleration is to think of the unmodeled deformation variables
described in the :ref:`Computation chapter<soPrimal>`. Imagine two bodies pressed together, leading to deformation at
the contact. Now pull the bodies apart very quickly; the motion of the deformation as it settles into its undeformed
state is the reference acceleration.
This acceleration is defined by two numbers, a stiffness :math:`k` and damping :math:`b` which can be set directly or
re-parameterized as the time-constant and damping ratio of a mass-spring-damper system (a `harmonic oscillator
<https://en.wikipedia.org/wiki/Harmonic_oscillator>`__). The reference acceleration is controlled by the :at:`solref`
attribute.
There are two formats for this attribute, determined by the sign of the numbers. If both numbers are positive the
specification is considered to be in the :math:`(\text{timeconst}, \text{dampratio})` format. If negative it is in the
@@ -441,7 +453,7 @@ solref, solimp
If one of the two geoms has higher priority, its solref and solimp parameters are used. If both geoms have the same
priority, the weighted average is used. The weights are proportional to the solmix attributes, i.e., weight1 =
solmix1 / (solmix1 + solmix2) and similarly for weight2. There is one important exception to this weighted averaging
rule. If solref for either geom is non-positive, i.e., it relies on the new direct format introduced in MuJoCo 2.0,
rule. If solref for either geom is non-positive, i.e., it relies on the direct format,
then the element-wise minimum is used regardless of solmix. This is because averaging solref parameters in different
formats would be meaningless.
@@ -555,10 +567,44 @@ general guidelines and observations:
setup operation for the main PGS and Noslip PGS is the same, thus the setup cost is paid only once when both are
enabled.
.. _CActuator:
.. _CActuators:
Actuator shortcuts
~~~~~~~~~~~~~~~~~~
Actuators
~~~~~~~~~
This section describes various aspects of using actuators in MuJoCo. See the :ref:`Actuation model <geActuation>`
regarding the computational model.
.. _CActDisable:
Group disable
^^^^^^^^^^^^^
The :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` attribute, which can be changed at runtime by setting the
:ref:`mjOption.disableactuator<mjOption>` integer bitfield, allows the user to disable sets of actuators according to
their :ref:`group<actuator-general-group>`. This feature is convenient when one would like to use multiple types of
actuators for the same kinematic tree. For example consider a robot with firmware that supports mutiple control modes
e.g., torque-control and position-control. In this case, one can define both types of actuators in the same MJCF
model, assigning one type of actuator to group 0 and the other to group 1.
.. youtube:: H9qG9Zf2W44
:align: right
:width: 40%
The :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` MJCF attribute selects which groups are disabled by
default, and :ref:`mjOption.disableactuator<mjOption>` can be set at runtime to switch the active set. Note that the
total number of actuators ``mjModel.nu`` remains unchanged, as do the actuator indices, so it is up to the user to know
that the respective ``mjData.ctrl`` values of disabled actuators will be ignored and produce no force. `This example
model <https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
has three actuator groups which can be toggled at runtime in the :ref:`simulate<saSimulate>` interactive viewer.
See `example model
<https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
and associated screen-capture on the right.
.. _CActShortcuts:
Shortcuts
^^^^^^^^^
As explained in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo offers a flexible
actuator model with transmission, activation dynamics and force generation components that can be specified
@@ -593,8 +639,8 @@ explicitly.
.. _CForceRange:
Actuator force clamping
~~~~~~~~~~~~~~~~~~~~~~~
Force limits
^^^^^^^^^^^^
Actuator forces are usually limited between lower and upper bounds. These limits can be enforced in three ways:
@@ -633,62 +679,14 @@ Force clamping at joint input with :ref:`joint/actuatorfrcrange<body-joint-actua
The three clamping options above are non-exclusive and can be combined as required.
.. _CActRange:
Activation clamping
~~~~~~~~~~~~~~~~~~~
As described in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo supports actuators
with internal dynamics whose states are called "activations". One useful application of these stateful actuators is the
"integrated-velocity" actuator, implemented by the :ref:`intvelocity<actuator-intvelocity>` shortcut. Different from the
:ref:`pure velocity<actuator-velocity>` actuators, which implement direct feedback on transmission target's velocity,
*integrated-velocity* actuators couple an *integrator* with a *position-feedback* actuator. In this case the semantics
of the activation state are "the setpoint of the position actuator", and the semantics of the control signal are "the
velocity of the setpoint of the position actuator". Note that in real robotic systems this integrated-velocity actuator
is the most common implementation of actuators with velocity semantics, rather than pure feedback on velocity which is
often quite unstable (both in real life and in simulation).
In the case of integrated-velocity actuators, it is often desirable to *clamp* the activation state, since otherwise the
position target would keep integrating beyond the joint limits, leading to loss of controllabillity. To see the effect
of activation clamping, load the example model below:
.. code-block:: xml
<mujoco>
<default>
<joint axis="0 0 1" limited="true" range="-90 90" damping="0.3"/>
<geom size=".1 .1 .1" type="box"/>
</default>
<worldbody>
<body>
<joint name="joint1"/>
<geom/>
</body>
<body pos=".3 0 0">
<joint name="joint2"/>
<geom/>
</body>
</worldbody>
<actuator>
<general name="unclamped" joint="joint1" gainprm="1" biastype="affine"
biasprm="0 -1" dyntype="integrator"/>
<intvelocity name="clamped" joint="joint2" actrange="-1.57 1.57"/>
</actuator>
</mujoco>
Note that the :at:`actrange` attribute is always specified in native units (radians), even though the joint range
can be either in degrees (the default) or radians, depending on the :ref:`compiler/angle <compiler>` attribute.
.. _CLengthRange:
Actuator length range
~~~~~~~~~~~~~~~~~~~~~
Length range
^^^^^^^^^^^^
As of MuJoCo 2.0, the field mjModel.actuator_lengthrange contains the range of feasible actuator lengths (or more
precisely, lengths of the actuator's transmission). This is needed to simulate :ref:`muscle actuators <CMuscle>` as
explained below. Here we focus on what actuator_lengthrange means and how to set it.
The field ``mjModel.actuator_lengthrange`` contains the range of feasible actuator lengths (or more
precisely, lengths of the actuator's transmission). This is needed to simulate :ref:`muscle actuators <CMuscle>`.
Here we focus on what actuator_lengthrange means and how to set it.
Unlike all other fields of mjModel which are exact physical or geometric quantities, actuator_lengthrange is an
approximation. Intuitively it corresponds to the minimum and maximum length that the actuator's transmission can reach
@@ -732,12 +730,69 @@ practice length ranges will almost always be used with muscle actuators attached
joint limits defined in the model, effectively limiting the lengths of the muscle actuators. If you get a convergence
error in such a model, the most likely explanation is that you forgot to include joint limits.
.. _CActivation:
Stateful actuators
^^^^^^^^^^^^^^^^^^
As described in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo supports actuators
with internal dynamics whose states are called "activations".
.. _CActRange:
Activation limits
'''''''''''''''''
One useful application of stateful actuators is the
"integrated-velocity" actuator, implemented by the :ref:`intvelocity<actuator-intvelocity>` shortcut. Different from the
:ref:`pure velocity<actuator-velocity>` actuators, which implement direct feedback on transmission target's velocity,
*integrated-velocity* actuators couple an *integrator* with a *position-feedback* actuator. In this case the semantics
of the activation state are "the setpoint of the position actuator", and the semantics of the control signal are "the
velocity of the setpoint of the position actuator". Note that in real robotic systems this integrated-velocity actuator
is the most common implementation of actuators with velocity semantics, rather than pure feedback on velocity which is
often quite unstable (both in real life and in simulation).
In the case of integrated-velocity actuators, it is often desirable to *clamp* the activation state, since otherwise the
position target would keep integrating beyond the joint limits, leading to loss of controllabillity. To see the effect
of activation clamping, load the example model below:
.. collapse:: Example model with activation limits
.. code-block:: xml
<mujoco>
<default>
<joint axis="0 0 1" limited="true" range="-90 90" damping="0.3"/>
<geom size=".1 .1 .1" type="box"/>
</default>
<worldbody>
<body>
<joint name="joint1"/>
<geom/>
</body>
<body pos=".3 0 0">
<joint name="joint2"/>
<geom/>
</body>
</worldbody>
<actuator>
<general name="unclamped" joint="joint1" gainprm="1" biastype="affine"
biasprm="0 -1" dyntype="integrator"/>
<intvelocity name="clamped" joint="joint2" actrange="-1.57 1.57"/>
</actuator>
</mujoco>
Note that the :at:`actrange` attribute is always specified in native units (radians), even though the joint range
can be either in degrees (the default) or radians, depending on the :ref:`compiler/angle <compiler>` attribute.
.. _CMuscle:
Muscle actuators
~~~~~~~~~~~~~~~~
Muscles
'''''''
As of MuJoCo 2.0, we provide a set of tools for modeling biological muscles. Users who want to add muscles with minimum
We provide a set of tools for modeling biological muscles. Users who want to add muscles with minimum
effort can do so with a single line of XML in the actuator section:
.. code-block:: xml
@@ -937,13 +992,13 @@ be 0. This effect can be approximated by scaling down the muscle force and also
Tendon wrapping is also more limited in MuJoCo. We allow spheres and infinite cylinders as wrapping objects, and require
two wrapping objects to be separated by a fixed site in the tendon path. This is to avoid the need for iterative
computations of tendon paths. As of MuJoCo 2.0 we also allow "side sites" to be placed inside the sphere or cylinder,
computations of tendon paths. We also allow "side sites" to be placed inside the sphere or cylinder,
which causes an inverse wrap: the tendon path is constrained to pass through the object instead of going around it. This
can replace torus wrapping objects used in OpenSim to keep the tendon path within a given area. Overall, tendon wrapping
is the most challenging part of converting an OpenSim model to a MuJoCo model, and requires some manual work. On the
bright side, there is a small number of high-quality OpenSim models in use, so once they are converted we are done.
Below we illustrate the four types of tendon wrapping available in MuJoCo 2.0. Note that the curved sections of the
Below we illustrate the four types of tendon wrapping available. Note that the curved sections of the
wrapping tendons are rendered as straight, but the geometry pipeline works with the actual curves and computes their
lengths and moments analytically:
@@ -981,7 +1036,7 @@ Cameras
Besides the default, user-controllable, free camera, "fixed" cameras can be attached to the kinematic tree.
Extrinsics
By default, camera frames are attached to the parent body. The optional :ref:`mode<body-camera-mode>` and
By default, camera frames are attached to the containing body. The optional :ref:`mode<body-camera-mode>` and
:ref:`target<body-camera-target>` attributes can be used to specify camera that track (move with) or target (look at)
a body or subtree. Cameras look towards the negative Z axis of the camera frame, while positive X and Y correspond to
*right* and *up* in the image plane, respectively.
@@ -1003,8 +1058,7 @@ Intrinsics
Composite objects
~~~~~~~~~~~~~~~~~
Composite objects were introduced in MuJoCo 2.0, along with solver optimizations to speed up the simulation of such
objects. They are not new model elements. Instead, they are (large) collections of existing elements designed to
Composite objects are not new model elements. Instead, they are (large) collections of existing elements designed to
simulate particle systems, ropes, cloth, and soft bodies. These collections are generated by the model compiler
automatically. The user configures the automatic generator on a high level, using the new XML element
:ref:`composite <body-composite>` and its attributes and sub-elements, as described in the XML reference
@@ -1026,7 +1080,7 @@ of these equality constraints can be adjusted by the user, thereby adjusting the
composite objects.
In addition to setting up the physics, the composite object generator creates suitable rendering. 2D and 3D objects
can be rendered as :ref:`skins <asset-skin>` which are also new in MuJoCo 2.0. The skin is generated
can be rendered as :ref:`skins <asset-skin>`. The skin is generated
automatically, and can be textured as well as subdivided using bi-cubic interpolation. The actual physics and in
particular the collision detection are based on the element bodies and their geoms, while the skin is purely a
visualization object. Yet in most situations we prefer to look at the skin representation. To facilitate this, the
@@ -1034,13 +1088,12 @@ generator places all geoms, sites and tendons in group 3 whose visualization is
a 2D grid for example, you will see a continuous flexible surface and not a collection of spheres connected with
tendons. However when fine-tuning the model and trying to understand the physics behind it, it is useful to be able to
render the spheres and tendons. To switch the rendering style, disable the rendering of skins and enable group 3 for
geoms and tendons (note that starting with MuJoCo 2.0 we have added a group property to sites, tendons and joints in
addition to geoms).
geoms and tendons.
We have designed the composite object generator to have intuitive high-level controls as much as possible, but at the
same time it exposes a large number of options that interact with each other and can profoundly affect the resulting
physics. So at some point users should read the :ref:`reference documentation <body-composite>` carefully.
As a quick start though, MuJoCo 2.0 comes with an example of each composite object type. Below we go over these
As a quick start though, MuJoCo comes with an example of each composite object type. Below we go over these
examples and explain the less obvious aspects. In all examples we have a static scene which is included in the model,
followed by a single composite object. The static scene has a mocap body (large capsule) that can be moved around with
the mouse to probe the behavior of the system. The XML snippets below are just the definition of the composite object;
@@ -1277,6 +1330,8 @@ into the engine in future releases.
.. code-block:: xml
<option timestep=".001"/>
<extension>
<plugin plugin="mujoco.elasticity.solid"/>
</extension>
@@ -1435,7 +1490,7 @@ mocap bodies around:
The key thing to understand about mocap bodies is that the simulator treats them as being fixed. We are causing them
to move from one simulation time step to the next by updating their position and orientation directly, but as far as
the physics model is concerned their position and orientation are constant. So what happens if we make contact with a
regular dynamic body, as in the composite object examples provided with the MuJoCo 2.0 distribution (recall that in
regular dynamic body, as in the composite object examples provided with the MuJoCo distribution (recall that in
those example we have a capsule probe which is a mocap body that we move with the mouse). A contact between two
regular bodies will experience penetration as well as relative velocity, while contact with a mocap body is missing
the relative velocity component because the simulator does not know that the mocap body itself is moving. So the
@@ -1587,7 +1642,7 @@ the same yet the damping component of the force is integrated implicitly.
Restitution
~~~~~~~~~~~
MuJoCo 2.0 introduced a new mechanism for specifying :at:`solref` as explained in :ref:`Solver parameters <CSolver>`.
Another mechanism exists for specifying :at:`solref`, as explained in :ref:`Solver parameters <CSolver>`.
When both numbers are non-positive, they are interpreted as (-stiffness, -damping) and scaled by the constraint
impedance. To achieve perfect restitution for contacts and other constraints, set stiffness to some reasonably large
value and damping to zero. Below is an example of a sphere bouncing on a plane with restitution coefficient of 1, so
doc/overview.rst
+12 -7
View File
@@ -944,15 +944,20 @@ can be obtained as:
int qposadr = -1, qveladr = -1;
// make sure we have a floating body: it has a single free joint
if( bodyid>=0 && m->body_jntnum[bodyid]==1 &&
m->jnt_type[m->body_jntadr[bodyid]]==mjJNT_FREE )
{
// extract the addresses from the joint specification
qposadr = m->jnt_qposadr[m->body_jntadr[bodyid]];
qveladr = m->jnt_dofadr[m->body_jntadr[bodyid]];
}
if( bodyid>=0 && m->body_jntnum[bodyid]==1 && m->jnt_type[m->body_jntadr[bodyid]]==mjJNT_FREE ) {
// extract the addresses from the joint specification
qposadr = m->jnt_qposadr[m->body_jntadr[bodyid]];
qveladr = m->jnt_dofadr[m->body_jntadr[bodyid]];
}
Now if everything went well (i.e., "myfloatingbody" was indeed a floating body), qposadr and qveladr are the addresses
in qpos and qvel where the data for our floating body/joint lives. The position data is 7 numbers (3D position followed
by unit quaternion) while the velocity data is 6 numbers (3D linear velocity followed by 3D angular velocity). These
numbers can now be set to the desired pose and velocity of the body.
The semantics of free joints are as follows. The linear postions of free joints are in the global frame, as are
linear velocities. The orientation of a free joint (the quaternion) is also in the global frame. However, the rotational
velocities of a free joint are in the local body frame. This is not so much a design decision but rather correct
use of the topology of quaternions. Angular velocities live in the quaternion tangent space, which is defined locally
for a certain orientation, so frame-local angular velocities are a natural parameterization.
Accelerations are defined in the same space as the corresponding velocities.
doc/programming/extension.rst
+15 -11
View File
@@ -250,20 +250,24 @@ A future version of this section will include:
Currently, there are three directories of first-party plugins:
* **elasticity:** The plugins in the `elasticity/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__
directory are passive forces based on continuum mechanics for 1-dimensional and 3-dimensional bodies. The
1D model is invariant under rotations and captures the large deformation of elastic cables, decoupling twisting and
bending strains. The 3D solid is a
`Saint Venant-Kirchhoff <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__
model discretized with piecewise linear finite elements, which is suitable for large deformations with small strains.
See also :ref:`composite <CComposite>` and :ref:`deformable <CDeformable>` objects.
* **elasticity:** The plugins in the `elasticity/
<https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__ directory are passive forces based on
continuum mechanics for 1-dimensional and 3-dimensional bodies. The 1D model is invariant under rotations and captures
the large deformation of elastic cables, decoupling twisting and bending strains. The 3D solid is a `Saint
Venant-Kirchhoff <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__ model
discretized with piecewise linear finite elements, which is suitable for large deformations with small strains. See
also :ref:`composite <CComposite>` and :ref:`deformable <CDeformable>` objects. For more information, please see the
`README <https://github.com/google-deepmind/mujoco/blob/main/plugin/elasticity/README.md>`__.
* **sensor:** The plugins in the `sensor/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sensor>`__
directory implement custom sensors. Currently the sole sensor plugin is the touch grid sensor, see the
`README <https://github.com/google-deepmind/mujoco/blob/main/plugin/sensor/README.md>`__ for details.
* **sdf:** The plugins in the `sdf/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sdf>`__
directory specify custom shapes in a mesh-free manner, by defining methods computing a signed distance field and its
gradient at query points. This shape then acts as a new geom type in the collision table at the top of
`engine_collision_driver.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__.
* **sdf:** The plugins in the `sdf/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sdf>`__ directory
specify custom shapes in a mesh-free manner, by defining methods computing a signed distance field and its gradient at
query points. This shape then acts as a new geom type in the collision table at the top of `engine_collision_driver.c
<https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__. For more information
concerning the available SDFs and how to write your own implicit geometry, please see the `README
<https://github.com/google-deepmind/mujoco/blob/main/plugin/sdf/README.md>`__. The rest of this section will give more
detail concerning the collision algorithm and the plugin engine interface.
Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent.
Because SDFs are non-convex, multiple starting points are required in order to converge to multiple local minima.
doc/programming/index.rst
+2 -1
View File
@@ -35,7 +35,7 @@ Thread
to process Tasks asynchronously. To enable use in MuJoCo, create a ThreadPool and assign it to the thread_pool field
in mjData.
UI framework
The UI framework (new in MuJoCo 2.0) is written in C. UI elements are rendered in OpenGL. It has its own event
The UI framework is written in C. UI elements are rendered in OpenGL. It has its own event
mechanism and abstract hooks for keyboard and mouse input. The code samples use it with GLFW, but it can also be used
with other window libraries.
@@ -275,5 +275,6 @@ now lazily resolved at runtime after the switch to GLAD, the "nogl" libraries ar
simulation
visualization
ui
samples
extension
doc/programming/samples.rst
+29 -28
View File
@@ -73,33 +73,6 @@ Where the command line arguments are
`start /affinity <https://learn.microsoft.com/en-us/windows-server/administration/windows-commands/start>`__
command on Windows (processor affinity cannot be specified through documented API means on macOS).
.. _saCompile:
`compile <https://github.com/google-deepmind/mujoco/blob/main/sample/compile.cc>`_
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This code sample evokes the built-in parser and compiler. It implements all possible model conversions from (MJCF, URDF,
MJB) format to (MJCF, MJB, TXT) format. Models saved as MJCF use a canonical subset of our format as described in the
:doc:`../modeling` chapter, and therefore MJCF-to-MJCF conversion will generally result in a different file.
The TXT format is a human-readable road-map to the model. It cannot be loaded by MuJoCo, but can be a very useful aid
during model development. It is in one-to-one correspondence with the compiled mjModel. Note also that one can use the
function :ref:`mj_printData` to create a text file which is in one-to-one correspondence
with mjData, although this is not done by the code sample.
.. _saBasic:
`basic <https://github.com/google-deepmind/mujoco/blob/main/sample/basic.cc>`_
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This code sample is a minimal interactive simulator. The model file must be provided as command-line argument. It
opens an OpenGL window using the platform-independent GLFW library, and renders the simulation state at 60 fps while
advancing the simulation in real-time. Press Backspace to reset the simulation. The mouse can be used to control the
camera: left drag to rotate, right drag to translate in the vertical plane, shift right drag to translate in the
horizontal plane, scroll or middle drag to zoom.
The :ref:`Visualization` programming guide below explains how visualization works. This code sample is a minimal
illustration of the concepts in that guide.
.. _saSimulate:
`simulate <https://github.com/google-deepmind/mujoco/blob/main/simulate>`_
@@ -108,10 +81,11 @@ illustration of the concepts in that guide.
This code sample is a fully-featured interactive simulator. It opens an OpenGL window using the platform-independent
GLFW library, and renders the simulation state in it. There is built-in help, simulation statistics, profiler, sensor
data plots. The model file can be specified as a command-line argument, or loaded at runtime using drag-and-drop
functionality. As of MuJoCo 2.0, this code sample uses the native UI to render various controls, and provides an
functionality. This code sample uses the native UI to render various controls, and provides an
illustration of how the new UI framework is intended to be used. Below is a screen-capture of ``simulate`` in action:
.. youtube:: 0ORsj_E17B0
:width: 95%
:align: center
Interaction is done with the mouse; built-in help with a summary of available commands is available by pressing the
@@ -140,6 +114,33 @@ power settings, the OS may reduce the CPU frequency; this is because :ref:`simul
the time in order to slow down to realtime. This results in inaccurate timings. To avoid this problem, change the
Windows power plan so that the minimum processor state is 100%.
.. _saCompile:
`compile <https://github.com/google-deepmind/mujoco/blob/main/sample/compile.cc>`_
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This code sample evokes the built-in parser and compiler. It implements all possible model conversions from (MJCF, URDF,
MJB) format to (MJCF, MJB, TXT) format. Models saved as MJCF use a canonical subset of our format as described in the
:doc:`../modeling` chapter, and therefore MJCF-to-MJCF conversion will generally result in a different file.
The TXT format is a human-readable road-map to the model. It cannot be loaded by MuJoCo, but can be a very useful aid
during model development. It is in one-to-one correspondence with the compiled mjModel. Note also that one can use the
function :ref:`mj_printData` to create a text file which is in one-to-one correspondence
with mjData, although this is not done by the code sample.
.. _saBasic:
`basic <https://github.com/google-deepmind/mujoco/blob/main/sample/basic.cc>`_
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This code sample is a minimal interactive simulator. The model file must be provided as command-line argument. It
opens an OpenGL window using the platform-independent GLFW library, and renders the simulation state at 60 fps while
advancing the simulation in real-time. Press Backspace to reset the simulation. The mouse can be used to control the
camera: left drag to rotate, right drag to translate in the vertical plane, shift right drag to translate in the
horizontal plane, scroll or middle drag to zoom.
The :ref:`Visualization` programming guide below explains how visualization works. This code sample is a minimal
illustration of the concepts in that guide.
.. _saRecord:
`record <https://github.com/google-deepmind/mujoco/blob/main/sample/record.cc>`_
doc/programming/ui.rst
+105
View File
@@ -0,0 +1,105 @@
.. _UI:
User Interface
--------------
MuJoCo has a native UI framework. Its use is illustrated in the :ref:`simulate.cc <saSimulate>` viewer. It is
designed to be fast in terms of updating and rendering, easy to use both for the developer and for the user,
cross-platform, and integrated with the native MuJoCo renderer. In order to achieve these design goals, we have omitted
many features and customization options that are available in other UI frameworks, and instead focused on efficiency and
automation.
.. _uiDesign:
Design overview
~~~~~~~~~~~~~~~
Native OpenGL rendering
We do not use any helper tools or libraries. Instead we provide C code for rendering all UI elements directly in
OpenGL. We support multiple UIs, each of which is a virtual rectangle that can be taller than the visible window. The
elements of each UI are rendered offscreen in auxiliary OpenGL buffers, via minimal updates, only when changes are
necessary. At each screen refresh we then copy the pixels from these auxiliary buffers to the window framebuffer, and
also implement vertical scroll bars when the window is smaller than the UI. This copy operation is done on the GPU and
is very fast.
Platform abstraction
The software design has 3 layers: OpenGL rendering of UI elements working in conjunction with the MuJoCo renderer
(which is fully cross-platform); abstract functions for access to windows, keyboard and mouse, defined as pure virtual
functions in the ``PlatformUIAdapter`` class; and an implementation of those functions in the derived class
``GlfwAdapter``. `GLFW <https://www.glfw.org/>`__ itself is cross-platform. Nevertheless, we have opted for this
layered design in order to separate generic from platform-specific functionality. If GLFW needs to be replaced with
another similar framework for some reason, only ``GlfwAdapter`` will need to be rewritten.
Themes and appearance
Individual UI elements do not allow customization in terms of appearance or layout. Instead, we use themes for colors
and spacing, and arrange all UI elements automatically. Several built-in themes are provided and the user can design
custom themes, however the entire UI uses a single theme for all elements. Appearance is minimalist: mostly colored
rectangles with text. Bitmaps and other custom decorations are not supported. The UI element types are check boxes,
radio button groups, selection lists, sliders, text edit boxes, static text, buttons, separators. These elements are
grouped into sections which can be expanded and collapsed.
Layout and rectangles
Each UI is one virtual rectangle, whose width is determined by the theme and whose height is determined by the
sections, items within each section, and also the expand/collapse state of each section. The sizes and auxiliary
buffers for these virtual rectangles are handled automatically when the UI is updated. Each UI has a visible rectangle
on the screen, and in addition there are other rectangles -- for 3D rendering, 2D figures, and possibly custom OpenGL
rendering. All these visible rectangles are saved (in :ref:`mjuiState`) and are used to determine where mouse events
should be directed. The rectangle layout is updated by a callback provided by the user.
Static allocation and creation
Rather than allocating and deallocating a large number of objects corresponding to UI elements and linking them
together, we create a single C struct (type :ref:`mjUI`) with static allocation supporting some maximum number of
sections and elements; and then keep a record of how many are in use. UI creation is simplified by helper functions
whose input is a C struct (type :ref:`mjuiDef`) that is essentially a table where each row describes one UI element
(see below). This makes it possible to construct elaborate user interfaces with surprisingly little C code.
Programmatic UI creation is also possible, for example when populating a UI with sliders corresponding to MuJoCo model
joints.
Minimal state
The UI is designed to be as stateless as possible, so as to simplify development. This has two aspects. First, instead
of replicating user data within the UI elements, we store pointers to user data. For example, we might create a UI
slider and set its data pointer to ``mjData* d->qpos+7``. This slider will visualize as well as control the 7th scalar
component of the qpos vector of a MuJoCo model. Thus, when the simulation is updated, we have to remember to update
the UI as well. And furthermore we have to disable UI editing when the simulation is being updated. But the advantage
is that the UI becomes easier to construct, and there is no danger of discrepancies between user data and the UI.
Second, the UI elements themselves are mostly stateless. Instead we keep track of a minimal set of global states, in
particular mouse and keyboard state, section expand/collapse, contents of the text box being edited if any.
Automated enable and disable
While each UI item can be set in enabled or disabled state directly, we also provide automation as follows. Each UI
item can be assigned an integer category. Then a :ref:`mjfItemEnable` callback determines whether each category should
be enabled or disabled, based on some program-specific conditions. For example, sliders that can change the values of
MuJoCo model joints should be disabled when the simulation state is being updated.
.. _uiAPI:
Main API
~~~~~~~~
Click on the links below for detailed API reference of the main UI data structures and functions.
**Main data structures:**
- :ref:`mjUI<mjUI>`: An entire UI.
- :ref:`mjuiState<mjuiState>`: Global UI state.
- :ref:`mjuiDef<mjuiDef>`: One entry in the definition table used for UI construction.
**Main functions:**
- :ref:`mjui_update<mjui_update>`: Main UI update function.
- :ref:`mjui_render<mjui_render>`: Renders the UI.
- :ref:`mjui_event<mjui_event>`: Low-level event handler.
- :ref:`mjui_add<mjui_add>`: Helper function used to construct a UI.
doc/python.rst
+1 -1
View File
@@ -215,7 +215,7 @@ Optionally, ``viewer.launch_passive`` accepts the following keyword arguments.
viewer.sync()
...
- ``show_left_ui_panel`` and ``show_right_ui_panel``: Boolean arguments indicating whether UI panels should be visible
- ``show_left_ui`` and ``show_right_ui``: Boolean arguments indicating whether UI panels should be visible
or hidden when the viewer is launched. Note that regardless of the values specified, the user can still toggle the
visibility of these panels after launch by pressing Tab or Shift+Tab.
include/mujoco/mjdata.h
+8 -4
View File
@@ -100,7 +100,7 @@ struct mjContact_ { // result of collision detection functions
// contact parameters set by near-phase collision function
mjtNum dist; // distance between nearest points; neg: penetration
mjtNum pos[3]; // position of contact point: midpoint between geoms
mjtNum frame[9]; // normal is in [0-2]
mjtNum frame[9]; // normal is in [0-2], points from geom[0] to geom[1]
// contact parameters set by mj_collideGeoms
mjtNum includemargin; // include if dist<includemargin=margin-gap
@@ -115,8 +115,8 @@ struct mjContact_ { // result of collision detection functions
// contact descriptors set by mj_collideXXX
int dim; // contact space dimensionality: 1, 3, 4 or 6
int geom1; // id of geom 1
int geom2; // id of geom 2
int geom1; // id of geom 1; deprecated, use geom[0]
int geom2; // id of geom 2; deprecated, use geom[1]
int geom[2]; // geom ids; -1 for flex
int flex[2]; // flex ids; -1 for geom
int elem[2]; // element ids; -1 for geom or flex vertex
@@ -320,7 +320,11 @@ struct mjData_ {
mjtNum* qfrc_bias; // C(qpos,qvel) (nv x 1)
// computed by mj_fwdVelocity/mj_passive
mjtNum* qfrc_passive; // passive force (nv x 1)
mjtNum* qfrc_spring; // passive spring force (nv x 1)
mjtNum* qfrc_damper; // passive damper force (nv x 1)
mjtNum* qfrc_gravcomp; // passive gravity compensation force (nv x 1)
mjtNum* qfrc_fluid; // passive fluid force (nv x 1)
mjtNum* qfrc_passive; // total passive force (nv x 1)
// computed by mj_sensorVel/mj_subtreeVel if needed
mjtNum* subtree_linvel; // linear velocity of subtree com (nbody x 3)
include/mujoco/mjmacro.h
+3
View File
@@ -34,6 +34,9 @@
#define mjDISABLED(x) (m->opt.disableflags & (x))
#define mjENABLED(x) (m->opt.enableflags & (x))
// is actuator disabled
#define mjACTUATORDISABLED(i) (m->opt.disableactuator & (1 << m->actuator_group[i]))
// annotation for functions that accept printf-like variadic arguments
#ifndef mjPRINTFLIKE
#if defined(__GNUC__)
include/mujoco/mjmodel.h
+5 -1
View File
@@ -433,6 +433,7 @@ struct mjOption_ { // physics options
int mpr_iterations; // maximum number of MPR solver iterations
int disableflags; // bit flags for disabling standard features
int enableflags; // bit flags for enabling optional features
int disableactuator; // bit flags for disabling actuators by group id
// sdf collision settings
int sdf_initpoints; // number of starting points for gradient descent
@@ -505,6 +506,7 @@ struct mjVisual_ { // visualization options
float framewidth; // bodyframe axis width
float constraint; // constraint width
float slidercrank; // slidercrank width
float frustum; // frustum zfar plane
} scale;
struct { // color of decor elements
@@ -530,6 +532,7 @@ struct mjVisual_ { // visualization options
float constraint[4]; // constraint
float slidercrank[4]; // slidercrank
float crankbroken[4]; // used when crank must be stretched/broken
float frustum[4]; // camera frustum
} rgba;
};
typedef struct mjVisual_ mjVisual;
@@ -661,7 +664,7 @@ struct mjModel_ {
int* body_treeid; // id of body's kinematic tree; -1: static (nbody x 1)
int* body_geomnum; // number of geoms (nbody x 1)
int* body_geomadr; // start addr of geoms; -1: no geoms (nbody x 1)
mjtByte* body_simple; // 1: diagonal M; 2: diag M, no rotations (nbody x 1)
mjtByte* body_simple; // 1: diag M; 2: diag M, sliders only (nbody x 1)
mjtByte* body_sameframe; // inertial frame is same as body frame (nbody x 1)
mjtNum* body_pos; // position offset rel. to parent body (nbody x 3)
mjtNum* body_quat; // orientation offset rel. to parent body (nbody x 4)
@@ -839,6 +842,7 @@ struct mjModel_ {
mjtNum* flex_edgedamping; // edge damping (nflex x 1)
mjtByte* flex_edgeequality; // is edge equality constraint defined (nflex x 1)
mjtByte* flex_rigid; // are all verices in the same body (nflex x 1)
mjtByte* flexedge_rigid; // are both edge vertices in same body (nflexedge x 1)
mjtByte* flex_centered; // are all vertex coordinates (0,0,0) (nflex x 1)
mjtByte* flex_flatskin; // render flex skin with flat shading (nflex x 1)
int* flex_bvhadr; // address of bvh root; -1: no bvh (nflex x 1)
include/mujoco/mjplugin.h
+4
View File
@@ -17,6 +17,7 @@
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjvisualize.h>
@@ -117,6 +118,9 @@ struct mjpPlugin_ {
// called during compilation for marching cubes
mjtNum (*sdf_staticdistance)(const mjtNum point[3], const mjtNum* attributes);
// convert attributes and provide defaults if not present
void (*sdf_attribute)(mjtNum attribute[], const char* name[], const char* value[]);
// bounding box of implicit surface
void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
};
include/mujoco/mjui.h
+1 -1
View File
@@ -18,7 +18,7 @@
#include <mujoco/mjrender.h>
#define mjMAXUISECT 10 // maximum number of sections
#define mjMAXUIITEM 100 // maximum number of items per section
#define mjMAXUIITEM 200 // maximum number of items per section
#define mjMAXUITEXT 300 // maximum number of chars in edittext and other
#define mjMAXUINAME 40 // maximum number of chars in name
#define mjMAXUIMULTI 35 // maximum number of radio/select items in group
include/mujoco/mjvisualize.h
+3
View File
@@ -422,6 +422,7 @@ struct mjvSceneState_ {
int nskin;
int nflex;
int nflexvert;
int nflextexcoord;
int nskinvert;
int nskinface;
int nskinbone;
@@ -486,6 +487,7 @@ struct mjvSceneState_ {
mjtNum* cam_fovy;
mjtNum* cam_ipd;
float* cam_intrinsic;
float* cam_sensorsize;
mjtByte* light_directional;
@@ -511,6 +513,7 @@ struct mjvSceneState_ {
int* flex_shell;
int* flex_shellnum;
int* flex_shelldataadr;
int* flex_texcoordadr;
int* flex_bvhadr;
int* flex_bvhnum;
mjtNum* flex_radius;
include/mujoco/mjxmacro.h
+9 -3
View File
@@ -43,6 +43,7 @@
X( int, mpr_iterations ) \
X( int, disableflags ) \
X( int, enableflags ) \
X( int, disableactuator ) \
X( int, sdf_initpoints ) \
X( int, sdf_iterations )
@@ -86,7 +87,7 @@
X ( nflexelemdata ) \
X ( nflexshelldata ) \
X ( nflexevpair ) \
X ( nflextexcoord ) \
XMJV( nflextexcoord ) \
XMJV( nmesh ) \
X ( nmeshvert ) \
X ( nmeshnormal ) \
@@ -276,7 +277,7 @@
X ( int, cam_targetbodyid, ncam, 1 ) \
X ( int, cam_resolution, ncam, 2 ) \
XMJV( float, cam_sensorsize, ncam, 2 ) \
X ( float, cam_intrinsic, ncam, 4 ) \
XMJV( float, cam_intrinsic, ncam, 4 ) \
X ( mjtNum, cam_pos, ncam, 3 ) \
X ( mjtNum, cam_quat, ncam, 4 ) \
X ( mjtNum, cam_poscom0, ncam, 3 ) \
@@ -329,7 +330,7 @@
XMJV( int, flex_shelldataadr, nflex, 1 ) \
X ( int, flex_evpairadr, nflex, 1 ) \
X ( int, flex_evpairnum, nflex, 1 ) \
X ( int, flex_texcoordadr, nflex, 1 ) \
XMJV( int, flex_texcoordadr, nflex, 1 ) \
X ( int, flex_vertbodyid, nflexvert, 1 ) \
X ( int, flex_edge, nflexedge, 2 ) \
XMJV( int, flex_elem, nflexelemdata, 1 ) \
@@ -345,6 +346,7 @@
X ( mjtNum, flex_edgedamping, nflex, 1 ) \
X ( mjtByte, flex_edgeequality, nflex, 1 ) \
X ( mjtByte, flex_rigid, nflex, 1 ) \
X ( mjtByte, flexedge_rigid, nflexedge, 1 ) \
X ( mjtByte, flex_centered, nflex, 1 ) \
XMJV( mjtByte, flex_flatskin, nflex, 1 ) \
XMJV( int, flex_bvhadr, nflex, 1 ) \
@@ -617,6 +619,10 @@
X ( mjtNum, cvel, nbody, 6 ) \
X ( mjtNum, cdof_dot, nv, 6 ) \
X ( mjtNum, qfrc_bias, nv, 1 ) \
X ( mjtNum, qfrc_spring, nv, 1 ) \
X ( mjtNum, qfrc_damper, nv, 1 ) \
X ( mjtNum, qfrc_gravcomp, nv, 1 ) \
X ( mjtNum, qfrc_fluid, nv, 1 ) \
X ( mjtNum, qfrc_passive, nv, 1 ) \
X ( mjtNum, subtree_linvel, nbody, 3 ) \
X ( mjtNum, subtree_angmom, nbody, 3 ) \
include/mujoco/mujoco.h
+2 -2
View File
@@ -336,10 +336,10 @@ MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, i
// Compute cvel, cdof_dot.
MJAPI void mj_comVel(const mjModel* m, mjData* d);
// Compute qfrc_passive from spring-dampers, viscosity and density.
// Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.
MJAPI void mj_passive(const mjModel* m, mjData* d);
// subtree linear velocity and angular momentum
// Sub-tree linear velocity and angular momentum: compute subtree_linvel, subtree_angmom.
MJAPI void mj_subtreeVel(const mjModel* m, mjData* d);
// RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term.
introspect/functions.py
+2 -2
View File
@@ -1666,7 +1666,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
),
),
doc='Compute qfrc_passive from spring-dampers, viscosity and density.',
doc='Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.', # pylint: disable=line-too-long
)),
('mj_subtreeVel',
FunctionDecl(
@@ -1686,7 +1686,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
),
),
),
doc='subtree linear velocity and angular momentum',
doc='Sub-tree linear velocity and angular momentum: compute subtree_linvel, subtree_angmom.', # pylint: disable=line-too-long
)),
('mj_rne',
FunctionDecl(
introspect/structs.py
+78 -6
View File
@@ -276,6 +276,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=ValueType(name='int'),
doc='bit flags for enabling optional features',
),
StructFieldDecl(
name='disableactuator',
type=ValueType(name='int'),
doc='bit flags for disabling actuators by group id',
),
StructFieldDecl(
name='sdf_initpoints',
type=ValueType(name='int'),
@@ -578,6 +583,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=ValueType(name='float'),
doc='slidercrank width',
),
StructFieldDecl(
name='frustum',
type=ValueType(name='float'),
doc='frustum zfar plane',
),
),
),
doc='',
@@ -762,6 +772,14 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='used when crank must be stretched/broken',
),
StructFieldDecl(
name='frustum',
type=ArrayType(
inner_type=ValueType(name='float'),
extents=(4,),
),
doc='camera frustum',
),
),
),
doc='',
@@ -1311,7 +1329,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=PointerType(
inner_type=ValueType(name='mjtByte'),
),
doc='1: diagonal M; 2: diag M, no rotations (nbody x 1)',
doc='1: diag M; 2: diag M, sliders only (nbody x 1)',
),
StructFieldDecl(
name='body_sameframe',
@@ -2426,6 +2444,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='are all verices in the same body (nflex x 1)',
),
StructFieldDecl(
name='flexedge_rigid',
type=PointerType(
inner_type=ValueType(name='mjtByte'),
),
doc='are both edge vertices in same body (nflexedge x 1)',
),
StructFieldDecl(
name='flex_centered',
type=PointerType(
@@ -3846,7 +3871,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
inner_type=ValueType(name='mjtNum'),
extents=(9,),
),
doc='normal is in [0-2]',
doc='normal is in [0-2], points from geom[0] to geom[1]',
),
StructFieldDecl(
name='includemargin',
@@ -3906,12 +3931,12 @@ STRUCTS: Mapping[str, StructDecl] = dict([
StructFieldDecl(
name='geom1',
type=ValueType(name='int'),
doc='id of geom 1',
doc='id of geom 1; deprecated, use geom[0]',
),
StructFieldDecl(
name='geom2',
type=ValueType(name='int'),
doc='id of geom 2',
doc='id of geom 2; deprecated, use geom[1]',
),
StructFieldDecl(
name='geom',
@@ -4673,12 +4698,40 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='C(qpos,qvel) (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='qfrc_spring',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='passive spring force (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='qfrc_damper',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='passive damper force (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='qfrc_gravcomp',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='passive gravity compensation force (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='qfrc_fluid',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='passive fluid force (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='qfrc_passive',
type=PointerType(
inner_type=ValueType(name='mjtNum'),
),
doc='passive force (nv x 1)', # pylint: disable=line-too-long
doc='total passive force (nv x 1)', # pylint: disable=line-too-long
),
StructFieldDecl(
name='subtree_linvel',
@@ -6190,6 +6243,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
type=ValueType(name='int'),
doc='',
),
StructFieldDecl(
name='nflextexcoord',
type=ValueType(name='int'),
doc='',
),
StructFieldDecl(
name='nskinvert',
type=ValueType(name='int'),
@@ -6557,6 +6615,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='',
),
StructFieldDecl(
name='cam_intrinsic',
type=PointerType(
inner_type=ValueType(name='float'),
),
doc='',
),
StructFieldDecl(
name='cam_sensorsize',
type=PointerType(
@@ -6718,6 +6783,13 @@ STRUCTS: Mapping[str, StructDecl] = dict([
),
doc='',
),
StructFieldDecl(
name='flex_texcoordadr',
type=PointerType(
inner_type=ValueType(name='int'),
),
doc='',
),
StructFieldDecl(
name='flex_bvhadr',
type=PointerType(
@@ -8403,7 +8475,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
name='item',
type=ArrayType(
inner_type=ValueType(name='mjuiItem'),
extents=(100,),
extents=(200,),
),
doc='preallocated array of items',
),
mjx/mujoco/mjx/_src/collision_driver.py
+21 -6
View File
@@ -336,16 +336,28 @@ def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet:
def ncon(m: Model) -> int:
"""Returns the number of contacts computed in MJX given a model."""
if m.opt.disableflags & DisableBit.CONTACT:
return 0
candidates = collision_candidates(m)
max_count = _max_contact_points(m)
count = sum([
len(v) * get_collision_fn(k[0:2]).ncon for k, v in candidates.items() # pytype: disable=attribute-error
])
count = 0
for k, v in candidates.items():
fn = get_collision_fn(k[0:2])
if fn is None:
continue
count += len(v) * fn.ncon # pytype: disable=attribute-error
return min(max_count, count) if max_count > -1 else count
def collision(m: Model, d: Data) -> Data:
"""Collides geometries."""
ncon_ = ncon(m)
if ncon_ == 0:
return d.replace(contact=Contact.zero(), ncon=0)
candidate_set = collision_candidates(m)
contacts = []
@@ -354,7 +366,7 @@ def collision(m: Model, d: Data) -> Data:
contacts.append(_collide_geoms(m, d, geom_types, candidates))
if not contacts:
return d.replace(contact=Contact.zero(), ncon=0)
raise RuntimeError('No contacts found.')
contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts)
@@ -364,10 +376,13 @@ def collision(m: Model, d: Data) -> Data:
_, idx = jax.lax.top_k(-contact.dist, k=max_contact_points)
contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact)
ncon_ = contact.dist.shape[0]
if ncon_ != contact.dist.shape[0]:
raise RuntimeError('Number of contacts does not match ncon.')
# TODO(robotics-simulation): move this logic to device_put
ns = d.ne + d.nf + d.nl
contact = contact.replace(efc_address=np.arange(ns, ns + ncon_ * 4, 4))
# TODO(robotics-simulation): add support for other friction dimensions
contact = contact.replace(efc_address=np.arange(ns, ns + d.ncon * 4, 4))
contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32))
return d.replace(contact=contact, ncon=ncon_)
mjx/mujoco/mjx/_src/collision_driver_test.py
+28 -2
View File
@@ -24,11 +24,14 @@ import jax
import jax.numpy as jp
import mujoco
from mujoco import mjx
from mujoco.mjx._src import collision_driver
from mujoco.mjx._src import test_util
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import Contact
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import Model
# pylint: emable=g-importing-member
# pylint: enable=g-importing-member
import numpy as np
@@ -289,7 +292,7 @@ class ConvexTest(absltest.TestCase):
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="box"/>
</body>
<body pos="0.1 1.0 0.495" euler="0.1 -0.1 45">
<body pos="0.1 1.0 0.495" euler="0.1 -0.1 0">
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1 0.1" type="box"/>
</body>
@@ -450,6 +453,29 @@ class BodyPairFilterTest(absltest.TestCase):
self.assertEqual(dx.contact.pos.shape[0], 1)
class NconTest(parameterized.TestCase):
"""Tests ncon."""
def test_ncon(self):
m = test_util.load_test_file('ant.xml')
d = mujoco.MjData(m)
d.qpos[2] = 0.0
mx = mjx.device_put(m)
ncon = collision_driver.ncon(mx)
self.assertEqual(ncon, 4)
def test_disable_contact(self):
m = test_util.load_test_file('ant.xml')
d = mujoco.MjData(m)
d.qpos[2] = 0.0
m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
mx = mjx.device_put(m)
ncon = collision_driver.ncon(mx)
self.assertEqual(ncon, 0)
class TopKContactTest(absltest.TestCase):
"""Tests top-k contacts."""
mjx/mujoco/mjx/_src/constraint.py
+186 -143
View File
@@ -14,13 +14,12 @@
# ==============================================================================
"""Core non-smooth constraint functions."""
from typing import Tuple
from typing import Optional, Tuple
import jax
from jax import numpy as jp
import mujoco
from mujoco.mjx._src import math
from mujoco.mjx._src import scan
from mujoco.mjx._src import support
# pylint: disable=g-importing-member
from mujoco.mjx._src.dataclasses import PyTreeNode
@@ -35,16 +34,15 @@ import numpy as np
class _Efc(PyTreeNode):
"""Support data for creating constraint matrices."""
J: jax.Array
R: jax.Array
aref: jax.Array
pos: jax.Array
pos_norm: jax.Array
invweight: jax.Array
solref: jax.Array
solimp: jax.Array
frictionloss: jax.Array
@classmethod
def zero(cls, m: Model) -> '_Efc':
z = jp.empty((0,))
return _Efc(J=jp.empty((0, m.nv)), R=z, aref=z, frictionloss=z)
def _kbi(
m: Model,
@@ -84,22 +82,18 @@ def _kbi(
return k, b, imp # corresponds to K, B, I of efc_KBIP
def _instantiate_connect(m: Model, d: Data) -> _Efc:
"""Returns jacobians and supporting data for connect equality constraints."""
def _instantiate_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for connect equality constraints."""
if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0:
return _Efc.zero(m)
ids = np.nonzero(m.eq_type == EqType.CONNECT)[0]
connect_id = np.nonzero(m.eq_type == EqType.CONNECT)[0]
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
return None
if connect_id.size == 0:
return _Efc.zero(m)
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
body1id, body2id = m.eq_obj1id[connect_id], m.eq_obj2id[connect_id]
data = m.eq_data[connect_id]
solref, solimp = m.eq_solref[connect_id], m.eq_solimp[connect_id]
def fn(data, id1, id2, solref, solimp):
@jax.vmap
def fn(data, id1, id2):
anchor1, anchor2 = data[0:3], data[3:6]
# find global points
pos1 = d.xmat[id1] @ anchor1 + d.xpos[id1]
@@ -113,35 +107,31 @@ def _instantiate_connect(m: Model, d: Data) -> _Efc:
jacp2, _ = support.jac(m, d, pos2, id2)
j = (jacp1 - jacp2).T
# impedance, inverse constraint mass, reference acceleration
k, b, imp = _kbi(m, solref, solimp, math.norm(cpos))
invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0]
r = jp.maximum(invweight * (1 - imp) / imp, mujoco.mjMINVAL).repeat(3)
aref = -b * (j @ d.qvel) - k * imp * cpos
return j, cpos, jp.repeat(math.norm(cpos), 3)
return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r))
# concatenate to drop connect grouping dimension
j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0]
invweight = jp.repeat(invweight, 3)
solref = jp.tile(m.eq_solref[ids], (3, 1))
solimp = jp.tile(m.eq_solimp[ids], (3, 1))
frictionloss = jp.zeros_like(pos_norm)
efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp)
return jax.tree_map(jp.concatenate, efcs)
return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
def _instantiate_weld(m: Model, d: Data) -> _Efc:
"""Returns jacobians and supporting data for connect weld constraints."""
def _instantiate_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for weld equality constraints."""
if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0:
return _Efc.zero(m)
ids = np.nonzero(m.eq_type == EqType.WELD)[0]
weld_id = np.nonzero(m.eq_type == EqType.WELD)[0]
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
return None
if weld_id.size == 0:
return _Efc.zero(m)
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
body1id, body2id = m.eq_obj1id[weld_id], m.eq_obj2id[weld_id]
data = m.eq_data[weld_id]
solref, solimp = m.eq_solref[weld_id], m.eq_solimp[weld_id]
def fn(data, id1, id2, solref, solimp):
@jax.vmap
def fn(data, id1, id2):
anchor1, anchor2 = data[0:3], data[3:6]
relpose, torquescale = data[6:10], data[10]
@@ -168,119 +158,159 @@ def _instantiate_weld(m: Model, d: Data) -> _Efc:
jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr)
j = jp.concatenate((jacdifp.T, jacdifr.T))
pos = jp.concatenate((cpos, crot))
pos = jp.concatenate((cpos, crot * torquescale))
# impedance, inverse constraint mass, reference acceleration
k, b, imp = _kbi(m, solref, solimp, math.norm(pos.at[3:].mul(torquescale)))
invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
r = jp.maximum(invweight * (1 - imp) / imp, mujoco.mjMINVAL).repeat(3)
aref = -b * (j @ d.qvel) - k * imp * pos
return j, pos, jp.repeat(math.norm(pos), 6)
return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r))
# concatenate to drop weld grouping dimension
j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
invweight = jp.repeat(invweight, 3)
solref = jp.tile(m.eq_solref[ids], (6, 1))
solimp = jp.tile(m.eq_solimp[ids], (6, 1))
frictionloss = jp.zeros_like(pos_norm)
efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp)
return jax.tree_map(jp.concatenate, efcs)
return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
def _instantiate_friction(m: Model, d: Data) -> _Efc:
def _instantiate_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for joint equality constraints."""
ids = np.nonzero(m.eq_type == EqType.JOINT)[0]
if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
return None
id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
dofadr1, dofadr2 = m.jnt_dofadr[id1], m.jnt_dofadr[id2]
qposadr1, qposadr2 = m.jnt_qposadr[id1], m.jnt_qposadr[id2]
@jax.vmap
def fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2):
pos1, pos2 = d.qpos[qposadr1], d.qpos[qposadr2]
ref1, ref2 = m.qpos0[qposadr1], m.qpos0[qposadr2]
pos2, ref2 = pos2 * (id2 > -1), ref2 * (id2 > -1)
dif = pos2 - ref2
dif_power = jp.power(dif, jp.arange(0, 5))
deriv = jp.dot(data[1:5], dif_power[:4] * jp.arange(1, 5))
j = jp.zeros((m.nv)).at[dofadr1].set(1.0).at[dofadr2].set(-deriv)
pos = pos1 - ref1 - jp.dot(data[:5], dif_power)
return j, pos
j, pos = fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2)
invweight = m.dof_invweight0[dofadr1] + m.dof_invweight0[dofadr2] * (id2 > -1)
solref, solimp = m.eq_solref[ids], m.eq_solimp[ids]
frictionloss = jp.zeros_like(pos)
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
def _instantiate_friction(m: Model, d: Data) -> Optional[_Efc]:
# TODO(robotics-team): implement _instantiate_friction
del d
return _Efc.zero(m)
del m, d
return None
def _instantiate_limit(m: Model, d: Data) -> _Efc:
"""Returns jacobians and supporting data for joint limits."""
def _instantiate_limit_ball(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for ball joint limits."""
if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any():
return _Efc.zero(m)
ids = np.nonzero((m.jnt_type == JointType.BALL) & m.jnt_limited)[0]
def fn(jnt_typs, jnt_range, solref, solimp, margin, qpos, dofs, invweight0):
js, rs, arefs = [], [], []
qpos_i, dof_i = 0, 0
if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
return None
for i in range(len(jnt_typs)):
jnt_typ = JointType(jnt_typs[i])
jnt_range = m.jnt_range[ids]
jnt_margin = m.jnt_margin[ids]
qposadr = np.array([np.arange(q, q + 4) for q in m.jnt_qposadr[ids]])
dofadr = np.array([np.arange(d, d + 3) for d in m.jnt_dofadr[ids]])
if jnt_typ == JointType.FREE:
return None # omit constraint rows for free joints
elif jnt_typ == JointType.BALL:
axis, angle = math.quat_to_axis_angle(qpos[qpos_i : qpos_i + 4])
dist = jp.amax(jnt_range[i]) - angle
j = jp.sum(
jax.vmap(jp.multiply)(dofs[dof_i : dof_i + 3], -axis), axis=0
)
elif jnt_typ in (JointType.HINGE, JointType.SLIDE):
dist_min = qpos[qpos_i] - jnt_range[i, 0]
dist_max = jnt_range[i, 1] - qpos[qpos_i]
dist = jp.minimum(dist_min, dist_max)
j = dofs[dof_i] * ((dist_min < dist_max) * 2 - 1)
else:
raise RuntimeError(f'unrecognized joint type: {jnt_typ}')
@jax.vmap
def fn(jnt_range, jnt_margin, qposadr, dofadr):
axis, angle = math.quat_to_axis_angle(d.qpos[qposadr])
j = jp.zeros(m.nv).at[dofadr].set(-axis)
pos = jp.amax(jnt_range) - angle - jnt_margin
active = pos < 0
return j * active, pos * active
dist = dist - margin[i]
k, b, imp = _kbi(m, solref[i], solimp[i], dist)
r = jp.maximum(invweight0[dof_i] * (1 - imp) / imp, mujoco.mjMINVAL)
aref = -b * (j @ d.qvel) - k * imp * dist
j, aref = j * (dist < 0), aref * (dist < 0)
js, rs, arefs = js + [j], rs + [r], arefs + [aref]
dof_i, qpos_i = dof_i + jnt_typ.dof_width(), qpos_i + jnt_typ.qpos_width()
j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
invweight = m.dof_invweight0[m.jnt_dofadr[ids]]
solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
frictionloss = jp.zeros_like(pos)
return jp.stack(js), jp.stack(rs), jp.stack(arefs)
j, r, aref = scan.flat(
m,
fn,
'jjjjjqvv',
'jjj',
m.jnt_type,
m.jnt_range,
m.jnt_solref,
m.jnt_solimp,
m.jnt_margin,
d.qpos,
jp.eye(m.nv),
m.dof_invweight0,
)
return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r))
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
def _instantiate_contact(m: Model, d: Data) -> _Efc:
"""Returns jacobians and supporitng data for contacts."""
def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for slide and hinge joint limits."""
slide_hinge = np.isin(m.jnt_type, (JointType.SLIDE, JointType.HINGE))
ids = np.nonzero(slide_hinge & m.jnt_limited)[0]
if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
return None
jnt_range = m.jnt_range[ids]
jnt_margin = m.jnt_margin[ids]
qposadr = m.jnt_qposadr[ids]
dofadr = m.jnt_dofadr[ids]
@jax.vmap
def fn(jnt_range, jnt_margin, qposadr, dofadr):
dist_min = d.qpos[qposadr] - jnt_range[0]
dist_max = jnt_range[1] - d.qpos[qposadr]
j = jp.zeros(m.nv).at[dofadr].set((dist_min < dist_max) * 2 - 1)
pos = jp.minimum(dist_min, dist_max) - jnt_margin
active = pos < 0
return j * active, pos * active
j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
invweight = m.dof_invweight0[dofadr]
solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
frictionloss = jp.zeros_like(pos)
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
"""Calculates constraint rows for contacts."""
if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0:
return _Efc.zero(m)
def fn(contact: Contact):
dist = contact.dist - contact.includemargin
k, b, imp = _kbi(m, contact.solref, contact.solimp, dist)
return None
@jax.vmap
def fn(c: Contact):
dist = c.dist - c.includemargin
geom_bodyid = jp.array(m.geom_bodyid)
body1, body2 = geom_bodyid[contact.geom1], geom_bodyid[contact.geom2]
diff = support.jac_dif_pair(m, d, contact.pos, body1, body2)
body1, body2 = geom_bodyid[c.geom1], geom_bodyid[c.geom2]
diff = support.jac_dif_pair(m, d, c.pos, body1, body2)
t = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0]
# rotate Jacobian differences to contact frame
diff_con = contact.frame @ diff.T
diff_con = c.frame @ diff.T
# TODO(robotics-simulation): add support for other friction dimensions
# 4 pyramidal friction directions
js, rs = [], []
for diff_tan, friction in zip(diff_con[1:], contact.friction[:2]):
js, invweights = [], []
for diff_tan, friction in zip(diff_con[1:], c.friction[:2]):
for f in (friction, -friction):
js.append(diff_con[0] + diff_tan * f)
rs.append((t + f * f * t) * 2 * f * f * (1 - imp) / imp)
invweights.append((t + f * f * t) * 2 * f * f)
j, r = jp.stack(js), jp.stack(rs)
r = jp.maximum(r, mujoco.mjMINVAL)
aref = -b * (j @ d.qvel) - k * imp * dist
mask_fn = jax.vmap(lambda x, mask=(dist < 0): x * mask)
j, aref = jax.tree_map(mask_fn, (j, aref))
active = dist < 0
j, invweight = jp.stack(js) * active, jp.stack(invweights)
pos = jp.repeat(dist, 4) * active
solref, solimp = jp.tile(c.solref, (4, 1)), jp.tile(c.solimp, (4, 1))
return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r))
return j, invweight, pos, solref, solimp
return jax.tree_map(jp.concatenate, jax.vmap(fn)(d.contact))
res = fn(d.contact)
# remove contact grouping dimension:
j, invweight, pos, solref, solimp = jax.tree_map(jp.concatenate, res)
frictionloss = jp.zeros_like(pos)
return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
@@ -291,18 +321,19 @@ def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]:
if m.opt.disableflags & DisableBit.EQUALITY:
ne = 0
else:
ne_weld = (m.eq_type == EqType.WELD).sum()
ne_connect = (m.eq_type == EqType.CONNECT).sum()
ne = ne_weld * 6 + ne_connect * 3
ne_weld = (m.eq_type == EqType.WELD).sum()
ne_joint = (m.eq_type == EqType.JOINT).sum()
ne = ne_connect * 3 + ne_weld * 6 + ne_joint
nf = 0
if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any():
if m.opt.disableflags & DisableBit.LIMIT:
nl = 0
else:
nl = (m.jnt_type != JointType.FREE).sum()
nl = int(m.jnt_limited.sum())
if (m.opt.disableflags & DisableBit.CONTACT):
if m.opt.disableflags & DisableBit.CONTACT:
nc = 0
else:
nc = d.ncon * 4
@@ -318,23 +349,35 @@ def make_constraint(m: Model, d: Data) -> Data:
d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)})
if m.opt.disableflags & DisableBit.CONSTRAINT:
efc = _Efc.zero(m)
efcs = ()
else:
efcs = (
_instantiate_connect(m, d),
_instantiate_weld(m, d),
efcs = tuple(efc for efc in (
_instantiate_equality_connect(m, d),
_instantiate_equality_weld(m, d),
_instantiate_equality_joint(m, d),
_instantiate_friction(m, d),
_instantiate_limit(m, d),
_instantiate_limit_ball(m, d),
_instantiate_limit_slide_hinge(m, d),
_instantiate_contact(m, d),
)
efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
) if efc is not None)
d = d.replace(
efc_J=efc.J,
efc_D=1 / efc.R,
efc_aref=efc.aref,
efc_frictionloss=efc.frictionloss,
nefc=efc.aref.shape[0],
)
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, nefc=0)
return d
efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
@jax.vmap
def fn(efc):
k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_norm)
r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL)
aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos
return aref, r
aref, r = fn(efc)
d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
d = d.replace(efc_frictionloss=efc.frictionloss, nefc=r.shape[0])
return d
mjx/mujoco/mjx/_src/constraint_test.py
+39 -4
View File
@@ -24,11 +24,12 @@ from mujoco.mjx._src import constraint
from mujoco.mjx._src import test_util
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import SolverType
# pylint: enable=g-importing-member
import numpy as np
def _assert_eq(a, b, name, step, fname, atol=1e-3, rtol=1e-3):
def _assert_eq(a, b, name, step, fname, atol=5e-3, rtol=5e-3):
err_msg = f'mismatch: {name} at step {step} in {fname}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
@@ -36,7 +37,7 @@ def _assert_eq(a, b, name, step, fname, atol=1e-3, rtol=1e-3):
class ConstraintTest(parameterized.TestCase):
@parameterized.parameters(enumerate(test_util.TEST_FILES))
def testconstraints(self, seed, fname):
def test_constraints(self, seed, fname):
"""Test constraints."""
np.random.seed(seed)
@@ -79,6 +80,40 @@ class ConstraintTest(parameterized.TestCase):
fname,
)
_JNT_RANGE = """
<mujoco>
<worldbody>
<body pos="0 0 1">
<joint type="slide" axis="1 0 0" range="-1.8 1.8" solreflimit=".08 1"
damping="5e-4"/>
<geom type="box" size="0.2 0.15 0.1" mass="1"/>
<body>
<joint axis="0 1 0" damping="2e-6"/>
<geom type="capsule" fromto="0 0 0 0 0 1" size="0.045" mass=".1"/>
</body>
</body>
</worldbody>
</mujoco>
"""
def test_jnt_range(self):
"""Tests that mixed joint ranges are respected."""
# TODO(robotics-simulation): also test ball
m = mujoco.MjModel.from_xml_string(self._JNT_RANGE)
m.opt.solver = SolverType.CG.value
d = mujoco.MjData(m)
d.qpos = np.array([2.0, 15.0])
mx = mjx.device_put(m)
dx = mjx.device_put(d)
efc = jax.jit(constraint._instantiate_limit_slide_hinge)(mx, dx)
# first joint is outside the joint range
np.testing.assert_array_almost_equal(efc.J[0, 0], -1.0)
# second joint has no range, so only one efc row
self.assertEqual(efc.J.shape[0], 1)
def test_disable_refsafe(self):
m = test_util.load_test_file('ant.xml')
@@ -112,7 +147,7 @@ class ConstraintTest(parameterized.TestCase):
self.assertEqual(dx.efc_J.shape[0], 0)
def test_disable_equality(self):
m = test_util.load_test_file('weld.xml')
m = test_util.load_test_file('equality.xml')
d = mujoco.MjData(m)
m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY
@@ -137,7 +172,7 @@ class ConstraintTest(parameterized.TestCase):
m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT
mx, dx = mjx.device_put(m), mjx.device_put(d)
efc = constraint._instantiate_contact(mx, dx)
self.assertEqual(efc.J.shape[0], 0)
self.assertIsNone(efc)
if __name__ == '__main__':
mjx/mujoco/mjx/_src/device.py
+10 -6
View File
@@ -40,6 +40,10 @@ _MJ_TYPE_ATTR = {
mujoco.MjModel.opt,
mujoco.MjOption.integrator,
),
mujoco.mjtSolver: (
mujoco.MjModel.opt,
mujoco.MjOption.solver,
),
}
_TYPE_MAP = {
@@ -103,10 +107,6 @@ def _option_derived(value: types.Option) -> Dict[str, Any]:
def _validate(m: mujoco.MjModel):
"""Validates that an mjModel is compatible with MJX."""
if m.opt.solver not in set(types.SolverType):
name = mujoco.mjtSolver(m.opt.solver).name
warnings.warn(f'Solver {name} is not supported, reverting to CG.')
m.opt.solver = mujoco.mjtSolver.mjSOL_CG.value
# check enum types
for mj_type, attrs in _MJ_TYPE_ATTR.items():
@@ -116,8 +116,9 @@ def _validate(m: mujoco.MjModel):
typs = set(val) if isinstance(val, Iterable) else {val}
unsupported_typs = typs - set(_TYPE_MAP[mj_type])
if unsupported_typs:
raise NotImplementedError(f'{unsupported_typs} not implemented.')
unsupported = [mj_type(t) for t in unsupported_typs] # pylint: disable=too-many-function-args
if unsupported:
raise NotImplementedError(f'{unsupported} not implemented.')
# check condim
if any(dim != 3 for dim in m.geom_condim) or any(
@@ -125,6 +126,9 @@ def _validate(m: mujoco.MjModel):
):
raise NotImplementedError('Only condim=3 is supported.')
if m.ntendon:
raise NotImplementedError('Tendons are not supported.')
# check collision geom types
candidate_set = collision_driver.collision_candidates(m)
for g1, g2, *_ in candidate_set:
mjx/mujoco/mjx/_src/device_test.py
+35 -12
View File
@@ -110,56 +110,55 @@ class ValidateInputTest(absltest.TestCase):
def test_solver(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option solver="Newton"/><worldbody/></mujoco>'
'<mujoco><option solver="PGS"/><worldbody/></mujoco>'
)
with self.assertWarns(UserWarning):
mx = mjx.device_put(m)
self.assertEqual(mx.opt.solver, mujoco.mjtSolver.mjSOL_CG)
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
def test_integrator(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_cone(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_trn(self):
m = test_util.load_test_file('ant.xml')
m.actuator_trntype[0] = mujoco.mjtTrn.mjTRN_SITE
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_dyn(self):
m = test_util.load_test_file('ant.xml')
m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_gain(self):
m = test_util.load_test_file('ant.xml')
m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_bias(self):
m = test_util.load_test_file('ant.xml')
m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_condim(self):
m = test_util.load_test_file('ant.xml')
for i in [1, 4, 6]:
m.geom_condim[0] = i
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_geoms(self):
m = mujoco.MjModel.from_xml_string("""
@@ -177,7 +176,31 @@ class ValidateInputTest(absltest.TestCase):
</mujoco>
""")
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
mjx.device_put(m)
def test_tendon(self):
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<worldbody>
<body name="left_thigh" pos="0 0.1 -0.04">
<joint axis="0 1 0" name="left_hip_y" type="hinge"/>
<geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
<body name="left_shin" pos="0 -0.01 -0.403">
<joint axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
<geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
</body>
</body>
</worldbody>
<tendon>
<fixed name="left_hipknee">
<joint coef="-1" joint="left_hip_y"/>
<joint coef="1" joint="left_knee"/>
</fixed>
</tendon>
</mujoco>
""")
with self.assertRaises(NotImplementedError):
mjx.device_put(m)
if __name__ == '__main__':
mjx/mujoco/mjx/_src/forward.py
+1 -5
View File
@@ -37,7 +37,6 @@ 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 SolverType
# pylint: enable=g-importing-member
import numpy as np
@@ -333,10 +332,7 @@ def forward(m: Model, d: Data) -> Data:
d = d.replace(qacc=d.qacc_smooth)
return d
if m.opt.solver == SolverType.CG:
d = named_scope(solver.cg_solve)(m, d)
else:
raise NotImplementedError(f'solver {m.opt.solver} not implemented.')
d = named_scope(solver.solve)(m, d)
return d
mjx/mujoco/mjx/_src/forward_test.py
+56 -34
View File
@@ -14,8 +14,6 @@
# ==============================================================================
"""Tests for forward functions."""
import itertools
from absl.testing import absltest
from absl.testing import parameterized
import jax
@@ -38,13 +36,12 @@ def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3):
class ForwardTest(parameterized.TestCase):
@parameterized.parameters(enumerate(test_util.TEST_FILES))
def test_forward(self, seed, fname):
@parameterized.parameters(
filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
)
def test_forward(self, fname):
"""Test mujoco mj forward function matches mujoco_mjx forward function."""
if fname in ('weld.xml',):
return
np.random.seed(seed)
np.random.seed(test_util.TEST_FILES.index(fname))
m = test_util.load_test_file(fname)
d = mujoco.MjData(m)
@@ -62,36 +59,20 @@ class ForwardTest(parameterized.TestCase):
_assert_attr_eq(d, dx, 'qfrc_smooth', i, fname)
_assert_attr_eq(d, dx, 'qacc_smooth', i, fname)
@parameterized.parameters(itertools.product(test_util.TEST_FILES, (0, 1)))
def test_step(self, fname, integrator_type):
@parameterized.parameters(
filter(lambda s: s not in ('equality.xml',), test_util.TEST_FILES)
)
def test_step(self, fname):
"""Test mujoco mj step matches mujoco_mjx step."""
if fname in (
'mixed_joint_pendulum.xml',
'ball_pendulum.xml',
'convex.xml',
'humanoid.xml',
'triple_pendulum.xml', # TODO(b/301485081)
'weld.xml',
):
# skip models with big constraint violations at step 0 or too slow to run
return
np.random.seed(integrator_type)
np.random.seed(test_util.TEST_FILES.index(fname))
m = test_util.load_test_file(fname)
step_jit_fn = jax.jit(forward.step)
m.opt.integrator = integrator_type
int_typ = 'euler' if integrator_type == 0 else 'rk4'
test_name = f'{fname} - {int_typ}'
steps = 100 if int_typ == 'euler' else 30
dt = m.opt.timestep
m.opt.timestep = dt if int_typ == 'euler' else dt * 3
mx = mjx.device_put(m)
d = mujoco.MjData(m)
# give the system a little kick to ensure we have non-identity rotations
d.qvel = np.random.normal(m.nv) * 0.05
for i in range(steps):
for i in range(100):
# in order to avoid re-jitting, reuse the same mj_data shape
qpos, qvel = d.qpos, d.qvel
d = mujoco.MjData(m)
@@ -101,10 +82,51 @@ class ForwardTest(parameterized.TestCase):
mujoco.mj_step(m, d)
dx = step_jit_fn(mx, dx)
_assert_attr_eq(d, dx, 'qpos', i, test_name, atol=1e-2)
_assert_attr_eq(d, dx, 'qvel', i, test_name, atol=1e-2)
_assert_attr_eq(d, dx, 'act', i, test_name)
_assert_attr_eq(d, dx, 'time', i, test_name)
_assert_attr_eq(d, dx, 'qvel', i, fname, atol=1e-2)
_assert_attr_eq(d, dx, 'qpos', i, fname, atol=1e-2)
_assert_attr_eq(d, dx, 'act', i, fname)
_assert_attr_eq(d, dx, 'time', i, fname)
def test_rk4(self):
m = mujoco.MjModel.from_xml_string("""
<mujoco>
<option integrator="RK4">
<flag constraint="disable"/>
</option>
<worldbody>
<light pos="0 0 1"/>
<geom type="plane" size="1 1 .01" pos="0 0 -1"/>
<body pos="0.15 0 0">
<joint type="hinge" axis="0 1 0"/>
<geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
<body pos="0.1 0 0">
<joint type="slide" axis="1 0 0" stiffness="200"/>
<geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
</body>
</body>
</worldbody>
</mujoco>
""")
step_jit_fn = jax.jit(forward.step)
mx = mjx.device_put(m)
d = mujoco.MjData(m)
# give the system a little kick to ensure we have non-identity rotations
d.qvel = np.random.normal(m.nv) * 0.05
for i in range(100):
# in order to avoid re-jitting, reuse the same mj_data shape
qpos, qvel = d.qpos, d.qvel
d = mujoco.MjData(m)
d.qpos, d.qvel = qpos, qvel
dx = mjx.device_put(d)
mujoco.mj_step(m, d)
dx = step_jit_fn(mx, dx)
_assert_attr_eq(d, dx, 'qvel', i, 'test_rk4', atol=1e-2)
_assert_attr_eq(d, dx, 'qpos', i, 'test_rk4', atol=1e-2)
_assert_attr_eq(d, dx, 'act', i, 'test_rk4')
_assert_attr_eq(d, dx, 'time', i, 'test_rk4')
def test_disable_eulerdamp(self):
m = test_util.load_test_file('ant.xml')
mjx/mujoco/mjx/_src/passive_test.py
+3 -3
View File
@@ -26,7 +26,7 @@ from mujoco import mjx
import numpy as np
def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
def _assert_attr_eq(a, b, attr, step, fname, atol=1e-4, rtol=1e-4):
err_msg = f'mismatch: {attr} at step {step} in {fname}'
a, b = getattr(a, attr), getattr(b, attr)
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
@@ -34,7 +34,7 @@ def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
class PassiveTest(parameterized.TestCase):
@parameterized.parameters(enumerate(('ant.xml', 'mixed_joint_pendulum.xml')))
@parameterized.parameters(enumerate(('ant.xml', 'pendula.xml')))
def test_stiffness_damping(self, seed, fname):
"""Tests stiffness and damping on Ant."""
np.random.seed(seed)
@@ -60,7 +60,7 @@ class PassiveTest(parameterized.TestCase):
_assert_attr_eq(d, dx, 'qfrc_passive', i, fname)
@parameterized.parameters(
itertools.product(range(3), ('triple_pendulum.xml',))
itertools.product(range(3), ('pendula.xml',))
)
def test_fluid(self, seed, fname):
np.random.seed(seed)
mjx/mujoco/mjx/_src/scan.py
+92 -58
View File
@@ -162,7 +162,7 @@ def flat(
) -> Y:
r"""Scan a function across bodies or actuators.
Scan group data according to type and batch shape then calls vmap(f) on it.\
Scan group data according to type and batch shape then calls vmap(f) on it.
Args:
m: an mjx model
@@ -340,48 +340,88 @@ def body_tree(
IndexError: if function output shape does not match out_types shape
"""
_check_input(m, args, in_types)
depth_fn = lambda i, p=m.body_parentid: int(i > 0) and 1 + depth_fn(p[i])
typ_body_id = {
'j': m.jnt_bodyid,
'v': m.dof_bodyid,
'q': _q_bodyid(m),
}
key_parents = {}
# build up groupings of bodies and type ids using (level, (jnt_type,)) keys
key_typ_ids, key_body_ids = {}, {}
for body_id in np.arange(m.nbody, dtype=np.int32):
depth = depth_fn(body_id)
# group together bodies that will be processed together. grouping key:
# 1) the tree depth: parent bodies are processed first, so that they are
# available as carry input to child bodies (or reverse if reverse=True)
# 2) the types of arguments passed to f, both carry and *args:
# * for 'b' arguments, there is no extra grouping
# * for 'j' arguments, we group by joint type
# * for 'q' arguments, we group by q width
# * for 'v' arguments, we group by dof width
depths = np.zeros(m.nbody, dtype=np.int32)
# create grouping key
if any(t in 'jqv' for t in in_types + out_types):
jnts = np.nonzero(typ_body_id['j'] == body_id)[0]
jnts_p = np.nonzero(typ_body_id['j'] == m.body_parentid[body_id])[0]
key = depth, tuple(m.jnt_type[jnts])
parent_key = depth - 1, tuple(m.jnt_type[jnts_p])
else:
key, parent_key = (depth, ()), (depth - 1, ())
# map key => body id
key_body_ids = {}
for body_id in range(m.nbody):
parent_id = -1
if body_id > 0:
parent_id = m.body_parentid[body_id]
depths[body_id] = 1 + depths[parent_id]
# create grouping key: depth, carry, args
key = (depths[body_id],)
for i, t in enumerate(out_types + in_types):
id_ = parent_id if i < len(out_types) else body_id
if t == 'b':
continue
elif t == 'j':
key += (tuple(m.jnt_type[np.nonzero(m.jnt_bodyid == id_)[0]]))
elif t == 'v':
key += (len(np.nonzero(m.dof_bodyid == id_)[0]),)
elif t == 'q':
key += (len(np.nonzero(_q_bodyid(m) == id_)[0]),)
key_parents[key] = parent_key
body_ids = key_body_ids.get(key, np.array([], dtype=np.int32))
key_body_ids[key] = np.append(body_ids, body_id)
# add ids per type
for t in set(in_types + out_types):
out = key_typ_ids.setdefault(key, {})
id_ = body_id if t == 'b' else np.nonzero(typ_body_id[t] == body_id)[0]
id_ = np.expand_dims(id_, axis=0)
out[t] = np.concatenate((out[t], id_)) if t in out else id_
# find parent keys of each key. a key may have multiple parents if the
# carry output keys of distinct parents are the same. e.g.:
# - depth 0 body 1 (slide joint)
# -- depth 1 body 1 (hinge joint)
# - depth 0 body 2 (ball joint)
# -- depth 1 body 2 (hinge joint)
# given a scan with 'j' in the in_types, we would group depth 0 bodies
# separately but we may group depth 1 bodies together
key_parents = {}
key_typ_ids = list(sorted(key_typ_ids.items(), reverse=reverse))
for key, body_ids in key_body_ids.items():
body_ids = body_ids[body_ids != 0] # ignore worldbody, has no parent
if body_ids.size == 0:
continue
# find any key which has a body id that is a parent of these body_ids
pids = m.body_parentid[body_ids]
parents = {k for k, v in key_body_ids.items() if np.isin(v, pids).any()}
key_parents[key] = list(sorted(parents))
# key => take indices
key_in_take, key_y_take = {}, {}
for key, body_ids in key_body_ids.items():
for i, typ in enumerate(in_types + out_types):
if typ == 'b':
ids = body_ids
elif typ == 'j':
ids = np.stack([np.nonzero(m.jnt_bodyid == b)[0] for b in body_ids])
elif typ == 'v':
ids = np.stack([np.nonzero(m.dof_bodyid == b)[0] for b in body_ids])
elif typ == 'q':
ids = np.stack([np.nonzero(_q_bodyid(m) == b)[0] for b in body_ids])
else:
raise ValueError(f'Unknown in_type: {typ}')
if i < len(in_types):
key_in_take.setdefault(key, []).append(ids)
else:
key_y_take.setdefault(key, []).append(np.hstack(ids))
# use this grouping to take the right data subsets and call vmap(f)
keys = sorted(key_body_ids, reverse=reverse)
key_y = {}
for key, typ_ids in key_typ_ids:
for key in keys:
carry = None
if reverse:
child_keys = [k for k, v in key_parents.items() if v == key]
child_keys = [k for k, v in key_parents.items() if key in v]
for child_key in child_keys:
y = key_y[child_key]
@@ -394,39 +434,33 @@ def body_tree(
y = jax.tree_map(index_sum, y)
carry = y if carry is None else jax.tree_map(jp.add, carry, y)
else:
parent_key = key_parents[key]
y = key_y.get(parent_key)
elif key in key_parents:
ys = [key_y[p] for p in key_parents[key]]
y = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
body_ids = np.concatenate([key_body_ids[p] for p in key_parents[key]])
parent_ids = m.body_parentid[key_body_ids[key]]
take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
carry = jax.tree_map(take_fn, y)
if y is not None:
body_ids = key_body_ids[parent_key]
parent_ids = m.body_parentid[key_body_ids[key]]
take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
carry = jax.tree_map(take_fn, y)
f_args = [_take(arg, typ_ids[typ]) for arg, typ in zip(args, in_types)]
f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])]
key_y[key] = _nvmap(f, carry, *f_args)
# slice None results from the final output
key_typ_ids = [(k, v) for k, v in key_typ_ids if key_y[k] is not None]
keys = [k for k in keys if key_y[k] is not None]
# concatenate back to a single tree and drop the grouping dimension
ys = [key_y[key] for key, _ in key_typ_ids]
f_ret_is_seq = isinstance(ys[0], (list, tuple))
ys = ys if f_ret_is_seq else [[y] for y in ys]
ys = [
[v if typ == 'b' else jp.concatenate(v) for v, typ in zip(y, out_types)]
for y in ys
]
ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
# concatenate ys, drop grouping dimensions, put back in order
y = []
for i, typ in enumerate(out_types):
y_typ = [key_y[key] for key in keys]
if len(out_types) > 1:
y_typ = [y_[i] for y_ in y_typ]
if typ != 'b':
y_typ = jax.tree_map(jp.concatenate, y_typ)
y_typ = jax.tree_map(lambda *x: jp.concatenate(x), *y_typ)
y_take = np.argsort(np.concatenate([key_y_take[key][i] for key in keys]))
_check_output(y_typ, y_take, typ, i)
y.append(_take(y_typ, y_take))
# put concatenated results back into body order
reordered_ys = []
for i, (y, typ) in enumerate(zip(ys, out_types)):
ids = np.concatenate([np.hstack(v[typ]) for _, v in key_typ_ids])
take_ids = _index(ids, np.sort(ids))
_check_output(y, take_ids, typ, i)
reordered_ys.append(_take(y, take_ids))
y = reordered_ys if f_ret_is_seq else reordered_ys[0]
y = y[0] if len(out_types) == 1 else y
return y
mjx/mujoco/mjx/_src/smooth_test.py
+6 -6
View File
@@ -27,12 +27,12 @@ from mujoco.mjx._src.types import DisableBit
import numpy as np
def _assert_eq(a, b, name, step, fname, atol=1e-5, rtol=1e-5):
def _assert_eq(a, b, name, step, fname, atol=5e-4, rtol=5e-4):
err_msg = f'mismatch: {name} at step {step} in {fname}'
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5):
def _assert_attr_eq(a, b, attr, step, fname, atol=5e-4, rtol=5e-4):
err_msg = f'mismatch: {attr} at step {step} in {fname}'
a, b = getattr(a, attr), getattr(b, attr)
np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol)
@@ -43,7 +43,7 @@ class SmoothTest(parameterized.TestCase):
@parameterized.parameters(enumerate(test_util.TEST_FILES))
def test_smooth(self, seed, fname):
"""Tests mujoco mj smooth functions match mujoco_mjx smooth functions."""
if fname in ('convex.xml', 'weld.xml'):
if fname in ('convex.xml', 'equality.xml'):
return
np.random.seed(seed)
@@ -101,7 +101,7 @@ class SmoothTest(parameterized.TestCase):
# factor_m
dx = factor_m_fn(mx, dx, dx.qM)
_assert_attr_eq(d, dx, 'qLD', i, fname, atol=1e-3)
_assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3, rtol=1e-4)
_assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3)
# com_vel
dx = com_vel_jit_fn(mx, dx)
@@ -110,14 +110,14 @@ class SmoothTest(parameterized.TestCase):
# rne
dx = rne_jit_fn(mx, dx)
_assert_attr_eq(d, dx, 'qfrc_bias', i, fname, atol=1e-4)
_assert_attr_eq(d, dx, 'qfrc_bias', i, fname)
# mul_m (auxilliary function, not part of smooth step)
vec = np.random.random(m.nv)
mjx_vec = mul_m_jit_fn(mx, dx, jp.array(vec))
mj_vec = np.zeros(m.nv)
mujoco.mj_mulM(m, d, mj_vec, vec)
_assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname, atol=1e-4)
_assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname)
# transmission
dx = transmission_jit_fn(mx, dx)
mjx/mujoco/mjx/_src/solver.py
+56 -37
View File
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""CG and Newton solvers."""
"""Constraint solvers."""
from typing import Optional
@@ -26,11 +26,12 @@ from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import DisableBit
from mujoco.mjx._src.types import Model
from mujoco.mjx._src.types import SolverType
# pylint: enable=g-importing-member
class _CGContext(PyTreeNode):
"""Data updated during each cg solver iteration.
class _Context(PyTreeNode):
"""Data updated during each solver iteration.
Attributes:
qacc: acceleration (from Data) (nv,)
@@ -44,7 +45,7 @@ class _CGContext(PyTreeNode):
search: linesearch vector (nv,)
gauss: gauss Cost
cost: constraint + Gauss cost
prev_cost: cost from previous cg iter
prev_cost: cost from previous iter
solver_niter: number of solver iterations
"""
@@ -63,17 +64,17 @@ class _CGContext(PyTreeNode):
solver_niter: jax.Array
@classmethod
def create(cls, m: Model, d: Data, grad: bool = True) -> '_CGContext':
def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context':
jaref = d.efc_J @ d.qacc - d.efc_aref
# TODO(robotics-team): determine nv at which sparse mul is faster
M = smooth.dense_m(m, d) if m.nv < 100 else None # pylint: disable=invalid-name
ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
nv_0 = jp.zeros((m.nv,))
ctx = _CGContext(
ctx = _Context(
qacc=d.qacc,
qfrc_constraint=d.qfrc_constraint,
Jaref=jaref,
efc_force=-jaref * d.efc_D,
efc_force=jp.zeros(d.nefc),
M=M,
Ma=ma,
grad=nv_0,
@@ -84,9 +85,9 @@ class _CGContext(PyTreeNode):
prev_cost=0.0,
solver_niter=0,
)
ctx = _cg_update_constraint(m, d, ctx)
ctx = _update_constraint(m, d, ctx)
if grad:
ctx = _cg_update_gradient(m, d, ctx)
ctx = _update_gradient(m, d, ctx)
ctx = ctx.replace(search=-ctx.Mgrad) # start with preconditioned gradient
return ctx
@@ -110,7 +111,8 @@ class _LSPoint(PyTreeNode):
@classmethod
def create(
cls,
ctx: _CGContext,
d: Data,
ctx: _Context,
alpha: jax.Array,
jv: jax.Array,
quad: jax.Array,
@@ -119,8 +121,8 @@ class _LSPoint(PyTreeNode):
"""Creates a linesearch point with first and second derivatives."""
# roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c
# TODO(robotics-team): change this to support equality, friction constraints
active = (ctx.Jaref + alpha * jv) < 0
# TODO(robotics-team): change this to support friction constraints
active = ((ctx.Jaref + alpha * jv) < 0).at[:d.ne + d.nf].set(True)
quad = jax.vmap(jp.multiply)(quad, active) # only active
quad_total = quad_gauss + jp.sum(quad, axis=0)
@@ -131,7 +133,7 @@ class _LSPoint(PyTreeNode):
class _LSContext(PyTreeNode):
"""Data updated during each cg line search iteration.
"""Data updated during each line search iteration.
Attributes:
lo: low point bounding the line search interval
@@ -162,29 +164,30 @@ 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 _cg_update_constraint(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
"""Updates constraint force and resulting cost given latst CG iteration.
def _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
"""Updates constraint force and resulting cost given latst solver iteration.
Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
Args:
m: model defining constraints
d: data which contains latest qacc and smooth terms
ctx: current CG context
ctx: current solver context
Returns:
context with new constraint force and costs
"""
del m
# TODO(robotics-team): add equality, friction constraints
# also consider moving to _constraint.py to match mujoco layout
# TODO(robotics-team): add friction constraints
jaref = ctx.Jaref * (ctx.Jaref < 0) # non-negative constraints
efc_force = -jaref * d.efc_D
# only count active constraints
active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
efc_force = d.efc_D * -ctx.Jaref * active
qfrc_constraint = d.efc_J.T @ efc_force
gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
cost = 0.5 * jp.sum(jaref * jaref * d.efc_D) + gauss
cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active) + gauss
ctx = ctx.replace(
qfrc_constraint=qfrc_constraint,
@@ -197,22 +200,34 @@ def _cg_update_constraint(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
return ctx
def _cg_update_gradient(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
"""Updates grad and M / grad given latest CG iteration.
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
Args:
m: model defining constraints
d: data which contains latest smooth terms
ctx: current CG contet
ctx: current solver context
Returns:
context with new grad and M / grad
Raises:
NotImplementedError: for unsupported solver type
"""
grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint
mgrad = smooth.solve_m(m, d, grad)
if m.opt.solver == SolverType.CG:
mgrad = smooth.solve_m(m, d, grad)
elif m.opt.solver == SolverType.NEWTON:
active = (ctx.Jaref < 0).at[:d.ne + d.nf].set(True)
h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
h = smooth.dense_m(m, d) + h
h_ = jax.scipy.linalg.cho_factor(h)
mgrad = jax.scipy.linalg.cho_solve(h_, grad)
else:
raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
ctx = ctx.replace(grad=grad, Mgrad=mgrad)
@@ -223,13 +238,13 @@ def _rescale(m: Model, value: jax.Array) -> jax.Array:
return value / (m.stat.meaninertia * max(1, m.nv))
def _cg_search(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
"""Performs a zoom linesearch to find optimal search step size.
Args:
m: model defining search options and other needed terms
d: data with inertia matrix and other needed terms
ctx: current CG context
ctx: current solver context
Returns:
updated context with new qacc, Ma, Jaref
@@ -250,7 +265,7 @@ def _cg_search(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv))
quad = (quad * d.efc_D).T
point_fn = lambda alpha: _LSPoint.create(ctx, alpha, jv, quad, quad_gauss)
point_fn = lambda alpha: _LSPoint.create(d, ctx, alpha, jv, quad, quad_gauss)
def cond(ctx: _LSContext) -> jax.Array:
done = ctx.ls_iter >= m.opt.ls_iterations
@@ -308,10 +323,10 @@ def _cg_search(m: Model, d: Data, ctx: _CGContext) -> _CGContext:
return ctx
def cg_solve(m: Model, d: Data) -> Data:
def solve(m: Model, d: Data) -> Data:
"""Finds forces that satisfy constraints using conjugate gradient descent."""
def cond(ctx: _CGContext) -> jax.Array:
def cond(ctx: _Context) -> jax.Array:
improvement = _rescale(m, ctx.prev_cost - ctx.cost)
gradient = _rescale(m, math.norm(ctx.grad))
@@ -321,11 +336,11 @@ def cg_solve(m: Model, d: Data) -> Data:
return ~done
def body(ctx: _CGContext) -> _CGContext:
ctx = _cg_search(m, d, ctx)
def body(ctx: _Context) -> _Context:
ctx = _linesearch(m, d, ctx)
prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad # pylint: disable=invalid-name
ctx = _cg_update_constraint(m, d, ctx)
ctx = _cg_update_gradient(m, d, ctx)
ctx = _update_constraint(m, d, ctx)
ctx = _update_gradient(m, d, ctx)
# polak-ribiere:
beta = jp.dot(ctx.grad, ctx.Mgrad - prev_Mgrad)
@@ -339,12 +354,16 @@ def cg_solve(m: Model, d: Data) -> Data:
# warmstart:
qacc = d.qacc_smooth
if not m.opt.disableflags & DisableBit.WARMSTART:
warm = _CGContext.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
smth = _CGContext.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 = jax.lax.while_loop(cond, body, _CGContext.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,
mjx/mujoco/mjx/_src/test_util.py
+2 -8
View File
@@ -24,16 +24,10 @@ import numpy as np
TEST_FILES: List[str] = [
'ant.xml',
'ball_pendulum.xml',
'cherry_pendulum.xml',
'convex.xml',
'equality.xml',
'humanoid.xml',
'mixed_joint_pendulum.xml',
'single_pendulum.xml',
'slide_pendulum.xml',
'triple_pendulum.xml',
'triple_pendulum_free.xml',
'weld.xml',
'pendula.xml',
]
_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
mjx/mujoco/mjx/_src/types.py
+13 -11
View File
@@ -131,8 +131,9 @@ class SolverType(enum.IntEnum):
Attributes:
CG: Conjugate gradient (primal)
"""
# unsupported: PGS, NEWTON
# unsupported: PGS
CG = mujoco.mjtSolver.mjSOL_CG
NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
class EqType(enum.IntEnum):
@@ -145,7 +146,8 @@ class EqType(enum.IntEnum):
"""
CONNECT = mujoco.mjtEq.mjEQ_CONNECT
WELD = mujoco.mjtEq.mjEQ_WELD
# unsupported: JOINT, TENDON, DISTANCE
JOINT = mujoco.mjtEq.mjEQ_JOINT
# unsupported: TENDON, DISTANCE
class TrnType(enum.IntEnum):
@@ -380,7 +382,7 @@ class Model(PyTreeNode):
nexclude: int
neq: int
nnumeric: int
nM: int
nM: int # pylint:disable=invalid-name
opt: Option
stat: Statistic
qpos0: jax.Array
@@ -419,14 +421,14 @@ class Model(PyTreeNode):
dof_bodyid: np.ndarray
dof_jntid: np.ndarray
dof_parentid: np.ndarray
dof_Madr: np.ndarray
dof_Madr: np.ndarray # pylint:disable=invalid-name
dof_solref: jax.Array
dof_solimp: jax.Array
dof_frictionloss: jax.Array
dof_armature: jax.Array
dof_damping: jax.Array
dof_invweight0: jax.Array
dof_M0: jax.Array
dof_M0: jax.Array # pylint:disable=invalid-name
geom_type: np.ndarray
geom_contype: np.ndarray
geom_conaffinity: np.ndarray
@@ -635,14 +637,14 @@ class Data(PyTreeNode):
crb: jax.Array
actuator_length: jax.Array
actuator_moment: jax.Array
qM: jax.Array
qLD: jax.Array
qLDiagInv: jax.Array
qLDiagSqrtInv: jax.Array
qM: jax.Array # pylint:disable=invalid-name
qLD: jax.Array # pylint:disable=invalid-name
qLDiagInv: jax.Array # pylint:disable=invalid-name
qLDiagSqrtInv: jax.Array # pylint:disable=invalid-name
contact: Contact
efc_J: jax.Array
efc_J: jax.Array # pylint:disable=invalid-name
efc_frictionloss: jax.Array
efc_D: jax.Array
efc_D: jax.Array # pylint:disable=invalid-name
# position, velocity dependent:
actuator_velocity: jax.Array
cvel: jax.Array
mjx/mujoco/mjx/benchmark/benchmark.py
+26 -10
View File
@@ -33,21 +33,34 @@ _PATHS = {
'shadow_hand': 'benchmark/model/shadow_hand/scene_right.xml',
}
_BATCH_SIZE = {
('humanoid', 'TPU v5 lite'): 1024,
('barkour', 'TPU v5 lite'): 1024,
('shadow_hand', 'TPU v5 lite'): 1024,
('humanoid', 'Tesla V100-SXM2-16GB'): 8192,
('barkour', 'Tesla V100-SXM2-16GB'): 8192,
('shadow_hand', 'Tesla V100-SXM2-16GB'): 4096,
('humanoid', 'cpu'): 64,
('barkour', 'tpu_v5e'): 1024,
('humanoid', 'tpu_v5e'): 1024,
('shadow_hand', 'tpu_v5e'): 1024,
('barkour', 'gpu_a100'): 8192,
('humanoid', 'gpu_a100'): 8192,
('shadow_hand', 'gpu_a100'): 4096,
('barkour', 'cpu'): 64,
('humanoid', 'cpu'): 64,
('shadow_hand', 'cpu'): 64,
}
_SOLVER_CONFIG = {
('barkour', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 4, 6),
('humanoid', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 6, 6),
('shadow_hand', 'tpu_v5e'): (mujoco.mjtSolver.mjSOL_CG, 8, 6),
('humanoid', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
('barkour', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
('shadow_hand', 'gpu_a100'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
('barkour', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
('humanoid', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
('shadow_hand', 'cpu'): (mujoco.mjtSolver.mjSOL_NEWTON, 1, 4),
}
flags.DEFINE_string('model', 'humanoid', 'Model to benchmark')
flags.DEFINE_string('device', 'cpu', 'Device benchmark is running on')
flags.DEFINE_enum('device', 'cpu', ('cpu', 'tpu_v5e', 'gpu_a100'),
'Device benchmark is running on')
def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float:
@@ -95,6 +108,9 @@ def _run(state: benchmark.State):
f = epath.resource_path('mujoco.mjx') / _PATHS[FLAGS.model]
m = mujoco.MjModel.from_xml_path(f.as_posix())
m.opt.solver, m.opt.iterations, m.opt.ls_iterations = _SOLVER_CONFIG[
(FLAGS.model, FLAGS.device)
]
m = mjx.device_put(m)
def init(rng):
@@ -106,7 +122,7 @@ def _run(state: benchmark.State):
def step(d):
return mjx.step(m, d)
batch_size = _BATCH_SIZE[(FLAGS.model, jax.devices()[0].device_kind)]
batch_size = _BATCH_SIZE[(FLAGS.model, FLAGS.device)]
_measure_fn(state, init, step, batch_size=batch_size)
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml
+1 -1
View File
@@ -2,7 +2,7 @@
<compiler angle="radian" meshdir="." texturedir="assets" autolimits="true"/>
<option timestep="0.002" iterations="4" ls_iterations="6" solver="CG">
<option timestep="0.002" iterations="1" ls_iterations="4">
<flag eulerdamp="disable"/>
</option>
mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.xml
+3 -3
View File
@@ -14,7 +14,7 @@
-->
<mujoco model="Humanoid">
<option timestep="0.005" solver="CG" iterations="6" ls_iterations="6">
<option timestep="0.005" iterations="1" ls_iterations="4">
<flag eulerdamp="disable"/>
</option>
@@ -195,7 +195,7 @@
<pair geom1="foot2_right" geom2="floor"/>
</contact>
<tendon>
<!-- <tendon>
<fixed name="hamstring_right" limited="true" range="-0.3 2">
<joint joint="hip_y_right" coef=".5"/>
<joint joint="knee_right" coef="-.5"/>
@@ -204,7 +204,7 @@
<joint joint="hip_y_left" coef=".5"/>
<joint joint="knee_left" coef="-.5"/>
</fixed>
</tendon>
</tendon> -->
<actuator>
<motor name="abdomen_y" gear="40" joint="abdomen_y"/>
mjx/mujoco/mjx/benchmark/model/shadow_hand/right_hand.xml
+1 -1
View File
@@ -1,7 +1,7 @@
<mujoco model="right_shadow_hand">
<compiler angle="radian" meshdir="assets" autolimits="true"/>
<option impratio="10" solver="CG" iterations="8" ls_iterations="6">
<option impratio="10" iterations="1" ls_iterations="4">
<flag eulerdamp="disable"/>
</option>
mjx/mujoco/mjx/test_data/ball_pendulum.xml
-23
View File
@@ -1,23 +0,0 @@
<mujoco model="ball_pendulum">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<compiler autolimits="true"/>
<default>
<geom contype="0" conaffinity="0"/>
<joint damping="10"/>
</default>
<option solver="CG"/>
<worldbody>
<body>
<joint axis="1 0 0" type="ball" range="0 10"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0.3 0.4 0.5">
<joint axis="1 0 0" type="hinge" range="-20 20"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge" range="-30 30"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/cherry_pendulum.xml
-20
View File
@@ -1,20 +0,0 @@
<mujoco model="cherry_pendulum">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<default>
<geom contype="0" conaffinity="0"/>
</default>
<worldbody>
<body>
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
</body>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/weld.xml → mjx/mujoco/mjx/test_data/equality.xml
+17 -4
View File
@@ -47,12 +47,25 @@
<freejoint/>
<geom class="free"/>
</body>
<body name="box5" pos="4 0 0">
<geom class="free"/>
<joint name="joint1" axis="1 0 0" type="hinge" />
</body>
<body name="box6" pos="4 0 0">
<geom class="free"/>
<joint name="joint2" axis="1 0 0" type="hinge" />
</body>
</worldbody>
<equality>
<connect name="connect anchor" body1="box1" body2="beam1" anchor="0 0 -1" />
<weld name="weld anchor weak torques" body1="box2" body2="beam2" torquescale="0.002" anchor="0 -2 0"/>
<weld name="weld relpose" body1="box3" body2="beam3" relpose="0 0 0 1 -.3 0 0"/>
<weld name="weld relpose+anchor" body1="box4" body2="beam4" relpose="0 0 0 1 -.3 0 0" anchor="0 0 -1"/>
<connect name="connect anchor" body1="box1" body2="beam1" anchor="0 0 -1" />
<weld name="weld anchor weak torques" body1="box2" body2="beam2" torquescale="0.002" anchor="0 -2 0"/>
<weld name="weld relpose" body1="box3" body2="beam3" relpose="0 0 0 1 -.3 0 0"/>
<weld name="weld relpose+anchor" body1="box4" body2="beam4" relpose="0 0 0 1 -.3 0 0" anchor="0 0 -1"/>
<joint name="joint" joint1="joint1" joint2="joint2" polycoef="0 -1 0.1 0.15 0.2" />
</equality>
</mujoco>
mjx/mujoco/mjx/test_data/humanoid.xml
+1 -1
View File
@@ -5,7 +5,7 @@
<geom conaffinity="0" condim="3" contype="0" material="geom"/>
<motor ctrllimited="true" ctrlrange="-.4 .4"/>
</default>
<option iterations="8" timestep="0.003" solver="CG"/>
<option iterations="8" timestep="0.003"/>
<size nkey="5" nuser_geom="1"/>
<visual>
<map fogend="5" fogstart="3"/>
mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml
-23
View File
@@ -1,23 +0,0 @@
<mujoco model="revolute">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<compiler autolimits="true"/>
<default>
<geom contype="0" conaffinity="0"/>
<joint damping="20"/>
</default>
<worldbody>
<body>
<joint axis="1 0 0" type="ball" range="0 10"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0.3 0.4 0.5">
<joint axis="1 0 0" type="hinge" range="-20 20"/>
<joint axis="0 1 0" type="hinge" range="-20 20"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge" range="-30 30"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/pendula.xml
+102
View File
@@ -0,0 +1,102 @@
<!-- For validating dynamics of joints:
* free, ball, slide, hinge joints
* stacked joints (e.g. hinge + slide, ball + slide, etc)
* n-link kinematic chains
* limits, armature, damping
-->
<mujoco model="pendula">
<compiler autolimits="true"/>
<option timestep="0.02">
<flag contact="disable" />
</option>
<default>
<geom type="box" pos=".1 .2 .3" size=".1 .2 .3"/>
<joint damping="0.25" stiffness="0.1"/>
</default>
<worldbody>
<!-- a single free body -->
<body pos="0 0 0">
<freejoint/>
<geom/>
</body>
<!-- a single ball joint with a limit -->
<body pos="0.5 0 0">
<joint type="ball" range="0 35"/>
<geom/>
</body>
<!-- a single slide joint with a limit -->
<body pos="1.0 0 0">
<joint type="slide" axis="0.1 0.2 0.3" range="-1 1"/>
<geom/>
</body>
<!-- a single hinge joint with a limit -->
<body pos="1.5 0 0">
<joint type="hinge" axis="0.1 0.2 0.3" range="-35 50"/>
<geom/>
</body>
<!-- stacked joint: hinge + slide -->
<body pos="2.0 0 0">
<joint type="hinge" axis="0.1 0.2 0.3"/>
<joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
<geom/>
</body>
<!-- stacked joint: slide + ball -->
<body pos="2.5 0 0">
<joint type="slide" axis="0.4 0.5 0.6" range="0 1"/>
<joint type="ball"/>
<geom/>
</body>
<!-- triple pendulum of hinges -->
<body pos="3.0 0 0">
<joint axis="0.1 0.2 0.3" type="hinge"/>
<geom/>
<body pos="0 0 -0.8">
<joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
<geom/>
<body pos="0 -0.7 0">
<joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
<geom/>
</body>
</body>
</body>
<!-- cherry pendulum: two bodies attached to same parent body -->
<body pos="3.5 0 0">
<joint type="ball" damping="0.5" />
<geom/>
<body pos="0 0 -0.8">
<joint axis="0.4 0.5 0.6" type="hinge" armature="0.02" range="-20 20"/>
<geom/>
</body>
<body pos="0 -0.7 0">
<joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
<geom/>
</body>
</body>
<!-- falling pendulum -->
<body pos="4.0 0 0">
<freejoint/>
<geom/>
<body pos="0 0 -0.8">
<joint axis="0.4 0.5 0.6" type="slide" armature="0.02" range="-0.4 0.6"/>
<geom/>
<body pos="0 -0.7 0">
<joint axis="0.7 0.8 0.9" type="hinge" damping="0.75" range="-30 30"/>
<geom/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/single_pendulum.xml
-13
View File
@@ -1,13 +0,0 @@
<mujoco model="pendulum">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<worldbody>
<body>
<joint name="slider" axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
</body>
</worldbody>
<!-- Tests that a single actuator doesn't get mangled in a physics step. -->
<actuator>
<motor name="slide" joint="slider" gear="10" ctrllimited="true" ctrlrange="-1 1"/>
</actuator>
</mujoco>
mjx/mujoco/mjx/test_data/slide_pendulum.xml
-20
View File
@@ -1,20 +0,0 @@
<mujoco model="slide_pendulum">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<default>
<geom contype="0" conaffinity="0"/>
</default>
<worldbody>
<body>
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="slide"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/triple_pendulum.xml
-21
View File
@@ -1,21 +0,0 @@
<mujoco model="triple_pendulum">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<compiler autolimits="true"/>
<default>
<geom contype="0" conaffinity="0"/>
</default>
<worldbody>
<body>
<joint axis="1 0 0" type="hinge" armature="0.01" range="-10 10"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge" armature="0.02" range="-20 20"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge" armature="0.03" range="-30 30"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/test_data/triple_pendulum_free.xml
-24
View File
@@ -1,24 +0,0 @@
<mujoco model="triple_pendulum_free">
<option timestep="0.02" solver="CG" iterations="6" ls_iterations="6"/>
<default>
<geom contype="0" conaffinity="0"/>
</default>
<worldbody>
<body>
<joint type="free"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
</body>
<body>
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="2" type="sphere"/>
<body pos="0 0.5 0">
<joint axis="1 0 0" type="hinge"/>
<geom pos="0 0.5 0" size=".15" mass="3" type="sphere"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/viewer.py
+1 -5
View File
@@ -14,8 +14,8 @@
# ==============================================================================
"""An example integration of MJX with the MuJoCo viewer."""
from collections.abc import Sequence
import time
from typing import Sequence
from absl import app
from absl import flags
@@ -38,10 +38,6 @@ def main(argv: Sequence[str]) -> None:
m = mujoco.MjModel.from_xml_path(_MODEL_PATH.value)
d = mujoco.MjData(m)
# Override the solver option to CG since that is currently the only one
# supported by MJX.
m.opt.solver = mujoco.mjtSolver.mjSOL_CG
mx = mjx.device_put(m)
dx = mjx.make_data(mx)
mjx/pyproject.toml
+1 -1
View File
@@ -31,7 +31,7 @@ dependencies = [
"etils[epath]",
"jax",
"jaxlib",
"mujoco>=3.0.1",
"mujoco>=3.0.1.dev0",
"scipy",
"trimesh",
]
mjx/requirements.txt
+2 -2
View File
@@ -33,8 +33,8 @@ jaxlib==0.4.18; python_version >= '3.9' \
--hash=sha256:6cb20bbbdafd90e71ad0deb9295519a0175c108c8c557b84fb9fe94f751daee4 \
--hash=sha256:116a0d6aedd3e856b52493d7e392fb1b40952b84fb72448fde1c1ab5687db667 \
--hash=sha256:9593ff69f424947567e206f3e356b2a2df55ca68e6d815d5adc6cae308e8f652
pip==23.2.1 \
--hash=sha256:7ccf472345f20d35bdc9d1841ff5f313260c2c33fe417f48c30ac46cccabf5be
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
pytest==7.4.2 \
--hash=sha256:1d881c6124e08ff0a1bb75ba3ec0bfd8b5354a01c194ddd5a0a870a48d99b002
pytest-xdist==3.3.1 \
mjx/tutorial.ipynb
+440 -248
View File
@@ -187,11 +187,9 @@
},
"source": [
"# Training a Policy with MJX\n",
"MJX is an implementation of MuJoCo written in [jax](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate training RL policies with MJX.\n",
"MJX is an implementation of MuJoCo written in [JAX](https://jax.readthedocs.io/en/latest/index.html), enabling large batch training on GPU/TPU. In this notebook, we demonstrate how to train RL policies with MJX.\n",
"\n",
"First, we implement an environment `State` so that we can plug into the [brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and etc. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n",
"\n",
"`MjxEnv` is an implementation of the `brax.envs.base.Env` class that initializes `mjx.Model` and `mjx.Data` objects. Inheriting from `brax.envs.base.Env` allows us to use the training agents implemented in brax. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`."
"First, we implement an environment `State` so that we can plug into the [Brax](https://github.com/google/brax) environment API. `State` holds the observation, reward, metrics, and environment info. Notably `State.pipeline_state` holds a `mjx.Data` object, which is analogous to `mjData` in MuJoCo.\n"
]
},
{
@@ -202,7 +200,7 @@
},
"outputs": [],
"source": [
"#@title State and MjxEnv\n",
"#@title State\n",
"\n",
"@struct.dataclass\n",
"class State(Base):\n",
@@ -223,8 +221,30 @@
" reward: jax.Array\n",
" done: jax.Array\n",
" metrics: Dict[str, jax.Array] = struct.field(default_factory=dict)\n",
" info: Dict[str, Any] = struct.field(default_factory=dict)\n",
" info: Dict[str, Any] = struct.field(default_factory=dict)\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "acpXtDLNXLV9"
},
"source": [
"\n",
"Next, we implement `MjxEnv`, an environment class we'll use through the notebook. `MjxEnv` initializes a `mjx.Model` and `mjx.Data` object. Notice that `MjxEnv` calls `mjx.step` for every `pipeline_step`, which is analgous to `mujoco.mj_step`.\n",
"\n",
"`MjxEnv` also inherits from `brax.envs.base.Env` which allows us to use the training agents implemented in brax."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ccujYeJ5XOhx"
},
"outputs": [],
"source": [
"#@title MjxEnv\n",
"\n",
"class MjxEnv(Env):\n",
" \"\"\"API for driving an MJX system for training and inference in brax.\"\"\"\n",
@@ -293,7 +313,7 @@
" \"\"\"Returns 6d spatial transform and 6d velocity for all bodies.\"\"\"\n",
" x = Transform(pos=data.xpos[1:, :], rot=data.xquat[1:, :])\n",
" cvel = Motion(vel=data.cvel[1:, 3:], ang=data.cvel[1:, :3])\n",
" offset = data.xipos[1:, :] - data.subtree_com[\n",
" offset = data.xpos[1:, :] - data.subtree_com[\n",
" self.model.body_rootid[np.arange(1, self.model.nbody)]]\n",
" xd = Transform.create(pos=offset).vmap().do(cvel)\n",
" return x, xd\n"
@@ -305,7 +325,7 @@
"id": "iPlFu4CiIgBN"
},
"source": [
"Now we can define environment implementations for Humanoid and the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0). The environments define the reward and stepping logic. `reset` initializes a `State`, and `step` steps through the physics step and reward logic."
"Finally we can implement a real environment. We choose to first implement the Humanoid environment. Notice that `reset` initializes a `State`, and `step` steps through the physics step and reward logic. The reward and stepping logic train the Humanoid to run forwards."
]
},
{
@@ -336,6 +356,9 @@
" )\n",
" mj_model = mujoco.MjModel.from_xml_path(\n",
" (path / 'humanoid.xml').as_posix())\n",
" mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
" mj_model.opt.iterations = 6\n",
" mj_model.opt.ls_iterations = 6\n",
"\n",
" physics_steps_per_control_step = 5\n",
" kwargs['physics_steps_per_control_step'] = kwargs.get(\n",
@@ -444,11 +467,376 @@
"envs.register_environment('humanoid', Humanoid)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "P1K6IznI2y83"
},
"source": [
"## Visualize a Rollout\n",
"\n",
"Let's instantiate the environment and visualize a short rollout.\n",
"\n",
"NOTE: Since episodes terminates early if the torso is below the healthy z-range, the only relevant contacts for this task are between the feet and the plane. We turn off other contacts."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "EhKLFK54C1CH"
},
"outputs": [],
"source": [
"# instantiate the environment\n",
"env_name = 'humanoid'\n",
"env = envs.get_environment(env_name)\n",
"\n",
"# define the jit reset/step functions\n",
"jit_reset = jax.jit(env.reset)\n",
"jit_step = jax.jit(env.step)\n",
"\n",
"# instantiate the renderer\n",
"renderer = mujoco.Renderer(env.model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9f2ME2WbA5Ip"
},
"outputs": [],
"source": [
"#@title Define a render utility function\n",
"\n",
"def get_image(state: State, camera: str) -\u003e np.ndarray:\n",
" \"\"\"Renders the environment state.\"\"\"\n",
" d = mujoco.MjData(env.model)\n",
" # write the mjx.Data into an mjData object\n",
" mjx.device_get_into(d, state.pipeline_state)\n",
" mujoco.mj_forward(env.model, d)\n",
" # use the mjData object to update the renderer\n",
" renderer.update_scene(d, camera=camera)\n",
" return renderer.render()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ph8u-v2Q2xLS"
},
"outputs": [],
"source": [
"# initialize the state\n",
"state = jit_reset(jax.random.PRNGKey(0))\n",
"rollout = [state]\n",
"images = [get_image(state, camera='side')]\n",
"\n",
"# grab a trajectory\n",
"for i in range(10):\n",
" ctrl = -0.1 * jp.ones(env.sys.nu)\n",
" state = jit_step(state, ctrl)\n",
" rollout.append(state)\n",
" images.append(get_image(state, camera='side'))\n",
"\n",
"media.show_video(images, fps=1.0 / env.dt)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BQDG6NQ1CbZD"
},
"source": [
"## Train Humanoid Policy\n",
"\n",
"Let's finally train a policy with PPO to make the Humanoid run forwards. Training takes about 13-14 minutes on a Tesla V100 GPU."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "xLiddQYPApBw"
},
"outputs": [],
"source": [
"train_fn = functools.partial(\n",
" ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
" episode_length=1000, normalize_observations=True, action_repeat=1,\n",
" unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
" discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
" batch_size=1024, seed=0)\n",
"\n",
"\n",
"x_data = []\n",
"y_data = []\n",
"ydataerr = []\n",
"times = [datetime.now()]\n",
"\n",
"max_y, min_y = 13000, 0\n",
"def progress(num_steps, metrics):\n",
" times.append(datetime.now())\n",
" x_data.append(num_steps)\n",
" y_data.append(metrics['eval/episode_reward'])\n",
" ydataerr.append(metrics['eval/episode_reward_std'])\n",
"\n",
" plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
" plt.ylim([min_y, max_y])\n",
"\n",
" plt.xlabel('# environment steps')\n",
" plt.ylabel('reward per episode')\n",
" plt.title(f'y={y_data[-1]:.3f}')\n",
"\n",
" plt.errorbar(\n",
" x_data, y_data, yerr=ydataerr)\n",
" plt.show()\n",
"\n",
"make_inference_fn, params, _= train_fn(environment=env, progress_fn=progress)\n",
"\n",
"print(f'time to jit: {times[1] - times[0]}')\n",
"print(f'time to train: {times[-1] - times[1]}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "YYIch0HEApBx"
},
"source": [
"## Save and Load Policy\n",
"\n",
"We can save and load the policy using the brax model API."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Z8gI6qH6ApBx"
},
"outputs": [],
"source": [
"#@title Save Model\n",
"model_path = '/tmp/mjx_brax_policy'\n",
"model.save_params(model_path, params)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "h4reaWgxApBx"
},
"outputs": [],
"source": [
"#@title Load Model and Define Inference Function\n",
"params = model.load_params(model_path)\n",
"\n",
"inference_fn = make_inference_fn(params)\n",
"jit_inference_fn = jax.jit(inference_fn)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0G357XIfApBy"
},
"source": [
"## Visualize Policy\n",
"\n",
"Finally we can visualize the policy."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "osYasMw4ApBy"
},
"outputs": [],
"source": [
"eval_env = envs.get_environment(env_name)\n",
"\n",
"jit_reset = jax.jit(eval_env.reset)\n",
"jit_step = jax.jit(eval_env.step)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "d-UhypudApBy"
},
"outputs": [],
"source": [
"# initialize the state\n",
"rng = jax.random.PRNGKey(0)\n",
"state = jit_reset(rng)\n",
"rollout = [state]\n",
"images = [get_image(state, camera='side')]\n",
"\n",
"# grab a trajectory\n",
"n_steps = 500\n",
"render_every = 2\n",
"\n",
"for i in range(n_steps):\n",
" act_rng, rng = jax.random.split(rng)\n",
" ctrl, _ = jit_inference_fn(state.obs, act_rng)\n",
" state = jit_step(state, ctrl)\n",
" rollout.append(state)\n",
" if i % render_every == 0:\n",
" images.append(get_image(state, camera='side'))\n",
"\n",
" if state.done:\n",
" break\n",
"\n",
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "zR-heox6LARK"
},
"source": [
"# MJX Policy in MuJoCo\n",
"\n",
"Note that we can also perform the physics step using the original MuJoCo python bindings to show that the policy trained in MJX works in MuJoCo."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "w6ixFi4dApBy"
},
"outputs": [],
"source": [
"mj_model = eval_env.model\n",
"mj_data = mujoco.MjData(mj_model)\n",
"\n",
"renderer = mujoco.Renderer(mj_model)\n",
"ctrl = jp.zeros(mj_model.nu)\n",
"\n",
"images = []\n",
"for i in range(n_steps):\n",
" act_rng, rng = jax.random.split(rng)\n",
"\n",
" obs = eval_env._get_obs(mjx.device_put(mj_data), ctrl)\n",
" ctrl, _ = jit_inference_fn(obs, act_rng)\n",
"\n",
" mj_data.ctrl = ctrl\n",
" for _ in range(eval_env._physics_steps_per_control_step):\n",
" mujoco.mj_step(mj_model, mj_data) # Physics step using MuJoCo mj_step.\n",
"\n",
" if i % render_every == 0:\n",
" renderer.update_scene(mj_data, camera='side')\n",
" images.append(renderer.render())\n",
"\n",
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "65mIPj6DQNNa"
},
"source": [
"# Domain Randomization\n",
"\n",
"We might also want to include randomization over certain `mjModel` parameters while training a policy. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "h8mhzKjHQuoL"
},
"outputs": [],
"source": [
"def domain_randomize(sys, rng):\n",
" \"\"\"Randomizes the mjx.Model.\"\"\"\n",
" @jax.vmap\n",
" def rand(rng):\n",
" _, key = jax.random.split(rng, 2)\n",
" # friction\n",
" friction = jax.random.uniform(key, (1,), minval=0.6, maxval=1.4)\n",
" friction = sys.geom_friction.at[:, 0].set(friction)\n",
" # actuator\n",
" _, key = jax.random.split(key, 2)\n",
" gain_range = (-10, -5)\n",
" param = jax.random.uniform(\n",
" key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
" ) + sys.actuator_gainprm[:, 0]\n",
" gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
" bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
" return friction, gain, bias\n",
"\n",
" friction, gain, bias = rand(rng)\n",
"\n",
" in_axes = jax.tree_map(lambda x: None, sys)\n",
" in_axes = in_axes.tree_replace({\n",
" 'geom_friction': 0,\n",
" 'actuator_gainprm': 0,\n",
" 'actuator_biasprm': 0,\n",
" })\n",
"\n",
" sys = sys.tree_replace({\n",
" 'geom_friction': friction,\n",
" 'actuator_gainprm': gain,\n",
" 'actuator_biasprm': bias,\n",
" })\n",
"\n",
" return sys, in_axes"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gnsZo-GWSYYj"
},
"source": [
"If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`, which returns a batched `mjModel` along with a dictionary specifying the axes that are batched."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1K45Kp2ASV9s"
},
"outputs": [],
"source": [
"rng = jax.random.PRNGKey(0)\n",
"rng = jax.random.split(rng, 10)\n",
"batched_sys, _ = domain_randomize(env.sys, rng)\n",
"\n",
"print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
"print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
"\n",
"print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
"print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "efnxNOnpQFuC"
},
"source": [
"## Quadruped Env\n",
"\n",
"Let's define a quadruped environment that takes advantage of the domain randomization function. Here we use the [Barkour v0 Quadruped](https://github.com/google-deepmind/mujoco_menagerie/tree/main/google_barkour_v0) and an environment that trains a joystick policy."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "y79PoJOCIl-O"
},
"outputs": [],
@@ -483,13 +871,14 @@
" # Penalize non-zero roll and pitch angles. L2 penalty.\n",
" orientation=-5.0,\n",
" # L2 regularization of joint torques, |tau|^2.\n",
" torques=-0.0002,\n",
" # torques=-0.0002,\n",
" torques=-0.002,\n",
" # Penalize the change in the action and encourage smooth\n",
" # actions. L2 regularization |action - last_action|^2\n",
" action_rate=-0.3,\n",
" action_rate=-0.1,\n",
" # Encourage long swing steps. However, it does not\n",
" # encourage high clearances.\n",
" feet_air_time=0.1,\n",
" feet_air_time=0.2,\n",
" # Encourage no motion at zero command, L2 regularization\n",
" # |q - q_default|^2.\n",
" stand_still=-0.5,\n",
@@ -525,6 +914,9 @@
" )\n",
" mj_model = mujoco.MjModel.from_xml_path(\n",
" (path / 'barkour_v0_mjx.xml').as_posix())\n",
" mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
" mj_model.opt.iterations = 4\n",
" mj_model.opt.ls_iterations = 6\n",
"\n",
" physics_steps_per_control_step = 10\n",
" kwargs['physics_steps_per_control_step'] = kwargs.get(\n",
@@ -550,7 +942,7 @@
" self.uppers = self._default_ap_pose + jp.array([0.2, 0.8, 0.8] * 4)\n",
"\n",
" def sample_command(self, rng: jax.Array) -\u003e jax.Array:\n",
" lin_vel_x = [-0.6, 1.5] # min max [m/s]\n",
" lin_vel_x = [-0.6, 1.0] # min max [m/s]\n",
" lin_vel_y = [-0.8, 0.8] # min max [m/s]\n",
" ang_vel_yaw = [-0.7, 0.7] # min max [rad/s]\n",
"\n",
@@ -857,171 +1249,19 @@
"envs.register_environment('barkour', BarkourEnv)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "P1K6IznI2y83"
},
"source": [
"## Visualize a Rollout\n",
"\n",
"Let's visualize an environment rollout by running a few env steps with sinusoidal actuation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "EhKLFK54C1CH"
"id": "pi_yrcz-Qp3W"
},
"outputs": [],
"source": [
"# instantiate the environment\n",
"\n",
"env_name = 'barkour' # @param ['barkour', 'humanoid']\n",
"env_name = 'barkour'\n",
"env = envs.get_environment(env_name)\n",
"\n",
"# define the jit reset/step functions\n",
"jit_reset = jax.jit(env.reset)\n",
"jit_step = jax.jit(env.step)\n",
"\n",
"# instantiate the renderer\n",
"renderer = mujoco.Renderer(env.model)\n",
"\n",
"def get_image(state: State):\n",
" d = mujoco.MjData(env.model)\n",
" # write the mjx.Data into an mjData object\n",
" mjx.device_get_into(d, state.pipeline_state)\n",
" mujoco.mj_forward(env.model, d)\n",
" if env_name == 'barkour':\n",
" camera='track'\n",
" elif env_name == 'humanoid':\n",
" camera = 'side'\n",
" else:\n",
" raise NotImplementedError(env_name)\n",
" # use the mjData object to update the renderer\n",
" renderer.update_scene(d, camera=camera)\n",
" return renderer.render()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Ph8u-v2Q2xLS"
},
"outputs": [],
"source": [
"# initialize the state\n",
"state = jit_reset(jax.random.PRNGKey(0))\n",
"rollout = [state]\n",
"images = [get_image(state)]\n",
"\n",
"# grab a trajectory\n",
"n_steps = 500\n",
"render_every = 5\n",
"\n",
"for i in range(n_steps):\n",
" ctrl = jp.ones(env.sys.nu) * jp.sin(2.0 * jp.pi * i / 500)\n",
" state = jit_step(state, ctrl)\n",
" rollout.append(state)\n",
" if i % render_every == 0:\n",
" images.append(get_image(state))\n",
"\n",
"media.show_video(images, fps=1.0 / env.dt / render_every)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "65mIPj6DQNNa"
},
"source": [
"# Domain Randomization\n",
"\n",
"We have all the pieces to train a policy, but we might also want to include randomization over certain `mjModel` parameters. In MJX, we can easily create a batch of environments with randomized values populated in `mjx.Model`. Below, we show a function that randomizes friction and actuator gain/bias."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "h8mhzKjHQuoL"
},
"outputs": [],
"source": [
"def domain_randomize(sys, rng):\n",
" \"\"\"Randomizes the mjx.Model.\"\"\"\n",
" @jax.vmap\n",
" def rand(rng):\n",
" _, key = jax.random.split(rng, 2)\n",
" # friction\n",
" friction = jax.random.uniform(key, (1,), minval=0.8, maxval=1.2)\n",
" friction = sys.geom_friction.at[:, 0].set(friction)\n",
" # actuator\n",
" _, key = jax.random.split(key, 2)\n",
" gain_range = (-5, 5)\n",
" param = jax.random.uniform(\n",
" key, (1,), minval=gain_range[0], maxval=gain_range[1]\n",
" ) + sys.actuator_gainprm[:, 0]\n",
" gain = sys.actuator_gainprm.at[:, 0].set(param)\n",
" bias = sys.actuator_biasprm.at[:, 1].set(-param)\n",
" return friction, gain, bias\n",
"\n",
" friction, gain, bias = rand(rng)\n",
"\n",
" in_axes = jax.tree_map(lambda x: None, sys)\n",
" in_axes = in_axes.tree_replace({\n",
" 'geom_friction': 0,\n",
" 'actuator_gainprm': 0,\n",
" 'actuator_biasprm': 0,\n",
" })\n",
"\n",
" sys = sys.tree_replace({\n",
" 'geom_friction': friction,\n",
" 'actuator_gainprm': gain,\n",
" 'actuator_biasprm': bias,\n",
" })\n",
"\n",
" return sys, in_axes"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gnsZo-GWSYYj"
},
"source": [
"If we wanted 10 environments with randomized friction and actuator params, we can call `domain_randomize`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "1K45Kp2ASV9s"
},
"outputs": [],
"source": [
"rng = jax.random.PRNGKey(0)\n",
"rng = jax.random.split(rng, 10)\n",
"batched_sys, _ = domain_randomize(env.sys, rng)\n",
"\n",
"print('Single env friction shape: ', env.sys.geom_friction.shape)\n",
"print('Batched env friction shape: ', batched_sys.geom_friction.shape)\n",
"\n",
"print('Friction on geom 0: ', env.sys.geom_friction[0, 0])\n",
"print('Random frictions on geom 0: ', batched_sys.geom_friction[:, 0, 0])"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "LdqPYBlPVWwc"
},
"source": [
"Note that we'll pass the `domain_randomize` function into the brax trainer; the brax training algorithms will call the randomizer over the environment batch dimension."
"# re-instantiate the renderer\n",
"renderer = mujoco.Renderer(env.model)"
]
},
{
@@ -1030,9 +1270,9 @@
"id": "nxaNFP9mA23H"
},
"source": [
"# Train Policy\n",
"## Train Policy\n",
"\n",
"Let's finally train a policy with PPO. Training takes about 12-13 minutes on a Tesla V100 GPU for both environments."
"To train a policy with domain randomization, we pass in the domain randomization function into the brax train function; brax will call the domain randomization function when rolling out episodes. Training the quadruped takes about 14 minutes on a Tesla V100 GPU."
]
},
{
@@ -1043,59 +1283,28 @@
},
"outputs": [],
"source": [
"make_networks_factory = functools.partial(\n",
" ppo_networks.make_ppo_networks,\n",
" policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
"train_fn = functools.partial(\n",
" ppo.train,\n",
" num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n",
" episode_length=1000, normalize_observations=True,\n",
" action_repeat=1, unroll_length=20, num_minibatches=8, gae_lambda=0.95,\n",
" num_updates_per_batch=4, discounting=0.99, learning_rate=3e-4,\n",
" entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n",
" network_factory=make_networks_factory,\n",
" num_resets_per_eval=10,\n",
" randomization_fn=domain_randomize, seed=0)\n",
"\n",
"if env_name == 'barkour':\n",
" make_networks_factory = functools.partial(\n",
" ppo_networks.make_ppo_networks,\n",
" policy_hidden_layer_sizes=(128, 128, 128, 128))\n",
" train_fn = functools.partial(\n",
" ppo.train,\n",
" num_timesteps=60_000_000, num_evals=3, reward_scaling=1,\n",
" episode_length=1000, normalize_observations=True,\n",
" action_repeat=1, unroll_length=20, num_minibatches=32, gae_lambda=0.95,\n",
" num_updates_per_batch=4, discounting=0.99, learning_rate=3.0e-4,\n",
" entropy_cost=1e-2, num_envs=8192, batch_size=1024,\n",
" network_factory=make_networks_factory,\n",
" num_resets_per_eval=10,\n",
" randomization_fn=domain_randomize, seed=0)\n",
"elif env_name == 'humanoid':\n",
" train_fn = functools.partial(\n",
" ppo.train, num_timesteps=30_000_000, num_evals=5, reward_scaling=0.1,\n",
" episode_length=1000, normalize_observations=True, action_repeat=1,\n",
" unroll_length=10, num_minibatches=32, num_updates_per_batch=8,\n",
" discounting=0.97, learning_rate=3e-4, entropy_cost=1e-3, num_envs=2048,\n",
" batch_size=1024, seed=0)\n",
"else:\n",
" raise NotImplementedError(\n",
" f'env_name: {env_name} is not implemented in this notebook.')\n",
"\n",
"\n",
"max_y = {'barkour':30, 'humanoid': 13000}[env_name]\n",
"min_y = {'barkour': -15}.get(env_name, 0)\n",
"\n",
"x_data = []\n",
"y_data = []\n",
"ydataerr = []\n",
"times = [datetime.now()]\n",
"max_y, min_y = 30, 0\n",
"\n",
"def progress(num_steps, metrics):\n",
" times.append(datetime.now())\n",
" x_data.append(num_steps)\n",
" y_data.append(metrics['eval/episode_reward'])\n",
" ydataerr.append(metrics['eval/episode_reward_std'])\n",
"\n",
" plt.xlim([0, train_fn.keywords['num_timesteps'] * 1.25])\n",
" plt.ylim([min_y, max_y])\n",
"\n",
" plt.xlabel('# environment steps')\n",
" plt.ylabel('reward per episode')\n",
" plt.title(f'y={y_data[-1]:.3f}')\n",
"\n",
" plt.errorbar(\n",
" x_data, y_data, yerr=ydataerr)\n",
" plt.show()\n",
"\n",
"# Reset environments since internals may be overwritten by tracers due to\n",
"# Reset environments since internals may be overwritten by tracers from the\n",
"# domain randomization function.\n",
"env = envs.get_environment(env_name)\n",
"eval_env = envs.get_environment(env_name)\n",
@@ -1107,30 +1316,6 @@
"print(f'time to train: {times[-1] - times[1]}')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "-dRmWHsHB_1K"
},
"source": [
"# Save and Load Policy\n",
"\n",
"We can save and load the policy using the brax model API."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "u9sdN4Xa5JoH"
},
"outputs": [],
"source": [
"#@title Save Model\n",
"model_path = '/tmp/mjx_brax_policy'\n",
"model.save_params(model_path, params)"
]
},
{
"cell_type": "code",
"execution_count": null,
@@ -1139,7 +1324,9 @@
},
"outputs": [],
"source": [
"#@title Load Model and Define Inference Function\n",
"# Save and reload params.\n",
"model_path = '/tmp/mjx_brax_quadruped_policy'\n",
"model.save_params(model_path, params)\n",
"params = model.load_params(model_path)\n",
"\n",
"inference_fn = make_inference_fn(params)\n",
@@ -1152,9 +1339,9 @@
"id": "L01IrN4oCIkC"
},
"source": [
"# Visualize Policy\n",
"## Visualize Policy\n",
"\n",
"Finally we can visualize the policy. For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
"For the Barkour Quadruped, the joystick commands can be set through `x_vel`, `y_vel`, and `ang_vel`. `x_vel` and `y_vel` define the linear forward and sideways velocities with respect to the quadruped torso. `ang_vel` defines the angular velocity of the torso in the z direction."
]
},
{
@@ -1183,7 +1370,7 @@
"# @markdown Commands **only used for Barkour Env**:\n",
"x_vel = 1.0 #@param {type: \"number\"}\n",
"y_vel = 0.0 #@param {type: \"number\"}\n",
"ang_vel = -0.0 #@param {type: \"number\"}\n",
"ang_vel = -0.5 #@param {type: \"number\"}\n",
"\n",
"the_command = jp.array([x_vel, y_vel, ang_vel])\n",
"\n",
@@ -1192,7 +1379,7 @@
"state = jit_reset(rng)\n",
"state.info['command'] = the_command\n",
"rollout = [state]\n",
"images = [get_image(state)]\n",
"images = [get_image(state, camera='track')]\n",
"\n",
"# grab a trajectory\n",
"n_steps = 500\n",
@@ -1204,7 +1391,7 @@
" state = jit_step(state, ctrl)\n",
" rollout.append(state)\n",
" if i % render_every == 0:\n",
" images.append(get_image(state))\n",
" images.append(get_image(state, camera='track'))\n",
"\n",
"media.show_video(images, fps=1.0 / eval_env.dt / render_every)"
]
@@ -1216,7 +1403,12 @@
"gpuClass": "premium",
"gpuType": "V100",
"private_outputs": true,
"provenance": [],
"provenance": [
{
"file_id": "1brcF4_qCRS2ASc-QQw1rsEwl5IjzGvq2",
"timestamp": 1697763780236
}
],
"toc_visible": true
},
"kernelspec": {
model/plugin/elasticity/floppy_flex.xml
+3 -3
View File
@@ -22,7 +22,7 @@
<compiler autolimits="true"/>
<option solver="CG" tolerance="1e-6" timestep=".001" integrator="implicitfast"/>
<option solver="CG" tolerance="1e-6" timestep=".001"/>
<size memory="100M"/>
@@ -32,13 +32,13 @@
<worldbody>
<flexcomp type="grid" count="24 4 4" spacing=".1 .1 .1" pos=".1 0 1.5"
radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="7">
radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="25">
<contact condim="3" solref="0.01 1" solimp=".95 .99 .0001" selfcollide="none"/>
<edge damping="1"/>
<plugin plugin="mujoco.elasticity.solid">
<config key="poisson" value="0.2"/>
<!--Units are in Pa (SI)-->
<config key="young" value="5e4"/>
<config key="damping" value="0.002"/>
</plugin>
</flexcomp>
plugin/elasticity/README.md
+1 -1
View File
@@ -1,5 +1,5 @@
<h1>
<a href="#"><img alt="MuJoCo" src="banner.png" width="100%"/></a>
<a href="#"><img alt="MuJoCo" src="../../banner.png" width="100%"/></a>
</h1>
## Elasticity plugins
plugin/elasticity/elasticity.cc
+4 -2
View File
@@ -64,8 +64,10 @@ int CreateStencils(std::vector<T>& elements,
elements[t].edges[e] = it->second;
}
if (!edgeidx.empty()) {
assert(elements[t].edges[e] == edgeidx[T::kNumEdges*t+e]);
if (!edgeidx.empty()) { // SHOULD NOT OCCUR
if (elements[t].edges[e] != edgeidx[T::kNumEdges*t+e]) {
mju_error("edge ordering is incoherent between flex and plugin");
}
}
}
}
plugin/elasticity/elasticity.h
+55
View File
@@ -48,6 +48,13 @@ inline void UpdateSquaredLengths(std::vector<mjtNum>& len,
}
}
inline void UpdateSquaredLengthsFlex(std::vector<mjtNum>& len,
const mjtNum* flexedge_length) {
for (int e = 0; e < len.size(); e++) {
len[e] = flexedge_length[e]*flexedge_length[e];
}
}
struct Stencil2D {
static constexpr int kNumEdges = 3;
static constexpr int kNumVerts = 3;
@@ -78,6 +85,54 @@ void inline GradSquaredLengths(mjtNum gradient[T::kNumEdges][2][3],
}
}
template <typename T>
inline void ComputeForce(mjtNum* qfrc_passive,
const std::vector<T>& elements,
const std::vector<mjtNum>& metric,
const std::vector<mjtNum>& elongationglob,
const mjtNum* xpos) {
for (int t = 0; t < elements.size(); t++) {
const int* v = elements[t].vertices;
// compute length gradient with respect to dofs
mjtNum gradient[T::kNumEdges][2][3];
GradSquaredLengths<T>(gradient, xpos, v);
// extract elongation of edges belonging to this element
mjtNum elongation[T::kNumEdges];
for (int e = 0; e < T::kNumEdges; e++) {
int idx = elements[t].edges[e];
elongation[e] = elongationglob[idx];
}
// we now multiply the elongations by the precomputed metric tensor,
// notice that if metric=diag(1/reference) then this would yield a
// mass-spring model
// compute local force
mjtNum force[T::kNumVerts*3] = {0};
int offset = T::kNumEdges*T::kNumEdges;
for (int ed1 = 0; ed1 < T::kNumEdges; ed1++) {
for (int ed2 = 0; ed2 < T::kNumEdges; ed2++) {
for (int i = 0; i < 2; i++) {
for (int x = 0; x < 3; x++) {
force[3 * T::edge[ed2][i] + x] +=
elongation[ed1] * gradient[ed2][i][x] *
metric[offset * t + T::kNumEdges * ed1 + ed2];
}
}
}
}
// insert into global force
for (int i = 0; i < T::kNumVerts; i++) {
for (int x = 0; x < 3; x++) {
qfrc_passive[3*v[i]+x] -= force[3*i+x];
}
}
}
}
// compute metric tensor of edge lengths inner product
template <typename T>
void inline MetricTensor(std::vector<mjtNum>& metric, int idx, mjtNum mu,
plugin/elasticity/membrane.cc
+41 -40
View File
@@ -84,10 +84,12 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
mjtNum E = strtod(mj_getPluginConfig(m, instance, "young"), nullptr);
mjtNum thick =
strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
mjtNum damp =
strtod(mj_getPluginConfig(m, instance, "damping"), nullptr);
std::vector<int> face, edge;
String2Vector(mj_getPluginConfig(m, instance, "face"), face);
String2Vector(mj_getPluginConfig(m, instance, "edge"), edge);
return Membrane(m, d, instance, nu, E, thick, face, edge);
return Membrane(m, d, instance, nu, E, thick, damp, face, edge);
} else {
mju_warning("Invalid parameter specification in shell plugin");
return std::nullopt;
@@ -96,9 +98,10 @@ std::optional<Membrane> Membrane::Create(const mjModel* m, mjData* d,
// plugin constructor
Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
mjtNum E, mjtNum thick, const std::vector<int>& simplex,
mjtNum E, mjtNum thick, mjtNum damp,
const std::vector<int>& simplex,
const std::vector<int>& edgeidx)
: thickness(thick) {
: f0(-1), damping(damp), thickness(thick) {
// count plugin bodies
nv = ne = 0;
for (int i = 1; i < m->nbody; i++) {
@@ -152,50 +155,48 @@ Membrane::Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu,
// compute metric tensor
MetricTensor<Stencil2D>(metric, t, mu, la, basis);
}
// allocate array
ne = edges.size();
reference.assign(ne, 0);
deformed.assign(ne, 0);
previous.assign(ne, 0);
elongation.assign(ne, 0);
// compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
// save previous lengths
previous = reference;
}
void Membrane::Compute(const mjModel* m, mjData* d, int instance) {
for (int t = 0; t < nt; t++) {
int* v = elements[t].vertices;
mjtNum kD = damping / m->opt.timestep;
// compute length gradient with respect to dofs
mjtNum gradient[kNumEdges][2][3];
GradSquaredLengths<Stencil2D>(gradient, d->xpos+3*i0, v);
// update edge lengths
if (f0 < 0) {
UpdateSquaredLengths(deformed, edges, d->xpos+3*i0);
} else {
UpdateSquaredLengthsFlex(deformed,
d->flexedge_length + m->flex_edgeadr[f0]);
}
// compute elongation
mjtNum elongation[kNumEdges];
for (int e = 0; e < kNumEdges; e++) {
int idx = elements[t].edges[e] + m->flex_edgeadr[f0];
mjtNum deformed = d->flexedge_length[idx]*d->flexedge_length[idx];
mjtNum reference = m->flexedge_length0[idx]*m->flexedge_length0[idx];
elongation[e] = deformed - reference;
}
// we add generalized Rayleigh damping as decribed in Section 5.2 of
// Kharevych et al., "Geometric, Variational Integrators for Computer
// Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
// we now multiply the elongations by the precomputed metric tensor,
// notice that if metric=diag(1/reference) then this would yield a
// mass-spring model
for (int idx = 0; idx < ne; idx++) {
elongation[idx] = deformed[idx] - reference[idx] +
( deformed[idx] - previous[idx] ) * kD;
}
// compute local force
mjtNum force[kNumVerts*3] = {0};
int offset = kNumEdges*kNumEdges;
for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
for (int i = 0; i < 2; i++) {
for (int x = 0; x < 3; x++) {
force[3 * Stencil2D::edge[ed2][i] + x] +=
elongation[ed1] * gradient[ed2][i][x] *
metric[offset * t + kNumEdges * ed1 + ed2];
}
}
}
}
// compute gradient of elastic energy and insert into passive force
ComputeForce<Stencil2D>(d->qfrc_passive + m->body_dofadr[i0], elements,
metric, elongation, d->xpos + 3 * i0);
// insert into global force
for (int i = 0; i < kNumVerts; i++) {
for (int x = 0; x < 3; x++) {
d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
}
}
// update stored lengths
if (kD > 0) {
previous = deformed;
}
}
@@ -208,7 +209,7 @@ void Membrane::RegisterPlugin() {
plugin.name = "mujoco.elasticity.membrane";
plugin.capabilityflags |= mjPLUGIN_PASSIVE;
const char* attributes[] = {"face", "edge", "young", "poisson", "thickness"};
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/elasticity/membrane.h
+7 -2
View File
@@ -52,13 +52,18 @@ class Membrane {
std::vector<std::pair<int, int> > edges; // edge to vertex map (ne x 2)
// precomputed quantities
std::vector<mjtNum> metric; // geom-induced metric (nt x 9)
std::vector<mjtNum> metric; // geom-induced metric (nt x 9)
std::vector<mjtNum> reference; // reference lengths (ne x 1)
std::vector<mjtNum> deformed; // deformed lengths (ne x 1)
std::vector<mjtNum> previous; // previous-step lengths (ne x 1)
std::vector<mjtNum> elongation; // edge elongation (ne x 1)
mjtNum damping;
mjtNum thickness;
private:
Membrane(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
mjtNum thick, const std::vector<int>& simplex,
mjtNum thick, mjtNum damp, const std::vector<int>& simplex,
const std::vector<int>& edgeidx);
};
plugin/elasticity/solid.cc
+21 -55
View File
@@ -165,75 +165,41 @@ Solid::Solid(const mjModel* m, mjData* d, int instance, mjtNum nu, mjtNum E,
reference.assign(ne, 0);
deformed.assign(ne, 0);
previous.assign(ne, 0);
elongation.assign(ne, 0);
// compute edge lengths at equilibrium
// compute edge lengths at equilibrium (m->flexedge_length0 not yet available)
UpdateSquaredLengths(reference, edges, m->body_pos+3*i0);
// save previous lengths
previous = reference;
}
void Solid::Compute(const mjModel* m, mjData* d, int instance) {
// update edges if no flex
mjtNum kD = damping / m->opt.timestep;
// update edge lengths
if (f0 < 0) {
UpdateSquaredLengths(deformed, edges, d->xpos+3*i0);
} else {
UpdateSquaredLengthsFlex(deformed,
d->flexedge_length + m->flex_edgeadr[f0]);
}
// loop over all elements
for (int t = 0; t < nt; t++) {
int* v = elements[t].vertices;
// we add generalized Rayleigh damping as decribed in Section 5.2 of
// Kharevych et al., "Geometric, Variational Integrators for Computer
// Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
// compute length gradient with respect to dofs
mjtNum gradient[kNumEdges][2][3];
GradSquaredLengths<Stencil3D>(gradient, d->xpos+3*i0, v);
// we add generalized Rayleigh damping as decribed in Section 5.2 of
// Kharevych et al., "Geometric, Variational Integrators for Computer
// Animation" http://multires.caltech.edu/pubs/DiscreteLagrangian.pdf
// compute elongation
mjtNum elongation[kNumEdges];
for (int e = 0; e < kNumEdges; e++) {
if (f0 < 0) {
int idx = elements[t].edges[e];
mjtNum kD = damping / m->opt.timestep;
elongation[e] = deformed[idx] - reference[idx] +
( deformed[idx] - previous[idx] ) * kD;
} else {
int idx = elements[t].edges[e] + m->flex_edgeadr[f0];
mjtNum deformed = d->flexedge_length[idx]*d->flexedge_length[idx];
mjtNum reference = m->flexedge_length0[idx]*m->flexedge_length0[idx];
elongation[e] = deformed - reference;
}
}
// we now multiply the elongations by the precomputed metric tensor,
// notice that if metric=diag(1/reference) then this would yield a
// mass-spring model
// compute local force
mjtNum force[kNumVerts*3] = {0};
int offset = kNumEdges*kNumEdges;
for (int ed1 = 0; ed1 < kNumEdges; ed1++) {
for (int ed2 = 0; ed2 < kNumEdges; ed2++) {
for (int i = 0; i < 2; i++) {
for (int x = 0; x < 3; x++) {
force[3 * Stencil3D::edge[ed2][i] + x] +=
elongation[ed1] * gradient[ed2][i][x] *
metric[offset * t + kNumEdges * ed1 + ed2];
}
}
}
}
// insert into global force
for (int i = 0; i < kNumVerts; i++) {
for (int x = 0; x < 3; x++) {
d->qfrc_passive[m->body_dofadr[i0]+3*v[i]+x] -= force[3*i+x];
}
}
for (int idx = 0; idx < ne; idx++) {
elongation[idx] = deformed[idx] - reference[idx] +
( deformed[idx] - previous[idx] ) * kD;
}
// compute gradient of elastic energy and insert into passive force
ComputeForce<Stencil3D>(d->qfrc_passive + m->body_dofadr[i0], elements,
metric, elongation, d->xpos + 3 * i0);
// update stored lengths
if (f0 < 0) {
if (kD > 0) {
previous = deformed;
}
}
plugin/elasticity/solid.h
+1
View File
@@ -54,6 +54,7 @@ class Solid {
std::vector<mjtNum> reference; // reference lengths (ne x 1)
std::vector<mjtNum> deformed; // deformed lengths (ne x 1)
std::vector<mjtNum> previous; // previous-step lengths (ne x 1)
std::vector<mjtNum> elongation; // edge elongation (ne x 1)
mjtNum damping;
plugin/sdf/README.md
+100
View File
@@ -0,0 +1,100 @@
<h1>
<a href="#"><img alt="MuJoCo" src="../../banner.png" width="100%"/></a>
</h1>
## Signed distance function (SDF) plugins
These are first-party plugins that implement implicit geometries using SDFs. They can be applied to **geoms** and
**meshes** (in the **asset** section). Sample models can be found in [this folder](../../model/plugin/sdf/).
### Bolt
Implemented in [bolt.cc](bolt.cc). Example usage in [nutbolt.xml](../../model/plugin/sdf/nutbolt.xml).
This plugin implements a bolt with a hexagonal head, similar to https://www.shadertoy.com/view/XtffzX.
Parameters:
- `radius` [m]: bolt radius (default `0.26`).
### Bowl
Implemented in [bowl.cc](bowl.cc). Example usage in [bowl.xml](../../model/plugin/sdf/bowl.xml).
The plugin implements a cut hollow sphere from https://www.shadertoy.com/view/7tVXRt.
Parameters:
- `height` [m]: location of the cut plane (default `0.4`).
- `radius` [m]: radius of the sphere (default `1`).
- `thickness` [m]: thickness of the bowl (default `0.02`).
### Gear
Implemented in [gear.cc](gear.cc). Example usage in [gear.xml](../../model/plugin/sdf/gear.xml).
The plugin implements a 3D extrusion of the 2D gear geometry from https://www.shadertoy.com/view/3lG3WR.
Parameters:
- `alpha` [m]: initial angle of rotation of the gear (default `0`).
- `diameter` [m]: gear diameter (default `2.8`).
- `teeth` []: number of teeth (default `25`).
### Nut
Implemented in [nut.cc](nut.cc). Example usage in [nutbolt.xml](../../model/plugin/sdf/nutbolt.xml).
This plugin implements a hexagonal nut identical to the bolt head from https://www.shadertoy.com/view/XtffzX.
Parameters:
- `radius` [m]: nut radius (default `0.26`).
### Torus
Implemented in [torus.cc](torus.cc). Example usage in [torus.xml](../../model/plugin/sdf/torus.xml).
This plugin implements a torus.
Parameters:
- `radius1` [m]: major radius (default `0.35`).
- `radius1` [m]: minor radius (default `0.15`).
### How to make your own SDF
Create your `MySDF.h` and `MySDF.cc` files in the SDF folder, where this README is located. Implement your SDF using the
following interface:
```
struct MySDFAttribute {
static constexpr int nattribute =
/* insert the number of attributes */;
static constexpr char const* names[nattribute] =
/* an array of attributes with the same order as the attribute array in your SDF class */;
static constexpr mjtNum defaults[nattribute] =
/* an array of default values for your attributes */;
};
class MySDF {
public:
// creates a new MySDF instance or returns null on failure.
static std::optional<MySDF> Create(const mjModel* m, mjData* d, int instance);
MySDF(MySDF&&) = default;
~MySDF() = default;
// functions that return the SDF and its gradient at a query point
mjtNum Distance(const mjtNum point[3]) const;
void Gradient(mjtNum grad[3], const mjtNum point[3]) const;
// a call to this needs to be added to register.cc
static void RegisterPlugin();
// an array of attributes with the same order as in the struct above
mjtNum attribute[MySDFAttribute::nattribute];
private:
Torus(const mjModel* m, mjData* d, int instance);
};
```
plugin/sdf/bolt.cc
+21 -9
View File
@@ -13,8 +13,10 @@
// limitations under the License.
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <optional>
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
@@ -75,7 +77,13 @@ std::optional<Bolt> Bolt::Create(
// plugin constructor
Bolt::Bolt(const mjModel* m, mjData* d, int instance) {
radius = strtod(mj_getPluginConfig(m, instance, "radius"), nullptr);
SdfDefault<BoltAttribute> defattribute;
for (int i=0; i < BoltAttribute::nattribute; i++) {
attribute[i] = defattribute.GetDefault(
BoltAttribute::names[i],
mj_getPluginConfig(m, instance, BoltAttribute::names[i]));
}
}
// add new element in the vector storing iteration counts
@@ -96,19 +104,19 @@ void Bolt::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
// sdf
mjtNum Bolt::Distance(const mjtNum point[3]) const {
return distance(point, &radius);
return distance(point, attribute);
}
// gradient of sdf
void Bolt::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
mjtNum eps = 1e-8;
mjtNum dist0 = distance(point, &radius);
mjtNum dist0 = distance(point, attribute);
mjtNum pointX[3] = {point[0]+eps, point[1], point[2]};
mjtNum distX = distance(pointX, &radius);
mjtNum distX = distance(pointX, attribute);
mjtNum pointY[3] = {point[0], point[1]+eps, point[2]};
mjtNum distY = distance(pointY, &radius);
mjtNum distY = distance(pointY, attribute);
mjtNum pointZ[3] = {point[0], point[1], point[2]+eps};
mjtNum distZ = distance(pointZ, &radius);
mjtNum distZ = distance(pointZ, attribute);
grad[0] = (distX - dist0) / eps;
grad[1] = (distY - dist0) / eps;
@@ -123,9 +131,8 @@ void Bolt::RegisterPlugin() {
plugin.name = "mujoco.sdf.bolt";
plugin.capabilityflags |= mjPLUGIN_SDF;
const char* attributes[] = {"radius"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nattribute = BoltAttribute::nattribute;
plugin.attributes = BoltAttribute::names;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -177,6 +184,11 @@ void Bolt::RegisterPlugin() {
aabb[3] = aabb[4] = .6;
aabb[5] = 1;
};
plugin.sdf_attribute =
+[](mjtNum attribute[], const char* name[], const char* value[]) {
SdfDefault<BoltAttribute> defattribute;
defattribute.GetDefaults(attribute, name, value);
};
mjp_registerPlugin(&plugin);
}
plugin/sdf/bolt.h
+7 -1
View File
@@ -28,6 +28,12 @@ namespace mujoco::plugin::sdf {
// this plugin implements a modification of the signed distance function
// from https://www.shadertoy.com/view/XtffzX of a bolt with a hexagonal head
struct BoltAttribute {
static constexpr int nattribute = 1;
static constexpr char const* names[nattribute] = {"radius"};
static constexpr mjtNum defaults[nattribute] = {0.26};
};
class Bolt {
public:
// Creates a new Bolt instance (allocated with `new`) or
@@ -45,7 +51,7 @@ class Bolt {
static void RegisterPlugin();
mjtNum radius;
mjtNum attribute[BoltAttribute::nattribute];
private:
Bolt(const mjModel* m, mjData* d, int instance);
plugin/sdf/bowl.cc
+25 -14
View File
@@ -12,12 +12,14 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <sstream>
#include <cstdint>
#include <optional>
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "bowl.h"
namespace mujoco::plugin::sdf {
@@ -51,9 +53,16 @@ std::optional<Bowl> Bowl::Create(
// plugin constructor
Bowl::Bowl(const mjModel* m, mjData* d, int instance) {
radius = strtod(mj_getPluginConfig(m, instance, "radius"), nullptr);
height = strtod(mj_getPluginConfig(m, instance, "height"), nullptr);
thick = strtod(mj_getPluginConfig(m, instance, "thickness"), nullptr);
SdfDefault<BowlAttribute> defattribute;
for (int i=0; i < BowlAttribute::nattribute; i++) {
attribute[i] = defattribute.GetDefault(
BowlAttribute::names[i],
mj_getPluginConfig(m, instance, BowlAttribute::names[i]));
}
mjtNum height = attribute[0];
mjtNum radius = attribute[1];
width = mju_sqrt(radius*radius - height*height);
}
@@ -75,8 +84,7 @@ void Bowl::Visualize(const mjModel* m, mjData* d, const mjvOption* opt, mjvScene
// sdf
mjtNum Bowl::Distance(const mjtNum point[3]) const {
mjtNum attributes[3]= {height, radius, thick};
return distance(point, attributes);
return distance(point, attribute);
}
// gradient of sdf
@@ -98,16 +106,15 @@ void Bowl::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
// grad[1] = - grad_dist * point[1] / pnorm;
// grad[2] = - grad_dist * point[2] / pnorm;
// }
mjtNum attributes[3]= {height, radius, thick};
mjtNum eps = 1e-8;
mjtNum dist0 = distance(point, attributes);
mjtNum dist0 = distance(point, attribute);
mjtNum pointX[3] = {point[0]+eps, point[1], point[2]};
mjtNum distX = distance(pointX, attributes);
mjtNum distX = distance(pointX, attribute);
mjtNum pointY[3] = {point[0], point[1]+eps, point[2]};
mjtNum distY = distance(pointY, attributes);
mjtNum distY = distance(pointY, attribute);
mjtNum pointZ[3] = {point[0], point[1], point[2]+eps};
mjtNum distZ = distance(pointZ, attributes);
mjtNum distZ = distance(pointZ, attribute);
grad[0] = (distX - dist0) / eps;
grad[1] = (distY - dist0) / eps;
@@ -122,9 +129,8 @@ void Bowl::RegisterPlugin() {
plugin.name = "mujoco.sdf.bowl";
plugin.capabilityflags |= mjPLUGIN_SDF;
const char* attributes[] = {"radius", "height", "thickness"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nattribute = BowlAttribute::nattribute;
plugin.attributes = BowlAttribute::names;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -177,6 +183,11 @@ void Bowl::RegisterPlugin() {
aabb[0] = aabb[1] = aabb[2] = 0;
aabb[3] = aabb[4] = aabb[5] = radius + thick;
};
plugin.sdf_attribute =
+[](mjtNum attribute[], const char* name[], const char* value[]) {
SdfDefault<BowlAttribute> defattribute;
defattribute.GetDefaults(attribute, name, value);
};
mjp_registerPlugin(&plugin);
}
plugin/sdf/bowl.h
+8 -3
View File
@@ -25,6 +25,13 @@
namespace mujoco::plugin::sdf {
struct BowlAttribute {
static constexpr int nattribute = 3;
static constexpr char const* names[nattribute] = {"height", "radius",
"thickness"};
static constexpr mjtNum defaults[nattribute] = {.4, 1., .02};
};
class Bowl {
public:
// Creates a new Bowl instance (allocated with `new`) or
@@ -42,9 +49,7 @@ class Bowl {
static void RegisterPlugin();
mjtNum radius;
mjtNum height;
mjtNum thick;
mjtNum attribute[BowlAttribute::nattribute];
mjtNum width;
private:
plugin/sdf/gear.cc
+45 -28
View File
@@ -12,14 +12,16 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <cstdio>
#include <sstream>
#include <cstdint>
#include <cstdlib>
#include <optional>
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mjvisualize.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "gear.h"
namespace mujoco::plugin::sdf {
@@ -50,10 +52,11 @@ static mjtNum mod(mjtNum x, mjtNum y) {
return x - y * floor(x/y);
}
static mjtNum distance2D(const mjtNum p[3], const mjtNum attributes[3]) {
static mjtNum distance2D(const mjtNum p[3],
const mjtNum attributes[GearAttribute::nattribute]) {
// see https://www.shadertoy.com/view/3lG3WR
mjtNum D = 2.8; // should be an attribute
mjtNum N = 25; // should be an attribute
mjtNum D = attributes[1];
mjtNum N = attributes[2];
mjtNum psi = 3.096e-5 * N * N -6.557e-3 * N + 0.551; // pressure angle
mjtNum alpha = attributes[0];
@@ -78,7 +81,10 @@ static mjtNum distance2D(const mjtNum p[3], const mjtNum attributes[3]) {
mjtNum h = 2.2 / Pd;
mjtNum innerR = Ro - h - 0.4;
mjtNum innerR = Ro - h - 0.14*D;
if (attributes[4] >= 0.0) {
innerR = attributes[4] / 2.0;
}
// Early exit
if (innerR - rho > 0.0)
@@ -126,33 +132,41 @@ static mjtNum distance2D(const mjtNum p[3], const mjtNum attributes[3]) {
fib - (alphaStride - shift));
cogs = Intersection(baseWalls, cogs);
cogs = smoothIntersection(gearOuter, cogs, 0.01);
cogs = smoothIntersection(gearOuter, cogs, 0.0035*D);
cogs = smoothUnion(gearLowBase, cogs, Rb - Ro + h);
cogs = Subtraction(cogs, crownBase);
return extrusion(p, cogs, .1);
return cogs;
}
static mjtNum distance(const mjtNum p[3], const mjtNum attributes[3]) {
return extrusion(p, distance2D(p, attributes), .1) - .005;
static mjtNum distance(const mjtNum p[3],
const mjtNum attributes[GearAttribute::nattribute]) {
return extrusion(p, distance2D(p, attributes), attributes[3]/2.);
}
} // namespace
// factory function
std::optional<Gear> Gear::Create(
const mjModel* m, mjData* d, int instance) {
if (CheckAttr("alpha", m, instance)) {
return Gear(m, d, instance);
std::optional<Gear> Gear::Create(const mjModel* m, mjData* d, int instance) {
if (CheckAttr("alpha", m, instance) && CheckAttr("diameter", m, instance) &&
CheckAttr("teeth", m, instance) &&
CheckAttr("innerdiameter", m, instance)) {
return Gear(m, d, instance);
} else {
mju_warning("Invalid parameter specification in Gear plugin");
return std::nullopt;
mju_warning("Invalid parameter specification in Gear plugin");
return std::nullopt;
}
}
// plugin constructor
Gear::Gear(const mjModel* m, mjData* d, int instance) {
alpha = strtod(mj_getPluginConfig(m, instance, "alpha"), nullptr);
SdfDefault<GearAttribute> defattribute;
for (int i=0; i < GearAttribute::nattribute; i++) {
attribute[i] = defattribute.GetDefault(
GearAttribute::names[i],
mj_getPluginConfig(m, instance, GearAttribute::names[i]));
}
}
// plugin computation
@@ -173,21 +187,20 @@ void Gear::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
// sdf
mjtNum Gear::Distance(const mjtNum point[3]) const {
return distance(point, &alpha);
return distance(point, attribute);
}
// gradient of sdf
void Gear::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
mjtNum attributes[1]= {alpha};
mjtNum eps = 1e-8;
mjtNum dist0 = distance(point, attributes);
mjtNum dist0 = distance(point, attribute);
mjtNum pointX[3] = {point[0]+eps, point[1], point[2]};
mjtNum distX = distance(pointX, attributes);
mjtNum distX = distance(pointX, attribute);
mjtNum pointY[3] = {point[0], point[1]+eps, point[2]};
mjtNum distY = distance(pointY, attributes);
mjtNum distY = distance(pointY, attribute);
mjtNum pointZ[3] = {point[0], point[1], point[2]+eps};
mjtNum distZ = distance(pointZ, attributes);
mjtNum distZ = distance(pointZ, attribute);
grad[0] = (distX - dist0) / eps;
grad[1] = (distY - dist0) / eps;
@@ -202,9 +215,8 @@ void Gear::RegisterPlugin() {
plugin.name = "mujoco.sdf.gear";
plugin.capabilityflags |= mjPLUGIN_SDF;
const char* attributes[] = {"alpha"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nattribute = GearAttribute::nattribute;
plugin.attributes = GearAttribute::names;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -253,8 +265,13 @@ void Gear::RegisterPlugin() {
plugin.sdf_aabb =
+[](mjtNum aabb[6], const mjtNum* attributes) {
aabb[0] = aabb[1] = aabb[2] = 0;
aabb[3] = aabb[4] = 1.7;
aabb[5] = .11;
aabb[3] = aabb[4] = attributes[1] / 2. * 1.25;
aabb[5] = attributes[3] / 2. * 1.1;
};
plugin.sdf_attribute =
+[](mjtNum attribute[], const char* name[], const char* value[]) {
SdfDefault<GearAttribute> defattribute;
defattribute.GetDefaults(attribute, name, value);
};
mjp_registerPlugin(&plugin);
plugin/sdf/gear.h
+9 -2
View File
@@ -16,7 +16,6 @@
#define MUJOCO_PLUGIN_SDF_GEAR_H_
#include <optional>
#include <vector>
#include <mujoco/mjdata.h>
#include <mujoco/mjmodel.h>
@@ -26,6 +25,14 @@
namespace mujoco::plugin::sdf {
struct GearAttribute {
static constexpr int nattribute = 5;
static constexpr char const* names[nattribute] = {"alpha", "diameter",
"teeth", "thickness",
"innerdiameter"};
static constexpr mjtNum defaults[nattribute] = { 0, 2.8, 25, .2 , -1};
};
class Gear {
public:
// Creates a new Gear instance (allocated with `new`) or
@@ -43,7 +50,7 @@ class Gear {
static void RegisterPlugin();
mjtNum alpha;
mjtNum attribute[GearAttribute::nattribute];
private:
Gear(const mjModel* m, mjData* d, int instance);
plugin/sdf/nut.cc
+23 -10
View File
@@ -13,12 +13,15 @@
// limitations under the License.
#include <cmath>
#include <sstream>
#include <cstdint>
#include <cstdlib>
#include <optional>
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "nut.h"
namespace mujoco::plugin::sdf {
@@ -74,7 +77,13 @@ std::optional<Nut> Nut::Create(
// plugin constructor
Nut::Nut(const mjModel* m, mjData* d, int instance) {
radius = strtod(mj_getPluginConfig(m, instance, "radius"), nullptr);
SdfDefault<NutAttribute> defattribute;
for (int i=0; i < NutAttribute::nattribute; i++) {
attribute[i] = defattribute.GetDefault(
NutAttribute::names[i],
mj_getPluginConfig(m, instance, NutAttribute::names[i]));
}
}
// plugin computation
@@ -95,19 +104,19 @@ void Nut::Visualize(const mjModel* m, mjData* d, const mjvOption* opt,
// sdf
mjtNum Nut::Distance(const mjtNum point[3]) const {
return distance(point, &radius);
return distance(point, attribute);
}
// gradient of sdf
void Nut::Gradient(mjtNum grad[3], const mjtNum point[3]) const {
mjtNum eps = 1e-8;
mjtNum dist0 = distance(point, &radius);
mjtNum dist0 = distance(point, attribute);
mjtNum pointX[3] = {point[0]+eps, point[1], point[2]};
mjtNum distX = distance(pointX, &radius);
mjtNum distX = distance(pointX, attribute);
mjtNum pointY[3] = {point[0], point[1]+eps, point[2]};
mjtNum distY = distance(pointY, &radius);
mjtNum distY = distance(pointY, attribute);
mjtNum pointZ[3] = {point[0], point[1], point[2]+eps};
mjtNum distZ = distance(pointZ, &radius);
mjtNum distZ = distance(pointZ, attribute);
grad[0] = (distX - dist0) / eps;
grad[1] = (distY - dist0) / eps;
@@ -122,9 +131,8 @@ void Nut::RegisterPlugin() {
plugin.name = "mujoco.sdf.nut";
plugin.capabilityflags |= mjPLUGIN_SDF;
const char* attributes[] = {"radius"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nattribute = NutAttribute::nattribute;
plugin.attributes = NutAttribute::names;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -176,6 +184,11 @@ void Nut::RegisterPlugin() {
aabb[3] = aabb[4] = .6;
aabb[5] = 1;
};
plugin.sdf_attribute =
+[](mjtNum attribute[], const char* name[], const char* value[]) {
SdfDefault<NutAttribute> defattribute;
defattribute.GetDefaults(attribute, name, value);
};
mjp_registerPlugin(&plugin);
}
plugin/sdf/nut.h
+7 -1
View File
@@ -28,6 +28,12 @@ namespace mujoco::plugin::sdf {
// this plugin implements a modification of the signed distance function
// from https://www.shadertoy.com/view/XtffzX of hexagonal nut
struct NutAttribute {
static constexpr int nattribute = 1;
static constexpr char const* names[nattribute] = {"radius"};
static constexpr mjtNum defaults[nattribute] = {0.26};
};
class Nut {
public:
// Creates a new Nut instance (allocated with `new`) or
@@ -45,7 +51,7 @@ class Nut {
static void RegisterPlugin();
mjtNum radius;
mjtNum attribute[NutAttribute::nattribute];
private:
Nut(const mjModel* m, mjData* d, int instance);
plugin/sdf/sdf.cc
+3 -1
View File
@@ -13,10 +13,12 @@
// limitations under the License.
#include <algorithm>
#include <cctype>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <mujoco/mjplugin.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
plugin/sdf/sdf.h
+40
View File
@@ -15,6 +15,9 @@
#ifndef MUJOCO_PLUGIN_SDF_SDF_H_
#define MUJOCO_PLUGIN_SDF_SDF_H_
#include <map>
#include <stdexcept>
#include <string>
#include <vector>
#include <mujoco/mujoco.h>
@@ -40,6 +43,43 @@ inline mjtNum Fract(mjtNum x) {
// reads numeric attributes
bool CheckAttr(const char* name, const mjModel* m, int instance);
// converts attributes to numeric or returns default if not present
template <typename T>
class SdfDefault {
public:
SdfDefault() {
for (int i = 0; i < T::nattribute; i++) {
default_[T::names[i]] = T::defaults[i];
}
}
// get a single default value
mjtNum GetDefault(const char* name, const char* value) {
if (std::string(value).empty()) {
return default_[name];
}
try {
mjtNum num = std::stod(value);
return num;
} catch (const std::invalid_argument& e) {
mju_error("invalid attribute value for '%s'", name);
return 0;
}
}
// populate attribute array
void GetDefaults(mjtNum* attribute, const char* names[],
const char* values[]) {
for (int i = 0; i < default_.size(); i++) {
attribute[i] = GetDefault(names[i], values[i]);
}
}
private:
std::map<std::string, mjtNum> default_;
};
// stores the history of gradient descent iterations
class SdfVisualizer {
public:
SdfVisualizer();
plugin/sdf/torus.cc
+20 -9
View File
@@ -12,12 +12,14 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <sstream>
#include <cstdint>
#include <optional>
#include <utility>
#include <mujoco/mjplugin.h>
#include <mujoco/mjtnum.h>
#include <mujoco/mujoco.h>
#include "sdf.h"
#include "torus.h"
namespace mujoco::plugin::sdf {
@@ -43,19 +45,24 @@ std::optional<Torus> Torus::Create(
// plugin constructor
Torus::Torus(const mjModel* m, mjData* d, int instance) {
radius[0] = strtod(mj_getPluginConfig(m, instance, "radius1"), nullptr);
radius[1] = strtod(mj_getPluginConfig(m, instance, "radius2"), nullptr);
SdfDefault<TorusAttribute> defattribute;
for (int i=0; i < TorusAttribute::nattribute; i++) {
attribute[i] = defattribute.GetDefault(
TorusAttribute::names[i],
mj_getPluginConfig(m, instance, TorusAttribute::names[i]));
}
}
// sdf
mjtNum Torus::Distance(const mjtNum point[3]) const {
return distance(point, radius);
return distance(point, attribute);
}
// gradient of sdf
void Torus::Gradient(mjtNum grad[3], const mjtNum p[3]) const {
mjtNum len_xy = mju_sqrt(p[0]*p[0] + p[1]*p[1]);
mjtNum q = len_xy - radius[0];
mjtNum q = len_xy - attribute[0];
mjtNum grad_q[2] = { p[0] / len_xy, p[1] / len_xy };
mjtNum len_qz = mju_sqrt(q*q + p[2]*p[2]);
grad[0] = q*grad_q[0] / mjMAX(len_qz, mjMINVAL);
@@ -71,9 +78,8 @@ void Torus::RegisterPlugin() {
plugin.name = "mujoco.sdf.torus";
plugin.capabilityflags |= mjPLUGIN_SDF;
const char* attributes[] = {"radius1", "radius2", "axis"};
plugin.nattribute = sizeof(attributes) / sizeof(attributes[0]);
plugin.attributes = attributes;
plugin.nattribute = TorusAttribute::nattribute;
plugin.attributes = TorusAttribute::names;
plugin.nstate = +[](const mjModel* m, int instance) { return 0; };
plugin.init = +[](const mjModel* m, mjData* d, int instance) {
@@ -117,8 +123,13 @@ void Torus::RegisterPlugin() {
aabb[3] = aabb[4] = attributes[0] + attributes[1];
aabb[5] = attributes[1];
};
plugin.sdf_attribute =
+[](mjtNum attribute[], const char* name[], const char* value[]) {
SdfDefault<TorusAttribute> defattribute;
defattribute.GetDefaults(attribute, name, value);
};
mjp_registerPlugin(&plugin);
mjp_registerPlugin(&plugin);
}
} // namespace mujoco::plugin::sdf
plugin/sdf/torus.h
+7 -1
View File
@@ -24,6 +24,12 @@
namespace mujoco::plugin::sdf {
struct TorusAttribute {
static constexpr int nattribute = 2;
static constexpr char const* names[nattribute] = {"radius1", "radius2"};
static constexpr mjtNum defaults[nattribute] = { .35, .15 };
};
class Torus {
public:
// Creates a new Torus instance or returns null on failure.
@@ -36,7 +42,7 @@ class Torus {
static void RegisterPlugin();
mjtNum radius[2];
mjtNum attribute[TorusAttribute::nattribute];
private:
Torus(const mjModel* m, mjData* d, int instance);
python/build_requirements.txt
+2 -2
View File
@@ -41,8 +41,8 @@ numpy==1.26.0; python_version >= '3.9' \
--hash=sha256:c78a22e95182fb2e7874712433eaa610478a3caf86f28c621708d35fa4fd6e7f \
--hash=sha256:86f737708b366c36b76e953c46ba5827d8c27b7a8c9d0f471810728e5a2fe57c \
--hash=sha256:020cdbee66ed46b671429c7265cf00d8ac91c046901c55684954c3958525dab2
pip==23.2.1 \
--hash=sha256:7ccf472345f20d35bdc9d1841ff5f313260c2c33fe417f48c30ac46cccabf5be
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
PyOpenGL==3.1.7 \
--hash=sha256:a6ab19cf290df6101aaf7470843a9c46207789855746399d0af92521a0a92b7a
pytest==7.4.2 \
python/make_sdist_requirements.txt
+2 -2
View File
@@ -2,8 +2,8 @@ absl-py==2.0.0 \
--hash=sha256:9a28abb62774ae4e8edbe2dd4c49ffcd45a6a848952a5eccc6a49f3f0fc1e2f3
build==1.0.3 \
--hash=sha256:589bf99a67df7c9cf07ec0ac0e5e2ea5d4b37ac63301c4986d1acb126aa83f8f
pip==23.2.1 \
--hash=sha256:7ccf472345f20d35bdc9d1841ff5f313260c2c33fe417f48c30ac46cccabf5be
pip==23.3.1 \
--hash=sha256:55eb67bb6171d37447e82213be585b75fe2b12b359e993773aca4de9247a052b
setuptools==68.2.2 \
--hash=sha256:b454a35605876da60632df1a60f736524eb73cc47bbc9f3f1ef1b644de74fd2a
python/mujoco/CMakeLists.txt
-1
View File
@@ -265,7 +265,6 @@ target_sources(
structs_header
INTERFACE indexer_xmacro.h
indexers.h
mjdata_meta.h
structs.h
)
set_target_properties(structs_header PROPERTIES PUBLIC_HEADER structs.h)
python/mujoco/bindings_test.py
+89 -4
View File
@@ -74,10 +74,21 @@ TEST_XML_SENSOR = r"""
"""
TEST_XML_PLUGIN = r"""
<mujoco model="test">
<mujoco>
<option gravity="0 0 0"/>
<extension>
<plugin plugin="mujoco.elasticity.cable"/>
</extension>
<worldbody>
<composite type="cable" curve="s" count="41 1 1" size="1" offset="0 0 1" initial="none">
<plugin plugin="mujoco.elasticity.cable">
<config key="twist" value="1e6"/>
<config key="bend" value="1e9"/>
</plugin>
<joint kind="main" damping="2"/>
<geom type="capsule" size=".005" density="1"/>
</composite>
</worldbody>
</mujoco>
"""
@@ -1048,10 +1059,24 @@ Euler integrator, semi-implicit in velocity.
while data_instances:
d = data_instances.pop()
self.assertEqual(sys.getrefcount(d), 2)
del d
while model_instances:
m = model_instances.pop()
self.assertEqual(sys.getrefcount(m), 2)
# This test is disabled on PyPy as it uses sys.getrefcount
# However PyPy is not officially supported by MuJoCo
@absltest.skipIf(sys.implementation.name == 'pypy',
reason='requires sys.getrefcount')
def test_mjdata_holds_ref_to_model(self):
data = mujoco.MjData(mujoco.MjModel.from_xml_string('<mujoco/>'))
model = data.model
# references: one in `data.model, one in `model`, one in the temporary ref
# passed to getrefcount.
self.assertEqual(sys.getrefcount(data.model), 3)
del data
self.assertEqual(sys.getrefcount(model), 2)
def test_can_initialize_mjv_structs(self):
self.assertIsInstance(mujoco.MjvScene(), mujoco.MjvScene)
self.assertIsInstance(mujoco.MjvCamera(), mujoco.MjvCamera)
@@ -1110,6 +1135,41 @@ Euler integrator, semi-implicit in velocity.
bodyexclude=0,
geomid=geomid)
def test_mj_multi_ray(self):
nray = 3
geom1 = np.zeros(1, np.int32)
pnt = np.array([-0.3, 0, 0.1])
vec = np.array([[1, 0, 0], [0, 0, 1], [0, 0, -1]], np.float64)
dist_ex = np.array([0.2, -1, 0.1])
geom_ex = np.array([1, -1, 0])
geomid = np.zeros(nray, np.int32)
dist = np.zeros(nray, np.float64)
mujoco.mj_forward(self.model, self.data)
mujoco.mj_multiRay(
m=self.model,
d=self.data,
pnt=pnt,
vec=vec.flatten(),
geomgroup=None,
flg_static=1,
bodyexclude=-1,
geomid=geomid,
dist=dist,
nray=nray,
cutoff=mujoco.mjMAXVAL)
for i in range(0, 3):
self.assertEqual(
dist[i],
mujoco.mj_ray(
self.model, self.data, pnt, vec[i], None, 1, -1, geom1
),
)
self.assertEqual(geomid[i], geom1)
self.assertEqual(geomid[i], geom_ex[i])
self.assertAlmostEqual(dist[i], dist_ex[i])
def test_inverse_fd_none(self):
eps = 1e-6
flg_centered = 0
@@ -1287,6 +1347,34 @@ Euler integrator, semi-implicit in velocity.
self.assertEqual(data.geom(3).xpos[2], 4)
self.assertEqual(data.geom(4).xpos[2], 5)
def test_load_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)
def test_copy_mjdata_with_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data1 = mujoco.MjData(model)
self.assertIs(data1.model, model)
mujoco.mj_step(model, data1)
data2 = copy.copy(data1)
mujoco.mj_step(model, data1)
mujoco.mj_step(model, data2)
np.testing.assert_array_equal(data1.qpos, data2.qpos)
self.assertIs(data1.model, data2.model)
def test_deepcopy_mjdata_with_plugin(self):
model = mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
data1 = mujoco.MjData(model)
self.assertIs(data1.model, model)
mujoco.mj_step(model, data1)
data2 = copy.deepcopy(data1)
mujoco.mj_step(model, data1)
mujoco.mj_step(model, data2)
np.testing.assert_array_equal(data1.qpos, data2.qpos)
self.assertIsNot(data1.model, data2.model)
self.assertNotEqual(data1.model._address, data2.model._address)
def _assert_attributes_equal(self, actual_obj, expected_obj, attr_to_compare):
for name in attr_to_compare:
actual_value = getattr(actual_obj, name)
@@ -1300,9 +1388,6 @@ Euler integrator, semi-implicit in velocity.
self.fail("Attribute '{}' differs from expected value: {}".format(
name, str(e)))
def test_load_plugin(self):
mujoco.MjModel.from_xml_string(TEST_XML_PLUGIN)
if __name__ == '__main__':
absltest.main()
python/mujoco/functions.cc
+24 -3
View File
@@ -23,6 +23,8 @@
#include <Eigen/Core>
#include <mujoco/mjxmacro.h>
#include <mujoco/mujoco.h>
#include "errors.h"
#include "function_traits.h"
#include "functions.h"
#include "private.h"
@@ -37,6 +39,7 @@ PYBIND11_MODULE(_functions, pymodule) {
namespace py = ::pybind11;
namespace traits = python_traits;
using EigenVectorI = Eigen::Vector<int, Eigen::Dynamic>;
using EigenVectorX = Eigen::Vector<mjtNum, Eigen::Dynamic>;
using EigenArrayXX = Eigen::Array<
mjtNum, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
@@ -567,7 +570,25 @@ PYBIND11_MODULE(_functions, pymodule) {
Def<traits::mj_versionString>(pymodule);
// Ray collision
Def<traits::mj_multiRay>(pymodule);
Def<traits::mj_multiRay>(
pymodule,
[](const raw::MjModel* m, raw::MjData* d, const mjtNum(*pnt)[3],
Eigen::Ref<const EigenVectorX> vec,
std::optional<Eigen::Ref<const Eigen::Vector<mjtByte, mjNGROUP>>>
geomgroup,
mjtByte flg_static, int bodyexclude, Eigen::Ref<EigenVectorI> geomid,
Eigen::Ref<EigenVectorX> dist, int nray, mjtNum cutoff) {
if (dist.size() != nray || geomid.size() != nray) {
throw py::type_error("dist and geomid should be of size nray");
}
if (vec.size() != 3 * nray) {
throw py::type_error("vec should be of size 3*nray");
}
InterceptMjErrors(::mj_multiRay)(
m, d, &(*pnt)[0], vec.data(),
geomgroup.has_value() ? geomgroup->data() : nullptr, flg_static,
bodyexclude, geomid.data(), dist.data(), nray, cutoff);
});
Def<traits::mj_ray>(
pymodule,
[](const raw::MjModel* m, const raw::MjData* d, const mjtNum(*pnt)[3],
@@ -1367,7 +1388,7 @@ PYBIND11_MODULE(_functions, pymodule) {
}
#undef MJ_M
#define MJ_M(x) d.metadata().x
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) data->x
#define X(type, name, nr, nc) \
@@ -1379,7 +1400,7 @@ PYBIND11_MODULE(_functions, pymodule) {
}
MJDATA_ARENA_POINTERS_PRIMAL
if (d.metadata().is_dual) {
if (mj_isDual(d.model().get())) {
MJDATA_ARENA_POINTERS_DUAL
}
#undef X
python/mujoco/indexers.cc
+13 -21
View File
@@ -13,16 +13,13 @@
// limitations under the License.
#include <algorithm>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <utility>
#include <variant>
#include <vector>
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include "util/crossplatform.h"
@@ -76,7 +73,6 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
// the MuJoCo category of the entity itself, e.g. nbody indicates that the
// field belongs to a body.
// T: Scalar data type of the field.
// M: Either raw::MjModel or MjDataMetadata.
//
// Args:
// base_ptr: Pointer to the first entry in the entire field.
@@ -87,35 +83,35 @@ IDToName MakeIDToName(int count, IntPtr name_offsets, CharPtr names) {
// additional dimension of the size len(qvel) of the particular joint.
// m: Used for dereferencing MjSize.
// owner: The base object whose lifetime is tied to the returned array.
template <auto MjSize, typename T, typename M>
template <auto MjSize, typename T>
py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
const M& m, py::handle owner) {
const raw::MjModel& m, py::handle owner) {
int offset;
if (MjSize == &M::nq) {
if (MjSize == &raw::MjModel::nq) {
offset = m.jnt_qposadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_qposadr[index+1] : m.nq) - offset);
} else if (MjSize == &M::nv) {
} else if (MjSize == &raw::MjModel::nv) {
offset = m.jnt_dofadr[index];
shape.insert(
shape.begin(),
((index < m.njnt-1) ? m.jnt_dofadr[index+1] : m.nv) - offset);
} else if (MjSize == &M::nhfielddata) {
} else if (MjSize == &raw::MjModel::nhfielddata) {
offset = m.hfield_adr[index];
shape.insert(shape.begin(), m.hfield_ncol[index]);
shape.insert(shape.begin(), m.hfield_nrow[index]);
} else if (MjSize == &M::ntexdata) {
} else if (MjSize == &raw::MjModel::ntexdata) {
offset = m.tex_adr[index];
shape.insert(shape.begin(), m.tex_width[index]);
shape.insert(shape.begin(), m.tex_height[index]);
} else if (MjSize == &M::nsensordata) {
} else if (MjSize == &raw::MjModel::nsensordata) {
offset = m.sensor_adr[index];
shape.insert(shape.begin(), m.sensor_dim[index]);
} else if (MjSize == &M::nnumericdata) {
} else if (MjSize == &raw::MjModel::nnumericdata) {
offset = m.numeric_adr[index];
shape.insert(shape.begin(), m.numeric_size[index]);
} else if (MjSize == &M::ntupledata) {
} else if (MjSize == &raw::MjModel::ntupledata) {
offset = m.tuple_adr[index];
shape.insert(shape.begin(), m.tuple_size[index]);
} else {
@@ -135,9 +131,7 @@ py::array_t<T> MakeArray(T* base_ptr, int index, std::vector<int>&& shape,
}
} // namespace
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
NameToIDMappings::NameToIDMappings(const M& m)
NameToIDMappings::NameToIDMappings(const raw::MjModel& m)
: body(MakeNameToID(m.nbody, m.name_bodyadr, m.names)),
jnt(MakeNameToID(m.njnt, m.name_jntadr, m.names)),
geom(MakeNameToID(m.ngeom, m.name_geomadr, m.names)),
@@ -160,9 +154,7 @@ NameToIDMappings::NameToIDMappings(const M& m)
tuple(MakeNameToID(m.ntuple, m.name_tupleadr, m.names)),
key(MakeNameToID(m.nkey, m.name_keyadr, m.names)) {}
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
IDToNameMappings::IDToNameMappings(const M& m)
IDToNameMappings::IDToNameMappings(const raw::MjModel& m)
: body(MakeIDToName(m.nbody, m.name_bodyadr, m.names)),
jnt(MakeIDToName(m.njnt, m.name_jntadr, m.names)),
geom(MakeIDToName(m.ngeom, m.name_geomadr, m.names)),
@@ -223,7 +215,7 @@ MJMODEL_VIEW_GROUPS
MJMODEL_VIEW_GROUPS
#undef XGROUP
MjDataIndexer::MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
MjDataIndexer::MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
py::handle owner)
: d_(d),
m_(m),
@@ -369,7 +361,7 @@ MJMODEL_KEYFRAME
#define X(type, prefix, var, dim0, dim1) \
py::array_t<type> XGROUP::var() { \
if (!var##_.has_value()) { \
var##_.emplace(MakeArray<&MjDataMetadata::dim0>( \
var##_.emplace(MakeArray<&raw::MjModel::dim0>( \
d_->prefix##var, index_, MAKE_SHAPE(dim1), *m_, owner_)); \
} \
return *var##_; \
python/mujoco/indexers.h
+6 -11
View File
@@ -25,7 +25,6 @@
#include <absl/container/flat_hash_map.h>
#include <mujoco/mjxmacro.h>
#include "indexer_xmacro.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
@@ -35,9 +34,7 @@ using NameToID = absl::flat_hash_map<std::string, int>;
using IDToName = std::vector<std::string>;
struct NameToIDMappings {
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
explicit NameToIDMappings(const M& m);
explicit NameToIDMappings(const raw::MjModel& m);
NameToID body;
NameToID jnt;
@@ -63,9 +60,7 @@ struct NameToIDMappings {
};
struct IDToNameMappings {
// M is either a raw::MjModel or MjDataMetadata.
template <typename M>
explicit IDToNameMappings(const M& m);
explicit IDToNameMappings(const raw::MjModel& m);
IDToName body;
IDToName jnt;
@@ -163,7 +158,7 @@ class MjModelIndexer {
class MjDataGroupedViewsBase {
public:
MjDataGroupedViewsBase(int index, std::string_view name, raw::MjData* d,
const MjDataMetadata* m,
const raw::MjModel* m,
pybind11::handle owner)
: index_(index), name_(name), d_(d), m_(m), owner_(owner) {}
@@ -175,7 +170,7 @@ class MjDataGroupedViewsBase {
int index_;
std::string name_;
raw::MjData* d_;
const MjDataMetadata* m_;
const raw::MjModel* m_;
pybind11::handle owner_;
};
@@ -201,7 +196,7 @@ MJDATA_VIEW_GROUPS
// (e.g. a particular geom or joint) either by name or by ID.
class MjDataIndexer {
public:
MjDataIndexer(raw::MjData* d, const MjDataMetadata* m,
MjDataIndexer(raw::MjData* d, const raw::MjModel* m,
pybind11::handle owner);
#define XGROUP(MjDataGroupedViews, field, nfield, FIELD_XMACROS) \
@@ -213,7 +208,7 @@ class MjDataIndexer {
private:
raw::MjData* d_;
const MjDataMetadata* m_;
const raw::MjModel* m_;
pybind11::handle owner_;
NameToIDMappings name_to_id_;
IDToNameMappings id_to_name_;
python/mujoco/mjdata_meta.h
-108
View File
@@ -1,108 +0,0 @@
// Copyright 2021 DeepMind Technologies Limited
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef MUJOCO_PYTHON_MJDATA_META_H_
#define MUJOCO_PYTHON_MJDATA_META_H_
#include <cstring>
#include <memory>
#include <mujoco/mujoco.h>
#include <mujoco/mjxmacro.h>
#include "raw.h"
#include "util/crossplatform.h"
namespace mujoco::python {
namespace _impl {
template <typename T> class MjWrapper;
} // namespace _impl
#define MJDATA_METADATA \
X( int, jnt_qposadr, njnt ) \
X( int, jnt_dofadr, njnt ) \
X( int, hfield_nrow, nhfield ) \
X( int, hfield_ncol, nhfield ) \
X( int, hfield_adr, nhfield ) \
X( int, tex_height, ntex ) \
X( int, tex_width, ntex ) \
X( int, tex_adr, ntex ) \
X( int, sensor_dim, nsensor ) \
X( int, sensor_adr, nsensor ) \
X( int, numeric_adr, nnumeric ) \
X( int, numeric_size, nnumeric ) \
X( int, tuple_adr, ntuple ) \
X( int, tuple_size, ntuple ) \
X( int, name_bodyadr, nbody ) \
X( int, name_jntadr, njnt ) \
X( int, name_geomadr, ngeom ) \
X( int, name_siteadr, nsite ) \
X( int, name_camadr, ncam ) \
X( int, name_lightadr, nlight ) \
X( int, name_meshadr, nmesh ) \
X( int, name_skinadr, nskin ) \
X( int, name_hfieldadr, nhfield ) \
X( int, name_texadr, ntex ) \
X( int, name_matadr, nmat ) \
X( int, name_pairadr, npair ) \
X( int, name_excludeadr, nexclude ) \
X( int, name_eqadr, neq ) \
X( int, name_tendonadr, ntendon ) \
X( int, name_actuatoradr, nu ) \
X( int, name_sensoradr, nsensor ) \
X( int, name_numericadr, nnumeric ) \
X( int, name_textadr, ntext ) \
X( int, name_tupleadr, ntuple ) \
X( int, name_keyadr, nkey ) \
X( char, names, nnames )
// A subset of mjModel fields that are required to reconstruct an MjDataWrapper.
struct MjDataMetadata {
public:
friend class _impl::MjWrapper<raw::MjData>;
explicit MjDataMetadata(const raw::MjModel* m)
:
#define X(var) var(m->var),
MJMODEL_INTS
#undef X
#define X(dtype, var, n) \
var([](dtype* src, int len) { \
dtype* dst = new dtype[len]; \
std::memcpy(dst, src, len * sizeof(dtype)); \
return dst; \
}(m->var, m->n)),
MJDATA_METADATA
#undef X
is_dual(mj_isDual(m)) {
}
#define X(var) decltype(raw::MjModel::var) var;
MJMODEL_INTS
#undef X
#define X(type, var, n) std::shared_ptr<type[]> var;
MJDATA_METADATA
#undef X
bool is_dual;
private:
MjDataMetadata() = default;
MjDataMetadata(const MjDataMetadata& other) = default;
MjDataMetadata(MjDataMetadata&& other) = default;
};
} // namespace mujoco::python
#endif // MUJOCO_PYTHON_MJDATA_META_H_
python/mujoco/mjpython/mjpython.py
+28 -1
View File
@@ -25,6 +25,8 @@ import ctypes
import importlib.util
import os
import platform
import re
import subprocess
import sys
if platform.system() != 'Darwin':
@@ -48,7 +50,32 @@ def main(argv):
# Conda doesn't create a separate shared library for Python.
# We instead use the Python binary itself, which can be dlopened just as well.
os.environ['MJPYTHON_LIBPYTHON'] = get_executable_path()
libpython_path = get_executable_path()
os.environ['MJPYTHON_LIBPYTHON'] = libpython_path
# In some installations (e.g. CommandLineTools), the Python interpreter loads
# dylibs from @executable_path-relative paths. This will not resolve
# correctly since @executable_path will be the directory containing the
# mjpython binary when we execve. We therefore preemptively resolve all
# @executable_path-relative paths now and add them to
# DYLD_FALLBACK_LIBRARY_PATH.
libpython_dir = os.path.dirname(libpython_path)
dyld_fallback_paths = (
os.environ.get('DYLD_FALLBACK_LIBRARY_PATH', '').split(':'))
pattern = re.compile(r'@executable_path/(.+) \(offset \d+\)\Z')
otool_out = subprocess.run(
['otool', '-l', libpython_path],
capture_output=True,
check=True,
).stdout.decode()
for line in otool_out.split('\n'):
m = pattern.search(line)
if m is not None:
new_path = os.path.dirname(os.path.join(libpython_dir, m.group(1)))
if new_path not in dyld_fallback_paths:
dyld_fallback_paths.insert(0, new_path)
os.environ['DYLD_FALLBACK_LIBRARY_PATH'] = ':'.join(dyld_fallback_paths)
# argv[0] is currently the path to this script.
# Replace it with sys.executable to preserve e.g. virtualenv path.
python/mujoco/structs.cc
+88 -77
View File
@@ -37,10 +37,11 @@
#include "errors.h"
#include "function_traits.h"
#include "indexers.h"
#include "mjdata_meta.h"
#include "private.h"
#include "raw.h"
#include "serialization.h"
#include <pybind11/cast.h>
#include <pybind11/detail/common.h>
#include <pybind11/numpy.h>
#include <pybind11/operators.h>
#include <pybind11/pybind11.h>
@@ -147,7 +148,8 @@ MjVisualRgbaWrapper::MjWrapper()
X(rangefinder),
X(constraint),
X(slidercrank),
X(crankbroken) {}
X(crankbroken),
X(frustum) {}
MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner)
: WrapperBase(ptr, owner),
@@ -172,7 +174,8 @@ MjVisualRgbaWrapper::MjWrapper(raw::MjVisualRgba* ptr, py::handle owner)
X(rangefinder),
X(constraint),
X(slidercrank),
X(crankbroken) {}
X(crankbroken),
X(frustum) {}
#undef X
MjVisualRgbaWrapper::MjWrapper(const MjVisualRgbaWrapper& other)
@@ -451,7 +454,8 @@ void MjModelWrapper::Serialize(std::ostream& output) const {
WriteBytes(output, buffer.data(), model_size);
}
MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
std::unique_ptr<MjModelWrapper> MjModelWrapper::Deserialize(
std::istream& input) {
CheckInput(input, "mjModel");
char serializationVersion = ReadChar(input);
@@ -477,7 +481,7 @@ MjModelWrapper MjModelWrapper::Deserialize(std::istream& input) {
if (!model) {
throw py::value_error("Invalid serialized mjModel.");
}
return MjModelWrapper(model);
return std::unique_ptr<MjModelWrapper>(new MjModelWrapper(model));
}
// ==================== MJCONTACT ==============================================
@@ -552,13 +556,13 @@ MjDataWrapper* MjDataWrapper::FromRawPointer(raw::MjData* m) noexcept {
}
}
MjDataWrapper::MjWrapper(const MjModelWrapper& model)
: WrapperBase(InterceptMjErrors(mj_makeData)(model.get()),
MjDataWrapper::MjWrapper(MjModelWrapper* model)
: WrapperBase(InterceptMjErrors(mj_makeData)(model->get()),
&MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) model.get()->x
#define MJ_M(x) model->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(model.get()->dim0, dim1), ptr_->var, owner_)),
var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
@@ -569,8 +573,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(model.get()),
indexer_(ptr_, &metadata_, owner_) {
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -585,9 +590,9 @@ MjDataWrapper::MjWrapper(const MjModelWrapper& model)
MjDataWrapper::MjWrapper(const MjDataWrapper& other)
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) other.metadata_.x
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
@@ -599,8 +604,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(other.metadata_),
indexer_(ptr_, &metadata_, owner_) {
model_(other.model_),
model_ref_(other.model_ref_),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -615,9 +621,9 @@ MjDataWrapper::MjWrapper(const MjDataWrapper& other)
MjDataWrapper::MjWrapper(MjDataWrapper&& other)
: WrapperBase(other.ptr_, other.owner_),
#undef MJ_M
#define MJ_M(x) other.metadata_.x
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.metadata_.dim0, dim1), ptr_->var, \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
@@ -629,8 +635,9 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(other.metadata_),
indexer_(ptr_, &metadata_, owner_) {
model_(other.model_),
model_ref_(std::move(other.model_ref_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -644,12 +651,13 @@ MjDataWrapper::MjWrapper(MjDataWrapper&& other)
other.ptr_ = nullptr;
}
MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
: WrapperBase(d, &MjDataCapsuleDestructor),
MjDataWrapper::MjWrapper(const MjDataWrapper& other, MjModelWrapper* model)
: WrapperBase(other.Copy(), &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) metadata.x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(metadata.dim0, dim1), ptr_->var, owner_)),
#define MJ_M(x) other.model_->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(other.model_->get()->dim0, dim1), ptr_->var, \
owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
@@ -660,8 +668,39 @@ MjDataWrapper::MjWrapper(MjDataMetadata&& metadata, raw::MjData* d)
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
metadata_(std::move(metadata)),
indexer_(ptr_, &metadata_, owner_) {
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
is_newly_inserted = MjDataRawPointerMap().insert({ptr_, this}).second;
}
if (!is_newly_inserted) {
throw UnexpectedError(
"MjDataRawPointerMap already contains this raw mjData*");
}
}
MjDataWrapper::MjWrapper(MjModelWrapper* model, raw::MjData* d)
: WrapperBase(d, &MjDataCapsuleDestructor),
#undef MJ_M
#define MJ_M(x) model->get()->x
#define X(dtype, var, dim0, dim1) \
var(InitPyArray(X_ARRAY_SHAPE(model->get()->dim0, dim1), ptr_->var, owner_)),
MJDATA_POINTERS
#undef MJ_M
#define MJ_M(x) (x)
#undef X
contact(MjContactList(ptr_->contact, NConMax(ptr_), &ptr_->ncon, owner_)),
#define X(dtype, var, dim0, dim1) var(InitPyArray(ptr_->var, owner_)),
MJDATA_VECTOR
#undef X
model_(model),
model_ref_(py::cast(model_)),
indexer_(ptr_, model_->get(), owner_) {
bool is_newly_inserted = false;
{
py::gil_scoped_acquire gil;
@@ -692,19 +731,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
// TODO: Replace this custom serialization with a protobuf
WriteChar(output, kSerializationVersion);
// Write all size fields
#define X(var) WriteInt(output, this->metadata_.var);
MJMODEL_INTS
#undef X
WriteInt(output, this->metadata_.is_dual);
#define X(dtype, var, n) \
WriteBytes(output, this->metadata_.var.get(), \
this->metadata_.n * sizeof(dtype));
MJDATA_METADATA
#undef X
model_->Serialize(output);
// Write struct and scalar fields
#define X(var) WriteBytes(output, &ptr_->var, sizeof(ptr_->var))
@@ -727,15 +754,15 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
// Write buffer contents
{
MJDATA_POINTERS_PREAMBLE((&this->metadata_))
MJDATA_POINTERS_PREAMBLE((this->model_->get()))
#define X(type, name, nr, nc) \
WriteBytes(output, ptr_->name, sizeof(type)*(this->metadata_.nr)*(nc));
WriteBytes(output, ptr_->name, sizeof(type)*(this->model_->get()->nr)*(nc));
MJDATA_POINTERS
#undef X
#undef MJ_M
#define MJ_M(x) this->metadata_.x
#define MJ_M(x) this->model_->get()->x
#undef MJ_D
#define MJ_D(x) this->ptr_->x
#define X(type, name, nr, nc) \
@@ -745,7 +772,7 @@ void MjDataWrapper::Serialize(std::ostream& output) const {
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
if (this->metadata_.is_dual) {
if (mj_isDual(this->model_->get())) {
MJDATA_ARENA_POINTERS_DUAL
}
if (this->ptr_->nisland) {
@@ -767,28 +794,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
throw py::value_error("Incompatible serialization version.");
}
// Read all size and address fields
MjDataMetadata metadata;
raw::MjModel m{0};
// Read the model that was used to create the mjData.
std::unique_ptr<MjModelWrapper> m_wrapper =
MjModelWrapper::Deserialize(input);
raw::MjModel& m = *m_wrapper->get();
#define X(var) \
metadata.var = ReadInt(input); \
CheckInput(input, "mjData"); \
m.var = metadata.var;
MJMODEL_INTS
#undef X
metadata.is_dual = ReadInt(input);
#define X(dtype, var, n) \
metadata.var.reset(new dtype[metadata.n]); \
ReadBytes(input, metadata.var.get(), metadata.n * sizeof(dtype)); \
CheckInput(input, "mjData"); \
m.var = metadata.var.get();
MJDATA_METADATA
#undef X
bool is_dual = mj_isDual(&m);
raw::MjData* d = mj_makeData(&m);
if (!d) {
@@ -839,7 +850,7 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
MJDATA_ARENA_POINTERS_CONTACT
MJDATA_ARENA_POINTERS_PRIMAL
if (metadata.is_dual) {
if (is_dual) {
MJDATA_ARENA_POINTERS_DUAL
}
if (d->nisland) {
@@ -859,18 +870,12 @@ MjDataWrapper MjDataWrapper::Deserialize(std::istream& input) {
throw py::value_error("Invalid serialized mjData.");
}
return MjDataWrapper(std::move(metadata), d);
return MjDataWrapper(m_wrapper.release(), d);
}
raw::MjData* MjDataWrapper::Copy() const {
raw::MjModel m{0};
#define X(var) m.var = this->metadata_.var;
MJMODEL_INTS
#undef X
#define X(dtype, var, n) m.var = this->metadata_.var.get();
MJDATA_METADATA
#undef X
return InterceptMjErrors(mj_copyData)(NULL, &m, this->ptr_);
const raw::MjModel* m = model_->get();
return InterceptMjErrors(mj_copyData)(NULL, m, this->ptr_);
}
// ==================== MJSTATISTIC ============================================
@@ -1488,6 +1493,7 @@ PYBIND11_MODULE(_structs, m) {
X(framewidth);
X(constraint);
X(slidercrank);
X(frustum);
#undef X
py::class_<MjVisualRgbaWrapper> mjVisualRgba(mjVisual, "Rgba");
@@ -1522,6 +1528,7 @@ PYBIND11_MODULE(_structs, m) {
X(constraint);
X(slidercrank);
X(crankbroken);
X(frustum);
#undef X
#define X(var) \
@@ -1886,15 +1893,19 @@ This is useful for example when the MJB is not available as a file on disk.)"));
// ==================== MJDATA ===============================================
py::class_<MjDataWrapper> mjData(m, "MjData");
mjData.def(py::init<const MjModelWrapper&>());
mjData.def(py::init<MjModelWrapper*>());
mjData.def_property_readonly("_address", [](const MjDataWrapper& d) {
return reinterpret_cast<std::uintptr_t>(d.get());
});
mjData.def_property_readonly("model", [](const MjDataWrapper& d) {
return &d.model();
});
mjData.def("__copy__", [](const MjDataWrapper& other) {
return MjDataWrapper(other);
});
mjData.def("__deepcopy__", [](const MjDataWrapper& other, py::dict) {
return MjDataWrapper(other);
MjModelWrapper* model_copy = new MjModelWrapper(other.model());
return MjDataWrapper(other, model_copy);
});
mjData.def(py::pickle(
[](const MjDataWrapper& d) { // __getstate__
@@ -1923,7 +1934,7 @@ This is useful for example when the MJB is not available as a file on disk.)"));
#undef X
#undef MJ_M
#define MJ_M(x) d.metadata().x
#define MJ_M(x) d.model().get()->x
#undef MJ_D
#define MJ_D(x) d.get()->x
#define X(dtype, var, dim0, dim1) \
python/mujoco/structs.h
+14 -6
View File
@@ -29,10 +29,10 @@
#include <mujoco/mujoco.h>
#include <mujoco/mjxmacro.h>
#include "indexers.h"
#include "mjdata_meta.h"
#include "raw.h"
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/pytypes.h>
namespace mujoco::python {
namespace _impl {
@@ -232,6 +232,7 @@ class MjWrapper<raw::MjVisualRgba> : public WrapperBase<raw::MjVisualRgba> {
X(constraint);
X(slidercrank);
X(crankbroken);
X(frustum);
#undef X
};
@@ -446,7 +447,8 @@ class MjWrapper<raw::MjModel> : public WrapperBase<raw::MjModel> {
MjModelIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
static MjWrapper<raw::MjModel> Deserialize(std::istream& input);
static std::unique_ptr<MjWrapper<raw::MjModel>> Deserialize(
std::istream& input);
static MjWrapper LoadXMLFile(
const std::string& filename,
@@ -563,12 +565,14 @@ struct is_mj_struct_list<raw::MjContact> {
template <>
class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
public:
explicit MjWrapper(const MjModelWrapper& model);
explicit MjWrapper(MjModelWrapper* model);
MjWrapper(const MjWrapper& other);
MjWrapper(MjWrapper&&);
// Used for deepcopy
MjWrapper(const MjWrapper& other, MjModelWrapper* model);
~MjWrapper();
const MjDataMetadata& metadata() const { return metadata_; }
const MjModelWrapper& model() const { return *model_; }
MjDataIndexer& indexer() { return indexer_; }
void Serialize(std::ostream& output) const;
@@ -597,10 +601,14 @@ class MjWrapper<raw::MjData>: public WrapperBase<raw::MjData> {
protected:
// Internal constructor which takes ownership of given mjData pointer.
// Used for deserialization.
explicit MjWrapper(MjDataMetadata&& metadata, raw::MjData* d);
explicit MjWrapper(MjModelWrapper* model, raw::MjData* d);
raw::MjData* Copy() const;
MjDataMetadata metadata_;
// A reference to the model that was used to create this mjData.
MjModelWrapper* model_;
// A py::object pointing to the same model as model_, to make sure Python
// doesn't doesn't garbage collect it until this mjData is released.
pybind11::object model_ref_;
MjDataIndexer indexer_;
};
python/tutorial.ipynb
+216 -4
View File
@@ -72,7 +72,7 @@
},
"outputs": [],
"source": [
"#@title Check if installation was successful\n",
"#@title Set up rendering, check installation\n",
"\n",
"from google.colab import files\n",
"\n",
@@ -921,7 +921,7 @@
"id": "g1MKUEL_eSCM"
},
"source": [
"Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium. "
"Below is a model of a chaotic pendulum, similar to [this one](https://www.exploratorium.edu/exhibits/chaotic-pendulum) in the San Francisco Exploratorium."
]
},
{
@@ -1242,7 +1242,7 @@
"free_body_MJCF = \"\"\"\n",
"\u003cmujoco\u003e\n",
" \u003casset\u003e\n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\" \n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" rgb1=\".1 .2 .3\"\n",
" rgb2=\".2 .3 .4\" width=\"300\" height=\"300\" mark=\"edge\" markrgb=\".2 .3 .4\"/\u003e\n",
" \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\"\n",
" reflectance=\".2\"/\u003e\n",
@@ -1447,7 +1447,7 @@
" \u003cworldbody\u003e\n",
" \u003clight name=\"light\" pos=\"-.2 0 1\"/\u003e\n",
" \u003cgeom name=\"ground\" type=\"plane\" size=\".5 .5 10\" material=\"grid\"\n",
" zaxis=\"-.3 0 1\" friction=\".1\"/\u003e \n",
" zaxis=\"-.3 0 1\" friction=\".1\"/\u003e\n",
" \u003ccamera name=\"y\" pos=\"-.1 -.6 .3\" xyaxes=\"1 0 0 0 1 2\"/\u003e\n",
" \u003cbody pos=\"0 0 .1\"\u003e\n",
" \u003cjoint/\u003e\n",
@@ -1941,6 +1941,218 @@
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "Zzzugf-qPExb"
},
"source": [
"## Camera control\n",
"\n",
"Cameras can be controlled dynamically in order to achieve cinematic effects. Run the three cells below to see the difference between rendering from a static and moving camera.\n",
"\n",
"The camera-control code smoothly transitions between two trajectories, one orbiting a fixed point, the other tracking a moving object. Parameter values in the code were obtained by iterating quickly on low-res videos."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "-SW-K9WuPGrp"
},
"outputs": [],
"source": [
"#@title Load the \"dominos\" model\n",
"dominos_xml = \"\"\"\n",
"\u003cmujoco\u003e\n",
" \u003casset\u003e\n",
" \u003ctexture type=\"skybox\" builtin=\"gradient\" rgb1=\".3 .5 .7\" rgb2=\"0 0 0\" width=\"32\" height=\"512\"/\u003e\n",
" \u003ctexture name=\"grid\" type=\"2d\" builtin=\"checker\" width=\"512\" height=\"512\" rgb1=\".1 .2 .3\" rgb2=\".2 .3 .4\"/\u003e\n",
" \u003cmaterial name=\"grid\" texture=\"grid\" texrepeat=\"2 2\" texuniform=\"true\" reflectance=\".2\"/\u003e\n",
" \u003c/asset\u003e\n",
"\n",
" \u003cstatistic meansize=\".01\"/\u003e\n",
"\n",
" \u003cvisual\u003e\n",
" \u003cglobal offheight=\"2160\" offwidth=\"3840\"/\u003e\n",
" \u003cquality offsamples=\"8\"/\u003e\n",
" \u003c/visual\u003e\n",
"\n",
" \u003cdefault\u003e\n",
" \u003cgeom type=\"box\" solref=\".005 1\"/\u003e\n",
" \u003cdefault class=\"static\"\u003e\n",
" \u003cgeom rgba=\".3 .5 .7 1\"/\u003e\n",
" \u003c/default\u003e\n",
" \u003c/default\u003e\n",
"\n",
" \u003coption timestep=\"5e-4\"/\u003e\n",
"\n",
" \u003cworldbody\u003e\n",
" \u003clight pos=\".3 -.3 .8\" mode=\"trackcom\" diffuse=\"1 1 1\" specular=\".3 .3 .3\"/\u003e\n",
" \u003clight pos=\"0 -.3 .4\" mode=\"targetbodycom\" target=\"box\" diffuse=\".8 .8 .8\" specular=\".3 .3 .3\"/\u003e\n",
" \u003cgeom name=\"floor\" type=\"plane\" size=\"3 3 .01\" pos=\"-0.025 -0.295 0\" material=\"grid\"/\u003e\n",
" \u003cgeom name=\"ramp\" pos=\".25 -.45 -.03\" size=\".04 .1 .07\" euler=\"-30 0 0\" class=\"static\"/\u003e\n",
" \u003ccamera name=\"top\" pos=\"-0.37 -0.78 0.49\" xyaxes=\"0.78 -0.63 0 0.27 0.33 0.9\"/\u003e\n",
"\n",
" \u003cbody name=\"ball\" pos=\".25 -.45 .1\"\u003e\n",
" \u003cfreejoint name=\"ball\"/\u003e\n",
" \u003cgeom name=\"ball\" type=\"sphere\" size=\".02\" rgba=\".65 .81 .55 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".26 -.3 .03\" euler=\"0 0 -90.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".0015 .015 .03\" rgba=\"1 .5 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".26 -.27 .04\" euler=\"0 0 -81.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".002 .02 .04\" rgba=\"1 1 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".24 -.21 .06\" euler=\"0 0 -63.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".003 .03 .06\" rgba=\".5 1 .5 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".2 -.16 .08\" euler=\"0 0 -45.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".004 .04 .08\" rgba=\".5 1 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".15 -.12 .1\" euler=\"0 0 -27.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".005 .05 .1\" rgba=\".5 .5 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody pos=\".09 -.1 .12\" euler=\"0 0 -9.0\"\u003e\n",
" \u003cfreejoint/\u003e\n",
" \u003cgeom size=\".006 .06 .12\" rgba=\"1 .5 1 1\"/\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody name=\"seasaw_wrapper\" pos=\"-.23 -.1 0\" euler=\"0 0 30\"\u003e\n",
" \u003cgeom size=\".01 .01 .015\" pos=\"0 .05 .015\" class=\"static\"/\u003e\n",
" \u003cgeom size=\".01 .01 .015\" pos=\"0 -.05 .015\" class=\"static\"/\u003e\n",
" \u003cgeom type=\"cylinder\" size=\".01 .0175\" pos=\"-.09 0 .0175\" class=\"static\"/\u003e\n",
" \u003cbody name=\"seasaw\" pos=\"0 0 .03\"\u003e\n",
" \u003cjoint axis=\"0 1 0\"/\u003e\n",
" \u003cgeom type=\"cylinder\" size=\".005 .039\" zaxis=\"0 1 0\" rgba=\".84 .15 .33 1\"/\u003e\n",
" \u003cgeom size=\".1 .02 .005\" pos=\"0 0 .01\" rgba=\".84 .15 .33 1\"/\u003e\n",
" \u003c/body\u003e\n",
" \u003c/body\u003e\n",
"\n",
" \u003cbody name=\"box\" pos=\"-.3 -.14 .05501\" euler=\"0 0 -30\"\u003e\n",
" \u003cfreejoint name=\"box\"/\u003e\n",
" \u003cgeom name=\"box\" size=\".01 .01 .01\" rgba=\".0 .7 .79 1\"/\u003e\n",
" \u003c/body\u003e\n",
" \u003c/worldbody\u003e\n",
"\u003c/mujoco\u003e\n",
"\"\"\"\n",
"model = mujoco.MjModel.from_xml_string(dominos_xml)\n",
"data = mujoco.MjData(model)\n",
"renderer = mujoco.Renderer(model, height=1024, width=1440)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "a2WruafiPhPk"
},
"outputs": [],
"source": [
"#@title Render from fixed camera\n",
"duration = 2.5 # (seconds)\n",
"framerate = 60 # (Hz)\n",
"\n",
"# Simulate and display video.\n",
"frames = []\n",
"mujoco.mj_resetData(model, data) # Reset state and time.\n",
"while data.time \u003c duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) \u003c data.time * framerate:\n",
" renderer.update_scene(data, camera='top')\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "Kie3y-27bQ3J"
},
"outputs": [],
"source": [
"#@title Render from moving camera\n",
"duration = 3 # (seconds)\n",
"\n",
"# find time when box is thrown (speed \u003e 2cm/s)\n",
"throw_time = 0.0\n",
"mujoco.mj_resetData(model, data)\n",
"while data.time \u003c duration and not throw_time:\n",
" mujoco.mj_step(model, data)\n",
" box_speed = np.linalg.norm(data.joint('box').qvel[:3])\n",
" if box_speed \u003e 0.02:\n",
" throw_time = data.time\n",
"assert throw_time \u003e 0\n",
"\n",
"def mix(time, t0=0.0, width=1.0):\n",
" \"\"\"Sigmoidal mixing function.\"\"\"\n",
" t = (time - t0) / width\n",
" s = 1 / (1 + np.exp(-t))\n",
" return 1 - s, s\n",
"\n",
"def unit_cos(t):\n",
" \"\"\"Unit cosine sigmoid from (0,0) to (1,1).\"\"\"\n",
" return 0.5 - np.cos(np.pi*np.clip(t, 0, 1))/2\n",
"\n",
"def orbit_motion(t):\n",
" \"\"\"Return orbit trajectory.\"\"\"\n",
" distance = 0.9\n",
" azimuth = 140 + 100 * unit_cos(t)\n",
" elevation = -30\n",
" lookat = data.geom('floor').xpos.copy()\n",
" return distance, azimuth, elevation, lookat\n",
"\n",
"def track_motion():\n",
" \"\"\"Return box-track trajectory.\"\"\"\n",
" distance = 0.08\n",
" azimuth = 280\n",
" elevation = -10\n",
" lookat = data.geom('box').xpos.copy()\n",
" return distance, azimuth, elevation, lookat\n",
"\n",
"def cam_motion():\n",
" \"\"\"Return sigmoidally-mixed {orbit, box-track} trajectory.\"\"\"\n",
" d0, a0, e0, l0 = orbit_motion(data.time / throw_time)\n",
" d1, a1, e1, l1 = track_motion()\n",
" mix_time = 0.3\n",
" w0, w1 = mix(data.time, throw_time, mix_time)\n",
" return w0*d0+w1*d1, w0*a0+w1*a1, w0*e0+w1*e1, w0*l0+w1*l1\n",
"\n",
"# Make a camera.\n",
"cam = mujoco.MjvCamera()\n",
"mujoco.mjv_defaultCamera(cam)\n",
"\n",
"# Simulate and display video.\n",
"framerate = 60 # (Hz)\n",
"slowdown = 4 # 4x slow-down\n",
"mujoco.mj_resetData(model, data)\n",
"frames = []\n",
"while data.time \u003c duration:\n",
" mujoco.mj_step(model, data)\n",
" if len(frames) \u003c data.time * framerate * slowdown:\n",
" cam.distance, cam.azimuth, cam.elevation, cam.lookat = cam_motion()\n",
" renderer.update_scene(data, cam)\n",
" pixels = renderer.render()\n",
" frames.append(pixels)\n",
"media.show_video(frames, fps=framerate)"
]
}
],
"metadata": {
sample/testspeed.cc
+10 -6
View File
@@ -233,12 +233,15 @@ int main(int argc, char** argv) {
// profiler, top-level
printf(" Internal profiler%s, %ss per step\n", nthread > 1 ? " for thread 0" : "", mu_str);
mjtNum tstep = d[0]->timer[mjTIMER_STEP].duration/d[0]->timer[mjTIMER_STEP].number;
int number = d[0]->timer[mjTIMER_STEP].number;
mjtNum tstep = number ? d[0]->timer[mjTIMER_STEP].duration/number : 0.0;
mjtNum components = 0, total = 0;
for (int i=0; i <= mjTIMER_ADVANCE; i++) {
if (d[0]->timer[i].number > 0) {
mjtNum istep = d[0]->timer[i].duration/d[0]->timer[i].number;
std::printf(" %17s : %6.1f (%6.2f %%)\n", mjTIMERSTRING[i], istep, 100*istep/tstep);
int number = d[0]->timer[i].number;
mjtNum istep = number ? d[0]->timer[i].duration/number : 0.0;
mjtNum percent = number ? 100*istep/tstep : 0.0;
std::printf(" %17s : %6.1f (%6.2f %%)\n", mjTIMERSTRING[i], istep, percent);
// save step time, add up timing of components
if (i == 0) total = istep;
@@ -276,9 +279,10 @@ int main(int argc, char** argv) {
// components of mjTIMER_POS_COLLISION
if (i == mjTIMER_POS_COLLISION) {
for (int j : {mjTIMER_COL_BROAD, mjTIMER_COL_MID, mjTIMER_COL_NARROW}) {
mjtNum jstep = d[0]->timer[j].duration/d[0]->timer[j].number;
std::printf(" %-11s : %6.1f (%6.2f %%)\n",
mjTIMERSTRING[j]+4, jstep, 100*jstep/tstep);
int number = d[0]->timer[j].number;
mjtNum jstep = number ? d[0]->timer[j].duration/number : 0.0;
mjtNum percent = number ? 100*jstep/tstep : 0.0;
std::printf(" %-11s : %6.1f (%6.2f %%)\n", mjTIMERSTRING[j]+4, jstep, percent);
}
}
}
simulate/simulate.cc
+95 -26
View File
@@ -130,8 +130,8 @@ const mjuiDef defFile[] = {
const char help_content[] =
"Space\n"
"+ -\n"
"Right arrow\n"
"Left arrow\n"
"Left / Right arrow\n"
"Tab / Shift-Tab\n"
"[ ]\n"
"Esc\n"
"Double-click\n"
@@ -148,25 +148,25 @@ const char help_content[] =
"F3\n"
"F4\n"
"F5\n"
"UI right hold\n"
"UI right-button hold\n"
"UI title double-click";
const char help_title[] =
"Play / Pause\n"
"Speed up / down\n"
"Step forward\n"
"Step back\n"
"Speed Up / Down\n"
"Step Back / Forward\n"
"Toggle Left / Right UI\n"
"Cycle cameras\n"
"Free camera\n"
"Select\n"
"Select parent\n"
"Center\n"
"Center camera\n"
"Tracking camera\n"
"Zoom\n"
"View rotate\n"
"View translate\n"
"Object rotate\n"
"Object translate\n"
"View Orbit\n"
"View Pan\n"
"Object Rotate\n"
"Object Translate\n"
"Help\n"
"Info\n"
"Profiler\n"
@@ -716,6 +716,16 @@ void MakePhysicsSection(mj::Simulate* sim, int oldstate) {
{mjITEM_EDITNUM, "Friction", 2, &(opt->o_friction), "5"},
{mjITEM_END}
};
mjuiDef defDisableActuator[] = {
{mjITEM_SEPARATOR, "Actuator Group Enable", 1},
{mjITEM_CHECKBYTE, "Act Group 0", 2, sim->enableactuator+0, " 0"},
{mjITEM_CHECKBYTE, "Act Group 1", 2, sim->enableactuator+1, " 1"},
{mjITEM_CHECKBYTE, "Act Group 2", 2, sim->enableactuator+2, " 2"},
{mjITEM_CHECKBYTE, "Act Group 3", 2, sim->enableactuator+3, " 3"},
{mjITEM_CHECKBYTE, "Act Group 4", 2, sim->enableactuator+4, " 4"},
{mjITEM_CHECKBYTE, "Act Group 5", 2, sim->enableactuator+5, " 5"},
{mjITEM_END}
};
// add physics
mjui_add(&sim->ui0, defPhysics);
@@ -736,9 +746,11 @@ void MakePhysicsSection(mj::Simulate* sim, int oldstate) {
defFlag[0].pdata = sim->enable + i;
mjui_add(&sim->ui0, defFlag);
}
// add contact override
mjui_add(&sim->ui0, defOverride);
// add actuator group enable/disable
mjui_add(&sim->ui0, defDisableActuator);
}
@@ -1002,7 +1014,7 @@ void MakeJointSection(mj::Simulate* sim, int oldstate) {
// add scalar joints, exit if UI limit reached
int itemcnt = 0;
for (int i=0; i < sim->jnt_type_.size() && itemcnt<mjMAXUIITEM; i++)
for (int i=0; i < sim->jnt_type_.size() && itemcnt<mjMAXUIITEM; i++) {
if ((sim->jnt_type_[i]==mjJNT_HINGE || sim->jnt_type_[i]==mjJNT_SLIDE)) {
// skip if joint group is disabled
if (!sim->opt.jointgroup[mjMAX(0, mjMIN(mjNGROUP-1, sim->jnt_group_[i]))]) {
@@ -1035,6 +1047,7 @@ void MakeJointSection(mj::Simulate* sim, int oldstate) {
mjui_add(&sim->ui1, defSlider);
itemcnt++;
}
}
}
// make control section of UI
@@ -1051,15 +1064,21 @@ void MakeControlSection(mj::Simulate* sim, int oldstate) {
// add section
mjui_add(&sim->ui1, defControl);
defSlider[0].state = 2;
// add controls, exit if UI limit reached (Clear button already added)
int itemcnt = 1;
for (int i=0; i < sim->actuator_ctrlrange_.size() && itemcnt<mjMAXUIITEM; i++) {
// skip if actuator group is disabled
if (!sim->opt.actuatorgroup[mjMAX(0, mjMIN(mjNGROUP-1, sim->actuator_group_[i]))]) {
// skip if actuator vis group is disabled
int group = sim->actuator_group_[i];
if (!sim->opt.actuatorgroup[mjMAX(0, mjMIN(mjNGROUP-1, group))]) {
continue;
}
// grey out if actuator group is disabled
if (group >= 0 && group <= 30 && sim->m_->opt.disableactuator & (1 << group)) {
defSlider[0].state = 0;
} else {
defSlider[0].state = 2;
}
// set data and name
if (!sim->is_passive_) {
@@ -1189,19 +1208,27 @@ void CopyCamera(mj::Simulate* sim) {
void UpdateSettings(mj::Simulate* sim, const mjModel* m) {
// physics flags
for (int i=0; i<mjNDISABLE; i++) {
int new_value = ((m->opt.disableflags & (1<<i)) !=0);
int new_value = ((m->opt.disableflags & (1<<i)) != 0);
if (sim->disable[i] != new_value) {
sim->disable[i] = new_value;
sim->pending_.ui_update_physics = true;
}
}
for (int i=0; i<mjNENABLE; i++) {
int new_value = ((m->opt.enableflags & (1<<i)) !=0);
int new_value = ((m->opt.enableflags & (1<<i)) != 0);
if (sim->enable[i] != new_value) {
sim->enable[i] = new_value;
sim->pending_.ui_update_physics = true;
}
}
for (int i=0; i<mjNGROUP; i++) {
int enabled = ((m->opt.disableactuator & (1<<i)) == 0);
if (sim->enableactuator[i] != enabled) {
sim->enableactuator[i] = enabled;
sim->pending_.ui_update_physics = true;
sim->pending_.ui_remake_ctrl = true;
}
}
// camera
int old_camera = sim->camera;
@@ -1254,7 +1281,7 @@ int UiPredicate(int category, void* userdata) {
return sim->m_ || sim->is_passive_;
case 3: // require model and nkey
return !sim->is_passive_ && sim->nkey_;
return (sim->m_ || sim->is_passive_) && sim->nkey_;
case 4: // require model and paused
return sim->m_ && !sim->run;
@@ -1404,17 +1431,39 @@ void UiEvent(mjuiState* state) {
// update disable flags in mjOption
opt->disableflags = 0;
for (int i=0; i<mjNDISABLE; i++)
for (int i=0; i<mjNDISABLE; i++) {
if (sim->disable[i]) {
opt->disableflags |= (1<<i);
}
}
// update enable flags in mjOption
opt->enableflags = 0;
for (int i=0; i<mjNENABLE; i++)
for (int i=0; i<mjNENABLE; i++) {
if (sim->enable[i]) {
opt->enableflags |= (1<<i);
}
}
// update disableactuator bitflag in mjOption
bool group_changed = false;
for (int i=0; i<mjNGROUP; i++) {
if ((!sim->enableactuator[i]) != (opt->disableactuator & (1<<i))) {
group_changed = true;
if (!sim->enableactuator[i]) {
// disable actuator group i
opt->disableactuator |= (1<<i);
} else {
// enable actuator group i
opt->disableactuator &= ~(1<<i);
}
}
}
// remake control section if actuator disable group changed
if (group_changed) {
sim->pending_.ui_remake_ctrl = true;
}
}
// rendering section
@@ -1461,10 +1510,7 @@ void UiEvent(mjuiState* state) {
// remake control section if actuator group changed
if (it->name[0]=='A' && it->name[1]=='c') {
sim->ui1.nsect = SECT_CONTROL;
MakeControlSection(sim, sim->ui1.sect[SECT_CONTROL].state);
sim->ui1.nsect = NSECT1;
UiModify(&sim->ui1, state, &sim->platform_ui->mjr_context());
sim->pending_.ui_remake_ctrl = true;
}
}
@@ -1625,6 +1671,18 @@ void UiEvent(mjuiState* state) {
sim->speed_changed = true;
}
break;
case mjKEY_TAB: // toggle left/right UI
if (!state->shift) {
// toggle left UI
sim->ui0_enable = !sim->ui0_enable;
UiModify(&sim->ui0, state, &sim->platform_ui->mjr_context());
} else {
// toggle right UI
sim->ui1_enable = !sim->ui1_enable;
UiModify(&sim->ui1, state, &sim->platform_ui->mjr_context());
}
break;
}
return;
@@ -1778,7 +1836,7 @@ void Simulate::Sync() {
range.emplace(m_->actuator_ctrlrange[2*i], m_->actuator_ctrlrange[2*i + 1]);
}
if (actuator_ctrlrange_[i] != range) {
pending_.ui_update_ctrl = true;
pending_.ui_remake_ctrl = true;
actuator_ctrlrange_[i].swap(range);
}
}
@@ -1839,6 +1897,7 @@ void Simulate::Sync() {
X(mpr_iterations);
X(disableflags);
X(enableflags);
X(disableactuator);
X(sdf_initpoints);
X(sdf_iterations);
@@ -2387,6 +2446,16 @@ void Simulate::Render() {
pending_.ui_update_joint = false;
}
if (pending_.ui_remake_ctrl) {
if (this->ui1_enable && this->ui1.sect[SECT_CONTROL].state) {
this->ui1.nsect = SECT_CONTROL;
MakeControlSection(this, this->ui1.sect[SECT_CONTROL].state);
this->ui1.nsect = NSECT1;
UiModify(&this->ui1, &this->uistate, &this->platform_ui->mjr_context());
}
pending_.ui_remake_ctrl = false;
}
if (pending_.ui_update_ctrl) {
if (this->ui1_enable && this->ui1.sect[SECT_CONTROL].state) {
mjui_update(SECT_CONTROL, -1, &this->ui1, &this->uistate, &this->platform_ui->mjr_context());

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