From 8f9c690c8593df7531597aa1a3589febb731ff0a Mon Sep 17 00:00:00 2001 From: Erik Frey Date: Wed, 18 Oct 2023 03:09:54 +0100 Subject: [PATCH] Add MJX and bump version number to 3.0.0. Co-authored-by: Baruch Tabanpour PiperOrigin-RevId: 574327508 Change-Id: Ia9b62fbc929c6869dfcec87636b2e10d405a1060 --- CMakeLists.txt | 2 +- dist/mujoco.rc | 8 +- dist/simulate.rc | 8 +- doc/APIreference/APIglobals.rst | 2 +- doc/changelog.rst | 101 +- doc/images/mjx/mujoco_vs_mjx_large_scene.png | Bin 0 -> 26229 bytes doc/index.rst | 1 + doc/mjx.rst | 312 + doc/unity.rst | 4 +- include/mujoco/mujoco.h | 2 +- mjx/MANIFEST.in | 2 + mjx/README.md | 54 + mjx/cuda_requirements.txt | 17 + mjx/mujoco/mjx/__init__.py | 34 + mjx/mujoco/mjx/_src/__init__.py | 14 + mjx/mujoco/mjx/_src/collision_base.py | 64 + mjx/mujoco/mjx/_src/collision_convex.py | 724 + mjx/mujoco/mjx/_src/collision_driver.py | 373 + mjx/mujoco/mjx/_src/collision_driver_test.py | 499 + mjx/mujoco/mjx/_src/collision_primitive.py | 126 + mjx/mujoco/mjx/_src/constraint.py | 340 + mjx/mujoco/mjx/_src/constraint_test.py | 144 + mjx/mujoco/mjx/_src/dataclasses.py | 151 + mjx/mujoco/mjx/_src/device.py | 282 + mjx/mujoco/mjx/_src/device_test.py | 184 + mjx/mujoco/mjx/_src/forward.py | 356 + mjx/mujoco/mjx/_src/forward_test.py | 125 + mjx/mujoco/mjx/_src/io.py | 105 + mjx/mujoco/mjx/_src/io_test.py | 42 + mjx/mujoco/mjx/_src/math.py | 356 + mjx/mujoco/mjx/_src/math_test.py | 202 + mjx/mujoco/mjx/_src/mesh.py | 232 + mjx/mujoco/mjx/_src/mesh_test.py | 119 + mjx/mujoco/mjx/_src/passive.py | 141 + mjx/mujoco/mjx/_src/passive_test.py | 114 + mjx/mujoco/mjx/_src/scan.py | 432 + mjx/mujoco/mjx/_src/scan_test.py | 227 + mjx/mujoco/mjx/_src/smooth.py | 466 + mjx/mujoco/mjx/_src/smooth_test.py | 166 + mjx/mujoco/mjx/_src/solver.py | 356 + mjx/mujoco/mjx/_src/solver_test.py | 130 + mjx/mujoco/mjx/_src/support.py | 81 + mjx/mujoco/mjx/_src/support_test.py | 81 + mjx/mujoco/mjx/_src/test_util.py | 340 + mjx/mujoco/mjx/_src/test_util_test.py | 38 + mjx/mujoco/mjx/_src/types.py | 660 + mjx/mujoco/mjx/benchmark/__init__.py | 14 + mjx/mujoco/mjx/benchmark/benchmark.py | 116 + .../mjx/benchmark/model/barkour_v0/README.md | 35 + .../model/barkour_v0/assets/abduction.stl | Bin 0 -> 429334 bytes .../barkour_v0/assets/barkour_v0_mjx.xml | 243 + .../model/barkour_v0/assets/body.stl | Bin 0 -> 1169584 bytes .../model/barkour_v0/assets/foot.stl | Bin 0 -> 23384 bytes .../model/barkour_v0/assets/handle.stl | Bin 0 -> 739084 bytes .../model/barkour_v0/assets/head.stl | Bin 0 -> 267584 bytes .../model/barkour_v0/assets/head_mount.stl | Bin 0 -> 163984 bytes .../model/barkour_v0/assets/joystick.gif | Bin 0 -> 4161040 bytes .../barkour_v0/assets/lower_leg_1to1.stl | Bin 0 -> 523084 bytes .../model/barkour_v0/assets/powercable.stl | Bin 0 -> 1199984 bytes .../model/barkour_v0/assets/upper_left_1.stl | Bin 0 -> 530834 bytes .../model/barkour_v0/assets/upper_left_2.stl | Bin 0 -> 628484 bytes .../model/barkour_v0/assets/upper_left_3.stl | Bin 0 -> 1500284 bytes .../model/barkour_v0/assets/upper_right_1.stl | Bin 0 -> 1106034 bytes .../model/barkour_v0/assets/upper_right_2.stl | Bin 0 -> 670284 bytes .../model/barkour_v0/assets/upper_right_3.stl | Bin 0 -> 1521284 bytes .../mjx/benchmark/model/humanoid/README.md | 26 + .../mjx/benchmark/model/humanoid/humanoid.png | Bin 0 -> 397686 bytes .../mjx/benchmark/model/humanoid/humanoid.xml | 256 + .../mjx/benchmark/model/shadow_hand/README.md | 28 + .../model/shadow_hand/assets/f_distal_pst.obj | 10772 +++++ .../model/shadow_hand/assets/f_knuckle.obj | 1200 + .../model/shadow_hand/assets/f_middle.obj | 908 + .../model/shadow_hand/assets/f_proximal.obj | 2206 + .../model/shadow_hand/assets/forearm_0.obj | 2220 + .../model/shadow_hand/assets/forearm_1.obj | 37554 ++++++++++++++++ .../shadow_hand/assets/forearm_collision.obj | 924 + .../shadow_hand/assets/lf_metacarpal.obj | 2996 ++ .../shadow_hand/assets/mounting_plate.obj | 3936 ++ .../model/shadow_hand/assets/palm.obj | 23598 ++++++++++ .../shadow_hand/assets/th_distal_pst.obj | 9528 ++++ .../model/shadow_hand/assets/th_middle.obj | 2544 ++ .../model/shadow_hand/assets/th_proximal.obj | 840 + .../model/shadow_hand/assets/wrist.obj | 3708 ++ .../model/shadow_hand/right_hand.xml | 327 + .../model/shadow_hand/scene_right.xml | 31 + .../model/shadow_hand/shadow_hand.png | Bin 0 -> 1315458 bytes .../integration_test/collision_driver_test.py | 91 + .../mjx/integration_test/forward_test.py | 71 + .../mjx/integration_test/smooth_test.py | 74 + mjx/mujoco/mjx/test_data/ant.xml | 82 + mjx/mujoco/mjx/test_data/ball_pendulum.xml | 23 + mjx/mujoco/mjx/test_data/cherry_pendulum.xml | 20 + mjx/mujoco/mjx/test_data/convex.xml | 46 + mjx/mujoco/mjx/test_data/humanoid.xml | 109 + .../mjx/test_data/meshes/dodecahedron.stl | Bin 0 -> 1884 bytes mjx/mujoco/mjx/test_data/meshes/pyramid.stl | Bin 0 -> 384 bytes .../mjx/test_data/meshes/tetrahedron.stl | Bin 0 -> 284 bytes .../mjx/test_data/mixed_joint_pendulum.xml | 23 + mjx/mujoco/mjx/test_data/single_pendulum.xml | 9 + mjx/mujoco/mjx/test_data/slide_pendulum.xml | 20 + mjx/mujoco/mjx/test_data/triple_pendulum.xml | 21 + .../mjx/test_data/triple_pendulum_free.xml | 24 + mjx/mujoco/mjx/test_data/weld.xml | 58 + mjx/mujoco/mjx/viewer.py | 84 + mjx/pyproject.toml | 43 + mjx/requirements.txt | 161 + mjx/tutorial.ipynb | 1271 + python/build_requirements.txt | 20 +- python/mujoco/CMakeLists.txt | 4 +- python/mujoco/mjpython/Info.plist | 8 +- python/pyproject.toml | 6 +- sample/CMakeLists.txt | 2 +- simulate/CMakeLists.txt | 2 +- src/engine/engine_support.c | 4 +- .../Editor/Bindings/MujocoBinaryRetriever.cs | 4 +- unity/Runtime/Bindings/MjBindings.cs | 2 +- unity/package.json | 2 +- 117 files changed, 114430 insertions(+), 81 deletions(-) create mode 100644 doc/images/mjx/mujoco_vs_mjx_large_scene.png create mode 100644 doc/mjx.rst create mode 100644 mjx/MANIFEST.in create mode 100644 mjx/README.md create mode 100644 mjx/cuda_requirements.txt create mode 100644 mjx/mujoco/mjx/__init__.py create mode 100644 mjx/mujoco/mjx/_src/__init__.py create mode 100644 mjx/mujoco/mjx/_src/collision_base.py create mode 100644 mjx/mujoco/mjx/_src/collision_convex.py create mode 100644 mjx/mujoco/mjx/_src/collision_driver.py create mode 100644 mjx/mujoco/mjx/_src/collision_driver_test.py create mode 100644 mjx/mujoco/mjx/_src/collision_primitive.py create mode 100644 mjx/mujoco/mjx/_src/constraint.py create mode 100644 mjx/mujoco/mjx/_src/constraint_test.py create mode 100644 mjx/mujoco/mjx/_src/dataclasses.py create mode 100644 mjx/mujoco/mjx/_src/device.py create mode 100644 mjx/mujoco/mjx/_src/device_test.py create mode 100644 mjx/mujoco/mjx/_src/forward.py create mode 100644 mjx/mujoco/mjx/_src/forward_test.py create mode 100644 mjx/mujoco/mjx/_src/io.py create mode 100644 mjx/mujoco/mjx/_src/io_test.py create mode 100644 mjx/mujoco/mjx/_src/math.py create mode 100644 mjx/mujoco/mjx/_src/math_test.py create mode 100644 mjx/mujoco/mjx/_src/mesh.py create mode 100644 mjx/mujoco/mjx/_src/mesh_test.py create mode 100644 mjx/mujoco/mjx/_src/passive.py create mode 100644 mjx/mujoco/mjx/_src/passive_test.py create mode 100644 mjx/mujoco/mjx/_src/scan.py create mode 100644 mjx/mujoco/mjx/_src/scan_test.py create mode 100644 mjx/mujoco/mjx/_src/smooth.py create mode 100644 mjx/mujoco/mjx/_src/smooth_test.py create mode 100644 mjx/mujoco/mjx/_src/solver.py create mode 100644 mjx/mujoco/mjx/_src/solver_test.py create mode 100644 mjx/mujoco/mjx/_src/support.py create mode 100644 mjx/mujoco/mjx/_src/support_test.py create mode 100644 mjx/mujoco/mjx/_src/test_util.py create mode 100644 mjx/mujoco/mjx/_src/test_util_test.py create mode 100644 mjx/mujoco/mjx/_src/types.py create mode 100644 mjx/mujoco/mjx/benchmark/__init__.py create mode 100644 mjx/mujoco/mjx/benchmark/benchmark.py create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/abduction.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/body.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/foot.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/handle.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head_mount.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/joystick.gif create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/lower_leg_1to1.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/powercable.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_1.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_2.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_3.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_1.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_2.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_3.stl create mode 100644 mjx/mujoco/mjx/benchmark/model/humanoid/README.md create mode 100644 mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.png create mode 100644 mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.xml create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/README.md create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_distal_pst.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_knuckle.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_middle.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_proximal.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_0.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_1.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_collision.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/lf_metacarpal.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/mounting_plate.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/palm.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_distal_pst.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_middle.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_proximal.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/wrist.obj create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/right_hand.xml create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/scene_right.xml create mode 100644 mjx/mujoco/mjx/benchmark/model/shadow_hand/shadow_hand.png create mode 100644 mjx/mujoco/mjx/integration_test/collision_driver_test.py create mode 100644 mjx/mujoco/mjx/integration_test/forward_test.py create mode 100644 mjx/mujoco/mjx/integration_test/smooth_test.py create mode 100644 mjx/mujoco/mjx/test_data/ant.xml create mode 100644 mjx/mujoco/mjx/test_data/ball_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/cherry_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/convex.xml create mode 100644 mjx/mujoco/mjx/test_data/humanoid.xml create mode 100644 mjx/mujoco/mjx/test_data/meshes/dodecahedron.stl create mode 100644 mjx/mujoco/mjx/test_data/meshes/pyramid.stl create mode 100644 mjx/mujoco/mjx/test_data/meshes/tetrahedron.stl create mode 100644 mjx/mujoco/mjx/test_data/mixed_joint_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/single_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/slide_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/triple_pendulum.xml create mode 100644 mjx/mujoco/mjx/test_data/triple_pendulum_free.xml create mode 100644 mjx/mujoco/mjx/test_data/weld.xml create mode 100644 mjx/mujoco/mjx/viewer.py create mode 100644 mjx/pyproject.toml create mode 100644 mjx/requirements.txt create mode 100644 mjx/tutorial.ipynb diff --git a/CMakeLists.txt b/CMakeLists.txt index bc6c5126..285250b6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON) project( mujoco - VERSION 2.3.8 + VERSION 3.0.0 DESCRIPTION "MuJoCo Physics Simulator" HOMEPAGE_URL "https://mujoco.org" ) diff --git a/dist/mujoco.rc b/dist/mujoco.rc index 6804238a..2f472f99 100644 --- a/dist/mujoco.rc +++ b/dist/mujoco.rc @@ -1,6 +1,6 @@ 1 VERSIONINFO -FILEVERSION 2,3,8,0 -PRODUCTVERSION 2,3,8,0 +FILEVERSION 3,0,0,0 +PRODUCTVERSION 3,0,0,0 FILEOS 0x4 FILETYPE 0x1 { @@ -9,9 +9,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "2.3.8" + VALUE "ProductVersion", "3.0.0" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "2.3.8" + VALUE "FileVersion", "3.0.0" VALUE "InternalName", "mujoco.dll" VALUE "OriginalFilename", "mujoco.dll" VALUE "CompanyName", "Google DeepMind" diff --git a/dist/simulate.rc b/dist/simulate.rc index a7a8ed35..7cfeae8c 100644 --- a/dist/simulate.rc +++ b/dist/simulate.rc @@ -1,8 +1,8 @@ MUJOCO ICON "mujoco.ico" 1 VERSIONINFO -FILEVERSION 2,3,8,0 -PRODUCTVERSION 2,3,8,0 +FILEVERSION 3,0,0,0 +PRODUCTVERSION 3,0,0,0 FILEOS 0x4 FILETYPE 0x1 { @@ -11,9 +11,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "2.3.8" + VALUE "ProductVersion", "3.0.0" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "2.3.8" + VALUE "FileVersion", "3.0.0" VALUE "InternalName", "simulate.exe" VALUE "OriginalFilename", "simulate.exe" VALUE "CompanyName", "Google DeepMind" diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 8a6502d1..f307715f 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - Maximum number of UI rectangles. Defined in `mjui.h `_. * - ``mjVERSION_HEADER`` - - 238 + - 300 - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning but for the compiled library. diff --git a/doc/changelog.rst b/doc/changelog.rst index d80ab46f..183dc496 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,17 +2,28 @@ Changelog ========= -Upcoming version (not yet released) ------------------------------------ +Version 3.0.0 (October 18, 2023) +-------------------------------- New features ^^^^^^^^^^^^ +1. Added simulation on GPU and TPU via the new :doc:`mjx` (MJX) Python module. Python users can now + natively run MuJoCo simulations at millions of steps per second on Google TPU or their own accelerator hardware. + + - MJX is designed to work with on-device reinforcement learning algorithms. This Colab notebook demonstrates using + MJX along with reinforcement learning to train humanoid and quadruped robots to locomote: |colab| + - The MJX API is compatible with MuJoCo but is missing some features in this release. See the outline of + :ref:`MJX feature parity ` for more details. + +.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg + :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb + .. youtube:: QewlEqIZi1o :align: right :width: 240px -1. Added new signed distance field (SDF) collision primitive. SDFs can take any shape and are not constrained to be +2. Added new signed distance field (SDF) collision primitive. SDFs can take any shape and are not constrained to be convex. Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent. - Added new SDF plugin for defining implicit geometries. The plugin must define methods computing an SDF and its @@ -22,7 +33,7 @@ New features :align: right :width: 240px -2. Added new low-level model element called ``flex``, used to define deformable objects. These +3. Added new low-level model element called ``flex``, used to define deformable objects. These `simplicial complexes `__ can be of dimension 1, 2 or 3, corresponding to stretchable lines, triangles or tetrahedra. Two new MJCF elements are used to define flexes. The top-level :ref:`deformable` section contains the low-level flex definition. @@ -36,19 +47,19 @@ New features functionality is currently available both via :ref:`deformable` and :ref:`composite`, and both are modifiable by the first-party `elasticity plugins `__. We expect some of - this functionallity to be unified in the future. + this functionality to be unified in the future. .. youtube:: Vc1tq0fFvQA :align: right :width: 240px -3. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints +4. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints and degrees-of-freedom that do not interact. The only solver which currently supports islands is :ref:`CG`. Island discovery can be activated using a new :ref:`enable flag`. If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island, - see video. Island discovery is currently disabled for models that have deformable objects (see prevous item). + see video. Island discovery is currently disabled for models that have deformable objects (see previous item). -4. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine +5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded: - Island constraint resolution, if island discovery is :ref:`enabled` and the @@ -60,7 +71,7 @@ New features Engine-internal threading is a work in progress and currently only available in first-party code via the :ref:`testspeed` utility, exposed with the ``npoolthread`` flag. -5. Added capability to initialize :ref:`composite` particles from OBJ files. Fixes :github:issue:`642` +6. Added capability to initialize :ref:`composite` particles from OBJ files. Fixes :github:issue:`642` and :github:issue:`674`. General @@ -69,32 +80,32 @@ General .. admonition:: Breaking API changes :class: attention - 6. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``. + 7. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``. **Migration:** These macros have been replaced by new functions :ref:`mj_markStack` and :ref:`mj_freeStack`. These functions manage the :ref:`mjData stack` in a fully encapsulated way (i.e., without introducing a local variable at the call site). - 7. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an + 8. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an arbitrary number of bytes and has an additional argument for specifying the alignment of the returned pointer. **Migration:** The functionality for allocating ``mjtNum`` arrays is now available via :ref:`mj_stackAllocNum`. - 8. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``, + 9. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``, and ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s. - 9. Changed :ref:`mjData.solver`, the array used to collect solver diagnostic information. - This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix. - Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. - If the solver does not use islands, only row 0 is filled. + 10. Changed :ref:`mjData.solver`, the array used to collect solver diagnostic information. + This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix. + Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. + If the solver does not use islands, only row 0 is filled. - - The new constant :ref:`mjNISLAND` was set to 20. - - :ref:`mjNSOLVER` was reduced from 1000 to 200. - - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. - - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer - vectors of length ``mjNISLAND``. + - The new constant :ref:`mjNISLAND` was set to 20. + - :ref:`mjNSOLVER` was reduced from 1000 to 200. + - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. + - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer + vectors of length ``mjNISLAND``. - 10. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute. + 11. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute. **Migration:** @@ -105,39 +116,39 @@ General :ref:`conaffinity` attributes in the model and then setting them globally to ``0`` using |br| `` ``. - 11. Removed the :at:`rope` and :at:`cloth` composite objects. + 12. Removed the :at:`rope` and :at:`cloth` composite objects. **Migration:** Users should use the :at:`cable` and :at:`shell` elasticity plugins. - 12. Added :ref:`mjData.eq_active` user input variable, for enabling/disabling the state of equality + 13. Added :ref:`mjData.eq_active` user input variable, for enabling/disabling the state of equality constraints. Renamed ``mjModel.eq_active`` to :ref:`mjModel.eq_active0`, which now has the semantic of "initial value of ``mjData.eq_active``". Fixes :github:issue:`876`. **Migration:** Replace uses of ``mjModel.eq_active`` with ``mjData.eq_active``. - 13. Changed the default of :ref:`autolimits` from "false" to "true". This is a minor breaking + 14. Changed the default of :ref:`autolimits` from "false" to "true". This is a minor breaking change. The potential breakage applies to models which have elements with "range" defined and "limited" not set. Such models cannot be loaded since version 2.2.2 (July 2022). -14. Added a new :ref:`dyntype`, ``filterexact``, which updates first-order filter states with +15. Added a new :ref:`dyntype`, ``filterexact``, which updates first-order filter states with the exact formula rather than with Euler integration. -15. Added an actuator attribute, :ref:`actearly`, which uses semi-implicit integration for +16. Added an actuator attribute, :ref:`actearly`, which uses semi-implicit integration for actuator forces: using the next step's actuator state to compute the current actuator forces. -16. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to +17. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to :ref:`actuatorfrcrange` and :ref:`actuatorfrclimited`, respectively. -17. Added the flag :ref:`eulerdamp`, which disables implicit integration of joint damping in the +18. Added the flag :ref:`eulerdamp`, which disables implicit integration of joint damping in the Euler integrator. See the :ref:`Numerical Integration` section for more details. -18. Added the flag :ref:`invdiscrete`, which enables discrete-time inverse dynamics for all +19. Added the flag :ref:`invdiscrete`, which enables discrete-time inverse dynamics for all :ref:`integrators` other than ``RK4``. See the flag documentation for more details. -19. Added :ref:`ls_iterations` and :ref:`ls_tolerance` options for adjusting +20. Added :ref:`ls_iterations` and :ref:`ls_tolerance` options for adjusting linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning. -20. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to +21. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to mesh assets. Fixes :github:issue:`409`. -21. Added camera :ref:`resolution` attribute and :ref:`camprojection` +22. Added camera :ref:`resolution` attribute and :ref:`camprojection` sensor. If camera resolution is set to positive values, the camera projection sensor will report the location of a target site, projected onto the camera image, in pixel coordinates. -22. Added :ref:`camera` calibration attributes: +23. Added :ref:`camera` calibration attributes: - The new attributes are :ref:`resolution`, :ref:`focal`, :ref:`focalpixel`, :ref:`principal`, @@ -146,21 +157,21 @@ General attributes are specified. See the following `example model `__. - Note that these attributes only take effect for offline rendering and do not affect interactive visualisation. -23. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values +24. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values ``mjDEPTH_ZERONEAR`` and ``mjDEPTH_ZEROFAR``, which can be used to set the new ``readDepthMap`` attribute in :ref:`mjrContext` to control how the depth returned by :ref:`mjr_readPixels` is mapped from ``znear`` to ``zfar``. Contribution :github:pull:`978` by `Levi Burner `__. -24. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the +25. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the `WriteReadCompare `__ test. -25. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`. +26. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`. Python bindings ^^^^^^^^^^^^^^^ -26. Fixed :github:issue:`870` where calling ``update_scene`` with an invalid camera name used the default camera. -27. Added ``user_scn`` to the :ref:`passive viewer` handle, which allows users to add custom +27. Fixed :github:issue:`870` where calling ``update_scene`` with an invalid camera name used the default camera. +28. Added ``user_scn`` to the :ref:`passive viewer` handle, which allows users to add custom visualization geoms (:github:issue:`1023`). -28. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch`` +29. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch`` and ``viewer.launch_passive``, which allow users to launch a viewer with UI panels hidden. Simulate @@ -170,11 +181,11 @@ Simulate :align: right :width: 240px -29. Added **state history** mechanism to :ref:`simulate` and the managed +30. Added **state history** mechanism to :ref:`simulate` and the managed :ref:`Python viewer`. State history can be viewed by scrubbing the History slider and (more precisely) with the left and right arrow keys. See screen capture: -30. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by +31. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by `Levi Burner `__. Documentation @@ -184,17 +195,17 @@ Documentation :align: right :width: 240px -31. Added :doc:`detailed documentation ` of fluid force modeling, and an illustrative example model +32. Added :doc:`detailed documentation ` of fluid force modeling, and an illustrative example model showing `tumbling cards `__ using the ellipsoid-based fluid model. Bug fixes ^^^^^^^^^ -32. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of +33. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of midphase collision trees. -33. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints. +34. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints. Version 2.3.7 (July 20, 2023) diff --git a/doc/images/mjx/mujoco_vs_mjx_large_scene.png b/doc/images/mjx/mujoco_vs_mjx_large_scene.png new file mode 100644 index 0000000000000000000000000000000000000000..d9ef4b543143031b196c8fb9a8eb21fccc4010af GIT binary patch literal 26229 zcmce;1yojT_cwSW0xBg2A*moG(w!wb#Z(k5g3SX$?(S7Vn@nqV zUZIvZt=MbnE;YZiRBKpj=!t#X;gZPP_TcqBcR4}JE{vCyWy}>vO=Ukm*K*X0(+WwD z(Gmt(FqN>~B4OOPVk4%=^^GW>`Bq#BZ{{4M*?2?^>zjeMQKfQr=gu?f5-0ez2)qr>SlGC@N{Pvy^YNFTUH`-3+xM~ zp0tD*>J<6!ebuK3_~g8`q`Dmn#Y~O-LUaBg;s_sN*-OdZ$6B~RK*UH}7xr)pJ|eNd zuVyc5Wnp1xX^#@MHGE=kXh7xk)ZT=61rptN=U8uJ5 z9kVQ_t}qa$||h1ZeIaIwWm>e@n|s?JS1I{(A9n#9$y3Z9lQ zc~@8d36EWJI=Q&XwrCFG5b-nU=s$-F=tTYT8e3bJGJevCY#$tL9OMm|R~{^8#tFK- z_B=jd^EmjmCG0d-?qEKY{~{pZ+-l9KCk|%43#EY7xZnkB?ERj^ljk)xHN_`pYU=7F z8yg!94Gmpob_+U3hgwyxToYcWLY;zJIJ6N`w=VL!@396E-Y);)^gQ6G&~uZZ1sw}d z)7r(;vzo(R^E@LXqw_+%lUk|u^_MSSUb%X8YH11kj+`7`qQ?#iOU)5`V`HPLq5nQ# z|ECrMy{2FiXJ_ZnRujTIYZDWbGfGNIgZYNoCnqPPE(`HQIN>)IG|Ft}Tn_tKxHuh_ z`em;xiFw)xdVXv2aCiUu?OSqI)^a8Folq9Lh?eK=Mu+!_GCzLwPhWz<^~ZJBo#mf6;i$U0y0#dup!e^o z0)v8vsyz-L_NjZWZu8myP$+j;@@v$4vcEakL3vtU&i|H4{hp!W&3b?QPmfD1hhNWS zM8?Ei=%VSd6%&gYmgm1QRN-W`vpUxL3>}MD(Am0bj!r%y_5FKq#T4;8r?qj{Ru;}r z7RLwM?#J5$G6{mwcYQInHq#msgxm(J+<4&P2FsGm#Ke-PC&wi01_7;4QZq6bs}ENN zP}9@X>l+(SKerpT$FQ09(m%@8mQ`0Dx>=(w%0!6mb-bj*7A+%q{Hw>=xiqI5#+u-^ z%7^^TK~ExUY`e;%NBK50;-O@G_iSv~C@3h-Vqmc8)}jt}*G>y^I-A&uab#k-L%dG5 z5>r2X$g^9J^TQ$e0*65A(Idj^*RKat2p#0>EWq|H59ayKG>0i`Ylr&#V_6IrCRt5X zLwuGlGgTgKw!VpsWHas{LsrVsksA#Stu2l(0@l`?n-umzzcWGT3K7xF(vtKb&ja!f&Tt{ixKXfl@SR=MS}71ar@=LZ~|)a-dGOj zRekw*{ueC9SHZfO^MajJGmUg4NNHc(jwsoIT>i)({XwsdxW%Flm2-yVCJ zfS_)Ds{UzT#>1*+N@4eRjw{1u%W-oH3k9%g`2__G3=DP6&Hi0o3VQnb=W%hXmJi;( zd9%Jfpd(<`Rq0~)rM{ko&(@$Sz9BjJ2DjC?jE2TFLqo&8y}hm?GgZs#eIp+qpT)kV z-gIRoM$E@6qpRF@;dJ{XCQ??pZtKl{4c{IQYd6wg>d)@5$Z*-3$Dt4`ZCUsp^p%BG z5)%{i(v>S>Mn<$(uU_pad%h5=jrQwj&nGD#3H{!_zKuF8O6wn8*O3*3t)z`W_OCxj zt2If~zx4FvaO2mnvy+fkdwYAkiY*@ByLXRy&@qfc2%@Rkuu5Ch-Q7JPIGBck;bm{H z3Y5QJ2ZdLx#*h78k#D8h7!BkV6i}?&?^2M#c?lYF~H+#6e#g)y{-kB4h~T zZ*Lgf*XyvdbZRKPyu4r$g2Te_IgHw@&$~qFEp)_ZX;)oK>6SDxp^q1E`k0#f9Byeh z-Ebc6()fl!c~Qs;25(?wbVfl_Q|8Vcl!Jpq<>6`t9ET^|u|LPF1m@=EkaKHbV8FR| z0Lhf;aWPJ8ySReF#V1dm%+AeiwDU~7dh@2i7hm{?t>mdYvI0;!*0;BBa&lhrt-Aw9 zNymKye z%OfRWP`_-}s&>B@nGLvJKUpc6n1^#ICLuBN^CzU$3%4F+P3~>XFdx&wmf;KUp(0wu zAWg%Go1LHEgv((P4Jl#5i<|0Zrlv4rYreV_2))Yv`B)Vd6{uv!{Eo|CqdANvA3V6h z#~1CkRyA;QlBc*Dsig4zkIBhCwY7JWlauM0m<(XgY3S&-zEX$rqFafJW3E+h;q~?P z66hpmeEsN4pS8tha3~GTLxqkgwX({jIvf$14bydfS)7u>j@4;L z!F&pw*0hE9CecxrHOD)c5?0sh_3}Pjj`0f$3OfHt@;`@%=Xib&n~2!}GHNg>Psbe7 z(YJ38Z0+o#xy%U3`R%D#Snyd`SRnqyAs0Zz&z~HwIfZ7CL3U{jA{uZzo}Siwl%@7; z%+eUv4Nh}2iA?q0GaXJj|#_?y|FdVg;A0`jqFO8f`OqS zEa7V0na@w^F^zvjv+4WU9e?`xadX^lEl;oU!pUBX0C&bGTU(wr(4LCk4#5!3e>lsiHuwL4Hc;EB_3A6q{Ceywz5psK2h zgvT<@ZhQ*TNccvvG!x720J2r$$;Zi|>!#TNig@@(fGAZxTRL;83E!34J-i{HPGks=W`KObx5 z*l%0!hl}J_F0=miC7sil#YI!$qjk}Y%uJX8SzcZq^8|);6tf{D8^M#k8Rn6UK2rxw z@#fv9Po7-Qt=hTP97;}2OO2M9AMDQVPc!o%6kO#pr0pVBi;$2wfcJMk!M zWxUF*BKT2!#emylFT=sk%AwBwBotY6-QmKg%>JfrMjRZ*9SsKTd2YLwP)<=FK72q$ zxevW-svh4LyMKRbW=0fFGaRLknh|sE*RNl%@9a2+(#9VqbPA>L(bET{D}SF;GzTr~{4i2qVR#pbQe2E5UBe%t@>x#!Q6x|NK zu9>C~JCc>qz(BoPZ?rd2QCkrE18%FgOGc-0!b5l2o_slrsy^O%9OYohc}+_WkAoNo z@KB6Y1Z|(dYa1AK442tF+4<@&r*2%H(c0(dFKfzudur}cke~mxwKc_Vp(B`>?aPO^ z046mv<9FED8RNOei`bw^$zQd z1f_?Vo?fyY`LC6cD8R(r2d+`4+1q8tofO8_j$}ENnk&OaGqB!9ZSUBay#SnFhh&$L zk+GFyCt=a4)qMZ{{X(mWagyO{7bX{M-7zKBMWX6-twX7_mkWxmCMXJq+Vb?9oo7OL z2FvU|ZEY$;Hozfft>wG%CO%$PRu&h+9=9VTmfPZ5nBe9)`h(`&KEA%Bf`w!R z1XE;H$9Z!_)a>l}b<$FN!d9YdACkypfndqdUjF&>C-YhX zkP%fKkH`>G9+)@q(NC~g2+$eFYoDt3KaY>^6fLqsoPD@nkN+|(Y~gD-b#H&a9xHL_ zMic4Fz`*j^1AbF$OVS0;UEq8ESQO5vlK)yH^>K?JM9C;nWfcPTHDp zH=n4Ehp6;^_3G7@B@fYX69^JZA#b>T4ceJauNL>e;VK=iUR9HT!$yRoB!b6EbdVL= zxxky;oo^^j;WA6W&cOl1_aZ=4Aj$<$Hyqc-+aaOF^-$qa3gh(FEQf(i00|Z{nh!(+ zlCy!INXg03LT~|(we1wz)l~8z?GYJ`C+z83D(?|t-g(z5)Ga5L^PcJ_FIfx#FQ|AO zxfO#?6q<>oM0G3|C>~)*N=xfQ1xC^hfTRj0J2FNavSHg`TCE?e9FR;<9c410nAG?5 zL_kc~+S)Sj)yKSj+XN-Q6Lz(9wS0;B=Qi`BDAC;(}WYlFQ3uu8aC{?#(SN zEHuBr>x;k)0QO`;u5X|~D<~+`Lc%3s(`%5g7A7t5?AZSKDN#vH4Pgnu@|#UKZ9zQQ z95&0%15IS1vRWFU#@~B-&R@7N4ZJ*)XLE|Ta0Rc~u9L^Mze`XWn8&x)RuXP=wTRJ= zo%Qv;P%8QG3lc#2=@dDQgTN6!J+d?EOjwCi9q2lIR$!;KwYjOujH0}FaXn(pCnw#^ zn$KfQF*YrY-m?7XNAZEMC1^iZBY7Np4Of=*%m zuJ?CAyLlB8^8heFr@-_8P%ZIb;yR$EX+W7EJ@;j*03r*_tKekq=9tb^&UJ!Q)W~N~)<1-p9}7^JxsL zVO@5t*jgVe)9HXr0NB5ul)&1ZApWp^F!?^VT%xe7oSdAmjsmDWFfv|zE*Dosf`Yev8H`S|#bW++c}K`Dwls_IhV zD($WITrC&ct>P@~036qGw7(Tk*P&&8OiM=x;#9odEfB8u7er251lB5NHcO#a+bw06 zOs}oUJ2-GfMn-nDo*X%Ieg~l2|JZ6{)UH$U@#CdUt)Ut(uLVuXX6`=7cyiu0nrDxj zbB`ft^X-2sEiNv`StPC%r-#;@?$Jx>JF+8^8_=w4N8>V5Qr#@2UyBH7&|rFEHa2;W zjd%H(jVb`&E+mP@wl{%JvTr87`u;F&fjlY=Xca8g`=-GCla+u7MU$4(Ko z=}Lsn473d4A%%sUAWOs%sW_Sp<==!5a^D#;X09tKF79`6iU}gmH_!N;K2=m2G9G4n z7Z7=$Zltxz@{EZ|N}hX^85t&WviI83W#L?y2p$7)X@LH`5?{vn>b~i+^5U-(0QbtDL2k9s!Ma3z1b6^08?wjY~ zmQh@)*--rkt5^11p^iu@#@uN{y4eaH9I>yt&R@I!Ou^8TchV*ym-+; z{y?I}WUQPEYCp0rK)=fbMSy|i=d(kdc~(;+A`wRM0a=!z_jZTTplbm3b_ZFC`MA2U zuy-hmJh3Qmdj-_1XV0DiGb=l5R58{8>Rk(LK1eGK=G;9cRys#VN1#}x0~M_E#oFH) zu>xVHGOHky-1s0@E$@q==K*YcZn=Q}OU-ma&m1CKGCU!dO&al(?(EhQ9CH3QYc;0{ zsi`#D&QlmbF+@E*g<;hp`{tVsFfZtX1Y8P=-e9Q>H84os?ZrOOh7?s*{p)mfM@wzC zt3lq#=Y{Ixb$U4QbGRsT>&z5`-KH$63lOdO)3gK`c? zmMtjlWn@sP2keYh*oB3w>$)Yqk!M3r)?<*>Xw^rr=jk!cb)C}ROmTz#L zl*f{WmGu>{lkli0e<+;5w28*bOeU&@*x1<2$I2rSo8j|kR#3p`>FKFNL@0H1bfV&8 zY-XCmw2WV+)1BK!_L;@@x>$MP&Q+r(KxtWHS2&4<3gy?+koQNuBU2`oDb zP)mfGTR(3%U}y6>aR)syZ)Bx|ir4j*$pS7;usF5yN0}!Z(~V5pmG59}V8U}*Rdu>l z%*=QoY?+!e0$RjP7!$M{<%P}I1OCG6yiNr=#Cxc~@BmX&(}0i=2hZxJ{s9wCcKs$C zgyw?-P~o&j3rht0E5K5d@k+ju3D1hHLu*(mM19NBs(1^eyBj&9K|w7oEp_$v^>g2q zqY$b_#%KE+iibcAPxhl*umXtw1)|=Ps^)A>+pji`PQjku5X+Hb+|0gjpc`n>mZ&~@ zWc%xfB1~^-OwS5rGc_hc5W1)1lxHCG0%Zji0TeQ24UGrZ)~r3BOQ;;yBmV;vN64r z>Ioo?0l7f|I)H`vtdZR@vgq)m>lvw4Ww^jie^HM+b_0|~P(fh1A+TAT zvED%LgW}r)63sejeB%tj&Y%1EoB_a=o}GP@i76N!?h9IT_?tIC)cv;(JO!LqQ$c(i z92~@48iSMpI}b_m2H0MphWO!;BkKqz2MgTUzDyHYBmx2gSVY)7pt=Y;`}$RKwQ@_Y zgvFZ@INE4ZOMCl=_wP~i^75civ<8!KXda}1aR8A>4mJcBQV{5}V5$JP%iINUnDA?D zf)cKN@ZbS2Bp*2c$WrPxocmpjh2H??1LBBK->UG6*RL-DK$Qj2TFVhHCK)qiWeG+$ z3tD5GutyB2Gv$se4?sym0t=)9pFsV8z~g)cT!n>& z0o|i#_B{ai0ffx=m6dUDRKQYGUW*A0$+iOg@pGhv1@sn3apAx}yG>H%6NQ05*Uue9 zLXkiyIS62|mb^C8XRUX>3-)}Jwew7%9JobG`*A$_ctM+Yg5W(aCOQy@FB20JL#pHh zjZiUok<4}X?p;}1Ep>u2PRzp(WS6gmyTUv-wzkYz+g5=!L8-Fsk>z1?Tvoff#JB(Qf6h@6w1Z zglzTp+iaK$NF$ryZsl^8j{$qx95PNU+g5|(iACvN`;JZj96tUhFz(tCgkpn(gENzx zY;0_h?}Onm@OWoXKdc+fLQ<19&eQ_cgGr>RMi?i!_#r|)$Se*<1d6Y& zgeU4`KFU@rFv4}L7@J$WrJB>w*||pNbxfDw?1NNA!K3wh^O52ZkR`yUZ3SKQYiDPT zscEI_wrNz2diEp4)@U&(|FLD<5!aHYn6kHG5tXT$YXtTYlV<5fI6Mvp1U8!r z=fxVPidZu8`nXP9PVLM%N+_HFoY}2U-UaIN{XuKf-@O0=5aXCAB({OY!pV%^vH@A2 zYHDigXzlcb1qEj3T(<p%~tJzy4b^2jz zZ?#(^4!PV%D1i-LVMISg`t;bM5Q4z|0vxQD3 zntd*RMCcF8K7#$LpO{EQqR?OD=$T>bTQmg~6)|1iNy8$Wn|*^e&J}5C0e}bzNJ&T1 zxWN+n0c;F#W-z6QHn}R&LNQyD3nB;z6rT-S3BfyJBgO%BSIe;ipe1}-+E%pOn|>Lv zCFIWU08q9%eA~c!EH6Juw$Nli_74ntQOJLwIXnTN=IqS5it1g;t&!8#rT|^rF`9+p zhAlJsFd+fa)QZ=^acfHi-4<|;@8Fvt_JnJzENz6eH3NuBC#Te5f?o=s_Qd&XWM*W5 zpiB!oGFT0Oroafppb`)uwX(7TBudPtcb28<*Cmh|AY<^3WJ2OqtvR?V>WA|Lz$Ju= z)H@mH=T4_mOC#`{oSGBYOP4OGmD%28sa*GI-L$GXe!DyF76xPQ`QV2)Qew5Bb9z98 zpm2!WV33sI_6L9&BUp?|$%LScEj~+Dzu~dHr;}Ad@EH|6Q zFz0F+h}J5pX}W_mEnc zTQy27+x>7!_u#fax{}aaqgadq+@KH~Jn4(S*>G9D3miMtxbqAybOAtL!Dl?Xpa5*x zkMKk-Fg6iyQmfoP32@T(+X=7=KvT8_*2MuR>!WP!yyi}?vOJJ&5)LCA&`PJ*yiNrm zD}%U^4_FZDEEG6WVm@0Mu!QGXy!L~(%8*jZ3^8Bkk03;`-YivhqZFlcq5hBI4;K;{5( z8?l+88{*Ws=pqdfi*}oFqDK<~-9h_@`;K~@9;3ilW(S>OzB5rCA_J-@ zcy@RwaE~=L$*9GH-UQ zon^^LY6EodxhUFnGZD zFq#i(Y3CwZ&enpVO9EP$yzo8^*b;iZX^My&g+lbzckhTnyJ@d<*Bmz z1AGGJ#FF{fJ!L&jMof#?ZD%DBTIg}W3EtSK=iv&-1|Xrafz}nq^Sg0H+nPkds;s1h z2h(>dxwA^G0*Myer!nqgn^4DQI=xOfK&nL{u4JJ}w+|3hkeKt-n3%(Iv$CMT;#Z7$ zjvL;oTF}UT1d9L&KB~`B6vPs6pUaH|Ug#J5*6DT;COa^zmH}i87T2m39jVm31*Gg8 zHg!1M>pd9gq5}hVcfV#=>s!*cH&69D^G5yub}!Mr!t$QNd6c z8K2(n_I42+4b32=-(`7`<3c#qt2=<+E82|;0OXG(dLBG=d-yO?==$}#)ukn~&aYoT zkE}Qzf|b1jy6_I@o+F)ceEO6eM!_EaS?T~)Jk?hocmVm@0Y5SA;=n855N@un0A_qKlhlc%DTng0jab^o+0Egm`+4&Y%5^oxDTmWc_+8_-njP=PTAga`DCoI7_;_vC0F71fu6 z07T6g+M1&<&3&luNXCRdAqBP^0w{M#g@cBdt25W^CTr0Univ0uet7==NTmR0y#{i6 z{-cTTKmPFlvZD~W`G09-{Qu<kQ4=aAEi(bB}X^|eZET14^spBa-jzfJUa z7-L-g)~bT@*Cpxg-)hS(BQIQsvy@ErQ#g031f{E05m^4^%vI!DVt2m_3D0~)>9&fb zD@y(TlGLt?nAo(CCp#)~BT=`tLrL+^?Z@;J?CG3c^;m8-UaC>TK(1AMWgnJq_hX2) zWvJ{mmw&{=D}UX=xq0ka-hJuiA}+N7)12Qp-zPOz#$p8c6c(z#b2?pWG&335V|vN* zX2GY|KJJjCK4Lspwh;xd^2PdGm#>bOvSLeK_NiUNl}INxxwuB0O%zj=zp2TH?|(2 zS!qrd9sP5eC05H+%KEa%MP9dXTUgpZ_q?+EjxMokM0TZ{lj&Ae2G#Fthg1leH6;Na zuJ7(HG?7}(0!QvJgsx4f!%#NB3?Oa}qjX}1v~ZJ~o3LOA66M)g9xKPEyU=U$2^`_Q z!$S_Uex|HPW)X!WYVDYKjko5FP>cgPsX19HLQ=nH$T9$8?tVSxNiF34FVKATG9+X! zRG_2dR$E(}0XQ;ki6V(WZuRu^hyhchK(L3F6fgkpLc<@FE&Zdt4QN2UNW(42-cBtm+M)R_oqEUaM9D&7}2gSJ;hoCITm4J0eb7^Q#>e}Wze+6?#>NUNuN7PSA} zxPiD-V8(O&aQxjr`hfJdBTYVr{ghQKh0|6fqJDWa^;2B(-dpn zGwaJ^_^8(5hm3y`bKoHvWJ?OUrhM1v3x8jhJ4nD;@q``)aXq+X&WH#(0Jh!SX%pjb zJ#Fnsw*9I1yRclLX%uSDp?p3dixixa3vMpazXydpDpAUK;X(gZLWTW}!@r;9Cg}F7 z&aLh2rF#y8BDkWM|&Ju#O1LU+nJsJ-ZMDSM`G6m zm9>$W_NQtvgq({IX^YI_vxk+9{AWy+RwU8|o&)H*FI8V5H3viL>Dk260=t_Q(Qr82 zbpD>Hd_hvZ2{$hxYH9ePg6gM#1|;}%SySql{{k6}43U zus1R?gC6z_sf0G8D<4-M<0K>9k8KA zzW+VdSiyT0dI|PBws-os&C=Q>Re$3Fk`;Gnkr9SZ71}yJ0Le*zkX&!*Bj}T!~!Pcu{_B-S%tzw@W+q0k;Rhbn3T7qi?v|exsR>D=7%vAp>+YFYN zEwK_aLta#A{W#dxQ6iQ7*E3_&{?%&bl z!@k^qp;Ec!_s==#tUlhLa3DqX?_C{iZT%+zc;_Ds-@3)c$nfV>35TtAcd$NHuekL1 zMCvaL(hnmt``Yb)4$H@mCs2)+#FnC#(%X6eBcy|~&!E!OD)0QW$PAy! zg>4YXbArnyb@&@~!MWbulc^+EZk9u-RM#l{b8TJcgbSwA_HIS`c#+n>Q))Jy(DD^KC?)b>G>GaIsuXiw047=$YDOn3LUkpns1@q02+&p3ZkLS3dG6f%~6o> z5}>;s43*9-^#U+EWWTRoZ8oqS=g;HbzSs{)j4lW=Xy0X;o}M1?Cqr8&fEGpzd+={9 zbbbP{3;1VK2q~6#yAd_0HPFD1wDz8zuijnq#7)1KYYrL?2_?7@KL??TXnv+v^)VBWqG(L2s|~U zaTTPwL z8zIl!^qti?*(7q^T}4`FG*1%#;>ZYfgcFHyWYFG`l{8@Ii<0;aLFzCQNv~c7;Diw! z9sT>JOCKMbIM!c8F}04XS`iWugeOKyDD#wh96BQHx!?_fhn?7(55`WLQ9*>h_v!d{7Ddu3=Zae`drXU7 zg9T%0nN$4AObiwazz8&<1BcSwg|QHa@b4-C9mY229%12;eEat8Ceqf5w0-O7Ot`js z{NdVOng4@p;|dL!7U`9YfcsE^Fka|W>q@GR?4|}kaHm##(b3)4r3X`EoC6!juho=Y zKp@uJ+Z%HrNi!bQsI;ss8hU!k#?mHmvXQRRsHmua@DC|X&C#YbSd66rNaNE>0I>mi z+=Le)5OEh2{J5~(%uMS4GYZ)aR{p<2AzK(Ngts*`+2WJucSz!fHgMlWP33J;yRUiu z$RBQPZ*LqJ*bU5hrkMbJIPg>0k$V@@7fnQ@##O>n4A!P<#9i6+yFQ72Q}7fxQ_!0C z`rW&)Kw+g@-`Lv(<+aD19TG)ZYWK8UqKlDovERb=^p$<^Mg9{r zFqy$agjYs@ev3o<&Hoer-9|zKl5#s3VPGC#2N!35YXM#JQH3gY;>!}>?trRdzi&?C z3Z1{hG7izB9{Od9?k};;cm)KHf7S+<;lBX*9MKVkn;+fm ze(I`9koZ!W?OkUAm2%Q=zPbOBK~np^OkI3@iCm%O=p75A@_%dp%MyL_%je#gz5Y|l zLe!G7(UDTFUSxNROX}r6Wvh~QG&EiLe_W`ZZlvQ*D~IuZzE(KH(c<5dg`gnEgd$vf*n%ut&*ry@yI_C+&Yu?}RiPHbV1*>*#k%fznwy{?g13hG^eFjWnI~7+WzN0G24^@yI2jkfe~n8sPC`Zgh{An&StRSx0}8Kzz}<1BCx5UKbBFOa z`;NB#%b3k8mqj!iqMMEMSy+^9{`%W2ZYHuY7pqAjOLi0&)@8n{JU)vSI+?tGR+Y$z z_1cEluk$gT9@k>H(BIESm%XXk!M`d2-@ZaKH(e<{CSLnk+N6Ki6~&auyzndQ=isL8 zL+E7&EEJB-;cy|TQRKweY*3MyJ9ba}O*RIi*oyJazF(&7 zR6LmOVGRBqtzR3tNXbMqc|>Xv-5`V8^1br`uW1kd?}1%-eS_#HXZ9JBlV3r59Vg_d z;xA|9PIoMs65Ug$|3D_qxc$%ylkIQ!WRkSd=)LTOZ+gY~RwZjT{)ziCR){~7?Zn2t zQ9`Uc+a)9PF^;5thxTp=YJiakv2vREP3X%{o$gVxIXml>XwdK1_@_DV%Ko1EU8ULS zN~`!gtvnOW1x_L;#;D5rXbEHGC@-adAJCOkc(m~N-C8A4?!kUw1ba zlG&V+p&(RciHn-kXuNyB%Mly!}S=JXW zJ1g?7szg&BOJ`9T)d#^nyRsjQPPP9Y*B9Uvdm<+iN0e?Cg2l~0UEtfQ!PisI6GP@z z*6(WooD}I^NBUskog>7Q0(Dg&s=2xOG6~7cA#)=zI^pH6(K04r#)E@uP{iSz=Htt* zE${vdj>04Tx8Y*Bx%~@65YoVm5fAs1TvXPpZcM*pM(WKqM?F9db0BYj% z^TzmwM*mOj2RF8*|BXhJbF}MhACx+Ggs~@6h3&nlHNxkyMy_IUsPTKkXZ{k5j*i#7 zI-z!y&Aok$;`Kb@1$iFYSGGVK`~OBTdf_fMh#-f~2WXr7tB*Xnts37#kpCOQB(fPVWb+VxQ% zpb;e{BV%zP_ikptv)~5P&Q(gK&t}cwF?& zvZSJohsa{fcH3%8RqWN#Zp~>sHcKVaXa#k6?zF^dZA0-5R!yy!-HV<#@ zig`DvlD{WY(bC*G^NjfQt3%p~0y{phHUmeVq30JBq&5#dg|~)=7F>)Rx0aXZOFA_P zQSkn1H~ifG4DX%Su?ABtCEs4!J$yop3D#z)S(Ok*x{%ur&QOTV!-rsbh+gq}A8|kz zhps9^Hhnr_`n8Db9FxbTkkf-Rw^fhM`n9&pKU@3JTCXs%>c4qd1@9F2P<`guOZaqs zb2zd0uF`NWb*jYBR`N!iT$2IZ{3e_ORRU_zbnN# z^sDdWLvv**v>vSRhc)|-j*iY~SaBxP>wz;a>PUTfJv4mTXq1;YN;|Z8uJg<3xWKyk z30>vGtLrDtXLtkYHX|m_&TJ@&5`~4&x+VCHyPb&fPl?LmiyW~c?|~t2q2@EqOxF1# z4RV~a^&vdt0SK6emi`yX^lFvPY~R0sw+6jRA6#N^sJ?(NlSF?3UM39!vriw|AW+_q zt*Wj<+wy(*#{ptuV>dt!Oe}?WE3db1+GIz`Y2UkdqY_{E%$f)O86T8|7e>r{yZ9k# zZk`<^3ArBsHGw_zrM(x)Q#omFR zOjKap*#vDBA<$WdjNiq@#qm{EPEHdzkHK%>5^B24MI!BvD6#?9pR)_qMwqESc_xQj zm(nJdpRaXFzmJT%v6I7hQMqRjsmmg-E_?1%UfYuk)EZGFzKcpsP`9#J%(K#8$JfNVp{W!sD>|m!)o=Xlp(bcY@`;JyY@8!vxPw zd?Om-L+*>hl)7d;8YR>G+4ya}OR)==_m1j`&%D1ET9J=2KW|QW zw$NXJdMA=Q$Z1S?v{qH$0Y-@4D;oat!I4$*FQ*6SsaXW9XV3~yJlQl7kA*HyV!Q=? zftcoH=oh|ycbvm|Y-=5R5>)MJ|LNVHBffL1N3Y{W#l#w@r0n0|@^RsgzX&4WIYdL< zu^zzVJf5~Z#IXX908dza*tX3>Jr%}MtLS_U^?v1rao0gZZqmTO@9lvHxGpdyJbz9X zRp-%nibCxToTAHad4<-#4Ou3GJ77vBkDsT&uiRzvb{fACFz#%E#n`?1m1XL1#MTiH ziia;vcXt^}&9(D2UU(fG9cZY+)3vK6>YPLxbHaaJptv@B#3t!HEx$f$Pl}@Y?7$I2 z*!Jwyx9IODqy$whEd-qQ&AejIe@~e%u@q5rcf`)?Il;NlNDL9v$j1@&-`ZPa@fm!a zS-^owFn~eyWKg+vlK-?B>A2YAh%FFYN1uy?X4=u=UB4;C|^!VFK()Aq6eJS*Vt=AHpZRp=Vtx8Z-Ct9h=aAX>-4Cb+sH@j$bietXKVF>a8*~aXWdMHdiqfXHS?KzW+*s z@N>lP;f9T^S=y(Rq3@N7H8$U5tRBrx>a)__vc!#bGFxvtN1S*&iC3aw&-}a#P9*y4 zt>A*|M8VIj8rkLn7j~hRZ-Ej~0-r&0wSa4IO?Q(d0n*!gffb;c}@g?VM z>(urp_N~$mcTlKbe$)!rLhlR}bMvh$pXzKe0wCpN)SJ+f_!ro=(`j1imn=UQWAN+tjP9Tg`VU658=@BJs7@v4oiu zM>X^9eD}SZpP!ca$yRXDG`wq83GEDLExU71Lr`??4i)Xq`sg?2It^r~UkCggYOA;W zi><;)Sn*ZfIXDz4ESe^{@4Ot-ejv>wI``H6E44&D+sBJ~uCxMC$ZQx~)6B7--b9_i@bW74wqB z8Z%_)U-DnJy(x41&BZ0XG_tR$y~FIu(EzL;|26#MSC2f?rF80^S^e=$Tly3)#=dLaNChxP>AVW>uA7*~vgc+K7fV_&#gN{$ohdqeZ%2hj)8{!5kVN!BiVkJIKpD6~qi=rb1$Un|Ht zpekJSJ!gJGUl^K2Fi$d#q2FfBSANcCptS?p{g$L_i~?8SAAaFb@vI)ei6ni%X-?n$ zd766b-M1EuOCJfpy>d?In7k16$>0wEP4t{CW}FwTF;+=$qu!-N^{dagdc>M3zTR6L7SMMVcC=Yl91 zAvsvDnsBImr^1c;N%HDD$7Iar%x7zL+@iJ%QErZ+W2$&nQvvF#MU)zSo@a%%DP2X4 z7w&$eBqw3kV*cyyMwVFCKY2-_LL&&y`H3@@tae+cw%3X{rSB~;U+v7r*K;op*lPa9 zYZcrOO8RSx{j!oPt>%Kxl?zc56zsl+F6Vvv^Pib@+EY zpfn@yQV$mEio!#B&9_lhYSpCTkIgz)>^3h7?5IG#qM8fkq|;uo?HYl%#-Q1AYHFiM z^nsQZxrT-YEBuo(j+jtrUr@hB4^YNr=96u@`KjfcX;S#N+jG|zALC|?BIXf;xSvEeY=iElt?=>-D8jz!>3dRXXY(Cu0uBF_*l;L*b3wK z>H9Lw7j2$|2lq&@s<3NnOl^L73IV;2x>oFg`626yjjY7`N2|V;T7iPDi9C>)z3C&I}FTrl=0m{MC@3@)FY9&s49wjS6YK{Gx{Gf$bGsy+y>% zTi92vJd>v1oF!1kZdc&rk0(k`mGo&wm+1a!eI=ek(Xz1kPWB+S!>O=I5r5+`|lSqCcoSuAtG#H%~B{!Qfd6C zgzJCJq4MCRkE#e-zOq={+t=QeMg-K~Ia3>hB)2cDTTceKj#OWYn3Rz@Lz_TC!n$O- zqBXndX?C{2mNMA8Otf-qfhtm-M;x%vU^JP9|0My-sjF9#Pcubr(yZB{yWSjBRF z(5S*g4qb}0=Z^0$J(E4fu_#F?gZuS1Z%`{Fcz=cY-&QHM57U5+^TvGSQ-l z%khfz@OuM=&ONuP$NS`Q^r607Ec^|}B_&r(ffyw0s^y8Ue3KNVX6Mj<_39!Ps>}X; zPU9Vm%xA+!><6x79Rd#|TW zR8YyuuG{3i9O3%KDO{Q%{jV8ROWt@TEOY086^)qSm?i)ZBh8G`BN>_`2o?QxP6w4)23U#n?IAp)?a?$^0 z-k_-1uO`d_ zZAD#rTyKXOov@YAm*?K;iTwx;$_i3qZgOcanEHdywK7!i!2gPmetjnC)m`Hj8O<80 z>%oi7<{!>RpF1Gso*6weJ$fgrwK8O>(g8cbVI%R#8~0tFjdk23oaU#E;b;YA=j*;z zvASFI&tk_~`HuHrsLFP8S-2T|qknGAduk08+Pi(~;;NSt?Cpoq(cu#Ikm3sOuGIEEk;l4%%&2??1!C7?5b|#5Yyrx1aw$^ zB_!+nxP=(r{dnXjy0f@yo2;-N?D|D~slU10Eze=%C;#M0-_5}lEj-Nbh=rd2 zqqZ{-gtF`V_}Ioanj!l-mUdY}36o)vmfKcJwz5=|CKV!-bqoo~9WAnk%AP25XB)+B z7a7|yqmnRDc9QpZ-S7L)^FGhN&tHSN#>|}ST<4rQzwhTW$0$i!7}Ex=`L*kT5ki;@ ztfSLKta7hB-^Zf6-PE5QzRxYlxs+nM*Cjqjr?q+UWR#!7y+4vw8uFz%fi|=Qd)?=I zpDfR%!T@3xY42Vb*66^MqjMt>^&DNnXUZ<$t?pU5(%#)!IcrcmP*X%FI@N_`2e@Q} z=LPgP-B8?>FR!RU=PB4fb4Rp#Iyup2CP~`NdgrEXQEpbvbg{MC1Y+iqu6cjq+#CJm z)#3$5r(L`*iO%_X^buX zI85Q$RDZi|5I?->W%`idYXxHbF|&C3q4NU zJBB{|l+1D57jhtYAB{ctz^#JGpYyiJDV4!8Id>ib( zil6G#`4$VG2hoE0>i8X-EaRq5;}-Zy?DA{9JyKd-gKmC1rhGj3K5kOD#8szNUDBz< zJLyD2+s5qY*scDmij9!ab;%AIk|8RTniiblv*jYkgh}s88@$wTpQD~wcK=KBx1}S# zO~m=Yu+J$6$Pc<@ZXy?R^wGCWS-wkH;aozrZH8tAZN%W7DeU~bN>e;1(N5VXJpWi(+WdyxnG<4K}bTY2$R{e7V?ras!IAl+2f-Zm8 zhb)s64_s+Op0QW|+>I128zyQ{BGaYYdU>1h3iUuyGOJ@eps(oZ$QnL>u zjRh<)`Rw-bR>@6#ohVSScs8m;q7T=z zsbVLd=1}^J!}qHEvS>R(NQoiDECyKqMjihnoaM)*A;p`zOSh9Q1miQ+NTa^G(RRk9 zqPUoE_wL=&dKZrOu4Z1SE8z~uxA}UvKXThFP|)XWrHou zA)04+YMhZH{*|eSSmF$Wf-(OVdNPEDrHo<;1t^n`@`Cn}S`j8h`|vYbg-1nR2tTC# zh`h_;80>E82(i~2;zfN@`}WlVb%#(Rki8=7!2WRakH;ERa+QQW^DevRg4F(*O->^U zWw5-g`uL>&Pk4*y-7!9B0x`JkR~*w!B{D5HBz-1fQjIRlXc4nFwIn z`={kpjG9<-=J-o`V^}+s#mV%Df>RoGJA_4BmwIcQx?X+1PI2dn-1kw)Y8&r)vDwg# zUr)2>HibCZ3e$Xo5+?eDVFAU|D=^qF%~81z^OuXMPsM_H)R;!1f?}NgtEJldQLbgf znhe%1f$oDA=^E7UY0X5rzD;c9S4T1>hQAktS6-8W4U%jVGcl5C!Kl51Ui=f%1-uur zyG@5fx)+`)`LL>050C3r9o!*{J5H+*aoX$ddswIBtZen?4(xslJuz;KKm<1X%rvb<@kr7RhqPeCN-qHBE;vec7@zILx#r<4=LusyPbB2n)D$*X8pe&9{n z!iw;5P4;oB>oy9%YgUYnoSs0;WJ=iP!*#b)=qhI2BDPd*#%_lJQtACG!Z;6u)+n1t zvjL*555gVe%Zq3C9KzuzcbeIUXQJ`Xgc>usZAFLzJ8{Q#lHyGGOl`tN;}!ilEZHrE zlU%n7TuOtxqKkQ|H<1(^ok~L(w|vc#sa}~^TE=(=?de5#yK)#UlE_Y88&V0%?3hrL4iHNB8PdyYN#x zJtxLPNs*ofNtq^lbp-v^3$S-~TV<|t#9o5;qsKoP66Nk){*!4D}HOWsN%@-J}fBtjAv zd(y>U0y3(=Wz<`nsHjZd+KOxUFjXK_>hHHc_x#00&cb1lCb?DJb_`?H*7RB4QrCN- zr?!z78?RcVls{oejO6Lul&yAC8sVC%v|cL0$UWCtttk|h48P+wHMvnv@z1aF&#;U< zV0jaEu7Ufhhs+TEkXZ4Xaam`-ZgT|Nf{=7G^4!^*Ws3QTDZ#>wZI|m5g#xB$)B730 zrPup`qnqUi;*IBaA+vlc`HtUm3Zsd zrzGa~hYugZU98W1?lG9!QLMQwPD!b=L)u+dRu+E3J}~~&v}7bA2t*Ib#Dpu-^Ypv- z@3V7qRI>};zEuFAlc1p*rviCn%CyDal#KwDn$cJ z2`f+>(4Q)NdBgcV+xMTuXdkG~%zAH|+=~AUq#uG(gI|wiD-P(ohVWmMv6hy9fAa+y(dkXMT-$V-&S$p;$D{|9OPbnljfXu1? zJQ_BzkzYhm;Y1YF_mm(tbvG#~41y(Yph%ofQ_h9A1BlN%us-$6kixRJ2eQKXlkVTA zfY)RgR9jx&-fm5ps}#30?1t#(=4PV1yGovQ28hpu0^0B=PM=PM7Xk={MT#+?t(Ml- zF^s8qF4viSwWFQ*L+eYB`;mp(30>-RNH{R9h)ftT;&V+HWk9LTgM#)Ox>FI+Frv6Y zOqr16F(Z+9kQ1LE^ZEB;RjnIP-plY9n`6f!3k%f|m+y3%$s|A@?;%%xyQQUt`V$0Q z65x-h19T0-)8G`_i6q-ik`rc}CZQH`_s>70AtBb5ZYTw8xR;NQ=T{6Dkc;(@InzGd zA|CJ_hr^)DLRPLKfF~FnNT#OTkPD9nj=uC7`VEtIoy73>b1gL&e@8Y}JsL>|R}=LS zxmpl(uwK2I;>Mr}w6(DbgV;C%oCdCZNfpJQWz?SCLU=LftAB4L zlU_i72W4F1N!Ot0+A1S6P{@&bbdTHHv*UF-h|uN0h8qErrgC`tQ!meDJ$ohyiEZkg zJMWOo^Gj>$OAmAw0HitP6L<$?7X+o+s;4RSI zyLar@TxSzxrm(JdH)g@zL}~pJ=Lq^-No_Gs08Lqd8?G=JwQtJ^9{YN(!>|lWL63vK zznM-_p1pA-IVxj)`uudfa^sh)%F>uI#B>XugQ#nr1SwFm=tGZg?~wWY{q?6!&ez)Y zAYK;^{Rr-j={&H8Ielwg&+xne&)XtcgM+L>Xgh8mZ7eEF{*=n=rqsF&41A(5mO ztDtLX50yt-@wMN8b4MVT7eI`NpNe0^9ZnC?!$Z|In4>2U0TRZ~+tT!}sNC0ZY3Yz@ zM7#%xmBj&QUnnPQgUuU(q!Fo^*2E z*#xX2)K$X)@?-&IaUQO7s1PEmOV05yhxPgH|r!0Zjo{oa)UnfZ@>Jfzx~V6|Gfib9rpV4hqY?f%A;lhN*|az za$mh3-aYNrX=uw~@>1B!8c@jrM)O@H+o`^i(}_^`CG(p34SNP~iHevjh0Jnk(5CmG z>Y%5c@5D+%NnggMK9!2znqVo`opcopW_6r^Ptb}@(0r@_+7E=Et&iYap3B%6{<@Hx zpYH)(ad90Y9c)mPa-B&)gL8IWA0nYiFgji9*w3#NIxm`Ml^zG%vPi?nV3p7)U0q#P zzVZ4FP)ET(wRLum2M;0igktB|UB$1tL1QWhnwh&98Hu&OzGyCbC>jOg_6_^;fS6MY zR3dGU_F-?mtQIhgYb>c)tE8&PId zR8?IgDI=#NyBHDoR8&@i3rrA1{-SeG*F~d1c)A%TLD3#?#9Tj`nD`-31BK$RgCKS)(F!K3WdDX?*4u%z|KPL?d`8w5yD!~?}G_Y z&whv0wt=Kjcc~UfCM@V{nNtSEug8R5Gs$2icZYOje}(sEP{DL@-hc=m1tWU@%k$X1 zIySXg{VYl*VYkfk%56hDQJ$s79itVA_a68f;C-f*JGRk50Gf*S&D_KL9ztKtKEsEFk^~KN?!F)&dwQQSgFX`A-x50uLpwBoV~QBI&OIV(o(eAApu_Q z#r5jy>YIcN7;Vmgce|jZqzR_qE8vQq_kIJQY%hb6Q&iNh@uHxB2#Y8Py>RyA_ryZ1 zSYF?^PUA&!aphManq#5+x-p7fKyqmeELwv1=RaZlJeEpHfo~@GP}69={Hsiwa6=cl zM4)FQZGX9DWn8<2u7!cC2~Qv>K$Ze(a-Kd#%-k?Jh~_*)HoeWvB+YGo z0V5R$^?94p3&=!eW@mH5Anj1$CHfuEV-MFDya00ke$U1j<5g98fc2q+o-7JN($r6C zIX44utBi!q{Q)yB7|6M8m{{`cJ?|^jQ4W;I^ Wd|77b<4O4P5!4}busoWciTEEyMfDH> literal 0 HcmV?d00001 diff --git a/doc/index.rst b/doc/index.rst index b9556e9f..0aba3e4c 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -14,6 +14,7 @@ programming/index.rst APIreference/index.rst python + MJX unity models changelog diff --git a/doc/mjx.rst b/doc/mjx.rst new file mode 100644 index 00000000..f0fc1532 --- /dev/null +++ b/doc/mjx.rst @@ -0,0 +1,312 @@ +========== +MuJoCo XLA +========== + +Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the +`mjx `__ directory. MJX allows MuJoCo to run on compute +hardware supported by the `XLA `__ compiler via the +`JAX `__ framework. MJX runs on a +`all platforms supported by JAX `__: Nvidia +and AMD GPUs, Apple Silicon, and `Google Cloud TPUs `__. + +The MJX API is consistent with the main simulation functions in the MuJoCo API, although it is currently missing some +features. While the :ref:`API documentation ` is applicable to both libraries, we indicate features +unsupported by MJX in the :ref:`notes ` below. + +MJX is distributed as a separate package called ``mujoco-mjx`` on `PyPI `__. +Although it depends on the main ``mujoco`` package for model compilation and visualization, it is a re-implementation of +MuJoCo that uses the same algorithms as the MuJoCo implementation. However, in order to properly leverage JAX, MJX +deliberately diverges from the MuJoCo API in a few places, see below. + +MJX is a successor to the `generalized physics pipeline `__ +in Google's `Brax `__ physics and reinforcement learning library. MJX was built +by core contributors to both MuJoCo and Brax, who will together continue to support both Brax (for its reinforcement +learning algorithms and included environments) and MJX (for its physics algorithms). A future version of Brax will +depend on the ``mujoco-mjx`` package, and Brax's existing +`generalized pipeline `__ will be deprecated. This change +will be largely transparent to users of Brax. + +.. _MjxNotebook: + +Tutorial notebook +================= + +The following IPython notebook demonstrates the use of MJX along with reinforcement learning to train humanoid and +quadruped robots to locomote: |colab|. + +.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg + :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb + +.. _MjxInstallation: + +Installation +============ + +The recommended way to install this package is via `PyPI `__: + +.. code-block:: shell + + pip install mujoco-mjx + +A copy of the MuJoCo library is provided as part of this package's depdendencies and does **not** need to be downloaded +or installed separately. + +.. _MjxUsage: + +Basic usage +=========== + +Once installed, the package can be imported via ``from mujoco import mjx``. Structs, functions, and enums are available +directly from the top-level ``mjx`` module. + +.. _MjxStructs: + +Structs +------- + +Before running MJX functions on an accelerator device, structs must be copied onto the device via the ``mjx.device_put`` +function. Placing an :ref:`mjModel` on device yields an ``mjx.Model``. Placing an :ref:`mjData` on device yields +an ``mjx.Data``: + +.. code-block:: python + + model = mujoco.MjModel.from_xml_string("...") + data = mujoco.MjData(model) + mjx_model = mjx.device_put(model) + mjx_data = mjx.device_put(data) + +These MJX variants mirror their MuJoCo counterparts but have three key differences: + +#. Fields in ``mjx.Model`` and ``mjx.Data`` are JAX arrays copied onto device, instead of numpy arrays. +#. Some fields are missing from ``mjx.Model`` and ``mjx.Data`` for features that are + :ref:`unsupported ` in MJX. +#. Arrays in ``mjx.Model`` and ``mjx.Data`` support adding batch dimensions. Batch dimensions are a natural way to + express domain randomization (in the case of ``mjx.Model``) or high-throughput simulation for reinforcement learning + (in the case of ``mjx.Data``). + + +Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An ``mjx.Data`` may be created by calling +``mjx.make_data``, which mirrors the :ref:`mj_makeData` function in MuJoCo: + +.. code-block:: python + + model = mujoco.MjModel.from_xml_string("...") + 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``. + +.. _MjxFunctions: + +Functions +--------- + +MuJoCo functions are exposed as MJX functions of the same name, but following +`PEP 8 `__-compliant names. Most of the :ref:`main simulation ` and +some of the :ref:`sub-components ` for forward simulation are available from the top-level ``mjx`` module. + +MJX functions are not `JIT compiled `__ by default -- we +leave it to the user to JIT MJX functions, or JIT their own functions that reference MJX functions. See the +:ref:`minimal example ` below. + +.. _MjxEnums: + +Enums and constants +------------------- + +MJX enums are available as ``mjx.EnumType.ENUM_VALUE``, for example ``mjx.JointType.FREE``. Enums for unsupported MJX +features are omitted from the MJX enum declaration. MJX declares no constants but references MuJoCo constants directly. + +.. _MjxExample: + +Minimal example +--------------- + +.. code-block:: python + + # Throw a ball at 100 different velocities. + + import jax + import mujoco + from mujoco import mjx + + XML=r""" + + + + + + + + + """ + + model = mujoco.MjModel.from_xml_string(XML) + mjx_model = mjx.device_put(model) + + @jax.vmap + def batched_step(vel): + mjx_data = mjx.make_data(mjx_model) + qvel = mjx_data.qvel.at[0].set(vel) + mjx_data = mjx_data.replace(qvel=qvel) + pos = mjx.step(mjx_model, mjx_data).qpos[0] + return pos + + vel = jax.numpy.arange(0.0, 1.0, 0.01) + pos = jax.jit(batched_step)(vel) + print(pos) + +.. _MjxFeatureParity: + +Feature Parity +============== + +MJX supports most of the main simulation features of MuJoCo, with a few exceptions. MJX will raise an exception if +asked to copy to device an :ref:`mjModel` with field values referencing unsupported features. + +The following features are **fully supported** in MJX: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - Dynamics + - :ref:`Forward ` + * - :ref:`Joint ` + - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE`` + * - :ref:`Transmission ` + - ``TRN_JOINT`` + * - :ref:`Actuation ` + - ``DYN_NONE``, ``DYN_INTEGRATOR``, ``DYN_FILTER``, ``GAIN_FIXED``, ``GAIN_AFFINE``, ``BIAS_NONE``, + ``BIAS_AFFINE`` + * - :ref:`Geom ` + - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` + * - :ref:`Constraint ` + - ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL`` + * - :ref:`Integrator ` + - ``EULER``, ``RK4`` + * - :ref:`Cone ` + - ``PYRAMIDAL`` + * - :ref:`Condim ` + - 3 + * - :ref:`Solver ` + - ``CG`` + * - Fluid Model + - :ref:`flInertia` + +The following features are **in development** and coming soon: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - Dynamics + - :ref:`Inverse ` + * - :ref:`Transmission ` + - ``TRN_TENDON`` + * - :ref:`Geom ` + - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF`` + * - :ref:`Integrator ` + - ``IMPLICIT``, ``IMPLICITFAST`` + * - :ref:`Cone ` + - ``ELLIPTIC`` + * - :ref:`Condim ` + - 1, 4, 6 + * - :ref:`Solver ` + - ``NEWTON`` + * - Fluid Model + - :ref:`flEllipsoid` + * - :ref:`Tendons ` + - :ref:`Spatial `, :ref:`Fixed ` + +The following features are **unsupported**: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - :ref:`Transmission ` + - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``, ``MUSCLE`` + * - :ref:`Solver ` + - ``PGS`` + * - :ref:`Callbacks ` + - ``mjDYN_USER``, ``mjGAIN_USER``, ``mjBIAS_USER``, ``mjSENS_USER`` + +.. _MjxSharpBits: + +🔪 MJX - The Sharp Bits 🔪 +========================== + +GPUs and TPUs have unique performance tradeoffs that MJX is subject to. MJX specializes in simulating big batches of +parallel identical physics scenes using algorithms that can be efficiently vectorized on +`SIMD hardware `__. This specialization is useful +for machine learning workloads such as `reinforcement learning `__ +that require massive data throughput. + +There are certain workflows that MJX is ill-suited for: + +Single scene simulation + Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been + carefully optimized for CPU. MJX works best when simulating thousands or tens of thousands of scenes in parallel. + +Large, complex scenes with many contacts + Accelerators exhibit poor performance for + `branching code `__. + Branching is used in broad-phase collision detection, when identifying potential collisions between large numbers of + bodies in a scene. MJX ships with a simple branchless broad-phase algorithm (see performance tuning) but it is not as + powerful as the one in MuJoCo. + + To see how this affects simulation, let us consider a physics scene with increasing numbers of physics bodies. We + simulate a scene with a variable number of humanoids (from 1 to 10) and then compare MJX's performance on an Nvidia + A100 GPU to MuJoCo on a 12-core workstation: + + .. figure:: images/mjx/mujoco_vs_mjx_large_scene.png + :width: 658px + :align: center + + Notice that as we increase the number of humanoids (which increases the number of potential contacts in a scene), MJX + performance degrades more rapidly than MuJoCo. At the limit, for such a large scene, MuJoCo performance nearly + matches MJX. + +Scenes with collisions between meshes with many vertices + MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of + `mesh collision algorithms `__. + These algorithms work well for smaller meshes (with hundreds of vertices) but suffer with large meshes. With careful + tuning, MJX can simulate scenes with mesh collisions well -- see the MJX + `shadow hand `__ + config for an example. + +.. _MjxPerformance: + +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. + +:ref:`contact-pair` element + Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider + during each step. Enabling only an explicit list of valid contacts can have a dramatic effect on simulation + performance in MJX. Doing this well often requires an understanding of the task -- for example, the + `OpenAI Gym Humanoid `__ task resets when + the humanoid starts to fall, so full contact with the floor is not needed. + +:ref:`option-flag` element + Disabling ``eulerdamp`` can help performance and is often not needed for stability. diff --git a/doc/unity.rst b/doc/unity.rst index 9fe13ef9..2c756da0 100644 --- a/doc/unity.rst +++ b/doc/unity.rst @@ -30,14 +30,14 @@ _____ The MuJoCo app needs to be run at least once before the native library can be used, in order to register the library as a trusted binary. Then, copy the dynamic library file from -``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.2.3.8.dylib`` (it can be +``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.0.0.dylib`` (it can be found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``. Linux _____ Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from -``~/.mujoco/mujoco-2.3.8/lib/libmujoco.so.2.3.8`` and rename it as ``libmujoco.so``. +``~/.mujoco/mujoco-3.0.0/lib/libmujoco.so.3.0.0`` and rename it as ``libmujoco.so``. Windows _______ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 61d63325..1f98efca 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -24,7 +24,7 @@ extern "C" { #endif // header version; should match the library version as returned by mj_version() -#define mjVERSION_HEADER 238 +#define mjVERSION_HEADER 300 // needed to define size_t, fabs and log10 #include diff --git a/mjx/MANIFEST.in b/mjx/MANIFEST.in new file mode 100644 index 00000000..b963ce6e --- /dev/null +++ b/mjx/MANIFEST.in @@ -0,0 +1,2 @@ +recursive-include mujoco/mjx/test_data * +recursive-include mujoco/mjx/benchmark *.obj *.stl *.xml diff --git a/mjx/README.md b/mjx/README.md new file mode 100644 index 00000000..516271fc --- /dev/null +++ b/mjx/README.md @@ -0,0 +1,54 @@ +# MuJoCo XLA (MJX) + +[![PyPI Python Version][pypi-versions-badge]][pypi] +[![PyPI version][pypi-badge]][pypi] + +[pypi-versions-badge]: https://img.shields.io/pypi/pyversions/mujoco-mjx +[pypi-badge]: https://badge.fury.io/py/mujoco-mjx.svg +[pypi]: https://pypi.org/project/mujoco-mjx/ + +This package is a re-implementation of the +[MuJoCo physics engine](https://github.com/google-deepmind/mujoco) in +[JAX](https://github.com/google/jax). This library is developed and maintained +by Google DeepMind, and is kept up-to-date with the latest developments in +MuJoCo itself. + +The `mujoco-mjx` package is API-compatible with MuJoCo, but is missing some +features found in MuJoCo. See our +[documentation](https://mujoco.readthedocs.io/en/stable/mjx.html) for more +details concerning feature parity. + +## Installation + +The recommended way to install this package is via [PyPI](https://pypi.org/project/mujoco-mjx/): + +```sh +pip install mujoco-mjx +``` + +## Usage + +Once installed, the package can be imported via `from mujoco import mjx`. Please +consult our [documentation](https://mujoco.readthedocs.io/en/stable/mjx.html) +for further detail on the package's API. + +We recommend going through the tutorial notebook which covers the basics of +MuJoCo using Python: +[![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) + +## Versioning + +The `major.minor.micro` portion of the version number matches the version of +MuJoCo that this library provides. Optionally, if we release updates to MJX that +target the same version of MuJoCo, a `.postN` suffix is added, for example +`3.0.1.post2` represents the second update to MJX for MuJoCo 3.0.1. + +## License and Disclaimer + +Copyright 2023 DeepMind Technologies Limited + +MuJoCo and its libraries are licensed under the Apache License, +Version 2.0. You may obtain a copy of the License at +https://www.apache.org/licenses/LICENSE-2.0. + +This is not an officially supported Google product. diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt new file mode 100644 index 00000000..3898472f --- /dev/null +++ b/mjx/cuda_requirements.txt @@ -0,0 +1,17 @@ +-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html + +jax[cuda12_local]==0.4.13; python_version == '3.8' \ + --hash=sha256:03bfe6749dfe647f16f15f6616638adae6c4a7ca7167c75c21961ecfd3a3baaa +jax[cuda12_local]==0.4.18; python_version >= '3.9' \ + --hash=sha256:c3ab72ea2f1c5d8ccf2561e79f6562fb2964629f3e55b3ac1c11c48b64c20336 +jaxlib==0.4.13+cuda12.cudnn89; python_version == '3.8' \ + --hash=sha256:31372f41dc28ecb11a5cc5573ae632e8ee9cfb63788edc8ad8a77e8d3279f569 +jaxlib==0.4.18+cuda12.cudnn89; python_version >= '3.9' \ + --hash=sha256:14f74ff081882ea091c121e355051b35932e39cb7ff7242b88a87f3690f3ca90 \ + --hash=sha256:7c87dc2d68257b02e83c04be88a3c447373ee7077d65f43545bcbda5bfe2231d \ + --hash=sha256:4d16e9c7592e1aaca0b3d28d2c8beba415a2721bb7001f2947728247951a250d \ + --hash=sha256:759c08c69f4a5b1e6b39c3e4eff908a04ce3b2b483bb594ed624407c7d12d110 \ + --hash=sha256:a7a04dbe1851cd50d07691282116aee49a2f0be7838e55b76d7ada86db06be62 \ + --hash=sha256:2bf842db3d58c8c6c52fbc8ed3fabefd7b91a21746cd59d3eaf3522eea229b53 \ + --hash=sha256:35d265ef9bb3835a14580cbaa9402060f117e46056f80e0996405fff3964667a \ + --hash=sha256:0e4352f24d629e912965e6435e140c1b06086243a098651f2d01b75f3738b51c diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py new file mode 100644 index 00000000..3ec6312a --- /dev/null +++ b/mjx/mujoco/mjx/__init__.py @@ -0,0 +1,34 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Public API for MJX.""" + +# pylint:disable=g-importing-member +from mujoco.mjx._src.collision_driver import collision +from mujoco.mjx._src.constraint import make_constraint +from mujoco.mjx._src.device import device_get_into +from mujoco.mjx._src.device import device_put +from mujoco.mjx._src.forward import forward +from mujoco.mjx._src.forward import step +from mujoco.mjx._src.io import make_data +from mujoco.mjx._src.passive import passive +from mujoco.mjx._src.smooth import com_pos +from mujoco.mjx._src.smooth import com_vel +from mujoco.mjx._src.smooth import crb +from mujoco.mjx._src.smooth import factor_m +from mujoco.mjx._src.smooth import kinematics +from mujoco.mjx._src.smooth import mul_m +from mujoco.mjx._src.smooth import rne +from mujoco.mjx._src.smooth import transmission +from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/__init__.py b/mjx/mujoco/mjx/_src/__init__.py new file mode 100644 index 00000000..4ad51d69 --- /dev/null +++ b/mjx/mujoco/mjx/_src/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/mjx/mujoco/mjx/_src/collision_base.py b/mjx/mujoco/mjx/_src/collision_base.py new file mode 100644 index 00000000..c843dceb --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_base.py @@ -0,0 +1,64 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collision base.""" + +import dataclasses +from typing import Dict, List, Optional, Tuple + +import jax +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import GeomType +# pylint: enable=g-importing-member + +Contact = Tuple[jax.Array, jax.Array, jax.Array] + + +@dataclasses.dataclass +class Candidate: + geom1: int + geom2: int + ipair: int + geomp: int # priority geom + dim: int + + +CandidateSet = Dict[ + Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]], + List[Candidate], +] + + +class GeomInfo(PyTreeNode): + """Collision info for a geom.""" + + pos: jax.Array + mat: jax.Array + size: jax.Array + face: Optional[jax.Array] = None + vert: Optional[jax.Array] = None + edge: Optional[jax.Array] = None + facenorm: Optional[jax.Array] = None + + +class SolverParams(PyTreeNode): + """Contact solver params.""" + + friction: jax.Array + solref: jax.Array + solreffriction: jax.Array + solimp: jax.Array + margin: jax.Array + gap: jax.Array diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py new file mode 100644 index 00000000..62d357f9 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -0,0 +1,724 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Convex collisions.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import math +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Contact +from mujoco.mjx._src.collision_base import GeomInfo +# pylint: enable=g-importing-member + + +def _closest_segment_point_plane( + a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Gets the closest point between a line segment and a plane. + + Args: + a: first line segment point + b: second line segment point + p0: point on plane + plane_normal: plane normal + + Returns: + closest point between the line segment and the plane + """ + # Parametrize a line segment as S(t) = a + t * (b - a), plug it into the plane + # equation dot(n, S(t)) - d = 0, then solve for t to get the line-plane + # intersection. We then clip t to be in [0, 1] to be on the line segment. + n = plane_normal + d = jp.sum(p0 * n) # shortest distance from origin to plane + denom = jp.sum(n * (b - a)) + t = (d - jp.sum(n * a)) / (denom + 1e-6 * (denom == 0.0)) + t = jp.clip(t, 0, 1) + segment_point = a + t * (b - a) + + return segment_point + + +def _closest_triangle_point( + p0: jax.Array, p1: jax.Array, p2: jax.Array, pt: jax.Array +) -> jax.Array: + """Gets the closest point between a triangle and a point in space. + + Args: + p0: triangle point + p1: triangle point + p2: triangle point + pt: point to test + + Returns: + closest point on the triangle w.r.t point pt + """ + # Parametrize the triangle s.t. a point inside the triangle is + # Q = p0 + u * e0 + v * e1, when 0 <= u <= 1, 0 <= v <= 1, and + # 0 <= u + v <= 1. Let e0 = (p1 - p0) and e1 = (p2 - p0). + # We analytically minimize the distance between the point pt and Q. + e0 = p1 - p0 + e1 = p2 - p0 + a = e0.dot(e0) + b = e0.dot(e1) + c = e1.dot(e1) + d = pt - p0 + # The determinant is 0 only if the angle between e1 and e0 is 0 + # (i.e. the triangle has overlapping lines). + det = a * c - b * b + u = (c * e0.dot(d) - b * e1.dot(d)) / det + v = (-b * e0.dot(d) + a * e1.dot(d)) / det + inside = (0 <= u) & (u <= 1) & (0 <= v) & (v <= 1) & (u + v <= 1) + closest_p = p0 + u * e0 + v * e1 + d0 = (closest_p - pt).dot(closest_p - pt) + + # If the closest point is outside the triangle, it must be on an edge, so we + # check each triangle edge for a closest point to the point pt. + closest_p1, d1 = math.closest_segment_point_and_dist(p0, p1, pt) + closest_p = jp.where((d0 < d1) & inside, closest_p, closest_p1) + min_d = jp.where((d0 < d1) & inside, d0, d1) + + closest_p2, d2 = math.closest_segment_point_and_dist(p1, p2, pt) + closest_p = jp.where(d2 < min_d, closest_p2, closest_p) + min_d = jp.minimum(min_d, d2) + + closest_p3, d3 = math.closest_segment_point_and_dist(p2, p0, pt) + closest_p = jp.where(d3 < min_d, closest_p3, closest_p) + + return closest_p + + +def _closest_segment_triangle_points( + a: jax.Array, + b: jax.Array, + p0: jax.Array, + p1: jax.Array, + p2: jax.Array, + triangle_normal: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Gets the closest points between a line segment and triangle. + + Args: + a: first line segment point + b: second line segment point + p0: triangle point + p1: triangle point + p2: triangle point + triangle_normal: normal of triangle + + Returns: + closest point on the triangle w.r.t the line segment + """ + # The closest triangle point is either on the edge or within the triangle. + # First check triangle edges for the closest point. + # TODO(robotics-simulation): consider vmapping over closest point functions + seg_pt1, tri_pt1 = math.closest_segment_to_segment_points(a, b, p0, p1) + d1 = (seg_pt1 - tri_pt1).dot(seg_pt1 - tri_pt1) + seg_pt2, tri_pt2 = math.closest_segment_to_segment_points(a, b, p1, p2) + d2 = (seg_pt2 - tri_pt2).dot(seg_pt2 - tri_pt2) + seg_pt3, tri_pt3 = math.closest_segment_to_segment_points(a, b, p0, p2) + d3 = (seg_pt3 - tri_pt3).dot(seg_pt3 - tri_pt3) + + # Next, handle the case where the closest triangle point is inside the + # triangle. Either the line segment intersects the triangle or a segment + # endpoint is closest to a point inside the triangle. + seg_pt4 = _closest_segment_point_plane(a, b, p0, triangle_normal) + tri_pt4 = _closest_triangle_point(p0, p1, p2, seg_pt4) + d4 = (seg_pt4 - tri_pt4).dot(seg_pt4 - tri_pt4) + + # Get the point with minimum distance from the line segment point to the + # triangle point. + distance = jp.array([[d1, d2, d3, d4]]) + min_dist = jp.amin(distance) + mask = (distance == min_dist).T + seg_pt = jp.array([seg_pt1, seg_pt2, seg_pt3, seg_pt4]) * mask + tri_pt = jp.array([tri_pt1, tri_pt2, tri_pt3, tri_pt4]) * mask + seg_pt = jp.sum(seg_pt, axis=0) / jp.sum(mask) + tri_pt = jp.sum(tri_pt, axis=0) / jp.sum(mask) + + return seg_pt, tri_pt + + +def _manifold_points( + poly: jax.Array, poly_mask: jax.Array, poly_norm: jax.Array +) -> jax.Array: + """Chooses four points on the polygon with approximately maximal area.""" + dist_mask = jp.where(poly_mask, 0.0, -1e6) + a_idx = jp.argmax(dist_mask) + a = poly[a_idx] + # choose point b furthest from a + b_idx = (((a - poly) ** 2).sum(axis=1) + dist_mask).argmax() + b = poly[b_idx] + # choose point c furthest along the axis orthogonal to (a-b) + ab = jp.cross(poly_norm, a - b) + ap = a - poly + c_idx = (jp.abs(ap.dot(ab)) + dist_mask).argmax() + c = poly[c_idx] + # choose point d furthest from the other two triangle edges + ac = jp.cross(poly_norm, a - c) + bc = jp.cross(poly_norm, b - c) + bp = b - poly + dist_bp = jp.abs(bp.dot(bc)) + dist_mask + dist_ap = jp.abs(ap.dot(ac)) + dist_mask + d_idx = jp.concatenate([dist_bp, dist_ap]).argmax() % poly.shape[0] + return jp.array([a_idx, b_idx, c_idx, d_idx]) + + +def _project_pt_onto_plane( + pt: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Projects a point onto a plane along the plane normal.""" + dist = (pt - plane_pt).dot(plane_normal) + return pt - dist * plane_normal + + +def _project_poly_onto_plane( + poly: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Projects a polygon onto a plane using the plane normal.""" + return jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])( + poly, plane_pt, math.normalize(plane_normal) + ) + + +def _project_poly_onto_poly_plane( + poly1: jax.Array, norm1: jax.Array, poly2: jax.Array, norm2: jax.Array +) -> jax.Array: + """Projects poly1 onto the poly2 plane along poly1's normal.""" + d = poly2[0].dot(norm2) + denom = norm1.dot(norm2) + t = (d - poly1.dot(norm2)) / (denom + 1e-6 * (denom == 0.0)) + new_poly = poly1 + t.reshape(-1, 1) * norm1 + return new_poly + + +def _point_in_front_of_plane( + plane_pt: jax.Array, plane_normal: jax.Array, pt: jax.Array +) -> jax.Array: + """Checks if a point is strictly in front of a plane.""" + return (pt - plane_pt).dot(plane_normal) > 1e-6 + + +def _clip_edge_to_planes( + edge_p0: jax.Array, + edge_p1: jax.Array, + plane_pts: jax.Array, + plane_normals: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Clips an edge against side planes. + + We return two clipped points, and a mask to include the new edge or not. + + Args: + edge_p0: the first point on the edge + edge_p1: the second point on the edge + plane_pts: side plane points + plane_normals: side plane normals + + Returns: + new_ps: new edge points that are clipped against side planes + mask: a boolean mask, True if an edge point is a valid clipped point and + False otherwise + """ + p0, p1 = edge_p0, edge_p1 + p0_in_front = jax.vmap(jp.dot)(p0 - plane_pts, plane_normals) > 1e-6 + p1_in_front = jax.vmap(jp.dot)(p1 - plane_pts, plane_normals) > 1e-6 + + # Get candidate clipped points along line segment (p0, p1) by clipping against + # all clipping planes. + candidate_clipped_ps = jax.vmap( + _closest_segment_point_plane, in_axes=[None, None, 0, 0] + )(p0, p1, plane_pts, plane_normals) + + def clip_edge_point(p0, p1, p0_in_front, clipped_ps): + @jax.vmap + def choose_edge_point(in_front, clipped_p): + return jp.where(in_front, clipped_p, p0) + + # Pick the clipped point if p0 is in front of the clipping plane. Otherwise + # keep p0 as the edge point. + new_edge_ps = choose_edge_point(p0_in_front, clipped_ps) + + # Pick the clipped point that is most along the edge direction. + # This degenerates to picking the original point p0 if p0 is *not* in front + # of any clipping planes. + dists = jp.dot(new_edge_ps - p0, p1 - p0) + new_edge_p = new_edge_ps[jp.argmax(dists)] + return new_edge_p + + # Clip each edge point. + new_p0 = clip_edge_point(p0, p1, p0_in_front, candidate_clipped_ps) + new_p1 = clip_edge_point(p1, p0, p1_in_front, candidate_clipped_ps) + clipped_pts = jp.array([new_p0, new_p1]) + + # Keep the original points if both points are in front of any of the clipping + # planes, rather than creating a new clipped edge. If the entire subject edge + # is in front of any clipping plane, we need to grab an edge from the clipping + # polygon instead. + both_in_front = p0_in_front & p1_in_front + mask = ~jp.any(both_in_front) + new_ps = jp.where(mask, clipped_pts, jp.array([p0, p1])) + # Mask out crossing clipped edge points. + mask = jp.where((p0 - p1).dot(new_ps[0] - new_ps[1]) < 0, False, mask) + return new_ps, jp.array([mask, mask]) + + +def _clip( + clipping_poly: jax.Array, + subject_poly: jax.Array, + clipping_normal: jax.Array, + subject_normal: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Clips a subject polygon against a clipping polygon. + + A parallelized clipping algorithm for convex polygons. The result is a set of + vertices on the clipped subject polygon in the subject polygon plane. + + Args: + clipping_poly: the polygon that we use to clip the subject polygon against + subject_poly: the polygon that gets clipped + clipping_normal: normal of the clipping polygon + subject_normal: normal of the subject polygon + + Returns: + clipped_pts: points on the clipped polygon + mask: True if a point is in the clipping polygon, False otherwise + """ + # Get clipping edge points, edge planes, and edge normals. + clipping_p0 = jp.roll(clipping_poly, 1, axis=0) + clipping_plane_pts = clipping_p0 + clipping_p1 = clipping_poly + clipping_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])( + clipping_p1 - clipping_p0, + clipping_normal, + ) + + # Get subject edge points, edge planes, and edge normals. + subject_edge_p0 = jp.roll(subject_poly, 1, axis=0) + subject_plane_pts = subject_edge_p0 + subject_edge_p1 = subject_poly + subject_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])( + subject_edge_p1 - subject_edge_p0, + subject_normal, + ) + + # Clip all edges of the subject poly against clipping side planes. + clipped_edges0, masks0 = jax.vmap( + _clip_edge_to_planes, in_axes=[0, 0, None, None] + )( + subject_edge_p0, + subject_edge_p1, + clipping_plane_pts, + clipping_plane_normals, + ) + + # Project the clipping poly onto the subject plane. + clipping_p0_s = _project_poly_onto_poly_plane( + clipping_p0, clipping_normal, subject_poly, subject_normal + ) + clipping_p1_s = _project_poly_onto_poly_plane( + clipping_p1, clipping_normal, subject_poly, subject_normal + ) + + # Clip all edges of the clipping poly against subject planes. + clipped_edges1, masks1 = jax.vmap( + _clip_edge_to_planes, in_axes=[0, 0, None, None] + )(clipping_p0_s, clipping_p1_s, subject_plane_pts, subject_plane_normals) + + # Merge the points and reshape. + clipped_edges = jp.concatenate([clipped_edges0, clipped_edges1]) + masks = jp.concatenate([masks0, masks1]) + clipped_points = clipped_edges.reshape((-1, 3)) + mask = masks.reshape(-1) + + return clipped_points, mask + + +def _create_contact_manifold( + clipping_poly: jax.Array, + subject_poly: jax.Array, + clipping_norm: jax.Array, + subject_norm: jax.Array, + sep_axis: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Creates a contact manifold between two convex polygons. + + The polygon faces are expected to have a counter clockwise winding order so + that clipping plane normals point away from the polygon center. + + Args: + clipping_poly: the reference polygon to clip the contact against. + subject_poly: the subject polygon to clip contacts onto. + clipping_norm: the clipping polygon normal. + subject_norm: the subject polygon normal. + sep_axis: the separating axis + + Returns: + tuple of dist, pos, and normal + """ + # Clip the subject (incident) face onto the clipping (reference) face. + # The incident points are clipped points on the subject polygon. + poly_incident, mask = _clip( + clipping_poly, subject_poly, clipping_norm, subject_norm + ) + # The reference points are clipped points on the clipping polygon. + poly_ref = _project_poly_onto_plane( + poly_incident, clipping_poly[0], clipping_norm + ) + behind_clipping_plane = _point_in_front_of_plane( + clipping_poly[0], -clipping_norm, poly_incident + ) + mask = mask & behind_clipping_plane + + # Choose four contact points. + best = _manifold_points(poly_ref, mask, clipping_norm) + contact_pts = jp.take(poly_ref, best, axis=0) + mask_pts = jp.take(mask, best, axis=0) + penetration_dir = jp.take(poly_incident, best, axis=0) - contact_pts + penetration = penetration_dir.dot(-clipping_norm) + + dist = jp.where(mask_pts, -penetration, jp.ones_like(penetration)) + pos = contact_pts + normal = -jp.stack([sep_axis] * 4, 0) + return dist, pos, normal + + +def _sat_hull_hull( + faces_a: jax.Array, + faces_b: jax.Array, + vertices_a: jax.Array, + vertices_b: jax.Array, + normals_a: jax.Array, + normals_b: jax.Array, + unique_edges_a: jax.Array, + unique_edges_b: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Runs the Separating Axis Test for a pair of hulls. + + Given two convex hulls, the Separating Axis Test finds a separating axis + between all edge pairs and face pairs. Edge pairs create a single contact + point and face pairs create a contact manifold (up to four contact points). + We return both the edge and face contacts. Valid contacts can be checked with + dist < 0. Resulting edge contacts should be preferred over face contacts. + + Args: + faces_a: An ndarray of hull A's polygon faces. + faces_b: An ndarray of hull B's polygon faces. + vertices_a: Vertices for hull A. + vertices_b: Vertices for hull B. + normals_a: Normal vectors for hull A's polygon faces. + normals_b: Normal vectors for hull B's polygon faces. + unique_edges_a: Unique edges for hull A. + unique_edges_b: Unique edges for hull B. + + Returns: + tuple of dist, pos, and normal + """ + # get the separating axes + edge_dir_a = unique_edges_a[:, 0] - unique_edges_a[:, 1] + edge_dir_b = unique_edges_b[:, 0] - unique_edges_b[:, 1] + edge_dir_a_r = jp.tile(edge_dir_a, reps=(unique_edges_b.shape[0], 1)) + edge_dir_b_r = jp.repeat(edge_dir_b, repeats=unique_edges_a.shape[0], axis=0) + edge_edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r) + edge_edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))( + edge_edge_axes + ) + + axes = jp.concatenate([normals_a, normals_b, edge_edge_axes]) + + # for each separating axis, get the support + @jax.vmap + def get_support(axis): + support_a = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_a) + support_b = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_b) + dist1 = support_a.max() - support_b.min() + dist2 = support_b.max() - support_a.min() + sign = jp.where(dist1 > dist2, -1, 1) + dist = jp.minimum(dist1, dist2) + dist = jp.where(~jp.all(axis == 0.0), dist, 1e6) # degenerate axis + return dist, sign + + support, sign = get_support(axes) + + # choose the best separating axis + best_idx = jp.argmin(support) + best_sign = sign[best_idx] + best_axis = axes[best_idx] + is_edge_contact = best_idx >= (normals_a.shape[0] + normals_b.shape[0]) + + # get the (reference) face most aligned with the separating axis + dist_a = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_a) + dist_b = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_b) + a_max = dist_a.argmax() + b_max = dist_b.argmax() + a_min = dist_a.argmin() + b_min = dist_b.argmin() + + ref_face = jp.where(best_sign > 0, faces_a[a_max], faces_b[b_max]) + ref_face_norm = jp.where(best_sign > 0, normals_a[a_max], normals_b[b_max]) + incident_face = jp.where(best_sign > 0, faces_b[b_min], faces_a[a_min]) + incident_face_norm = jp.where( + best_sign > 0, normals_b[b_min], normals_a[a_min] + ) + + dist, pos, normal = _create_contact_manifold( + ref_face, + incident_face, + ref_face_norm, + incident_face_norm, + -best_sign * best_axis, + ) + + # For edge contacts, we use the clipped face point, mainly for performance + # reasons. For small penetration, the clipped face point is roughly the edge + # contact point. + idx = dist.argmin() + dist = jp.where( + is_edge_contact, + jp.array([dist[idx], 1, 1, 1]), + dist, + ) + pos = jp.where(is_edge_contact, jp.tile(pos[idx], (4, 1)), pos) + + return dist, pos, normal + + +def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contacts between a plane and a convex object.""" + vert = convex.vert + + # get points in the convex frame + plane_pos = convex.mat.T @ (plane.pos - convex.pos) + n = convex.mat.T @ plane.mat[:, 2] + support = (plane_pos - vert) @ n + idx = _manifold_points(vert, support > 0, n) + pos = vert[idx] + + # convert to world frame + pos = convex.pos + pos @ convex.mat.T + n = plane.mat[:, 2] + + frame = jp.stack([math.make_frame(n)] * 4, axis=0) + unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1 + dist = jp.where(unique, -support[idx], 1) + return dist, pos, frame + + +def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contact between a sphere and a convex object.""" + faces = jp.take(convex.vert, convex.face, axis=0) + normals = convex.facenorm + + # Put sphere in convex frame. + sphere_pos = convex.mat.T @ (sphere.pos - convex.pos) + + # Get support from face normals. + @jax.vmap + def get_support(faces, normal): + pos = sphere_pos - normal * sphere.size[0] + return jp.dot(pos - faces[0], normal) + + support = get_support(faces, normals) + + # Pick the face with minimal penetration as long as it has support. + support = jp.where(support >= 0, -1e12, support) + best_idx = support.argmax() + face = faces[best_idx] + normal = normals[best_idx] + + # Get closest point between the polygon face and the sphere center point. + # Project the sphere center point onto poly plane. If it's inside polygon + # edge normals, then we're done. + pt = _project_pt_onto_plane(sphere_pos, face[0], normal) + edge_p0 = jp.roll(face, 1, axis=0) + edge_p1 = face + edge_normals = jax.vmap(jp.cross, in_axes=[0, None])( + edge_p1 - edge_p0, + normal, + ) + edge_dist = jax.vmap( + lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm) + )(edge_p0, edge_normals) + inside = jp.all(edge_dist <= 0) # lte to handle degenerate edges + + # If the point is outside edge normals, project onto the closest edge plane + # that the point is in front of. + degenerate_edge = jp.all(edge_normals == 0, axis=1) + behind = edge_dist < 0.0 + edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist) + idx = edge_dist.argmin() + edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt) + + pt = jp.where(inside, pt, edge_pt) + + # Get the normal, dist, and contact position. + n, d = math.normalize_with_norm(pt - sphere_pos) + spt = sphere_pos + n * sphere.size[0] + dist = d - sphere.size[0] + pos = (pt + spt) * 0.5 + + # Go back to world frame. + n = convex.mat @ n + pos = convex.mat @ pos + convex.pos + + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contacts between a capsule and a convex object.""" + # Get convex transformed normals, faces, and vertices. + faces = jp.take(convex.vert, convex.face, axis=0) + normals = convex.facenorm + + # Put capsule in convex frame. + cap_pos = convex.mat.T @ (cap.pos - convex.pos) + axis, length = cap.mat[:, 2], cap.size[1] + axis = convex.mat.T @ axis + seg = axis * length + cap_pts = jp.array([ + cap_pos - seg, + cap_pos + seg, + ]) + + # Get support from face normals. + @jax.vmap + def get_support(face, normal): + pts = cap_pts - normal * cap.size[0] + sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts) + return sup.min() + + support = get_support(faces, normals) + has_support = jp.all(support < 0) + + # Pick the face with minimal penetration as long as it has support. + support = jp.where(support >= 0, -1e12, support) + best_idx = support.argmax() + face = faces[best_idx] + normal = normals[best_idx] + + # Clip the edge against side planes and create two contact points against the + # face. + edge_p0 = jp.roll(face, 1, axis=0) + edge_p1 = face + edge_normals = jax.vmap(jp.cross, in_axes=[0, None])( + edge_p1 - edge_p0, + normal, + ) + cap_pts_clipped, mask = _clip_edge_to_planes( + cap_pts[0], cap_pts[1], edge_p0, edge_normals + ) + cap_pts_clipped = cap_pts_clipped - normal * cap.size[0] + face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])( + cap_pts_clipped, face[0], normal + ) + # Create variables for the face contact. + pos = (cap_pts_clipped + face_pts) * 0.5 + norm = jp.stack([normal] * 2, 0) + penetration = jp.where( + mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1 + ) + + # Get a potential edge contact. + edge_closest, cap_closest = jax.vmap( + math.closest_segment_to_segment_points, in_axes=[0, 0, None, None] + )(edge_p0, edge_p1, cap_pts[0], cap_pts[1]) + e_idx = ((edge_closest - cap_closest) ** 2).sum(axis=1).argmin() + cap_closest_pt, edge_closest_pt = cap_closest[e_idx], edge_closest[e_idx] + edge_axis = cap_closest_pt - edge_closest_pt + edge_axis, edge_dist = math.normalize_with_norm(edge_axis) + edge_pos = ( + edge_closest_pt + (cap_closest_pt - edge_axis * cap.size[0]) + ) * 0.5 + edge_norm = edge_axis + edge_penetration = cap.size[0] - edge_dist + has_edge_contact = edge_penetration > 0 + + # Get the contact info. + pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos) + n = -jp.where(has_edge_contact, norm.at[0].set(edge_norm), norm) + + # Go back to world frame. + pos = convex.pos + pos @ convex.mat.T + n = n @ convex.mat.T + + dist = -jp.where( + has_edge_contact, penetration.at[0].set(edge_penetration), penetration + ) + frame = jax.vmap(math.make_frame)(n) + return dist, pos, frame + + +def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact: + """Calculates contacts between two convex objects.""" + if c1.face is None or c2.face is None or c1.vert is None or c2.vert is None: + raise AssertionError('Mesh info missing.') + # pad face vertices so that we can broadcast between geom1 and geom2 + s1, s2 = c1.face.shape[-1], c2.face.shape[-1] + if s1 < s2: + face = jp.pad(c1.face, ((0, 0), (0, s2 - s1)), 'edge') + c1 = c1.replace(face=face) + elif s2 < s1: + face = jp.pad(c2.face, ((0, 0), (0, s1 - s2)), 'edge') + c2 = c2.replace(face=face) + + # ensure that the first object has fewer verts + swapped = c1.vert.shape[0] > c2.vert.shape[0] + if swapped: + c1, c2 = c2, c1 + + faces1 = jp.take(c1.vert, c1.face, axis=0) + faces2 = jp.take(c2.vert, c2.face, axis=0) + + to_local_pos = c2.mat.T @ (c1.pos - c2.pos) + to_local_mat = c2.mat.T @ c1.mat + + faces1 = to_local_pos + faces1 @ to_local_mat.T + normals1 = c1.facenorm @ to_local_mat.T + normals2 = c2.facenorm + + vertices1 = to_local_pos + c1.vert @ to_local_mat.T + vertices2 = c2.vert + + unique_edges1 = jp.take(vertices1, c1.edge, axis=0) + unique_edges2 = jp.take(vertices2, c2.edge, axis=0) + + dist, pos, normal = _sat_hull_hull( + faces1, + faces2, + vertices1, + vertices2, + normals1, + normals2, + unique_edges1, + unique_edges2, + ) + + # Go back to world frame. + pos = c2.pos + pos @ c2.mat.T + normal = normal @ c2.mat.T + normal = -normal if swapped else normal + + frame = jax.vmap(math.make_frame)(normal) + return dist, pos, frame + + +# store ncon as function attributes +plane_convex.ncon = 4 +sphere_convex.ncon = 1 +capsule_convex.ncon = 2 +convex_convex.ncon = 4 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py new file mode 100644 index 00000000..b5f58130 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -0,0 +1,373 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collide geometries.""" + +from typing import Callable, Dict, Optional, Sequence, Tuple, Union + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_base +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Candidate +from mujoco.mjx._src.collision_base import CandidateSet +from mujoco.mjx._src.collision_base import GeomInfo +from mujoco.mjx._src.collision_base import SolverParams +from mujoco.mjx._src.collision_convex import capsule_convex +from mujoco.mjx._src.collision_convex import convex_convex +from mujoco.mjx._src.collision_convex import plane_convex +from mujoco.mjx._src.collision_convex import sphere_convex +from mujoco.mjx._src.collision_primitive import capsule_capsule +from mujoco.mjx._src.collision_primitive import plane_capsule +from mujoco.mjx._src.collision_primitive import plane_sphere +from mujoco.mjx._src.collision_primitive import sphere_capsule +from mujoco.mjx._src.collision_primitive import sphere_sphere +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 GeomType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +# pair-wise collision functions +_COLLISION_FUNC = { + (GeomType.PLANE, GeomType.SPHERE): plane_sphere, + (GeomType.PLANE, GeomType.CAPSULE): plane_capsule, + (GeomType.PLANE, GeomType.BOX): plane_convex, + (GeomType.PLANE, GeomType.MESH): plane_convex, + (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere, + (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule, + (GeomType.SPHERE, GeomType.BOX): sphere_convex, + (GeomType.SPHERE, GeomType.MESH): sphere_convex, + (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule, + (GeomType.CAPSULE, GeomType.BOX): capsule_convex, + (GeomType.CAPSULE, GeomType.MESH): capsule_convex, + (GeomType.BOX, GeomType.BOX): convex_convex, + (GeomType.BOX, GeomType.MESH): convex_convex, + (GeomType.MESH, GeomType.MESH): convex_convex, +} + + +def get_collision_fn( + key: Tuple[Union[GeomType, mujoco.mjtGeom], Union[GeomType, mujoco.mjtGeom]] +) -> Optional[Callable[[GeomInfo, GeomInfo], collision_base.Contact]]: + """Returns a collision function given a pair of geom types.""" + return _COLLISION_FUNC.get(key, None) + + +def _add_candidate( + result: CandidateSet, + m: Union[Model, mujoco.MjModel], + g1: int, + g2: int, + ipair: int = -1, +): + """Adds a candidate to test for collision.""" + t1, t2 = m.geom_type[g1], m.geom_type[g2] + if t1 > t2: + t1, t2, g1, g2 = t2, t1, g2, g1 + + def mesh_key(i): + convex_data = [[None] * m.ngeom] * 3 + if isinstance(m, Model): + convex_data = [m.geom_convex_face, m.geom_convex_vert, m.geom_convex_edge] + key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data) + return key + + k1, k2 = mesh_key(g1), mesh_key(g2) + + candidates = {(c.geom1, c.geom2) for c in result.get((t1, t2, k1, k2), [])} + if (g1, g2) in candidates: + return + + if ipair > -1: + candidate = Candidate(g1, g2, ipair, -1, m.pair_dim[ipair]) + elif m.geom_priority[g1] != m.geom_priority[g2]: + gp = g1 if m.geom_priority[g1] > m.geom_priority[g2] else g2 + candidate = Candidate(g1, g2, -1, gp, m.geom_condim[gp]) + else: + dim = max(m.geom_condim[g1], m.geom_condim[g2]) + candidate = Candidate(g1, g2, -1, -1, dim) + + result.setdefault((t1, t2, k1, k2), []).append(candidate) + + +def _pair_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params for pair geoms.""" + ipair = jp.array([c.ipair for c in candidates]) + friction = jp.clip(m.pair_friction[ipair], a_min=mujoco.mjMINMU) + solref = m.pair_solref[ipair] + solreffriction = m.pair_solreffriction[ipair] + solimp = m.pair_solimp[ipair] + margin = m.pair_margin[ipair] + gap = m.pair_gap[ipair] + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _priority_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params from priority geoms.""" + geomp = jp.array([c.geomp for c in candidates]) + friction = m.geom_friction[geomp][:, jp.array([0, 0, 1, 2, 2])] + solref = m.geom_solref[geomp] + solreffriction = jp.zeros(geomp.shape + (mujoco.mjNREF,)) + solimp = m.geom_solimp[geomp] + g = jp.array([(c.geom1, c.geom2) for c in candidates]) + margin = jp.amax(m.geom_margin[g.T], axis=0) + gap = jp.amax(m.geom_gap[g.T], axis=0) + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _dynamic_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params for dynamic geoms.""" + g1 = jp.array([c.geom1 for c in candidates]) + g2 = jp.array([c.geom2 for c in candidates]) + + friction = jp.maximum(m.geom_friction[g1], m.geom_friction[g2]) + # copy friction terms for the full geom pair + friction = friction[:, jp.array([0, 0, 1, 2, 2])] + + minval = jp.array(mujoco.mjMINVAL) + solmix1, solmix2 = m.geom_solmix[g1], m.geom_solmix[g2] + mix = solmix1 / (solmix1 + solmix2) + mix = jp.where((solmix1 < minval) & (solmix2 < minval), 0.5, mix) + mix = jp.where((solmix1 < minval) & (solmix2 >= minval), 0.0, mix) + mix_fn = jax.vmap(lambda a, b, m: m * a + (1 - m) * b) + + solref1, solref2 = m.geom_solref[g1], m.geom_solref[g2] + solref = jp.minimum(solref1, solref2) + s_mix = mix_fn(solref1, solref2, mix) + solref = jp.where((solref1[0] > 0) & (solref2[0] > 0), s_mix, solref) + solreffriction = jp.zeros(g1.shape + (mujoco.mjNREF,)) + solimp = mix_fn(m.geom_solimp[g1], m.geom_solimp[g2], mix) + margin = jp.maximum(m.geom_margin[g1], m.geom_margin[g2]) + gap = jp.maximum(m.geom_gap[g1], m.geom_gap[g2]) + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _pair_info( + m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int] +) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]: + """Returns geom pair info for calculating collision.""" + g1, g2 = jp.array(geom1), jp.array(geom2) + info1 = GeomInfo( + d.geom_xpos[g1], + d.geom_xmat[g1], + m.geom_size[g1], + ) + info2 = GeomInfo( + d.geom_xpos[g2], + d.geom_xmat[g2], + m.geom_size[g2], + ) + in_axes1 = in_axes2 = jax.tree_map(lambda x: 0, info1) + if m.geom_convex_face[geom1[0]] is not None: + info1 = info1.replace( + face=jp.stack([m.geom_convex_face[i] for i in geom1]), + vert=jp.stack([m.geom_convex_vert[i] for i in geom1]), + edge=jp.stack([m.geom_convex_edge[i] for i in geom1]), + facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom1]), + ) + in_axes1 = in_axes1.replace(face=0, vert=0, edge=0, facenorm=0) + if m.geom_convex_face[geom2[0]] is not None: + info2 = info2.replace( + face=jp.stack([m.geom_convex_face[i] for i in geom2]), + vert=jp.stack([m.geom_convex_vert[i] for i in geom2]), + edge=jp.stack([m.geom_convex_edge[i] for i in geom2]), + facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom2]), + ) + in_axes2 = in_axes2.replace(face=0, vert=0, edge=0, facenorm=0) + return info1, info2, [in_axes1, in_axes2] + + +def _body_pair_filter( + m: Union[Model, mujoco.MjModel], b1: int, b2: int +) -> bool: + """Filters body pairs for collision.""" + dsbl_filterparent = m.opt.disableflags & DisableBit.FILTERPARENT + weld1 = m.body_weldid[b1] + weld2 = m.body_weldid[b2] + parent_weld1 = m.body_weldid[m.body_parentid[weld1]] + parent_weld2 = m.body_weldid[m.body_parentid[weld2]] + + if weld1 == weld2: + # filter out self-collisions + return True + + if ( + not dsbl_filterparent + and weld1 != 0 + and weld2 != 0 + and (weld1 == parent_weld2 or weld2 == parent_weld1) + ): + # filter out parent-child collisions + return True + + return False + + +def _collide_geoms( + m: Model, + d: Data, + geom_types: Tuple[GeomType, GeomType], + candidates: Sequence[Candidate], +) -> Contact: + """Collides a geom pair.""" + fn = get_collision_fn(geom_types) + if not fn: + return Contact.zero() + + # group sol params by different candidate types + typ_cands = {} + for c in candidates: + typ = (c.ipair > -1, c.geomp > -1) + typ_cands.setdefault(typ, []).append(c) + + geom1, geom2, params = [], [], [] + for (pair, priority), candidates in typ_cands.items(): + geom1.extend([c.geom1 for c in candidates]) + geom2.extend([c.geom2 for c in candidates]) + if pair: + params.append(_pair_params(m, candidates)) + elif priority: + params.append(_priority_params(m, candidates)) + else: + params.append(_dynamic_params(m, candidates)) + + # call contact function + g1, g2, in_axes = _pair_info(m, d, geom1, geom2) + res = jax.vmap(fn, in_axes=in_axes)(g1, g2) + dist, pos, frame = jax.tree_map(jp.concatenate, res) + + params = jax.tree_map(lambda *x: jp.concatenate(x), *params) + geom1, geom2 = jp.array(geom1), jp.array(geom2) + # repeat params by the number of contacts per geom pair + n_repeat = dist.shape[-1] // geom1.shape[0] + geom1, geom2, params = jax.tree_map( + lambda x: jp.repeat(x, n_repeat, axis=0), + (geom1, geom2, params), + ) + + con = Contact( + dist=dist, + pos=pos, + frame=frame, + includemargin=params.margin - params.gap, + friction=params.friction, + solref=params.solref, + solreffriction=params.solreffriction, + solimp=params.solimp, + geom1=geom1, + geom2=geom2, + dim=np.array([]), + efc_address=np.array([]), + ) + return con + + +def _max_contact_points(m: Model) -> int: + """Returns the maximum number of contact points when set as a numeric.""" + for i in range(m.nnumeric): + name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0] + if name == 'max_contact_points': + return int(m.numeric_data[m.numeric_adr[i]]) + + return -1 + + +def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: + """Returns candidates for collision checking.""" + candidate_set = {} + + for ipair in range(m.npair): + g1, g2 = m.pair_geom1[ipair], m.pair_geom2[ipair] + _add_candidate(candidate_set, m, g1, g2, ipair) + + body_pairs = [] + exclude_signature = set(m.exclude_signature) + for b1 in range(m.nbody): + for b2 in range(b1, m.nbody): + signature = (b1 << 16) + (b2) + if signature in exclude_signature: + continue + if _body_pair_filter(m, b1, b2): + continue + body_pairs.append((b1, b2)) + + for b1, b2 in body_pairs: + start1 = m.body_geomadr[b1] + end1 = m.body_geomadr[b1] + m.body_geomnum[b1] + for g1 in range(start1, end1): + start2 = m.body_geomadr[b2] + end2 = m.body_geomadr[b2] + m.body_geomnum[b2] + for g2 in range(start2, end2): + mask = m.geom_contype[g1] & m.geom_conaffinity[g2] + mask |= m.geom_contype[g2] & m.geom_conaffinity[g1] + if mask != 0: + _add_candidate(candidate_set, m, g1, g2) + + return candidate_set + + +def ncon(m: Model) -> int: + """Returns the number of contacts computed in MJX given a model.""" + 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 + ]) + return min(max_count, count) if max_count > -1 else count + + +def collision(m: Model, d: Data) -> Data: + """Collides geometries.""" + candidate_set = collision_candidates(m) + + contacts = [] + for key, candidates in candidate_set.items(): + geom_types = key[0:2] + contacts.append(_collide_geoms(m, d, geom_types, candidates)) + + if not contacts: + return d.replace(contact=Contact.zero(), ncon=0) + + contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts) + + max_contact_points = _max_contact_points(m) + if max_contact_points > -1 and contact.dist.shape[0] > max_contact_points: + # get top-k contacts + _, 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] + ns = d.ne + d.nf + d.nl + # 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_) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py new file mode 100644 index 00000000..14e14fb7 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -0,0 +1,499 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests the collision driver.""" + +import dataclasses +from typing import Dict, Optional, Tuple + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import jax.numpy as jp +import mujoco +from mujoco import mjx +# 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 Model +# pylint: emable=g-importing-member +import numpy as np + + +def _assert_attr_eq(mjx_d, mj_d, attr, name, atol): + if attr == 'efc_address': + # we do not test efc_address since it gets set in constraint logic + return + err_msg = f'mismatch: {attr} in run: {name}' + mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr) + if attr == 'frame': + mj_d = mj_d.reshape((-1, 3, 3)) + if mjx_d.shape != mj_d.shape: + raise AssertionError(f'{attr} shape mismatch: {mjx_d.shape}, {mj_d.shape}') + np.testing.assert_allclose(mjx_d, mj_d, err_msg=err_msg, atol=atol) + + +def _collide( + mjcf: str, assets: Optional[Dict[str, str]] = None +) -> Tuple[mujoco.MjModel, mujoco.MjData, Model, Data]: + m = mujoco.MjModel.from_xml_string(mjcf, assets or {}) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + return d, dx + + +class SphereCollisionTest(parameterized.TestCase): + _SPHERE_PLANE = """ + + + + + + + + + + """ + + _SPHERE_SPHERE = """ + + + + + + + + + + + + + """ + + _SPHERE_CAP = """ + + + + + + + + + + + + + """ + + @parameterized.parameters( + ('sphere_plane', _SPHERE_PLANE), + ('sphere_sphere', _SPHERE_SPHERE), + ('sphere_cap', _SPHERE_CAP), + ) + def test_sphere(self, name, mjcf): + d, dx = _collide(mjcf) + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5) + + _SPHERE_CONVEX = """ + + + + + + + + + + + + + """ + + def test_sphere_convex(self): + d, dx = _collide(self._SPHERE_CONVEX) + + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, 'sphere_convex', 1e-4) + + +class CapsuleCollisionTest(parameterized.TestCase): + _CAP_PLANE = """ + + + + + + + + + + """ + + _CAP_CAP = """ + + + + + + + + + + + + + """ + + @parameterized.parameters( + ('capsule_plane', _CAP_PLANE), + ('capsule_capsule', _CAP_CAP), + ) + def test_capsule(self, name, mjcf): + d, dx = _collide(mjcf) + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4) + + _PARALLEL_CAP = """ + + + + + + + + + + + + + """ + + def test_parallel_capsules(self): + """Tests that two parallel capsules are colliding at the midpoint.""" + _, dx = _collide(self._PARALLEL_CAP) + + np.testing.assert_allclose(dx.contact.dist, -0.05) + np.testing.assert_allclose( + dx.contact.pos[0], + np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]), + atol=1e-5, + ) + np.testing.assert_allclose( + dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5 + ) + + _CAP_BOX = """ + + + + + + + + + + + + + """ + + def test_capsule_convex(self): + """Tests a capsule-convex collision for a face contact.""" + d, dx = _collide(self._CAP_BOX) + + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4) + + _CAP_EDGE_BOX = """ + + + + + + + + + + + + + """ + + def test_capsule_convex_edge(self): + """Tests a capsule-convex collision for an edge contact.""" + d, dx = _collide(self._CAP_EDGE_BOX) + + c = dx.contact + self.assertEqual(c.pos.shape[0], 2) + self.assertGreater(c.dist[1], 0) + # extract the contact point with penetration + c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = c.replace(dim=c.dim[np.array([0])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) + + +class ConvexTest(absltest.TestCase): + """Tests the convex contact functions.""" + + _BOX_PLANE = """ + + + + + + + + + + """ + + def test_box_plane(self): + """Tests box collision with a plane.""" + d, dx = _collide(self._BOX_PLANE) + + np.testing.assert_array_less(dx.contact.dist[:2], 0) + np.testing.assert_array_less(-dx.contact.dist[2:], 0) + # extract the contact points with penetration + c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) + c = c.replace(dim=c.dim[np.array([0, 1])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2) + + _BOX_BOX = """ + + + + + + + + + + + + + """ + + def test_box_box(self): + """Tests a face contact for a box-box collision.""" + d, dx = _collide(self._BOX_BOX) + c = dx.contact + + self.assertEqual(c.pos.shape[0], 4) + self.assertTrue((c.dist < 0).all()) + np.testing.assert_array_almost_equal(c.pos[:, 2], np.array([0.39] * 4), 2) + np.testing.assert_array_almost_equal( + c.frame[:, 0, :], np.array([[0.0, 0.0, 1.0]] * 4) + ) + np.testing.assert_array_almost_equal( + c.frame.reshape((-1, 9)), d.contact.frame[:4, :] + ) + + _BOX_BOX_EDGE = """ + + + + + + + + + + + + + """ + + def test_box_box_edge(self): + """Tests an edge contact for a box-box collision.""" + d, dx = _collide(self._BOX_BOX_EDGE) + + # Only one contact point. + np.testing.assert_array_less(dx.contact.dist[:1], 0) + np.testing.assert_array_less(-dx.contact.dist[1:], 0) + # extract the contact point with penetration + c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = c.replace(dim=c.dim[np.array([0])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) + + _CONVEX_CONVEX = """ + + + + + + + + + + + + + + + + + """ + + def test_convex_convex(self): + """Tests generic convex-convex collision.""" + directory = epath.resource_path('mujoco.mjx') + assets = { + 'meshes/tetrahedron.stl': ( + directory / 'test_data' / 'meshes/tetrahedron.stl' + ).read_bytes(), + 'meshes/dodecahedron.stl': ( + directory / 'test_data' / 'meshes/dodecahedron.stl' + ).read_bytes(), + } + _, dx = _collide(self._CONVEX_CONVEX, assets=assets) + c = dx.contact + + # Only one contact point for an edge contact. + self.assertLess(c.dist[0], 0) + self.assertTrue((c.dist[1:] > 0).all()) + np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1])) + + +class BodyPairFilterTest(absltest.TestCase): + """Tests that certain body pairs get filtered.""" + + _SELF_COLLISION = """ + + + + + + + + + + """ + + def test_filter_self_collision(self): + """Tests that self collisions get filtered.""" + d, dx = _collide(self._SELF_COLLISION) + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 0) + + _PARENT_CHILD = """ + + + + + + + + + + + + + """ + + def test_filter_parent_child(self): + """Tests that parent-child collisions get filtered.""" + m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 0) + + def test_disable_filter_parent_child(self): + """Tests that filterparent flag disables parent-child filtering.""" + m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) + m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_FILTERPARENT + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + # one collision between parent-child spheres + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 1) + + +class TopKContactTest(absltest.TestCase): + """Tests top-k contacts.""" + + _CAPSULES = """ + + + + + + + + + + + + + + + + + + + + """ + + def test_top_k_contacts(self): + m = mujoco.MjModel.from_xml_string(self._CAPSULES) + mx_top_k = mjx.device_put(m) + mx_all = mx_top_k.replace( + nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([]) + ) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx_all, dx) + + dx_all = collision_jit_fn(mx_all, dx) + dx_top_k = collision_jit_fn(mx_top_k, dx) + + self.assertEqual(dx_all.ncon, 3) + self.assertEqual(dx_top_k.ncon, 2) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py new file mode 100644 index 00000000..8c2dd529 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_primitive.py @@ -0,0 +1,126 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collision primitives.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import math +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Contact +from mujoco.mjx._src.collision_base import GeomInfo +# pylint: enable=g-importing-member + + +def _plane_sphere( + plane_normal: jax.Array, + plane_pos: jax.Array, + sphere_pos: jax.Array, + radius: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Returns the penetration and contact point between a plane and sphere.""" + cdist = jp.dot(sphere_pos - plane_pos, plane_normal) + dist = cdist - radius + pos = sphere_pos - plane_normal * (radius + 0.5 * dist) + return dist, pos + + +def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact: + """Calculates contact between a plane and a sphere.""" + n = plane.mat[:, 2] + dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact: + """Calculates two contacts between a capsule and a plane.""" + n, axis = plane.mat[:, 2], cap.mat[:, 2] + # align contact frames with capsule axis + b, b_norm = math.normalize_with_norm(axis - n * jp.dot(n, axis)) + y, z = jp.array([0.0, 1.0, 0.0]), jp.array([0.0, 0.0, 1.0]) + b = jp.where(b_norm < 0.5, jp.where((-0.5 < n[1]) & (n[1] < 0.5), y, z), b) + frame = jp.array([[n, b, jp.cross(n, b)]]) + segment = axis * cap.size[1] + contacts = [] + for offset in [segment, -segment]: + dist, pos = _plane_sphere(n, plane.pos, cap.pos + offset, cap.size[0]) + dist = jp.expand_dims(dist, axis=0) + pos = jp.expand_dims(pos, axis=0) + contacts.append((dist, pos, frame)) + return jax.tree_map(lambda *x: jp.concatenate(x), *contacts) + + +def _sphere_sphere( + pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array +) -> Contact: + """Returns the penetration, contact point, and normal between two spheres.""" + n, dist = math.normalize_with_norm(pos2 - pos1) + n = jp.where(dist == 0.0, jp.array([1.0, 0.0, 0.0]), n) + dist = dist - (radius1 + radius2) + pos = pos1 + n * (radius1 + dist * 0.5) + return dist, pos, n + + +def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Contact: + """Calculates contact between two spheres.""" + dist, pos, n = _sphere_sphere(s1.pos, s1.size[0], s2.pos, s2.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact: + """Calculates one contact between a sphere and a capsule.""" + axis, length = cap.mat[:, 2], cap.size[1] + segment = axis * length + pt = math.closest_segment_point( + cap.pos - segment, cap.pos + segment, sphere.pos + ) + dist, pos, n = _sphere_sphere(sphere.pos, sphere.size[0], pt, cap.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact: + """Calculates one contact between two capsules.""" + axis1, length1, axis2, length2 = ( + cap1.mat[:, 2], + cap1.size[1], + cap2.mat[:, 2], + cap2.size[1], + ) + seg1, seg2 = axis1 * length1, axis2 * length2 + pt1, pt2 = math.closest_segment_to_segment_points( + cap1.pos - seg1, + cap1.pos + seg1, + cap2.pos - seg2, + cap2.pos + seg2, + ) + radius1, radius2 = cap1.size[0], cap2.size[0] + dist, pos, n = _sphere_sphere(pt1, radius1, pt2, radius2) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + +# store ncon as function attributes +plane_sphere.ncon = 1 +plane_capsule.ncon = 2 +sphere_sphere.ncon = 1 +sphere_capsule.ncon = 1 +capsule_capsule.ncon = 1 diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py new file mode 100644 index 00000000..bd0bbaa6 --- /dev/null +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -0,0 +1,340 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Core non-smooth constraint functions.""" + +from typing import 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 +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 EqType +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +class _Efc(PyTreeNode): + J: jax.Array + R: jax.Array + aref: 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, + solref: jax.Array, + solimp: jax.Array, + pos: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Calculates stiffness, damping, and impedance of a constraint.""" + timeconst, dampratio = solref + + if not m.opt.disableflags & DisableBit.REFSAFE: + timeconst = jp.maximum(timeconst, 2 * m.opt.timestep) * (timeconst > 0) + + dmin, dmax, width, mid, power = solimp + + dmin = jp.clip(dmin, mujoco.mjMINIMP, mujoco.mjMAXIMP) + dmax = jp.clip(dmax, mujoco.mjMINIMP, mujoco.mjMAXIMP) + width = jp.maximum(0, width) + mid = jp.clip(mid, mujoco.mjMINIMP, mujoco.mjMAXIMP) + power = jp.maximum(1, power) + + # See https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters + k = 1 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) + b = 2 / (dmax * timeconst) + # TODO(robotics-simulation): check various solparam settings in model gen test + k = jp.where(dampratio <= 0, -dampratio / (dmax * dmax), k) + b = jp.where(timeconst <= 0, -timeconst / dmax, b) + + imp_x = jp.abs(pos) / width + imp_a = (1.0 / jp.power(mid, power - 1)) * jp.power(imp_x, power) + imp_b = 1 - (1.0 / jp.power(1 - mid, power - 1)) * jp.power(1 - imp_x, power) + imp_y = jp.where(imp_x < mid, imp_a, imp_b) + imp = dmin + imp_y * (dmax - dmin) + imp = jp.clip(imp, dmin, dmax) + imp = jp.where(imp_x > 1.0, dmax, imp) + + 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.""" + + if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0: + return _Efc.zero(m) + + connect_id = np.nonzero(m.eq_type == EqType.CONNECT)[0] + + if connect_id.size == 0: + return _Efc.zero(m) + + 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): + anchor1, anchor2 = data[0:3], data[3:6] + # find global points + pos1 = d.xmat[id1] @ anchor1 + d.xpos[id1] + pos2 = d.xmat[id2] @ anchor2 + d.xpos[id2] + + # compute position error + cpos = pos1 - pos2 + + # compute Jacobian difference (opposite of contact: 0 - 1) + jacp1, _ = support.jac(m, d, pos1, id1) + 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 _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp) + + return jax.tree_map(jp.concatenate, efcs) + + +def _instantiate_weld(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporting data for connect weld constraints.""" + + if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0: + return _Efc.zero(m) + + weld_id = np.nonzero(m.eq_type == EqType.WELD)[0] + + if weld_id.size == 0: + return _Efc.zero(m) + + 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): + anchor1, anchor2 = data[0:3], data[3:6] + relpose, torquescale = data[6:10], data[10] + + # find global points + pos1 = d.xmat[id1] @ anchor2 + d.xpos[id1] + pos2 = d.xmat[id2] @ anchor1 + d.xpos[id2] + + # compute position error + cpos = pos1 - pos2 + + # compute Jacobian difference (opposite of contact: 0 - 1) + jacp1, jacr1 = support.jac(m, d, pos1, id1) + jacp2, jacr2 = support.jac(m, d, pos2, id2) + jacdifp = jacp1 - jacp2 + jacdifr = (jacr1 - jacr2) * torquescale + + # compute orientation error: neg(q1) * q0 * relpose (axis components only) + quat = math.quat_mul(d.xquat[id1], relpose) + quat1 = math.quat_inv(d.xquat[id2]) + crot = math.quat_mul(quat1, quat)[1:] # copy axis components + + # correct rotation Jacobian: 0.5 * neg(q1) * (jac0-jac1) * q0 * relpose + jac_fn = lambda j: math.quat_mul(math.quat_mul_axis(quat1, j), quat)[1:] + jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr) + + j = jp.concatenate((jacdifp.T, jacdifr.T)) + pos = jp.concatenate((cpos, crot)) + + # 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 _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp) + + return jax.tree_map(jp.concatenate, efcs) + + +def _instantiate_friction(m: Model, d: Data) -> _Efc: + # TODO(robotics-team): implement _instantiate_friction + del d + return _Efc.zero(m) + + +def _instantiate_limit(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporting data for joint limits.""" + + if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any(): + return _Efc.zero(m) + + def fn(jnt_typs, jnt_range, solref, solimp, margin, qpos, dofs, invweight0): + js, rs, arefs = [], [], [] + qpos_i, dof_i = 0, 0 + + for i in range(len(jnt_typs)): + jnt_typ = JointType(jnt_typs[i]) + + 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}') + + 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() + + 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)) + + +def _instantiate_contact(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporitng data 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) + + 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) + t = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0] + + # rotate Jacobian differences to contact frame + diff_con = contact.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]): + 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) + + 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)) + + return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + return jax.tree_map(jp.concatenate, jax.vmap(fn)(d.contact)) + + +def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: + """Returns equality, friction, limit, and contact constraint counts.""" + if m.opt.disableflags & DisableBit.CONSTRAINT: + return 0, 0, 0, 0 + + 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 + + nf = 0 + + if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any(): + nl = 0 + else: + nl = (m.jnt_type != JointType.FREE).sum() + + if (m.opt.disableflags & DisableBit.CONTACT): + nc = 0 + else: + nc = d.ncon * 4 + + return ne, nf, nl, nc + + +def make_constraint(m: Model, d: Data) -> Data: + """Creates constraint jacobians and other supporting data.""" + + ns = sum(count_constraints(m, d)[:-1]) + # TODO(robotics-simulation): make device_put set nefc/efc_address instead + 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) + else: + efcs = ( + _instantiate_connect(m, d), + _instantiate_weld(m, d), + _instantiate_friction(m, d), + _instantiate_limit(m, d), + _instantiate_contact(m, d), + ) + efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs) + + 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], + ) + + return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py new file mode 100644 index 00000000..ce289df7 --- /dev/null +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -0,0 +1,144 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for constraint functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import constraint +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(a, b, name, step, fname, atol=1e-3, rtol=1e-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) + + +class ConstraintTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(test_util.TEST_FILES)) + def testconstraints(self, seed, fname): + """Test constraints.""" + np.random.seed(seed) + + # exclude convex.xml since convex contacts are not exactly equivalent + if fname == 'convex.xml': + return + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + for i in range(100): + dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel)) + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + nnz_filter = dx.efc_J.any(axis=1) + + mj_efc_j = d.efc_J.reshape((-1, m.nv)) + mjx_efc_j = dx.efc_J[nnz_filter] + _assert_eq(mj_efc_j, mjx_efc_j, 'efc_J', i, fname) + + mjx_efc_d = dx.efc_D[nnz_filter] + _assert_eq(d.efc_D, mjx_efc_d, 'efc_D', i, fname) + + mjx_efc_aref = dx.efc_aref[nnz_filter] + _assert_eq(d.efc_aref, mjx_efc_aref, 'efc_aref', i, fname) + + mjx_efc_frictionloss = dx.efc_frictionloss[nnz_filter] + _assert_eq( + d.efc_frictionloss, + mjx_efc_frictionloss, + 'efc_frictionloss', + i, + fname, + ) + + def test_disable_refsafe(self): + m = test_util.load_test_file('ant.xml') + + timeconst = m.opt.timestep / 4.0 # timeconst < 2 * timestep + solimp = jp.array([timeconst, 1.0]) + solref = jp.array([0.8, 0.99, 0.001, 0.2, 2]) + pos = jp.ones(3) + + m.opt.disableflags = m.opt.disableflags | DisableBit.REFSAFE + mx = mjx.device_put(m) + k, *_ = constraint._kbi(mx, solimp, solref, pos) + self.assertEqual(k, 1 / (0.99**2 * timeconst**2)) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.REFSAFE + mx = mjx.device_put(m) + k, *_ = constraint._kbi(mx, solimp, solref, pos) + self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2)) + + def test_disableconstraint(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = constraint.make_constraint(mx, dx) + self.assertGreater(dx.efc_J.shape[0], 1) + + m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT + mx = mjx.device_put(m) + dx = constraint.make_constraint(mx, dx) + self.assertEqual(dx.efc_J.shape[0], 0) + + def test_disable_equality(self): + m = test_util.load_test_file('weld.xml') + d = mujoco.MjData(m) + + m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = constraint.make_constraint(mx, dx) + self.assertEqual(dx.efc_J.shape[0], 0) + + def test_disable_contact(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + d.qpos[2] = 0.0 + mujoco.mj_forward(m, d) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = dx.tree_replace( + {'contact.frame': dx.contact.frame.reshape((-1, 3, 3))} + ) + efc = constraint._instantiate_contact(mx, dx) + self.assertIsNotNone(efc) + + 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) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py new file mode 100644 index 00000000..0936bec6 --- /dev/null +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -0,0 +1,151 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Wrapper that automatically registers dataclass as a Jax PyTree.""" + +import copy +import dataclasses + +import typing +from typing import Dict, Optional, Sequence, TypeVar +import jax +import numpy as np + +_T = TypeVar('_T') + + +def dataclass(clz: _T) -> _T: + """Wraps a dataclass with metadata for which fields are pytrees. + + This is based off flax.struct.dataclass, but instead of using field + descriptors to specify which fields are pytrees, we follow a simple rule: + a leaf field is a pytree node if and only if it's a jax.Array + + Args: + clz: the class to register as a dataclass + + Returns: + the resulting dataclass, registered with Jax + """ + data_clz = dataclasses.dataclass(frozen=True)(clz) + meta_fields, data_fields = [], [] + for field in dataclasses.fields(data_clz): + if any(( + field.type is jax.Array, + dataclasses.is_dataclass(field.type), + jax.Array in typing.get_args(field.type), + )): + data_fields.append(field) + else: + meta_fields.append(field) + + def replace(self, **updates): + """"Returns a new object replacing the specified fields with new values.""" + return dataclasses.replace(self, **updates) + + data_clz.replace = replace + + def iterate_clz_with_keys(x): + # numpy arrays are not hashable, so convert them to tuples for jit cache + to_tup = lambda x: tuple(x) if len(x.shape) == 1 else tuple(map(to_tup, x)) + + def to_meta(field, obj): + val = getattr(obj, field.name) + return to_tup(val) if isinstance(val, np.ndarray) else val + + def to_data(field, obj): + return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name)) + + data = tuple(to_data(f, x) for f in data_fields) + meta = tuple(to_meta(f, x) for f in meta_fields) + return data, meta + + def clz_from_iterable(meta, data): + + def from_meta(field, meta): + if field.type is np.ndarray: + return (field.name, np.array(meta)) + else: + return (field.name, meta) + + from_data = lambda field, meta: (field.name, meta) + + meta_args = tuple(from_meta(f, m) for f, m in zip(meta_fields, meta)) + data_args = tuple(from_data(f, m) for f, m in zip(data_fields, data)) + + return data_clz(**dict(meta_args + data_args)) + + jax.tree_util.register_pytree_with_keys( + data_clz, iterate_clz_with_keys, clz_from_iterable + ) + + return data_clz + + +TNode = TypeVar('TNode', bound='PyTreeNode') + + +class PyTreeNode: + """Base class for dataclasses that should act like a JAX pytree node. + + This base class additionally avoids type checking errors when using PyType. + """ + + def __init_subclass__(cls): + dataclass(cls) + + def __init__(self, *args, **kwargs): + # stub for pytype + raise NotImplementedError + + def replace(self: TNode, **overrides) -> TNode: + # stub for pytype + raise NotImplementedError + + def tree_replace( + self, params: Dict[str, Optional[jax.typing.ArrayLike]] + ) -> 'PyTreeNode': + new = self + for k, v in params.items(): + new = _tree_replace(new, k.split('.'), v) + return new + + +def _tree_replace( + base: PyTreeNode, + attr: Sequence[str], + val: Optional[jax.typing.ArrayLike], +) -> PyTreeNode: + """Sets attributes in a struct.dataclass with values.""" + if not attr: + return base + + # special case for List attribute + if len(attr) > 1 and isinstance(getattr(base, attr[0]), list): + lst = copy.deepcopy(getattr(base, attr[0])) + + for i, g in enumerate(lst): + if not hasattr(g, attr[1]): + continue + v = val if not hasattr(val, '__iter__') else val[i] + lst[i] = _tree_replace(g, attr[1:], v) + + return base.replace(**{attr[0]: lst}) + + if len(attr) == 1: + return base.replace(**{attr[0]: val}) + + return base.replace( + **{attr[0]: _tree_replace(getattr(base, attr[0]), attr[1:], val)} + ) diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py new file mode 100644 index 00000000..3c3c213f --- /dev/null +++ b/mjx/mujoco/mjx/_src/device.py @@ -0,0 +1,282 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Get and put mujoco data on/off device.""" + +import copy +import dataclasses +from typing import Any, Dict, Iterable, List, Union, overload +import warnings + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import mesh +from mujoco.mjx._src import types + +_MJ_TYPE_ATTR = { + mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,), + mujoco.mjtDyn: (mujoco.MjModel.actuator_dyntype,), + mujoco.mjtEq: (mujoco.MjModel.eq_type,), + mujoco.mjtGain: (mujoco.MjModel.actuator_gaintype,), + mujoco.mjtTrn: (mujoco.MjModel.actuator_trntype,), + mujoco.mjtCone: ( + mujoco.MjModel.opt, + mujoco.MjOption.cone, + ), + mujoco.mjtIntegrator: ( + mujoco.MjModel.opt, + mujoco.MjOption.integrator, + ), +} + +_TYPE_MAP = { + mujoco._structs._MjContactList: types.Contact, # pylint: disable=protected-access + mujoco.MjData: types.Data, + mujoco.MjModel: types.Model, + mujoco.MjOption: types.Option, + mujoco.MjStatistic: types.Statistic, + mujoco.mjtBias: types.BiasType, + mujoco.mjtCone: types.ConeType, + mujoco.mjtDisableBit: types.DisableBit, + mujoco.mjtDyn: types.DynType, + mujoco.mjtEq: types.EqType, + mujoco.mjtGain: types.GainType, + mujoco.mjtIntegrator: types.IntegratorType, + mujoco.mjtSolver: types.SolverType, + mujoco.mjtTrn: types.TrnType, +} + +_TRANSFORMS = { + (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Model, 'actuator_trnid'): lambda x: x[:, 0], + (types.Contact, 'frame'): ( + lambda x: x.reshape(x.shape[:-1] + (3, 3)) # pylint: disable=g-long-lambda + if x is not None and x.shape[0] else jp.zeros((0, 3, 3)) + ), +} + +_INVERSE_TRANSFORMS = { + (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Contact, 'frame'): ( + lambda x: x.reshape(x.shape[:-2] + (9,)) # pylint: disable=g-long-lambda + if x is not None and x.shape[0] else jp.zeros((0, 9)) + ), +} + +_DERIVED = mesh.DERIVED.union( + # efc_J is dense in MJX, sparse in MJ. ignore for now. + {(types.Data, 'efc_J'), (types.Option, 'has_fluid_params')} +) + + +def _model_derived(value: mujoco.MjModel) -> Dict[str, Any]: + return {k: jax.device_put(v) for k, v in mesh.get(value).items()} + + +def _data_derived(value: mujoco.MjData) -> Dict[str, Any]: + return {'efc_J': jax.device_put(value.efc_J)} + + +def _option_derived(value: types.Option) -> Dict[str, Any]: + has_fluid = ( + value.density > 0 or value.viscosity > 0 or (value.wind != 0.0).any() + ) + return {'has_fluid_params': has_fluid} + + +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(): + val = m + for attr in attrs: + val = attr.fget(val) # pytype: disable=attribute-error + + 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.') + + # check condim + if any(dim != 3 for dim in m.geom_condim) or any( + dim != 3 for dim in m.pair_dim + ): + raise NotImplementedError('Only condim=3 is supported.') + + # check collision geom types + candidate_set = collision_driver.collision_candidates(m) + for g1, g2, *_ in candidate_set: + g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) + if g1 == mujoco.mjtGeom.mjGEOM_PLANE and g2 in ( + mujoco.mjtGeom.mjGEOM_PLANE, + mujoco.mjtGeom.mjGEOM_HFIELD, + ): + # MuJoCo does not collide planes with other planes or hfields + continue + if collision_driver.get_collision_fn((g1, g2)) is None: + raise NotImplementedError(f'({g1}, {g2}) collisions not implemented.') + + # TODO(erikfrey): warn for high solver iterations, nefc, etc. + + # mjNDISABLE is not a DisableBit flag, so must be explicitly ignored + disablebit_members = set(mujoco.mjtDisableBit.__members__.values()) - { + mujoco.mjtDisableBit.mjNDISABLE} + unsupported_disable = disablebit_members - { + mujoco.mjtDisableBit(t.value) for t in types.DisableBit + } + for f in unsupported_disable: + if f & m.opt.disableflags: + warnings.warn(f'Ignoring disable flag {f.name}.') + + # mjNENABLE is not an EnableBit flag, so must be explicitly ignored + unsupported_enable = set(mujoco.mjtEnableBit.__members__.values()) - { + mujoco.mjtEnableBit.mjNENABLE + } + for f in unsupported_enable: + if f & m.opt.enableflags: + warnings.warn(f'Ignoring enable flag {f.name}.') + + +@overload +def device_put(value: mujoco.MjData) -> types.Data: + ... + + +@overload +def device_put(value: mujoco.MjModel) -> types.Model: + ... + + +def device_put(value): + """Places mujoco data onto a device. + + Args: + value: a mujoco struct to transfer + + Returns: + on-device MJX struct reflecting the input value + """ + clz = _TYPE_MAP.get(type(value)) + if clz is None: + raise NotImplementedError(f'{type(value)} is not supported for device_put.') + + if isinstance(value, mujoco.MjModel): + _validate(value) # type: ignore + + init_kwargs = {} + for f in dataclasses.fields(clz): # type: ignore + if (clz, f.name) in _DERIVED: + continue + + field_value = getattr(value, f.name) + if (clz, f.name) in _TRANSFORMS: + field_value = _TRANSFORMS[(clz, f.name)](field_value) + + if f.type is jax.Array: + field_value = jax.device_put(field_value) + elif type(field_value) in _TYPE_MAP.keys(): + field_value = device_put(field_value) + + init_kwargs[f.name] = copy.copy(field_value) + + derived_kwargs = {} + if isinstance(value, mujoco.MjModel): + derived_kwargs = _model_derived(value) + elif isinstance(value, mujoco.MjData): + derived_kwargs = _data_derived(value) + elif isinstance(value, mujoco.MjOption): + derived_kwargs = _option_derived(value) + + return clz(**init_kwargs, **derived_kwargs) # type: ignore + + +@overload +def device_get_into( + result: Union[mujoco.MjData, List[mujoco.MjData]], value: types.Data +): + ... + + +def device_get_into(result, value): + """Transfers data off device into a mujoco MjData. + + Data on device often has a batch dimension which adds (N,) to the beginning + of each array shape where N = batch size. + + If result is a single MjData, arrays are copied over with the batch dimension + intact. If result is a list, the list must be length N and will be populated + with distinct MjData structs where the batch dimension is stripped. + + Args: + result: struct (or list of structs) to transfer into + value: device value to transfer + + Raises: + RuntimeError: if result length doesn't match data batch size + """ + + value = jax.device_get(value) + + if isinstance(result, list): + array_shapes = [s.shape for s in jax.tree_util.tree_flatten(value)[0]] + + if any(len(s) < 1 or s[0] != array_shapes[0][0] for s in array_shapes): + raise ValueError('unrecognizable batch dimension in value') + + batch_size = array_shapes[0][0] + + if len(result) != batch_size: + raise ValueError( + f"result length ({len(result)}) doesn't match value batch size" + f' ({batch_size})' + ) + + for i in range(batch_size): + value_i = jax.tree_map(lambda x, i=i: x[i], value) + device_get_into(result[i], value_i) + + else: + if isinstance(result, mujoco.MjData): + mujoco._functions._realloc_con_efc( # pylint: disable=protected-access + result, ncon=value.ncon, nefc=value.nefc + ) + + for f in dataclasses.fields(value): # type: ignore + if (type(value), f.name) in _DERIVED: + continue + + field_value = getattr(value, f.name) + + if (type(value), f.name) in _INVERSE_TRANSFORMS: + field_value = _INVERSE_TRANSFORMS[(type(value), f.name)](field_value) + + if type(field_value) in _TYPE_MAP.values(): + device_get_into(getattr(result, f.name), field_value) + continue + + try: + setattr(result, f.name, field_value) + except AttributeError: + getattr(result, f.name)[:] = field_value diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py new file mode 100644 index 00000000..bdc4e2d4 --- /dev/null +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -0,0 +1,184 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for moving mujoco structs on and off device.""" + +import dataclasses + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import device +from mujoco.mjx._src import test_util +from mujoco.mjx._src import types +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(testcase, a, b, attr=None, name=None): + if (type(a), attr) in device._DERIVED: + return + + if attr: + a, b = getattr(a, attr), getattr(b, attr) + + if isinstance(a, PyTreeNode): + for field in dataclasses.fields(a): + _assert_eq(testcase, a, b, field.name, type(a).__name__) + return + + typ = {'Model': types.Model, 'Data': types.Data, + 'Contact': types.Contact}.get(name) + if (typ, attr) in device._TRANSFORMS: + b = device._TRANSFORMS[(typ, attr)](b) + + err_msg = f'mismatch: {attr} in {name}' + if not hasattr(b, 'shape') or not b.shape: + testcase.assertEqual(a, b, err_msg) + return + + a, b = np.array(a), np.array(b) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=1e-8) + + +class DeviceTest(parameterized.TestCase): + + @parameterized.parameters(test_util.TEST_FILES) + def testdevice_put(self, fname): + """Test putting MjData and MjModel on device.""" + m = test_util.load_test_file(fname) + # advance state to ensure non-zero fields + d = mujoco.MjData(m) + for _ in range(10): + mujoco.mj_step(m, d) + + _assert_eq(self, mjx.device_put(d), d) + _assert_eq(self, mjx.device_put(m), m) + + @parameterized.parameters(test_util.TEST_FILES) + def testdevice_get(self, fname): + """Test getting MjData from a device.""" + m = test_util.load_test_file(fname) + mx = device.device_put(m) + dx = mjx.make_data(mx) + d = mujoco.MjData(m) + device.device_get_into(d, dx) + _assert_eq(self, dx, d) + + @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + def testdevice_get_batched(self, fname): + """Test getting MjData from a device.""" + m = test_util.load_test_file(fname) + mx = device.device_put(m) + batch_size = 32 + + # create mjx_data and batch it + dx = mjx.make_data(mx) + dx = jax.tree_map( + lambda x: jp.repeat(x, batch_size).reshape((batch_size,) + x.shape), + dx, + ) + ds = [mujoco.MjData(m) for _ in range(batch_size - 1)] + + with self.assertRaises(ValueError): + device.device_get_into(ds, dx) + + ds = [mujoco.MjData(m) for _ in range(batch_size)] + device.device_get_into(ds, dx) + dx = jax.device_get(dx) # faster indexing for testing + for i in range(batch_size): + _assert_eq(self, jax.tree_map(lambda x, i=i: x[i], dx), ds[i]) + + +class ValidateInputTest(absltest.TestCase): + + def test_solver(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertWarns(UserWarning): + mx = mjx.device_put(m) + self.assertEqual(mx.opt.solver, mujoco.mjtSolver.mjSOL_CG) + + def test_integrator(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_cone(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertRaises(NotImplementedError): + _ = 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) + + 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) + + 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) + + 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) + + 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) + + def test_geoms(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + """) + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py new file mode 100644 index 00000000..42faeca2 --- /dev/null +++ b/mjx/mujoco/mjx/_src/forward.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Forward step functions.""" + +import functools +from typing import Optional, Sequence + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import constraint +from mujoco.mjx._src import math +from mujoco.mjx._src import passive +from mujoco.mjx._src import scan +from mujoco.mjx._src import smooth +from mujoco.mjx._src import solver +from mujoco.mjx._src import support +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import BiasType +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import DynType +from mujoco.mjx._src.types import GainType +from mujoco.mjx._src.types import 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 + +# RK4 tableau +_RK4_A = np.array([ + [0.5, 0.0, 0.0], + [0.0, 0.5, 0.0], + [0.0, 0.0, 1.0], +]) +_RK4_B = np.array([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]) + + +def named_scope(fn, name: str = ''): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with jax.named_scope(name or getattr(fn, '__name__')): + res = fn(*args, **kwargs) + return res + + return wrapper + + +@named_scope +def _position(m: Model, d: Data) -> Data: + """Position-dependent computations.""" + # TODO(robotics-simulation): tendon + d = smooth.kinematics(m, d) + d = smooth.com_pos(m, d) + d = smooth.crb(m, d) + d = smooth.factor_m(m, d, d.qM) + d = collision_driver.collision(m, d) + d = constraint.make_constraint(m, d) + d = smooth.transmission(m, d) + return d + + +@named_scope +def _velocity(m: Model, d: Data) -> Data: + """Velocity-dependent computations.""" + d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel) + d = smooth.com_vel(m, d) + d = passive.passive(m, d) + d = smooth.rne(m, d) + return d + + +@named_scope +def _actuation(m: Model, d: Data) -> Data: + """Actuation-dependent computations.""" + if not m.nu or m.opt.disableflags & DisableBit.ACTUATION: + return d.replace( + act_dot=jp.zeros((m.na,)), + qfrc_actuator=jp.zeros((m.nv,)), + ) + + ctrl = d.ctrl + if not m.opt.disableflags & DisableBit.CLAMPCTRL: + ctrlrange = jp.where( + m.actuator_ctrllimited[:, None], + m.actuator_ctrlrange, + jp.array([-jp.inf, jp.inf]), + ) + ctrl = jp.clip(ctrl, ctrlrange[:, 0], ctrlrange[:, 1]) + + # act_dot for stateful actuators + def get_act_dot(dyn_typ, dyn_prm, ctrl, act): + if dyn_typ == DynType.NONE: + act_dot = jp.array(0.0) + elif dyn_typ == DynType.INTEGRATOR: + act_dot = ctrl + elif dyn_typ == DynType.FILTER: + act_dot = (ctrl - act) / jp.clip(dyn_prm[0], mujoco.mjMINVAL) + else: + raise NotImplementedError(f'dyntype {dyn_typ.name} not implemented.') + return act_dot + + act_dot = jp.zeros((m.na,)) + if m.na: + act_dot = scan.flat( + m, + get_act_dot, + 'uuua', + 'a', + m.actuator_dyntype, + m.actuator_dynprm, + ctrl, + d.act, + group_by='u', + ) + + ctrl_act = ctrl + if m.na: + act_last_dim = d.act[m.actuator_actadr + m.actuator_actnum - 1] + ctrl_act = jp.where(m.actuator_actadr == -1, ctrl, act_last_dim) + + def get_force(*args): + gain_t, gain_p, bias_t, bias_p, len_, vel, ctrl_act = args + + typ, prm = GainType(gain_t), gain_p + if typ == GainType.FIXED: + gain = prm[0] + elif typ == GainType.AFFINE: + gain = prm[0] + prm[1] * len_ + prm[2] * vel + else: + raise RuntimeError(f'unrecognized gaintype {typ.name}.') + + typ, prm = BiasType(bias_t), bias_p + bias = jp.array(0.0) + if typ == BiasType.AFFINE: + bias = prm[0] + prm[1] * len_ + prm[2] * vel + + return gain * ctrl_act + bias + + force = scan.flat( + m, + get_force, + 'uuuuuuu', + 'u', + m.actuator_gaintype, + m.actuator_gainprm, + m.actuator_biastype, + m.actuator_biasprm, + d.actuator_length, + d.actuator_velocity, + ctrl_act, + group_by='u', + ) + forcerange = jp.where( + m.actuator_forcelimited[:, None], + m.actuator_forcerange, + jp.array([-jp.inf, jp.inf]), + ) + force = jp.clip(force, forcerange[:, 0], forcerange[:, 1]) + + qfrc_actuator = d.actuator_moment.T @ force + + # clamp qfrc_actuator + actfrcrange = jp.where( + m.jnt_actfrclimited[:, None], + m.jnt_actfrcrange, + jp.array([-jp.inf, jp.inf]), + ) + ids = sum( + ([i] * JointType(j).dof_width() for i, j in enumerate(m.jnt_type)), [] + ) + actfrcrange = jp.take(actfrcrange, jp.array(ids), axis=0) + qfrc_actuator = jp.clip(qfrc_actuator, actfrcrange[:, 0], actfrcrange[:, 1]) + + d = d.replace(act_dot=act_dot, qfrc_actuator=qfrc_actuator) + return d + + +@named_scope +def _acceleration(m: Model, d: Data) -> Data: + """Add up all non-constraint forces, compute qacc_smooth.""" + qfrc_applied = d.qfrc_applied + support.xfrc_accumulate(m, d) + qfrc_smooth = d.qfrc_passive - d.qfrc_bias + d.qfrc_actuator + qfrc_applied + qacc_smooth = smooth.solve_m(m, d, qfrc_smooth) + d = d.replace(qfrc_smooth=qfrc_smooth, qacc_smooth=qacc_smooth) + return d + + +@named_scope +def _integrate_pos( + jnt_typs: Sequence[str], qpos: jax.Array, qvel: jax.Array, dt: jax.Array +) -> jax.Array: + """Integrate position given velocity.""" + qs, qi, vi = [], 0, 0 + + for jnt_typ in jnt_typs: + if jnt_typ == JointType.FREE: + pos = qpos[qi : qi + 3] + dt * qvel[vi : vi + 3] + quat = math.quat_integrate( + qpos[qi + 3 : qi + 7], qvel[vi + 3 : vi + 6], dt + ) + qs.append(jp.concatenate([pos, quat])) + qi, vi = qi + 7, vi + 6 + elif jnt_typ == JointType.BALL: + quat = math.quat_integrate(qpos[qi : qi + 4], qvel[vi : vi + 3], dt) + qs.append(quat) + qi, vi = qi + 4, vi + 3 + elif jnt_typ in (JointType.HINGE, JointType.SLIDE): + pos = qpos[qi] + dt * qvel[vi] + qs.append(pos[None]) + qi, vi = qi + 1, vi + 1 + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + return jp.concatenate(qs) if qs else jp.empty((0,)) + + +@named_scope +def _advance( + m: Model, + d: Data, + act_dot: jax.Array, + qacc: jax.Array, + qvel: Optional[jax.Array] = None, +) -> Data: + """Advance state and time given activation derivatives and acceleration.""" + act = d.act + if m.na: + act = d.act + act_dot * m.opt.timestep + actrange = jp.where( + m.actuator_actlimited[:, None], + m.actuator_actrange, + jp.array([-jp.inf, jp.inf]), + ) + fn = lambda act, actrange: jp.clip(act, actrange[0], actrange[1]) + act = scan.flat(m, fn, 'au', 'a', act, actrange, group_by='u') + + # advance velocities + d = d.replace(qvel=d.qvel + qacc * m.opt.timestep) + + # advance positions with qvel if given, d.qvel otherwise (semi-implicit) + qvel = d.qvel if qvel is None else qvel + integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep) + qpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d.qpos, qvel) + + # advance time + time = d.time + m.opt.timestep + + return d.replace(act=act, qpos=qpos, time=time) + + +@named_scope +def _euler(m: Model, d: Data) -> Data: + """Euler integrator, semi-implicit in velocity.""" + # integrate damping implicitly + qacc = d.qacc + if not m.opt.disableflags & DisableBit.EULERDAMP: + # TODO(robotics-simulation): can this be done with a smaller perf hit + mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping) + dh = smooth.factor_m(m, d, mh) + qfrc = d.qfrc_smooth + d.qfrc_constraint + qacc = smooth.solve_m(m, dh, qfrc) + return _advance(m, d, d.act_dot, qacc) + + +@named_scope +def _rungekutta4(m: Model, d: Data) -> Data: + """Runge-Kutta explicit order 4 integrator.""" + d_t0 = d + # pylint: disable=invalid-name + A, B = _RK4_A, _RK4_B + C = jp.tril(A).sum(axis=0) # C(i) = sum_j A(i,j) + T = d.time + C * m.opt.timestep + # pylint: enable=invalid-name + + kqvel = d.qvel # intermediate RK solution + # RK solutions sum + qvel, qacc, act_dot = jax.tree_map( + lambda k: B[0] * k, (kqvel, d.qacc, d.act_dot) + ) + integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep) + + def f(carry, x): + qvel, qacc, act_dot, kqvel, d = carry + a, b, t = x # tableau numbers + dqvel, dqacc, dact_dot = jax.tree_map( + lambda k: a * k, (kqvel, d.qacc, d.act_dot) + ) + # get intermediate RK solutions + kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel) + kact = d_t0.act + dact_dot * m.opt.timestep + kqvel = d_t0.qvel + dqacc * m.opt.timestep + d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t) + d = forward(m, d) + + qvel += b * kqvel + qacc += b * d.qacc + act_dot += b * d.act_dot + + return (qvel, qacc, act_dot, kqvel, d), None + + abt = jp.vstack([jp.diag(A), B[1:4], T]).T + out, _ = jax.lax.scan(f, (qvel, qacc, act_dot, kqvel, d), abt, unroll=3) + qvel, qacc, act_dot, *_ = out + + d = _advance(m, d_t0, act_dot, qacc, qvel) + return d + + +@named_scope +def forward(m: Model, d: Data) -> Data: + """Forward dynamics.""" + d = _position(m, d) + d = _velocity(m, d) + d = _actuation(m, d) + d = _acceleration(m, d) + + if d.efc_J.size == 0: + 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.') + + return d + + +@named_scope +def step(m: Model, d: Data) -> Data: + """Advance simulation.""" + d = forward(m, d) + + if m.opt.integrator == IntegratorType.EULER: + d = _euler(m, d) + elif m.opt.integrator == IntegratorType.RK4: + d = _rungekutta4(m, d) + else: + raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') + + return d diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py new file mode 100644 index 00000000..0857db2f --- /dev/null +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -0,0 +1,125 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for forward functions.""" + +import itertools + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import forward +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3): + 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) + + +class ForwardTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(test_util.TEST_FILES)) + def test_forward(self, seed, fname): + """Test mujoco mj forward function matches mujoco_mjx forward function.""" + if fname in ('weld.xml',): + return + + np.random.seed(seed) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) * 0.05 + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + + _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): + """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) + 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): + # 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, '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) + + def test_disable_eulerdamp(self): + m = test_util.load_test_file('ant.xml') + m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP + + d = mujoco.MjData(m) + mx = mjx.device_put(m) + self.assertTrue((mx.dof_damping > 0).any()) + dx = mjx.device_put(d) + dx = jax.jit(forward.forward)(mx, dx) + + dx = dx.replace(qvel=jp.ones_like(dx.qvel), qacc=jp.ones_like(dx.qacc)) + dx = jax.jit(forward._euler)(mx, dx) + np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py new file mode 100644 index 00000000..67e50a7e --- /dev/null +++ b/mjx/mujoco/mjx/_src/io.py @@ -0,0 +1,105 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Functions to initialize, load, or save data.""" + +from jax import numpy as jp +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import constraint +# 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 Model +# pylint: enable=g-importing-member +import numpy as np + + +def make_data(m: Model) -> Data: + """Allocate and initialize Data.""" + + # create first d to get num contacts and nc + d = Data( + solver_niter=jp.array(0, dtype=jp.int32), + ne=0, + nf=0, + nl=0, + nefc=0, + ncon=0, + time=jp.zeros((), dtype=jp.float32), + qpos=m.qpos0, + qvel=jp.zeros(m.nv, dtype=jp.float32), + act=jp.zeros(m.na, dtype=jp.float32), + qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32), + ctrl=jp.zeros(m.nu, dtype=jp.float32), + qfrc_applied=jp.zeros(m.nv, dtype=jp.float32), + xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32), + eq_active=jp.zeros(m.neq, dtype=jp.int32), + qacc=jp.zeros(m.nv, dtype=jp.float32), + act_dot=jp.zeros(m.na, dtype=jp.float32), + xpos=jp.zeros((m.nbody, 3), dtype=jp.float32), + xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), + xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), + xipos=jp.zeros((m.nbody, 3), dtype=jp.float32), + ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), + xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32), + xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32), + geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32), + geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32), + subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32), + cdof=jp.zeros((m.nv, 6), dtype=jp.float32), + cinert=jp.zeros((m.nbody, 10), dtype=jp.float32), + actuator_length=jp.zeros(m.nu, dtype=jp.float32), + actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), + crb=jp.zeros((m.nbody, 10), dtype=jp.float32), + qM=jp.zeros(m.nM, dtype=jp.float32), + qLD=jp.zeros(m.nM, dtype=jp.float32), + qLDiagInv=jp.zeros(m.nv, dtype=jp.float32), + qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32), + contact=Contact.zero(), + efc_J=jp.zeros((), dtype=jp.float32), + efc_frictionloss=jp.zeros((), dtype=jp.float32), + efc_D=jp.zeros((), dtype=jp.float32), + actuator_velocity=jp.zeros(m.nu, dtype=jp.float32), + cvel=jp.zeros((m.nbody, 6), dtype=jp.float32), + cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32), + qfrc_bias=jp.zeros(m.nv, dtype=jp.float32), + qfrc_passive=jp.zeros(m.nv, dtype=jp.float32), + efc_aref=jp.zeros((), dtype=jp.float32), + actuator_force=jp.zeros(m.nu, dtype=jp.float32), + qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32), + qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32), + qacc_smooth=jp.zeros(m.nv, dtype=jp.float32), + qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32), + qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32), + efc_force=jp.zeros((), dtype=jp.float32), + ) + + # get contact data with correct shapes + ncon = collision_driver.ncon(m) + d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon) + d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)}) + + ne, nf, nl, nc = constraint.count_constraints(m, d) + d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc) + ns = ne + nf + nl + d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)}) + d = d.replace( + efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32), + efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32), + efc_D=jp.zeros(d.nefc, dtype=jp.float32), + efc_aref=jp.zeros(d.nefc, dtype=jp.float32), + efc_force=jp.zeros(d.nefc, dtype=jp.float32), + ) + + return d diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py new file mode 100644 index 00000000..269509a8 --- /dev/null +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -0,0 +1,42 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for io functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from mujoco import mjx +from mujoco.mjx._src import test_util + + +class IoTest(parameterized.TestCase): + + @parameterized.parameters(test_util.TEST_FILES) + def test_make_data(self, fname): + """Test that data created by make_data matches data returned by step.""" + + m = test_util.load_test_file(fname) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + dx_step = mjx.step(mx, dx) + + _, dx_treedef = jax.tree_util.tree_flatten(dx) + _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step) + + self.assertEqual(dx_treedef, dx_step_treedef) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/math.py b/mjx/mujoco/mjx/_src/math.py new file mode 100644 index 00000000..e96e6ee0 --- /dev/null +++ b/mjx/mujoco/mjx/_src/math.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Some useful math functions.""" + +from typing import Optional, Tuple, Union + +import jax +from jax import numpy as jp + + +def norm( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> jax.Array: + """Calculates a linalg.norm(x) that's safe for gradients at x=0. + + Avoids a poorly defined gradient for jnp.linal.norm(0) see + https://github.com/google/jax/issues/3058 for details + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + Norm of the array x. + """ + + is_zero = jp.allclose(x, 0.0) + # temporarily swap x with ones if is_zero, then swap back + x = jp.where(is_zero, jp.ones_like(x), x) + n = jp.linalg.norm(x, axis=axis) + n = jp.where(is_zero, 0.0, n) + return n + + +def normalize_with_norm( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> Tuple[jax.Array, jax.Array]: + """Normalizes an array. + + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + A tuple of (normalized array x, the norm). + """ + n = norm(x, axis=axis) + x = x / (n + 1e-6 * (n == 0.0)) + return x, n + + +def normalize( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> jax.Array: + """Normalizes an array. + + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + normalized array x + """ + return normalize_with_norm(x, axis=axis)[0] + + +def rotate(vec: jax.Array, quat: jax.Array) -> jax.Array: + """Rotates a vector vec by a unit quaternion quat. + + Args: + vec: (3,) a vector + quat: (4,) a quaternion + + Returns: + ndarray(3) containing vec rotated by quat. + """ + if len(vec.shape) != 1: + raise ValueError('vec must have no batch dimensions.') + s, u = quat[0], quat[1:] + r = 2 * (jp.dot(u, vec) * u) + (s * s - jp.dot(u, u)) * vec + r = r + 2 * s * jp.cross(u, vec) + return r + + +def quat_inv(q: jp.ndarray) -> jp.ndarray: + """Calculates the inverse of quaternion q. + + Args: + q: (4,) quaternion [w, x, y, z] + + Returns: + The inverse of q, where qmult(q, inv_quat(q)) = [1, 0, 0, 0]. + """ + return q * jp.array([1, -1, -1, -1]) + + +def quat_sub(u: jax.Array, v: jax.Array) -> jax.Array: + """Subtracts two quaternions (u - v) as a 3D velocity.""" + q = quat_mul(quat_inv(v), u) + axis, angle = quat_to_axis_angle(q) + return axis * angle + + +def quat_mul(u: jax.Array, v: jax.Array) -> jax.Array: + """Multiplies two quaternions. + + Args: + u: (4,) quaternion (w,x,y,z) + v: (4,) quaternion (w,x,y,z) + + Returns: + A quaternion u * v. + """ + return jp.array([ + u[0] * v[0] - u[1] * v[1] - u[2] * v[2] - u[3] * v[3], + u[0] * v[1] + u[1] * v[0] + u[2] * v[3] - u[3] * v[2], + u[0] * v[2] - u[1] * v[3] + u[2] * v[0] + u[3] * v[1], + u[0] * v[3] + u[1] * v[2] - u[2] * v[1] + u[3] * v[0], + ]) + + +def quat_mul_axis(q: jax.Array, axis: jax.Array) -> jax.Array: + """Multiplies a quaternion and an axis. + + Args: + q: (4,) quaternion (w,x,y,z) + axis: (3,) axis (x,y,z) + + Returns: + A quaternion q * axis + """ + return jp.array([ + -q[1] * axis[0] - q[2] * axis[1] - q[3] * axis[2], + q[0] * axis[0] + q[2] * axis[2] - q[3] * axis[1], + q[0] * axis[1] + q[3] * axis[0] - q[1] * axis[2], + q[0] * axis[2] + q[1] * axis[1] - q[2] * axis[0], + ]) + + +# TODO(erikfrey): benchmark this against brax's quat_to_3x3 +def quat_to_mat(q: jax.Array) -> jax.Array: + """Converts a quaternion into a 9-dimensional rotation matrix.""" + q = jp.outer(q, q) + + return jp.array([ + [ + q[0, 0] + q[1, 1] - q[2, 2] - q[3, 3], + 2 * (q[1, 2] - q[0, 3]), + 2 * (q[1, 3] + q[0, 2]), + ], + [ + 2 * (q[1, 2] + q[0, 3]), + q[0, 0] - q[1, 1] + q[2, 2] - q[3, 3], + 2 * (q[2, 3] - q[0, 1]), + ], + [ + 2 * (q[1, 3] - q[0, 2]), + 2 * (q[2, 3] + q[0, 1]), + q[0, 0] - q[1, 1] - q[2, 2] + q[3, 3], + ], + ]) + + +def quat_to_axis_angle(q: jax.Array) -> Tuple[jax.Array, jax.Array]: + """Converts a quaternion into axis and angle.""" + axis, sin_a_2 = normalize_with_norm(q[1:]) + angle = 2 * jp.arctan2(sin_a_2, q[0]) + angle = jp.where(angle > jp.pi, angle - 2 * jp.pi, angle) + + return axis, angle + + +def axis_angle_to_quat(axis: jax.Array, angle: jax.Array) -> jax.Array: + """Provides a quaternion that describes rotating around axis by angle. + + Args: + axis: (3,) axis (x,y,z) + angle: () float angle to rotate by + + Returns: + A quaternion that rotates around axis by angle + """ + s, c = jp.sin(angle * 0.5), jp.cos(angle * 0.5) + return jp.insert(axis * s, 0, c) + + +def quat_integrate(q: jax.Array, v: jax.Array, dt: jax.Array) -> jax.Array: + """Integrates a quaternion given angular velocity and dt.""" + v, norm_ = normalize_with_norm(v) + angle = dt * norm_ + q_res = axis_angle_to_quat(v, angle) + q_res = quat_mul(q, q_res) + return normalize(q_res) + + +def inert_mul(i: jax.Array, v: jax.Array) -> jax.Array: + """Multiply inertia by motion, producing force. + + Args: + i: (10,) inertia (inertia matrix, position, mass) + v: (6,) spatial motion + + Returns: + resultant force + """ + tri_id = jp.array([[0, 3, 4], [3, 1, 5], [4, 5, 2]]) # cinert inr order + inr, pos, mass = i[tri_id], i[6:9], i[9] + ang = jp.dot(inr, v[:3]) + jp.cross(pos, v[3:]) + vel = mass * v[3:] - jp.cross(pos, v[:3]) + return jp.concatenate((ang, vel)) + + +def transform_motion(vel: jax.Array, offset: jax.Array, rotmat: jax.Array): + """Transform spatial motion. + + Args: + vel: (6,) spatial motion (3 angular, 3 linear) + offset: (3,) translation + rotmat: (3, 3) rotation + + Returns: + 6d spatial velocity + """ + # TODO(robotics-simulation): are quaternions faster here + ang, vel = vel[:3], vel[3:] + vel = rotmat.T @ (vel - jp.cross(offset, ang)) + ang = rotmat.T @ ang + return jp.concatenate([ang, vel]) + + +def motion_cross(u, v): + """Cross product of two motions. + + Args: + u: (6,) spatial motion + v: (6,) spatial motion + + Returns: + resultant spatial motion + """ + ang = jp.cross(u[:3], v[:3]) + vel = jp.cross(u[3:], v[:3]) + jp.cross(u[:3], v[3:]) + return jp.concatenate((ang, vel)) + + +def motion_cross_force(v, f): + """Cross product of a motion and force. + + Args: + v: (6,) spatial motion + f: (6,) force + + Returns: + resultant force + """ + ang = jp.cross(v[:3], f[:3]) + jp.cross(v[3:], f[3:]) + vel = jp.cross(v[:3], f[3:]) + return jp.concatenate((ang, vel)) + + +def orthogonals(a: jax.Array) -> Tuple[jax.Array, jax.Array]: + """Returns orthogonal vectors `b` and `c`, given a vector `a`.""" + y, z = jp.array([0, 1, 0]), jp.array([0, 0, 1]) + b = jp.where((-0.5 < a[1]) & (a[1] < 0.5), y, z) + b = b - a * a.dot(b) + # normalize b. however if a is a zero vector, zero b as well. + b = normalize(b) * jp.any(a) + return b, jp.cross(a, b) + + +def make_frame(a: jax.Array) -> jax.Array: + """Makes a right-handed 3D frame given a direction.""" + a = normalize(a) + b, c = orthogonals(a) + return jp.array([a, b, c]) + + +# Geometry. + + +def closest_segment_point( + a: jax.Array, b: jax.Array, pt: jax.Array +) -> jax.Array: + """Returns the closest point on the a-b line segment to a point pt.""" + ab = b - a + t = jp.dot(pt - a, ab) / (jp.dot(ab, ab) + 1e-6) + return a + jp.clip(t, 0.0, 1.0) * ab + + +def closest_segment_point_and_dist( + a: jax.Array, b: jax.Array, pt: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Returns closest point on the line segment and the distance squared.""" + closest = closest_segment_point(a, b, pt) + dist = (pt - closest).dot(pt - closest) + return closest, dist + + +def closest_segment_to_segment_points( + a0: jax.Array, a1: jax.Array, b0: jax.Array, b1: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Returns closest points between two line segments.""" + # Gets the closest segment points by first finding the closest points + # between two lines. Points are then clipped to be on the line segments + # and edge cases with clipping are handled. + dir_a, len_a = normalize_with_norm(a1 - a0) + dir_b, len_b = normalize_with_norm(b1 - b0) + + # Segment mid-points. + half_len_a = len_a * 0.5 + half_len_b = len_b * 0.5 + a_mid = a0 + dir_a * half_len_a + b_mid = b0 + dir_b * half_len_b + + # Translation between two segment mid-points. + trans = a_mid - b_mid + + # Parametrize points on each line as follows: + # point_on_a = a_mid + t_a * dir_a + # point_on_b = b_mid + t_b * dir_b + # and analytically minimize the distance between the two points. + dira_dot_dirb = dir_a.dot(dir_b) + dira_dot_trans = dir_a.dot(trans) + dirb_dot_trans = dir_b.dot(trans) + denom = 1 - dira_dot_dirb * dira_dot_dirb + + orig_t_a = (-dira_dot_trans + dira_dot_dirb * dirb_dot_trans) / (denom + 1e-6) + orig_t_b = dirb_dot_trans + orig_t_a * dira_dot_dirb + t_a = jp.clip(orig_t_a, -half_len_a, half_len_a) + t_b = jp.clip(orig_t_b, -half_len_b, half_len_b) + + best_a = a_mid + dir_a * t_a + best_b = b_mid + dir_b * t_b + + # Resolve edge cases where both closest points are clipped to the segment + # endpoints by recalculating the closest segment points for the current + # clipped points, and then picking the pair of points with smallest + # distance. An example of this edge case is when lines intersect but line + # segments don't. + new_a, d1 = closest_segment_point_and_dist(a0, a1, best_b) + new_b, d2 = closest_segment_point_and_dist(b0, b1, best_a) + best_a = jp.where(d1 < d2, new_a, best_a) + best_b = jp.where(d1 < d2, best_b, new_b) + + return best_a, best_b diff --git a/mjx/mujoco/mjx/_src/math_test.py b/mjx/mujoco/mjx/_src/math_test.py new file mode 100644 index 00000000..c2678321 --- /dev/null +++ b/mjx/mujoco/mjx/_src/math_test.py @@ -0,0 +1,202 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for math.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax.numpy as jp +from mujoco.mjx._src import math +import numpy as np + + +def _get_rand_point(seed=None): + if seed is not None: + np.random.seed(seed) + verts = np.random.randn(1, 3) + return verts[0, :] + + +def _get_rand_line_segment(seed=None): + if seed is not None: + np.random.seed(seed) + verts = np.random.randn(2, 3) + return verts[0, :], verts[1, :] + + +def _get_rand_unit(seed: int): + np.random.seed(seed) + theta = np.random.random(1) * 2 * np.pi + a = (np.random.random(1) - 0.5) * 2.0 + phi = np.arccos(a) + x = np.sin(phi) * np.cos(theta) + y = np.sin(phi) * np.sin(theta) + z = np.cos(phi) + return jp.array([x, y, z]).squeeze() + + +class OrthoganalsTest(parameterized.TestCase): + """Tests the orthogonals function.""" + + @parameterized.parameters(range(30)) + def test_orthogonals(self, i): + a = _get_rand_unit(i) + b, c = math.orthogonals(a) + np.testing.assert_almost_equal(jp.linalg.norm(a), 1) + np.testing.assert_almost_equal(jp.linalg.norm(b), 1) + np.testing.assert_almost_equal(jp.linalg.norm(c), 1) + self.assertAlmostEqual(np.abs(a.dot(b)), 0, 6) + self.assertAlmostEqual(np.abs(b.dot(c)), 0, 6) + self.assertAlmostEqual(np.abs(a.dot(c)), 0, 6) + + +def _minimize(fn, sample_fn, lb, ub, tol, max_iter=20, seed=42): + """Minimize a function using the cross-entropy method.""" + assert lb.shape == ub.shape, "bounds need to have the same shape" + np.random.seed(seed) + + i, n = 0, 1_000 + mu = (ub + lb) * 0.5 + sigma = (ub - lb) * 0.5 + size = lb.shape[0] + val, prev_val = fn(mu), None + + while prev_val is None or np.abs(val - prev_val) > tol: + params = sample_fn(mu, sigma, n, size, lb, ub) + vals = np.array([fn(p) for p in params]) + if val < vals.min(): # early exit + return mu + idx = vals.argsort() + best_idx = idx[: int(n * 0.05)] + mu = params[best_idx].mean(axis=0) + sigma = params[best_idx].std(axis=0) + 1e-10 + + prev_val = val + val = fn(mu) + + i += 1 + if i == max_iter: + break + + return mu + + +def _closest_segment_to_segment_points(a0, a1, b0, b1): + dir_a = a1 - a0 + len_a = np.sqrt(dir_a.dot(dir_a)) + half_len_a = len_a / 2 + dir_a = dir_a / len_a + + dir_b = b1 - b0 + len_b = np.sqrt(dir_b.dot(dir_b)) + half_len_b = len_b / 2 + dir_b = dir_b / len_b + + a_mid = a0 + dir_a * half_len_a + b_mid = b0 + dir_b * half_len_b + + # Parametrize both line segments. + def fn(t): + best_a = a_mid + dir_a * t[0] + best_b = b_mid + dir_b * t[1] + return (best_a - best_b).dot(best_a - best_b) + + def sample_fn(mu, sigma, n, size, lb, ub): + params = np.random.normal(mu, sigma, size=(n, size)) + params = np.clip(params, lb, ub) + return params + + lb = np.array([-half_len_a, -half_len_b]) + ub = np.array([half_len_a, half_len_b]) + ta, tb = _minimize(fn, sample_fn, lb, ub, tol=1e-4) + best_a = a_mid + dir_a * ta + best_b = b_mid + dir_b * tb + return best_a, best_b + + +class ClosestSegmentSegmentPointsTest(parameterized.TestCase): + """Tests for closest segment-to-segment points.""" + + def test_closest_segments_points(self): + a0 = jp.array([0.73432405, 0.12372768, 0.20272314]) + a1 = jp.array([1.10600128, 0.88555209, 0.65209485]) + b0 = jp.array([0.85599262, 0.61736299, 0.9843583]) + b1 = jp.array([1.84270939, 0.92891793, 1.36343326]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [1.09063, 0.85404, 0.63351], 5) + self.assertSequenceAlmostEqual(best_b, [0.99596, 0.66156, 1.03813], 5) + + def test_intersecting_segments(self): + """Tests segments that intersect.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([-1.0, 0.0, 0.0]), jp.array([1.0, 0.0, 0.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5) + + def test_intersecting_lines(self): + """Tests that intersecting lines get clipped.""" + a0, a1 = jp.array([0.2, 0.2, 0.0]), jp.array([1.0, 1.0, 0.0]) + b0, b1 = jp.array([0.2, 0.4, 0.0]), jp.array([1.0, 2.0, 0.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.3, 0.3, 0.0], 2) + self.assertSequenceAlmostEqual(best_b, [0.2, 0.4, 0.0], 2) + + def test_parallel_segments(self): + """Tests that parallel segments have closest points at the midpoint.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([1.0, 0.0, -1.0]), jp.array([1.0, 0.0, 1.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.0], 5) + + def test_parallel_offset_segments(self): + """Tests that offset parallel segments are close at segment endpoints.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([1.0, 0.0, 1.0]), jp.array([1.0, 0.0, 3.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 1.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 1.0], 5) + + def test_zero_length_segments(self): + """Test that zero length segments don't return NaNs.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, -1.0]) + b0, b1 = jp.array([1.0, 0.0, 0.1]), jp.array([1.0, 0.0, 0.1]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, -1.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.1], 5) + + def test_overlapping_segments(self): + """Tests that perfectly overlapping segments intersect at the midpoints.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5) + + params = list(zip(np.repeat(np.arange(10), 10), np.tile(np.arange(10), 10))) + + @parameterized.parameters(*params) + def test_closest_segment_to_segment_points(self, i, j): + a0, a1 = _get_rand_line_segment(i) + b0, b1 = _get_rand_line_segment(j) + expected = _closest_segment_to_segment_points(a0, a1, b0, b1) + ans = math.closest_segment_to_segment_points(a0, a1, b0, b1) + expected_dist = (expected[0] - expected[1]).dot(expected[0] - expected[1]) + test_dist = (ans[0] - ans[1]).dot(ans[0] - ans[1]) + self.assertAlmostEqual(expected_dist, test_dist, 4) + + +if __name__ == "__main__": + absltest.main() diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py new file mode 100644 index 00000000..eaa5d3bf --- /dev/null +++ b/mjx/mujoco/mjx/_src/mesh.py @@ -0,0 +1,232 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Mesh processing.""" + +import itertools +from typing import Dict, Optional, Sequence, Tuple + +import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np +from scipy import spatial +import trimesh + + +_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1))) +# pyformat: disable +# Rectangular box faces using a counter-clockwise winding order convention. +_BOX_FACES = [ + 0, 4, 5, 1, # left + 0, 2, 6, 4, # bottom + 6, 7, 5, 4, # front + 2, 3, 7, 6, # right + 1, 5, 7, 3, # top + 0, 1, 3, 2, # back +] +# pyformat: enable +_MAX_HULL_FACE_VERTICES = 20 +_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {} +_DERIVED_ARGS = [ + 'geom_convex_face', + 'geom_convex_vert', + 'geom_convex_edge', + 'geom_convex_facenormal', +] +DERIVED = {(Model, d) for d in _DERIVED_ARGS} + + +def _box(size: np.ndarray): + """Creates a mesh for a box with rectangular faces.""" + box_corners = np.array(_BOX_CORNERS) + vert = box_corners * size.reshape(-1, 3) + face = np.array([_BOX_FACES]).reshape(-1, 4) + return vert, face + + +def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray: + """Calculates face normals given vertices and face indexes.""" + assert len(vert.shape) == 2 and len(face.shape) == 2, ( + f'vert and face should have dim of 2, got {len(vert.shape)} and ' + f'{len(face.shape)}' + ) + face_vert = vert[face, :] + # use CCW winding order convention + edge0 = face_vert[:, 1, :] - face_vert[:, 0, :] + edge1 = face_vert[:, -1, :] - face_vert[:, 0, :] + face_norm = np.cross(edge0, edge1) + face_norm = face_norm / np.linalg.norm(face_norm, axis=1).reshape((-1, 1)) + return face_norm + + +def _get_unique_edges(vert: np.ndarray, face: np.ndarray) -> np.ndarray: + """Returns unique edges. + + Args: + vert: (n_vert, 3) vertices + face: (n_face, n_vert) face index array + + Returns: + edges: tuples of vertex indexes for each edge + """ + r_face = np.roll(face, 1, axis=1) + edges = np.concatenate(np.array([face, r_face]).T) + + # do a first pass to remove duplicates + edges.sort(axis=1) + edges = np.unique(edges, axis=0) + edges = edges[edges[:, 0] != edges[:, 1]] # get rid of edges from padded face + + # get normalized edge directions + edge_vert = vert.take(edges, axis=0) + edge_dir = edge_vert[:, 0] - edge_vert[:, 1] + norms = np.sqrt(np.sum(edge_dir**2, axis=1)) + edge_dir = edge_dir / norms.reshape((-1, 1)) + + # get the first unique edge for all pairwise comparisons + diff1 = edge_dir[:, None, :] - edge_dir[None, :, :] + diff2 = edge_dir[:, None, :] + edge_dir[None, :, :] + matches = (np.linalg.norm(diff1, axis=-1) < 1e-6) | ( + np.linalg.norm(diff2, axis=-1) < 1e-6 + ) + matches = np.tril(matches).sum(axis=-1) + unique_edge_idx = np.where(matches == 1)[0] + + return edges[unique_edge_idx] + + +def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray: + """Calculates the convex hull for a set of points on a plane.""" + # project points onto the closest axis plane + best_axis = np.abs(np.eye(3).dot(normal)).argmax() + axis = np.eye(3)[best_axis] + d = points.dot(axis).reshape((-1, 1)) + axis_points = points - d * axis + axis_points = axis_points[:, list({0, 1, 2} - {best_axis})] + + # get the polygon face, and make the points ccw wrt the face normal + c = spatial.ConvexHull(axis_points) + order_ = np.where(axis.dot(normal) > 0, 1, -1) + order_ *= np.where(best_axis == 1, -1, 1) + hull_point_idx = c.vertices[::order_] + assert (axis_points - c.points).sum() == 0 + + return hull_point_idx + + +def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray: + """Merges coplanar facets.""" + if not tm.facets: + return tm.faces.copy() # no facets + if not tm.faces.shape[0]: + raise ValueError('Mesh has no faces.') + + # Get faces. + face_idx = set(range(tm.faces.shape[0])) - set(np.concatenate(tm.facets)) + face_idx = np.array(list(face_idx)) + faces = tm.faces[face_idx] if face_idx.shape[0] > 0 else np.array([]) + + # Get facets. + facets = [] + for i, facet in enumerate(tm.facets): + point_idx = np.unique(tm.faces[facet]) + points = tm.vertices[point_idx] + normal = tm.facets_normal[i] + + # convert triangulated facet to a polygon + hull_point_idx = _convex_hull_2d(points, normal) + face = point_idx[hull_point_idx] + + # resize faces that exceed max polygon vertices + every = face.shape[0] // _MAX_HULL_FACE_VERTICES + 1 + face = face[::every] + facets.append(face) + + # Pad facets so that they can be stacked. + max_len = max(f.shape[0] for f in facets) if facets else faces.shape[1] + assert max_len <= _MAX_HULL_FACE_VERTICES + for i, f in enumerate(facets): + if f.shape[0] < max_len: + f = np.pad(f, (0, max_len - f.shape[0]), 'edge') + facets[i] = f + + if not faces.shape[0]: + assert facets + return np.array(facets) # no faces, return facets + + # Merge faces and facets. + faces = np.pad(faces, ((0, 0), (0, max_len - faces.shape[1])), 'edge') + return np.concatenate([faces, facets]) + + +def _get_faces_verts( + m: mujoco.MjModel, +) -> Tuple[Sequence[np.ndarray], Sequence[np.ndarray]]: + """Extracts mesh faces and vertices from MjModel.""" + verts, faces = [], [] + for i in range(m.nmesh): + last = (i + 1) >= m.nmesh + face_start = m.mesh_faceadr[i] + face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0] + face = m.mesh_face[face_start:face_end] + faces.append(face) + + vert_start = m.mesh_vertadr[i] + vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0] + vert = m.mesh_vert[vert_start:vert_end] + verts.append(vert) + return verts, faces + + +def _geom_mesh_kwargs( + vert: np.ndarray, face: np.ndarray +) -> Dict[str, np.ndarray]: + """Generates convex mesh attributes for mjx.Model.""" + tm = trimesh.Trimesh(vertices=vert, faces=face) + tm_convex = trimesh.convex.convex_hull(tm) + vert = np.array(tm_convex.vertices) + face = _merge_coplanar(tm_convex) + return { + 'geom_convex_face': face, + 'geom_convex_vert': vert, + 'geom_convex_edge': _get_unique_edges(vert, face), + 'geom_convex_facenormal': _get_face_norm(vert, face), + } + + +def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]: + """Derives geom mesh attributes for mjx.Model from MjModel.""" + kwargs = {k: [] for k in _DERIVED_ARGS} + verts, faces = _get_faces_verts(m) + for geomid in range(m.ngeom): + dataid = m.geom_dataid[geomid] + typ = m.geom_type[geomid] + if typ == GeomType.BOX: + vert, face = _box(m.geom_size[geomid]) + elif dataid >= 0: + vert, face = verts[dataid], faces[dataid] + else: + kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS} + continue + + key = (hash(vert.data.tobytes()), hash(face.data.tobytes())) + if key not in _CONVEX_CACHE: + _CONVEX_CACHE[key] = _geom_mesh_kwargs(vert, face) + + kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS} + + return kwargs diff --git a/mjx/mujoco/mjx/_src/mesh_test.py b/mjx/mujoco/mjx/_src/mesh_test.py new file mode 100644 index 00000000..7be8197a --- /dev/null +++ b/mjx/mujoco/mjx/_src/mesh_test.py @@ -0,0 +1,119 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for mesh.py.""" + +from absl.testing import absltest +from mujoco.mjx._src import mesh +import numpy as np + + +class GeomMeshKwargsTest(absltest.TestCase): + + def test_pyramid(self): + """Tests that a triangulated pyramid converts to merged coplanar faces.""" + vert = np.array([ + [-0.025, 0.05, 0.05], + [-0.025, -0.05, -0.05], + [-0.025, -0.05, 0.05], + [-0.025, 0.05, -0.05], + [0.075, 0.0, 0.0], + ]) + face = np.array( + [[0, 1, 2], [0, 3, 1], [0, 4, 3], [0, 2, 4], [2, 1, 4], [1, 3, 4]] + ) + h = mesh._geom_mesh_kwargs(vert, face) + + # get index of vertices in h['geom_convex_vert'] for vertices in vert + dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile( + h['geom_convex_vert'], (vert.shape[0], 1) + ) + dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1)) + vidx = np.argmin(dist, axis=0) + + # check verts + np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx]) + + # check face vertices + map_ = {v: k for k, v in enumerate(vidx)} + h_face = np.vectorize(map_.get)(h['geom_convex_face']) + face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()]) + expected_face_verts = sorted([ + (0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)]) + self.assertSequenceEqual( + face_verts, + expected_face_verts, + ) + + # check edges + unique_edge = np.vectorize(map_.get)(h['geom_convex_edge']) + unique_edge = np.array(sorted(unique_edge.tolist())) + np.testing.assert_array_equal( + unique_edge, + np.array([[0, 2], [0, 3], [0, 4], [1, 4], [2, 4], [3, 4]]), + ) + + # face normals + self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3)) + + +class ConvexHull2DTest(absltest.TestCase): + + def test_convex_hull_2d_axis1(self): + """Tests for the correct winding order of a polgyon with +y normal.""" + pts = np.array([ + [-0.04634297, -0.06652775, 0.05853534], + [-0.01877651, -0.08309858, -0.05236476], + [0.02362804, -0.08010745, 0.05499557], + [0.04066505, -0.09034877, -0.01354446], + [-0.07255043, -0.06837638, -0.00781699], + ]) + normal = np.array([-0.18467607, -0.97768016, 0.10018111]) + idx = mesh._convex_hull_2d(pts, normal) + expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0]) + expected /= np.linalg.norm(expected) + np.testing.assert_array_almost_equal(normal, expected) + + def test_convex_hull_2d_axis2(self): + """Tests for the correct winding order for a polgyon with +z normal.""" + pts = np.array([ + [0.08607829, -0.03881998, -0.03291714], + [-0.01877651, -0.08309858, -0.05236476], + [0.05470364, 0.00027677, -0.08371042], + [-0.01010019, -0.02708892, -0.0957297], + [0.04066505, -0.09034877, -0.01354446], + ]) + normal = np.array([0.3839915, -0.60171936, -0.70034587]) + idx = mesh._convex_hull_2d(pts, normal) + expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0]) + expected /= np.linalg.norm(expected) + np.testing.assert_array_almost_equal(normal, expected) + + +class UniqueEdgesTest(absltest.TestCase): + + def test_tetrahedron_edges(self): + """Tests unique edges for a tetrahedron.""" + vert = np.array( + [[-0.1, 0.0, -0.1], [0.0, 0.1, 0.1], [0.1, 0.0, -0.1], [0.0, -0.1, 0.1]] + ) + face = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 1], [2, 1, 3]]) + idx = mesh._get_unique_edges(vert, face) + np.testing.assert_array_equal( + idx, np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]]) + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py new file mode 100644 index 00000000..9b0b41f0 --- /dev/null +++ b/mjx/mujoco/mjx/_src/passive.py @@ -0,0 +1,141 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Passive forces.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +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.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def _inertia_box_fluid_model( + m: Model, + inertia: jax.Array, + mass: jax.Array, + root_com: jax.Array, + xipos: jax.Array, + ximat: jax.Array, + cvel: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Fluid forces based on inertia-box approximation.""" + box = jp.repeat(inertia[None, :], 3, axis=0) + box *= jp.ones((3, 3)) - 2 * jp.eye(3) + box = 6.0 * jp.clip(jp.sum(box, axis=-1), a_min=1e-12) + box = jp.sqrt(box / jp.maximum(mass, 1e-12)) * (mass > 0.0) + + # transform to local coordinate frame + offset = xipos - root_com + lvel = math.transform_motion(cvel, offset, ximat) + lwind = ximat.T @ m.opt.wind + lvel = lvel.at[3:].add(-lwind) + + # set viscous force and torque + diam = jp.mean(box, axis=-1) + lfrc_ang = lvel[:3] * -jp.pi * diam**3 * m.opt.viscosity + lfrc_vel = lvel[3:] * -3.0 * jp.pi * diam * m.opt.viscosity + + # add lift and drag force and torque + scale_vel = jp.array([box[1] * box[2], box[0] * box[2], box[0] * box[1]]) + scale_ang = jp.array([ + box[0] * (box[1] ** 4 + box[2] ** 4), + box[1] * (box[0] ** 4 + box[2] ** 4), + box[2] * (box[0] ** 4 + box[1] ** 4), + ]) + lfrc_vel -= 0.5 * m.opt.density * scale_vel * jp.abs(lvel[3:]) * lvel[3:] + lfrc_ang -= ( + 1.0 * m.opt.density * scale_ang * jp.abs(lvel[:3]) * lvel[:3] / 64.0 + ) + + # rotate to global orientation: lfrc -> bfrc + force, torque = ximat @ lfrc_vel, ximat @ lfrc_ang + + return force, torque + + +def passive(m: Model, d: Data) -> Data: + """Adds all passive forces.""" + if m.opt.disableflags & DisableBit.PASSIVE: + return d + + # joint-level springs + def fn(jnt_typs, stiffness, qpos_spring, qpos): + qpos_i = 0 + qfrcs = [] + for i in range(len(jnt_typs)): + jnt_typ = JointType(jnt_typs[i]) + q = qpos[qpos_i : qpos_i + jnt_typ.qpos_width()] + qs = qpos_spring[qpos_i : qpos_i + jnt_typ.qpos_width()] + qfrc = jp.zeros(jnt_typ.dof_width()) + if jnt_typ == JointType.FREE: + qfrc = qfrc.at[:3].set(-stiffness[i] * (q[:3] - qs[:3])) + qfrc = qfrc.at[3:6].set(-stiffness[i] * math.quat_sub(q[3:7], qs[3:7])) + elif jnt_typ == JointType.BALL: + qfrc = -stiffness[i] * math.quat_sub(q, qs) + elif jnt_typ in ( + JointType.SLIDE, + JointType.HINGE, + ): + qfrc = -stiffness[i] * (q - qs) + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + qfrcs.append(qfrc) + qpos_i += jnt_typ.qpos_width() + return jp.concatenate(qfrcs) + + qfrc_passive = scan.flat( + m, + fn, + 'jjqq', + 'v', + m.jnt_type, + m.jnt_stiffness, + m.qpos_spring, + d.qpos, + ) + + # dof-level dampers + qfrc_passive -= m.dof_damping * d.qvel + + # TODO(robotics-simulation): body-level gravity compensation + + # body-level viscosity, lift and drag + if m.opt.has_fluid_params: + force, torque = jax.vmap( + _inertia_box_fluid_model, in_axes=(None, 0, 0, 0, 0, 0, 0) + )( + m, + m.body_inertia, + m.body_mass, + d.subtree_com[jp.array(m.body_rootid)], + d.xipos, + d.ximat, + d.cvel, + ) + qfrc_target = jax.vmap(support.apply_ft, in_axes=(None, None, 0, 0, 0, 0))( + m, d, force, torque, d.xipos, jp.arange(m.nbody) + ) + qfrc_passive += jp.sum(qfrc_target, axis=0) + + d = d.replace(qfrc_passive=qfrc_passive) + return d diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py new file mode 100644 index 00000000..49264014 --- /dev/null +++ b/mjx/mujoco/mjx/_src/passive_test.py @@ -0,0 +1,114 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests passive forces.""" + +import itertools + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import jax.numpy as jp +import mujoco +from mujoco import mjx +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5): + 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) + + +class PassiveTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(('ant.xml', 'mixed_joint_pendulum.xml'))) + def test_stiffness_damping(self, seed, fname): + """Tests stiffness and damping on Ant.""" + np.random.seed(seed) + path = epath.resource_path('mujoco.mjx') / 'test_data' + path /= fname + m = mujoco.MjModel.from_xml_string(path.read_text()) + + # set stiffness/damping + m.jnt_stiffness = np.random.uniform(size=m.njnt) + m.dof_damping = np.random.uniform(size=m.nv) + d = mujoco.MjData(m) + d.qvel = np.random.random(m.nv) # random kick + + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + passive_jit_fn = jax.jit(mjx.passive) + + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) + + @parameterized.parameters( + itertools.product(range(3), ('triple_pendulum.xml',)) + ) + def test_fluid(self, seed, fname): + np.random.seed(seed) + path = epath.resource_path('mujoco.mjx') / 'test_data' + path /= fname + m = mujoco.MjModel.from_xml_string(path.read_text()) + + # set density/viscosity/wind + m.opt.density = np.random.uniform() + m.opt.viscosity = np.random.uniform() + m.opt.wind = np.random.uniform() + + passive_jit_fn = jax.jit(mjx.passive) + + mx = mjx.device_put(m) + d = mujoco.MjData(m) + d.qvel = np.random.random(m.nv) # random kick + + for i in range(100): + mujoco.mj_step(m, d) + dx = mjx.device_put(d) + mujoco.mj_passive(m, d) + dx = passive_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) + + def test_disable_passive(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + """) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + dx = dx.replace(qvel=jp.ones(mx.nv)) + + passive_jit_fn = jax.jit(mjx.passive) + dx = passive_jit_fn(mx, dx) + np.testing.assert_equal(dx.qfrc_passive, np.zeros(mx.nv)) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py new file mode 100644 index 00000000..b54655c0 --- /dev/null +++ b/mjx/mujoco/mjx/_src/scan.py @@ -0,0 +1,432 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Scan across data ordered by body joint types and kinematic tree order.""" + +from typing import Any, Callable, TypeVar + +import jax +from jax import numpy as jp +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import TrnType +# pylint: enable=g-importing-member +import numpy as np + + +Y = TypeVar('Y') + + +# TODO(erikfrey): re-check if this really helps perf +def _take(obj: Y, idx: np.ndarray) -> Y: + """Takes idxs on any pytree given to it. + + XLA executes x[jp.array([1, 2, 3])] slower than x[1:4], so we detect when + take indices are contiguous, and convert them to slices. + + Args: + obj: an input pytree + idx: indices to take + + Returns: + obj pytree with leaves taken by idxs + """ + + if isinstance(obj, np.ndarray): + return obj[idx] + + def take(x): + # TODO(erikfrey): if this helps perf, add support for striding too + if ( + len(idx.shape) == 1 + and idx.size > 0 + and (idx == np.arange(idx[0], idx[0] + idx.size)).all() + and (idx > 0).all() + ): + x = x[idx[0] : idx[-1] + 1] + else: + x = x.take(jp.array(idx), axis=0, mode='wrap') + return x + + return jax.tree_map(take, obj) + + +def _q_bodyid(m: Model) -> np.ndarray: + """Returns the bodyid for each qpos adress.""" + q_bodyids = [np.array([], dtype=np.int32)] + for jnt_type, jnt_bodyid in zip(m.jnt_type, m.jnt_bodyid): + width = {JointType.FREE: 7, JointType.BALL: 4}.get(jnt_type, 1) + q_bodyids.append(np.repeat(jnt_bodyid, width)) + return np.concatenate(q_bodyids) + + +def _q_jointid(m: Model) -> np.ndarray: + """Returns the jointid for each qpos adress.""" + q_jointid = [np.array([], dtype=np.int32)] + for i, jnt_type in enumerate(m.jnt_type): + width = {JointType.FREE: 7, JointType.BALL: 4}.get(jnt_type, 1) + q_jointid.append(np.repeat(i, width)) + return np.concatenate(q_jointid) + + +def _index(haystack: np.ndarray, needle: np.ndarray) -> np.ndarray: + """Returns indexes in haystack for elements in needle.""" + idx = np.argsort(haystack) + sorted_haystack = haystack[idx] + sorted_idx = np.searchsorted(sorted_haystack, needle) + idx = np.take(idx, sorted_idx, mode='clip') + idx[haystack[idx] != needle] = -1 + + return idx + + +def _nvmap(f: Callable[..., Y], *args) -> Y: + """A vmap that accepts numpy arrays. + + Numpy arrays are statically vmapped, and the elements are passed to f as + static arguments. The implication is that all the elements of numpy array + arguments must be the same. + + Args: + f: function to be mapped over + *args: args to be mapped along, passed to f + + Returns: + the result of vmapping f over args + + Raises: + RuntimeError: if numpy arg elements do not match + """ + for arg in args: + if isinstance(arg, np.ndarray) and not np.all(arg == arg[0]): + raise RuntimeError(f'numpy arg elements do not match: {arg}') + + np_args = [a[0] if isinstance(a, np.ndarray) else None for a in args] + args = [a if n is None else None for n, a in zip(np_args, args)] + in_axes = [None if a is None else 0 for a in args] + + def outer_f(*args, np_args=np_args): + args = [a if n is None else n for n, a in zip(args, np_args)] + return f(*args) + + return jax.vmap(outer_f, in_axes=in_axes)(*args) + + +def _check_input(m: Model, args: Any, in_types: str) -> None: + """Checks that scan input has the right shape.""" + size = {'b': m.nbody, 'j': m.njnt, 'q': m.nq, 'v': m.nv, 'u': m.nu, 'a': m.na} + for idx, (a, typ) in enumerate(zip(args, in_types)): + if len(a) != size[typ]: + raise IndexError( + ( + f'f argument "{idx}" with type "{typ}" has length "{len(a)}"' + f' which does not match the in_types[{idx}] expected length of ' + f'"{size[typ]}".' + ) + ) + + +def _check_output( + y: jax.Array, take_ids: np.ndarray, typ: str, idx: int +) -> None: + """Checks that scan output has the right shape.""" + if y.shape[0] != take_ids.shape[0]: + raise IndexError( + ( + f'f output "{idx}" with type "{typ}" has shape "{y.shape[0]}" ' + f'which does not match the out_types[{idx}] expected size of' + f' "{take_ids.shape[0]}".' + ) + ) + + +def flat( + m: Model, + f: Callable[..., Y], + in_types: str, + out_types: str, + *args, + group_by: str = 'j', +) -> Y: + r"""Scan a function across bodies or actuators. + + Scan group data according to type and batch shape then calls vmap(f) on it.\ + + Args: + m: an mjx model + f: a function to be scanned with the following type signature: + def f(key, *args) -> y + where + ``key`` gives grouping key for this function instance + ``*args`` are input arguments with types matching ``in_types`` + ``y`` is an output arguments with types matching ``out_type`` + in_types: string specifying the type of each input arg: + 'b': split according to bodies + 'j': split according to joint types + 'q': split according to generalized coordinates (len(qpos)) + 'v': split according to degrees of freedom (len(qvel)) + 'u': split according to actuators + 'a': split according to actuator activations + out_types: string specifying the types the output dimension matches + *args: the input arguments corresponding to ``in_types`` + group_by: the type to group by, either joints or actuators + + Returns: + The stacked outputs of ``f`` matching the model's order. + + Raises: + IndexError: if function output shape does not match out_types shape + """ + _check_input(m, args, in_types) + + if group_by not in {'j', 'u'}: + raise NotImplementedError(f'group by type "{group_by}" not implemented.') + + def key_j(ids): + if any(t in 'jqv' for t in in_types + out_types): + return tuple(m.jnt_type[ids]) + return () + + def key_u(ids_u, ids_j): + return ( + m.actuator_biastype[ids_u], + m.actuator_gaintype[ids_u], + m.actuator_dyntype[ids_u], + m.actuator_trntype[ids_u], + m.jnt_type[ids_j], + ) + + def type_ids_j(m, i): + return { + 'b': i, + 'j': np.nonzero(m.jnt_bodyid == i)[0], + 'v': np.nonzero(m.dof_bodyid == i)[0], + 'q': np.nonzero(_q_bodyid(m) == i)[0], + } + + def type_ids_u(m, i): + typ_ids = { + 'u': i, + 'a': m.actuator_actadr[i], + 'j': ( + m.actuator_trnid[i] + if m.actuator_trntype[i] == TrnType.JOINT + else np.array(-1) + ), + } + # v/q associated with joint transmissions + typ_ids.update({ + 'v': np.nonzero(m.dof_jntid == typ_ids['j'])[0], + 'q': np.nonzero(_q_jointid(m) == typ_ids['j'])[0], + }) + return typ_ids + + # build up a grouping of type take-ids in body/actuator order + key_typ_ids, order = {}, [] + all_types = set(in_types + out_types) + n_items = {'j': m.nbody, 'u': m.nu}[group_by] + for i in np.arange(n_items, dtype=np.int32): + typ_ids = type_ids_j(m, i) if group_by == 'j' else type_ids_u(m, i) + + # create grouping key + key = ( + key_j(typ_ids['j']) + if group_by == 'j' + else key_u(typ_ids['u'], typ_ids['j']) + ) + order.append((key, typ_ids)) + + # add ids per type to the corresponding group + for t in all_types: + out = key_typ_ids.setdefault(key, {}) + val = np.expand_dims(typ_ids[t], axis=0) + out[t] = np.concatenate((out[t], val)) if t in out else val + + key_typ_ids = list(sorted(key_typ_ids.items())) + + # use this grouping to take the right data subsets and call vmap(f) + ys = [] + for _, typ_ids in key_typ_ids: + # only execute f if we would actually take something from the result + if any(typ_ids[v].size > 0 for v in out_types): + f_args = [_take(arg, typ_ids[typ]) for arg, typ in zip(args, in_types)] + y = _nvmap(f, *f_args) + ys.append(y) + else: + ys.append(None) + + # remove None results from the final output + key_typ_ids = [v for y, v in zip(ys, key_typ_ids) if y is not None] + ys = [y for y in ys if y is not None] + ys_keys = set([k for k, *_ in key_typ_ids]) + order = [o for k, o in order if k in ys_keys] + + # get the original input order + order = [[o[t] for o in order] for t in all_types] + order = [ + np.concatenate(o) if isinstance(o[0], np.ndarray) else np.array(o) + for o in order + ] + order = dict(zip(all_types, order)) + + # concatenate back to a single tree and drop the grouping dimension + f_ret_is_seq = isinstance(ys[0], (list, tuple)) + ys = ys if f_ret_is_seq else [[y] for y in ys] + flat_ = {'j': 'b', 'u': 'uaj'}[group_by] + ys = [ + [v if typ in flat_ 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) + + # put concatenated results back in order + reordered_ys = [] + for i, (y, typ) in enumerate(zip(ys, out_types)): + _check_output(y, order[typ], typ, i) + ids = np.concatenate([np.hstack(v[typ]) for _, v in key_typ_ids]) + input_order = order[typ][np.where(order[typ] != -1)] + reordered_ys.append(_take(y, _index(ids, input_order))) + y = reordered_ys if f_ret_is_seq else reordered_ys[0] + + return y + + +def body_tree( + m: Model, + f: Callable[..., Y], + in_types: str, + out_types: str, + *args, + reverse: bool = False, +) -> Y: + r"""Scan ``f`` across bodies in tree order, carrying results up/down the tree. + + This function groups bodies according to level and attached joints, then calls + vmap(f) on them.\ + + Args: + m: an mjx mjmodel + f: a function to be scanned with the following type signature: + def f(y, *args) -> y + where + ``y`` is the carry value and return value + ``*args`` are input arguments with types matching ``in_types`` + in_types: string specifying the type of each input arg: + 'b': split according to bodies + 'j': split according to joint types + 'q': split according to generalized coordinates (len(qpos)) + 'v': split according to degrees of freedom (len(qvel)) + out_types: string specifying the types the output dimension matches + *args: the input arguments corresponding to ``in_types`` + reverse: if True, scans up the body tree from leaves to root, otherwise + root to leaves + + Returns: + The stacked outputs of ``f`` matching the model's body order. + + Raises: + 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) + + # 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, ()) + + 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_ + + key_typ_ids = list(sorted(key_typ_ids.items(), reverse=reverse)) + + # use this grouping to take the right data subsets and call vmap(f) + key_y = {} + for key, typ_ids in key_typ_ids: + carry = None + + if reverse: + child_keys = [k for k, v in key_parents.items() if v == key] + + for child_key in child_keys: + y = key_y[child_key] + body_ids = key_body_ids[key] + parent_ids = m.body_parentid[key_body_ids[child_key]] + id_map = _index(body_ids, parent_ids) + + def index_sum(x, i=id_map, s=body_ids.size): + return jax.ops.segment_sum(x, i, s) + + 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) + + 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)] + 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] + + # 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) + + # 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] + + return y diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py new file mode 100644 index 00000000..9b845de1 --- /dev/null +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -0,0 +1,227 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for scan functions.""" + +from absl.testing import absltest +from jax import numpy as jp +import mujoco +from mujoco import mjx +# pylint: disable=g-importing-member +from mujoco.mjx._src import scan +from mujoco.mjx._src.types import JointType +# pylint: enable=g-importing-member +import numpy as np + + +class ScanTest(absltest.TestCase): + + _MULTI_DOF_XML = """ + + + + + + + + + + + + + + + + + + + """ + + def test_flat_empty(self): + """Test scanning over just world body.""" + m = mujoco.MjModel.from_xml_string(""" + + + + """) + m = mjx.device_put(m) + + def fn(body_id): + return body_id + 1 + + b_in = jp.array([1]) + b_expect = jp.array([2]) + b_out = scan.flat(m, fn, 'b', 'b', b_in) + + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + def test_flat_joints(self): + """Tests scanning over bodies with joints of different types.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) + m = mjx.device_put(m) + + # we will test two functions: + # 1) j_fn receives jnt_types as a jp array + # 2) s_fn receives jnt_types as a static np array and can switch on it + j_fn = lambda jnt_pos, val: val + jp.sum(jnt_pos) + s_fn = lambda jnt_types, val: val + sum(jnt_types) + + b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]]) + b_expect = jp.array([[0, 0], [1, 1], [3, 3], [8, 8]]) + b_out = scan.flat(m, j_fn, 'jb', 'b', m.jnt_pos, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.flat(m, s_fn, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # None should be omitted from the results + def no_free(jnt_types, val): + if tuple(jnt_types) == (JointType.FREE,): + return None + return val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) + b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # we should not call functions for which we know we will discard the results + def no_world(jnt_types, val): + if jnt_types.size == 0: + self.fail('world has no dofs, should not be called') + return val + sum(jnt_types) + v_in = jp.ones((m.nv, 1)) + scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in) + + def test_body_tree(self): + """Tests tree scanning over bodies with different joint counts.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) + m = mjx.device_put(m) + + # we will test two functions: + # 1) j_fn receives jnt_pos which is a jp array + # 2) s_fn receives jnt_types which is a static np array + def j_fn(carry, jnt_pos, val): + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + jp.sum(jnt_pos) + + def s_fn(carry, jnt_types, val): + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + sum(jnt_types) + + b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]]) + b_expect = jp.array([[0, 0], [1, 1], [4, 4], [9, 9]]) + + b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # and reverse too: + b_expect = jp.array([[12, 12], [12, 12], [3, 3], [8, 8]]) + b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in, reverse=True) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in, reverse=True) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # None should be omitted from the results + def no_free(carry, jnt_types, val): + if tuple(jnt_types) == (JointType.FREE,): + return None + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) + b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + _MULTI_ACT_XML = """ + + + """ + + def testscan_actuators(self): + """Tests scanning over actuators.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML) + m = mjx.device_put(m) + + fn = lambda *args: args + args = ( + m.actuator_gear, + m.jnt_type, + jp.arange(m.nq), + jp.arange(m.nv), + jp.array([1.4, 1.1]), + ) + gear, jnt_typ, qadr, vadr, act = scan.flat( + m, fn, 'ujqva', 'ujqva', *args, group_by='u' + ) + + np.testing.assert_array_equal(gear, m.actuator_gear) + np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid]) + np.testing.assert_array_equal(act, jp.array([1.4, 1.1])) + expected_vadr = np.concatenate( + [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid] + ) + np.testing.assert_array_equal(vadr, expected_vadr) + expected_qadr = np.concatenate( + [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid] + ) + np.testing.assert_array_equal(qadr, expected_qadr) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py new file mode 100644 index 00000000..1193a1f8 --- /dev/null +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -0,0 +1,466 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Core smooth dynamics functions.""" + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math +from mujoco.mjx._src import scan +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def kinematics(m: Model, d: Data) -> Data: + """Converts position/velocity from generalized coordinates to maximal.""" + + def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat): + # calculate joint anchors, axes, body pos and quat in global frame + # also normalize qpos while we're at it + + if carry is not None: + _, _, _, parent_pos, parent_quat, _ = carry + pos = parent_pos + math.rotate(pos, parent_quat) + quat = math.quat_mul(parent_quat, quat) + + anchors, axes = [], [] + + qpos_i = 0 + for i, jnt_typ in enumerate(jnt_typs): + if jnt_typ == JointType.FREE: + anchor, axis = qpos[qpos_i : qpos_i + 3], jp.array([0.0, 0.0, 1.0]) + else: + anchor = math.rotate(jnt_pos[i], quat) + pos + axis = math.rotate(jnt_axis[i], quat) + anchors, axes = anchors + [anchor], axes + [axis] + + if jnt_typ == JointType.FREE: + pos = qpos[qpos_i : qpos_i + 3] + quat = math.normalize(qpos[qpos_i + 3 : qpos_i + 7]) + qpos = qpos.at[qpos_i + 3 : qpos_i + 7].set(quat) + qpos_i += 7 + elif jnt_typ == JointType.BALL: + qloc = math.normalize(qpos[qpos_i : qpos_i + 4]) + qpos = qpos.at[qpos_i : qpos_i + 4].set(qloc) + quat = math.quat_mul(quat, qloc) + pos = anchor - math.rotate(jnt_pos[i], quat) # off-center rotation + qpos_i += 4 + elif jnt_typ == JointType.HINGE: + angle = qpos[qpos_i] - qpos0[qpos_i] + qloc = math.axis_angle_to_quat(jnt_axis[i], angle) + quat = math.quat_mul(quat, qloc) + pos = anchor - math.rotate(jnt_pos[i], quat) # off-center rotation + qpos_i += 1 + elif jnt_typ == JointType.SLIDE: + pos += axis * (qpos[qpos_i] - qpos0[qpos_i]) + qpos_i += 1 + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + anchor = jp.stack(anchors) if anchors else jp.empty((0, 3)) + axis = jp.stack(axes) if axes else jp.empty((0, 3)) + mat = math.quat_to_mat(quat) + + return qpos, anchor, axis, pos, quat, mat + + qpos, xanchor, xaxis, xpos, xquat, xmat = scan.body_tree( + m, + fn, + 'jjjqqbb', + 'qjjbbb', + m.jnt_type, + m.jnt_pos, + m.jnt_axis, + d.qpos, + m.qpos0, + m.body_pos, + m.body_quat, + ) + + @jax.vmap + def local_to_global(pos1, quat1, pos2, quat2): + pos = pos1 + math.rotate(pos2, quat1) + mat = math.quat_to_mat(math.quat_mul(quat1, quat2)) + return pos, mat + + # TODO(erikfrey): confirm that quats are more performant for mjx than mats + xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat) + geom_xpos, geom_xmat = local_to_global( + xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat + ) + + d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos) + d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat) + d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat) + + return d + + +def com_pos(m: Model, d: Data) -> Data: + """Maps inertias and motion dofs to global frame centered at subtree-CoM.""" + + # calculate center of mass of each subtree + def subtree_sum(carry, xipos, body_mass): + pos, mass = xipos * body_mass, body_mass + if carry is not None: + subtree_pos, subtree_mass = carry + pos, mass = pos + subtree_pos, mass + subtree_mass + return pos, mass + + pos, mass = scan.body_tree( + m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True + ) + cond = jp.tile(mass < jp.array(mujoco.mjMINVAL), (3, 1)).T + subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass)) + d = d.replace(subtree_com=subtree_com) + + # map inertias to frame centered at subtree_com + @jax.vmap + def inert_com(inert, ximat, off, mass): + h = jp.cross(off, -jp.eye(3)) + inert = ximat @ jp.diag(inert) @ ximat.T + h @ h.T * mass + # cinert is triu(inert), mass * off, mass + inert = inert[(jp.array([0, 1, 2, 0, 0, 1]), jp.array([0, 1, 2, 1, 2, 2]))] + return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)]) + + root_com = subtree_com[jp.array(m.body_rootid)] + offset = d.xipos - root_com + cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass) + d = d.replace(cinert=cinert) + + # map motion dofs to global frame centered at subtree_com + def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis): + cdofs = [] + + dof_com_fn = lambda a, o: jp.concatenate([a, jp.cross(a, o)]) + + for i, jnt_typ in enumerate(jnt_typs): + offset = root_com - xanchor[i] + if jnt_typ == JointType.FREE: + cdofs.append(jp.eye(3, 6, 3)) # free translation + cdofs.append(jax.vmap(dof_com_fn, in_axes=(0, None))(xmat.T, offset)) + elif jnt_typ == JointType.BALL: + cdofs.append(jax.vmap(dof_com_fn, in_axes=(0, None))(xmat.T, offset)) + elif jnt_typ == JointType.HINGE: + cdof = dof_com_fn(xaxis[i], offset) + cdofs.append(jp.expand_dims(cdof, 0)) + elif jnt_typ == JointType.SLIDE: + cdof = jp.concatenate((jp.zeros((3,)), xaxis[i])) + cdofs.append(jp.expand_dims(cdof, 0)) + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + cdof = jp.concatenate(cdofs) if cdofs else jp.empty((0, 6)) + + return cdof + + cdof = scan.flat( + m, + cdof_fn, + 'jbbjj', + 'v', + m.jnt_type, + root_com, + d.xmat, + d.xanchor, + d.xaxis, + ) + d = d.replace(cdof=cdof) + + return d + + +def crb(m: Model, d: Data) -> Data: + """Runs composite rigid body inertia algorithm.""" + + def crb_fn(crb_child, crb_body): + if crb_child is not None: + crb_body += crb_child + return crb_body + + crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d.cinert, reverse=True) + crb_body = crb_body.at[0].set(0.0) + d = d.replace(crb=crb_body) + + # TODO(erikfrey): do centralized take fn? + crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0) + crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof) + + dof_i, dof_j, diag = [], [], [] + for i in range(m.nv): + diag.append(len(dof_i)) + j = i + while j > -1: + dof_i, dof_j = dof_i + [i], dof_j + [j] + j = m.dof_parentid[j] + + crb_codf_i = jp.take(crb_cdof, jp.array(dof_i), axis=0) + cdof_j = jp.take(d.cdof, jp.array(dof_j), axis=0) + qm = jax.vmap(jp.dot)(crb_codf_i, cdof_j) + + # add armature to diagonal + qm = qm.at[jp.array(diag)].add(m.dof_armature) + + d = d.replace(qM=qm) + + return d + + +def factor_m( + m: Model, + d: Data, + qM: jax.Array, # pylint:disable=invalid-name +) -> Data: + """Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd.""" + + # build up indices for where we will do backwards updates over qLD + # TODO(erikfrey): do fewer updates by combining non-overlapping ranges + dof_madr = jp.array(m.dof_Madr) + updates = {} + madr_ds = [] + for i in range(m.nv): + madr_d = madr_ij = m.dof_Madr[i] + j = i + while True: + madr_ds.append(madr_d) + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + madr_j_range = tuple(m.dof_Madr[j : j + 2]) + updates.setdefault(madr_j_range, []).append((madr_d, madr_ij)) + + qld = qM + + for (out_beg, out_end), vals in sorted(updates.items(), reverse=True): + madr_d, madr_ij = jp.array(vals).T + + @jax.vmap + def off_diag_fn(madr_d, madr_ij, qld=qld, width=out_end - out_beg): + qld_row = jax.lax.dynamic_slice(qld, (madr_ij,), (width,)) + return -(qld_row[0] / qld[madr_d]) * qld_row + + qld_update = jp.sum(off_diag_fn(madr_d, madr_ij), axis=0) + qld = qld.at[out_beg:out_end].add(qld_update) + # TODO(erikfrey): determine if this minimum value guarding is necessary: + # qld = qld.at[dof_madr].set(jp.maximum(qld[dof_madr], _MJ_MINVAL)) + + qld_diag = qld[dof_madr] + qld = (qld / qld[jp.array(madr_ds)]).at[dof_madr].set(qld_diag) + + d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag) + + return d + + +def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array: + """Computes sparse backsubstitution: x = inv(L'*D*L)*y .""" + + updates_i, updates_j = {}, {} + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + updates_i.setdefault(i, []).append((madr_ij, j)) + updates_j.setdefault(j, []).append((madr_ij, i)) + + # x <- inv(L') * x + for j, vals in sorted(updates_j.items(), reverse=True): + madr_ij, i = jp.array(vals).T + x = x.at[j].add(-jp.sum(d.qLD[madr_ij] * x[i])) + + # x <- inv(D) * x + x = x * d.qLDiagInv + + # x <- inv(L) * x + for i, vals in sorted(updates_i.items()): + madr_ij, j = jp.array(vals).T + x = x.at[i].add(-jp.sum(d.qLD[madr_ij] * x[j])) + + return x + + +def dense_m(m: Model, d: Data) -> jax.Array: + """Reconstitute dense mass matrix from qM.""" + + is_, js, madr_ijs = [], [], [] + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij] + + i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs)) + + mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM[madr_ij]) + + # diagonal, upper triangular, lower triangular + mat = jp.diag(d.qM[jp.array(m.dof_Madr)]) + mat + mat.T + + return mat + + +def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array: + """Multiply vector by inertia matrix.""" + + diag_mul = d.qM[jp.array(m.dof_Madr)] * vec + + is_, js, madr_ijs = [], [], [] + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij] + + i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs)) + + out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j]) + out = out.at[j].add(d.qM[madr_ij] * vec[i]) + + return out + + +def com_vel(m: Model, d: Data) -> Data: + """Computes cvel, cdof_dot.""" + + # forward scan down tree: accumulate link center of mass velocity + def fn(parent, jnt_typs, cdof, qvel): + cvel = jp.zeros((6,)) if parent is None else parent[0] + + cross_fn = jax.vmap(math.motion_cross, in_axes=(None, 0)) + cdof_x_qvel = jax.vmap(jp.multiply)(cdof, qvel) + + dof_beg = 0 + cdof_dots = [] + for jnt_typ in jnt_typs: + dof_end = dof_beg + JointType(jnt_typ).dof_width() + if jnt_typ == JointType.FREE: + cvel += jp.sum(cdof_x_qvel[:3], axis=0) + cdof_ang_dot = cross_fn(cvel, cdof[3:]) + cvel += jp.sum(cdof_x_qvel[3:], axis=0) + cdof_dots.append(jp.concatenate((jp.zeros((3, 6)), cdof_ang_dot))) + else: + cdof_dots.append(cross_fn(cvel, cdof[dof_beg:dof_end])) + cvel += jp.sum(cdof_x_qvel[dof_beg:dof_end], axis=0) + dof_beg = dof_end + + cdof_dot = jp.concatenate(cdof_dots) if cdof_dots else jp.empty((0, 6)) + return cvel, cdof_dot + + cvel, cdof_dot = scan.body_tree( + m, + fn, + 'jvv', + 'bv', + m.jnt_type, + d.cdof, + d.qvel, + ) + + d = d.replace(cvel=cvel, cdof_dot=cdof_dot) + + return d + + +def rne(m: Model, d: Data) -> Data: + """Computes inverse dynamics using the recursive Newton-Euler algorithm.""" + # forward scan over tree: accumulate link center of mass acceleration + def cacc_fn(cacc, cdof_dot, qvel): + if cacc is None: + if m.opt.disableflags & DisableBit.GRAVITY: + cacc = jp.zeros((6,)) + else: + cacc = jp.concatenate((jp.zeros((3,)), -m.opt.gravity)) + + cacc += jp.sum(jax.vmap(jp.multiply)(cdof_dot, qvel), axis=0) + + return cacc + + cacc = scan.body_tree(m, cacc_fn, 'vv', 'b', d.cdof_dot, d.qvel) + + def frc(cinert, cacc, cvel): + frc = math.inert_mul(cinert, cacc) + frc += math.motion_cross_force(cvel, math.inert_mul(cinert, cvel)) + + return frc + + loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel) + + # backward scan up tree: accumulate body forces + def cfrc_fn(cfrc_child, cfrc): + if cfrc_child is not None: + cfrc += cfrc_child + return cfrc + + cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True) + qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)]) + + d = d.replace(qfrc_bias=qfrc_bias) + + return d + + +def transmission(m: Model, d: Data) -> Data: + """Computes actuator/transmission lengths and moments.""" + if not m.nu: + return d + + def fn(gear, jnt_typ, m_i, m_j, qpos): + # handles joint transmissions only + if jnt_typ == JointType.FREE: + length = jp.zeros(1) + moment = gear + m_i = jp.repeat(m_i, 6) + m_j = m_j + jp.arange(6) + elif jnt_typ == JointType.BALL: + axis, _ = math.quat_to_axis_angle(qpos) + length = jp.dot(axis, gear[:3])[None] + moment = gear[:3] + m_i = jp.repeat(m_i, 3) + m_j = m_j + jp.arange(3) + elif jnt_typ in (JointType.SLIDE, JointType.HINGE): + length = qpos * gear[0] + moment = gear[:1] + m_i, m_j = m_i[None], m_j[None] + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + return length, moment, m_i, m_j + + length, m_val, m_i, m_j = scan.flat( + m, + fn, + 'ujujq', + 'uvvv', + m.actuator_gear, + m.jnt_type, + jp.arange(m.nu), + jp.array(m.jnt_dofadr), + d.qpos, + group_by='u', + ) + moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val) + d = d.replace(actuator_length=length.squeeze(), actuator_moment=moment) + return d diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py new file mode 100644 index 00000000..4cdfc6ce --- /dev/null +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -0,0 +1,166 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for smooth dynamics functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(a, b, name, step, fname, atol=1e-5, rtol=1e-5): + 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): + 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) + + +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'): + return + + np.random.seed(seed) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + + kinematics_jit_fn = jax.jit(mjx.kinematics) + com_pos_jit_fn = jax.jit(mjx.com_pos) + crb_jit_fn = jax.jit(mjx.crb) + factor_m_fn = jax.jit(mjx.factor_m) + com_vel_jit_fn = jax.jit(mjx.com_vel) + rne_jit_fn = jax.jit(mjx.rne) + mul_m_jit_fn = jax.jit(mjx.mul_m) + transmission_jit_fn = jax.jit(mjx.transmission) + + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + + # kinematics + dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + _assert_attr_eq(d, dx, 'xanchor', i, fname) + _assert_attr_eq(d, dx, 'xaxis', i, fname) + _assert_attr_eq(d, dx, 'xpos', i, fname) + _assert_attr_eq(d, dx, 'xquat', i, fname) + _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname) + _assert_attr_eq(d, dx, 'xipos', i, fname) + _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname) + _assert_attr_eq(d, dx, 'geom_xpos', i, fname) + _assert_eq( + d.geom_xmat.reshape((-1, 3, 3)), + dx.geom_xmat, + 'geom_xmat', + i, + fname, + ) + + # com_pos + dx = com_pos_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'subtree_com', i, fname) + _assert_attr_eq(d, dx, 'cinert', i, fname) + _assert_attr_eq(d, dx, 'cdof', i, fname) + + # crb + dx = crb_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'crb', i, fname) + _assert_attr_eq(d, dx, 'qM', i, fname) + + # 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) + + # com_vel + dx = com_vel_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'cvel', i, fname) + _assert_attr_eq(d, dx, 'cdof_dot', i, fname) + + # rne + dx = rne_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'qfrc_bias', i, fname, atol=1e-4) + + # 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) + + # transmission + dx = transmission_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'actuator_length', i, fname) + _assert_attr_eq(d, dx, 'actuator_moment', i, fname) + + +class DisableGravityTest(absltest.TestCase): + + def test_disabled(self): + m = mujoco.MjModel.from_xml_string(""" + + + """) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + # test with gravity + step_jit_fn = jax.jit(mjx.step) + dx = step_jit_fn(mx, dx) + np.testing.assert_array_almost_equal( + dx.qpos, np.array([0.0, 0.0, -9.81e-4, 1.0, 0.0, 0.0, 0.0]), decimal=7 + ) + + # test with gravity disabled + mx = mx.tree_replace( + {'opt.disableflags': mx.opt.disableflags | DisableBit.GRAVITY} + ) + dx = mjx.device_put(d) + step_jit_fn = jax.jit(mjx.step) + dx = step_jit_fn(mx, dx) + np.testing.assert_equal( + dx.qpos, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py new file mode 100644 index 00000000..0aefd027 --- /dev/null +++ b/mjx/mujoco/mjx/_src/solver.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""CG and Newton solvers.""" + +from typing import Optional + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math +from mujoco.mjx._src import smooth +# pylint: disable=g-importing-member +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 +# pylint: enable=g-importing-member + + +class _CGContext(PyTreeNode): + """Data updated during each cg solver iteration. + + Attributes: + qacc: acceleration (from Data) (nv,) + qfrc_constraint: constraint force (from Data) (nv,) + Jaref: Jac*qacc - aref (nefc,) + efc_force: constraint force in constraint space (nefc,) + M: dense mass matrix, populated for nv < 100 (nv, nv) + Ma: M*qacc (nv,) + grad: gradient of master cost (nv,) + Mgrad: M / grad (nv,) + search: linesearch vector (nv,) + gauss: gauss Cost + cost: constraint + Gauss cost + prev_cost: cost from previous cg iter + solver_niter: number of solver iterations + """ + + qacc: jax.Array + qfrc_constraint: jax.Array + Jaref: jax.Array # pylint: disable=invalid-name + efc_force: jax.Array + M: Optional[jax.Array] + Ma: jax.Array # pylint: disable=invalid-name + grad: jax.Array + Mgrad: jax.Array # pylint: disable=invalid-name + search: jax.Array + gauss: jax.Array + cost: jax.Array + prev_cost: jax.Array + solver_niter: jax.Array + + @classmethod + def create(cls, m: Model, d: Data, grad: bool = True) -> '_CGContext': + 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( + qacc=d.qacc, + qfrc_constraint=d.qfrc_constraint, + Jaref=jaref, + efc_force=-jaref * d.efc_D, + M=M, + Ma=ma, + grad=nv_0, + Mgrad=nv_0, + search=nv_0, + gauss=0.0, + cost=jp.inf, + prev_cost=0.0, + solver_niter=0, + ) + ctx = _cg_update_constraint(m, d, ctx) + if grad: + ctx = _cg_update_gradient(m, d, ctx) + ctx = ctx.replace(search=-ctx.Mgrad) # start with preconditioned gradient + + return ctx + + +class _LSPoint(PyTreeNode): + """Line search evaluation point. + + Attributes: + alpha: step size that reduces f(x + alpha * p) given search direction p + cost: line search cost + deriv_0: first derivative of quadratic + deriv_1: second derivative of quadratic + """ + + alpha: jax.Array + cost: jax.Array + deriv_0: jax.Array + deriv_1: jax.Array + + @classmethod + def create( + cls, + ctx: _CGContext, + alpha: jax.Array, + jv: jax.Array, + quad: jax.Array, + quad_gauss: jax.Array, + ) -> '_LSPoint': + """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 + quad = jax.vmap(jp.multiply)(quad, active) # only active + quad_total = quad_gauss + jp.sum(quad, axis=0) + + cost = alpha * alpha * quad_total[2] + alpha * quad_total[1] + quad_total[0] + deriv_0 = 2 * alpha * quad_total[2] + quad_total[1] + deriv_1 = 2 * quad_total[2] + return _LSPoint(alpha=alpha, cost=cost, deriv_0=deriv_0, deriv_1=deriv_1) + + +class _LSContext(PyTreeNode): + """Data updated during each cg line search iteration. + + Attributes: + lo: low point bounding the line search interval + hi: high point bounding the line search interval + swap: True if low or hi was swapped in the line search iteration + ls_iter: number of linesearch iterations + """ + + lo: _LSPoint + hi: _LSPoint + swap: jax.Array + ls_iter: jax.Array + + +def _while_loop_scan(cond_fun, body_fun, init_val, max_iter): + """Scan-based implementation (jit ok, reverse-mode autodiff ok).""" + def _iter(val): + next_val = body_fun(val) + next_cond = cond_fun(next_val) + return next_val, next_cond + + def _fun(tup, it): + val, cond = tup + # When cond is met, we start doing no-ops. + return jax.lax.cond(cond, _iter, lambda x: (x, False), val), it + + init = (init_val, cond_fun(init_val)) + 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. + + 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 + + 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 + + jaref = ctx.Jaref * (ctx.Jaref < 0) # non-negative constraints + efc_force = -jaref * d.efc_D + 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 + + ctx = ctx.replace( + qfrc_constraint=qfrc_constraint, + gauss=gauss, + cost=cost, + prev_cost=ctx.cost, + efc_force=efc_force, + ) + + return ctx + + +def _cg_update_gradient(m: Model, d: Data, ctx: _CGContext) -> _CGContext: + """Updates grad and M / grad given latest CG 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 + + Returns: + context with new grad and M / grad + """ + + grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint + mgrad = smooth.solve_m(m, d, grad) + + ctx = ctx.replace(grad=grad, Mgrad=mgrad) + + return ctx + + +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: + """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 + + Returns: + updated context with new qacc, Ma, Jaref + """ + smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv) + gtol = m.opt.tolerance * m.opt.ls_tolerance * smag + + # compute Mv, Jv + mv = smooth.mul_m(m, d, ctx.search) if ctx.M is None else ctx.M @ ctx.search + jv = d.efc_J @ ctx.search + + # prepare quadratics + quad_gauss = jp.stack(( + ctx.gauss, + jp.dot(ctx.search, ctx.Ma) - jp.dot(ctx.search, d.qfrc_smooth), + 0.5 * jp.dot(ctx.search, mv), + )) + quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv)) + quad = (quad * d.efc_D).T + + point_fn = lambda alpha: _LSPoint.create(ctx, alpha, jv, quad, quad_gauss) + + def cond(ctx: _LSContext) -> jax.Array: + done = ctx.ls_iter >= m.opt.ls_iterations + done |= ~ctx.swap # if we did not adjust the interval + done |= (ctx.lo.deriv_0 < 0) & (ctx.lo.deriv_0 > -gtol) + done |= (ctx.hi.deriv_0 > 0) & (ctx.hi.deriv_0 < gtol) + + return ~done + + def body(ctx: _LSContext) -> _LSContext: + # always compute new bracket boundaries and a midpoint + lo, hi = ctx.lo, ctx.hi + lo_next = point_fn(lo.alpha - lo.deriv_0 / lo.deriv_1) + hi_next = point_fn(hi.alpha - hi.deriv_0 / hi.deriv_1) + mid = point_fn(0.5 * (lo.alpha + hi.alpha)) + + # we swap lo/hi if: + # 1) they are not correctly at a bracket boundary (e.g. lo.deriv_0 > 0), OR + # 2) if moving to next or mid narrows the bracket + swap_lo_next = (lo.deriv_0 > 0) | (lo.deriv_0 < lo_next.deriv_0) + lo = jax.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next) + swap_lo_mid = (mid.deriv_0 < 0) & (lo.deriv_0 < mid.deriv_0) + lo = jax.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid) + + swap_hi_next = (hi.deriv_0 < 0) | (hi.deriv_0 > hi_next.deriv_0) + hi = jax.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next) + swap_hi_mid = (mid.deriv_0 > 0) & (hi.deriv_0 > mid.deriv_0) + hi = jax.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid) + + swap = swap_lo_next | swap_lo_mid | swap_hi_next | swap_hi_mid + + ctx = ctx.replace(lo=lo, hi=hi, swap=swap, ls_iter=ctx.ls_iter + 1) + + return ctx + + # initialize interval + p0 = point_fn(jp.array(0.0)) + lo = point_fn(p0.alpha - p0.deriv_0 / p0.deriv_1) + lesser_fn = lambda x, y: jp.where(lo.deriv_0 < p0.deriv_0, x, y) + hi = jax.tree_map(lesser_fn, p0, lo) + lo = jax.tree_map(lesser_fn, lo, p0) + ls_ctx = _LSContext(lo=lo, hi=hi, swap=jp.array(True), ls_iter=0) + ls_ctx = _while_loop_scan(cond, body, ls_ctx, m.opt.ls_iterations) + + # move to new solution if improved + lo, hi = ls_ctx.lo, ls_ctx.hi + improved = (lo.cost < p0.cost) | (hi.cost < p0.cost) + alpha = jp.where(lo.cost < hi.cost, lo.alpha, hi.alpha) + qacc = ctx.qacc + improved * ctx.search * alpha + ma = ctx.Ma + improved * mv * alpha + jaref = ctx.Jaref + improved * jv * alpha + + ctx = ctx.replace(qacc=qacc, Ma=ma, Jaref=jaref) + + return ctx + + +def cg_solve(m: Model, d: Data) -> Data: + """Finds forces that satisfy constraints using conjugate gradient descent.""" + + def cond(ctx: _CGContext) -> jax.Array: + improvement = _rescale(m, ctx.prev_cost - ctx.cost) + gradient = _rescale(m, math.norm(ctx.grad)) + + done = ctx.solver_niter >= m.opt.iterations + done |= improvement < m.opt.tolerance + done |= gradient < m.opt.tolerance + + return ~done + + def body(ctx: _CGContext) -> _CGContext: + ctx = _cg_search(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) + + # polak-ribiere: + beta = jp.dot(ctx.grad, ctx.Mgrad - prev_Mgrad) + beta = beta / jp.maximum(mujoco.mjMINVAL, jp.dot(prev_grad, prev_Mgrad)) + beta = jp.maximum(0, beta) + search = -ctx.Mgrad + beta * ctx.search + ctx = ctx.replace(search=search, solver_niter=ctx.solver_niter + 1) + + return ctx + + # 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) + 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)) + + d = d.replace( + qacc_warmstart=ctx.qacc, + qacc=ctx.qacc, + qfrc_constraint=ctx.qfrc_constraint, + efc_force=ctx.efc_force, + ) + + return d diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py new file mode 100644 index 00000000..6a4f7792 --- /dev/null +++ b/mjx/mujoco/mjx/_src/solver_test.py @@ -0,0 +1,130 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for forward functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import mujoco +from mujoco import mjx +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-2, rtol=1e-2): + 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) + + +class Solver64Test(parameterized.TestCase): + """Tests solvers at 64 bit precision.""" + + def setUp(self): + super().setUp() + jax.config.update('jax_enable_x64', True) + + def tearDown(self): + super().tearDown() + jax.config.update('jax_enable_x64', False) + + @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) + def test_cg(self, seed, fname): + """Test mjx cg solver matches mujoco cg solver at 64 bit precision.""" + f = epath.resource_path('mujoco.mjx') / 'test_data' / fname + m = mujoco.MjModel.from_xml_string(f.read_text()) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + + jax.config.update('jax_enable_x64', True) + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + np.random.seed(seed) + d.qvel = 0.01 * np.random.random(m.nv) + + for i in range(100): + # in order to avoid re-jitting, reuse the same mj_data shape + save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth + d = mujoco.MjData(m) + d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + # at 64 bits the solutions returned by the two solvers are quite close + self.assertLessEqual(dx.solver_niter[0], d.solver_niter[0]) + _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname) + _assert_attr_eq(d, dx, 'qacc', i, fname) + + +class SolverTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) + def test_cg(self, seed, fname): + """Test mjx cg solver is close to mj at 32 bit precision. + + Args: + seed: int + fname: file to test + + At lower float resolution there's wiggle room in valid forces that satisfy + constraints. So instead let's mainly validate that mjx is finding solutions + with as good cost as mujoco, even if the resulting forces/accelerations + are not quite the same. + """ + f = epath.resource_path('mujoco.mjx') / 'test_data' / fname + m = mujoco.MjModel.from_xml_string(f.read_text()) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + np.random.seed(seed) + d.qvel = 0.01 * np.random.random(m.nv) + + for i in range(100): + # in order to avoid re-jitting, reuse the same mj_data shape + save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth + d = mujoco.MjData(m) + d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + def cost(qacc): + jaref = np.zeros(d.nefc) + mujoco.mj_mulJacVec(m, d, jaref, qacc) + jaref -= d.efc_aref + cost = np.array([0.0]) + mujoco.mj_constraintUpdate(m, d, jaref, cost, 0) + return cost[0] + + cost_mj, cost_mjx = cost(d.qacc), cost(dx.qacc) + + self.assertLessEqual( + cost_mjx, + cost_mj * 1.01, + msg=f'mismatch: {fname} at step {i}, cost too high', + ) + _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname, atol=1e-1, rtol=1e-1) + _assert_attr_eq(d, dx, 'qacc', i, fname, atol=1e-1, rtol=1e-1) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py new file mode 100644 index 00000000..cfbdfadf --- /dev/null +++ b/mjx/mujoco/mjx/_src/support.py @@ -0,0 +1,81 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Engine support functions.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import scan +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def jac( + m: Model, d: Data, point: jax.Array, body_id: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Compute pair of (NV, 3) Jacobians of global point attached to body.""" + fn = lambda carry, b: b if carry is None else b + carry + mask = (jp.arange(m.nbody) == body_id) * 1 + mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True) + mask = mask[jp.array(m.dof_bodyid)] > 0 + + offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]] + jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof) + jacp = jax.vmap(jp.multiply)(jacp, mask) + jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask) + + return jacp, jacr + + +def jac_dif_pair( + m: Model, + d: Data, + pos: jax.Array, + body_1: jax.Array, + body_2: jax.Array, +) -> jax.Array: + """Compute Jacobian difference for two body points.""" + jacp2, _ = jac(m, d, pos, body_2) + jacp1, _ = jac(m, d, pos, body_1) + return jacp2 - jacp1 + + +def apply_ft( + m: Model, + d: Data, + force: jax.Array, + torque: jax.Array, + point: jax.Array, + body_id: jax.Array, +) -> jax.Array: + """Apply Cartesian force and torque.""" + jacp, jacr = jac(m, d, point, body_id) + return jacp @ force + jacr @ torque + + +def xfrc_accumulate(m: Model, d: Data) -> jax.Array: + """Accumulate xfrc_applied into a qfrc.""" + qfrc = jax.vmap(apply_ft, in_axes=(None, None, 0, 0, 0, 0))( + m, + d, + d.xfrc_applied[:, :3], + d.xfrc_applied[:, 3:], + d.xipos, + jp.arange(m.nbody), + ) + return jp.sum(qfrc, axis=0) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py new file mode 100644 index 00000000..fe88fc84 --- /dev/null +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -0,0 +1,81 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for support.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import support +from mujoco.mjx._src import test_util +import numpy as np + + +class SupportTest(parameterized.TestCase): + + @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + def test_jac(self, fname): + np.random.seed(0) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + mx = mjx.device_put(m) + dx = mjx.device_put(d) + point = np.random.randn(3) + body = np.random.choice(m.nbody) + jacp, jacr = jax.jit(support.jac)(mx, dx, point, body) + + jacp_expected, jacr_expected = np.zeros((3, m.nv)), np.zeros((3, m.nv)) + mujoco.mj_jac(m, d, jacp_expected, jacr_expected, point, body) + np.testing.assert_almost_equal(jacp, jacp_expected.T, 6) + np.testing.assert_almost_equal(jacr, jacr_expected.T, 6) + + def test_xfrc_accumulate(self): + """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate.""" + np.random.seed(0) + + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + mx = mjx.device_put(m) + dx = mjx.device_put(d) + self.assertFalse((dx.xipos == 0.0).all()) + + xfrc = np.random.rand(*dx.xfrc_applied.shape) + + d.xfrc_applied[:] = xfrc + dx = dx.replace(xfrc_applied=jp.array(xfrc)) + + qfrc = jax.jit(support.xfrc_accumulate)(mx, dx) + qfrc_expected = np.zeros(m.nv) + for i in range(1, m.nbody): + mujoco.mj_applyFT( + m, + d, + d.xfrc_applied[i, :3], + d.xfrc_applied[i, 3:], + d.xipos[i], + i, + qfrc_expected, + ) + + np.testing.assert_almost_equal(qfrc, qfrc_expected, 6) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py new file mode 100644 index 00000000..5ac50816 --- /dev/null +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -0,0 +1,340 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for testing.""" + +import sys +from typing import Dict, List, Tuple +from xml.etree import ElementTree as ET + +from etils import epath +import mujoco +import numpy as np + +TEST_FILES: List[str] = [ + 'ant.xml', + 'ball_pendulum.xml', + 'cherry_pendulum.xml', + 'convex.xml', + 'humanoid.xml', + 'mixed_joint_pendulum.xml', + 'single_pendulum.xml', + 'slide_pendulum.xml', + 'triple_pendulum.xml', + 'triple_pendulum_free.xml', + 'weld.xml', +] + +_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] +_JOINT_TYPES = ['free', 'hinge', 'slide', 'ball'] +_JOINT_AXES = ['1 0 0', '0 1 0', '0 0 1'] +_FRICTIONS = ['1.2 0.003 0.0002', '0.2 0.0001 0.0005'] +_KP_POS = ['1', '2'] +_KP_INTVEL = ['10000', '2000'] +_KV_VEL = ['123', '1'] +_PAIR_FRICTIONS = ['1.2 0.9 0.003 0.0002 0.0001'] +_SOLREFS = ['0.04 1.01', '0.05 1.02', '0.03 1.1', '0.015 1.0'] +_SOLIMPS = [ + '0.75 0.94 0.002 0.2 2', + '0.8 0.99 0.001 0.3 6', + '0.6 0.9 0.003 0.1 1', +] +_DIMS = ['3'] +_MARGINS = ['0.0', '0.01', '0.02'] +_GAPS = ['0.0', '0.005'] +_GEARS = ['20', '50', '100'] + + +def p(pct: int) -> bool: + assert 0 <= pct <= 100 + return np.random.uniform(low=0, high=100) < pct + + +def _make_joint(joint_type: str, name: str) -> Dict[str, str]: + """Returns attributes for a joint.""" + joint_attr = {'type': joint_type, 'name': name} + + if joint_type not in ('free', 'ball'): + joint_attr['axis'] = np.random.choice(_JOINT_AXES) + lb, ub = -np.random.uniform() * 90, np.random.uniform() * 90 + joint_attr['range'] = f'{lb:.2f} {ub:.2f}' + elif joint_type == 'ball': + joint_attr['axis'] = '1 0 0' + ub = np.random.uniform() * 90 + joint_attr['range'] = f'0.0 {ub:.2f}' + + if p(50) and joint_type != 'free': + lb, ub = -np.random.uniform(), np.random.uniform() + joint_attr['actuatorfrcrange'] = f'{lb:.2f} {ub:.2f}' + + if joint_type not in ('free',): + joint_attr['damping'] = '{:.2f}'.format(np.random.uniform() * 20) + joint_attr['stiffness'] = '{:.2f}'.format(np.random.uniform() * 20) + + return joint_attr + + +def _geom_solparams( + pair: bool = False, enable_contact: bool = True +) -> Dict[str, str]: + """Returns geom solver parameters.""" + params = { + 'contype': np.random.choice(['0', '1']) if enable_contact else '0', + 'conaffinity': np.random.choice(['0', '1']) if enable_contact else '0', + 'priority': np.random.choice(['-1', '2']), + 'solmix': np.random.choice(['0.0', '1.6']), + 'friction': np.random.choice(_FRICTIONS), + 'condim': np.random.choice(_DIMS), + } + pair_params = { + 'solreffriction': np.random.choice(_SOLREFS), + 'friction': np.random.choice(_PAIR_FRICTIONS), + 'condim': np.random.choice(_DIMS), + } + params = pair_params if pair else params + params.update({ + 'solimp': np.random.choice(_SOLIMPS), + 'solref': np.random.choice(_SOLREFS), + 'margin': np.random.choice(_MARGINS), + 'gap': np.random.choice(_GAPS), + }) + + return params + + +def _make_geom( + pos: str, size: float, name: str, enable_contact: bool = True +) -> Dict[str, str]: + """Returns attributes for a sphere geom.""" + attr = { + 'pos': pos, + 'type': 'sphere', + 'name': name, + 'size': f'{size:.2f}', + 'mass': '1', + } + attr.update(_geom_solparams(pair=False, enable_contact=enable_contact)) + + return attr + + +def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]: + """Returns attributes for an actuator.""" + attr = {'joint': joint} + if actuator_type == 'motor': + attr['gear'] = np.random.choice(_GEARS) + elif actuator_type == 'position': + attr['kp'] = np.random.choice(_KP_POS) + elif actuator_type == 'general': + attr['biastype'] = 'affine' + attr['gainprm'] = '35 0 0' + attr['biasprm'] = '0 -35 -0.65' + elif actuator_type == 'intvelocity': + attr['kp'] = np.random.choice(_KP_INTVEL) + lb, ub = -np.random.uniform(), np.random.uniform() + attr['actrange'] = f'{lb:.2f} {ub:.2f}' + elif actuator_type == 'velocity': + attr['kv'] = np.random.choice(_KV_VEL) + + if p(50) and actuator_type != 'intvelocity': + lb, ub = -np.random.uniform(), np.random.uniform() + attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}' + + if p(50): + lb, ub = -np.random.uniform(), np.random.uniform() + attr['forcerange'] = f'{lb*10:.2f} {ub*10:.2f}' + + return attr + + +def create_mjcf( + seed: int, + min_trees: int = 1, + max_trees: int = 1, + max_tree_depth: int = 5, + body_pos: Tuple[float, float, float] = (0.0, 0.0, -0.5), + geom_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0), + max_stacked_joints=4, + max_geoms_per_body=2, + max_contact_excludes=1, + max_contact_pairs=4, + disable_actuation_pct: int = 0, + add_actuators: bool = False, + root_always_free: bool = False, + enable_contact: bool = True, +) -> str: + """Creates a random MJCF for testing. + + Args: + seed: seed for rng + min_trees: minimum number of kinematic trees to generate + max_trees: maximum number of kinematic trees to generate + max_tree_depth: the maximum tree depth + body_pos: the default body position relative to the parent + geom_pos: the default geom position in the body frame + max_stacked_joints: maximum number of joints to stack for each body + max_geoms_per_body: maximum number of geoms per body + max_contact_excludes: maximum number of bodies to exlude from contact + max_contact_pairs: maximum number of explicit geom contact pairs in the xml + disable_actuation_pct: the percentage of time to disable actuation via the + disable flag + add_actuators: whether to add actuators + root_always_free: if True, the root body of each kinematic tree has a free + joint with the world + enable_contact: if False, disables all contacts via contype/conaffinity + + Returns: + an XML string for the MuJoCo config + Raises: + AssertionError when args are not in the correct ranges + """ + np.random.seed(seed) + + assert min_trees <= max_trees + assert max_tree_depth >= 1 + assert 0 <= disable_actuation_pct <= 100 + assert max_stacked_joints >= 1 + assert max_geoms_per_body >= 1 + assert max_contact_excludes >= 1 + assert max_contact_pairs >= 1 + + mjcf = ET.Element('mujoco') + opt = ET.SubElement(mjcf, 'option', {'timestep': '0.005', 'solver': 'CG'}) + world = ET.SubElement(mjcf, 'worldbody') + ET.SubElement(mjcf, 'compiler', {'autolimits': 'true'}) + + # disable flags + if p(disable_actuation_pct): + ET.SubElement(opt, 'flag', {'actuation': 'disable'}) + + ET.SubElement( + world, + 'geom', + { + 'name': 'plane', + 'type': 'plane', + 'contype': '1' if enable_contact else '0', + 'conaffinity': '1' if enable_contact else '0', + 'size': '40 40 40', + }, + ) + + # kinematic trees + tree_depth = np.random.randint(1, max_tree_depth + 1) + + def make_tree(body: ET.Element, depth: int) -> None: + if depth >= tree_depth: + return + + z_pos = np.random.uniform(low=-1, high=1) * 0.01 # small jitter + pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}' + n_bodies = len(list(mjcf.iter('body'))) + child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'}) + + n_joints = len(list(mjcf.iter('joint'))) + for nj in range(np.random.randint(1, max_stacked_joints + 1)): + joint_type = np.random.choice(_JOINT_TYPES) + if nj == 0 and depth == 0 and root_always_free: + joint_type = 'free' + + # free joint only allowed at top level + while joint_type == 'free' and (depth > 0 or nj > 0): + joint_type = np.random.choice(_JOINT_TYPES) + + joint_attr = _make_joint(joint_type, name=f'joint{n_joints + nj}') + ET.SubElement(child, 'joint', joint_attr) + + prev_joints = child.findall('joint') + had_ball_or_free = any( + [j.get('type') in ('ball', 'free') for j in prev_joints] + ) + if had_ball_or_free: + break # do not stack more joints + + n_geoms = len(list(mjcf.iter('geom'))) + for _ in range(np.random.randint(1, max_geoms_per_body + 1)): + pos = ('{:.2f} ' * 3).format(*geom_pos).strip() + size = 0.2 + np.random.uniform(low=-1, high=1) * 0.02 + geom_attr = _make_geom( + pos, size, name=f'geom{n_geoms}', enable_contact=enable_contact + ) + ET.SubElement(child, 'geom', geom_attr) + n_geoms += 1 + + make_tree(child, depth + 1) + + num_trees = np.random.randint(min_trees, max_trees + 1) + for _ in range(num_trees): + make_tree(world, 0) + + # actuators + if add_actuators: + actuator = ET.SubElement(mjcf, 'actuator') + n_joints = len(list(mjcf.iter('joint'))) + nu = np.random.randint(1, n_joints + 1) + actuators = [] + for i in range(nu): + actuator_type = np.random.choice(_ACTUATOR_TYPES) + attr = _make_actuator(actuator_type, joint=f'joint{i}') + actuators.append((actuator_type, attr)) + + np.random.shuffle(actuators) + for typ, attr in actuators: + ET.SubElement(actuator, typ, attr) + + # contact pairs + contact = ET.SubElement(mjcf, 'contact') + geoms = list(mjcf.iter('geom')) + geom_names = [geom.get('name') for geom in geoms] + n_geoms = len(geoms) + pairs = set() + for _ in range(min(max_contact_pairs, n_geoms * (n_geoms - 1) // 2)): + if p(80): + continue + + geom1, geom2 = np.random.choice(geom_names, replace=False, size=2) + if geom1 > geom2: + geom1, geom2 = geom2, geom1 + + if (geom1, geom2) in pairs: + continue + + pairs.add((geom1, geom2)) + attr = {'geom1': geom1, 'geom2': geom2} + attr.update(_geom_solparams(pair=True)) + ET.SubElement(contact, 'pair', attr) + + # exclude contacts + bodies = list(mjcf.iter('body')) + body_names = [b.get('name') for b in bodies] + n_bodies = len(bodies) + for _ in range(min(max_contact_excludes, (n_bodies * (n_bodies - 1) // 2))): + if p(50): + continue + + body1, body2 = np.random.choice(body_names, replace=False, size=2) + ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2}) + + # ElementTree.indent is not available before Python 3.9 + if sys.version_info.minor >= 9: + ET.indent(mjcf) + + return ET.tostring(mjcf).decode('utf-8') + + +def load_test_file(name: str) -> mujoco.MjModel: + """Loads a mujoco.MjModel based on the file name.""" + path = epath.resource_path('mujoco.mjx') / 'test_data' / name + m = mujoco.MjModel.from_xml_path(path.as_posix()) + return m diff --git a/mjx/mujoco/mjx/_src/test_util_test.py b/mjx/mujoco/mjx/_src/test_util_test.py new file mode 100644 index 00000000..e8d32dab --- /dev/null +++ b/mjx/mujoco/mjx/_src/test_util_test.py @@ -0,0 +1,38 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the test_util.""" + +from absl.testing import absltest +from etils import epath +from mujoco.mjx._src import test_util + + +class TestUtilTest(absltest.TestCase): + + def test_files_in_test_data_match(self): + directory = epath.resource_path('mujoco.mjx') / 'test_data' + files = set([f.name for f in directory.glob('*.xml')]) + self.assertSetEqual( + files, + set(test_util.TEST_FILES), + msg=( + '`_test_util.TEST_FILES` must match the files in the ' + 'test_data/*.xml directory' + ), + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py new file mode 100644 index 00000000..bd3f9fe1 --- /dev/null +++ b/mjx/mujoco/mjx/_src/types.py @@ -0,0 +1,660 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Base types used in MJX.""" + +import enum +from typing import Sequence + +import jax +import jax.numpy as jp +import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +# pylint: enable=g-importing-member +import numpy as np + + +class DisableBit(enum.IntFlag): + """Disable default feature bitflags. + + Attributes: + CONSTRAINT: entire constraint solver + EQUALITY: equality constraints + FRICTIONLOSS: joint and tendon frictionloss constraints + LIMIT: joint and tendon limit constraints + CONTACT: contact constraints + PASSIVE: passive forces + GRAVITY: gravitational forces + CLAMPCTRL: clamp control to specified range + WARMSTART: warmstart constraint solver + ACTUATION: apply actuation forces + REFSAFE: integrator safety: make ref[0]>=2*timestep + """ + CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT + EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY + LIMIT = mujoco.mjtDisableBit.mjDSBL_LIMIT + CONTACT = mujoco.mjtDisableBit.mjDSBL_CONTACT + PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE + GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY + CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL + WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART + ACTUATION = mujoco.mjtDisableBit.mjDSBL_ACTUATION + REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE + EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP + FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT + # unsupported: FRICTIONLOSS, SENSOR, MIDPHASE + + +class JointType(enum.IntEnum): + """Type of degree of freedom. + + Attributes: + FREE: global position and orientation (quat) (7,) + BALL: orientation (quat) relative to parent (4,) + SLIDE: sliding distance along body-fixed axis (1,) + HINGE: rotation angle (rad) around body-fixed axis (1,) + """ + FREE = mujoco.mjtJoint.mjJNT_FREE + BALL = mujoco.mjtJoint.mjJNT_BALL + SLIDE = mujoco.mjtJoint.mjJNT_SLIDE + HINGE = mujoco.mjtJoint.mjJNT_HINGE + + def dof_width(self) -> int: + return {0: 6, 1: 3, 2: 1, 3: 1}[self.value] + + def qpos_width(self) -> int: + return {0: 7, 1: 4, 2: 1, 3: 1}[self.value] + + +class IntegratorType(enum.IntEnum): + """Integrator mode. + + Attributes: + EULER: semi-implicit Euler + RK4: 4th-order Runge Kutta + """ + EULER = mujoco.mjtIntegrator.mjINT_EULER + RK4 = mujoco.mjtIntegrator.mjINT_RK4 + # unsupported: IMPLICIT, IMPLICITFAST + + +class GeomType(enum.IntEnum): + """Type of geometry. + + Attributes: + PLANE: plane + HFIELD: height field + SPHERE: sphere + CAPSULE: capsule + ELLIPSOID: ellipsoid + CYLINDER: cylinder + BOX: box + MESH: mesh + """ + + PLANE = mujoco.mjtGeom.mjGEOM_PLANE + HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD + SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE + CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE + ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID + CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER + BOX = mujoco.mjtGeom.mjGEOM_BOX + MESH = mujoco.mjtGeom.mjGEOM_MESH + # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE + + +class ConeType(enum.IntEnum): + """Type of friction cone. + + Attributes: + PYRAMIDAL: pyramidal + """ + PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL + # unsupported: ELLIPTIC + + +class SolverType(enum.IntEnum): + """Constraint solver algorithm. + + Attributes: + CG: Conjugate gradient (primal) + """ + # unsupported: PGS, NEWTON + CG = mujoco.mjtSolver.mjSOL_CG + + +class EqType(enum.IntEnum): + """Type of equality constraint. + + Attributes: + CONNECT: connect two bodies at a point (ball joint) + WELD: fix relative position and orientation of two bodies + JOINT: couple the values of two scalar joints with cubic + """ + CONNECT = mujoco.mjtEq.mjEQ_CONNECT + WELD = mujoco.mjtEq.mjEQ_WELD + # unsupported: JOINT, TENDON, DISTANCE + + +class TrnType(enum.IntEnum): + """Type of actuator transmission. + + Attributes: + JOINT: force on joint + """ + JOINT = mujoco.mjtTrn.mjTRN_JOINT + # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, SITE, BODY + + +class DynType(enum.IntEnum): + """Type of actuator dynamics. + + Attributes: + NONE: no internal dynamics; ctrl specifies force + INTEGRATOR: integrator: da/dt = u + """ + NONE = mujoco.mjtDyn.mjDYN_NONE + INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR + FILTER = mujoco.mjtDyn.mjDYN_FILTER + # unsupported: FILTEREXACT, MUSCLE, USER + + +class GainType(enum.IntEnum): + """Type of actuator gain. + + Attributes: + FIXED: fixed gain + AFFINE: const + kp*length + kv*velocity + """ + FIXED = mujoco.mjtGain.mjGAIN_FIXED + AFFINE = mujoco.mjtGain.mjGAIN_AFFINE + # unsupported: MUSCLE, USER + + +class BiasType(enum.IntEnum): + """Type of actuator bias. + + Attributes: + NONE: no bias + AFFINE: const + kp*length + kv*velocity + """ + NONE = mujoco.mjtBias.mjBIAS_NONE + AFFINE = mujoco.mjtBias.mjBIAS_AFFINE + # unsupported: MUSCLE, USER + + +class Option(PyTreeNode): + """Physics options. + + Attributes: + timestep: timestep + tolerance: main solver tolerance + ls_tolerance: CG/Newton linesearch tolerance + gravity: gravitational acceleration (3,) + wind: wind (for lift, drag and viscosity) + density: density of medium + viscosity: viscosity of medium + has_fluid_params: automatically set by mjx if wind/density/viscosity are + nonzero. Not used by mj + integrator: integration mode + cone: type of friction cone + solver: solver algorithm + integrator: integration mode + iterations: number of main solver iterations + ls_iterations: maximum number of CG/Newton linesearch iterations + disableflags: bit flags for disabling standard features + """ + timestep: jax.Array + tolerance: jax.Array + ls_tolerance: jax.Array + # unsupported: apirate, impratio, noslip_tolerance, mpr_tolerance + gravity: jax.Array + wind: jax.Array + density: jax.Array + viscosity: jax.Array + has_fluid_params: bool + # unsupported: magnetic, o_margin, o_solref, o_solimp + integrator: IntegratorType + cone: ConeType + # unsupported: jacobian + solver: SolverType + iterations: int + ls_iterations: int + # unsupported: noslip_iterations, mpr_iterations + disableflags: DisableBit + # unsupported: enableflags + + +class Statistic(PyTreeNode): + """Model statistics (in qpos0). + + Attributes: + meaninertia: mean diagonal inertia + """ + meaninertia: jax.Array + # unsupported: meanmass, meansize, extent, center + + +class Model(PyTreeNode): + """Static model of the scene that remains unchanged with each physics step. + + Attributes: + nq: number of generalized coordinates = dim(qpos) + nv: number of degrees of freedom = dim(qvel) + nu: number of actuators/controls = dim(ctrl) + na: number of activation states = dim(act) + nbody: number of bodies + njnt: number of joints + ngeom: number of geoms + nmesh: number of meshes + npair: number of predefined geom pairs + nexclude: number of excluded geom pairs + neq: number of equality constraints + nnumeric: number of numeric custom fields + nM: number of non-zeros in sparse inertia matrix + opt: physics options + stat: model statistics + qpos0: qpos values at default pose (nq,) + qpos_spring: reference pose for springs (nq,) + body_parentid: id of body's parent (nbody,) + body_rootid: id of root above body (nbody,) + body_weldid: id of body that this body is welded to (nbody,) + body_jntnum: number of joints for this body (nbody,) + body_jntadr: start addr of joints; -1: no joints (nbody,) + body_dofnum: number of motion degrees of freedom (nbody,) + body_dofadr: start addr of dofs; -1: no dofs (nbody,) + body_geomnum: number of geoms (nbody,) + body_geomadr: start addr of geoms; -1: no geoms (nbody,) + body_pos: position offset rel. to parent body (nbody, 3) + body_quat: orientation offset rel. to parent body (nbody, 4) + body_ipos: local position of center of mass (nbody, 3) + body_iquat: local orientation of inertia ellipsoid (nbody, 4) + body_mass: mass (nbody,) + body_subtreemass: mass of subtree starting at this body (nbody,) + body_inertia: diagonal inertia in ipos/iquat frame (nbody, 3) + body_invweight0: mean inv inert in qpos0 (trn, rot) (nbody, 2) + jnt_type: type of joint (mjtJoint) (njnt,) + jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) + jnt_dofadr: start addr in 'qvel' for joint's data (njnt,) + jnt_bodyid: id of joint's body (njnt,) + jnt_group: group for visibility (njnt,) + jnt_limited: does joint have limits (njnt,) + jnt_solref: constraint solver reference: limit (njnt, mjNREF) + jnt_solimp: constraint solver impedance: limit (njnt, mjNIMP) + jnt_pos: local anchor position (njnt, 3) + jnt_axis: local joint axis (njnt, 3) + jnt_stiffness: stiffness coefficient (njnt,) + jnt_range: joint limits (njnt, 2) + jnt_actfrcrange: range of total actuator force (njnt, 2) + jnt_margin: min distance for limit detection (njnt,) + dof_bodyid: id of dof's body (nv,) + dof_jntid: id of dof's joint (nv,) + dof_parentid: id of dof's parent; -1: none (nv,) + dof_Madr: dof address in M-diagonal (nv,) + dof_solref: constraint solver reference:frictionloss (nv, mjNREF) + dof_solimp: constraint solver impedance:frictionloss (nv, mjNIMP) + dof_frictionloss: dof friction loss (nv,) + dof_armature: dof armature inertia/mass (nv,) + dof_damping: damping coefficient (nv,) + dof_invweight0: diag. inverse inertia in qpos0 (nv,) + dof_M0: diag. inertia in qpos0 (nv,) + geom_type: geometric type (mjtGeom) (ngeom,) + geom_contype: geom contact type (ngeom,) + geom_conaffinity: geom contact affinity (ngeom,) + geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) + geom_bodyid: id of geom's body (ngeom,) + geom_priority: geom contact priority (ngeom,) + geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) + geom_solref: constraint solver reference: contact (ngeom, mjNREF) + geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) + geom_size: geom-specific size parameters (ngeom, 3) + geom_pos: local position offset rel. to body (ngeom, 3) + geom_quat: local orientation offset rel. to body (ngeom, 4) + geom_friction: friction for (slide, spin, roll) (ngeom, 3) + geom_margin: include in solver if dist 'Contact': + """Returns a contact filled with zeros.""" + return Contact( + dist=jp.zeros(shape), + pos=jp.zeros(shape + (3,)), + frame=jp.zeros(shape + (3, 3)), + includemargin=jp.zeros(shape), + friction=jp.zeros(shape + (5,)), + solref=jp.zeros(shape + (mujoco.mjNREF,)), + solreffriction=jp.zeros(shape + (mujoco.mjNREF,)), + solimp=jp.zeros(shape + (mujoco.mjNIMP,)), + dim=np.zeros(shape, dtype=np.int32), + geom1=jp.zeros(shape, dtype=jp.int32), + geom2=jp.zeros(shape, dtype=jp.int32), + efc_address=np.zeros(shape, dtype=np.int32), + ) + + +class Data(PyTreeNode): + """Dynamic state that updates each step. + + Attributes: + solver_niter: number of solver iterations, per island (mjNISLAND,) + ne: number of equality constraints + nf: number of friction constraints + nl: number of limit constraints + nefc: number of constraints + ncon: nubmer of contacts + time: simulation time + qpos: position (nq,) + qvel: velocity (nv,) + act: actuator activation (na,) + qacc_warmstart: acceleration used for warmstart (nv,) + ctrl: control (nu,) + qfrc_applied: applied generalized force (nv,) + xfrc_applied: applied Cartesian force/torque (nbody, 6) + eq_active: enable/disable constraints (neq,) + qacc: acceleration (nv,) + act_dot: time-derivative of actuator activation (na,) + xpos: Cartesian position of body frame (nbody, 3) + xquat: Cartesian orientation of body frame (nbody, 4) + xmat: Cartesian orientation of body frame (nbody, 3, 3) + xipos: Cartesian position of body com (nbody, 3) + ximat: Cartesian orientation of body inertia (nbody, 3, 3) + xanchor: Cartesian position of joint anchor (njnt, 3) + xaxis: Cartesian joint axis (njnt, 3) + geom_xpos: Cartesian geom position (ngeom, 3) + geom_xmat: Cartesian geom orientation (ngeom, 3, 3) + subtree_com: center of mass of each subtree (nbody, 3) + cdof: com-based motion axis of each dof (nv, 6) + cinert: com-based body inertia and mass (nbody, 10) + actuator_length: actuator lengths (nu,) + actuator_moment: actuator moments (nu, nv) + crb: com-based composite inertia and mass (nbody, 10) + qM: total inertia (sparse) (nM,) + qLD: L'*D*L factorization of M (sparse) (nM,) + qLDiagInv: 1/diag(D) (nv,) + qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,) + contact: list of all detected contacts (ncon,) + efc_J: constraint Jacobian (nefc, nv) + efc_frictionloss: frictionloss (friction) (nefc,) + efc_D: constraint mass (nefc,) + actuator_velocity: actuator velocities (nu,) + cvel: com-based velocity [3D rot; 3D tran] (nbody, 6) + cdof_dot: time-derivative of cdof (nv, 6) + qfrc_bias: C(qpos,qvel) (nv,) + qfrc_passive: passive force (nv,) + efc_aref: reference pseudo-acceleration (nefc,) + actuator_force: actuator force in actuation space (nu,) + qfrc_actuator: actuator force (nv,) + qfrc_smooth: net unconstrained force (nv,) + qacc_smooth: unconstrained acceleration (nv,) + qfrc_constraint: constraint force (nv,) + qfrc_inverse: net external force; should equal: (nv,) + qfrc_applied + J'*xfrc_applied + qfrc_actuator + efc_force: constraint force in constraint space (nefc,) + """ + # solver statistics: + solver_niter: jax.Array + # sizes (variable in MJ, constant in MJX) + ne: int + nf: int + nl: int + nefc: int + ncon: int + # global properties: + time: jax.Array + # state: + qpos: jax.Array + qvel: jax.Array + act: jax.Array + qacc_warmstart: jax.Array + # control: + ctrl: jax.Array + qfrc_applied: jax.Array + xfrc_applied: jax.Array + eq_active: jax.Array + # dynamics: + qacc: jax.Array + act_dot: jax.Array + # position dependent: + xpos: jax.Array + xquat: jax.Array + xmat: jax.Array + xipos: jax.Array + ximat: jax.Array + xanchor: jax.Array + xaxis: jax.Array + geom_xpos: jax.Array + geom_xmat: jax.Array + subtree_com: jax.Array + cdof: jax.Array + cinert: jax.Array + crb: jax.Array + actuator_length: jax.Array + actuator_moment: jax.Array + qM: jax.Array + qLD: jax.Array + qLDiagInv: jax.Array + qLDiagSqrtInv: jax.Array + contact: Contact + efc_J: jax.Array + efc_frictionloss: jax.Array + efc_D: jax.Array + # position, velocity dependent: + actuator_velocity: jax.Array + cvel: jax.Array + cdof_dot: jax.Array + qfrc_bias: jax.Array + qfrc_passive: jax.Array + efc_aref: jax.Array + # position, velcoity, control & acceleration dependent: + actuator_force: jax.Array + qfrc_actuator: jax.Array + qfrc_smooth: jax.Array + qacc_smooth: jax.Array + qfrc_constraint: jax.Array + qfrc_inverse: jax.Array + efc_force: jax.Array diff --git a/mjx/mujoco/mjx/benchmark/__init__.py b/mjx/mujoco/mjx/benchmark/__init__.py new file mode 100644 index 00000000..4ad51d69 --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/mjx/mujoco/mjx/benchmark/benchmark.py b/mjx/mujoco/mjx/benchmark/benchmark.py new file mode 100644 index 00000000..b0a59af0 --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/benchmark.py @@ -0,0 +1,116 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Run benchmarks on various devices.""" + +import sys +import time + +from absl import flags +from etils import epath +import google_benchmark as benchmark +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx + +FLAGS = flags.FLAGS + +_PATHS = { + 'humanoid': 'benchmark/model/humanoid/humanoid.xml', + 'barkour': 'benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml', + 'shadow_hand': 'benchmark/model/shadow_hand/scene_right.xml', +} + + +_BATCH_SIZE = { + ('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', 'cpu'): 64, + ('shadow_hand', 'cpu'): 64, +} + +flags.DEFINE_string('model', 'humanoid', 'Model to benchmark') +flags.DEFINE_string('device', 'cpu', 'Device benchmark is running on') + + +def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float: + """Reports jit time and op time for a function.""" + + step_count = 100 if FLAGS.device == 'cpu' else 1000 + + @jax.jit + def run_batch(seed: jp.ndarray): + rngs = jax.random.split(jax.random.PRNGKey(seed), batch_size) + init_state = jax.vmap(init_fn)(rngs) + + @jax.vmap + def run(state): + def step(state, _): + state = step_fn(state) + return state, () + + return jax.lax.scan(step, state, (), length=step_count) + + return run(init_state) + + # run once to jit + beg = time.perf_counter() + jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(0)) + first_t = time.perf_counter() - beg + + times = [] + while state: + beg = time.perf_counter() + batch = run_batch(jp.array(len(times))) + jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch) + times.append(time.perf_counter() - beg) + + op_time = jp.mean(jp.array(times)) + batch_sps = batch_size * step_count / op_time + + state.counters['jit_time'] = first_t - op_time + state.counters['batch_sps'] = batch_sps + + +@benchmark.option.unit(benchmark.kSecond) +def _run(state: benchmark.State): + """Benchmark a model.""" + + f = epath.resource_path('mujoco.mjx') / _PATHS[FLAGS.model] + m = mujoco.MjModel.from_xml_path(f.as_posix()) + m = mjx.device_put(m) + + def init(rng): + d = mjx.make_data(m) + qvel = 0.01 * jax.random.normal(rng, shape=(m.nv,)) + d = d.replace(qvel=qvel) + return d + + def step(d): + return mjx.step(m, d) + + batch_size = _BATCH_SIZE[(FLAGS.model, jax.devices()[0].device_kind)] + _measure_fn(state, init, step, batch_size=batch_size) + + +if __name__ == '__main__': + FLAGS(sys.argv) + benchmark.register(_run, name=FLAGS.model + '_' + FLAGS.device) + benchmark.main() diff --git a/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md b/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md new file mode 100644 index 00000000..b7cfdcef --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md @@ -0,0 +1,35 @@ +# Google Barkour v0 Joystick Policy + +## Overview + +This folder contains a training script for a flat-terrain joystick policy for the [Barkour v0 Quadruped](https://ai.googleblog.com/2023/05/barkour-benchmarking-animal-level.html) which demonstrates sim2real transfer. + +`barkour_joystick.py` contains the environment definition, while the [colab](https://colab.research.google.com/github/google/brax/blob/main/experimental/barkour_v0/barkour_v0_joystick.ipynb) shows how to train the policy. + +

+ +

+ +## MJCF Instructions + +The MuJoCo config in `assets/barkour_v0_mjx.xml` was copied from https://github.com/deepmind/mujoco_menagerie/google_barkour_v0. The following edits were made to the MJCF specifically for brax: + +* `meshdir` was changed from `assets` to `.`. +* `frictionloss` was removed. `damping` was changed to 0.5239. +* A custom `init_qpos` was added. +* A sphere geom `lowerLegFoot` was added to all feet. All other contacts were turned off. +* The compiler option was changed to `