Updated version numbers to cmake files

PiperOrigin-RevId: 604406364
Change-Id: I85ab41b048fff9484c81e858c2bbcfeb6066fd4b

.github/workflows/build.yml

@@ -164,6 +164,7 @@ jobs:
       if: ${{ runner.os != 'Windows' }}
       working-directory: build
       run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
+           cp lib/libactuator.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libelasticity.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libsensor.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp lib/libsdf.* ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
@@ -171,6 +172,7 @@ jobs:
       if: ${{ runner.os == 'Windows' }}
       working-directory: build
       run: mkdir -p ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
+           cp bin/Release/actuator.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp bin/Release/elasticity.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin &&
            cp bin/Release/sensor.dll ${{ matrix.tmpdir }}/mujoco_install/mujoco_plugin
     - name: Configure samples
@@ -219,7 +221,7 @@ jobs:
       working-directory: python/dist
       run: >
         source ${{ matrix.tmpdir }}/venv/bin/activate &&
-        pip install mujoco-*.whl
+        pip install --no-index mujoco-*.whl
     - name: Test Python bindings
       if: ${{ runner.os != 'Windows' }}
       shell: bash
@@ -228,23 +230,56 @@ jobs:
       run: >
         source ${{ matrix.tmpdir }}/venv/bin/activate &&
         pytest -v --pyargs mujoco
+    - name: Package MJX
+      if: ${{ runner.os != 'Windows' }}
+      shell: bash
+      working-directory: mjx
+      run:
+        source ${{ matrix.tmpdir }}/venv/bin/activate &&
+        python -m build .
+    - name: Install MJX
+      if: ${{ runner.os != 'Windows' }}
+      shell: bash
+      working-directory: mjx
+      run:
+        source ${{ matrix.tmpdir }}/venv/bin/activate &&
+        pip install --require-hashes -r requirements.txt &&
+        pip install --no-index dist/mujoco_mjx-*.whl
+    - name: Test MJX
+      if: ${{ runner.os != 'Windows' }}
+      shell: bash
+      run:
+        source ${{ matrix.tmpdir }}/venv/bin/activate &&
+        pytest -n auto -v --pyargs mujoco.mjx
     - name: Notify team chat
       shell: bash
       env:
-        GCHAT_API: ${{ secrets.GCHAT_API }}
-        JOB_URL: "${{ github.server_url }}/${{ github.repository }}/actions/runs/\
-          ${{ github.run_id }}"
-      if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API != '' }}
-      run: >
-        curl
-        "$GCHAT_API&threadKey=$GITHUB_SHA&messageReplyOption=REPLY_MESSAGE_FALLBACK_TO_NEW_THREAD"
-        -X POST
-        -H "Content-Type: application/json"
-        -d "{
-        'text': '<$JOB_URL|*FAILURE*>:
-        job \`${{ matrix.os }}${{ matrix.additional_label }}\`
-        commit \`$(echo $GITHUB_SHA | head -c6)\`\n
-        \`\`\`Author: ${{ github.event.head_commit.author.name }}
-        <${{ github.event.head_commit.author.email }}>
-        \n\n${{ github.event.head_commit.message }}\`\`\`
-        '}"
+        GCHAT_API_URL: ${{ secrets.GCHAT_API }}
+        JOB_URL: "${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}"
+        CHATMSG_AUTHOR_NAME: ${{ github.event.head_commit.author.name }}
+        CHATMSG_AUTHOR_EMAIL: ${{ github.event.head_commit.author.email }}
+        CHATMSG_COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
+        CHATMSG_JOB_ID: ${{ matrix.os }}${{ matrix.additional_label }}
+      if: ${{ failure() && github.event_name == 'push' && env.GCHAT_API_URL != '' }}
+      run: |
+        CHATMSG="$(cat <<-'EOF' | python3
+        import json
+        import os
+        env = lambda x: os.getenv(x, '')
+        data = dict(
+            result=env('JOB_URL'),
+            job=env('CHATMSG_JOB_ID'),
+            commit=env('GITHUB_SHA')[:6],
+            name=env('CHATMSG_AUTHOR_NAME').replace('```', ''),
+            email=env('CHATMSG_AUTHOR_EMAIL'),
+            msg=env('CHATMSG_COMMIT_MESSAGE').replace('```', '')
+        )
+        text = '<{result}|*FAILURE*>: job `{job}` commit `{commit}`\n```Author: {name}<{email}>\n\n{msg}```'.format(**data)
+        print(json.dumps({'text' : text}))
+        EOF
+        )" &&
+
+        curl "$GCHAT_API_URL&threadKey=$GITHUB_SHA&messageReplyOption=REPLY_MESSAGE_FALLBACK_TO_NEW_THREAD" \
+        -X POST \
+        -H "Content-Type: application/json" \
+        --data-raw "${CHATMSG}"

CMakeLists.txt

@@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON)
 
 project(
   mujoco
-  VERSION 2.3.8
+  VERSION 3.1.2
   DESCRIPTION "MuJoCo Physics Simulator"
   HOMEPAGE_URL "https://mujoco.org"
 )
@@ -83,6 +83,7 @@ target_include_directories(
 )
 
 add_subdirectory(plugin/elasticity)
+add_subdirectory(plugin/actuator)
 add_subdirectory(plugin/sensor)
 add_subdirectory(plugin/sdf)
 add_subdirectory(src/engine)

README.md

@@ -55,6 +55,10 @@ running on Google Colab:
  - For a more advanced example, see the LQR tutorial which creates an LQR
    controller to balance a humanoid on one leg using MuJoCo's dynamics
    derivatives: [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/LQR.ipynb)
+ - The MJX tutorial provides usage examples of
+   [MuJoCo XLA](https://mujoco.readthedocs.io/en/stable/mjx.html), a branch of MuJoCo written in
+   JAX:
+   [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb)
 
 ## Installation
 
@@ -160,7 +164,7 @@ These packages give users of various languages access to MuJoCo functionality:
   by [Manoj Velmurugan](https://github.com/vmanoj1996).
 - **Swift**: [swift-mujoco](https://github.com/liuliu/swift-mujoco)
 - **Java**: [mujoco-java](https://github.com/CommonWealthRobotics/mujoco-java)
-- **Julia**: [Lyceum](https://github.com/Lyceum/MuJoCo.jl) (unmaintained)
+- **Julia**: [MuJoCo.jl](https://github.com/JamieMair/MuJoCo.jl)
 
 
 ### Converters

cmake/MujocoDependencies.cmake

@@ -39,27 +39,27 @@ set(MUJOCO_DEP_VERSION_qhull
     CACHE STRING "Version of `qhull` to be fetched."
 )
 set(MUJOCO_DEP_VERSION_Eigen3
-    211c5dfc6741a5570ad007983c113ef4d144f9f3
+    7fd7a3f946e5ac152d28dad388cff8bfa1026925
     CACHE STRING "Version of `Eigen3` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_abseil
-    c2435f8342c2d0ed8101cb43adfd605fdc52dca2 # LTS 20230125.3
+    4a2c63365eff8823a5221db86ef490e828306f9d # LTS 20240116.0
     CACHE STRING "Version of `abseil` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_gtest
-    b796f7d44681514f58a683a3a71ff17c94edb0c1 # v1.13.0
+    f8d7d77c06936315286eb55f8de22cd23c188571 # v1.14.0
     CACHE STRING "Version of `gtest` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_benchmark
-    2dd015dfef425c866d9a43f2c67d8b52d709acb6 # v1.8.0
+    e45585a4b8e75c28479fa4107182c28172799640 # v1.8.3
     CACHE STRING "Version of `benchmark` to be fetched."
 )
 
 set(MUJOCO_DEP_VERSION_sdflib
-    492847fa81e46653114da48e8886730ccefed377
+    7c49cfba9bbec763b5d0f7b90b26555f3dde8088
     CACHE STRING "Version of `SdfLib` to be fetched."
 )
 
@@ -184,6 +184,9 @@ findorfetch(
   EXCLUDE_FROM_ALL
 )
 
+option(SDFLIB_USE_ASSIMP OFF)
+option(SDFLIB_USE_OPENMP OFF)
+option(SDFLIB_USE_ENOKI OFF)
 findorfetch(
   USE_SYSTEM_PACKAGE
   OFF
@@ -195,8 +198,6 @@ findorfetch(
   https://github.com/UPC-ViRVIG/SdfLib.git
   GIT_TAG
   ${MUJOCO_DEP_VERSION_sdflib}
-  PATCH_COMMAND
-  git apply --reject --whitespace=fix ${mujoco_SOURCE_DIR}/cmake/sdflib-optional-dependencies.patch
   TARGETS
   SdfLib
   EXCLUDE_FROM_ALL

cmake/sdflib-optional-dependencies.patch

@@ -1,585 +0,0 @@
-diff --git a/CMakeLists.txt b/CMakeLists.txt
-index 20551cf..0d2a364 100644
---- a/CMakeLists.txt
-+++ b/CMakeLists.txt
-@@ -70,15 +70,25 @@ add_custom_target(copyShaders ALL SOURCES ${SHADER_FILES})
- # Add dependencies
- add_subdirectory(libs)
- 
-+if(SDFLIB_USE_ENOKI)
-+    target_link_libraries(${PROJECT_NAME} PUBLIC enoki)
-+    target_link_libraries(${PROJECT_NAME} PUBLIC fcpw)
-+    target_compile_definitions(${PROJECT_NAME} PUBLIC -DENOKI_AVAILABLE)
-+endif()
-+
-+if(SDFLIB_USE_ASSIMP)
-+    target_link_libraries(${PROJECT_NAME} PUBLIC assimp)
-+    target_compile_definitions(${PROJECT_NAME} PUBLIC -DASSIMP_AVAILABLE)
-+endif()
-+
-+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
-+    target_link_libraries(${PROJECT_NAME} PUBLIC args)
-+    target_link_libraries(${PROJECT_NAME} PUBLIC stb_image)
-+endif()
-+    
- target_link_libraries(${PROJECT_NAME} PUBLIC glm)
--target_link_libraries(${PROJECT_NAME} PUBLIC assimp)
--target_link_libraries(${PROJECT_NAME} PUBLIC args)
- target_link_libraries(${PROJECT_NAME} PUBLIC spdlog)
- target_link_libraries(${PROJECT_NAME} PUBLIC cereal)
--target_link_libraries(${PROJECT_NAME} PUBLIC enoki)
--target_link_libraries(${PROJECT_NAME} PUBLIC eigen)
--target_link_libraries(${PROJECT_NAME} PUBLIC fcpw)
--target_link_libraries(${PROJECT_NAME} PUBLIC stb_image)
- target_link_libraries(${PROJECT_NAME} PUBLIC icg)
- 
- if(CMAKE_CXX_COMPILER_ID MATCHES GNU)
-@@ -86,17 +96,22 @@ if(CMAKE_CXX_COMPILER_ID MATCHES GNU)
- endif()
- 
- # Add openMP
--if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
--    message("Enabling openmp llvm extension")
--    set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp:llvm")
--else()
--    find_package(OpenMP)
--    if(NOT OpenMP_CXX_FOUND)
--    	message(FATAL_ERROR  "OpenMP not found")
--    endif()
--    message("OpenMP version ${OpenMP_CXX_VERSION}")
--    target_link_libraries(${PROJECT_NAME} PUBLIC OpenMP::OpenMP_CXX)
--endif()                   
-+if(SDFLIB_USE_OPENMP)
-+    if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
-+        message("Enabling openmp llvm extension")
-+        set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp:llvm")
-+        target_compile_definitions(${PROJECT_NAME} PUBLIC -DOPENMP_AVAILABLE)
-+    else()
-+        find_package(OpenMP)
-+        if(OpenMP_CXX_FOUND)
-+        	message("OpenMP version ${OpenMP_CXX_VERSION}")
-+            target_link_libraries(${PROJECT_NAME} PUBLIC OpenMP::OpenMP_CXX)
-+            target_compile_definitions(${PROJECT_NAME} PUBLIC -DOPENMP_AVAILABLE)
-+        else()
-+            message("Disabling openmp")
-+        endif()
-+    endif()    
-+endif()
- 
- # Add executable
- if (NOT UNIX)
-@@ -160,6 +175,7 @@ if(SDFLIB_BUILD_DEBUG_APPS)
-     add_executable(GJKtest src/tools/GJKtest/main.cpp)
-     target_link_libraries(GJKtest PUBLIC ${PROJECT_NAME})
- 
-+    target_link_libraries(${PROJECT_NAME} PUBLIC eigen)
-     add_executable(CalculateInterpolationParameters src/tools/CalculateInterpolationParameters/main.cpp)
-     target_link_libraries(CalculateInterpolationParameters PUBLIC ${PROJECT_NAME})
- 
-diff --git a/include/SdfLib/ExactOctreeSdf.h b/include/SdfLib/ExactOctreeSdf.h
-index 79ad82d..06dd7ac 100644
---- a/include/SdfLib/ExactOctreeSdf.h
-+++ b/include/SdfLib/ExactOctreeSdf.h
-@@ -214,6 +214,8 @@ private:
- };
- }
- 
-+#ifdef OPENMP_AVAILABLE
- #include "ExactOctreeSdfDepthFirst.h"
-+#endif
- 
- #endif
-\ No newline at end of file
-diff --git a/include/SdfLib/InterpolationMethods.h b/include/SdfLib/InterpolationMethods.h
-index 077dfb2..f707d5b 100644
---- a/include/SdfLib/InterpolationMethods.h
-+++ b/include/SdfLib/InterpolationMethods.h
-@@ -4,7 +4,10 @@
- #include <array>
- 
- #include "utils/TriangleUtils.h"
-+
-+#ifdef ENOKI_AVAILABLE
- #include "enoki/array.h"
-+#endif
- 
- namespace sdflib
- {
-@@ -236,15 +239,6 @@ struct TriLinearInterpolation
- //         outCoeff[63] = 8 * inValues[0][0] + 4 * inValues[0][1] * nodeSize + 4 * inValues[0][2] * nodeSize + 4 * inValues[0][3] * nodeSize + -8 * inValues[1][0] + 4 * inValues[1][1] * nodeSize + -4 * inValues[1][2] * nodeSize + -4 * inValues[1][3] * nodeSize + -8 * inValues[2][0] + -4 * inValues[2][1] * nodeSize + 4 * inValues[2][2] * nodeSize + -4 * inValues[2][3] * nodeSize + 8 * inValues[3][0] + -4 * inValues[3][1] * nodeSize + -4 * inValues[3][2] * nodeSize + 4 * inValues[3][3] * nodeSize + -8 * inValues[4][0] + -4 * inValues[4][1] * nodeSize + -4 * inValues[4][2] * nodeSize + 4 * inValues[4][3] * nodeSize + 8 * inValues[5][0] + -4 * inValues[5][1] * nodeSize + 4 * inValues[5][2] * nodeSize + -4 * inValues[5][3] * nodeSize + 8 * inValues[6][0] + 4 * inValues[6][1] * nodeSize + -4 * inValues[6][2] * nodeSize + -4 * inValues[6][3] * nodeSize + -8 * inValues[7][0] + 4 * inValues[7][1] * nodeSize + 4 * inValues[7][2] * nodeSize + 4 * inValues[7][3] * nodeSize + 0.0f;
- //     }
- 
--//     inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart) 
--//     {
--//         return 0.0f
--//          + values[0] + values[1] * fracPart[0] + values[2] * fracPart[0] * fracPart[0] + values[3] * fracPart[0] * fracPart[0] * fracPart[0] + values[4] * fracPart[1] + values[5] * fracPart[0] * fracPart[1] + values[6] * fracPart[0] * fracPart[0] * fracPart[1] + values[7] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] + values[8] * fracPart[1] * fracPart[1] + values[9] * fracPart[0] * fracPart[1] * fracPart[1] + values[10] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[11] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[12] * fracPart[1] * fracPart[1] * fracPart[1] + values[13] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[14] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[15] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1]
--//          + values[16] * fracPart[2] + values[17] * fracPart[0] * fracPart[2] + values[18] * fracPart[0] * fracPart[0] * fracPart[2] + values[19] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] + values[20] * fracPart[1] * fracPart[2] + values[21] * fracPart[0] * fracPart[1] * fracPart[2] + values[22] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[23] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[24] * fracPart[1] * fracPart[1] * fracPart[2] + values[25] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[26] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[27] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[28] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[29] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[30] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[31] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2]
--//          + values[32] * fracPart[2] * fracPart[2] + values[33] * fracPart[0] * fracPart[2] * fracPart[2] + values[34] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[35] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[36] * fracPart[1] * fracPart[2] * fracPart[2] + values[37] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[38] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[39] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[40] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[41] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[42] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[43] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[44] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[45] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[46] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[47] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2]
--//          + values[48] * fracPart[2] * fracPart[2] * fracPart[2] + values[49] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[50] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[51] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[52] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[53] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[54] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[55] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[56] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[57] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[58] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[59] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[60] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[61] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[62] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[63] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2];
--//     }
--
- //     inline static void interpolateVertexValues(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart, float nodeSize, std::array<float, VALUES_PER_VERTEX>& outValues)
- //     {
- //         outValues[0] = 0.0f
-@@ -383,6 +377,7 @@ struct TriCubicInterpolation
-         outCoeff[63] = 8 * inValues[0][0] + 4 * inValues[0][1] + 4 * inValues[0][2] + 4 * inValues[0][3] + 2 * inValues[0][4] + 2 * inValues[0][5] + 2 * inValues[0][6] + 1 * inValues[0][7] + -8 * inValues[1][0] + 4 * inValues[1][1] + -4 * inValues[1][2] + -4 * inValues[1][3] + 2 * inValues[1][4] + 2 * inValues[1][5] + -2 * inValues[1][6] + 1 * inValues[1][7] + -8 * inValues[2][0] + -4 * inValues[2][1] + 4 * inValues[2][2] + -4 * inValues[2][3] + 2 * inValues[2][4] + -2 * inValues[2][5] + 2 * inValues[2][6] + 1 * inValues[2][7] + 8 * inValues[3][0] + -4 * inValues[3][1] + -4 * inValues[3][2] + 4 * inValues[3][3] + 2 * inValues[3][4] + -2 * inValues[3][5] + -2 * inValues[3][6] + 1 * inValues[3][7] + -8 * inValues[4][0] + -4 * inValues[4][1] + -4 * inValues[4][2] + 4 * inValues[4][3] + -2 * inValues[4][4] + 2 * inValues[4][5] + 2 * inValues[4][6] + 1 * inValues[4][7] + 8 * inValues[5][0] + -4 * inValues[5][1] + 4 * inValues[5][2] + -4 * inValues[5][3] + -2 * inValues[5][4] + 2 * inValues[5][5] + -2 * inValues[5][6] + 1 * inValues[5][7] + 8 * inValues[6][0] + 4 * inValues[6][1] + -4 * inValues[6][2] + -4 * inValues[6][3] + -2 * inValues[6][4] + -2 * inValues[6][5] + 2 * inValues[6][6] + 1 * inValues[6][7] + -8 * inValues[7][0] + 4 * inValues[7][1] + 4 * inValues[7][2] + 4 * inValues[7][3] + -2 * inValues[7][4] + -2 * inValues[7][5] + -2 * inValues[7][6] + 1 * inValues[7][7];
-     }
- 
-+#ifdef ENOKI_AVAILABLE
-     using vec4 = enoki::Array<float, 4>;
- 
-     inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart) 
-@@ -433,6 +428,16 @@ struct TriCubicInterpolation
- 
-         return sum;
-     }
-+#else
-+    inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart) 
-+    {
-+        return 0.0f
-+         + values[0] + values[1] * fracPart[0] + values[2] * fracPart[0] * fracPart[0] + values[3] * fracPart[0] * fracPart[0] * fracPart[0] + values[4] * fracPart[1] + values[5] * fracPart[0] * fracPart[1] + values[6] * fracPart[0] * fracPart[0] * fracPart[1] + values[7] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] + values[8] * fracPart[1] * fracPart[1] + values[9] * fracPart[0] * fracPart[1] * fracPart[1] + values[10] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[11] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[12] * fracPart[1] * fracPart[1] * fracPart[1] + values[13] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[14] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[15] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1]
-+         + values[16] * fracPart[2] + values[17] * fracPart[0] * fracPart[2] + values[18] * fracPart[0] * fracPart[0] * fracPart[2] + values[19] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] + values[20] * fracPart[1] * fracPart[2] + values[21] * fracPart[0] * fracPart[1] * fracPart[2] + values[22] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[23] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[24] * fracPart[1] * fracPart[1] * fracPart[2] + values[25] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[26] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[27] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[28] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[29] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[30] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[31] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2]
-+         + values[32] * fracPart[2] * fracPart[2] + values[33] * fracPart[0] * fracPart[2] * fracPart[2] + values[34] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[35] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[36] * fracPart[1] * fracPart[2] * fracPart[2] + values[37] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[38] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[39] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[40] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[41] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[42] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[43] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[44] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[45] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[46] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[47] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2]
-+         + values[48] * fracPart[2] * fracPart[2] * fracPart[2] + values[49] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[50] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[51] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[52] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[53] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[54] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[55] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[56] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[57] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[58] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[59] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[60] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[61] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[62] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[63] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2];
-+    }
-+#endif
- 
-     inline static glm::vec3 interpolateGradient(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart) 
-     {
-@@ -493,4 +498,4 @@ struct TriCubicInterpolation
- };
- }
- 
--#endif
-\ No newline at end of file
-+#endif
-diff --git a/include/SdfLib/TrianglesInfluence.h b/include/SdfLib/TrianglesInfluence.h
-index 3f3d33f..fc2ca52 100644
---- a/include/SdfLib/TrianglesInfluence.h
-+++ b/include/SdfLib/TrianglesInfluence.h
-@@ -1,7 +1,10 @@
- #ifndef TRIANGLES_INFLUENCE_H
- #define TRIANGLES_INFLUENCE_H
- 
-+#ifdef ENOKI_AVAILABLE
- #include <fcpw/fcpw.h>
-+#endif
-+
- #include "utils/Mesh.h"
- #include "utils/TriangleUtils.h"
- #include "OctreeSdfUtils.h"
-@@ -1008,6 +1011,7 @@ struct VHQueries
-     }
- };
- 
-+#ifdef ENOKI_AVAILABLE
- template<typename T>
- struct FCPWQueries
- {
-@@ -1118,6 +1122,8 @@ struct FCPWQueries
-     {
-     }
- };
-+#endif
-+
- }
- 
--#endif
-\ No newline at end of file
-+#endif
-diff --git a/include/SdfLib/utils/Mesh.h b/include/SdfLib/utils/Mesh.h
-index 28d5486..7d21e44 100644
---- a/include/SdfLib/utils/Mesh.h
-+++ b/include/SdfLib/utils/Mesh.h
-@@ -4,9 +4,11 @@
- #include <string>
- #include <vector>
- #include <glm/glm.hpp>
-+#ifdef ASSIMP_AVAILABLE
- #include <assimp/Importer.hpp>
- #include <assimp/scene.h>      
- #include <assimp/postprocess.h>
-+#endif
- #include "SdfLib/utils/UsefullSerializations.h"
- 
- namespace sdflib
-@@ -43,6 +45,23 @@ struct BoundingBox
-         return glm::length(glm::max(q,glm::vec3(0.0f))) + glm::min(glm::max(q.x, glm::max(q.y,q.z)),0.0f);
-     }
- 
-+    float getDistance(glm::vec3 point, glm::vec3& outGradient) const
-+    {
-+        glm::vec3 a = glm::abs(point) - getSize();
-+        int k = a[0] > a[1] ? 0 : 1;
-+        int l = a[2] > a[k] ? 2 : k;
-+        if (a[l] < 0) {
-+            outGradient[l] = point[l] / glm::abs(point[l]);
-+        } else {
-+            glm::vec3 b = glm::max(a, glm::vec3(0.0f));
-+            float c = glm::length(b);
-+            outGradient[0] = a[0] > 0 ? b[0] / c * point[0] / glm::abs(point[0]) : 0;
-+            outGradient[1] = a[1] > 0 ? b[1] / c * point[1] / glm::abs(point[1]) : 0;
-+            outGradient[2] = a[2] > 0 ? b[2] / c * point[2] / glm::abs(point[2]) : 0;
-+        }
-+        return getDistance(point);
-+    }
-+
-     template<class Archive>
-     void serialize(Archive & archive)
-     {
-@@ -54,8 +73,10 @@ class Mesh
- {
- public:
-     Mesh() {}
-+#ifdef ASSIMP_AVAILABLE
-     Mesh(std::string filePath);
-     Mesh(const aiMesh* mesh);
-+#endif
-     Mesh(glm::vec3* vertices, uint32_t numVertices,
-          uint32_t* indices, uint32_t numIndices);
- 
-@@ -74,7 +95,9 @@ public:
-     void computeNormals();
-     void applyTransform(glm::mat4 trans);
- private:
-+#ifdef ASSIMP_AVAILABLE
-     void initMesh(const aiMesh* mesh);
-+#endif
- 
-     std::vector<glm::vec3> mVertices;
-     std::vector<uint32_t> mIndices;
-@@ -83,4 +106,4 @@ private:
- };
- }
- 
--#endif
-\ No newline at end of file
-+#endif
-diff --git a/include/SdfLib/utils/TriangleUtils.h b/include/SdfLib/utils/TriangleUtils.h
-index 9f930ed..6ee2304 100644
---- a/include/SdfLib/utils/TriangleUtils.h
-+++ b/include/SdfLib/utils/TriangleUtils.h
-@@ -2,6 +2,7 @@
- #define TRIANGLE_UTILS_H
- 
- #include <glm/glm.hpp>
-+#include <algorithm>
- #include <vector>
- #include <array>
- #include <map>
-@@ -401,4 +402,4 @@ namespace TriangleUtils
- }
- }
- 
--#endif
-\ No newline at end of file
-+#endif
-diff --git a/libs/CMakeLists.txt b/libs/CMakeLists.txt
-index b48bf39..ea22b9a 100644
---- a/libs/CMakeLists.txt
-+++ b/libs/CMakeLists.txt
-@@ -14,37 +14,41 @@ if(NOT glm_lib_POPULATED)
- endif()
- 
- # assimp
--FetchContent_Declare(assimp_lib
--  GIT_REPOSITORY https://github.com/assimp/assimp.git
--  GIT_TAG 9519a62dd20799c5493c638d1ef5a6f484e5faf1 # 5.2.5
--)
--
--if(NOT assimp_lib)
--	FetchContent_Populate(assimp_lib)
--	
--	set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
--	set(BUILD_SHARED_LIBS OFF)
--	set(ASSIMP_BUILD_ASSIMP_TOOLS OFF)
--	set(ASSIMP_BUILD_TESTS OFF)
--	set(ASSIMP_INSTALL OFF)
--	set(ASSIMP_INJECT_DEBUG_POSTFIX OFF)
--	set(ASSIMP_BUILD_ASSIMP_VIEW OFF)
-+if(SDF_USE_ASSIMP)
-+	FetchContent_Declare(assimp_lib
-+	  GIT_REPOSITORY https://github.com/assimp/assimp.git
-+	  GIT_TAG 9519a62dd20799c5493c638d1ef5a6f484e5faf1 # 5.2.5
-+	)
- 	
--	add_subdirectory(${assimp_lib_SOURCE_DIR} ${assimp_lib_BINARY_DIR})
-+	if(NOT assimp_lib)
-+		FetchContent_Populate(assimp_lib)
-+		
-+		set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
-+		set(BUILD_SHARED_LIBS OFF)
-+		set(ASSIMP_BUILD_ASSIMP_TOOLS OFF)
-+		set(ASSIMP_BUILD_TESTS OFF)
-+		set(ASSIMP_INSTALL OFF)
-+		set(ASSIMP_INJECT_DEBUG_POSTFIX OFF)
-+		set(ASSIMP_BUILD_ASSIMP_VIEW OFF)
-+		
-+		add_subdirectory(${assimp_lib_SOURCE_DIR} ${assimp_lib_BINARY_DIR})
-+	endif()
- endif()
- 
- # args
--FetchContent_Declare(args_lib
--  GIT_REPOSITORY https://github.com/Taywee/args.git
--  GIT_TAG a48e1f880813b367d2354963a58dedbf2b708584 # 6.3.0
--)
--
--FetchContent_GetProperties(args_lib)
--if(NOT args_lib_POPULATED)
--	FetchContent_Populate(args_lib)
--	add_library(args INTERFACE)
--	target_include_directories(args INTERFACE ${args_lib_SOURCE_DIR})
--endif()
-+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
-+	FetchContent_Declare(args_lib
-+	  GIT_REPOSITORY https://github.com/Taywee/args.git
-+	  GIT_TAG a48e1f880813b367d2354963a58dedbf2b708584 # 6.3.0
-+	)
-+	
-+	FetchContent_GetProperties(args_lib)
-+	if(NOT args_lib_POPULATED)
-+		FetchContent_Populate(args_lib)
-+		add_library(args INTERFACE)
-+		target_include_directories(args INTERFACE ${args_lib_SOURCE_DIR})
-+	endif()
-+ endif()
- 
- # spdlog
- FetchContent_Declare(spdlog_lib
-@@ -76,47 +80,53 @@ if(NOT cereal_lib_POPULATED)
- endif()
- 
- # Enoki
--FetchContent_Declare(enoki_lib
--	GIT_REPOSITORY https://github.com/mitsuba-renderer/enoki.git
--	GIT_TAG 2a18afa
--)
--FetchContent_GetProperties(enoki_lib)
--if(NOT enoki_lib_POPULATED)
--	FetchContent_Populate(enoki_lib)
--	add_library(enoki INTERFACE)
--	add_subdirectory(${enoki_lib_SOURCE_DIR} ${enoki_lib_BINARY_DIR})
--	target_include_directories(enoki INTERFACE ${enoki_lib_SOURCE_DIR}/include)
--endif()
--
--# eigen
--FetchContent_Declare(eigen_lib
--GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
--GIT_TAG 46126273552afe13692929523d34006f54c19719 # 3.4
--)
-+if(SDFLIB_USE_ENOKI)
-+	FetchContent_Declare(enoki_lib
-+		GIT_REPOSITORY https://github.com/mitsuba-renderer/enoki.git
-+		GIT_TAG 2a18afa
-+	)
-+	FetchContent_GetProperties(enoki_lib)
-+	if(NOT enoki_lib_POPULATED)
-+		FetchContent_Populate(enoki_lib)
-+		add_library(enoki INTERFACE)
-+		add_subdirectory(${enoki_lib_SOURCE_DIR} ${enoki_lib_BINARY_DIR})
-+		target_include_directories(enoki INTERFACE ${enoki_lib_SOURCE_DIR}/include)
-+	endif()
- 
--FetchContent_GetProperties(eigen_lib)
--if(NOT eigen_lib_POPULATED)
--	FetchContent_Populate(eigen_lib)
--	add_library(eigen INTERFACE)
--	target_include_directories(eigen INTERFACE ${eigen_lib_SOURCE_DIR})
-+ 	# FCPW
-+	FetchContent_Declare(fcpw_lib
-+		GIT_REPOSITORY https://github.com/rohan-sawhney/fcpw.git
-+		GIT_TAG dd65ec2
-+	)
-+	
-+	FetchContent_GetProperties(fcpw_lib)
-+	if(NOT fcpw_lib_POPULATED)
-+		FetchContent_Populate(fcpw_lib)
-+		add_subdirectory(${fcpw_lib_SOURCE_DIR} ${fcpw_lib_BINARY_DIR})
-+		target_include_directories(fcpw INTERFACE ${fcpw_lib_SOURCE_DIR})
-+	endif()
- endif()
- 
--# FCPW
--FetchContent_Declare(fcpw_lib
--	GIT_REPOSITORY https://github.com/rohan-sawhney/fcpw.git
--	GIT_TAG dd65ec2
--)
--
--FetchContent_GetProperties(fcpw_lib)
--if(NOT fcpw_lib_POPULATED)
--	FetchContent_Populate(fcpw_lib)
--	add_subdirectory(${fcpw_lib_SOURCE_DIR} ${fcpw_lib_BINARY_DIR})
--	target_include_directories(fcpw INTERFACE ${fcpw_lib_SOURCE_DIR})
--endif()
-+# eigen
-+if(SDFLIB_BUILD_DEBUG_APPS)
-+	FetchContent_Declare(eigen_lib
-+	GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
-+	GIT_TAG 46126273552afe13692929523d34006f54c19719 # 3.4
-+	)
-+	
-+	FetchContent_GetProperties(eigen_lib)
-+	if(NOT eigen_lib_POPULATED)
-+		FetchContent_Populate(eigen_lib)
-+		add_library(eigen INTERFACE)
-+		target_include_directories(eigen INTERFACE ${eigen_lib_SOURCE_DIR})
-+	endif()
-+ endif()
- 
- # stb
--add_library(stb_image INTERFACE)
--target_include_directories(stb_image INTERFACE stb)
-+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
-+	add_library(stb_image INTERFACE)
-+	target_include_directories(stb_image INTERFACE stb)
-+ endif()
- 
- # icg
- add_library(icg INTERFACE)
-diff --git a/src/sdf/OctreeSdf.cpp b/src/sdf/OctreeSdf.cpp
-index ef8ed4d..0e1eb97 100644
---- a/src/sdf/OctreeSdf.cpp
-+++ b/src/sdf/OctreeSdf.cpp
-@@ -6,7 +6,9 @@
- #include "SdfLib/InterpolationMethods.h"
- #include "sdf/OctreeSdfDepthFirst.h"
- #include "sdf/OctreeSdfBreadthFirst.h"
-+#ifdef OPENMP_AVAILABLE
- #include "sdf/OctreeSdfBreadthFirstNoDelay.h"
-+#endif
- #include <array>
- #include <stack>
- 
-@@ -46,8 +48,11 @@ OctreeSdf::OctreeSdf(const Mesh& mesh, BoundingBox box,
-             break;
-         case OctreeSdf::InitAlgorithm::CONTINUITY:
-             //initOctreeWithContinuity<PerNodeRegionTrianglesInfluence<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
--            // initOctreeWithContinuity<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
-+#ifdef OPENMP_AVAILABLE
-             initOctreeWithContinuityNoDelay<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule, numThreads);
-+#else
-+            initOctreeWithContinuity<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
-+#endif            
-             break;
-         // case OctreeSdf::InitAlgorithm::GPU_IMPLEMENTATION:
-         //     Timer time;
-@@ -78,7 +83,7 @@ float OctreeSdf::getDistance(glm::vec3 sample) const
-        startArrayPos.y < 0 || startArrayPos.y >= mStartGridSize ||
-        startArrayPos.z < 0 || startArrayPos.z >= mStartGridSize)
-     {
--        return mBox.getDistance(sample) + glm::sqrt(3.0f) * mBox.getSize().x;
-+        return mBox.getDistance(sample) + mMinBorderValue;
-     }
- 
-     const OctreeNode* currentNode = &mOctreeData[startArrayPos.z * mStartGridXY + startArrayPos.y * mStartGridSize + startArrayPos.x];
-@@ -108,7 +113,7 @@ float OctreeSdf::getDistance(glm::vec3 sample, glm::vec3& outGradient) const
-        startArrayPos.y < 0 || startArrayPos.y >= mStartGridSize ||
-        startArrayPos.z < 0 || startArrayPos.z >= mStartGridSize)
-     {
--        return mBox.getDistance(sample) + mMinBorderValue;
-+        return mBox.getDistance(sample, outGradient) + mMinBorderValue;
-     }
- 
-     const OctreeNode* currentNode = &mOctreeData[startArrayPos.z * mStartGridXY + startArrayPos.y * mStartGridSize + startArrayPos.x];
-@@ -253,4 +258,4 @@ void OctreeSdf::getDepthDensity(std::vector<float>& depthsDensity)
-         size *= 0.125f;
-     }
- }
--}
-\ No newline at end of file
-+}
-diff --git a/src/sdf/OctreeSdfDepthFirst.h b/src/sdf/OctreeSdfDepthFirst.h
-index 53ee4b2..196d191 100644
---- a/src/sdf/OctreeSdfDepthFirst.h
-+++ b/src/sdf/OctreeSdfDepthFirst.h
-@@ -7,7 +7,10 @@
- #include "SdfLib/OctreeSdfUtils.h"
- #include <array>
- #include <stack>
-+#ifdef OPENMP_AVAILABLE
- #include <omp.h>
-+#endif    
-+
- 
- namespace sdflib
- {
-@@ -381,7 +384,9 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
-     };
- 
-     const uint32_t voxlesPerAxis = 1 << startDepth;
-+#ifdef OPENMP_AVAILABLE
-     if(numThreads < 2)
-+#endif
-     {
-         // Create the grid
-         mOctreeData.resize(voxlesPerAxis * voxlesPerAxis * voxlesPerAxis);
-@@ -402,7 +407,8 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
- 
-         mValueRange = mainThread.valueRange;
-     }
--    else
-+#ifdef OPENMP_AVAILABLE
-+    else 
-     {
-         std::vector<ThreadContext> threadsContext(numThreads, mainThread);
- 
-@@ -511,6 +517,7 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
-             }
-         #endif
-     }
-+#endif
- 	
- #ifdef SDFLIB_PRINT_STATISTICS
-     SPDLOG_INFO("Used an octree of max depth {}", maxDepth);
-@@ -544,4 +551,4 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
- }
- }
- 
--#endif
-\ No newline at end of file
-+#endif
-diff --git a/src/utils/Mesh.cpp b/src/utils/Mesh.cpp
-index b407d38..6fff5fe 100644
---- a/src/utils/Mesh.cpp
-+++ b/src/utils/Mesh.cpp
-@@ -5,6 +5,7 @@
- 
- namespace sdflib
- {
-+#ifdef ASSIMP_AVAILABLE
- Mesh::Mesh(std::string filePath)
- {
-     Assimp::Importer import;
-@@ -28,6 +29,7 @@ Mesh::Mesh(const aiMesh* mesh)
- {
-     initMesh(mesh);
- }
-+#endif
- 
- Mesh::Mesh(glm::vec3* vertices, uint32_t numVertices,
-          uint32_t* indices, uint32_t numIndices)
-@@ -39,7 +41,7 @@ Mesh::Mesh(glm::vec3* vertices, uint32_t numVertices,
-     std::memcpy(mIndices.data(), indices, sizeof(uint32_t) * numIndices);
- }
- 
--
-+#ifdef ASSIMP_AVAILABLE
- void Mesh::initMesh(const aiMesh* mesh)
- {
-     if(!(mesh->mPrimitiveTypes & aiPrimitiveType_TRIANGLE))
-@@ -83,6 +85,7 @@ void Mesh::initMesh(const aiMesh* mesh)
- 		computeNormals();
- 	}
- }
-+#endif
- 
- void Mesh::computeBoundingBox()
- {
-@@ -134,4 +137,4 @@ void Mesh::applyTransform(glm::mat4 trans)
- 
-     computeBoundingBox();
- }
--}
-\ No newline at end of file
-+}

dist/mujoco.rc

@@ -1,6 +1,6 @@
 1 VERSIONINFO
-FILEVERSION 2,3,8,0
-PRODUCTVERSION 2,3,8,0
+FILEVERSION 3,1,2,0
+PRODUCTVERSION 3,1,2,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -9,9 +9,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "2.3.8"
+      VALUE "ProductVersion", "3.1.2"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "2.3.8"
+      VALUE "FileVersion", "3.1.2"
       VALUE "InternalName", "mujoco.dll"
       VALUE "OriginalFilename", "mujoco.dll"
       VALUE "CompanyName", "Google DeepMind"

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,1,2,0
+PRODUCTVERSION 3,1,2,0
 FILEOS 0x4
 FILETYPE 0x1
 {
@@ -11,9 +11,9 @@ FILETYPE 0x1
     BLOCK "040904b0"
     {
       VALUE "ProductName", "MuJoCo"
-      VALUE "ProductVersion", "2.3.8"
+      VALUE "ProductVersion", "3.1.2"
       VALUE "FileDescription", "MuJoCo"
-      VALUE "FileVersion", "2.3.8"
+      VALUE "FileVersion", "3.1.2"
       VALUE "InternalName", "simulate.exe"
       VALUE "OriginalFilename", "simulate.exe"
       VALUE "CompanyName", "Google DeepMind"

doc/APIreference/APIglobals.rst

@@ -489,7 +489,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - 1000
      - Maximum number of textures allowed.
        Defined in `mjrender.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjrender.h>`_.
-   * - ``mjMAXTHREADS``
+   * - ``mjMAXTHREAD``
      - 128
      - Maximum number OS threads that can be used in a thread pool.
        Defined in `mjthread.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjthread.h>`_.
@@ -498,7 +498,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
      - Maximum number of UI sections.
        Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
    * - ``mjMAXUIITEM``
-     - 80
+     - 200
      - Maximum number of items per UI section.
        Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
    * - ``mjMAXUITEXT``
@@ -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 <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
    * - ``mjVERSION_HEADER``
-     - 238
+     - 312
      - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software
        version, so 210 corresponds to version 2.1. Defined in  mujoco.h. The API function :ref:`mj_version` returns a
        number with the same meaning but for the compiled library.

doc/APIreference/APItypes.rst

@@ -34,6 +34,8 @@ MuJoCo defines a large number of types:
   - Structs used by :ref:`engine plugins<tyPluginStructure>`.
 
 - Several :ref:`tyFunction` for user-defined callbacks.
+- :ref:`tyNotes` regarding specific data structures that require detailed description.
+
 
 
 .. _tyPrimitive:
@@ -502,6 +504,18 @@ These are the possible framebuffers. They are used as an argument to the functio
 .. mujoco-include:: mjtFramebuffer
 
 
+.. _mjtDepthMap:
+
+mjtDepthMap
+~~~~~~~~~~~
+
+These are the depth mapping options. They are used as a value for the ``readPixelDepth`` attribute of the
+:ref:`mjrContext` struct, to control how the depth returned by :ref:`mjr_readPixels` is mapped from
+``znear`` to ``zfar``.
+
+.. mujoco-include:: mjtDepthMap
+
+
 .. _mjtFontScale:
 
 mjtFontScale
@@ -897,6 +911,7 @@ This structure contains the custom OpenGL rendering context, with the ids of all
 User Interface
 ^^^^^^^^^^^^^^
 
+For a high-level description of the UI framework, see :ref:`UI`.
 The names of these struct types are prefixed with ``mjui``, except for the main :ref:`mjUI` struct itself.
 
 
@@ -905,7 +920,9 @@ The names of these struct types are prefixed with ``mjui``, except for the main
 mjuiState
 ~~~~~~~~~
 
-This structure contains the keyboard and mouse state used by the UI framework.
+This C struct represents the global state of the window, keyboard and mouse, input event descriptors, and all window
+rectangles (including the visible UI rectangles). There is only one ``mjuiState`` per application, even if there are
+multiple UIs. This struct would normally be defined as a global variable.
 
 .. mujoco-include:: mjuiState
 
@@ -955,7 +972,9 @@ This structure defines one section of the UI.
 mjuiDef
 ~~~~~~~
 
-This structure defines one entry in the definition table used for simplified UI construction.
+This structure defines one entry in the definition table used for simplified UI construction. It contains everything
+needed to define one UI item. Some translation is performed by the helper functions, so that multiple mjuiDefs can be
+defined as a static table.
 
 .. mujoco-include:: mjuiDef
 
@@ -965,7 +984,12 @@ This structure defines one entry in the definition table used for simplified UI
 mjUI
 ~~~~
 
-This structure defines the entire UI.
+This C struct represents an entire UI. The same application could have multiple UIs, for example on the left and the
+right of the window. This would normally be defined as a global variable. As explained earlier, it contains static
+allocation for a maximum number of supported UI sections (:ref:`mjuiSection<mjuiSection>`) each with a maximum number
+of supported items (:ref:`mjuiItem<mjuiItem>`). It also contains the color and spacing themes, enable/disable
+callback, virtual window descriptor, text edit state, mouse focus. Some of these fields are set only once when the UI
+is initialized, others change at runtime.
 
 .. mujoco-include:: mjUI
 
@@ -1195,9 +1219,9 @@ a frame at the center-of-mass of the local kinematic subtree (``mjData.subtree_c
 This choice increases the precision of kinematic computations for mechanisms that are distant from the global origin.
 
 ``cdof``:
-  These 6D motion vectors describe the instantaneous axis of a degree-of-freedom and are used by all Jacobian functions.
-  Therefore, the minimal computation required for analytic Jacobians is :ref:`mj_kinematics` followed by
-  :ref:`mj_comPos`.
+  These 6D motion vectors (3 rotation, 3 translation) describe the instantaneous axis of a degree-of-freedom and are
+  used by all Jacobian functions. The minimal computation required for analytic Jacobians is :ref:`mj_kinematics`
+  followed by :ref:`mj_comPos`.
 
 ``cinert``:
   These 10-vectors describe the inertial properties of a body in the c-frame and are used by the Composite Rigid Body

doc/APIreference/functions.rst

@@ -392,7 +392,9 @@ mj_objectAcceleration
 
 .. mujoco-include:: mj_objectAcceleration
 
-Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation.
+Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation. If acceleration or force
+sensors are not present in the model, :ref:`mj_rnePostConstraint` must be manually called in order to calculate
+:ref:`mjData`.cacc -- the total body acceleration, including contributions from the constraint solver.
 
 .. _mj_contactForce:
 
@@ -581,6 +583,17 @@ mj_RungeKutta
 
 Runge-Kutta explicit order-N integrator.
 
+.. _mj_implicit:
+
+mj_implicit
+~~~~~~~~~~~
+
+.. mujoco-include:: mj_implicit
+
+Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see
+:ref:`Numerical Integration<geIntegration>`), and advances simulation time. See `mjdata.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__ for fields computed by this function.
+
 .. _mj_invPosition:
 
 mj_invPosition
@@ -724,6 +737,15 @@ mj_camlight
 
 Compute camera and light positions and orientations.
 
+.. _mj_flex:
+
+mj_flex
+~~~~~~~
+
+.. mujoco-include:: mj_flex
+
+Compute flex-related quantities.
+
 .. _mj_tendon:
 
 mj_tendon
@@ -794,7 +816,7 @@ mj_passive
 
 .. mujoco-include:: mj_passive
 
-Compute qfrc_passive from spring-dampers, viscosity and density.
+Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.
 
 .. _mj_subtreeVel:
 
@@ -803,7 +825,10 @@ mj_subtreeVel
 
 .. mujoco-include:: mj_subtreeVel
 
-subtree linear velocity and angular momentum
+Sub-tree linear velocity and angular momentum: compute ``subtree_linvel``, ``subtree_angmom``.
+This function is triggered automatically if the subtree :ref:`velocity<sensor-subtreelinvel>` or
+:ref:`momentum<sensor-subtreeangmom>` sensors are present in the model.
+It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "vel".
 
 .. _mj_rne:
 
@@ -812,7 +837,8 @@ mj_rne
 
 .. mujoco-include:: mj_rne
 
-RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term.
+Recursive Newton Euler: compute :math:`M(q) \ddot q + C(q,\dot q)`. ``flg_acc=0`` removes the inertial term (i.e.
+assumes :math:`\ddot q = 0`).
 
 .. _mj_rnePostConstraint:
 
@@ -821,7 +847,21 @@ mj_rnePostConstraint
 
 .. mujoco-include:: mj_rnePostConstraint
 
-RNE with complete data: compute cacc, cfrc_ext, cfrc_int.
+Recursive Newton Euler with final computed forces and accelerations.
+Computes three body-level ``nv x 6`` arrays, all defined in the subtreecom-based
+:ref:`c-frame<tyNotesCom>` and arranged in ``[rotation(3), translation(3)]`` order.
+
+- ``cacc``: Body acceleration, required for :ref:`mj_objectAcceleration`.
+- ``cfrc_int``: Interaction force with the parent body.
+- ``cfrc_ext``: External force acting on the body.
+
+This function is triggered automatically if the following sensors are present in the model:
+:ref:`accelerometer<sensor-accelerometer>`, :ref:`force<sensor-force>`, :ref:`torque<sensor-torque>`,
+:ref:`framelinacc<sensor-framelinacc>`, :ref:`frameangacc<sensor-frameangacc>`.
+It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "acc".
+
+The computed force arrays ``cfrc_int`` and ``cfrc_ext`` currently suffer from a know bug, they do not take into account
+the effect of spatial tendons, see :github:issue:`832`.
 
 .. _mj_collision:
 
@@ -875,8 +915,8 @@ mj_constraintUpdate
 
 .. mujoco-include:: mj_constraintUpdate
 
-Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set \*cost = s(jar)
-where jar = Jac*qacc-aref.
+Compute ``efc_state``, ``efc_force``, ``qfrc_constraint``, and (optionally) cone Hessians.
+If ``cost`` is not ``NULL``, set ``*cost = s(jar)`` where ``jar = Jac*qacc - aref``.
 
 .. _Raycollisions:
 
@@ -948,6 +988,16 @@ mju_rayGeom
 
 Intersect ray with pure geom, return nearest distance or -1 if no intersection.
 
+.. _mju_rayFlex:
+
+mju_rayFlex
+~~~~~~~~~~~
+
+.. mujoco-include:: mju_rayFlex
+
+Intersect ray with flex, return nearest distance or -1 if no intersection,
+and also output nearest vertex id.
+
 .. _mju_raySkin:
 
 mju_raySkin
@@ -1926,6 +1976,8 @@ Find first rectangle containing mouse, -1: not found.
 UI framework
 ^^^^^^^^^^^^
 
+For a high-level description of the UI framework, see :ref:`UI`.
+
 .. _mjui_themeSpacing:
 
 mjui_themeSpacing
@@ -1951,7 +2003,12 @@ mjui_add
 
 .. mujoco-include:: mjui_add
 
-Add definitions to UI.
+This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs,
+each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added
+after the end of the last used section. There is also another version of this function
+(:ref:`mjui_addToSection<mjui_addToSection>`) which adds items to a specified section instead of adding them at the end
+of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by
+:ref:`mjMAXUISECT<glNumeric>` and :ref:`mjMAXUIITEM<glNumeric>`. Exceeding these maxima results in low-level errors.
 
 .. _mjui_addToSection:
 
@@ -1978,7 +2035,12 @@ mjui_update
 
 .. mujoco-include:: mjui_update
 
-Update specific section/item; -1: update all.
+This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers)
+changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user
+(``UiModify`` in :ref:`simulate.cc <saSimulate>`) which also recomputes the layout of all rectangles and associated
+auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user
+specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated
+(which is needed following major changes.)
 
 .. _mjui_event:
 
@@ -1987,7 +2049,10 @@ mjui_event
 
 .. mujoco-include:: mjui_event
 
-Handle UI event, return pointer to changed item, NULL if no change.
+This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item
+that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler
+implemented by the user (``UiEvent`` in :ref:`simulate.cc <saSimulate>`), and then some action is taken by user code
+depending on which UI item was modified and what the state of that item is after the event is handled.
 
 .. _mjui_render:
 
@@ -1996,7 +2061,9 @@ mjui_render
 
 .. mujoco-include:: mjui_render
 
-Copy UI image to current buffer.
+This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If
+there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the
+time, while :ref:`mjui_update` is called only when changes in the UI take place.
 
 .. _Errorandmemory:
 
@@ -2918,7 +2985,7 @@ mju_eig3
 
 .. mujoco-include:: mju_eig3
 
-Eigenvalue decomposition of symmetric 3x3 matrix.
+Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.
 
 .. _mju_boxQP:
 
@@ -2955,10 +3022,10 @@ outputs (optional):
 
 notes:
   The initial value of ``res`` is used to warmstart the solver.
-  ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output.
-  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output.
+  ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output.
+  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output.
   The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures.
-  Only the lower triangles of H and R and are read from and written to, respectively.
+  Only the lower triangles of H and R are read from and written to, respectively.
 
 .. _mju_boxQPmalloc:
 
@@ -3256,6 +3323,9 @@ These matrices and their dimensions are:
 - All outputs are optional (can be NULL).
 - ``eps`` is the finite-differencing epsilon.
 - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences.
+- Accuracy can be somewhat improved if solver :ref:`iterations<option-iterations>` are set to a
+  fixed (small) value and solver :ref:`tolerance<option-tolerance>` is set to 0. This insures that
+  all calls to the solver will perform exactly the same number of iterations.
 
 .. _mjd_inverseFD:
 

doc/APIreference/functions_override.rst

@@ -2,10 +2,6 @@
   This file contains each section text along with function doc overrides.  By default the docs use the function doc
   pulled from the header files.
 
-.. _Activation:
-
-The functions in this section are maintained for backward compatibility with the now-removed activation mechanism.
-
 .. _Virtualfilesystem:
 
 Virtual file system (VFS) enables the user to load all necessary files in memory, including MJB binary model files, XML
@@ -73,6 +69,12 @@ These functions can be used to print various quantities to the screen for debugg
 These are components of the simulation pipeline, called internally from :ref:`mj_step`, :ref:`mj_forward` and
 :ref:`mj_inverse`. It is unlikely that the user will need to call them.
 
+.. _mj_implicit:
+
+Integrates the simulation state using an implicit-in-velocity integrator (either "implicit" or "implicitfast", see
+:ref:`Numerical Integration<geIntegration>`), and advances simulation time. See `mjdata.h
+<https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjdata.h>`__ for fields computed by this function.
+
 .. _Subcomponents:
 
 These are sub-components of the simulation pipeline, called internally from the components above. It is very unlikely
@@ -90,10 +92,40 @@ Solve linear system :math:`M x = y` using factorization: :math:`x = (L^T D L)^{-
 
 Half of linear solve: :math:`x = \sqrt{D^{-1}} (L^T)^{-1} y`
 
+.. _mj_subtreeVel:
+
+Sub-tree linear velocity and angular momentum: compute ``subtree_linvel``, ``subtree_angmom``.
+This function is triggered automatically if the subtree :ref:`velocity<sensor-subtreelinvel>` or
+:ref:`momentum<sensor-subtreeangmom>` sensors are present in the model.
+It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "vel".
+
+.. _mj_rne:
+
+Recursive Newton Euler: compute :math:`M(q) \ddot q + C(q,\dot q)`. ``flg_acc=0`` removes the inertial term (i.e.
+assumes :math:`\ddot q = 0`).
+
+.. _mj_rnePostConstraint:
+
+Recursive Newton Euler with final computed forces and accelerations.
+Computes three body-level ``nv x 6`` arrays, all defined in the subtreecom-based
+:ref:`c-frame<tyNotesCom>` and arranged in ``[rotation(3), translation(3)]`` order.
+
+- ``cacc``: Body acceleration, required for :ref:`mj_objectAcceleration`.
+- ``cfrc_int``: Interaction force with the parent body.
+- ``cfrc_ext``: External force acting on the body.
+
+This function is triggered automatically if the following sensors are present in the model:
+:ref:`accelerometer<sensor-accelerometer>`, :ref:`force<sensor-force>`, :ref:`torque<sensor-torque>`,
+:ref:`framelinacc<sensor-framelinacc>`, :ref:`frameangacc<sensor-frameangacc>`.
+It is also triggered for :ref:`user sensors<sensor-user>` of :ref:`stage<sensor-user-needstage>` "acc".
+
+The computed force arrays ``cfrc_int`` and ``cfrc_ext`` currently suffer from a know bug, they do not take into account
+the effect of spatial tendons, see :github:issue:`832`.
+
 .. _mj_constraintUpdate:
 
-Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians. If cost is not NULL, set \*cost = s(jar)
-where jar = Jac*qacc-aref.
+Compute ``efc_state``, ``efc_force``, ``qfrc_constraint``, and (optionally) cone Hessians.
+If ``cost`` is not ``NULL``, set ``*cost = s(jar)`` where ``jar = Jac*qacc - aref``.
 
 .. _Support:
 
@@ -154,6 +186,12 @@ This function can be used to apply a Cartesian force and torque to a point on a
 mjData.qfrc_applied of all applied forces. Note that the function requires a pointer to this vector, because sometimes
 we want to add the result to a different vector.
 
+.. _mj_objectAcceleration:
+
+Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation. If acceleration or force
+sensors are not present in the model, :ref:`mj_rnePostConstraint` must be manually called in order to calculate
+mjData.cacc -- the total body acceleration, including contributions from the constraint solver.
+
 .. _mj_differentiatePos:
 
 This function subtracts two vectors in the format of qpos (and divides the result by dt), while respecting the
@@ -218,6 +256,43 @@ of how to use these functions.
 
 .. _UIframework:
 
+For a high-level description of the UI framework, see :ref:`UI`.
+
+.. _mjui_add:
+
+This is the helper function used to construct a UI. The second argument points to an array of :ref:`mjuiDef` structs,
+each corresponding to one item. The last (unused) item has its type set to -1, to mark termination. The items are added
+after the end of the last used section. There is also another version of this function
+(:ref:`mjui_addToSection<mjui_addToSection>`) which adds items to a specified section instead of adding them at the end
+of the UI. Keep in mind that there is a maximum preallocated number of sections and items per section, given by
+:ref:`mjMAXUISECT<glNumeric>` and :ref:`mjMAXUIITEM<glNumeric>`. Exceeding these maxima results in low-level errors.
+
+.. _mjui_update:
+
+This is the main UI update function. It needs to be called whenever the user data (pointed to by the item data pointers)
+changes, or when the UI state itself changes. It is normally called by a higher-level function implemented by the user
+(``UiModify`` in :ref:`simulate.cc <saSimulate>`) which also recomputes the layout of all rectangles and associated
+auxiliary buffers. The function updates the pixels in the offscreen OpenGL buffer. To perform minimal updates, the user
+specifies the section and the item that was modified. A value of -1 means all items and/or sections need to be updated
+(which is needed following major changes.)
+
+.. _mjui_event:
+
+This function is the low-level event handler. It makes the necessary changes in the UI and returns a pointer to the item
+that received the event (or ``NULL`` if no valid event was recorded). This is normally called within the event handler
+implemented by the user (``UiEvent`` in :ref:`simulate.cc <saSimulate>`), and then some action is taken by user code
+depending on which UI item was modified and what the state of that item is after the event is handled.
+
+
+.. _mjui_render:
+
+This function is called in the screen refresh loop. It copies the offscreen OpenGL buffer to the window framebuffer. If
+there are multiple UIs in the application, it should be called once for each UI. Thus ``mjui_render`` is called all the
+time, while :ref:`mjui_update` is called only when changes in the UI take place.
+
+
+
+
 .. _Errorandmemory:
 
 .. _Standardmath:
@@ -417,10 +492,10 @@ outputs (optional):
 
 notes:
   The initial value of ``res`` is used to warmstart the solver.
-  ``R`` must have allocatd size ``n*(n+7)``, but only ``nfree*nfree`` values are used in output.
-  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used in output.
+  ``R`` must have allocated size ``n*(n+7)``, but only ``nfree*nfree`` values are used as output.
+  ``index`` (if given) must have allocated size ``n``, but only ``nfree`` values are used as output.
   The convenience function :ref:`mju_boxQPmalloc` allocates the required data structures.
-  Only the lower triangles of H and R and are read from and written to, respectively.
+  Only the lower triangles of H and R are read from and written to, respectively.
 
 .. _mju_boxQPmalloc:
 
@@ -463,6 +538,9 @@ These matrices and their dimensions are:
 - All outputs are optional (can be NULL).
 - ``eps`` is the finite-differencing epsilon.
 - ``flg_centered`` denotes whether to use forward (0) or centered (1) differences.
+- Accuracy can be somewhat improved if solver :ref:`iterations<option-iterations>` are set to a
+  fixed (small) value and solver :ref:`tolerance<option-tolerance>` is set to 0. This insures that
+  all calls to the solver will perform exactly the same number of iterations.
 
 .. _mjd_inverseFD:
 

doc/XMLschema.rst

@@ -65,7 +65,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`linewidth<visual-global-linewidth>`                       | :ref:`glow<visual-global-glow>`                                 | :ref:`offwidth<visual-global-offwidth>`                         | :ref:`offheight<visual-global-offheight>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`realtime<visual-global-realtime>`                         | :ref:`ellipsoidinertia<visual-global-ellipsoidinertia>`         |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`realtime<visual-global-realtime>`                         | :ref:`ellipsoidinertia<visual-global-ellipsoidinertia>`         | :ref:`bvactive<visual-global-bvactive>`                         |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| visual |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -109,6 +109,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`framelength<visual-scale-framelength>`                    | :ref:`framewidth<visual-scale-framewidth>`                      | :ref:`constraint<visual-scale-constraint>`                      | :ref:`slidercrank<visual-scale-slidercrank>`                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`frustum<visual-scale-frustum>`                            |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| visual |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`rgba                         | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -124,7 +126,9 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`contacttorque<visual-rgba-contacttorque>`                 | :ref:`contactgap<visual-rgba-contactgap>`                       | :ref:`rangefinder<visual-rgba-rangefinder>`                     | :ref:`constraint<visual-rgba-constraint>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`slidercrank<visual-rgba-slidercrank>`                     | :ref:`crankbroken<visual-rgba-crankbroken>`                     |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`slidercrank<visual-rgba-slidercrank>`                     | :ref:`crankbroken<visual-rgba-crankbroken>`                     | :ref:`frustum<visual-rgba-frustum>`                             | :ref:`bv<visual-rgba-bv>`                                       | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`bvactive<visual-rgba-bvactive>`                           |                                                                 |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | mujoco |br| |L|                    |    | *no attributes*                                                                                                                                                                                                                                                              |
@@ -153,7 +157,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<asset-hfield-name>`                                  | :ref:`content_type<asset-hfield-content_type>`                  | :ref:`file<asset-hfield-file>`                                  | :ref:`nrow<asset-hfield-nrow>`                                  | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`ncol<asset-hfield-ncol>`                                  | :ref:`size<asset-hfield-size>`                                  |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`ncol<asset-hfield-ncol>`                                  | :ref:`size<asset-hfield-size>`                                  | :ref:`elevation<asset-hfield-elevation>`                        |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| asset |br| |_| |L|             |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -222,13 +226,13 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`wind<option-wind>`                                        | :ref:`magnetic<option-magnetic>`                                | :ref:`density<option-density>`                                  | :ref:`viscosity<option-viscosity>`                              | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`o_margin<option-o_margin>`                                | :ref:`o_solref<option-o_solref>`                                | :ref:`o_solimp<option-o_solimp>`                                | :ref:`integrator<option-integrator>`                            | |
+|                                    |    |    | :ref:`o_margin<option-o_margin>`                                | :ref:`o_solref<option-o_solref>`                                | :ref:`o_solimp<option-o_solimp>`                                | :ref:`o_friction<option-o_friction>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`cone<option-cone>`                                        | :ref:`jacobian<option-jacobian>`                                | :ref:`solver<option-solver>`                                    | :ref:`iterations<option-iterations>`                            | |
+|                                    |    |    | :ref:`integrator<option-integrator>`                            | :ref:`cone<option-cone>`                                        | :ref:`jacobian<option-jacobian>`                                | :ref:`solver<option-solver>`                                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`ls_iterations<option-ls_iterations>`                      | :ref:`noslip_iterations<option-noslip_iterations>`              | :ref:`mpr_iterations<option-mpr_iterations>`                    | :ref:`sdf_iterations<option-sdf_iterations>`                    | |
+|                                    |    |    | :ref:`iterations<option-iterations>`                            | :ref:`ls_iterations<option-ls_iterations>`                      | :ref:`noslip_iterations<option-noslip_iterations>`              | :ref:`mpr_iterations<option-mpr_iterations>`                    | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`sdf_initpoints<option-sdf_initpoints>`                    |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`sdf_iterations<option-sdf_iterations>`                    | :ref:`sdf_initpoints<option-sdf_initpoints>`                    | :ref:`actuatorgroupdisable<option-actuatorgroupdisable>`        |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| option |br| |_| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -481,6 +485,120 @@
 |                                    |    |    | :ref:`key<plugin-config-key>`                                   | :ref:`value<plugin-config-value>`                               |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| body |br| |_| |L|              |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`flexcomp                     | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <body-flexcomp>`                   |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<body-flexcomp-name>`                                 | :ref:`type<body-flexcomp-type>`                                 | :ref:`group<body-flexcomp-group>`                               | :ref:`dim<body-flexcomp-dim>`                                   | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`count<body-flexcomp-count>`                               | :ref:`spacing<body-flexcomp-spacing>`                           | :ref:`radius<body-flexcomp-radius>`                             | :ref:`rigid<body-flexcomp-rigid>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`mass<body-flexcomp-mass>`                                 | :ref:`inertiabox<body-flexcomp-inertiabox>`                     | :ref:`scale<body-flexcomp-scale>`                               | :ref:`file<body-flexcomp-file>`                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`point<body-flexcomp-point>`                               | :ref:`element<body-flexcomp-element>`                           | :ref:`texcoord<body-flexcomp-texcoord>`                         | :ref:`material<body-flexcomp-material>`                         | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`rgba<body-flexcomp-rgba>`                                 | :ref:`flatskin<body-flexcomp-flatskin>`                         | :ref:`pos<body-flexcomp-pos>`                                   | :ref:`quat<body-flexcomp-quat>`                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`axisangle<body-flexcomp-axisangle>`                       | :ref:`xyaxes<body-flexcomp-xyaxes>`                             | :ref:`zaxis<body-flexcomp-zaxis>`                               | :ref:`euler<body-flexcomp-euler>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`edge                         | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flexcomp-edge>`                   |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`equality<flexcomp-edge-equality>`                         | :ref:`solref<flexcomp-edge-solref>`                             | :ref:`solimp<flexcomp-edge-solimp>`                             | :ref:`stiffness<flexcomp-edge-stiffness>`                       | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`damping<flexcomp-edge-damping>`                           |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`contact                      | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flexcomp-contact>`                |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`contype<flexcomp-contact-contype>`                        | :ref:`conaffinity<flexcomp-contact-conaffinity>`                | :ref:`condim<flexcomp-contact-condim>`                          | :ref:`priority<flexcomp-contact-priority>`                      | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`friction<flexcomp-contact-friction>`                      | :ref:`solmix<flexcomp-contact-solmix>`                          | :ref:`solref<flexcomp-contact-solref>`                          | :ref:`solimp<flexcomp-contact-solimp>`                          | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`margin<flexcomp-contact-margin>`                          | :ref:`gap<flexcomp-contact-gap>`                                | :ref:`internal<flexcomp-contact-internal>`                      | :ref:`selfcollide<flexcomp-contact-selfcollide>`                | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`activelayers<flexcomp-contact-activelayers>`              |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`pin                          | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flexcomp-pin>`                    |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`id<flexcomp-pin-id>`                                      | :ref:`range<flexcomp-pin-range>`                                | :ref:`grid<flexcomp-pin-grid>`                                  | :ref:`gridrange<flexcomp-pin-gridrange>`                        | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flexcomp |br| |_2| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`plugin                       | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flexcomp-plugin>`                 |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`plugin<flexcomp-plugin-plugin>`                           | :ref:`instance<flexcomp-plugin-instance>`                       |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_3| plugin |br| |_3| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`config                       | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <plugin-config>`                   |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`key<plugin-config-key>`                                   | :ref:`value<plugin-config-value>`                               |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| mujoco |br| |L|                    |    | *no attributes*                                                                                                                                                                                                                                                              |
+| :ref:`deformable<deformable>`      |    |                                                                                                                                                                                                                                                                              |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| deformable |br| |_| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`flex                         | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <deformable-flex>`                 |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<deformable-flex-name>`                               | :ref:`group<deformable-flex-group>`                             | :ref:`dim<deformable-flex-dim>`                                 | :ref:`radius<deformable-flex-radius>`                           | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`material<deformable-flex-material>`                       | :ref:`rgba<deformable-flex-rgba>`                               | :ref:`flatskin<deformable-flex-flatskin>`                       | :ref:`body<deformable-flex-body>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`vertex<deformable-flex-vertex>`                           | :ref:`element<deformable-flex-element>`                         | :ref:`texcoord<deformable-flex-texcoord>`                       |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flex |br| |_2| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`contact                      | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flex-contact>`                    |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`contype<flex-contact-contype>`                            | :ref:`conaffinity<flex-contact-conaffinity>`                    | :ref:`condim<flex-contact-condim>`                              | :ref:`priority<flex-contact-priority>`                          | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`friction<flex-contact-friction>`                          | :ref:`solmix<flex-contact-solmix>`                              | :ref:`solref<flex-contact-solref>`                              | :ref:`solimp<flex-contact-solimp>`                              | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`margin<flex-contact-margin>`                              | :ref:`gap<flex-contact-gap>`                                    | :ref:`internal<flex-contact-internal>`                          | :ref:`selfcollide<flex-contact-selfcollide>`                    | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`activelayers<flex-contact-activelayers>`                  |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| flex |br| |_2| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`edge                         | ?  |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <flex-edge>`                       |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`stiffness<flex-edge-stiffness>`                           | :ref:`damping<flex-edge-damping>`                               |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| deformable |br| |_| |L|        |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`skin                         | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <deformable-skin>`                 |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<deformable-skin-name>`                               | :ref:`file<deformable-skin-file>`                               | :ref:`material<deformable-skin-material>`                       | :ref:`rgba<deformable-skin-rgba>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`inflate<deformable-skin-inflate>`                         | :ref:`vertex<deformable-skin-vertex>`                           | :ref:`texcoord<deformable-skin-texcoord>`                       | :ref:`face<deformable-skin-face>`                               | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`group<deformable-skin-group>`                             |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_2| skin |br| |_2| |L|            |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`bone                         | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <skin-bone>`                       |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`body<skin-bone-body>`                                     | :ref:`bindpos<skin-bone-bindpos>`                               | :ref:`bindquat<skin-bone-bindquat>`                             | :ref:`vertid<skin-bone-vertid>`                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`vertweight<skin-bone-vertweight>`                         |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | mujoco |br| |L|                    |    | *no attributes*                                                                                                                                                                                                                                                              |
 | :ref:`contact<contact>`            |    |                                                                                                                                                                                                                                                                              |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -543,6 +661,15 @@
 |                                    |    |    | :ref:`polycoef<equality-tendon-polycoef>`                       | :ref:`active<equality-tendon-active>`                           | :ref:`solref<equality-tendon-solref>`                           | :ref:`solimp<equality-tendon-solimp>`                           | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
+| |_| equality |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
+| :ref:`flex                         | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
+| <equality-flex>`                   |    |                                                                                                                                                                                                                                                                              |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`name<equality-flex-name>`                                 | :ref:`class<equality-flex-class>`                               | :ref:`flex<equality-flex-flex>`                                 | :ref:`active<equality-flex-active>`                             | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`solref<equality-flex-solref>`                             | :ref:`solimp<equality-flex-solimp>`                             |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | mujoco |br| |L|                    |    | *no attributes*                                                                                                                                                                                                                                                              |
 | :ref:`tendon<tendon>`              |    |                                                                                                                                                                                                                                                                              |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -653,7 +780,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`jointinparent<actuator-position-jointinparent>`           | :ref:`tendon<actuator-position-tendon>`                         | :ref:`slidersite<actuator-position-slidersite>`                 | :ref:`cranksite<actuator-position-cranksite>`                   | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`site<actuator-position-site>`                             | :ref:`refsite<actuator-position-refsite>`                       | :ref:`kp<actuator-position-kp>`                                 |                                                                 | |
+|                                    |    |    | :ref:`site<actuator-position-site>`                             | :ref:`refsite<actuator-position-refsite>`                       | :ref:`kp<actuator-position-kp>`                                 | :ref:`kv<actuator-position-kv>`                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| actuator |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -685,6 +812,8 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`cranksite<actuator-intvelocity-cranksite>`                | :ref:`site<actuator-intvelocity-site>`                          | :ref:`refsite<actuator-intvelocity-refsite>`                    | :ref:`kp<actuator-intvelocity-kp>`                              | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`kv<actuator-intvelocity-kv>`                              |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| actuator |br| |_| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
 | :ref:`damper                       | \* |    :class: mjcf-attributes                                                                                                                                                                                                                                                   |
@@ -754,13 +883,15 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`name<actuator-plugin-name>`                               | :ref:`class<actuator-plugin-class>`                             | :ref:`plugin<actuator-plugin-plugin>`                           | :ref:`instance<actuator-plugin-instance>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`group<actuator-plugin-group>`                             | :ref:`ctrllimited<actuator-plugin-ctrllimited>`                 | :ref:`forcelimited<actuator-plugin-forcelimited>`               | :ref:`ctrlrange<actuator-plugin-ctrlrange>`                     | |
+|                                    |    |    | :ref:`group<actuator-plugin-group>`                             | :ref:`ctrllimited<actuator-plugin-ctrllimited>`                 | :ref:`forcelimited<actuator-plugin-forcelimited>`               | :ref:`actlimited<actuator-plugin-actlimited>`                   | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`forcerange<actuator-plugin-forcerange>`                   | :ref:`lengthrange<actuator-plugin-lengthrange>`                 | :ref:`gear<actuator-plugin-gear>`                               | :ref:`cranklength<actuator-plugin-cranklength>`                 | |
+|                                    |    |    | :ref:`ctrlrange<actuator-plugin-ctrlrange>`                     | :ref:`forcerange<actuator-plugin-forcerange>`                   | :ref:`actrange<actuator-plugin-actrange>`                       | :ref:`lengthrange<actuator-plugin-lengthrange>`                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`joint<actuator-plugin-joint>`                             | :ref:`jointinparent<actuator-plugin-jointinparent>`             | :ref:`site<actuator-plugin-site>`                               | :ref:`tendon<actuator-plugin-tendon>`                           | |
+|                                    |    |    | :ref:`gear<actuator-plugin-gear>`                               | :ref:`cranklength<actuator-plugin-cranklength>`                 | :ref:`joint<actuator-plugin-joint>`                             | :ref:`jointinparent<actuator-plugin-jointinparent>`             | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               |                                                                 | |
+|                                    |    |    | :ref:`site<actuator-plugin-site>`                               | :ref:`dyntype<actuator-plugin-dyntype>`                         | :ref:`dynprm<actuator-plugin-dynprm>`                           | :ref:`tendon<actuator-plugin-tendon>`                           | |
+|                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+|                                    |    |    | :ref:`cranksite<actuator-plugin-cranksite>`                     | :ref:`slidersite<actuator-plugin-slidersite>`                   | :ref:`user<actuator-plugin-user>`                               | :ref:`actearly<actuator-plugin-actearly>`                       | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_2| plugin |br| |_2| |L|          |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -1313,7 +1444,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`gear<default-position-gear>`                              | :ref:`cranklength<default-position-cranklength>`                | :ref:`user<default-position-user>`                              | :ref:`group<default-position-group>`                            | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`kp<default-position-kp>`                                  |                                                                 |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`kp<default-position-kp>`                                  | :ref:`kv<default-position-kv>`                                  |                                                                 |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| default |br| |_| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |
@@ -1335,7 +1466,7 @@
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 |                                    |    |    | :ref:`actrange<default-intvelocity-actrange>`                   | :ref:`gear<default-intvelocity-gear>`                           | :ref:`cranklength<default-intvelocity-cranklength>`             | :ref:`user<default-intvelocity-user>`                           | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
-|                                    |    |    | :ref:`group<default-intvelocity-group>`                         | :ref:`kp<default-intvelocity-kp>`                               |                                                                 |                                                                 | |
+|                                    |    |    | :ref:`group<default-intvelocity-group>`                         | :ref:`kp<default-intvelocity-kp>`                               | :ref:`kv<default-intvelocity-kv>`                               |                                                                 | |
 |                                    |    |    +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
 +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 | |_| default |br| |_| |L|           |    | .. table::                                                                                                                                                                                                                                                                   |

doc/_static/gyroscopic.xml

@@ -0,0 +1,22 @@
+<mujoco>
+  <asset>
+    <texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
+    <texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
+    <material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
+  </asset>
+
+  <statistic extent="1"/>
+
+  <!-- switch the integrator to "implicit" for better gyroscopic stabillity -->
+  <option timestep="0.01" integrator="implicitfast"/>
+
+  <worldbody>
+    <geom pos="0 0 -.5" type="plane" size="100 100 .1" zaxis=".5 0 1" material="grid"/>
+    <body pos="0 0 .3">
+      <freejoint/>
+      <geom type="ellipsoid" size="0.15 0.1 0.07" euler="20 40 60"/>
+      <light pos="0 0 1.5" mode="track"/>
+      <camera pos="0 -1.498 0.1" xyaxes="1 0 0 0 0 1" mode="track"/>
+    </body>
+  </worldbody>
+</mujoco>

doc/_static/pendulum.xml

@@ -0,0 +1,27 @@
+<mujoco model="energy conserving pendulum">
+  <option integrator="RK4">
+    <flag constraint="disable" energy="enable"/>
+  </option>
+
+  <worldbody>
+    <light pos="0 0 2"/>
+    <geom pos="0 0 -.5" type="plane" size="1 1 .01"/>
+    <body pos="0 0 0">
+      <joint type="hinge" axis="0 1 0"/>
+      <geom type="cylinder" size="0.02" fromto="0 -.02 0 0 .02 0"/>
+      <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+      <body pos="0.1 0 0">
+        <joint type="slide" axis="1 0 0" stiffness="200"/>
+        <geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
+        <body pos=".1 0 0">
+          <joint type="ball"/>
+          <geom type="box" size=".02" fromto="0 0 0 0 .1 0"/>
+          <body pos="0 .1 0">
+            <joint axis="1 0 0"/>
+            <geom type="capsule" size="0.02" fromto="0 0 0 0 .1 0"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>

doc/changelog.rst

@@ -2,21 +2,210 @@
 Changelog
 =========
 
-Upcoming version (not yet released)
+Version 3.1.2 (February 05, 2024)
 -----------------------------------
 
+General
+^^^^^^^
+1. Improved the :ref:`discardvisual<compiler-discardvisual>` compiler flag, which now discards all visual-only assets.
+   See :ref:`discardvisual<compiler-discardvisual>` for details.
+2. Removed the :ref:`timer<mjtTimer>` for midphase colllision detection, it is now folded in with the narrowphase
+   timer. This is because timing the two phases seperately required fine-grained timers inside the collision
+   functions; these functions are so small and fast that the timer itself was incurring a measurable cost.
+3. Added the flag :ref:`bvactive<visual-global-bvactive>` to ``visual/global``, allowing users to turn off
+   visualisation of active bounding volumes (the red/green boxes in this :ref:`this changelog item<midphase>`). For
+   models with very high-resolution meshes, the computation required for this visualization can slow down simulation
+   speed. Fixes :github:issue:`1279`.
+
+   - Added color of :ref:`bounding volumes<visual-rgba-bv>` and :ref:`active bounding volumes<visual-rgba-bvactive>`
+     to :ref:`visual/rgba<visual-rgba>`.
+4. Height-field elevation data can now be specified directly in XML with the :ref:`elevation<asset-hfield-elevation>`
+   attribute (and not only with PNG files).
+   See `example model <https://github.com/google-deepmind/mujoco/blob/main/test/user/testdata/hfield_xml.xml>`__.
+
+MJX
+^^^
+5. Added :ref:`dyntype<actuator-general-dyntype>` ``filterexact``.
+6. Added :at:`site` transmission.
+7. Updated MJX colab tutorial with more stable quadruped environment.
+8. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes.
+9. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`.
+10. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`.
+11. Raise a not implemented error when nonzero frictionloss is present. Fixes :github:issue:`1344`.
+
+Python bindings
+^^^^^^^^^^^^^^^
+12. Improved the implementation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
+    code will require modification.
+
+    - Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
+    - Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
+    - Outputs are no longer squeezed and always have dim=3.
+13. The ``sync`` function for the :ref:`passive viewer<PyViewerPassive>` can now pick up changes to rendering flags in
+    ``user_scn``, as requested in :github:issue:`1190`.
+
+Bug fixes
+^^^^^^^^^
+14. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
+    :github:issue:`1270`.
+15. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
+    bound of the length range. Fixes :github:issue:`1342`.
+
+
+Version 3.1.1 (December 18, 2023)
+-----------------------------------
+
+Bug fixes
+^^^^^^^^^
+1. Fixed a bug (introduced in 3.1.0) where box-box collisions produced no contacts if one box was deeply embedded in the
+   other.
+2. Fixed a bug in :ref:`simulate<saSimulate>` where the "LOADING..." message was not showing correctly.
+3. Fixed a crash in the Python :ref:`passive viewer<PyViewerPassive>`, when used with models containing Flex objects.
+4. Fixed a bug in MJX where ``site_xmat`` was ignored in ``get_data`` and ``put_data``
+5. Fixed a bug in MJX where ``efc_address`` was sometimes incorrectly calculated in ``get_data``.
+
+
+Version 3.1.0 (December 12, 2023)
+---------------------------------
+
+General
+^^^^^^^
+1. Improved convergence of Signed Distance Function (SDF) collisions by using line search and a new objective function
+   for the optimization. This allows to decrease the number of initial points needed for finding the contacts and is
+   more robust for very small or large geom sizes.
+2. Added :ref:`frame<frame>` to MJCF, a :ref:`meta-element<meta-element>` which defines a pure coordinate transformation
+   on its direct children, without requiring a :ref:`body<body>`.
+3. Added the :at:`kv` attribute to the :ref:`position<actuator-position>` and :ref:`intvelocity<actuator-intvelocity>`
+   actuators, for specifying actuator-applied damping. This can be used to implement a PD controller with 0 reference
+   velocity. When using this attribute, it is recommended to use the implicitfast or implicit
+   :ref:`integrators<geIntegration>`.
+
+Plugins
+^^^^^^^
+4. Allow actuator plugins to use activation variables in ``mjData.act`` as their internal state, rather than
+   ``mjData.plugin_state``. Actuator plugins can now specify :ref:`callbacks<mjpPlugin>` that compute activation
+   variables, and they can be used with built-in :ref:`dyntype<actuator-plugin-dyntype>` actuator dynamics.
+5. Added the `pid <https://github.com/deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ actuator plugin, a
+   configurable PID controller that implements the Integral term, which is not available with native MuJoCo actuators.
+
+MJX
+^^^
+6.  Added ``site_xpos`` and ``site_xmat`` to MJX.
+7. Added ``put_data``, ``put_model``, ``get_data`` to replace ``device_put`` and ``device_get_into``, which will be
+   deprecated. These new functions correctly translate fields that are the result of intermediate calculations such as
+   ``efc_J``.
+
+Bug fixes
+^^^^^^^^^
+8. Fix bug in Cartesian actuation with movable refsite, as when using body-centric Cartesian actuators on a quadruped.
+   Before this fix such actuators could lead to non-conservation of momentum.
+9. Fix bug that prevented using flex with :ref:`simulate<saSimulate>`.
+10. Fix bug that prevented the use of elasticity plugins in combination with pinned flex vertices.
+11. Release Python wheels targeting macOS 10.16 to support x86_64 systems where ``SYSTEM_VERSION_COMPAT`` is set.
+    The minimum supported version is still 11.0, but we release these wheels to fix compatibility for those users. See
+    :github:issue:`1213`.
+
+Version 3.0.1 (November 15, 2023)
+---------------------------------
+
+General
+^^^^^^^
+1. Added sub-terms of total passive forces in ``mjData.qfrc_passive`` to :ref:`mjData`:
+   ``qfrc_{spring, damper, gravcomp, fluid}``. The sum of these vectors equals ``qfrc_passive``.
+
+.. youtube:: H9qG9Zf2W44
+   :align: right
+   :width: 240px
+
+2. Added :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` attribute and associated
+   :ref:`mjOption.disableactuator<mjOption>` integer bitfield, which can be used to disable sets of actuators at runtime
+   according to their :ref:`group<actuator-general-group>`. Fixes :github:issue:`1092`. See :ref:`CActDisable`.
+
+   - The first 6 actuator groups are toggleable in the :ref:`simulate<saSimulate>` viewer. See `example model
+     <https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
+     and associated screen-capture on the right.
+
+3. Increased ``mjMAXUIITEM`` (maximum number of UI elements per section in Simulate) to 200.
+
+MJX
+^^^
+4. Added support for Newton solver (``mjSOL_NEWTON`` in :ref:`mjtSolver`).  The Newton solver significantly speeds up
+   simulation on GPU:
+
+   .. list-table:: Steps-per-second, Conjugate Gradient vs. Newton on A100
+      :header-rows: 1
+      :align: left
+
+      * - Model
+        - CG
+        - Newton
+        - Speedup
+      * - `Humanoid <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/humanoid>`__
+        - 640,000
+        - 1,020,000
+        - **1.6 x**
+      * - `Barkour v0 <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/barkour_v0>`__
+        - 1,290,000
+        - 1,750,000
+        - **1.35 x**
+      * - `Shadow Hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
+        - 215,000
+        - 270,000
+        - **1.25 x**
+
+   Humanoid is the standard MuJoCo humanoid,
+   `Google Barkour <https://blog.research.google/2023/05/barkour-benchmarking-animal-level.html>`__ and the Shadow Hand
+   are both available in the :ref:`MuJoCo Menagerie<Menagerie>`.
+5. Added support for joint equality constraints (``mjEQ_JOINT`` in :ref:`mjtEq`).
+6. Fixed bug where mixed ``jnt_limited`` joints were not being constrained correctly.
+7. Made ``device_put`` type validation more verbose (fixes :github:issue:`1113`).
+8. Removed empty EFC rows from ``MJX``, for joints with no limits (fixes :github:issue:`1117`).
+9. Fixed bug in ``scan.body_tree`` that led to incorrect smooth dynamics for some kinematic tree layouts.
+
+Python bindings
+^^^^^^^^^^^^^^^
+
+10. Fix the macOS ``mjpython`` launcher to work with the Python interpreter from Apple Command Line Tools.
+11. Fixed a crash when copying instances of ``mujoco.MjData`` for models that use plugins. Introduced a ``model``
+    attribute to ``MjData`` which is reference to the model that was used to create that ``MjData`` instance.
+
+Simulate
+^^^^^^^^
+12. :ref:`simulate<saSimulate>`: correct handling of "Pause update", "Fullscreen" and "VSync" buttons.
+
+Documentation
+^^^^^^^^^^^^^
+
+.. youtube:: cE3s_IfO4g4
+   :align: right
+   :width: 240px
+
+13. Added cell to the `tutorial colab <https://github.com/google-deepmind/mujoco#getting-started>`__ providing an
+    example of procedural camera control:
+14. Added documentation for the :ref:`UI` framework.
+15. Fixed typos and supported fields in docs (fixes :github:issue:`1105` and :github:issue:`1106`).
+
+
+Bug fixes
+^^^^^^^^^
+16. Fixed bug relating to welds modified with :ref:`torquescale<equality-weld-torquescale>`.
+
+Version 3.0.0 (October 18, 2023)
+--------------------------------
+
 New features
 ^^^^^^^^^^^^
 
-.. youtube:: Vc1tq0fFvQA
-   :align: right
-   :width: 240px
+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.
 
-1. 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<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
-   If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
-   see video.
+   - 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 <MjxFeatureParity>` 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
@@ -28,19 +217,50 @@ New features
    - Added new SDF plugin for defining implicit geometries. The plugin must define methods computing an SDF and its
      gradient at query points. See the :ref:`documentation<exWriting>` for more details.
 
-3. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
-   pipeline.
-
 .. youtube:: ra2bTiZHGlw
    :align: right
    :width: 240px
 
-4. Added capability to initialize :ref:`composite<body-composite>` particles with arbitrary positions.
+3. Added new low-level model element called ``flex``, used to define deformable objects. These
+   `simplicial complexes <https://en.wikipedia.org/wiki/Simplicial_complex>`__ 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<deformable>` section contains the low-level flex definition.
+   The :ref:`flexcomp<body-flexcomp>` element, similar to :ref:`composite<body-composite>` is a convenience macro for
+   creating deformables, and supports the GMSH tetrahedral file format.
 
-5. Added `shell <https://github.com/deepmind/mujoco/blob/main/plugin/elasticity/shell.cc>`__ passive force plugin:
+   - Added `shell <https://github.com/deepmind/mujoco/blob/main/plugin/elasticity/shell.cc>`__ passive force plugin,
+     computing bending forces using a constant precomputed Hessian (cotangent operator).
 
-   - Collisions use spheres located at mesh vertices.
-   - Stretching as tendon constraints and bending using a constant precomputed Hessian (cotangent operator).
+   **Note**: This feature is still under development and subject to change. In particular, deformable object
+   functionality is currently available both via :ref:`deformable<CDeformable>` and :ref:`composite<CComposite>`,
+   and both are modifiable by the first-party
+   `elasticity plugins <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__. We expect some of
+   this functionality to be unified in the future.
+
+.. youtube:: Vc1tq0fFvQA
+   :align: right
+   :width: 240px
+
+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<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
+   If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
+   see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
+
+5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
+   pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
+
+   - Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
+     :ref:`CG solver<option-solver>` is selected. The
+     `22 humanoids <https://github.com/deepmind/mujoco/blob/main/model/humanoid/22_humanoids.xml>`__ model shows a 3x
+     speedup compared to the single threaded simulation.
+   - Inertia-related computations and collision detection will happen in parallel.
+
+   Engine-internal threading is a work in progress and currently only available in first-party code via the
+   :ref:`testspeed<saTestspeed>` utility, exposed with the ``npoolthread`` flag.
+
+6. Added capability to initialize :ref:`composite<body-composite>` particles from OBJ files. Fixes :github:issue:`642`
+   and :github:issue:`674`.
 
 General
 ^^^^^^^
@@ -48,72 +268,75 @@ 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<siStack>` 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<mjData>`, 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<mjData>`, 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<glNumeric>` was set to 20.
-      - :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
-      - Added :ref:`mjData.solver_nisland<mjData>`: 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<glNumeric>` was set to 20.
+       - :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
+       - Added :ref:`mjData.solver_nisland<mjData>`: 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:**
 
        - For models which have ``<option collision="all"/>``, delete the attribute.
        - For models which have ``<option collision="dynamic"/>``, delete all :ref:`pair<contact-pair>` elements.
-       - For models which have ``<option collision="pair"/>``, disable all dynamic collisions (determined
+       - For models which have ``<option collision="predefined"/>``, disable all dynamic collisions (determined
          via contype/conaffinity) by first deleting all :ref:`contype<body-geom-contype>` and
          :ref:`conaffinity<body-geom-conaffinity>` attributes in the model and then setting them globally to ``0`` using
          |br| ``<default> <geom contype="0" conaffinity="0"/> </default>``.
 
-   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<mjData>` user input variable, for enabling/disabling the state of equality
+   13. Added :ref:`mjData.eq_active<mjData>` user input variable, for enabling/disabling the state of equality
        constraints. Renamed ``mjModel.eq_active`` to :ref:`mjModel.eq_active0<mjModel>`, which now has the semantic of
-       "initial value of ``mjData.eq_active``".
-       Fixes `#876 <https://github.com/google-deepmind/mujoco/discussions/876>`__
+       "initial value of ``mjData.eq_active``". Fixes :github:issue:`876`.
 
        **Migration:** Replace uses of ``mjModel.eq_active`` with ``mjData.eq_active``.
 
-13. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
+   14. Changed the default of :ref:`autolimits<compiler-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).
+
+15. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
     the exact formula rather than with Euler integration.
-14. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
-    actuator forces: using the next step's actuator state to compute the current actuator forces at the current timestep.
-15. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
+16. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
+    actuator forces: using the next step's actuator state to compute the current actuator forces.
+17. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
     :ref:`actuatorfrcrange<body-joint-actuatorfrcrange>` and
     :ref:`actuatorfrclimited<body-joint-actuatorfrclimited>`, respectively.
-16. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
+18. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
     Euler integrator. See the :ref:`Numerical Integration<geIntegration>` section for more details.
-17. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
+19. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
     :ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
-18. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
+20. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
     linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning.
-19. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
-    mesh assets. Fixes `#409 <https://github.com/google-deepmind/mujoco/issues/409>`__ .
-20. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
+21. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
+    mesh assets. Fixes :github:issue:`409`.
+22. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-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.
-21. Added :ref:`camera<body-camera>` calibration attributes:
+23. Added :ref:`camera<body-camera>` calibration attributes:
 
     - The new attributes are :ref:`resolution<body-camera-resolution>`, :ref:`focal<body-camera-focal>`,
       :ref:`focalpixel<body-camera-focalpixel>`, :ref:`principal<body-camera-principal>`,
@@ -122,14 +345,22 @@ General
       attributes are specified. See the following
       `example model <https://github.com/deepmind/mujoco/blob/main/test/engine/testdata/vis_visualize/frustum.xml>`__.
     - Note that these attributes only take effect for offline rendering and do not affect interactive visualisation.
-22. Added multi-threaded constraint solving via :ref:`mj_island` and :ref:`mjThreadPool` to :ref:`testspeed<saTestspeed>`
-    exposed via npoolthread flag. The `22 humanoids <https://github.com/deepmind/mujoco/blob/main/model/humanoid/22_humanoids.xml>`__
-    model shows a 3x speedup compared to the single threaded simulation.
+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 <https://github.com/aftersomemath>`__.
+25. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the
+    `WriteReadCompare <https://github.com/google-deepmind/mujoco/blob/main/test/xml/xml_native_writer_test.cc>`__ test.
+26. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`.
+
 Python bindings
 ^^^^^^^^^^^^^^^
 
-23. Fixed `#870 <https://github.com/google-deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
-    camera name used the default camera.
+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<PyViewerPassive>` handle, which allows users to add custom
+    visualization geoms (:github:issue:`1023`).
+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
 ^^^^^^^^
@@ -138,20 +369,33 @@ Simulate
    :align: right
    :width: 240px
 
-24. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
+30. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
     :ref:`Python viewer<PyViewerManaged>`. State history can be viewed by scrubbing the History slider and (more
     precisely) with the left and right arrow keys. See screen capture:
 
-25. The ``LOADING...`` label is now shown correctly.
-    `Contribution <https://github.com/google-deepmind/mujoco/pull/1070>`__ by
+31. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by
     `Levi Burner <https://github.com/aftersomemath>`__.
 
+Documentation
+^^^^^^^^^^^^^
+
+.. youtube:: nljr0X79vI0
+   :align: right
+   :width: 240px
+
+32. Added :doc:`detailed documentation <computation/fluid>` of fluid force modeling, and an illustrative example model
+    showing `tumbling cards <https://github.com/google-deepmind/mujoco/blob/main/model/cards/cards.xml>`__ using the
+    ellipsoid-based fluid model.
+
 Bug fixes
 ^^^^^^^^^
 
-26. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
+33. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
     midphase collision trees.
 
+34. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints.
+
+
 Version 2.3.7 (July 20, 2023)
 -----------------------------
 
@@ -165,8 +409,8 @@ General
    :ref:`Cartesian actuator<actuator-general-refsite>` forces are realizable by individual motors at the joints.
    See :ref:`CForceRange` for details.
 #. Added an optional ``content_type`` attribute to hfield, texture, and mesh assets. This attribute supports a formatted
-   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`__ (previously known as MIME type)
-   string used to determine the type of the asset file without resorting to pulling the type from the file extension.
+   `Media Type <https://www.iana.org/assignments/media-types/media-types.xhtml>`_ (previously known as MIME type) string
+   used to determine the type of the asset file without resorting to pulling the type from the file extension.
 #. Added analytic derivatives for quaternion :ref:`subtraction<mjd_subQuat>` and :ref:`integration<mjd_quatIntegrate>`
    (rotation with an angular velocity). Derivatives are in the 3D tangent space.
 #. Added :ref:`mjv_connector` which has identical functionality to :ref:`mjv_makeConnector`, but with more convenient
@@ -178,13 +422,10 @@ Python bindings
 
 7. The :ref:`passive viewer<PyViewerPassive>` handle now exposes ``update_hfield``, ``update_mesh``, and
    ``update_texture`` methods to allow users to update renderable assets.
-   (`#812 <https://github.com/google-deepmind/mujoco/issues/812>`_,
-   `#958 <https://github.com/google-deepmind/mujoco/issues/958>`_,
-   `#965 <https://github.com/google-deepmind/mujoco/issues/965>`_)
-#. Allow a custom keyboard event callback to be specified in the :ref:`passive viewer<PyViewerPassive>`.
-   (`#766 <https://github.com/google-deepmind/mujoco/issues/766>`_)
-#. Fix GLFW crash when Python exits while the passive viewer is running.
-   (`#790 <https://github.com/google-deepmind/mujoco/issues/790>`_)
+   (Issues :github:issue:`812`, :github:issue:`958`, :github:issue:`965`).
+#. Allow a custom keyboard event callback to be specified in the :ref:`passive viewer<PyViewerPassive>`
+   (:github:issue:`766`).
+#. Fix GLFW crash when Python exits while the passive viewer is running (:github:issue:`790`).
 
 Models
 ^^^^^^
@@ -405,6 +646,8 @@ General
    :align: right
    :width: 350px
 
+.. _midphase:
+
 2. Added a collision mid-phase for pruning geoms in body pairs, see :ref:`documentation<coSelection>` for more details.
    This is based on static AABB bounding volume hierarchy (a BVH binary tree) in the body inertial frame. The GIF on
    the right is cut from `this longer video <https://youtu.be/e0babIM8hBo>`__.

doc/computation/fluid.rst

@@ -26,36 +26,55 @@ to positive values. These parameters correspond to the density :math:`\rho` and
 Inertia model
 -------------
 
-In this model, the shape of each body, for fluid dynamics purposes, is assumed to be the *equivalent inertia box*,
-which can also be visualized. Each forward-facing (relative to the linear velocity) face of the box experiences force
-along its normal direction. All faces also experience torque due to the angular velocity; this torque is obtained by
-integrating the force resulting from the rotation over the surface area. In this sub-section, let :math:`v` and
-:math:`\omega` denote the linear and angular body velocity in the body local frame (aligned with the equivalent
-inertia box), and :math:`s` the 3D vector of box sizes. When the contributions from all faces are added, the resulting
-force and torque applied to the body by a fluid of density :math:`\rho`, in local body coordinates, have the
-:math:`i`-th component
+In this model the shape of each body, for fluid dynamics purposes, is assumed to be the *equivalent inertia box*,
+which can also be visualized. For a body with mass :math:`\mathcal{M}` and inertia matrix :math:`\mathcal{I}`, the
+half-dimensions (i.e. half-width, half-depth and half-height) of the equivalent inertia box are
+
+.. math::
+   \begin{align*}
+   r_x = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{yy} + \mathcal{I}_{zz} - \mathcal{I}_{xx} \right)}  \\
+   r_y = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{zz} + \mathcal{I}_{xx} - \mathcal{I}_{yy} \right)}  \\
+   r_z = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{xx} + \mathcal{I}_{yy} - \mathcal{I}_{zz} \right)}
+   \end{align*}
+
+Let :math:`\mathbf{v}` and :math:`\boldsymbol{\omega}` denote the linear and angular body velocity of the body in
+the body-local frame (aligned with the equivalent inertia box). The force :math:`\mathbf{f}_{\text{inertia}}` and
+torque :math:`\mathbf{g}_{\text{inertia}}` exerted by the fluid onto the solid are the sum of the terms
+
+.. math::
+   \begin{align*}
+   \mathbf{f}_{\text{inertia}} &= \mathbf{f}_D + \mathbf{f}_V  \\
+   \mathbf{g}_{\text{inertia}} &= \mathbf{g}_D + \mathbf{g}_V
+   \end{align*}
+
+Here subscripts :math:`D` and :math:`V` denote quadratic Drag and Viscous resistance.
+
+The quadratic drag terms depend on the density :math:`\rho` of the fluid, scale quadratically with the velocity
+of the body, and are a valid approximation of the fluid forces at high Reynolds numbers.
+The torque is obtained by integrating the force resulting from the rotation over the surface area.
+The :math:`i`-th component of the force and torque can be written as
 
 .. math::
    \begin{aligned}
-   \text{density force} : \quad &- {1 \over 2} \rho s_j s_k |v_i| v_i \\
-   \text{density torque} : \quad &- {1 \over 64} \rho s_i \left(s_j^4 + s_k^4 \right) |\omega_i| \omega_i \\
+   f_{D, i} = \quad &- 2  \rho r_j r_k |v_i| v_i \\
+   g_{D, i} = \quad &- {1 \over 2} \rho r_i \left(r_j^4 + r_k^4 \right) |\omega_i| \omega_i \\
    \end{aligned}
 
-This model implicitly assumes high Reynolds numbers, with lift-to-drag ratio equal to the tangent of the angle of
-attack. One can also specify a non-zero :ref:`wind<option-wind>`, which is a 3D vector subtracted from the body linear
-velocity in the fluid dynamics computation.
-
-Each body also experiences a force and a torque proportional to the viscosity :math:`\beta` and opposite to its linear and
-angular velocity. Note that viscosity can be used independent of density, to make the simulation more damped. We use the
-formulas for a sphere at low Reynolds numbers, with diameter :math:`d` equal to the average of the equivalent inertia
-box sizes. The resulting 3D force and torque in local body coordinates are
+The viscous resistance terms depend on the fluid viscosity :math:`\beta`, scale linearly with the body velocity, and
+approximate the fluid forces at low Reynolds numbers. Note that viscosity can be used independent of density to make
+the simulation more damped. We use the formulas for the equivalent sphere with radius
+:math:`r_{eq} = (r_x + r_y + r_z) / 3`  at low Reynolds numbers. The resulting 3D force and torque in local
+body coordinates are
 
 .. math::
    \begin{aligned}
-   \text{viscosity force} : \quad &- 3 \beta \pi d v \\
-   \text{viscosity torque} : \quad &- \beta \pi d^3 \omega \\
+   f_{V, i} = \quad &- 6 \beta \pi r_{eq} v_i \\
+   g_{V, i} = \quad &- 8 \beta \pi r_{eq}^3 \omega_i \\
    \end{aligned}
 
+One can also affect these forces by specifing a non-zero :ref:`wind<option-wind>`, which is a 3D vector subtracted
+from the body linear velocity in the fluid dynamics computation.
+
 .. _flEllipsoid:
 
 Ellipsoid model
@@ -116,35 +135,36 @@ also disables the inertia-based model for the parent body. The
      - 1.0
 
 Elements of the model are a generalization of :cite:t:`andersen2005b` to 3 dimensions.
-The force :math:`\mathbf{f}_{\text{fluid}\rightarrow \text{solid}}` and torque
-:math:`\mathbf{g}_{\text{fluid} \rightarrow \text{solid}}` exerted by the fluid onto the solid are
+The force :math:`\mathbf{f}_{\text{ellipsoid}}` and torque
+:math:`\mathbf{g}_{\text{ellipsoid}}` exerted by the fluid onto the solid are
 the sum of of the terms
 
 .. math::
    \begin{align*}
-   \mathbf{f}_{\text{fluid} \rightarrow \text{solid}} &= \mathbf{f}_A + \mathbf{f}_D + \mathbf{f}_M + \mathbf{f}_K  \\
-   \mathbf{g}_{\text{fluid} \rightarrow \text{solid}} &= \mathbf{g}_A + \mathbf{g}_D
+   \mathbf{f}_{\text{ellipsoid}} &= \mathbf{f}_A + \mathbf{f}_D + \mathbf{f}_M + \mathbf{f}_K + \mathbf{f}_V  \\
+   \mathbf{g}_{\text{ellipsoid}} &= \mathbf{g}_A + \mathbf{g}_D + \mathbf{g}_V
    \end{align*}
 
-Where subscripts :math:`A`, :math:`D`, :math:`M` and :math:`K`, denote Added mass, viscous Drag, Magnus lift and
-Kutta lift, respectively. The :math:`D`, :math:`M` and :math:`K` terms are scaled by the respective
-:math:`C_D`, :math:`C_M` and :math:`C_K` coefficients above, while the added mass term cannot be scaled.
+Where subscripts :math:`A`, :math:`D`, :math:`M`, :math:`K` and  :math:`V` denote Added mass, viscous Drag, Magnus lift,
+Kutta lift and Viscous resistance, respectively. The :math:`D`, :math:`M` and :math:`K` terms are scaled by the respective
+:math:`C_D`, :math:`C_M` and :math:`C_K` coefficients above, the viscous resistance scales with the fluid viscosity
+:math:`\beta`, while the added mass term cannot be scaled.
 
 Notation
 ~~~~~~~~
 
 We describe the motion of the object in an inviscid, incompressible quiescent fluid of density :math:`\rho`. The
 arbitrarily-shaped object is described in the model as the equivalent ellipsoid of semi-axes
-:math:`\mathbf{d} = \{d_x, d_y, d_z\}`.
+:math:`\mathbf{r} = \{r_x, r_y, r_z\}`.
 The problem is described in a reference frame aligned with the sides of the ellipsoid and moving with it. The
 body has velocity :math:`\mathbf{v} = \{v_x, v_y, v_z\}` and angular velocity
 :math:`\boldsymbol{\omega} = \{\omega_x, \omega_y, \omega_z\}`. We will also use
 
 .. math::
    \begin{align*}
-       d_\text{max} &= \max(d_x, d_y, d_z) \\
-       d_\text{min} &= \min(d_x, d_y, d_z) \\
-       d_\text{mid} &= d_x + d_y + d_z - d_\text{max} - d_\text{min}
+       r_\text{max} &= \max(r_x, r_y, r_z) \\
+       r_\text{min} &= \min(r_x, r_y, r_z) \\
+       r_\text{mid} &= r_x + r_y + r_z - r_\text{max} - r_\text{min}
    \end{align*}
 
 The Reynolds number is the ratio between inertial and viscous forces within a flow and is defined as :math:`Re=u~l/\beta`, where
@@ -173,18 +193,149 @@ Finally, we use the subscripts :math:`i, j, k` to denote triplets of equations t
 Ellipsoid projection
 ~~~~~~~~~~~~~~~~~~~~
 
-We provide the following result without proof. For the derivation, contact the development team.
+We present the following result.
 
 .. admonition:: Lemma
    :class: note
 
-   Given an ellipsoid at the origin with semi-axes :math:`(d_x, d_y, d_z)` aligned
-   with the coordinate axes :math:`(x, y, z)`, and a unit vector :math:`\mathbf{u} = (u_x, u_y, u_z)`,
-   the area projected by the ellipsoid onto the plane normal to :math:`\mathbf{u}` is
+   Given an ellipsoid with semi-axes :math:`(r_x, r_y, r_z)` aligned with the coordinate axes :math:`(x, y, z)`, and a
+   unit vector :math:`\mathbf{u} = (u_x, u_y, u_z)`, the area projected by the ellipsoid onto the plane normal to
+   :math:`\mathbf{u}` is
 
    .. math::
-      A^{\mathrm{proj}}_{\mathbf{u}} = \pi \sqrt{\frac{d_y^4 d_z^4 u_x^2 + d_x^4 d_z^4 u_y^2 + d_x^4 d_y^4 u_z^2}{d_y^2 d_z^2 u_x^2 + d_x^2 d_z^2 u_y^2 + d_x^2 d_y^2 u_z^2}}
+      A^{\mathrm{proj}}_{\mathbf{u}} = \pi \sqrt{\frac{r_y^4 r_z^4 u_x^2 + r_z^4 r_x^4 u_y^2 + r_x^4 r_y^4 u_z^2}{r_y^2 r_z^2 u_x^2 + r_z^2 r_x^2 u_y^2 + r_x^2 r_y^2 u_z^2}}
 
+.. collapse:: Expand for derivation
+
+   .. admonition:: Derivation of lemma
+      :class: tip
+
+      **Area of an ellipse**
+         Any ellipse centered at the origin can be described in terms of a quadratic form a
+         :math:`\mathbf{x}^T Q \mathbf{x} = 1`, where :math:`Q` is a real, symmetric, positive-definite 2x2 matrix that
+         defines the orientation and  semi-axis lengths of the ellipse, and :math:`\mathbf{x} = (x, y)` are points on the
+         ellipse. The area of the ellipse is given by
+
+         .. math::
+            A = \frac{\pi}{\sqrt{\det Q}} .
+
+      **Ellipsoid cross-section**
+         We begin by computing the area of the ellipse formed by intersecting an ellipsoid centered at the origin with the
+         plane :math:`\Pi_{\mathbf{n}}` through the origin with unit normal :math:`\mathbf{n} = (n_x, n_y, n_z)`. Let
+         :math:`(r_x, r_y, r_z)` be the semi-axis lengths of the ellipsoid. Without loss of generality, it is sufficient to
+         assume that the axes of the ellipsoid are aligned with the coordinate axes. The ellipsoid can then be described as
+         :math:`\mathbf{x}^T Q \mathbf{x} = 1`, where
+         :math:`Q = \textrm{diag}\mathopen{}\left( \left. 1 \middle/ r_x^2 \right., \left. 1 \middle/ r_y^2 \right., \left. 1 \middle/ r_z^2 \right. \right)\mathclose{}`
+         and :math:`\mathbf{x} = (x, y, z)` are the points on the ellipsoid.
+
+         We proceed by rotating the plane :math:`\Pi_{\mathbf{n}}` together with the ellipsoid so that the normal of the
+         rotated plane points along the :math:`z` axis. This would then allow us to get the desired intersection by setting
+         the :math:`z` coordinate to zero. Writing :math:`\mathbf{\hat{z}}` for the unit vector along the :math:`z` axis, we
+         have
+
+         .. math::
+            \begin{align*}
+            \mathbf{n} \times \mathbf{\hat{z}} &= \sin\theta \, \mathbf{m}, \\
+            \mathbf{n} \cdot \mathbf{\hat{z}} &= \cos\theta ,
+            \end{align*}
+
+         where :math:`\mathbf{m}` is the unit vector that defines the rotation axis and :math:`\theta` is the rotation
+         angle. We can rearrange these to get quantities that we need to form a rotation quaternion, namely
+
+         .. math::
+            \begin{align*}
+            \cos\frac{\theta}{2}
+            &= \sqrt{\frac{1+\cos\theta}{2}}
+            &= \sqrt{\frac{1 + \mathbf{n} \cdot \mathbf{\hat{z}}}{2}}, \\
+            \sin\frac{\theta}{2}\,\mathbf{m}
+            &= \frac{\mathbf{n} \times \mathbf{\hat{z}}}{2\cos\frac{\theta}{2}}
+            &= \frac{\mathbf{n} \times \mathbf{\hat{z}}}{\sqrt{2 (1 + \mathbf{n} \cdot \mathbf{\hat{z}})}} .
+            \end{align*}
+
+         The rotation quaternion :math:`q = q_r + q_x \mathbf{i} + q_y \mathbf{j} + q_z \mathbf{k}` is therefore given by
+
+         .. math::
+            q_r = \sqrt{\frac{1 + n_z}{2}}, \quad
+            q_x = \frac{n_y}{\sqrt{2 \left(1+n_z\right)}},  \quad
+            q_y = \frac{-n_x}{\sqrt{2 \left(1+n_z\right)}}, \quad
+            q_z = 0 .
+
+         From this, the rotation matrix is given by
+
+         .. math::
+            \def\arraystretch{1.33}
+            \begin{align*}
+            R &= \begin{pmatrix}
+            1 - 2 q_y^2 - 2 q_z^2 & 2 \left(q_x q_y - q_r q_z\right) & 2 \left(q_x q_z + q_r q_y\right) \\
+            2 \left(q_x q_y + q_r q_z\right) & 1 - 2 q_x^2 - 2 q_z^2 & 2 \left(q_y q_z - q_r q_x\right) \\
+            2 \left(q_x q_z - q_r q_y\right) & 2 \left(q_y q_z + q_r q_x\right) & 1 - 2 q_x^2 - 2 q_y^2
+            \end{pmatrix} \\
+            &= \begin{pmatrix}
+            1 - \left. n_x^2 \middle/ \left( 1+n_z \right) \right. & \left. -n_x n_y \middle/ \left( 1+n_z \right) \right. & -n_x \\
+            \left. -n_x n_y \middle/ \left( 1+n_z \right) \right. & 1 - \left. n_y^2 \middle/ \left( 1+n_z \right) \right. & -n_y \\
+            n_x & n_y & 1 - \left. \left( n_x^2 + n_y^2 \right) \middle/ \left( \vphantom{n_x^2} 1+n_z \right) \right.
+            \end{pmatrix},
+            \end{align*}
+
+         and the rotated ellipsoid is described via the transformed quadratic form
+
+         .. math::
+            \mathbf{x}^T Q' \mathbf{x} = \mathbf{x}^T \left( R^T Q R \right) \mathbf{x} = 1 .
+
+         From the formula for ellipse area above, for the area of the ellipse at :math:`z=0`, we need
+
+         .. math::
+            \begin{align*}
+            Q'_{xx} &= \frac{1}{r_x^2} R_{xx}^2 + \frac{1}{r_y^2} R_{yx}^2 + \frac{1}{r_z^2} R_{zx}^2 , \\
+            Q'_{yy} &= \frac{1}{r_x^2} R_{xy}^2 + \frac{1}{r_y^2} R_{yy}^2 + \frac{1}{r_z^2} R_{zy}^2 , \\
+            Q'_{xy} &= \frac{1}{r_x^2} R_{xx} R_{xy} + \frac{1}{r_y^2} R_{yx} R_{yy} + \frac{1}{r_z^2} R_{zx} R_{zy} ,
+            \end{align*}
+
+         and the desired area is given by
+
+         .. math::
+            A^{\cap}_{\mathbf{n}}
+            = \frac{\pi}{\sqrt{\vphantom{Q'^2_{xy}} \det Q'}}
+            = \frac{\pi}{\sqrt{Q'_{xx} Q'_{yy} - Q'^2_{xy}}}
+            = \frac{\pi r_x r_y r_z}{\sqrt{r_x^2 n_x^2 + r_y^2 n_y^2 + r_z^2 n_z^2}},
+
+         where the superscript :math:`\cap` denotes that the area pertains to the ellipse at the *intersection*
+         with :math:`\Pi_{\mathbf{n}}`.
+
+      **Projected ellipse**
+         Let :math:`\mathbf{u} = (u_x, u_y, u_z)` be some unit vector (in our context, it is the direction of the velocity
+         of the fluid impinging on an ellipsoid) and let :math:`\Pi_{\mathbf{u}}` be the plane normal to :math:`\mathbf{u}`.
+         In general, the ellipse formed by projecting an ellipsoid :math:`\mathcal{E}` onto :math:`\Pi_{\mathbf{u}}`
+         (denoted :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}`) is different from the one formed by intersecting
+         :math:`\mathcal{E}` with :math:`\Pi_{\mathbf{u}}` (denoted :math:`\mathcal{E}^{\cap}_{\mathbf{u}}`).
+
+         An important property of :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}` is that :math:`\mathbf{u}` is tangent
+         tangent to the ellipsoid :math:`\mathcal{E}` at every point on :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}`.
+
+         We can regard :math:`\mathcal{E}` as the image of the unit sphere :math:`\mathcal{S}` under a stretching
+         transformation :math:`T = \mathrm{diag}(r_x, r_y, r_z)`. Furthermore, if :math:`\mathbf{\tilde{u}}` is a vector
+         tangent to :math:`\mathcal{S}`, then its image
+         :math:`\mathbf{u}=T\mathbf{\tilde{u}}=(r_x \tilde{u}_x, r_y \tilde{u}_y, r_z \tilde{u}_z)` is tangent to the
+         ellipsoid. The ellipse :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}` is therefore the image
+         under :math:`T` of the circle :math:`\mathcal{C}^{\cap}_{\mathbf{\tilde{u}}}` at the intersection between
+         :math:`\mathcal{S}` and :math:`\Pi_{\mathbf{\tilde{u}}}` (for spheres :math:`\mathcal{C}^{\cap}` and
+         :math:`\mathcal{C}^{\mathrm{proj}}` do coincide).
+
+         Let :math:`\mathbf{\tilde{v}}` and :math:`\mathbf{\tilde{w}}` be some orthogonal pair of vectors in the plane
+         :math:`\Pi_{\mathbf{\tilde{u}}}`, then :math:`\mathbf{\tilde{u}} = \mathbf{\tilde{v}} \times \mathbf{\tilde{w}}`.
+         Their images under :math:`T` are :math:`\mathbf{v} = (r_x \tilde{v}_x, r_y \tilde{v}_y, r_z \tilde {v}_z)` and
+         :math:`\mathbf{w} = (r_x \tilde{w}_x, r_y \tilde{w}_y, r_z \tilde {w}_z)` respectively, and they remain orthogonal
+         vectors in the plane of :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}`. A (non-unit) normal to the ellipse
+         :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}}` is therefore given by
+
+         .. math::
+            \mathbf{N} = \mathbf{v} \times \mathbf{w}
+            = (r_y r_z \tilde{u}_x, r_z r_x \tilde{u}_y, r_x r_y \tilde{u}_z)
+            = \left( \frac{r_y r_z}{r_x} u_x, \frac{r_z r_x}{r_y} u_y, \frac{r_x r_y}{r_z} u_z \right).
+
+         This shows that :math:`\mathcal{E}^{\mathrm{proj}}_{\mathbf{u}} = \mathcal{E}^{\cap}_{\mathbf{n}}`, where
+         :math:`\mathbf{n} = \mathbf{N} / \left\Vert\mathbf{N}\right\Vert`. Its area is given by the formula derived in the
+         previous section, leading to the result stated above.
 
 Added mass
 ~~~~~~~~~~
@@ -228,11 +379,11 @@ Here :math:`\circ` denotes an element-wise product, :math:`\dot{\mathbf{v}}` is
 :math:`\dot{\boldsymbol{\omega}}` is the angular acceleration. :math:`\mathbf{m}_A \circ \mathbf{v}` and
 :math:`\mathbf{I}_A \circ \boldsymbol{\omega}` are the virtual linear and angular momentum respectively.
 
-For an ellipsoid of semi-axis :math:`\mathbf{d} = \{d_x, d_y, d_z\}` and volume :math:`V = 4 \pi d_x d_y d_z / 3`, the
+For an ellipsoid of semi-axis :math:`\mathbf{r} = \{r_x, r_y, r_z\}` and volume :math:`V = 4 \pi r_x r_y r_z / 3`, the
 virtual inertia coefficients were derived by :cite:t:`tuckerman1925`. Let:
 
 .. math::
-   \kappa_i = \int_0^\infty \frac{d_i d_j d_k}{\sqrt{(d_i^2 + \lambda)^3 (d_j^2 + \lambda) (d_k^2 + \lambda)}} \textrm{d} \lambda
+   \kappa_i = \int_0^\infty \frac{r_i r_j r_k}{\sqrt{(r_i^2 + \lambda)^3 (r_j^2 + \lambda) (r_k^2 + \lambda)}} \textrm{d} \lambda
 
 
 It should be noted that these coefficients are non-dimensional (i.e. if all semi-axes are multiplied by the same scalar
@@ -244,7 +395,7 @@ the coefficients remain the same). The virtual masses of the ellipsoid are:
 And the virtual moments of inertia are:
 
 .. math::
-   I_{A, i} = \frac{\rho V}{5} \frac{(d_j^2 - d_k^2)^2 (\kappa_k-\kappa_j)}{2(d_j^2 - d_k^2) + (d_j^2 + d_k^2) (\kappa_j-\kappa_k)}
+   I_{A, i} = \frac{\rho V}{5} \frac{(r_j^2 - r_k^2)^2 (\kappa_k-\kappa_j)}{2(r_j^2 - r_k^2) + (r_j^2 + r_k^2) (\kappa_j-\kappa_k)}
 
 Viscous drag
 ~~~~~~~~~~~~
@@ -265,17 +416,22 @@ viscosity and scales with the second power of the velocity. It can be written as
 Where :math:`C_D` is a drag coefficient, and :math:`A_D` is a reference surface area (e.g. a measure of the projected
 area on the plane normal to the flow), and :math:`I_D` a reference moment of inertia.
 
+.. youtube:: nljr0X79vI0
+   :align: right
+   :width: 50%
+
 Even for simple shapes, the terms :math:`C_D`, :math:`A_D` and :math:`I_D` need to be tuned to the problem-specific
 physics and dynamical scales :cite:p:`duan2015`. For example, the drag coefficient :math:`C_D` generally decreases with
 increasing Reynolds numbers, and a single reference area :math:`A_D` may not be sufficient to account for the skin
 drag for highly irregular or slender bodies. For example, experimental fits are derived from problems ranging from
 falling playing cards :cite:p:`wang2004,andersen2005a,andersen2005b` to particle transport :cite:p:`loth2008,
-bagheri2016`.
+bagheri2016`. See screen capture of the
+`cards.xml <https://github.com/deepmind/mujoco/blob/main/model/card/cards.xml>`__ model on the right.
 
 We derive a formula for :math:`\mathbf{f}_\text{D}` based on two surfaces :math:`A^\text{proj}_\mathbf{v}` and
 :math:`A_\text{max}`. The first, :math:`A^\text{proj}_\mathbf{v}`, is the cylindrical projection of the body onto a
 plane normal to the velocity :math:`\mathbf{v}`. The second is the maximum projected surface
-:math:`A_\text{max} = 4 \pi d_{max} d_{min}`.
+:math:`A_\text{max} = 4 \pi r_{max} r_{min}`.
 
 .. math::
    \mathbf{f}_\text{D} = - \rho~ \big[  C_{D, \text{blunt}} ~ A^\text{proj}_\mathbf{v} ~ +
@@ -288,7 +444,7 @@ maximum swept ellipsoid obtained by the rotation of the body around the axis. Th
 moment of inertia are:
 
 .. math::
-   \mathbf{I}_{D,ii} = \frac{8\pi}{15} ~d_i ~\max(d_j, ~d_k)^4 .
+   \mathbf{I}_{D,ii} = \frac{8\pi}{15} ~r_i ~\max(r_j, ~r_k)^4 .
 
 Given this reference moment of inertia, the angular drag torque is computed as:
 
@@ -299,22 +455,21 @@ Given this reference moment of inertia, the angular drag torque is computed as:
 
 Here :math:`\mathbf{I}_\text{max}` is a vector with each entry equal to the maximal component of :math:`\mathbf{I}_D`.
 
-The viscosity :math:`\beta`
-For Reynolds numbers around or below :math:`O(10)`, the drag is best approximated as linear in the flow velocity
-(e.g. Stokes' law). For example, for a sphere the drag force :cite:p:`stokes1850` and torque :cite:p:`lamb1932` are:
+Finally the viscous resistance terms, also known as linear drag, well approvimate the fluid forces for Reynolds
+numbers around or below :math:`O(10)`. These are computed for the equivalent sphere with Stokes' law
+:cite:p:`stokes1850,lamb1932`:
 
 .. math::
    \begin{align*}
-   \mathbf{f}_\text{S} &= - 6 \pi r_D \rho ~ \beta \mathbf{v}\\
-   \mathbf{g}_\text{S} &= - 8 \pi r_D^3 \rho ~ \beta \boldsymbol{\omega}
+   \mathbf{f}_\text{V} &= - 6 \pi r_D \beta \mathbf{v}\\
+   \mathbf{g}_\text{V} &= - 8 \pi r_D^3 \beta \boldsymbol{\omega}
    \end{align*}
 
-Here, :math:`r_D` is the radius of the sphere and :math:`\beta` is the kinematic viscosity of the medium (e.g.
-:math:`1.48~\times 10^{-5}~m^2/s` for ambient-temperature air and :math:`0.89 \times 10^{-4}~m^2/s` for water). Here,
-for simplicity, we estimate the radius of the equivalent sphere as :math:`r_D = (d_x + d_y + d_z)/3`. To make a
-quantitative example, Stokes' law become accurate for room-temperature air if
-:math:`u\cdot l \lesssim 2 \times 10^{-4}~m^2/s`, where :math:`u` is the speed and :math:`l` a characteristic length of
-the body.
+Here, :math:`r_D = (r_x + r_y + r_z)/3` is the radius of the equivalent sphere and :math:`\beta` is the kinematic
+viscosity of the medium (e.g. :math:`1.48~\times 10^{-5}~m^2/s` for ambient-temperature air and
+:math:`0.89 \times 10^{-4}~m^2/s` for water). To make a quantitative example, Stokes' law become accurate for
+room-temperature air if :math:`u\cdot l \lesssim 2 \times 10^{-4}~m^2/s`, where :math:`u` is the speed and
+:math:`l` a characteristic length of the body.
 
 Viscous lift
 ~~~~~~~~~~~~
@@ -351,7 +506,7 @@ It's worth making an example. To reduce the number of variables, suppose a body
 sum of the force due to added mass and the force due to the Magnus effect along, for example, :math:`x` is:
 
 .. math::
-   \frac{f}{\pi \rho d_z} = v_y \omega_z \left(2 d_x \min\{d_x, d_z\} - (d_x + d_z)^2\right)
+   \frac{f}{\pi \rho r_z} = v_y \omega_z \left(2 r_x \min\{r_x, r_z\} - (r_x + r_z)^2\right)
 
 Note that the two terms have opposite signs.
 
@@ -365,8 +520,14 @@ in the surrounding flow a circulation of sufficient strength to hold the rear st
 This is the Kutta condition, a fluid dynamic phenomenon that can be observed for solid bodies with sharp corners, such
 as slender bodies or the trailing edges of airfoils.
 
-.. cssclass:: caption-small
 .. figure:: ../images/computation/kutta_cond_plate.svg
+   :class: only-light
+   :figwidth: 95%
+   :align: left
+
+.. cssclass:: caption-small
+.. figure:: ../images/computation/kutta_cond_plate_dark.svg
+   :class: only-dark
    :figwidth: 95%
    :align: left
 
@@ -376,13 +537,13 @@ as slender bodies or the trailing edges of airfoils.
    upward force acting on the plate.
 
 For a two-dimensional flow sketched in the figure above, the circulation due to the Kutta condition can be estimated as:
-:math:`\Gamma_\text{K} = C_K ~ d_x ~ \| \mathbf{v}\| ~ \sin(2\alpha)`,
+:math:`\Gamma_\text{K} = C_K ~ r_x ~ \| \mathbf{v}\| ~ \sin(2\alpha)`,
 where :math:`C_K` is a lift coefficient, and :math:`\alpha` is the angle between the velocity vector and its projection
 onto the surface. The lift force per unit length can be computed with the Kutta–Joukowski theorem as
 :math:`\mathbf{f}_K / L = \rho \Gamma_\text{K} \times \mathbf{v}`.
 
 In order to extend the lift force equation to three-dimensional motions, we consider the normal
-:math:`\mathbf{n}_{s, \mathbf{v}} = \{\frac{d_y d_z}{d_x}v_x, \frac{d_z d_x}{d_y}v_y, \frac{d_x d_x}{d_z}v_z\}`
+:math:`\mathbf{n}_{s, \mathbf{v}} = \{\frac{r_y r_z}{r_x}v_x, \frac{r_z r_x}{r_y}v_y, \frac{r_x r_x}{r_z}v_z\}`
 to the cross-section of the body which generates the body's projection :math:`A^\text{proj}_\mathbf{v}` onto a plane
 normal to the velocity given in the :ref:`lemma<flProjection>` above and the corresponding unit vector
 :math:`\hat{\mathbf{n}}_{s, \mathbf{v}}`.
@@ -403,25 +564,25 @@ Here, :math:`\hat{\mathbf{v}}` is the unit-normal along :math:`\mathbf{v}`. Note
 example, for spherical bodies :math:`\hat{\mathbf{n}}_{s, \mathbf{v}} \equiv \hat{\mathbf{v}}` and by construction
 :math:`\mathbf{f}_\text{K} = 0`.
 
-Let's unpack the relation with an example. Suppose a body with :math:`d_x = d_y` and :math:`d_z \ll d_x`. Note that the vector
+Let's unpack the relation with an example. Suppose a body with :math:`r_x = r_y` and :math:`r_z \ll r_x`. Note that the vector
 :math:`\hat{\mathbf{n}}_{s, \mathbf{v}} \times \hat{\mathbf{v}}` gives the direction of the circulation induced by the
-deflection of the flow by the solid body. Along :math:`z`, the circulation will be proportional to :math:`\frac{d_y d_z}{d_x}v_x v_y
-- \frac{d_z d_x}{d_y}v_x v_y = 0` (due to :math:`d_x = d_y`). Therefore, on the plane where the solid is blunt, the motion
+deflection of the flow by the solid body. Along :math:`z`, the circulation will be proportional to :math:`\frac{r_y r_z}{r_x}v_x v_y
+- \frac{r_z r_x}{r_y}v_x v_y = 0` (due to :math:`r_x = r_y`). Therefore, on the plane where the solid is blunt, the motion
 produces no circulation.
 
 Now, for simplicity, let :math:`v_x = 0`. In this case also the circulation along :math:`y`, proportional
-to :math:`\frac{d_y d_z}{d_x}v_x v_z - \frac{d_y d_x}{d_y}v_x v_z`, is zero. The only non-zero component of the circulation
-will be along :math:`x` and be proportional to :math:`\left(\frac{d_x d_z}{d_y} - \frac{d_x d_y}{d_z}\right) v_y v_z \approx
-\frac{d_x^2}{d_z} v_y v_z`.
+to :math:`\frac{r_y r_z}{r_x}v_x v_z - \frac{r_y r_x}{r_y}v_x v_z`, is zero. The only non-zero component of the circulation
+will be along :math:`x` and be proportional to :math:`\left(\frac{r_x r_z}{r_y} - \frac{r_x r_y}{r_z}\right) v_y v_z \approx
+\frac{r_x^2}{r_z} v_y v_z`.
 
 We would have :math:`\mathbf{v}_\parallel = \{v_x, 0, v_z\}` and
-:math:`\Gamma \propto \{d_z v_y v_z, ~ 0,~ - d_x v_x v_y \} / \|\mathbf{v}\|`.
+:math:`\Gamma \propto \{r_z v_y v_z, ~ 0,~ - r_x v_x v_y \} / \|\mathbf{v}\|`.
 The motion produces no circulation on the plane where the solid is blunt, and on the other two planes
 the circulation is
 :math:`\Gamma \propto r_\Gamma ~ \|\mathbf{v}\|~ \sin(2 \alpha) ~ = ~2 r_\Gamma ~\|\mathbf{v}\| ~\sin(\alpha)~\cos(\alpha)`
 with :math:`\alpha` the angle between the velocity and its projection on the body on the plane (e.g. on the plane
 orthogonal to :math:`x` we have :math:`\sin(\alpha) = v_y/\|\mathbf{v}\|` and
-:math:`\cos(\alpha) = v_z/\|\mathbf{v}\|`), and :math:`r_\Gamma`, the lift surface on the plane (e.g. :math:`d_z` for
+:math:`\cos(\alpha) = v_z/\|\mathbf{v}\|`), and :math:`r_\Gamma`, the lift surface on the plane (e.g. :math:`r_z` for
 the plane orthogonal to :math:`x`). Furthermore, the direction of the circulation is given by the cross product (because
 the solid boundary "rotates" the incoming flow velocity towards its projection on the body).
 

doc/computation/index.rst

@@ -143,47 +143,59 @@ Our notation is summarized in the table below. Additional notation specific to c
 When available, we also show the fields of main data structures :ref:`mjModel` and :ref:`mjData` corresponding to the
 mathematical notation.
 
-+-----------------+----------------+----------------+----------------------+
-| Symbol          | Size           | Description    | MuJoCo field         |
-+=================+================+================+======================+
-| :math:`n_Q`     |                | number of      | ``mjModel.nq``       |
-|                 |                | position       |                      |
-|                 |                | coordinates    |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`n_V`     |                | number of      | ``mjModel.nv``       |
-|                 |                | degrees of     |                      |
-|                 |                | freedom        |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`n_C`     |                | number of      | ``mjData.nefc``      |
-|                 |                | active         |                      |
-|                 |                | constraints    |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`q`       | :math:`n_Q`    | joint position | ``mjData.qpos``      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`v`       | :math:`n_V`    | joint velocity | ``mjData.qvel``      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`\tau`    | :math:`n_V`    | applied force: |                      |
-|                 |                | passive,       |                      |
-|                 |                | actuation,     |                      |
-|                 |                | external       |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`c(q, v)` | :math:`n_V`    | bias force:    | ``mjData.qfrc_bias`` |
-|                 |                | Coriolis,      |                      |
-|                 |                | centrifugal,   |                      |
-|                 |                | gravitational  |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`M(q)`    | :math:`n_V     | inertia in     | ``mjData.qM``        |
-|                 | \times n_V`    | joint space    |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`J(q)`    | :math:`n_C     | constraint     | ``mjData.efc_J``     |
-|                 | \times n_V`    | Jacobian       |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`r(q)`    | :math:`n_C`    | constraint     | ``mjData.efc_pos``   |
-|                 |                | residual       |                      |
-+-----------------+----------------+----------------+----------------------+
-| :math:`f(q, v,  | :math:`n_C`    | constraint     | ``mjData.efc_force`` |
-| \tau)`          |                | force          |                      |
-+-----------------+----------------+----------------+----------------------+
+.. list-table::
+   :widths: 2 2 7 4
+   :header-rows: 1
+
+   * - Symbol
+     - Size
+     - Description
+     - MuJoCo field
+   * - :math:`n_Q`
+     -
+     - number of position coordinates
+     - ``mjModel.nq``
+   * - :math:`n_V`
+     -
+     - number of degrees of freedom
+     - ``mjModel.nv``
+   * - :math:`n_C`
+     -
+     - number of active constraints
+     - ``mjData.nefc``
+   * - :math:`q`
+     - :math:`n_Q`
+     - joint position
+     - ``mjData.qpos``
+   * - :math:`v`
+     - :math:`n_V`
+     - joint velocity
+     - ``mjData.qvel``
+   * - :math:`\tau`
+     - :math:`n_V`
+     - applied force: passive, actuation, external
+     - ``mjData.qfrc_passive`` + ``mjData.qfrc_actuator`` + ``mjData.qfrc_applied``
+   * - :math:`c(q, v)`
+     - :math:`n_V`
+     - bias force: Coriolis, centrifugal, gravitational
+     - ``mjData.qfrc_bias``
+   * - :math:`M(q)`
+     - :math:`n_V \times n_V`
+     - inertia in joint space
+     - ``mjData.qM``
+   * - :math:`J(q)`
+     - :math:`n_C \times n_V`
+     - constraint
+       Jacobian
+     - ``mjData.efc_J``
+   * - :math:`r(q)`
+     - :math:`n_C`
+     - constraint residual
+     - ``mjData.efc_pos``
+   * - :math:`f(q, v,\tau)`
+     - :math:`n_C`
+     - constraint force
+     - ``mjData.efc_force``
 
 All model elements are enumerated at compile time and assembled into the above system-level vectors and matrices. In our
 earlier arm model :ref:`example <Examples>` the model has :math:`n_V = 13` degrees of freedom: 3 for the ball joint, one
@@ -259,7 +271,7 @@ the force outputs are stored in ``mjData.actuator_force``, and the activation st
 
 These three components of an actuator - transmission, activation dynamics, and force generation - determine how the
 actuator works. The user can set them independently for maximum flexibility, or use :ref:`Actuator shortcuts
-<CActuator>` which instantiate common actuator types.
+<CActShortcuts>` which instantiate common actuator types.
 
 .. _geTransmission:
 
@@ -284,28 +296,31 @@ is attached; the possible attachment object types are :at:`joint`, :at:`tendon`,
 :at:`slider-crank`
    :at:`slider-crank` `transmissions <https://en.wikipedia.org/wiki/Slider-crank_linkage>`_ transform a linear force to
    a torque, as in a piston-driven combustion engine. `This model
-   <https://github.com/google-deepmind/mujoco/tree/main/model/slider_crank>`_ contains pedagogical examples.
-   Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality constraints to
-   the rest of the system, but that would be less efficient.
+   <https://github.com/google-deepmind/mujoco/blob/main/model/slider_crank/slider_crank.xml>`__ contains pedagogical
+   examples. Slider-cranks can also be modeled explicitly by creating MuJoCo bodies and coupling them with equality
+   constraints, but that is both less efficient and less stable.
+
+:at:`body`
+   :el:`body` transmission corresponds to applying forces at contact points belonging to a body, in
+   order to model vacuum grippers and biomechanical adhesive appendages. For more information about adhesion, see the
+   :ref:`adhesion<actuator-adhesion>` actuator documentation. These transmission targets have a fixed zero length
+   :math:`l_i(q) = 0`.
 
 :at:`site`
-   :at:`site` transmission (without a :at:`refsite`, see below) and :at:`body` transmission targets have a fixed zero
-   length :math:`l_i(q) = 0`. They can therefore not be used to maintain a desired length, but can be used to apply
-   forces. Site transmissions correspond to applying a Cartsian force/torque at the site, and are useful for modeling
-   jets and propellors. :el:`body` transmissions correspond to applying forces at contact points belonging to a body, in
-   order to model vacuum grippers and biomechanical adhesive appendages. For more information about adhesion, see the
-   :ref:`adhesion<actuator-adhesion>` actuator documentation.
+   Site transmissions correspond to applying a Cartsian force/torque in the frame of a site. When a :at:`refsite` is not
+   defined (see below), these targets have a fixed zero length :math:`l_i(q) = 0` and are useful for modeling jets and
+   propellors: forces and torques which are fixed to the site frame.
 
-   If a :at:`site` transmission target is defined with the optional :at:`refsite` attribute, forces and torques are
-   applied in the frame of the reference site rather than the site's own frame. If a reference site is defined then
-   the length of the actuator is nonzero and corresponds to the pose difference of the two sites. This length can then
-   be controlled with a :el:`position` actuator, enabling Cartesian end-effector control. See the
-   :ref:`refsite<actuator-general-refsite>` documentation for more details.
+   If a :at:`site` transmission is defined with the optional :at:`refsite` attribute, forces and torques are applied in
+   the frame of the reference site rather than the site's own frame. If a reference site is defined, the length of the
+   actuator is nonzero and corresponds to the pose difference of the two sites, projected onto a chosen direction in the
+   reference frame. This length can then be controlled with a :el:`position` actuator, allowing for Cartesian
+   end-effector control. See the :ref:`refsite<actuator-general-refsite>` documentation for more details.
 
 .. _geActivation:
 
-Activation dynamics
-^^^^^^^^^^^^^^^^^^^
+Stateful actuators
+^^^^^^^^^^^^^^^^^^
 
 Some actuators such as pneumatic and hydraulic cylinders as well as biological muscles have an internal state called
 "activation". This is a true dynamic state, beyond the joint positions :math:`q` and velocities :math:`v`. Including
@@ -321,27 +336,33 @@ independent of the other actuators. The activation types currently implemented a
 .. math::
    \begin{aligned}
    \text{integrator}:  & & \dot{w}_i &= u_i \\
-   \text{filter}:      & & \dot{w}_i &= (u_i - w_i) / t \\
-   \text{filterexact}: & & \dot{w}_i &= (u_i - w_i) / t \\
+   \text{filter}:      & & \dot{w}_i &= (u_i - w_i) / \texttt{t} \\
+   \text{filterexact}: & & \dot{w}_i &= (u_i - w_i) / \texttt{t} \\
+   \text{muscle}:      & & \dot{w}_i &= \textrm{muscle}(u_i, w_i, l_i, \dot{l}_i)
    \end{aligned}
 
-where :math:`t` is an actuator-specific time constant stored in ``mjModel.actuator_dynprm``. In addition the type can
-be "user", in which case :math:`w_i` is computed by the user-defined callback :ref:`mjcb_act_dyn`. The type can also
-be "none" which corresponds to a regular actuator with no activation state. The dimensionality of :math:`w` equals
+where :math:`\texttt{t}` is an actuator-specific time-constant stored in ``mjModel.actuator_dynprm``. In addition, the
+type can be "user", in which case :math:`w_i` is computed by the user-defined callback :ref:`mjcb_act_dyn`. The type can
+also be "none" which corresponds to a regular actuator with no activation state. The dimensionality of :math:`w` equals
 the number of actuators whose activation type is different from "none".
 
+For more information regarding muscle activation dynamics, see :ref:`CMuscle`.
+
 For ``filterexact`` activation dynamics, Euler integration of :math:`\dot{w}` is replaced with the analytic integral:
 
 .. math::
    \begin{aligned}
-   \text{filter}:      & & w_{i+1} &= w_i + h (u_i - w_i) / t \\
-   \text{filterexact}: & & w_{i+1} &= w_i + (u_i - w_i) (1 - e^{-h / t}) \\
+   \text{filter}:      & & w_{i+1} &= w_i + h (u_i - w_i) / \texttt{t} \\
+   \text{filterexact}: & & w_{i+1} &= w_i + (u_i - w_i) (1 - e^{-h / \texttt{t}}) \\
    \end{aligned}
 
-The two expressions converge to the same value in the :math:`h \rightarrow 0` limit.
+The two expressions converge to the same value in the :math:`h \rightarrow 0` limit. Note that Euler-integrated filters
+diverge for :math:`\texttt{t} < h`, while exactly-integrated filters are stable for any positive :math:`\texttt{t}`.
 
-Note that Euler-integrated filters diverge for :math:`t < h`, while exactly-integrated filters are stable for any
-:math:`t > 0`.
+:ref:`actearly<actuator-general-actearly>`:
+  If the :ref:`actearly<actuator-general-actearly>` attribute is set to "true", ``mjData.actuator_force`` is computed
+  based on :math:`w_{i+1}` (the next activation), reducing the delay between changes to :math:`u` and their effects on
+  the acceleration by one time step (so the total dynamics are second-order rather than third order).
 
 .. _geActuatorForce:
 
@@ -379,10 +400,6 @@ This quantity is stored in ``mjData.qfrc_actuator``. It is added to the applied
 with any user-defined forces in joint or Cartesian coordinates (which are stored in ``mjData.qfrc_applied`` and
 ``mjData.xfrc_applied`` respectively).
 
-Optionally, the :ref:`actearly<actuator-general-actearly>` attribute on an actuator computes ``mjData.qfrc_actuator``
-based on the value of :math:`w_{i+1}` after integration, reducing the delay between changes to :math:`u` and
-:math:`t`.
-
 .. _gePassive:
 
 Passive forces
@@ -526,29 +543,48 @@ Fast implicit-in-velocity (``implicitfast``)
    derivatives are also the main source of asymmetry of :math:`D`, by dropping them and symmetrizing, we can use the
    faster Cholesky rather than LU decomposition.
 
-   .. tip::
-      The implicitfast integrator has similar computational cost to Euler, yet provides increased stability, and is
-      therefore a strict improvement. It is the recommended integrator and will become the default in a future version.
-
 4th-order Runge-Kutta (``RK4``)
    One advantage of our continuous-time formulation is that we can use higher order integrators such as Runge-Kutta or
    multistep methods. The only such integrator currently implemented is the fixed-step `4th-order Runge-Kutta method
-   <https://en.wikipedia.org/wiki/Runge–Kutta_methods#Derivation_of_the_Runge–Kutta_fourth-order_method>`_, though users
-   can easily implement other integrators by calling :ref:`mj_forward` and integrating accelerations themselves. We have
-   observed that for energy-conserving systems (`example
-   <https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/derivative/energy_conserving_pendulum.xml>`_) RK4
-   is qualitatively better than the single-step methods, both in terms of stability and accuracy, even when the timestep
-   is decreased by a factor of 4 (so the computational effort is identical).  In the presence of large velocity-
-   dependent forces, if the chosen single-step method integrates those forces implicitly, single-step methods can be
-   significantly more stable than RK4.
+   <https://en.wikipedia.org/wiki/Runge–Kutta_methods#Derivation_of_the_Runge–Kutta_fourth-order_method>`__, though
+   users can easily implement other integrators by calling :ref:`mj_forward` and integrating accelerations themselves.
+   We have observed that for energy-conserving systems (`example <../_static/pendulum.xml>`__), RK4 is qualitatively
+   better than the single-step methods, both in terms of stability and accuracy, even when the timestep is decreased by
+   a factor of 4 (so the computational effort is identical). In the presence of large velocity- dependent forces, if the
+   chosen single-step method integrates those forces implicitly, single-step methods can be significantly more stable
+   than RK4.
 
-.. note::
-   The accuracy and stability of all integrators can be improved by reducing the time step :math:`h` which is stored in
-   ``mjModel.opt.timestep``. Of course this also slows down the simulation. The time step is perhaps the most important
-   parameter that the user can adjust. If it is too large, the simulation will become unstable. If it is too small, CPU
-   time will be wasted without meaningful improvement in accuracy. There is always a comfortable range where the time
-   step is "just right", but that range is model-dependent.
+.. admonition:: Choosing timestep and integrator
+   :class: tip
 
+   :ref:`timestep<option-timestep>`
+    The accuracy and stability of all integrators can be improved by reducing the time step :math:`h`.
+    Of course a smaller time step also slows down the simulation. The time step is perhaps the single most important
+    parameter that the user can adjust. If it is too large, the simulation will become unstable. If it is too small, CPU
+    time will be wasted without meaningful improvement in accuracy. There is always a comfortable range where the time
+    step is "just right", but that range is model-dependent.
+
+   :ref:`integrator<option-integrator>`
+    Summary: The recommended integrator is ``implicitfast`` which usually has the best tradeoff of stabillity and
+    performance.
+
+    **Euler**:
+     Use ``Euler`` for compatibillity with older models and :ref:`MJX<Mjx>`. Specifically for MJX,
+     setting the :ref:`eulerdamp<option-flag-eulerdamp>` disable flag can :ref:`improve performance<MjxPerformance>`.
+    **implicitfast**:
+     The ``implicitfast`` integrator has similar computational cost to ``Euler``, yet provides
+     increased stability, and is therefore a strict improvement. It is the recommended integrator for most models.
+    **implicit**:
+     The benefit over ``implicitfast`` is the implicit integration of Coriolis and centripetal forces, including
+     gyroscopic forces. The most common case where integrating such forces implicitly leads to noticable improvement is
+     when free objects with assymetric inertia are spinning quickly. `gyroscopic.xml <../_static/gyroscopic.xml>`__
+     shows an ellipsoid rolling on an inclined plane which quickly diverges with ``implicitfast`` but is stable with
+     ``implicit``.
+    **RK4**:
+     This integrator is best for systems which are energy conserving, or almost energy-conserving. `pendulum.xml
+     <../_static/pendulum.xml>`__ shows a complicated pendulum mechanism which diverges quickly using ``Euler`` or
+     ``implicitfast`` yet conserves energy well under ``RK4``. Note that under ``implicit``, this model doesn't diverge
+     but rather loses energy.
 
 .. _geState:
 
@@ -623,6 +659,11 @@ MoCap poses: ``mocap_pos`` and ``mocap_quat``
   6D poses from a motion-capture device. The default values set by :ref:`mj_resetData` are the poses of the bodies at
   the default configuration.
 
+Equality constraint toggle: ``eq_active``
+  ``mjData.eq_active`` is a byte-valued array that allows the user to toggle the state of equality constraints at
+  runtime. The initial value of this array is ``mjModel.eq_active0`` which can be set in XML using the
+  :ref:`active<equality-connect-active>` attribute of :ref:`equality constraints<coEquality>`.
+
 User data: ``userdata``
   ``mjData.userdata`` acts as a user-defined memory space untouched by the engine. For example it can be used by
   callbacks. This is described in more detail in the :ref:`Programming chapter<siSimulation>`.
@@ -850,33 +891,25 @@ spatial frame and normal distance are given by the collision detector.
 In addition to the above quantities which are computed online, each contact has several parameters obtained from the
 model definition.
 
-+---------------------------+-----------------------------------+
-| Parameter                 | Description                       |
-+===========================+===================================+
-| ``condim``                | Dimensionality of the contact     |
-|                           | force/torque in the contact       |
-|                           | frame. It can be 1, 3, 4 or 6.    |
-+---------------------------+-----------------------------------+
-| ``friction``              | Vector of friction coefficients,  |
-|                           | with dimensionality ``condim-1``. |
-+---------------------------+-----------------------------------+
-| ``margin``                | The distance margin used to       |
-|                           | determine if the contact should   |
-|                           | be included in the global contact |
-|                           | array ``mjData.contact``.         |
-+---------------------------+-----------------------------------+
-| ``gap``                   | For custom computations it is     |
-|                           | sometimes convenient to include   |
-|                           | contacts in ``mjData.contact``    |
-|                           | but not generate contact forces.  |
-|                           | This is what ``gap`` does:        |
-|                           | contact forces are generated only |
-|                           | when the normal distance is below |
-|                           | margin-gap.                       |
-+---------------------------+-----------------------------------+
-| ``solref`` and ``solimp`` | :ref:`Solver <Solver>`            |
-|                           | parameters explained later.       |
-+---------------------------+-----------------------------------+
+.. list-table::
+   :widths: 1 5
+   :header-rows: 1
+
+   * - Parameter
+     - Description
+   * - ``condim``
+     - Dimensionality of the contact force/torque in the contact frame. |br| It can be 1, 3, 4 or 6.
+   * - ``friction``
+     - Vector of friction coefficients with dimensionality ``condim-1``.
+   * - ``margin``
+     - The distance margin used to determine if the contact should be included in the global contact array
+       ``mjData.contact``.
+   * - ``gap``
+     - For custom computations it is sometimes convenient to include contacts in ``mjData.contact`` but not generate
+       contact forces. This is what ``gap`` does: contact forces are generated only when the normal distance is below
+       (margin - gap).
+   * - ``solref`` and ``solimp``
+     - :ref:`Solver <Solver>` parameters, explained later.
 
 The contact friction cone can be either elliptic or pyramidal. This is a global setting determined by the choice of
 constraint solver: the elliptic solvers work with elliptic cones, while the pyramidal solvers work with pyramidal cones,
@@ -933,6 +966,12 @@ is :math:`E f`. The matrix of basis vectors is constructed as follows.
 .. image:: ../images/computation/contact_frame.svg
    :width: 700px
    :align: center
+   :class: only-light
+
+.. image:: ../images/computation/contact_frame_dark.svg
+   :width: 700px
+   :align: center
+   :class: only-dark
 
 The figure illustrates the full basis set corresponding to the case :math:`n = 6`. Otherwise we use only the first
 :math:`n` or :math:`2(n-1)` columns depending on the cone type. Elliptic cones are easier to understand. Since the
@@ -967,11 +1006,11 @@ to be solved numerically. In inverse dynamics, the problem becomes diagonal and
 
 The primal formulation is based on a generalization of the Gauss principle of least constraint. In its basic form, the
 Gauss principle states that if we have unconstrained dynamics :math:`M \dot{v} = \tau` and impose acceleration
-constraint :math:`J \dot{v} = a^*`, the resulting acceleration will be
+constraint :math:`J \dot{v} = \ar`, the resulting acceleration will be
 
 .. math::
    \dot{v} = \arg \min_x \left\| x-M^{-1} \tau \right\|^2_M \\
-   \textrm{subject to} \; J x = a^*
+   \textrm{subject to} \; J x = \ar
 
 where the weighted :math:`L_2` norm is the usual :math:`\|x\|^2_M = x^T M x`. Thus the constraint causes the smallest
 possible deviation from the unconstrained acceleration :math:`M^{-1}\tau`, where the metric for measuring deviations in
@@ -981,58 +1020,55 @@ will be done by generalizing both the cost function and the constraints in the G
 
 We will use the following notation beyond the notation introduced earlier:
 
-+----------------------+----------------------+----------------------+
-| Symbol               | Size                 | Description          |
-+======================+======================+======================+
-| :math:`z`            | :math:`n_C`          | constraint           |
-|                      |                      | deformations         |
-+----------------------+----------------------+----------------------+
-| :math:`\omega`       | :math:`n_C`          | velocity of          |
-|                      |                      | constraint           |
-|                      |                      | deformations         |
-+----------------------+----------------------+----------------------+
-| :math:`d`            | :math:`n_C`          | constraint impedance |
-+----------------------+----------------------+----------------------+
-| :math:`b`            | :math:`n_C`          | virtual constraint   |
-|                      |                      | damping              |
-+----------------------+----------------------+----------------------+
-| :math:`k`            | :math:`n_C`          | virtual constraint   |
-|                      |                      | stiffness            |
-+----------------------+----------------------+----------------------+
-| :math:`A(q)`         | :math:`n_C \times    | inverse inertia in   |
-|                      | n_C`                 | constraint space     |
-+----------------------+----------------------+----------------------+
-| :math:`R(q)`         | :math:`n_C \times    | diagonal regularizer |
-|                      | n_C`                 | in constraint space  |
-+----------------------+----------------------+----------------------+
-| :math:`a^*(q,v)`     | :math:`n_C`          | reference            |
-|                      |                      | acceleration in      |
-|                      |                      | constraint space     |
-+----------------------+----------------------+----------------------+
-| :math:`a^0(q, v,     | :math:`n_C`          | unconstrained        |
-| \tau)`               |                      | acceleration in      |
-|                      |                      | constraint space     |
-+----------------------+----------------------+----------------------+
-| :math:`a^1(q, v,     | :math:`n_C`          | constrained          |
-| \dot{v})`            |                      | acceleration in      |
-|                      |                      | constraint space     |
-+----------------------+----------------------+----------------------+
-| :math:`\mathcal{K}   |                      | product of all       |
-| (q)`                 |                      | contact friction     |
-|                      |                      | cones                |
-+----------------------+----------------------+----------------------+
-| :math:`\eta`         |                      | upper bounds on      |
-|                      |                      | friction loss forces |
-+----------------------+----------------------+----------------------+
-| :math:`\Omega(q)`    |                      | convex set of        |
-|                      |                      | admissible           |
-|                      |                      | constraint forces    |
-+----------------------+----------------------+----------------------+
-| :math:`\mathcal{E},  |                      | index sets for       |
-| \mathcal{F},         |                      | Equality, Friction   |
-| \mathcal{C}`         |                      | loss, Contact        |
-|                      |                      | constraints          |
-+----------------------+----------------------+----------------------+
+.. list-table::
+   :widths: 1 1 4
+   :header-rows: 1
+
+   * - Symbol
+     - Size
+     - Description
+   * - :math:`z`
+     - :math:`n_C`
+     - constraint deformations
+   * - :math:`\omega`
+     - :math:`n_C`
+     - velocity of constraint deformations
+   * - :math:`k`
+     - :math:`n_C`
+     - virtual constraint stiffness
+   * - :math:`b`
+     - :math:`n_C`
+     - virtual constraint damping
+   * - :math:`d`
+     - :math:`n_C`
+     - constraint impedance
+   * - :math:`A(q)`
+     - :math:`n_C \times n_C`
+     - inverse inertia in constraint space
+   * - :math:`R(q)`
+     - :math:`n_C \times n_C`
+     - diagonal regularizer in constraint space
+   * - :math:`\ar`
+     - :math:`n_C`
+     - reference acceleration in constraint space
+   * - :math:`\au(q, v, \tau)`
+     - :math:`n_C`
+     - unconstrained acceleration in constraint space
+   * - :math:`\ac(q, v, \dot{v})`
+     - :math:`n_C`
+     - constrained acceleration in constraint space
+   * - :math:`\mathcal{K}(q)`
+     -
+     - product of all contact friction cones
+   * - :math:`\eta`
+     -
+     - upper bounds on friction loss forces
+   * - :math:`\Omega(q)`
+     -
+     - convex set of admissible constraint forces
+   * - :math:`\mathcal{E}, \mathcal{F}, \mathcal{C}`
+     -
+     - index sets for Equality, Friction loss, Contact constraints
 
 The index sets will be used to refer to parts of vectors and matrices. For example, :math:`J_\mathcal{C}` is the
 sub-matrix of all rows of the Jacobian that correspond to contact constraints.
@@ -1048,7 +1084,7 @@ explain what it means and why it makes sense. That problem is
 .. math::
    (\dot{v}, \dot{\omega}) = \arg \min_{(x, y)}
                              \left\|x-M^{-1}(\tau-c)\right\|^2_M +
-                             \left\|y-a^*\right\|^{\text{Huber}(\eta)}_{R^{-1}} \\
+                             \left\|y-\ar\right\|^{\text{Huber}(\eta)}_{R^{-1}} \\
    \textrm{subject to} \;
    J_\mathcal{E} x_\mathcal{E} - y_\mathcal{E} = 0, \;
    J_\mathcal{F} x_\mathcal{F} - y_\mathcal{F} = 0, \;
@@ -1056,15 +1092,15 @@ explain what it means and why it makes sense. That problem is
    :label: eq:primal
 
 The new players here are the diagonal regularizer :math:`R > 0` which makes the constraints soft, and the reference
-acceleration :math:`a^*` which stabilizes the constraints. The latter is similar in spirit to Baumgarte stabilization,
+acceleration :math:`\ar` which stabilizes the constraints. The latter is similar in spirit to Baumgarte stabilization,
 but instead of adding a constraint force directly, it modifies the optimization problem whose solution is the constraint
-force. Since this problem is itself constrained, the relation between :math:`a^*` and :math:`f` is generally non-linear.
-The quantities :math:`R` and :math:`a^*` are computed from the solver :ref:`parameters <soParameters>` as described
+force. Since this problem is itself constrained, the relation between :math:`\ar` and :math:`f` is generally non-linear.
+The quantities :math:`R` and :math:`\ar` are computed from the solver :ref:`parameters <soParameters>` as described
 later. For now we assume they are given.
 
 The optimization variable :math:`x` stands for acceleration as in the Gauss principle, while :math:`y` is a slack
 variable in constraint space. It is needed to model soft constraints. If we forced the solution to reach
-:math:`y = a^*`, which we could do by taking the limit :math:`R \to 0`, we would obtain a hard constraint model. This
+:math:`y = \ar`, which we could do by taking the limit :math:`R \to 0`, we would obtain a hard constraint model. This
 limit is not allowed in MuJoCo, but nevertheless one can construct models that are phenomenologically hard.
 
 The symbol :math:`\mathcal{K}^*` denotes the dual to the friction cone. It is motivated by mathematical reverse
@@ -1095,7 +1131,7 @@ constraint space (which the constraint force aims to prevent), there is no posit
    \tilde{q} &= {q \brack z}, &
    \tilde{v} &= {v \brack \omega}, &
    \tilde{c} &= {c \brack 0}, \\
-   \tilde{\tau} &= {\tau \brack {R^{-1} a^*}}, &
+   \tilde{\tau} &= {\tau \brack {R^{-1} \ar}}, &
    \tilde{M} &= \left[\begin{array}{cc}
                       M & 0 \\
                       0 & R^{-1}
@@ -1114,20 +1150,20 @@ Unpacking all the tildes yields the explicit form of the original and the deform
 .. math::
    \begin{aligned}
    M \dot{v} + c &= \tau +J^T f \\
-   \dot{\omega} &= a^* - R f \\
+   \dot{\omega} &= \ar - R f \\
    \end{aligned}
 
-Thus :math:`R` has the meaning of inverse deformation inertia, while :math:`a^*` has the meaning of unforced deformation
+Thus :math:`R` has the meaning of inverse deformation inertia, while :math:`\ar` has the meaning of unforced deformation
 acceleration.
 
 Does MuJoCo keep these deformation variables as part of the system state and integrate their dynamics together with the
 joint positions and velocities? No, although such an option may be worth providing in the future. Recall that we defined
-the functional dependence of the regularizer and the reference acceleration as :math:`R(q)` and :math:`a^*(q, v)`. This
+the functional dependence of the regularizer and the reference acceleration as :math:`R(q)` and :math:`\ar(q, v)`. This
 makes problem :eq:`eq:primal` dependent only on :math:`(q, v, \tau)`, and so the original dynamics are not actually
 affected by the deformation dynamics. Since the general constraint model we developed up to now makes no assumptions
-about how :math:`R` and :math:`a^*` are computed, our choice is consistent and improves simulator efficiency.
+about how :math:`R` and :math:`\ar` are computed, our choice is consistent and improves simulator efficiency.
 Nevertheless, given that these quantities turned out to be related to the deformation dynamics, it may be more natural
-to define them as :math:`R(z)` and :math:`a^* (z, \omega)` and simulate the entire augmented system. Below we clarify
+to define them as :math:`R(z)` and :math:`\ar (z, \omega)` and simulate the entire augmented system. Below we clarify
 some of the benefits of such a simulation.
 
 When do the deformation dynamics "track" the original dynamics exactly? One can verify that this happens when the
@@ -1159,7 +1195,7 @@ we obtain the unconstrained problem
 
 .. math::
    \dot{v} = \arg \min_{x} \left\|x-M^{-1}(\tau-c)\right\|^2_M +
-             s \left( J x - a^* \right)
+             s \left( J x - \ar \right)
    :label: eq:reduced
 
 The function :math:`s(\cdot)` plays the role of a soft-constraint penalty. It can be shown to be convex and
@@ -1169,7 +1205,7 @@ Another appealing feature of the reduced formulation is that the inverse dynamic
 above problem is unconstrained and convex, the unique global minimum makes the gradient vanish. This yields the identity
 
 .. math::
-   M \dot{v} + c = \tau - J^T \nabla s \left( J \dot{v} - a^* \right)
+   M \dot{v} + c = \tau - J^T \nabla s \left( J \dot{v} - \ar \right)
 
 which is the analytical inverse dynamics in the presence of soft constraints. Comparing to the equations of motion
 :eq:`eq:motion`, we see that the constraint forces :math:`f` are given by the negative gradient of the function
@@ -1191,7 +1227,7 @@ Lagrange dual to the primal problem defined above is
 
 .. math::
    f = \arg\min_\lambda \frac{1}{2} \lambda^{T} \left( A+R \right) \lambda +
-       \lambda^T \left( a^0 - a^* \right) \\
+       \lambda^T \left( \au - \ar \right) \\
    \text{subject to} \; \lambda \in \Omega
    :label: eq:dual
 
@@ -1203,7 +1239,7 @@ where the inverse inertia in constraint space is
 and the unconstrained acceleration in constraint space is
 
 .. math::
-   a^0 = J M^{-1} (\tau-c) + \dot{J} v
+   \au = J M^{-1} (\tau-c) + \dot{J} v
 
 The constraint set :math:`\Omega` is as follows. :math:`\lambda_\mathcal{E}` is unconstrained, because it is the
 Lagrange multiplier for an equality constraint in the primal problem. For friction loss we have the box constraint
@@ -1218,28 +1254,28 @@ problem are described later.
 As mentioned earlier, MuJoCo's constraint model has uniquely-defined inverse dynamics, and we already saw one way to
 derive it in the reduced formulation above. Here we derive it again from the dual formulation. Recall that in inverse
 dynamics we have access to :math:`(q, v, \dot{v})` instead of :math:`(q, v, \tau)`, so the unconstrained acceleration
-:math:`a^0` is unknown. However we can compute the constrained acceleration
+:math:`\au` is unknown. However we can compute the constrained acceleration
 
 .. math::
-   a^1 = J \dot{v} + \dot{J} v
+   \ac = J \dot{v} + \dot{J} v
 
 Inverse dynamics can now be computed by solving the optimization problem
 
 .. math::
    f = \arg \min_\lambda \frac{1}{2} \lambda^{T} R \lambda +
-       \lambda^T \left( a^1 - a^* \right) \\
+       \lambda^T \left( \ac - \ar \right) \\
    \text{subject to} \; \lambda \in \Omega
 
 By comparing the KKT conditions for these two convex optimization problems, one can verify that their solutions coincide
 when
 
 .. math::
-   a^1 = a^0 + Af
+   \ac = \au + Af
    :label: eq:identity
 
 This key identity is essentially Newton's second law projected in constraint space. It is derived by moving the term
 :math:`c` in the equations of motion :eq:`eq:motion` to the right hand side, multiplying by :math:`J M^{-1}` from the
-left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, a^0, a^1`. In terms of
+left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, \au, \ac`. In terms of
 implementation, we do not actually compute the acceleration term :math:`\dot{J} v`. This is because our optimization
 problems depend on differences of constraint-space accelerations, and so this term would cancel out even if we were to
 compute it.
@@ -1295,6 +1331,12 @@ representations of the constraint Jacobian and related matrices.
    .. image:: ../images/computation/gPGS.svg
       :width: 500px
       :align: center
+      :class: only-light
+
+   .. image:: ../images/computation/gPGS_dark.svg
+      :width: 500px
+      :align: center
+      :class: only-dark
 
    When using pyramidal friction cones, the problem involves box constraints to which PGS has traditionally been
    applied. If we applied PGS directly to the conic constraints resulting from elliptic friction cones, it would get
@@ -1316,12 +1358,12 @@ representations of the constraint Jacobian and related matrices.
 Parameters
 ~~~~~~~~~~
 
-Here we explain how the quantities :math:`R, a^*` are computed from model parameters. For the chosen parameterization to
-make sense, we first need to understand how these quantities affect the dynamics. We focus on the unconstrained
-minimizer of :eq:`eq:dual`, namely
+Here we explain how the quantities :math:`R, \ar` are computed from model parameters. For the chosen
+parameterization to make sense, we first need to understand how these quantities affect the dynamics. We focus on the
+unconstrained minimizer of :eq:`eq:dual`, namely
 
 .. math::
-   f^+ = (A+R)^{-1} (a^* - a^0)
+   f^+ = (A+R)^{-1} (\ar - \au)
 
 If it happens that :math:`f^+ \in \Omega`, then :math:`f^+ = f` is the actual constraint force generated by our model.
 We focus on this case because it is common, in the sense that the subset of the constraints in :math:`\Omega` that are
@@ -1329,11 +1371,11 @@ active at any given time is usually small, and furthermore it is the only case t
 Substituting :math:`f^+` in the constraint dynamics :eq:`eq:identity` and rearranging terms yields
 
 .. math::
-   a^1 = A(A+R)^{-1} a^* + R (A+R)^{-1} a^0
+   \ac = A(A+R)^{-1} \ar + R (A+R)^{-1} \au
 
 Thus the constrained acceleration interpolates between the unconstrained and the reference acceleration. In particular,
-in the limit :math:`R \to 0` we have a hard constraint and :math:`a^1 = a^*`, while in the limit :math:`R \to \infty` we
-have have an infinitely soft constraint (i.e., no constraint) and :math:`a^1 = a^0`. It is then natural to introduce a
+in the limit :math:`R \to 0` we have a hard constraint and :math:`\ac = \ar`, while in the limit :math:`R \to \infty` we
+have have an infinitely soft constraint (i.e., no constraint) and :math:`\ac = \au`. It is then natural to introduce a
 model parameter which directly controls the interpolation. We call this parameter *impedance* and denote it :math:`d`.
 It is a vector with dimensionality :math:`n_C` satisfying :math:`0<d<1` element-wise. Once it is specified, we compute
 the diagonal elements of the regularizer as
@@ -1349,7 +1391,7 @@ approximation happened to be exact, and :math:`A` itself happened to be diagonal
 constraint would satisfy
 
 .. math::
-   a^1_i = d_i a^*_i + (1-d_i) a^0_i
+   \aci = d_i \ari + (1-d_i) \aui
 
 and so we would achieve the desired interpolation effect. This of course does not hold exactly in general, but the goal
 here is to construct a sensible and intuitive parameterization of the constraint model and get the scaling right.
@@ -1358,14 +1400,14 @@ Next we explain how the reference acceleration is computed. As already mentioned
 parameterized by *damping* and *stiffness* coefficients element-wise:
 
 .. math::
-   a^*_i = -b_i (J v)_i - k_i r_i
+   \ari = -b_i (J v)_i - k_i r_i
 
 Recall that :math:`r` is the position residual (which is zero for friction loss and friction dimensions of elliptic
 cones), while :math:`J v` is the joint velocity projected in constraint space; the indexing notation refers to one
 component of the projected velocity vector.
 
 To summarize, the user specifies the vectors of impedance coefficients :math:`0<d<1`, damping coefficients :math:`b > 0`
-and stiffness coefficients :math:`k > 0`. The quantities :math:`R, a^*` are then computed by MuJoCo as shown above, and
+and stiffness coefficients :math:`k > 0`. The quantities :math:`R, \ar` are then computed by MuJoCo as shown above, and
 the selected optimization algorithm is applied to solve problem :eq:`eq:dual`. As explained in the :ref:`solver
 parameters <CSolver>` section of the Modeling chapter, MuJoCo offers additional automation for setting :math:`d, b, k`
 so as to achieve critical damping, or model a soft contact layer by varying :math:`d` with distance.
@@ -1394,12 +1436,18 @@ approximations, no matter how accurate the approximation is. The figure below il
 where the pyramid is not even an approximation, but represents the same constraint set as the elliptic cone. We plot the
 contours of the penalty/shadow for the pyramidal (red) and elliptic (dashed blue) cones, for different friction
 coefficients varying from left to right. Mathematically, the penalty in the pyramidal case is a quadratic spline, while
-the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics - allowing
+the penalty in the elliptic case contains pieces that are quadratics minus square roots of quadratics -- allowing
 circular contours around the tip of the cone.
 
 .. image:: ../images/computation/softcontact.png
    :width: 600px
    :align: center
+   :class: only-light
+
+.. image:: ../images/computation/softcontact_dark.png
+   :width: 600px
+   :align: center
+   :class: only-dark
 
 In summary, elliptic and pyramidal friction cones define different soft-contact dynamics (although they are usually very
 close). The elliptic model is more principled and more consistent with physical intuition, and the corresponding solvers
@@ -1428,48 +1476,45 @@ others can be pruned quickly without a detailed check. MuJoCo has flexible mecha
 checked in detail. The decision process involves two stages: generation and filtering.
 
 Generation
-   First we generate a list of candidate geom pairs in one of two ways: "pair" or "dynamic". The user can also specify
-   "all" which merges both sources (and is the default). This is done via the setting ``mjModel.opt.collision``. "Pair"
-   refers to an explicit list of geom pairs defined with the :ref:`pair <contact-pair>` element in MJCF. It gives the
-   user full control, however it is a static mechanism (independent of the spatial arrangement of the geoms at runtime)
-   and can be tedious for large models. It is normally used to supplement the output of the "dynamic" mechanism. Dynamic
-   generation works with bodies rather than geoms; when a body pair is included this means that all geoms attached to
-   one body can collide with all geoms attached to the other body.
+   First we generate a list of candidate geom pairs by merging from two sources: pairs of bodies that might contain
+   colliding geoms and the explicit list of geom pairs defined with the :ref:`pair <contact-pair>` element in MJCF.
 
    The body pairs are generated via broad-phase collision detection based on a modified sweep-and-prune algorithm. The
    modification is that the axis for sorting is chosen as the principal eigenvector of the covariance matrix of all geom
-   centers - which maximizes the spread. Then, for each body pair, a mid-phase collision detection using a static
-   bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB) is performed. Each body is
-   equipped with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes,
+   centers -- which maximizes the spread. Then, for each body pair, mid-phase collision detection is performed using a
+   static bounding volume hierarchy (a BVH binary tree) of axis-aligned bounding boxes (AABB). Each body is equipped
+   with an AABB tree of its geoms, aligned with the body inertial or geom frames for all inner or leaf nodes,
    respectively.
 
-   Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude <contact-exclude>` element
-   in MJCF. Exclusion is applied when "dynamic" or "all" are selected, but not when "pair" is selected. At the end of
-   this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`, but can still be
-   pruned further before detailed collision checking.
+   Finally, the user can explicitly exclude certain body pairs using the :ref:`exclude <contact-exclude>` element in
+   MJCF. At the end of this step we have a list of geoms pairs that is typically much smaller than :math:`n (n-1)/2`,
+   but can still be pruned further before detailed collision checking.
 
 Filtering
    Next we apply four filters to the list generated in the previous step. Filters 1 and 2 are applied to all geom pairs.
-   Filters 3 and 4 are applied only to pairs generated by the "dynamic" mechanism, thereby allowing the user to bypass
+   Filters 3 and 4 are applied only to pairs generated by the body-pair mechanism, thereby allowing the user to bypass
    those filters by specifying geom pairs explicitly.
 
-   #. The types of the two geoms must correspond to a collision function that is capable of performing the detailed
+   1. The types of the two geoms must correspond to a collision function that is capable of performing the detailed
       check. This is usually the case but there are exceptions (for example plane-plane collisions are not supported),
       and furthermore the user may override the default table of collision functions with NULL pointers, effectively
       disabling collisions between certain geom types.
-   #. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a
+   2. A bounding sphere test is applied, taking into account the contact margin. If one of the geoms in the pair is a
       plane, this becomes a plane-sphere test.
-   #. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless
+   3. The two geoms cannot belong to the same body. Furthermore, they cannot belong to a parent and a child body, unless
       the parent is the world body. The motivation is to avoid permanent contacts within bodies and joints. Note that if
       several bodies are welded together in the sense that there are no joints between them, they are treated as a
       single body for the purposes of this test. The parent-filter test can be disabled by the user, while the same-body
       test cannot be disabled.
-   #. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and
+   4. The two geoms must be "compatible" in the following sense. Each geom has integer parameters ``contype`` and
       ``conaffinity``. The boolean expression below must be true for the test to pass:
-      ``(contype1 & conaffinity2) || (contype2 & conaffinity1)`` This requires the ``contype`` of one geom and the
-      ``conaffinity`` of the other geom to have a common bit set to 1. This is a powerful mechanism borrowed from the
-      Open Dynamics Engine. The default setting for all geoms is ``contype = conaffinity = 1`` which always passes the
-      test, so the user can ignore this mechanism if it is confusing at first.
+
+      ``(contype1 & conaffinity2) || (contype2 & conaffinity1)``
+
+      This requires the ``contype`` of one geom and the ``conaffinity`` of the other geom to have a common bit set to 1.
+      This is a powerful mechanism borrowed from Open Dynamics Engine. The default setting for all geoms is
+      ``contype = conaffinity = 1`` which always passes the test, so the user can ignore this mechanism if it is
+      confusing at first.
 
 .. _coChecking:
 
@@ -1490,12 +1535,12 @@ convex hull implicitly, however pre-computing that hull can substantially improv
 model compiler does that by default, using the `qhull <http://www.qhull.org/>`__ library.
 
 In order to model a non-convex object other than a height field, the user must decompose it into a union of convex geoms
-(which can be primitive shapes or meshes) and attach them to the same body. Tools such as the
-`HACD <https://github.com/kmammou/v-hacd>`__ library can be used outside MuJoCo to automate this process. Finally, all
-built-in collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose
-"triangle soup" collision detector for example. However we do not recommend such an approach. Pre-processing the
-geometry and representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both
-faster and more stable simulation.
+(which can be primitive shapes or meshes) and attach them to the same body. Open tools like the `CoACD library
+<https://github.com/SarahWeiii/CoACD>`__ can be used outside MuJoCo to automate this process. Finally, all built-in
+collision functions can be replaced with custom callbacks. This can be used to incorporate a general-purpose "triangle
+soup" collision detector for example. However we do not recommend such an approach. Pre-processing the geometry and
+representing it as a union of convex geoms takes some work, but it pays off at runtime and yields both faster and more
+stable simulation.
 
 .. _Pipeline:
 
@@ -1511,42 +1556,59 @@ be used to skip default steps and to enable optional steps respectively. Callbac
 Forward dynamics
 ~~~~~~~~~~~~~~~~
 
-The top-level function :ref:`mj_step` invokes the sequence of computations below. Alternatively one can call
-:ref:`mj_forward` which invokes only steps 2-21.
+The source file `engine_forward.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_forward.c>`__
+contains the high-level forward dynamics pipeline:
 
-#. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence
-   is detected, the state is automatically reset and the corresponding warning is raised.
-#. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites,
-   cameras and lights. It also normalizes all quaternions, just in case.
-#. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding
-   kinematic subtrees (to improve floating-point accuracy).
-#. Compute the actuator lengths and moment arms.
-#. Compute the composite rigid body inertias and construct the joint-space inertia matrix.
-#. Compute the sparse factorization of the joint-space inertia matrix.
-#. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection.
-#. Construct the constraint Jacobian and compute the constraint residuals.
-#. Compute the matrices and vectors needed by the constraint solvers.
-#. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
-   tendons.
-#. Compute sensor data that only depends on position, and the potential energy if enabled.
-#. Compute the tendon and actuator velocities.
-#. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at
-   the subtree centers of mass.
-#. Compute all passive forces: spring-dampers in joints and tendons, and fluid dynamics forces.
-#. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
-#. Compute the reference constraint acceleration.
-#. Compute the vector of Coriolis, centrifugal and gravitational forces.
-#. Compute the actuator forces and activation dynamics if defined.
-#. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces.
-#. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
-   constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics.
-#. Compute sensor data that depends on force and acceleration if enabled.
-#. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
-   automatically reset and the corresponding warning is raised.
-#. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
-   dynamics. This is an optional step, and is performed only when enabled.
-#. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
-   repeats the above sequence three more times, except for the optional computations which are performed only once.
+- The top-level function :ref:`mj_step` invokes the entire sequence of computations below.
+- :ref:`mj_forward` invokes only stages **2-22**, computing the continuous-time forward dynamics, ending with the
+  acceleration ``mjData.qacc``.
+- :ref:`mj_step1` invokes stages **1-18** and :ref:`mj_step2` invokes stages **19-25**, breaking :ref:`mj_step` into two
+  distinct phases. This allows the user to write controllers that depend on quantities derived from the positions and
+  velocities (but not forces, since those have not yet been computed). Note that the :ref:`mj_step1` → :ref:`mj_step2`
+  pipeline does not support the Runge Kutta integrator.
+
+1. Check the positions and velocities for invalid or unacceptably large real values indicating divergence. If divergence
+   is detected, the state is automatically reset and the corresponding warning is raised:
+   :ref:`mj_checkPos`, :ref:`mj_checkVel`
+2. Compute the forward kinematics. This yields the global positions and orientations of all bodies, geoms, sites,
+   cameras and lights. It also normalizes all quaternions: :ref:`mj_kinematics`, :ref:`mj_camLight`
+3. Compute the body inertias and joint axes, in global frames centered at the centers of mass of the corresponding
+   kinematic subtrees: :ref:`mj_comPos`
+4. Compute quantities related to :ref:`flex<deformable-flex>` objects: :ref:`mj_flex`
+5. Compute the actuator lengths and moment arms: :ref:`mj_tendon`
+6. Compute the composite rigid body inertias and joint-space inertia matrix: :ref:`mj_crb`
+7. Compute the sparse factorization of the joint-space inertia matrix: :ref:`mj_factorM`
+8. Construct the list of active contacts. This includes both broad-phase and near-phase collision detection:
+   :ref:`mj_collision`
+9. Construct the constraint Jacobian and compute the constraint residuals: :ref:`mj_makeConstraint`
+10. Compute the matrices and vectors needed by the constraint solvers: :ref:`mj_projectConstraint`
+11. Compute the tendon lengths and moment arms. This includes the computation of minimal-length paths for spatial
+    tendons: :ref:`mj_transmission`
+12. Compute sensor data that only depends on position, and the potential energy if enabled: :ref:`mj_sensorPos`,
+    :ref:`mj_energyPos`
+13. Compute the tendon, flex edge and actuator velocities: :ref:`mj_fwdVelocity`
+14. Compute the body velocities and rates of change of the joint axes, again in the global coordinate frames centered at
+    the subtree centers of mass: :ref:`mj_comVel`
+15. Compute passive forces -- spring-dampers in joints and tendons, and fluid forces: :ref:`mj_passive`
+16. Compute sensor data that depends on velocity, and the kinetic energy if enabled
+    (if required by sensors, call :ref:`mj_subtreeVel`): :ref:`mj_sensorVel`
+17. Compute the reference constraint acceleration: :ref:`mj_referenceConstraint`
+18. Compute the vector of Coriolis, centrifugal and gravitational forces: :ref:`mj_rne`
+19. Compute the actuator forces and activation dynamics if defined: :ref:`mj_fwdActuation`
+20. Compute the joint acceleration resulting from all forces except for the (still unknown) constraint forces:
+    :ref:`mj_fwdAcceleration`
+21. Compute the constraint forces with the selected solver, and update the joint acceleration so as to account for the
+    constraint forces. This yields the vector ``mjData.qacc`` which is the main output of forward dynamics:
+    :ref:`mj_fwdConstraint`
+22. Compute sensor data that depends on force and acceleration if enabled
+    (if required by sensors, call :ref:`mj_rnePostConstraint`): :ref:`mj_sensorAcc`
+23. Check the acceleration for invalid or unacceptably large real values. If divergence is detected, the state is
+    automatically reset and the corresponding warning is raised: :ref:`mj_checkAcc`
+24. Compare the results of forward and inverse dynamics, so as to diagnose poor solver convergence in the forward
+    dynamics. This is an optional step, and is performed only when enabled: :ref:`mj_compareFwdInv`
+25. Advance the simulation state by one time step, using the selected integrator. Note that the Runge-Kutta integrator
+    repeats the above sequence three more times, except for the optional computations which are performed only once:
+    one of :ref:`mj_Euler`, :ref:`mj_RungeKutta`, :ref:`mj_implicit`
 
 .. _piInverse:
 
@@ -1567,6 +1629,7 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
 #. Compute the tendon and actuator velocities.
 #. Compute the body velocities and joint axes rates of change.
 #. Compute sensor data that depends on velocity, and the kinetic energy if enabled.
+   If required by sensors, call :ref:`mj_subtreeVel`.
 #. Compute all passive forces.
 #. Compute the reference constraint acceleration.
 #. If the :ref:`invdiscrete<option-flag-invdiscrete>` flag is set and the :ref:`integrator<option-integrator>` is not
@@ -1574,6 +1637,7 @@ The top-level function :ref:`mj_inverse` invokes the following sequence of compu
 #. Compute the constraint force. This is done analytically, without using a numerical solver.
 #. Compute the inverse dynamics for the unconstrained system.
 #. Compute sensor data that depends on force and acceleration if enabled.
+   If required by sensors, call :ref:`mj_rnePostConstraint`.
 #. Compute the vector ``mjData.qfrc_inverse`` by combining all results. This is the main output of inverse dynamics. It
    equals the sum of external and actuation forces.
 

doc/conf.py

@@ -49,9 +49,15 @@ extensions = [
     'sphinx_favicon',
     'sphinx_reredirects',
     'sphinx_toolbox.collapse',
+    'sphinx_toolbox.github',
+    'sphinx_toolbox.sidebar_links',
     'mujoco_include',
 ]
 
+# GitHub-related options
+github_username = 'google-deepmind'
+github_repository = 'mujoco'
+
 # Bibtex references for sphinxcontrib.bibtex
 bibtex_bibfiles = ['references.bib']
 
@@ -74,6 +80,9 @@ exclude_patterns = [
 redirects = {
     # index.rst just contains the table of contents definition.
     'index': 'overview.html',
+    'computation': 'computation/index.html',
+    'programming': 'programming/index.html',
+    'APIreference': 'APIreference/index.html',
 }
 
 rst_prolog = """
@@ -182,8 +191,15 @@ favicons = [
 # -- Options for katex ------------------------------------------------------
 
 # See: https://sphinxcontrib-katex.readthedocs.io/en/0.4.1/macros.html
+# {ar au, ac} are {reference, unconstrained, constrained} acceleration, resp.
 latex_macros = r"""
     \def \d              #1{\operatorname{#1}}
+    \def \ar             {a_{\rm ref}}
+    \def \au             {a_0}
+    \def \ac             {a_1}
+    \def \ari            {a_{{\rm ref},i}}
+    \def \aui            {a_{0,i}}
+    \def \aci            {a_{1,i}}
 """
 
 # Translate LaTeX macros to KaTeX and add to options for HTML builder

doc/css/theme_overrides.css

@@ -19,9 +19,14 @@ body {
 }
 
 /* Make the search bar query text readable in dark theme. */
-body[data-theme=dark] {
+body[data-theme="dark"] {
   --color-sidebar-search-foreground: white;
 }
+@media (prefers-color-scheme: dark) {
+  body[data-theme="auto"] {
+    --color-sidebar-search-foreground: white;
+  }
+}
 
 .schema-small {
   font-size: 93%;

doc/images/computation/contact_frame_dark.svg

@@ -0,0 +1,296 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 559.62 244.73">
+  <defs>
+    <style>
+      .cls-1 {
+        font-size: 12px;
+      }
+
+      .cls-1, .cls-2, .cls-3, .cls-4, .cls-5, .cls-6 {
+        fill: #fff;
+      }
+
+      .cls-1, .cls-3, .cls-4, .cls-5 {
+        font-family: ArialMT, Arial;
+      }
+
+      .cls-1, .cls-3, .cls-4, .cls-5, .cls-7 {
+        isolation: isolate;
+      }
+
+      .cls-8 {
+        stroke-width: 2px;
+      }
+
+      .cls-8, .cls-9 {
+        stroke: #4c4c4c;
+      }
+
+      .cls-8, .cls-9, .cls-10 {
+        fill: none;
+      }
+
+      .cls-2, .cls-10 {
+        stroke: #fff;
+      }
+
+      .cls-3 {
+        font-size: 8px;
+      }
+
+      .cls-9 {
+        stroke-dasharray: 0 0 5 5;
+      }
+
+      .cls-4 {
+        font-size: 24px;
+      }
+
+      .cls-5 {
+        font-size: 9px;
+      }
+
+      .cls-6 {
+        stroke-width: 0px;
+      }
+    </style>
+  </defs>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="79.67" y2="174.04"/>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="72.1" y2="113.9"/>
+  <line class="cls-8" x1="28.52" y1="139.62" x2="28.52" y2="69.94"/>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(25.61 10.3)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(32.28 13.3)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(126 209.44)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(132.67 212.44)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(123.42 85.94)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(130.09 88.94)"><tspan x="0" y="0">3</tspan></text>
+  </g>
+  <g>
+    <path class="cls-10" d="M56.74,118.44c0-11,6.36-20,14.2-20s14.19,9,14.19,20"/>
+    <path class="cls-6" d="M85.13,117.67l2.7-1.64.06.09-1.75,4.41c-.34,1.53-.67,3.07-1,4.61-.34-1.54-.68-3.08-1-4.61l-1.81-4.41v-.09l2.8,1.64Z"/>
+  </g>
+  <g>
+    <path class="cls-10" d="M90.33,180.04c0-11-6.35-20-14.19-20s-14.19,9-14.19,20"/>
+    <path class="cls-6" d="M61.94,179.27l2.7-1.64.06.09-1.75,4.41c-.34,1.53-.68,3.07-1,4.6l-1-4.6-1.75-4.41.05-.09,2.69,1.64Z"/>
+  </g>
+  <g>
+    <path class="cls-10" d="M21.49,83.23c14.25,0,25.8-5.49,25.8-12.26s-11.55-12.26-25.8-12.26"/>
+    <path class="cls-6" d="M22.25,58.71l1.64,2.71-.09.06-4.41-1.75-4.6-1,4.6-1,4.41-1.79h.09l-1.64,2.77Z"/>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(0 62.46)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(6.67 65.46)"><tspan x="0" y="0">4</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(58.77 199.36)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(65.44 202.36)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(81.16 136.09)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(87.83 139.09)"><tspan x="0" y="0">6</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(18.39 122.2)"><tspan x="0" y="0">x</tspan></text>
+  <text class="cls-1" transform="translate(41.29 162.53)"><tspan x="0" y="0">y</tspan></text>
+  <text class="cls-1" transform="translate(50.31 137.92)"><tspan x="0" y="0">z</tspan></text>
+  <path class="cls-2" d="M32,139.57c0,1.92-1.56,3.48-3.48,3.48h0c-1.92,0-3.48-1.55-3.49-3.47h0c0-1.93,1.56-3.49,3.48-3.49,0,0,0,0,.01,0,1.92,0,3.48,1.56,3.48,3.48h0Z"/>
+  <g>
+    <line class="cls-10" x1="28.52" y1="48.07" x2="28.52" y2="22.19"/>
+    <path class="cls-6" d="M28.52,22.94l-2.71,1.64-.06-.09,1.75-4.41,1-4.61c.33,1.54.67,3.07,1,4.61l1.75,4.41-.05.09-2.68-1.64Z"/>
+  </g>
+  <g>
+    <line class="cls-10" x1="94.7" y1="184.06" x2="115.42" y2="198"/>
+    <path class="cls-6" d="M114.79,197.57l.14-3.16h.11l2.68,3.91c1.09,1.14,2.17,2.28,3.26,3.41l-4.39-1.73-4.63-1-.05-.09,2.88-1.34Z"/>
+  </g>
+  <g>
+    <line class="cls-10" x1="88.95" y1="104.87" x2="110.74" y2="92.01"/>
+    <path class="cls-6" d="M110.08,92.4l-2.75-1.46.05-.1,4.68-.74,4.48-1.46-3.45,3.21-2.9,3.75h-.11v-3.2Z"/>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(265.54 123.6)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(272.21 126.6)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(229.69 111.4)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(236.36 114.4)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(307.66 91.78)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(314.33 94.78)"><tspan x="0" y="0">3</tspan></text>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(188.06 133.82)"><tspan x="0" y="0">e</tspan></text>
+    <text class="cls-5" transform="translate(194.73 136.82)"><tspan x="0" y="0">4</tspan></text>
+  </g>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="303.2" y2="207.11"/>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="298.79" y2="149.1"/>
+  <line class="cls-8" x1="255.21" y1="174.82" x2="255.21" y2="105.14"/>
+  <path class="cls-2" d="M258.7,174.77c0,1.92-1.56,3.48-3.48,3.48,0,0,0,0-.01,0-1.92,0-3.48-1.56-3.48-3.48h0c0-1.92,1.56-3.48,3.48-3.48h0c1.92,0,3.48,1.55,3.49,3.47,0,0,0,0,0,.01Z"/>
+  <g>
+    <path class="cls-9" d="M301.33,92.38c7.4,10.56-7.74,29.27-33.83,41.8-26.09,12.53-53.24,14.14-60.65,3.59s7.74-29.27,33.83-41.8,53.24-14.14,60.65-3.59Z"/>
+    <g>
+      <line class="cls-10" x1="209.77" y1="140.36" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M210.38,140.82l-.33,3.15h-.11l-2.45-4.03-3.06-3.59,4.28,2,4.57,1.27.05.09-2.95,1.11Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="267.54" y1="136.09" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M267.33,136.81l-3.07.75v-.11l3-3.66,2.36-4.09c-.14,1.57-.28,3.14-.43,4.7l.33,4.73-.07.07-2.12-2.39Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="242.89" y1="100.97" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M243.01,101.72l-2.39,2.07-.08-.08,1-4.64.24-4.71c.59,1.46,1.17,2.92,1.76,4.38l2.45,4.06v.09l-2.98-1.17Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="300.66" y1="96.7" x2="255.21" y2="174.82"/>
+      <path class="cls-6" d="M300.33,97.35l-3.17.07v-.11l3.72-2.93c1.07-1.16,2.13-2.32,3.2-3.47l-1.44,4.49-.71,4.68-.08.06-1.52-2.79Z"/>
+    </g>
+  </g>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(70.45 14.72)"><tspan x="0" y="0">elliptic basis: E = I</tspan></text>
+    <text class="cls-5" transform="translate(166.83 17.72)"><tspan x="0" y="0">6</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(312.5 17.19)"><tspan x="0" y="0">pyramidal basis: E =</tspan></text>
+  <g>
+    <g>
+      <line class="cls-10" x1="255.97" y1="83.27" x2="255.97" y2="57.39"/>
+      <path class="cls-6" d="M255.97,58.15l-2.7,1.64-.06-.09,1.75-4.41c.34-1.53.67-3.07,1-4.6l1,4.6,1.75,4.41v.09l-2.74-1.64Z"/>
+    </g>
+    <g>
+      <path class="cls-10" d="M242.94,87.42c14.26,0,25.81-5.49,25.81-12.26s-11.55-12.22-25.81-12.22"/>
+      <path class="cls-6" d="M243.71,62.94l1.64,2.7-.09.06-4.41-1.76-4.6-1,4.6-1,4.41-1.75h.09l-1.64,2.75Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(239.43 76.17)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(246.1 79.17)"><tspan x="0" y="0">5</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <line class="cls-10" x1="255.97" y1="40.49" x2="255.97" y2="14.61"/>
+      <path class="cls-6" d="M255.97,15.37l-2.7,1.65-.06-.1,1.75-4.4c.34-1.54.67-3.07,1-4.61l1,4.61,1.75,4.4v.1l-2.74-1.65Z"/>
+    </g>
+    <g>
+      <path class="cls-10" d="M268.75,45.65c-14.25,0-25.81-5.49-25.81-12.26s11.56-12.26,25.81-12.26"/>
+      <path class="cls-6" d="M267.98,21.13l-1.64-2.7.09-.06,4.41,1.75,4.61,1-4.61,1-4.41,1.82-.09-.05,1.64-2.76Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(260.75 35.64)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(267.42 38.64)"><tspan x="0" y="0">6</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M323.24,213.31c0-11-6.36-20-14.2-20s-14.19,9-14.19,20"/>
+      <path class="cls-6" d="M294.85,212.55l2.7-1.65.07.1-1.75,4.4-1,4.61c-.34-1.54-.67-3.07-1-4.61l-1.75-4.4v-.1l2.73,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="309.01" y1="211" x2="309.01" y2="185.12"/>
+      <path class="cls-6" d="M309.01,185.88l-2.68,1.64-.06-.09,1.75-4.4,1-4.61c.33,1.54.67,3.07,1,4.61l1.75,4.4v.09l-2.76-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(301.82 221.38)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(308.49 224.38)"><tspan x="0" y="0">7</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M324.24,230.74c0-11.05,6.35-20,14.19-20s14.19,8.95,14.19,20"/>
+      <path class="cls-6" d="M352.62,229.94l2.71-1.65.06.09-1.75,4.41-1,4.61c-.33-1.54-.67-3.07-1-4.61l-1.75-4.41v-.09l2.73,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="338.39" y1="228.43" x2="338.39" y2="202.55"/>
+      <path class="cls-6" d="M338.39,203.31l-2.7,1.64-.06-.09,1.75-4.4c.34-1.54.67-3.07,1-4.61l1,4.61,1.75,4.4v.09l-2.74-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(331.21 238.81)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(337.88 241.81)"><tspan x="0" y="0">8</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M293.85,147.15c0-11,6.36-20,14.19-20s14.2,9,14.2,20"/>
+      <path class="cls-6" d="M322.24,146.39l2.7-1.65.06.09-1.75,4.41c-.34,1.54-.67,3.07-1,4.61l-1-4.61-1.75-4.41v-.09l2.74,1.65Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="308.01" y1="144.84" x2="308.01" y2="118.96"/>
+      <path class="cls-6" d="M308.01,119.72l-2.68,1.64-.06-.09,1.75-4.4,1-4.61c.33,1.53.67,3.07,1,4.61l1.75,4.4v.09l-2.76-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(300.82 155.22)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(307.49 158.22)"><tspan x="0" y="0">9</tspan></text>
+    </g>
+  </g>
+  <g>
+    <g>
+      <path class="cls-10" d="M353.08,128.06c0-11.05-6.36-20-14.2-20s-14.19,9-14.19,20"/>
+      <path class="cls-6" d="M324.69,127.29l2.7-1.64.06.09-1.75,4.41c-.33,1.54-.67,3.07-1,4.61-.34-1.54-.67-3.07-1-4.61l-1.77-4.41v-.09l2.76,1.64Z"/>
+    </g>
+    <g>
+      <line class="cls-10" x1="338.85" y1="125.75" x2="338.85" y2="99.87"/>
+      <path class="cls-6" d="M338.85,100.63l-2.71,1.64-.06-.09,1.75-4.41,1-4.6c.33,1.53.67,3.07,1,4.6l1.75,4.41-.05.09-2.68-1.64Z"/>
+    </g>
+    <g class="cls-7">
+      <text class="cls-1" transform="translate(331.66 136.13)"><tspan x="0" y="0">e</tspan></text>
+      <text class="cls-5" transform="translate(338.33 139.13)"><tspan x="0" y="0">10</tspan></text>
+    </g>
+  </g>
+  <text class="cls-1" transform="translate(378.74 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(373.33 65.99)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(390.33 68.99)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(378.74 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(378.74 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(407.55 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(407.55 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(402.13 65.99)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(416.12 68.99)"><tspan x="0" y="0">1</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(435.2 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(429.78 88.42)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(446.78 91.42)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(435.2 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(435.2 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(463 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(463 157.91)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(457.59 88.42)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(471.58 91.42)"><tspan x="0" y="0">2</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(517.79 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(512.37 157.17)"><tspan x="0" y="0">+m</tspan></text>
+    <text class="cls-3" transform="translate(529.37 160.17)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <text class="cls-1" transform="translate(517.79 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(517.79 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 44.29)"><tspan x="0" y="0">1</tspan></text>
+  <text class="cls-1" transform="translate(546.59 66.73)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 89.16)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 111.6)"><tspan x="0" y="0">0</tspan></text>
+  <text class="cls-1" transform="translate(546.59 134.03)"><tspan x="0" y="0">0</tspan></text>
+  <g class="cls-7">
+    <text class="cls-1" transform="translate(541.18 157.17)"><tspan x="0" y="0">-m</tspan></text>
+    <text class="cls-3" transform="translate(555.17 160.17)"><tspan x="0" y="0">5</tspan></text>
+  </g>
+  <text class="cls-4" transform="translate(485.39 98.39)"><tspan x="0" y="0">...</tspan></text>
+</svg>

doc/images/computation/gPGS_dark.svg

@@ -0,0 +1,102 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 454.98 171.24">
+  <defs>
+    <style>
+      .cls-1 {
+        font-family: ArialMT, Arial;
+        font-size: 12px;
+      }
+
+      .cls-1, .cls-2, .cls-3 {
+        fill: #fff;
+      }
+
+      .cls-1, .cls-4 {
+        isolation: isolate;
+      }
+
+      .cls-5 {
+        stroke-dasharray: 0 0 4 4;
+      }
+
+      .cls-5, .cls-2, .cls-6, .cls-7, .cls-3, .cls-8 {
+        stroke: #dce0df;
+      }
+
+      .cls-5, .cls-6, .cls-7, .cls-9, .cls-8, .cls-10 {
+        fill: none;
+      }
+
+      .cls-11 {
+        fill: #0c0c0c;
+      }
+
+      .cls-11, .cls-10 {
+        stroke: #666;
+      }
+
+      .cls-6 {
+        stroke-dasharray: 0 5;
+      }
+
+      .cls-7 {
+        stroke-dasharray: 0 0 5 5;
+      }
+
+      .cls-3 {
+        stroke-miterlimit: 4;
+        stroke-width: 1px;
+      }
+
+      .cls-9 {
+        stroke: #fff;
+        stroke-width: 3px;
+      }
+
+      .cls-12 {
+        fill: #dce0df;
+        stroke-width: 0px;
+      }
+    </style>
+  </defs>
+  <line class="cls-10" x1="68" y1="151.26" x2="133.33" y2="67.93"/>
+  <line class="cls-10" x1="68" y1="151.26" x2="2.67" y2="67.93"/>
+  <circle class="cls-3" cx="122.67" cy="118.59" r="2.67"/>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(0 45.93)"><tspan x="0" y="0">cone</tspan></text>
+    <text class="cls-1" transform="translate(0 60.33)"><tspan x="0" y="0">constraint</tspan></text>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(133.33 134.6)"><tspan x="0" y="0">unconstrained</tspan></text>
+    <text class="cls-1" transform="translate(133.33 149)"><tspan x="0" y="0">minimum</tspan></text>
+  </g>
+  <line class="cls-7" x1="122.67" y1="118.6" x2="122.67" y2="81.54"/>
+  <line class="cls-6" x1="125.33" y1="118.6" x2="93.61" y2="118.6"/>
+  <line class="cls-9" x1="93.14" y1="119.2" x2="122.67" y2="81.54"/>
+  <g>
+    <line class="cls-8" x1="104" y1="53.93" x2="104" y2="88.93"/>
+    <path class="cls-12" d="M104,88.17l2.7-1.64.06.09-1.75,4.4c-.33,1.54-.67,3.07-1,4.61-.34-1.54-.67-3.07-1-4.61l-1.75-4.4v-.09l2.74,1.64Z"/>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(85.81 35.63)"><tspan x="0" y="0">continuum of</tspan></text>
+    <text class="cls-1" transform="translate(85.81 50.03)"><tspan x="0" y="0">PGS local minima</tspan></text>
+  </g>
+  <line class="cls-10" x1="387.29" y1="170.93" x2="452.62" y2="87.6"/>
+  <line class="cls-10" x1="387.29" y1="170.93" x2="321.96" y2="87.6"/>
+  <path class="cls-11" d="M454.48,77.71c0,19.83-30.19,35.9-67.44,35.9s-67.42-16.08-67.42-35.9,30.2-35.9,67.45-35.9,67.41,16.06,67.41,35.9Z"/>
+  <path class="cls-2" d="M417.85,70.44c0,1.03-.83,1.86-1.86,1.86s-1.86-.83-1.86-1.86.83-1.86,1.86-1.86h0c1.02,0,1.85.82,1.86,1.84,0,0,0,.01,0,.02Z"/>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(402.51 10.3)"><tspan x="0" y="0">search</tspan></text>
+    <text class="cls-1" transform="translate(402.51 24.7)"><tspan x="0" y="0">ray</tspan></text>
+  </g>
+  <g class="cls-4">
+    <text class="cls-1" transform="translate(297.84 33.33)"><tspan x="0" y="0">search</tspan></text>
+    <text class="cls-1" transform="translate(297.84 47.73)"><tspan x="0" y="0">ellipsoid</tspan></text>
+  </g>
+  <g>
+    <line class="cls-8" x1="416.79" y1="68.24" x2="429.11" y2="25.33"/>
+    <path class="cls-12" d="M428.9,26.07l-3.05.83v-.1l2.89-3.75,2.25-4.15-.3,4.71.47,4.71-.07.08-2.19-2.33Z"/>
+  </g>
+  <line class="cls-5" x1="403.95" y1="112.94" x2="416.79" y2="68.24"/>
+  <line class="cls-8" x1="387.29" y1="170.93" x2="403.95" y2="112.94"/>
+</svg>

doc/images/computation/kutta_cond_plate_dark.svg

@@ -0,0 +1,124 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg id="Layer_1" data-name="Layer 1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 792.59 237.27">
+  <defs>
+    <style>
+      .cls-1 {
+        stroke-miterlimit: 4;
+        stroke-width: 16px;
+      }
+
+      .cls-1, .cls-2, .cls-3, .cls-4, .cls-5, .cls-6, .cls-7, .cls-8 {
+        fill: none;
+      }
+
+      .cls-1, .cls-5, .cls-6, .cls-8 {
+        stroke: #fff;
+      }
+
+      .cls-9 {
+        clip-path: url(#clippath);
+      }
+
+      .cls-2 {
+        stroke: #007b76;
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-3 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-3, .cls-7 {
+        stroke: #0076ba;
+      }
+
+      .cls-4, .cls-10, .cls-11, .cls-12 {
+        stroke-width: 0px;
+      }
+
+      .cls-5 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-6 {
+        stroke-miterlimit: 4;
+        stroke-width: 2px;
+      }
+
+      .cls-7, .cls-8 {
+        stroke-width: 2px;
+      }
+
+      .cls-13 {
+        font-family: ArialMT, Arial;
+        font-size: 32px;
+      }
+
+      .cls-13, .cls-14, .cls-15 {
+        isolation: isolate;
+      }
+
+      .cls-13, .cls-15, .cls-11 {
+        fill: #fff;
+      }
+
+      .cls-15 {
+        font-family: STIXGeneral-Italic, STIXGeneral;
+        font-size: 32.22px;
+        font-style: italic;
+      }
+
+      .cls-10 {
+        fill: #99195e;
+      }
+
+      .cls-12 {
+        fill: #0076ba;
+      }
+    </style>
+    <clipPath id="clippath">
+      <rect class="cls-4" x="520.61" y="100.37" width="16.82" height="14.57"/>
+    </clipPath>
+  </defs>
+  <g>
+    <path class="cls-1" d="M153.2,197.67L659.31,46.23"/>
+    <path class="cls-2" d="M653.5,38.16c-69.23-65.09-392.9,103.98-501.22,149.95"/>
+    <path class="cls-7" d="M792.58,52.6c-39.16-.63-79.53-1.41-113.82-20.52-11.37-6.34-21.66-14.59-33.61-19.79-20.67-9-42.42-8.25-64.46-5.66C410.68,26.63,262.71,144.22,96.39,182.57c-31.65,7.3-63.9,11.55-96.37,12.26"/>
+    <path class="cls-7" d="M788.09,72.98c-34.82-2.64-69.65,1.7-103.92,8.29-25.81,4.96-51.35,11.2-76.72,18.09-154.05,41.88-302.08,108.19-461.01,125.56-47.05,5.14-94.43,5.87-141.68,8.47"/>
+    <path class="cls-12" d="M787.39,47.75l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M783.33,67.36l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M29.24,188.93l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M30.24,227.67l-9.6,4.8,9.6,4.8v-9.6Z"/>
+    <path class="cls-12" d="M615.4,0l-9.6,4.8,9.6,4.8V0Z"/>
+    <path class="cls-12" d="M349.27,77.03l-6.8,8.3,10.72.46-3.92-8.76Z"/>
+    <path class="cls-12" d="M473.5,26.47l-7.5,7.68,10.64,1.39-3.15-9.07Z"/>
+    <path class="cls-3" d="M195.07,151.07l-1.58.58-.94.34"/>
+    <path class="cls-12" d="M191.84,147.14l-7.36,7.81,10.67,1.2-3.3-9.01Z"/>
+    <path class="cls-12" d="M613.28,92.82l-8.06,7.09,10.51,2.19-2.45-9.28Z"/>
+    <path class="cls-12" d="M442.34,145.91l-8.06,7.09,10.51,2.19-2.45-9.28Z"/>
+    <path class="cls-12" d="M212.2,210.07l-8.53,6.51,10.33,2.92-1.8-9.43Z"/>
+    <path class="cls-7" d="M553.05,59.87c11.3-6.46,36.25-12.77,22.33-20.73-17.28-9.89-181.61,49.61-169.38,59.08,8.52,6.6,46.77-6.02,72.27-14.91l.95-.33"/>
+    <path class="cls-12" d="M479.86,87.84l7.47-7.7-10.65-1.36,3.17,9.06Z"/>
+    <path class="cls-8" d="M405.92,122.4h219.45"/>
+    <path class="cls-11" d="M624.38,127.2l9.6-4.8-9.6-4.8v9.6Z"/>
+    <text class="cls-13" transform="translate(637.08 131.54)"><tspan x="0" y="0">v</tspan></text>
+    <path class="cls-6" d="M141.34,157.03l38.77-11.63.96-.29"/>
+    <path class="cls-11" d="M181.49,149.99l7.82-7.36-10.57-1.84,2.76,9.2Z"/>
+    <path class="cls-5" d="M141.31,157.86l-11.63-38.77-.29-.96"/>
+    <path class="cls-11" d="M134.28,117.7l-7.36-7.82-1.84,10.57,9.2-2.76Z"/>
+    <text class="cls-13" transform="translate(135.06 107.05)"><tspan x="0" y="0">y</tspan></text>
+    <text class="cls-13" transform="translate(181.5 131.54)"><tspan x="0" y="0">x</tspan></text>
+    <path class="cls-10" d="M672.67,54.95c2.03,1.97,2.03,5.17,0,7.14-2.03,1.97-5.32,1.97-7.35,0-2.03-1.97-2.03-5.17,0-7.14,2.03-1.97,5.32-1.97,7.35,0Z"/>
+    <path class="cls-10" d="M147.18,196.95c2.03,1.97,2.03,5.17,0,7.14s-5.32,1.97-7.35,0-2.03-5.17,0-7.14c2.03-1.97,5.32-1.97,7.35,0Z"/>
+    <path class="cls-8" d="M508.05,122.3c.72-19.73-3.57-30.4-12.89-31.99"/>
+    <text/>
+  </g>
+  <g class="cls-9">
+    <g class="cls-14">
+      <text class="cls-15" transform="translate(519.74 114.58)"><tspan x="0" y="0">α</tspan></text>
+    </g>
+  </g>
+</svg>

doc/includes/references.h

@@ -73,13 +73,17 @@ typedef enum mjtTimer_ {     // internal timers
   mjTIMER_POS_MAKE,          // make constraints
   mjTIMER_POS_PROJECT,       // project constraints
 
+  // breakdown of mj_collision
+  mjTIMER_COL_BROAD,         // broadphase
+  mjTIMER_COL_NARROW,        // narrowphase
+
   mjNTIMER                   // number of timers
 } mjtTimer;
 struct mjContact_ {                // result of collision detection functions
-  // contact parameters set by geom-specific collision detector
+  // contact parameters set by near-phase collision function
   mjtNum  dist;                    // distance between nearest points; neg: penetration
   mjtNum  pos[3];                  // position of contact point: midpoint between geoms
-  mjtNum  frame[9];                // normal is in [0-2]
+  mjtNum  frame[9];                // normal is in [0-2], points from geom[0] to geom[1]
 
   // contact parameters set by mj_collideGeoms
   mjtNum  includemargin;           // include if dist<includemargin=margin-gap
@@ -92,12 +96,16 @@ struct mjContact_ {                // result of collision detection functions
   mjtNum  mu;                      // friction of regularized cone, set by mj_makeConstraint
   mjtNum  H[36];                   // cone Hessian, set by mj_updateConstraint
 
-  // contact descriptors set by mj_collideGeoms
+  // contact descriptors set by mj_collideXXX
   int     dim;                     // contact space dimensionality: 1, 3, 4 or 6
-  int     geom1;                   // id of geom 1
-  int     geom2;                   // id of geom 2
+  int     geom1;                   // id of geom 1; deprecated, use geom[0]
+  int     geom2;                   // id of geom 2; deprecated, use geom[1]
+  int     geom[2];                 // geom ids; -1 for flex
+  int     flex[2];                 // flex ids; -1 for geom
+  int     elem[2];                 // element ids; -1 for geom or flex vertex
+  int     vert[2];                 // vertex ids;  -1 for geom or flex element
 
-  // flag set by mj_instantianteEquality
+  // flag set by mj_setContact or mj_instantiateContact
   int     exclude;                 // 0: include, 1: in gap, 2: fused, 3: no dofs
 
   // address computed by mj_instantiateContact
@@ -139,7 +147,7 @@ struct mjData_ {
 
   // memory utilization stats
   size_t  maxuse_stack;                      // maximum stack allocation in bytes
-  size_t  maxuse_threadstack[mjMAXTHREADS];  // maximum stack allocation per thread in bytes
+  size_t  maxuse_threadstack[mjMAXTHREAD];   // maximum stack allocation per thread in bytes
   size_t  maxuse_arena;                      // maximum arena allocation in bytes
   int     maxuse_con;                        // maximum number of contacts
   int     maxuse_efc;                        // maximum number of scalar constraints
@@ -155,12 +163,6 @@ struct mjData_ {
   int     solver_nnz[mjNISLAND];    // number of non-zeros in Hessian or efc_AR, per island
   mjtNum  solver_fwdinv[2];         // forward-inverse comparison: qfrc, efc
 
-  // collision statistics
-  int     nbodypair_broad;   // number of body pairs in collision according to the broad-phase
-  int     nbodypair_narrow;  // number of body pairs actually in collision in the narrow-phase
-  int     ngeompair_mid;     // number of geom pairs in collision according to the mid-phase
-  int     ngeompair_narrow;  // number of geom pairs actually in collision in the narrow-phase
-
   // variable sizes
   int     ne;                // number of equality constraints
   int     nf;                // number of friction constraints
@@ -234,17 +236,26 @@ struct mjData_ {
 
   // computed by mj_fwdPosition/mj_comPos
   mjtNum* subtree_com;       // center of mass of each subtree                   (nbody x 3)
-  mjtNum* cdof;              // com-based motion axis of each dof                (nv x 6)
+  mjtNum* cdof;              // com-based motion axis of each dof (rot:lin)      (nv x 6)
   mjtNum* cinert;            // com-based body inertia and mass                  (nbody x 10)
 
+  // computed by mj_fwdPosition/mj_flex
+  mjtNum* flexvert_xpos;     // Cartesian flex vertex positions                  (nflexvert x 3)
+  mjtNum* flexelem_aabb;     // flex element bounding boxes (center, size)       (nflexelem x 6)
+  int*    flexedge_J_rownnz; // number of non-zeros in Jacobian row              (nflexedge x 1)
+  int*    flexedge_J_rowadr; // row start address in colind array                (nflexedge x 1)
+  int*    flexedge_J_colind; // column indices in sparse Jacobian                (nflexedge x nv)
+  mjtNum* flexedge_J;        // flex edge Jacobian                               (nflexedge x nv)
+  mjtNum* flexedge_length;   // flex edge lengths                                (nflexedge x 1)
+
   // computed by mj_fwdPosition/mj_tendon
   int*    ten_wrapadr;       // start address of tendon's path                   (ntendon x 1)
   int*    ten_wrapnum;       // number of wrap points in path                    (ntendon x 1)
   int*    ten_J_rownnz;      // number of non-zeros in Jacobian row              (ntendon x 1)
   int*    ten_J_rowadr;      // row start address in colind array                (ntendon x 1)
   int*    ten_J_colind;      // column indices in sparse Jacobian                (ntendon x nv)
-  mjtNum* ten_length;        // tendon lengths                                   (ntendon x 1)
   mjtNum* ten_J;             // tendon Jacobian                                  (ntendon x nv)
+  mjtNum* ten_length;        // tendon lengths                                   (ntendon x 1)
   int*    wrap_obj;          // geom id; -1: site; -2: pulley                    (nwrap*2 x 1)
   mjtNum* wrap_xpos;         // Cartesian 3D points in all path                  (nwrap*2 x 3)
 
@@ -262,31 +273,37 @@ struct mjData_ {
   mjtNum* qLDiagSqrtInv;     // 1/sqrt(diag(D))                                  (nv x 1)
 
   // computed by mj_collisionTree
+  mjtNum*  bvh_aabb_dyn;     // global bounding box (center, size)               (nbvhdynamic x 6)
   mjtByte* bvh_active;       // volume has been added to collisions              (nbvh x 1)
 
   //-------------------- POSITION, VELOCITY dependent
 
   // computed by mj_fwdVelocity
+  mjtNum* flexedge_velocity; // flex edge velocities                             (nflexedge x 1)
   mjtNum* ten_velocity;      // tendon velocities                                (ntendon x 1)
   mjtNum* actuator_velocity; // actuator velocities                              (nu x 1)
 
   // computed by mj_fwdVelocity/mj_comVel
-  mjtNum* cvel;              // com-based velocity [3D rot; 3D tran]             (nbody x 6)
-  mjtNum* cdof_dot;          // time-derivative of cdof                          (nv x 6)
+  mjtNum* cvel;              // com-based velocity (rot:lin)                     (nbody x 6)
+  mjtNum* cdof_dot;          // time-derivative of cdof (rot:lin)                (nv x 6)
 
   // computed by mj_fwdVelocity/mj_rne (without acceleration)
   mjtNum* qfrc_bias;         // C(qpos,qvel)                                     (nv x 1)
 
   // computed by mj_fwdVelocity/mj_passive
-  mjtNum* qfrc_passive;      // passive force                                    (nv x 1)
+  mjtNum* qfrc_spring;       // passive spring force                             (nv x 1)
+  mjtNum* qfrc_damper;       // passive damper force                             (nv x 1)
+  mjtNum* qfrc_gravcomp;     // passive gravity compensation force               (nv x 1)
+  mjtNum* qfrc_fluid;        // passive fluid force                              (nv x 1)
+  mjtNum* qfrc_passive;      // total passive force                              (nv x 1)
 
   // computed by mj_sensorVel/mj_subtreeVel if needed
   mjtNum* subtree_linvel;    // linear velocity of subtree com                   (nbody x 3)
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum*   qH;              // L'*D*L factorization of modified M               (nM x 1)
-  mjtNum*   qHDiagInv;       // 1/diag(D) of modified M                          (nv x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
   int*    D_rownnz;          // non-zeros in each row                            (nv x 1)
@@ -441,9 +458,11 @@ typedef enum mjtGeom_ {           // type of geometric shape
   mjGEOM_ARROW1,                  // arrow without wedges
   mjGEOM_ARROW2,                  // arrow in both directions
   mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
+  mjGEOM_FLEX,                    // flex
   mjGEOM_SKIN,                    // skin
   mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle connecting a frame
+  mjGEOM_TRIANGLE,                // triangle
 
   mjGEOM_NONE         = 1001      // missing geom type
 } mjtGeom;
@@ -484,6 +503,7 @@ typedef enum mjtEq_ {             // type of equality constraint
   mjEQ_WELD,                      // fix relative position and orientation of two bodies
   mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
   mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
+  mjEQ_FLEX,                      // fix all edge lengths of a flex
   mjEQ_DISTANCE                   // unsupported, will cause an error if used
 } mjtEq;
 typedef enum mjtWrap_ {           // type of tendon wrap object
@@ -534,6 +554,7 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_SITE,                     // site
   mjOBJ_CAMERA,                   // camera
   mjOBJ_LIGHT,                    // light
+  mjOBJ_FLEX,                     // flex
   mjOBJ_MESH,                     // mesh
   mjOBJ_SKIN,                     // skin
   mjOBJ_HFIELD,                   // heightfield
@@ -549,7 +570,9 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_TEXT,                     // text
   mjOBJ_TUPLE,                    // tuple
   mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN                    // plugin instance
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT                       // number of object types
 } mjtObj;
 typedef enum mjtConstraint_ {     // type of constraint
   mjCNSTR_EQUALITY    = 0,        // equality constraint
@@ -645,6 +668,13 @@ typedef enum mjtLRMode_ {         // mode for actuator length range computation
   mjLRMODE_MUSCLEUSER,            // process muscle and user actuators
   mjLRMODE_ALL                    // process all actuators
 } mjtLRMode;
+typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
+  mjFLEXSELF_NONE   = 0,          // no self-collisions
+  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
+  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
+  mjFLEXSELF_SAP,                 // use SAP in midphase
+  mjFLEXSELF_AUTO                 // choose between BVH and SAP automatically
+} mjtFlexSelf;
 struct mjLROpt_ {                 // options for mj_setLengthRange()
   // flags
   int mode;                       // which actuators to process (mjtLRMode)
@@ -691,6 +721,7 @@ struct mjOption_ {                // physics options
   mjtNum o_margin;                // margin
   mjtNum o_solref[mjNREF];        // solref
   mjtNum o_solimp[mjNIMP];        // solimp
+  mjtNum o_friction[5];           // friction
 
   // discrete settings
   int integrator;                 // integration mode (mjtIntegrator)
@@ -703,6 +734,7 @@ struct mjOption_ {                // physics options
   int mpr_iterations;             // maximum number of MPR solver iterations
   int disableflags;               // bit flags for disabling standard features
   int enableflags;                // bit flags for enabling optional features
+  int disableactuator;            // bit flags for disabling actuators by group id
 
   // sdf collision settings
   int sdf_initpoints;             // number of starting points for gradient descent
@@ -718,9 +750,10 @@ struct mjVisual_ {                // visualization options
     float linewidth;              // line width for wireframe and ray rendering
     float glow;                   // glow coefficient for selected body
     float realtime;               // initial real-time factor (1: real time)
-    int offwidth;                 // width of offscreen buffer
-    int offheight;                // height of offscreen buffer
-    int ellipsoidinertia;         // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   offwidth;               // width of offscreen buffer
+    int   offheight;              // height of offscreen buffer
+    int   ellipsoidinertia;       // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   bvactive;               // visualize active bounding volumes (0: no, 1: yes)
   } global;
 
   struct {                        // rendering quality
@@ -771,6 +804,7 @@ struct mjVisual_ {                // visualization options
     float framewidth;             // bodyframe axis width
     float constraint;             // constraint width
     float slidercrank;            // slidercrank width
+    float frustum;                // frustum zfar plane
   } scale;
 
   struct {                        // color of decor elements
@@ -796,6 +830,9 @@ struct mjVisual_ {                // visualization options
     float constraint[4];          // constraint
     float slidercrank[4];         // slidercrank
     float crankbroken[4];         // used when crank must be stretched/broken
+    float frustum[4];             // camera frustum
+    float bv[4];                  // bounding volume
+    float bvactive[4];            // active bounding volume
   } rgba;
 };
 typedef struct mjVisual_ mjVisual;
@@ -816,12 +853,22 @@ struct mjModel_ {
   int nu;                         // number of actuators/controls = dim(ctrl)
   int na;                         // number of activation states = dim(act)
   int nbody;                      // number of bodies
-  int nbvh;                       // number of bounding volumes in all bodies
+  int nbvh;                       // number of total bounding volumes in all bodies
+  int nbvhstatic;                 // number of static bounding volumes (aabb stored in mjModel)
+  int nbvhdynamic;                // number of dynamic bounding volumes (aabb stored in mjData)
   int njnt;                       // number of joints
   int ngeom;                      // number of geoms
   int nsite;                      // number of sites
   int ncam;                       // number of cameras
   int nlight;                     // number of lights
+  int nflex;                      // number of flexes
+  int nflexvert;                  // number of vertices in all flexes
+  int nflexedge;                  // number of edges in all flexes
+  int nflexelem;                  // number of elements in all flexes
+  int nflexelemdata;              // number of element vertex ids in all flexes
+  int nflexshelldata;             // number of shell fragment vertex ids in all flexes
+  int nflexevpair;                // number of element-vertex pairs in all flexes
+  int nflextexcoord;              // number of vertices with texture coordinates
   int nmesh;                      // number of meshes
   int nmeshvert;                  // number of vertices in all meshes
   int nmeshnormal;                // number of normals in all meshes
@@ -909,7 +956,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // body is simple (has diagonal M)          (nbody x 1)
+  mjtByte*  body_simple;          // 1: diag M; 2: diag M, sliders only       (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)
@@ -920,16 +967,19 @@ struct mjModel_ {
   mjtNum*   body_inertia;         // diagonal inertia in ipos/iquat frame     (nbody x 3)
   mjtNum*   body_invweight0;      // mean inv inert in qpos0 (trn, rot)       (nbody x 2)
   mjtNum*   body_gravcomp;        // antigravity force, units of body weight  (nbody x 1)
+  mjtNum*   body_margin;          // MAX over all geom margins                (nbody x 1)
   mjtNum*   body_user;            // user data                                (nbody x nuser_body)
   int*      body_plugin;          // plugin instance id; -1: not in use       (nbody x 1)
+  int*      body_contype;         // OR over all geom contypes                (nbody x 1)
+  int*      body_conaffinity;     // OR over all geom conaffinities           (nbody x 1)
   int*      body_bvhadr;          // address of bvh root                      (nbody x 1)
   int*      body_bvhnum;          // number of bounding volumes               (nbody x 1)
 
   // bounding volume hierarchy
   int*      bvh_depth;            // depth in the bounding volume hierarchy   (nbvh x 1)
   int*      bvh_child;            // left and right children in tree          (nbvh x 2)
-  int*      bvh_geomid;           // geom id of the node; -1: non-leaf        (nbvh x 1)
-  mjtNum*   bvh_aabb;             // bounding box of node (center, size)      (nbvh x 6)
+  int*      bvh_nodeid;           // geom or elem id of node; -1: non-leaf    (nbvh x 1)
+  mjtNum*   bvh_aabb;             // local bounding box (center, size)        (nbvhstatic x 6)
 
   // joints
   int*      jnt_type;             // type of joint (mjtJoint)                 (njnt x 1)
@@ -1038,6 +1088,60 @@ struct mjModel_ {
   float*    light_diffuse;        // diffuse rgb (alpha=1)                    (nlight x 3)
   float*    light_specular;       // specular rgb (alpha=1)                   (nlight x 3)
 
+  // flexes: contact properties
+  int*      flex_contype;         // flex contact type                        (nflex x 1)
+  int*      flex_conaffinity;     // flex contact affinity                    (nflex x 1)
+  int*      flex_condim;          // contact dimensionality (1, 3, 4, 6)      (nflex x 1)
+  int*      flex_priority;        // flex contact priority                    (nflex x 1)
+  mjtNum*   flex_solmix;          // mix coef for solref/imp in contact pair  (nflex x 1)
+  mjtNum*   flex_solref;          // constraint solver reference: contact     (nflex x mjNREF)
+  mjtNum*   flex_solimp;          // constraint solver impedance: contact     (nflex x mjNIMP)
+  mjtNum*   flex_friction;        // friction for (slide, spin, roll)         (nflex x 3)
+  mjtNum*   flex_margin;          // detect contact if dist<margin            (nflex x 1)
+  mjtNum*   flex_gap;             // include in solver if dist<margin-gap     (nflex x 1)
+  mjtByte*  flex_internal;        // internal flex collision enabled          (nflex x 1)
+  int*      flex_selfcollide;     // self collision mode (mjtFlexSelf)        (nflex x 1)
+  int*      flex_activelayers;    // number of active element layers, 3D only (nflex x 1)
+
+  // flexes: other properties
+  int*      flex_dim;             // 1: lines, 2: triangles, 3: tetrahedra    (nflex x 1)
+  int*      flex_matid;           // material id for rendering                (nflex x 1)
+  int*      flex_group;           // group for visibility                     (nflex x 1)
+  int*      flex_vertadr;         // first vertex address                     (nflex x 1)
+  int*      flex_vertnum;         // number of vertices                       (nflex x 1)
+  int*      flex_edgeadr;         // first edge address                       (nflex x 1)
+  int*      flex_edgenum;         // number of edges                          (nflex x 1)
+  int*      flex_elemadr;         // first element address                    (nflex x 1)
+  int*      flex_elemnum;         // number of elements                       (nflex x 1)
+  int*      flex_elemdataadr;     // first element vertex id address          (nflex x 1)
+  int*      flex_shellnum;        // number of shells                         (nflex x 1)
+  int*      flex_shelldataadr;    // first shell data address                 (nflex x 1)
+  int*      flex_evpairadr;       // first evpair address                     (nflex x 1)
+  int*      flex_evpairnum;       // number of evpairs                        (nflex x 1)
+  int*      flex_texcoordadr;     // address in flex_texcoord; -1: none       (nflex x 1)
+  int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
+  int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
+  int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
+  int*      flex_elemlayer;       // element distance from surface, 3D only   (nflexelem x 1)
+  int*      flex_shell;           // shell fragment vertex ids (dim per frag) (nflexshelldata x 1)
+  int*      flex_evpair;          // (element, vertex) collision pairs        (nflexevpair x 2)
+  mjtNum*   flex_vert;            // vertex positions in local body frames    (nflexvert x 3)
+  mjtNum*   flex_xvert0;          // Cartesian vertex positions in qpos0      (nflexvert x 3)
+  mjtNum*   flexedge_length0;     // edge lengths in qpos0                    (nflexedge x 1)
+  mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
+  mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
+  mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
+  mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
+  mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)
+  mjtByte*  flex_rigid;           // are all verices in the same body         (nflex x 1)
+  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
+  mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
+  mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
+  int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)
+  int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
+  float*    flex_rgba;            // rgba when material is omitted            (nflex x 4)
+  float*    flex_texcoord;        // vertex texture coordinates               (nflextexcoord x 2)
+
   // meshes
   int*      mesh_vertadr;         // first vertex address                     (nmesh x 1)
   int*      mesh_vertnum;         // number of vertices                       (nmesh x 1)
@@ -1083,20 +1187,23 @@ struct mjModel_ {
   int*      skin_bonebodyid;      // body id of each bone                     (nskinbone x 1)
   int*      skin_bonevertid;      // mesh ids of vertices in each bone        (nskinbonevert x 1)
   float*    skin_bonevertweight;  // weights of vertices in each bone         (nskinbonevert x 1)
+  int*      skin_pathadr;         // address of asset path for skin; -1: none (nskin x 1)
 
   // height fields
-  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                  (nhfield x 4)
-  int*      hfield_nrow;          // number of rows in grid                   (nhfield x 1)
-  int*      hfield_ncol;          // number of columns in grid                (nhfield x 1)
-  int*      hfield_adr;           // address in hfield_data                   (nhfield x 1)
-  float*    hfield_data;          // elevation data                           (nhfielddata x 1)
+  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                    (nhfield x 4)
+  int*      hfield_nrow;          // number of rows in grid                     (nhfield x 1)
+  int*      hfield_ncol;          // number of columns in grid                  (nhfield x 1)
+  int*      hfield_adr;           // address in hfield_data                     (nhfield x 1)
+  float*    hfield_data;          // elevation data                             (nhfielddata x 1)
+  int*      hfield_pathadr;       // address of asset path for hfield; -1: none (nhfield x 1)
 
   // textures
-  int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
-  int*      tex_height;           // number of rows in texture image          (ntex x 1)
-  int*      tex_width;            // number of columns in texture image       (ntex x 1)
-  int*      tex_adr;              // address in rgb                           (ntex x 1)
-  mjtByte*  tex_rgb;              // rgb (alpha = 1)                          (ntexdata x 1)
+  int*      tex_type;             // texture type (mjtTexture)                  (ntex x 1)
+  int*      tex_height;           // number of rows in texture image            (ntex x 1)
+  int*      tex_width;            // number of columns in texture image         (ntex x 1)
+  int*      tex_adr;              // address in rgb                             (ntex x 1)
+  mjtByte*  tex_rgb;              // rgb (alpha = 1)                            (ntexdata x 1)
+  int*      tex_pathadr;         // address of asset path for texture; -1: none (ntex x 1)
 
   // materials
   int*      mat_texid;            // texture id; -1: none                     (nmat x 1)
@@ -1112,7 +1219,7 @@ struct mjModel_ {
   int*      pair_dim;             // contact dimensionality                   (npair x 1)
   int*      pair_geom1;           // id of geom1                              (npair x 1)
   int*      pair_geom2;           // id of geom2                              (npair x 1)
-  int*      pair_signature;       // (body1+1) << 16 + body2+1                (npair x 1)
+  int*      pair_signature;       // body1 << 16 + body2                      (npair x 1)
   mjtNum*   pair_solref;          // solver reference: contact normal         (npair x mjNREF)
   mjtNum*   pair_solreffriction;  // solver reference: contact friction       (npair x mjNREF)
   mjtNum*   pair_solimp;          // solver impedance: contact                (npair x mjNIMP)
@@ -1121,7 +1228,7 @@ struct mjModel_ {
   mjtNum*   pair_friction;        // tangent1, 2, spin, roll1, 2              (npair x 5)
 
   // excluded body pairs for collision detection
-  int*      exclude_signature;    // (body1+1) << 16 + body2+1                (nexclude x 1)
+  int*      exclude_signature;    // body1 << 16 + body2                      (nexclude x 1)
 
   // equality constraints
   int*      eq_type;              // constraint type (mjtEq)                  (neq x 1)
@@ -1241,6 +1348,7 @@ struct mjModel_ {
   int*      name_siteadr;         // site name pointers                       (nsite x 1)
   int*      name_camadr;          // camera name pointers                     (ncam x 1)
   int*      name_lightadr;        // light name pointers                      (nlight x 1)
+  int*      name_flexadr;         // flex name pointers                       (nflex x 1)
   int*      name_meshadr;         // mesh name pointers                       (nmesh x 1)
   int*      name_skinadr;         // skin name pointers                       (nskin x 1)
   int*      name_hfieldadr;       // hfield name pointers                     (nhfield x 1)
@@ -1322,6 +1430,15 @@ struct mjpPlugin_ {
   // called by mjv_updateScene (optional)
   void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance);
 
+  // methods specific to actuators (optional)
+
+  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
+  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
+
+  // updates the actuator plugin's entries in act_dot
+  // called after native act_dot is computed and before the compute callback
+  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
+
   // methods specific to signed distance fields (optional)
 
   // signed distance from the surface
@@ -1333,6 +1450,9 @@ struct mjpPlugin_ {
   // called during compilation for marching cubes
   mjtNum (*sdf_staticdistance)(const mjtNum point[3], const mjtNum* attributes);
 
+  // convert attributes and provide defaults if not present
+  void (*sdf_attribute)(mjtNum attribute[], const char* name[], const char* value[]);
+
   // bounding box of implicit surface
   void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
 };
@@ -1351,6 +1471,10 @@ typedef enum mjtFramebuffer_ {    // OpenGL framebuffer option
   mjFB_WINDOW         = 0,        // default/window buffer
   mjFB_OFFSCREEN                  // offscreen buffer
 } mjtFramebuffer;
+typedef enum mjtDepthMap_ {       // depth mapping for `mjr_readPixels`
+  mjDEPTH_ZERONEAR    = 0,        // standard depth map; 0: znear, 1: zfar
+  mjDEPTH_ZEROFAR     = 1         // reversed depth map; 1: znear, 0: zfar
+} mjtDepthMap;
 typedef enum mjtFontScale_ {      // font scale, used at context creation
   mjFONTSCALE_50      = 50,       // 50% scale, suitable for low-res rendering
   mjFONTSCALE_100     = 100,      // normal scale, suitable in the absence of DPI scaling
@@ -1454,6 +1578,9 @@ struct mjrContext_ {              // custom OpenGL context
 
   // pixel output format
   int     readPixelFormat;        // default color pixel format for mjr_readPixels
+
+  // depth output format
+  int     readDepthMap;           // depth mapping: mjDEPTH_ZERONEAR or mjDEPTH_ZEROFAR
 };
 typedef struct mjrContext_ mjrContext;
 typedef enum mjtTaskStatus_ {  // status values for mjTask
@@ -1727,6 +1854,7 @@ typedef enum mjtLabel_ {          // object labeling
   mjLABEL_TENDON,                 // tendon labels
   mjLABEL_ACTUATOR,               // actuator labels
   mjLABEL_CONSTRAINT,             // constraint labels
+  mjLABEL_FLEX,                   // flex labels
   mjLABEL_SKIN,                   // skin labels
   mjLABEL_SELECTION,              // selected object
   mjLABEL_SELPNT,                 // coordinates of selection point
@@ -1773,7 +1901,12 @@ typedef enum mjtVisFlag_ {        // flags enabling model element visualization
   mjVIS_SELECT,                   // selection point
   mjVIS_STATIC,                   // static bodies
   mjVIS_SKIN,                     // skin
-  mjVIS_MIDPHASE,                 // mid-phase bounding volume hierarchy
+  mjVIS_FLEXVERT,                 // flex vertices
+  mjVIS_FLEXEDGE,                 // flex edges
+  mjVIS_FLEXFACE,                 // flex element faces
+  mjVIS_FLEXSKIN,                 // flex smooth skin (disables the rest)
+  mjVIS_BODYBVH,                  // body bounding volume hierarchy
+  mjVIS_FLEXBVH,                  // flex bounding volume hierarchy
   mjVIS_MESHBVH,                  // mesh bounding volume hierarchy
   mjVIS_SDFITER,                  // iterations of SDF gradient descent
 
@@ -1800,6 +1933,7 @@ typedef enum mjtStereo_ {         // type of stereo rendering
 } mjtStereo;
 struct mjvPerturb_ {              // object selection and perturbation
   int      select;                // selected body id; non-positive: none
+  int      flexselect;            // selected flex id; negative: none
   int      skinselect;            // selected skin id; negative: none
   int      active;                // perturbation bitmask (mjtPertBit)
   int      active2;               // secondary perturbation bitmask (mjtPertBit)
@@ -1848,7 +1982,7 @@ struct mjvGeom_ {                 // abstract geom
   int      category;              // visual category
   int      texid;                 // texture id; -1: no texture
   int      texuniform;            // uniform cube mapping
-  int      texcoord;              // mesh geom has texture coordinates
+  int      texcoord;              // mesh or flex geom has texture coordinates
   int      segid;                 // segmentation id; -1: not shown
 
   // OpenGL info
@@ -1891,9 +2025,11 @@ struct mjvOption_ {                  // abstract visualization options
   mjtByte  jointgroup[mjNGROUP];     // joint visualization by group
   mjtByte  tendongroup[mjNGROUP];    // tendon visualization by group
   mjtByte  actuatorgroup[mjNGROUP];  // actuator visualization by group
+  mjtByte  flexgroup[mjNGROUP];      // flex visualization by group
   mjtByte  skingroup[mjNGROUP];      // skin visualization by group
   mjtByte  flags[mjNVISFLAG];        // visualization flags (indexed by mjtVisFlag)
-  int bvh_depth;                     // depth of the bounding volume hierarchy to be visualized
+  int      bvh_depth;                // depth of the bounding volume hierarchy to be visualized
+  int      flex_layer;               // element layer to be visualized for 3D flex
 };
 typedef struct mjvOption_ mjvOption;
 struct mjvScene_ {                // abstract scene passed to OpenGL renderer
@@ -1903,13 +2039,32 @@ struct mjvScene_ {                // abstract scene passed to OpenGL renderer
   mjvGeom* geoms;                 // buffer for geoms (ngeom)
   int*     geomorder;             // buffer for ordering geoms by distance to camera (ngeom)
 
+  // flex data
+  int      nflex;                 // number of flexes
+  int*     flexedgeadr;           // address of flex edges (nflex)
+  int*     flexedgenum;           // number of edges in flex (nflex)
+  int*     flexvertadr;           // address of flex vertices (nflex)
+  int*     flexvertnum;           // number of vertices in flex (nflex)
+  int*     flexfaceadr;           // address of flex faces (nflex)
+  int*     flexfacenum;           // number of flex faces allocated (nflex)
+  int*     flexfaceused;          // number of flex faces currently in use (nflex)
+  int*     flexedge;              // flex edge data (2*nflexedge)
+  float*   flexvert;              // flex vertices (3*nflexvert)
+  float*   flexface;              // flex faces vertices (9*sum(flexfacenum))
+  float*   flexnormal;            // flex face normals (9*sum(flexfacenum))
+  float*   flextexcoord;          // flex face texture coordinates (6*sum(flexfacenum))
+  mjtByte  flexvertopt;           // copy of mjVIS_FLEXVERT mjvOption flag
+  mjtByte  flexedgeopt;           // copy of mjVIS_FLEXEDGE mjvOption flag
+  mjtByte  flexfaceopt;           // copy of mjVIS_FLEXFACE mjvOption flag
+  mjtByte  flexskinopt;           // copy of mjVIS_FLEXSKIN mjvOption flag
+
   // skin data
   int      nskin;                 // number of skins
   int*     skinfacenum;           // number of faces in skin (nskin)
   int*     skinvertadr;           // address of skin vertices (nskin)
   int*     skinvertnum;           // number of vertices in skin (nskin)
-  float*   skinvert;              // skin vertex data (nskin)
-  float*   skinnormal;            // skin normal data (nskin)
+  float*   skinvert;              // skin vertex data (3*nskinvert)
+  float*   skinnormal;            // skin normal data (3*nskinvert)
 
   // OpenGL lights
   int      nlight;                // number of lights currently in buffer
@@ -1984,7 +2139,7 @@ struct mjvSceneState_ {
   int nbuffer;                     // size of the buffer in bytes
   void* buffer;                    // heap-allocated memory for all arrays in this struct
   int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene plugincache;            // scratch space for vis geoms inserted by plugins
+  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
 
   // fields in mjModel that are necessary to re-render a scene
   struct {
@@ -1993,6 +2148,7 @@ struct mjvSceneState_ {
     int na;
     int nbody;
     int nbvh;
+    int nbvhstatic;
     int njnt;
     int ngeom;
     int nsite;
@@ -2000,6 +2156,9 @@ struct mjvSceneState_ {
     int nlight;
     int nmesh;
     int nskin;
+    int nflex;
+    int nflexvert;
+    int nflextexcoord;
     int nskinvert;
     int nskinface;
     int nskinbone;
@@ -2013,6 +2172,7 @@ struct mjvSceneState_ {
     int nnames;
     int npaths;
     int nsensordata;
+    int narena;
 
     mjOption opt;
     mjVisual vis;
@@ -2036,7 +2196,7 @@ struct mjvSceneState_ {
 
     int* bvh_depth;
     int* bvh_child;
-    int* bvh_geomid;
+    int* bvh_nodeid;
     mjtNum* bvh_aabb;
 
     int* jnt_type;
@@ -2064,6 +2224,7 @@ struct mjvSceneState_ {
 
     mjtNum* cam_fovy;
     mjtNum* cam_ipd;
+    float* cam_intrinsic;
     float* cam_sensorsize;
 
     mjtByte* light_directional;
@@ -2076,6 +2237,28 @@ struct mjvSceneState_ {
     float* light_diffuse;
     float* light_specular;
 
+    mjtByte* flex_flatskin;
+    int* flex_dim;
+    int* flex_matid;
+    int* flex_group;
+    int* flex_vertadr;
+    int* flex_vertnum;
+    int* flex_elem;
+    int* flex_elemlayer;
+    int* flex_elemadr;
+    int* flex_elemnum;
+    int* flex_elemdataadr;
+    int* flex_shell;
+    int* flex_shellnum;
+    int* flex_shelldataadr;
+    int* flex_texcoordadr;
+    int* flex_bvhadr;
+    int* flex_bvhnum;
+    mjtNum* flex_radius;
+    float* flex_rgba;
+
+    int* hfield_pathadr;
+
     int* mesh_bvhadr;
     int* mesh_bvhnum;
     int* mesh_texcoordadr;
@@ -2102,6 +2285,9 @@ struct mjvSceneState_ {
     int* skin_bonebodyid;
     int* skin_bonevertid;
     float* skin_bonevertweight;
+    int* skin_pathadr;
+
+    int* tex_pathadr;
 
     int* mat_texid;
     mjtByte* mat_texuniform;
@@ -2200,6 +2386,7 @@ struct mjvSceneState_ {
     mjtNum* ten_length;
     mjtNum* wrap_xpos;
 
+    mjtNum* bvh_aabb_dyn;
     mjtByte* bvh_active;
     int* island_dofadr;
     int* island_dofind;
@@ -2207,8 +2394,11 @@ struct mjvSceneState_ {
     int* efc_island;
     int* tendon_efcadr;
 
+    mjtNum* flexvert_xpos;
+
     mjContact* contact;
     mjtNum* efc_force;
+    void* arena;
   } data;
 };
 typedef struct mjvSceneState_ mjvSceneState;
@@ -2271,6 +2461,7 @@ void mj_fwdAcceleration(const mjModel* m, mjData* d);
 void mj_fwdConstraint(const mjModel* m, mjData* d);
 void mj_Euler(const mjModel* m, mjData* d);
 void mj_RungeKutta(const mjModel* m, mjData* d, int N);
+void mj_implicit(const mjModel* m, mjData* d);
 void mj_invPosition(const mjModel* m, mjData* d);
 void mj_invVelocity(const mjModel* m, mjData* d);
 void mj_invConstraint(const mjModel* m, mjData* d);
@@ -2286,6 +2477,7 @@ void mj_checkAcc(const mjModel* m, mjData* d);
 void mj_kinematics(const mjModel* m, mjData* d);
 void mj_comPos(const mjModel* m, mjData* d);
 void mj_camlight(const mjModel* m, mjData* d);
+void mj_flex(const mjModel* m, mjData* d);
 void mj_tendon(const mjModel* m, mjData* d);
 void mj_transmission(const mjModel* m, mjData* d);
 void mj_crb(const mjModel* m, mjData* d);
@@ -2360,6 +2552,9 @@ mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
                   const mjtNum pnt[3], const mjtNum vec[3]);
 mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
                    const mjtNum pnt[3], const mjtNum vec[3], int geomtype);
+mjtNum mju_rayFlex(const mjModel* m, const mjData* d, int flex_layer, mjtByte flg_vert,
+                   mjtByte flg_edge, mjtByte flg_face, mjtByte flg_skin, int flexid,
+                   const mjtNum* pnt, const mjtNum* vec, int vertid[1]);
 mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
                    const mjtNum pnt[3], const mjtNum vec[3], int vertid[1]);
 void mjv_defaultCamera(mjvCamera* cam);
@@ -2394,7 +2589,8 @@ void mjv_applyPerturbForce(const mjModel* m, mjData* d, const mjvPerturb* pert);
 mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera* cam2);
 int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
                mjtNum aspectratio, mjtNum relx, mjtNum rely,
-               const mjvScene* scn, mjtNum selpnt[3], int geomid[1], int skinid[1]);
+               const mjvScene* scn, mjtNum selpnt[3],
+               int geomid[1], int flexid[1], int skinid[1]);
 void mjv_defaultOption(mjvOption* opt);
 void mjv_defaultFigure(mjvFigure* fig);
 void mjv_initGeom(mjvGeom* geom, int type, const mjtNum size[3],

doc/index.rst

@@ -14,6 +14,11 @@
    programming/index.rst
    APIreference/index.rst
    python
+   MJX <mjx>
    unity
    models
    changelog
+
+
+.. sidebar-links::
+   :github:

doc/mjx.rst

@@ -0,0 +1,359 @@
+.. _Mjx:
+
+================
+MuJoCo XLA (MJX)
+================
+
+Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the
+`mjx <https://github.com/google-deepmind/mujoco/tree/main/mjx>`__ directory.  MJX allows MuJoCo to run on compute
+hardware supported by the `XLA <https://www.tensorflow.org/xla>`__ compiler via the
+`JAX <https://github.com/google/jax#readme>`__ framework.  MJX runs on a
+`all platforms supported by JAX <https://jax.readthedocs.io/en/latest/installation.html#supported-platforms>`__: Nvidia
+and AMD GPUs, Apple Silicon, and `Google Cloud TPUs <https://cloud.google.com/tpu>`__.
+
+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 <Mainsimulation>` is applicable to both libraries, we indicate features
+unsupported by MJX in the :ref:`notes <MjxFeatureParity>` below.
+
+MJX is distributed as a separate package called ``mujoco-mjx`` on `PyPI <https://pypi.org/project/mujoco-mjx>`__.
+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 <https://github.com/google/brax/tree/main/brax/generalized>`__
+in Google's `Brax <https://github.com/google/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 <https://github.com/google/brax/tree/main/brax/generalized>`__ 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 <https://pypi.org/project/mujoco-mjx/>`__:
+
+.. 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 <mjxFeatureParity>` 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 ``mjx.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 <https://peps.python.org/pep-0008/>`__-compliant names.  Most of the :ref:`main simulation <Mainsimulation>` and
+some of the :ref:`sub-components <Subcomponents>` for forward simulation are available from the top-level ``mjx`` module.
+
+MJX functions are not `JIT compiled <https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html>`__ 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 <MjxExample>` 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"""
+   <mujoco>
+     <worldbody>
+       <body>
+         <freejoint/>
+         <geom size=".15" mass="1" type="sphere"/>
+       </body>
+     </worldbody>
+   </mujoco>
+   """
+
+   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: 2 5
+   :header-rows: 1
+
+   * - Category
+     - Feature
+   * - Dynamics
+     - :ref:`Forward <mj_forward>`
+   * - :ref:`Joint <mjtJoint>`
+     - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
+   * - :ref:`Transmission <mjtTrn>`
+     - ``TRN_JOINT``, ``TRN_SITE``
+   * - :ref:`Actuator Dynamics <mjtDyn>`
+     - ``NONE``, ``INTEGRATOR``, ``FILTER``, ``FILTEREXACT``
+   * - :ref:`Actuator Gain <mjtGain>`
+     - ``FIXED``, ``AFFINE``
+   * - :ref:`Actuator Bias <mjtBias>`
+     - ``NONE``, ``AFFINE``
+   * - :ref:`Geom <mjtGeom>`
+     - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
+   * - :ref:`Constraint <mjtConstraint>`
+     - ``EQUALITY``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
+   * - :ref:`Equality <mjtEq>`
+     - ``CONNECT``, ``WELD``, ``JOINT``
+   * - :ref:`Integrator <mjtIntegrator>`
+     - ``EULER``, ``RK4``
+   * - :ref:`Cone <mjtCone>`
+     - ``PYRAMIDAL``
+   * - :ref:`Condim <coContact>`
+     - 3
+   * - :ref:`Solver <mjtSolver>`
+     - ``CG``, ``NEWTON``
+   * - Fluid Model
+     - :ref:`flInertia`
+
+The following features are **in development** and coming soon:
+
+.. list-table::
+   :width: 90%
+   :align: left
+   :widths: 2 5
+   :header-rows: 1
+
+   * - Category
+     - Feature
+   * - Dynamics
+     - :ref:`Inverse <mj_inverse>`
+   * - :ref:`Transmission <mjtTrn>`
+     - ``TRN_TENDON``
+   * - :ref:`Actuator Dynamics <mjtDyn>`
+     - ``MUSCLE``
+   * - :ref:`Actuator Gain <mjtGain>`
+     - ``MUSCLE``
+   * - :ref:`Actuator Bias <mjtBias>`
+     - ``MUSCLE``
+   * - :ref:`Tendon Wrapping <mjtWrap>`
+     - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER``
+   * - :ref:`Geom <mjtGeom>`
+     - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``
+   * - :ref:`Constraint <mjtConstraint>`
+     - :ref:`Frictionloss <coFriction>`, ``CONTACT_FRICTIONLESS``, ``CONTACT_ELLIPTIC``, ``FRICTION_DOF``
+   * - :ref:`Integrator <mjtIntegrator>`
+     - ``IMPLICIT``, ``IMPLICITFAST``
+   * - :ref:`Cone <mjtCone>`
+     - ``ELLIPTIC``
+   * - :ref:`Condim <coContact>`
+     - 1, 4, 6
+   * - Fluid Model
+     - :ref:`flEllipsoid`
+   * - :ref:`Tendons <tendon>`
+     - :ref:`Spatial <tendon-spatial>`, :ref:`Fixed <tendon-fixed>`
+   * - :ref:`Equality <mjtEq>`
+     - ``TENDON``
+   * - :ref:`Sensors <mjtSensor>`
+     - All except ``PLUGIN``, ``USER``
+
+The following features are **unsupported**:
+
+.. list-table::
+   :width: 90%
+   :align: left
+   :widths: 2 5
+   :header-rows: 1
+
+   * - Category
+     - Feature
+   * - :ref:`Transmission <mjtTrn>`
+     - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_BODY``
+   * - :ref:`Actuator Dynamics <mjtDyn>`
+     - ``USER``
+   * - :ref:`Actuator Gain <mjtGain>`
+     - ``USER``
+   * - :ref:`Actuator Bias <mjtBias>`
+     - ``USER``
+   * - :ref:`Solver <mjtSolver>`
+     - ``PGS``
+   * - :ref:`Sensors <mjtSensor>`
+     - ``PLUGIN``, ``USER``
+   * - :ref:`Geom <mjtGeom>`
+     - ``SDF``
+
+.. _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 <https://en.wikipedia.org/wiki/Single_instruction,_multiple_data>`__.  This specialization is useful
+for machine learning workloads such as `reinforcement learning <https://en.wikipedia.org/wiki/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.
+
+Collisions between large meshes
+  MJX supports collisions between convex mesh geometries. However the convex collision algorithms
+  in MJX are implemented differently than in MuJoCo. MJX uses a branchless version of the
+  `Separating Axis Test <https://ubm-twvideo01.s3.amazonaws.com/o1/vault/gdc2013/slides/822403Gregorius_Dirk_TheSeparatingAxisTest.pdf>`__
+  (SAT) to determine if geometries are colliding with convex meshes, while MuJoCo uses the Minkowski Portal Refinement (MPR)
+  algorithm as implemented in `libccd <https://github.com/danfis/libccd>`__.
+  SAT works well for smaller meshes but suffers in both runtime and memory for larger meshes.
+
+  For
+  collisions between convex meshes and primitives (spheres, capsules, planes), use **3000 vertices or less** for your convex meshes.
+  For collisions between convex meshes and other convex meshes, use **30 vertices or less**.
+  With careful
+  tuning, MJX can simulate scenes with mesh collisions -- see the MJX
+  `shadow hand <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/shadow_hand>`__
+  config for an example. Speeding up mesh collision detection is an active area of development for MJX.
+
+Large, complex scenes with many contacts
+  Accelerators exhibit poor performance for
+  `branching code <https://aschrein.github.io/jekyll/update/2019/06/13/whatsup-with-my-branches-on-gpu.html#tldr>`__.
+  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 humanoid bodies,
+  varied from 1 to 10. We simulate this scene using CPU MuJoCo on an Apple M3 Max and a 64-core AMD 3995WX and time
+  it using :ref:`testspeed<saTestspeed>`, using ``2 x numcore`` threads. We time the MJX simulation on an Nvidia
+  A100 GPU using a batch size of 8192 and an 8-chip
+  `v5 TPU <https://cloud.google.com/blog/products/compute/announcing-cloud-tpu-v5e-and-a3-gpus-in-ga>`__
+  machine using a batch size of 16384. Note the vertical scale is logarithmic.
+
+  .. figure:: images/mjx/SPS.svg
+     :width: 95%
+     :align: center
+
+  The values for a single humanoid (leftmost datapoints) for the four timed architectures are **650K**, **1.8M**,
+  **950K** and **2.7M** steps per second, respectively. Note that as we increase the number of humanoids (which
+  increases the number of potential contacts in a scene), MJX throughput decreases more rapidly than MuJoCo.
+
+.. _MjxPerformance:
+
+Performance tuning
+==================
+
+For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values:
+
+:ref:`option` element
+  The ``iterations`` and ``ls_iterations`` attributes---which control solver and linesearch iterations, respectively---
+  should be brought down to just low enough that the simulation remains stable.  Accurate solver forces are not so
+  important in reinforcement learning in which domain randomization is often used to add noise to physics for sim-to-real.
+  The ``NEWTON`` :ref:`Solver <mjtSolver>` often delivers reasonable convergence with one solver iteration, and performs
+  well on GPU.  ``CG`` is currently a better choice for TPU.
+
+:ref:`contact-pair` element
+  Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider
+  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 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v4.py>`__ 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.
+
+:ref:`option-jacobian` element
+  Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized
+  hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as
+  they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or
+  greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the
+  Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10%
+  to 20%, as long as the dense matrices can fit on the device.

doc/modeling.rst

@@ -190,7 +190,7 @@ element, it must be undefined in the active defaults class.
 
 A final twist here is actuators. They are different because some of the actuator-related elements are actually
 shortcuts, and shortcuts interact with the defaults setting mechanism in a non-obvious way. This is explained in the
-:ref:`Actuator shortcuts <CActuator>` section below.
+:ref:`Actuator shortcuts <CActShortcuts>` section below.
 
 .. _CFrame:
 
@@ -198,22 +198,27 @@ Coordinate frames
 ~~~~~~~~~~~~~~~~~
 
 The positions and orientations of all elements defined in the kinematic tree are expressed in local coordinates,
-relative to the parent body for bodies, and relative to the body that owns the element for geoms, joints, sites, cameras
-and lights.
+relative to the parent body for bodies, and relative to the body that contains the element for geoms, joints, sites,
+cameras and lights.
 
 A related attribute is :ref:`compiler/angle<compiler-angle>`. It specifies whether angles in the MJCF file are expressed
 in degrees or radians (after compilation, angles are always expressed in radians).
 
+Positions are specified using
+
+:at:`pos`: :at-val:`real(3), "0 0 0"`
+   Position relative to parent.
+
 .. _COrientation:
 
 Frame orientations
-~~~~~~~~~~~~~~~~~~
+^^^^^^^^^^^^^^^^^^
 
 Several model elements have right-handed spatial frames associated with them. These are all the elements defined in the
 kinematic tree except for joints. A spatial frame is defined by its position and orientation. Specifying 3D positions is
 straightforward, but specifying 3D orientations can be challenging. This is why MJCF provides several alternative
-mechanisms. No matter which mechanism the user chooses, the frame orientation is always represented as a unit quaternion
-after compilation. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)`
+mechanisms. No matter which mechanism the user chooses, the frame orientation is always converted internally to a unit
+quaternion. Recall that a 3D rotation by angle :math:`a` around axis given by the unit vector :math:`(x, y, z)`
 corresponds to the quaternion :math:`(\cos(a/2), \: \sin(a/2) \cdot (x, y, z))`. Also recall that every 3D orientation
 can be uniquely specified by a single 3D rotation by some angle around some axis.
 
@@ -259,14 +264,15 @@ available in :ref:`option <option>`; it can be used to change all contact-relate
 experiment interactively with parameter settings or implement continuation methods for numerical optimization.
 
 Here we focus on a single scalar constraint. Using slightly different notation from the Computation chapter, let
-:math:`a_1` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in
+:math:`\ac` denote the acceleration, :math:`v` the velocity, :math:`r` the position or residual (defined as 0 in
 friction dimensions), :math:`k` and :math:`b` the stiffness and damping of the virtual spring used to define the
-reference acceleration :math:`a_{\rm ref} = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`a_0` the
+reference acceleration :math:`\ar = -b v - k r`. Let :math:`d` be the constraint impedance, and :math:`\au` the
 acceleration in the absence of constraint force. Our earlier analysis revealed that the dynamics in constraint space are
 approximately
 
 .. math::
-   a_1 + d \cdot (b v + k r) = (1 - d)\cdot a_0
+   \ac + d \cdot (b v + k r) = (1 - d)\cdot \au
+   :label: eq:constraint
 
 Again, the parameters that are under the user's control are :math:`d, b, k`. The remaining quantities are functions of
 the system state and are computed automatically at each time step.
@@ -278,15 +284,16 @@ Impedance
 
 We begin by explaining the constraint impedance :math:`d`.
 
-.. admonition:: Intuitive description
+.. admonition:: Intuitive description of the **impedance**
 
-   The *impedance* :math:`d \in (0, 1)` determines a constraint's **ability to generate force**.
+   The *impedance* :math:`d \in (0, 1)` corresponds to a constraint's **ability to generate force**.
    Small values of :math:`d` correspond to weak constraints while large values of :math:`d`
-   correspond to strong constraints. Impedance is set using the :at:`solimp` attribute.
+   correspond to strong constraints. The impedance affects the constraint at all times, in particular when the system is
+   at rest. Impedance is set using the :at:`solimp` attribute.
 
 Recall that :math:`d` must lie between 0 and 1; internally MuJoCo clamps it to the range [:ref:`mjMINIMP mjMAXIMP
 <glNumeric>`] which is currently set to [0.0001 0.9999]. It causes the solver to interpolate between the unforced
-acceleration :math:`a_0` and reference acceleration :math:`a_{\rm ref}`. The user can set :math:`d` to a constant, or
+acceleration :math:`\au` and reference acceleration :math:`\ar`. The user can set :math:`d` to a constant, or
 take advantage of its interpolating property and make it position-dependent, i.e., a function of the constraint
 violation :math:`r`. Position-dependent impedance can be used to model soft contact layers around objects, or define
 equality constraints that become stronger with larger violation (so as to approximate backlash, for example). The shape
@@ -311,7 +318,15 @@ of the function :math:`d(r)` is determined by the element-specific parameter vec
    units of :math:`\text{width}`. Note that when :math:`\text{power}` is 1, the function is linear regardless of the
    :math:`\text{midpoint}`.
 
-   |image0|
+   .. image:: images/modeling/impedance.png
+      :width: 600px
+      :align: center
+      :class: only-light
+
+   .. image:: images/modeling/impedance_dark.png
+      :width: 600px
+      :align: center
+      :class: only-dark
 
    These plots show the impedance :math:`d(r)` on the vertical axis, as a function of the constraint violation :math:`r`
    on the horizontal axis.
@@ -333,15 +348,23 @@ Reference
 ^^^^^^^^^
 
 Next we explain the setting of the stiffness :math:`k` and damping :math:`b` which control the reference acceleration
-:math:`a_{\rm ref}`.
+:math:`\ar`.
 
-.. admonition:: Intuitive description
+.. admonition:: Intuitive description of the **reference acceleration**
 
-   The *reference acceleration* :math:`a_{\rm ref}` determines **what the constraint is trying to achieve** (as opposed
-   to how well it can achieve it). This acceleration is defined by two numbers, a stiffness :math:`k` and damping
-   :math:`b` which can be set directly or re-parameterized as the time-constant and damping ratio of a
-   mass-spring-damper system (a `harmonic oscillator <https://en.wikipedia.org/wiki/Harmonic_oscillator>`__).
-   The reference acceleration is controlled by the :at:`solref` attribute.
+   The *reference acceleration* :math:`\ar` determines the **motion that constraint is trying to achieve** in order to
+   rectify violation. Imagine a body dropped onto the plane. Upon impact the constraint will generate a normal force
+   which attempts to rectify the penetration using a particular motion; this motion is the reference acceleration.
+
+   Another way of understanding the reference acceleration is to think of the unmodeled deformation variables
+   described in the :ref:`Computation chapter<soPrimal>`. Imagine two bodies pressed together, leading to deformation at
+   the contact. Now pull the bodies apart very quickly; the motion of the deformation as it settles into its undeformed
+   state is the reference acceleration.
+
+   This acceleration is defined by two numbers, a stiffness :math:`k` and damping :math:`b` which can be set directly or
+   re-parameterized as the time-constant and damping ratio of a mass-spring-damper system (a `harmonic oscillator
+   <https://en.wikipedia.org/wiki/Harmonic_oscillator>`__). The reference acceleration is controlled by the :at:`solref`
+   attribute.
 
 There are two formats for this attribute, determined by the sign of the numbers. If both numbers are positive the
 specification is considered to be in the :math:`(\text{timeconst}, \text{dampratio})` format. If negative it is in the
@@ -376,7 +399,12 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math:
    can go unstable. This is enforced internally, unless the :ref:`refsafe<option-flag-refsafe>` attribute of :ref:`flag
    <option-flag>` is set to false. The :math:`\text{dampratio}` parameter would normally be set to 1, corresponding to
    critical damping. Smaller values result in under-damped or bouncy constraints, while larger values result in
-   over-damped constraints.
+   over-damped constraints. Combining the above formula with :eq:`eq:constraint`, we can derive the following result.
+   If the reference acceleration is given using the positive number format and the impedance is constant
+   :math:`d = d_0 = d_\text{width}`, then the penetration depth at rest is
+
+   .. math::
+      r = \au \cdot (1 - d) \cdot \text{timeconst}^2 \cdot \text{dampratio}^2
 
    Next we describe the direct format where the two numbers are :math:`(-\text{stiffness}, -\text{damping})`. This
    allows direct control over restitution in particular. We still apply some scaling so that the same numbers can be
@@ -386,9 +414,15 @@ and the damping ratio is ignored. Equivalently, in the direct format, the :math:
    .. math::
       \begin{aligned}
       b &= \text{damping} / d_\text{width} \\
-      k &= \text{stiffness} / d_\text{width}^2 \\
+      k &= \text{stiffness} \cdot d(r) / d_\text{width}^2 \\
       \end{aligned}
 
+   Similarly to the above derivation, if the reference acceleration is given using the negative number format and the
+   impedance is constant, then the penetration depth at rest is
+
+   .. math::
+      r = \au \cdot (1 - d) \cdot \text{stiffness}
+
 .. tip::
    In the positive-value default format, the :math:`\text{timeconst}` parameter controls constraint **softness**.
    It is specified in units of time and means "how quickly is the constraint trying to resolve the violation". Larger
@@ -441,7 +475,7 @@ solref, solimp
    If one of the two geoms has higher priority, its solref and solimp parameters are used. If both geoms have the same
    priority, the weighted average is used. The weights are proportional to the solmix attributes, i.e., weight1 =
    solmix1 / (solmix1 + solmix2) and similarly for weight2. There is one important exception to this weighted averaging
-   rule. If solref for either geom is non-positive, i.e., it relies on the new direct format introduced in MuJoCo 2.0,
+   rule. If solref for either geom is non-positive, i.e., it relies on the direct format,
    then the element-wise minimum is used regardless of solmix. This is because averaging solref parameters in different
    formats would be meaningless.
 
@@ -555,10 +589,44 @@ general guidelines and observations:
    setup operation for the main PGS and Noslip PGS is the same, thus the setup cost is paid only once when both are
    enabled.
 
-.. _CActuator:
+.. _CActuators:
 
-Actuator shortcuts
-~~~~~~~~~~~~~~~~~~
+Actuators
+~~~~~~~~~
+
+This section describes various aspects of using actuators in MuJoCo. See the :ref:`Actuation model <geActuation>`
+regarding the computational model.
+
+.. _CActDisable:
+
+Group disable
+^^^^^^^^^^^^^
+
+The :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` attribute, which can be changed at runtime by setting the
+:ref:`mjOption.disableactuator<mjOption>` integer bitfield, allows the user to disable sets of actuators according to
+their :ref:`group<actuator-general-group>`. This feature is convenient when one would like to use multiple types of
+actuators for the same kinematic tree. For example consider a robot with firmware that supports mutiple control modes
+e.g., torque-control and position-control. In this case, one can define both types of actuators in the same MJCF
+model, assigning one type of actuator to group 0 and the other to group 1.
+
+.. youtube:: H9qG9Zf2W44
+   :align: right
+   :width: 40%
+
+The :ref:`actuatorgroupdisable<option-actuatorgroupdisable>` MJCF attribute selects which groups are disabled by
+default, and :ref:`mjOption.disableactuator<mjOption>` can be set at runtime to switch the active set. Note that the
+total number of actuators ``mjModel.nu`` remains unchanged, as do the actuator indices, so it is up to the user to know
+that the respective ``mjData.ctrl`` values of disabled actuators will be ignored and produce no force. `This example
+model <https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
+has three actuator groups which can be toggled at runtime in the :ref:`simulate<saSimulate>` interactive viewer.
+See `example model
+<https://github.com/google-deepmind/mujoco/blob/main/test/engine/testdata/actuation/actuator_group_disable.xml>`__
+and associated screen-capture on the right.
+
+.. _CActShortcuts:
+
+Shortcuts
+^^^^^^^^^
 
 As explained in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo offers a flexible
 actuator model with transmission, activation dynamics and force generation components that can be specified
@@ -593,8 +661,8 @@ explicitly.
 
 .. _CForceRange:
 
-Actuator force clamping
-~~~~~~~~~~~~~~~~~~~~~~~
+Force limits
+^^^^^^^^^^^^
 
 Actuator forces are usually limited between lower and upper bounds. These limits can be enforced in three ways:
 
@@ -633,62 +701,14 @@ Force clamping at joint input with :ref:`joint/actuatorfrcrange<body-joint-actua
 
 The three clamping options above are non-exclusive and can be combined as required.
 
-.. _CActRange:
-
-Activation clamping
-~~~~~~~~~~~~~~~~~~~
-
-As described in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo supports actuators
-with internal dynamics whose states are called "activations". One useful application of these stateful actuators is the
-"integrated-velocity" actuator, implemented by the :ref:`intvelocity<actuator-intvelocity>` shortcut. Different from the
-:ref:`pure velocity<actuator-velocity>` actuators, which implement direct feedback on transmission target's velocity,
-*integrated-velocity* actuators couple an *integrator* with a *position-feedback* actuator. In this case the semantics
-of the activation state are "the setpoint of the position actuator", and the semantics of the control signal are "the
-velocity of the setpoint of the position actuator". Note that in real robotic systems this integrated-velocity actuator
-is the most common implementation of actuators with velocity semantics, rather than pure feedback on velocity which is
-often quite unstable (both in real life and in simulation).
-
-In the case of integrated-velocity actuators, it is often desirable to *clamp* the activation state, since otherwise the
-position target would keep integrating beyond the joint limits, leading to loss of controllabillity. To see the effect
-of activation clamping, load the example model below:
-
-.. code-block:: xml
-
-   <mujoco>
-     <default>
-       <joint axis="0 0 1" limited="true" range="-90 90" damping="0.3"/>
-       <geom size=".1 .1 .1" type="box"/>
-     </default>
-
-     <worldbody>
-       <body>
-         <joint name="joint1"/>
-         <geom/>
-       </body>
-       <body pos=".3 0 0">
-         <joint name="joint2"/>
-         <geom/>
-       </body>
-     </worldbody>
-
-     <actuator>
-       <general name="unclamped" joint="joint1" gainprm="1" biastype="affine"
-         biasprm="0 -1" dyntype="integrator"/>
-       <intvelocity name="clamped" joint="joint2" actrange="-1.57 1.57"/>
-     </actuator>
-   </mujoco>
-
-Note that the :at:`actrange` attribute is always specified in native units (radians), even though the joint range
-can be either in degrees (the default) or radians, depending on the :ref:`compiler/angle <compiler>` attribute.
-
 .. _CLengthRange:
 
-Actuator length range
-~~~~~~~~~~~~~~~~~~~~~
+Length range
+^^^^^^^^^^^^
 
-As of MuJoCo 2.0, the field mjModel.actuator_lengthrange contains the range of feasible actuator lengths (or more
-precisely, lengths of the actuator's transmission). This is needed to simulate :ref:`muscle actuators <CMuscle>` as
-explained below. Here we focus on what actuator_lengthrange means and how to set it.
+The field ``mjModel.actuator_lengthrange`` contains the range of feasible actuator lengths (or more
+precisely, lengths of the actuator's transmission). This is needed to simulate :ref:`muscle actuators <CMuscle>`.
+Here we focus on what actuator_lengthrange means and how to set it.
 
 Unlike all other fields of mjModel which are exact physical or geometric quantities, actuator_lengthrange is an
 approximation. Intuitively it corresponds to the minimum and maximum length that the actuator's transmission can reach
@@ -732,12 +752,69 @@ practice length ranges will almost always be used with muscle actuators attached
 joint limits defined in the model, effectively limiting the lengths of the muscle actuators. If you get a convergence
 error in such a model, the most likely explanation is that you forgot to include joint limits.
 
+.. _CActivation:
+
+Stateful actuators
+^^^^^^^^^^^^^^^^^^
+
+As described in the :ref:`Actuation model <geActuation>` section of the Computation chapter, MuJoCo supports actuators
+with internal dynamics whose states are called "activations".
+
+.. _CActRange:
+
+Activation limits
+'''''''''''''''''
+
+One useful application of stateful actuators is the
+"integrated-velocity" actuator, implemented by the :ref:`intvelocity<actuator-intvelocity>` shortcut. Different from the
+:ref:`pure velocity<actuator-velocity>` actuators, which implement direct feedback on transmission target's velocity,
+*integrated-velocity* actuators couple an *integrator* with a *position-feedback* actuator. In this case the semantics
+of the activation state are "the setpoint of the position actuator", and the semantics of the control signal are "the
+velocity of the setpoint of the position actuator". Note that in real robotic systems this integrated-velocity actuator
+is the most common implementation of actuators with velocity semantics, rather than pure feedback on velocity which is
+often quite unstable (both in real life and in simulation).
+
+In the case of integrated-velocity actuators, it is often desirable to *clamp* the activation state, since otherwise the
+position target would keep integrating beyond the joint limits, leading to loss of controllabillity. To see the effect
+of activation clamping, load the example model below:
+
+.. collapse:: Example model with activation limits
+
+   .. code-block:: xml
+
+      <mujoco>
+      <default>
+         <joint axis="0 0 1" limited="true" range="-90 90" damping="0.3"/>
+         <geom size=".1 .1 .1" type="box"/>
+      </default>
+
+      <worldbody>
+         <body>
+            <joint name="joint1"/>
+            <geom/>
+         </body>
+         <body pos=".3 0 0">
+            <joint name="joint2"/>
+            <geom/>
+         </body>
+      </worldbody>
+
+      <actuator>
+         <general name="unclamped" joint="joint1" gainprm="1" biastype="affine"
+            biasprm="0 -1" dyntype="integrator"/>
+         <intvelocity name="clamped" joint="joint2" actrange="-1.57 1.57"/>
+      </actuator>
+      </mujoco>
+
+Note that the :at:`actrange` attribute is always specified in native units (radians), even though the joint range
+can be either in degrees (the default) or radians, depending on the :ref:`compiler/angle <compiler>` attribute.
+
 .. _CMuscle:
 
-Muscle actuators
-~~~~~~~~~~~~~~~~
+Muscles
+'''''''
 
-As of MuJoCo 2.0, we provide a set of tools for modeling biological muscles. Users who want to add muscles with minimum
+We provide a set of tools for modeling biological muscles. Users who want to add muscles with minimum
 effort can do so with a single line of XML in the actuator section:
 
 .. code-block:: xml
@@ -796,7 +873,15 @@ The advantage of the scaled quantities is that all muscles behave similarly in t
 captured by the Force-Length-Velocity (:math:`\text{\small FLV}`) function measured in many experimental papers. We
 approximate this function as follows:
 
-|image1|
+.. image:: images/modeling/musclemodel.png
+   :width: 650px
+   :align: center
+   :class: only-light
+
+.. image:: images/modeling/musclemodel_dark.png
+   :width: 650px
+   :align: center
+   :class: only-dark
 
 The function is in the form:
 
@@ -838,7 +923,15 @@ Before embarking on a mission to design more accurate :math:`\text{\small FLV}`
 operating range of the muscle has a bigger effect than the shape of the :math:`\text{\small FLV}` function, and in many
 cases this parameter is unknown. Below is a graphical illustration:
 
-|image2|
+.. image:: images/modeling/musclerange.png
+   :width: 500px
+   :align: center
+   :class: only-light
+
+.. image:: images/modeling/musclerange_dark.png
+   :width: 500px
+   :align: center
+   :class: only-dark
 
 This figure format is common in the biomechanics literature, showing the operating range of each muscle superimposed on
 the normalized :math:`\text{FL}` curve (ignore the vertical displacement). Our default range is shown in black. The blue
@@ -937,13 +1030,13 @@ be 0. This effect can be approximated by scaling down the muscle force and also
 
 Tendon wrapping is also more limited in MuJoCo. We allow spheres and infinite cylinders as wrapping objects, and require
 two wrapping objects to be separated by a fixed site in the tendon path. This is to avoid the need for iterative
-computations of tendon paths. As of MuJoCo 2.0 we also allow "side sites" to be placed inside the sphere or cylinder,
+computations of tendon paths. We also allow "side sites" to be placed inside the sphere or cylinder,
 which causes an inverse wrap: the tendon path is constrained to pass through the object instead of going around it. This
 can replace torus wrapping objects used in OpenSim to keep the tendon path within a given area. Overall, tendon wrapping
 is the most challenging part of converting an OpenSim model to a MuJoCo model, and requires some manual work. On the
 bright side, there is a small number of high-quality OpenSim models in use, so once they are converted we are done.
 
-Below we illustrate the four types of tendon wrapping available in MuJoCo 2.0. Note that the curved sections of the
+Below we illustrate the four types of tendon wrapping available. Note that the curved sections of the
 wrapping tendons are rendered as straight, but the geometry pipeline works with the actual curves and computes their
 lengths and moments analytically:
 
@@ -981,7 +1074,7 @@ Cameras
 Besides the default, user-controllable, free camera, "fixed" cameras can be attached to the kinematic tree.
 
 Extrinsics
-   By default, camera frames are attached to the parent body. The optional :ref:`mode<body-camera-mode>` and
+   By default, camera frames are attached to the containing body. The optional :ref:`mode<body-camera-mode>` and
    :ref:`target<body-camera-target>` attributes can be used to specify camera that track (move with) or target (look at)
    a body or subtree. Cameras look towards the negative Z axis of the camera frame, while positive X and Y correspond to
    *right* and *up* in the image plane, respectively.
@@ -1003,8 +1096,7 @@ Intrinsics
 Composite objects
 ~~~~~~~~~~~~~~~~~
 
-Composite objects were introduced in MuJoCo 2.0, along with solver optimizations to speed up the simulation of such
-objects. They are not new model elements. Instead, they are (large) collections of existing elements designed to
+Composite objects are not new model elements. Instead, they are (large) collections of existing elements designed to
 simulate particle systems, ropes, cloth, and soft bodies. These collections are generated by the model compiler
 automatically. The user configures the automatic generator on a high level, using the new XML element
 :ref:`composite <body-composite>` and its attributes and sub-elements, as described in the XML reference
@@ -1026,7 +1118,7 @@ of these equality constraints can be adjusted by the user, thereby adjusting the
 composite objects.
 
 In addition to setting up the physics, the composite object generator creates suitable rendering. 2D and 3D objects
-can be rendered as :ref:`skins <asset-skin>` which are also new in MuJoCo 2.0. The skin is generated
+can be rendered as :ref:`skins <asset-skin>`. The skin is generated
 automatically, and can be textured as well as subdivided using bi-cubic interpolation. The actual physics and in
 particular the collision detection are based on the element bodies and their geoms, while the skin is purely a
 visualization object. Yet in most situations we prefer to look at the skin representation. To facilitate this, the
@@ -1034,13 +1126,12 @@ generator places all geoms, sites and tendons in group 3 whose visualization is
 a 2D grid for example, you will see a continuous flexible surface and not a collection of spheres connected with
 tendons. However when fine-tuning the model and trying to understand the physics behind it, it is useful to be able to
 render the spheres and tendons. To switch the rendering style, disable the rendering of skins and enable group 3 for
-geoms and tendons (note that starting with MuJoCo 2.0 we have added a group property to sites, tendons and joints in
-addition to geoms).
+geoms and tendons.
 
 We have designed the composite object generator to have intuitive high-level controls as much as possible, but at the
 same time it exposes a large number of options that interact with each other and can profoundly affect the resulting
 physics. So at some point users should read the :ref:`reference documentation <body-composite>` carefully.
-As a quick start though, MuJoCo 2.0 comes with an example of each composite object type. Below we go over these
+As a quick start though, MuJoCo comes with an example of each composite object type. Below we go over these
 examples and explain the less obvious aspects. In all examples we have a static scene which is included in the model,
 followed by a single composite object. The static scene has a mocap body (large capsule) that can be moved around with
 the mouse to probe the behavior of the system. The XML snippets below are just the definition of the composite object;
@@ -1069,6 +1160,14 @@ has 1000 bodies (each with a geom), 3000 degrees of freedom and around 1000 acti
 takes around 1 ms on a single core of a modern processor. As with most other MuJoCo models, the soft constraints allow
 simulation at much larger timesteps (this model is stable at 30 ms timestep and even higher).
 
+Particles are also compatible with the passive forces 2D and 3D plugins, discussed in the :ref:`deformable
+<CDeformable>` section. However, collisions are limited to the particle themselves and not to the whole boundary of the
+skin that encloses them. This makes contacts very fast but does not guarantee that all penetrations can be avoided. For
+a more complete treatment, see again the :ref:`deformable <CDeformable>` section, which outlines how to use
+:ref:`flexcomp<body-flexcomp>` to create such an object. It is easy to port models create with composite particles to
+flex, see the folder `elasticity/ <https://github.com/google-deepmind/mujoco/tree/main/model/plugin/elasticity>`__ for
+several examples.
+
 **1D grid**.
 
 |image6| |image7|
@@ -1111,59 +1210,51 @@ coordinates. The plot on the right shows a cloth pinned to the world body at the
 capsule probe. The skin on the right is subdivided using bi-cubic interpolation, which increases visual quality in the
 absence of textures. When textures are present (left) the benefits of subdivision are less visible.
 
-**Rope and loop**.
+**Cable**.
 
-|image10| |image11|
+|coil|
 
 .. code-block:: xml
 
-   <body name="B10" pos="0 0 1">
-     <freejoint/>
-     <composite type="rope" count="21 1 1" spacing="0.04" offset="0 0 2">
-       <joint kind="main" damping="0.005"/>
-       <geom type="capsule" size=".01 .015" rgba=".8 .2 .1 1"/>
-     </composite>
-   </body>
+   <extension>
+      <plugin plugin="mujoco.elasticity.cable"/>
+   </extension>
 
-The remaining composite object types create kinematic trees of element bodies, and the parent body becomes the root of
-the tree. This is why :el:`composite` appears inside a moving body, and not inside the world body as in particle and
-grid objects. If it appeared inside the world body, the root of the composite object would not move. Unlike grids and
-particles, the orientation of the element bodies here can change. The kinematic tree is constructed using (mostly)
-hinge joints. In the case of rope and loop objects illustrated here, the tree is a chain. Note the naming of the
-parent body. This name must correspond to one of the automatically-generated names of the element bodies. This
-mechanism is used to specify where the composite object should attach to the parent. Compared to 1D grids, the rope
-and loop are less jittery and can use capsule and ellipsoid geoms in addition to spheres (thus filling the gaps for
-collision detection). However this comes at a price. Because we have long kinematic chains, the resulting differential
-equations become stiff and can no longer be integrated at large timesteps. The examples we provide illustrate
-comfortable timesteps where the models are stable. The rope can be easily tied into a knot using mouse perturbations,
-as shown in the left plot. Using a larger number of smaller elements makes knots and other manipulations even easier.
-The loop is similar to a rope but the first and last element bodies are connected with an equality constraint.
+   <worldbody>
+      <composite prefix="actuated" type="cable" curve="cos(s) sin(s) s" count="41 1 1"
+                 size="0.25 .1 4" offset="0.25 0 .05" initial="none">
+         <plugin plugin="mujoco.elasticity.cable">
+            <!--Units are in Pa (SI)-->
+            <config key="twist" value="5e8"/>
+            <config key="bend" value="15e8"/>
+            <config key="vmax" value="0"/>
+         </plugin>
+         <joint kind="main" damping="0.15" armature="0.01"/>
+         <geom type="capsule" size=".005" rgba=".8 .2 .1 1"/>
+      </composite>
+   </worldbody>
+
+The cable simulates an inextensible elastic 1D object having twist and bending stiffness. It is discretized using a
+sequence of capsules or boxes. Its stiffness and inertia properties are computed directly from the given parameters and
+the shape of the cross section, which allows for anisotropic behaviors, which can be found in e.g. belts or computer
+cables. It is a single kinematic tree, so it is exactly inextensible without the use of additional constraints, enabling
+the use of large time steps. The elastic model is geometrically exact and based on computing the Bishop or twist-free
+frame of the centerline, i.e., the line passing through the center of the cross section. The orientations of the geoms
+are expressed with respect to this frame and then decomposed into twist and bending components, hence different
+stiffnesses can be set independently. Moreover, it is possible to specify if the stress-free configuration is flat or
+curve, such as in the case of coil springs. The cable requires using a first-party :ref:`engine plugin<exPlugin>`, which
+may be integrated directly into the engine in the future.
+
+**Rope and loop**.
+
+The rope and loop are deprecated. It is recommended to use the cable for simulating inextensible elastic rods that are
+bent and twisted and 1D flex :ref:`deformable objects <CDeformable>` for extensible strings in a tensile loading
+scenario (e.g. a stretched rubber band).
 
 **Cloth**.
 
-|image12| |image13|
-
-.. code-block:: xml
-
-   <body name="B3_5" pos="0 0 1">
-     <freejoint/>
-     <composite type="cloth" count="9 9 1" spacing="0.05" flatinertia="0.01">
-       <joint kind="main" damping="0.001"/>
-       <skin material="matcarpet" texcoord="true" inflate="0.005" subgrid="2"/>
-       <geom type="capsule" size="0.015 0.01" rgba=".8 .2 .1 1"/>
-     </composite>
-   </body>
-
-The cloth type is an alternative to a 2D grid, and has somewhat different properties. Similar to rope vs. 1D grid, the
-cloth is less jittery than a 2D grid and can also fill collision holes better. This is done by using capsules or
-ellipsoids, and arranging them in the pattern shown on the right. The geom capsules are shown in red, the kinematic
-tree in thick blue, the equality-constrained tendons holding the cloth together in thin gray, and the joints in cyan.
-The element body corresponding to the parent body has a floating joint rendered as a cube, while the rest of the tree
-is constructed using pairs of hinge joints that form universal joints. Note the naming of the parent body: similar to
-rope, it must coincide with one of the automatically-generated element body names in the composite object. Explicit
-pinning is not possible. However if the parent is a static body, the cloth is essentially pinned but only at one
-point. Similar to rope, the cloth object involves long kinematic chains that require relatively small timesteps and
-some damping for stable integration. The parameters can be found in the XML model files in the software distribution.
+The cloth is deprecated. It is recommended to use 2D flex :ref:`deformable objects <CDeformable>` for simulating thin
+elastic structures.
 
 **Box**.
 
@@ -1222,6 +1313,99 @@ of the system making it softer or harder, damped or springy, etc. Note that box,
 involve long kinematic chains, and can be simulated at large timesteps - similar to particle and grid, and unlike rope
 and cloth.
 
+.. _CDeformable:
+
+Deformable objects
+~~~~~~~~~~~~~~~~~~
+
+The :ref:`composite objects <CComposite>` described earlier were intended to emulate soft bodies in what is effectively
+a rigid-body simulator. This was possible because MuJoCo constraints are soft, but nevertheless it was limited in
+functionality and modeling power. In MuJoCo 3.0 we have introduced true deformable objects involving new model elements.
+The :ref:`skin<deformable-skin>` described earlier was actually one such element, but it is merely used for
+visualization. We now have a related element :ref:`flex<deformable-flex>` which generates contact forces, constraint
+forces and passive forces as needed to model a wide range of deformable entities. Both skins and flexes are now defined
+within a new grouping element in the XML called :ref:`deformable<deformable>`. A flex is a low-level element that
+specifies everything needed at runtime, but is difficult to design at modeling time. To aid with modeling, we have
+further introduced the element :ref:`flexcomp<body-flexcomp>` which automates the creation of the low-level flex,
+similar to how :ref:`composite<body-composite>` automates the creation of (collections of) MuJoCo objects needed to
+emulate a soft body. Flexes may eventually supersede composites, but for now both are useful for somewhat different
+purposes.
+
+A flex is a collection of MuJoCo bodies that are connected with massless stretchable elements. These elements can be
+capsules (1D flex), triangles (2D flex), or tetrahedra (3D flex). In all cases we allow a radius, which makes the
+elements smooth and also volumetric in 1D and 2D. The primitive elements are illustrated below:
+
+.. image:: images/modeling/flexelem.png
+   :width: 600px
+   :align: center
+
+Thus far these look like geoms. But the key difference is that they deform: as the bodies (vertices) move independently
+of each other, the shape of the elements changes in real time. Collisions and contact forces are now generalized to
+handle these deformable geometric elements. Note that when two such elements collide, the contact no longer involves
+just two bodies, but can involve up to 8 bodies (if both elements are tetrahedra). Contact forces are computed as
+before, given the contact frame and relevant quantities expressed in that frame. But then the contact force is
+distributed among all interacting bodies. The notion of contact Jacobian is complicated because the contact point cannot
+be considered fixed in any body frame. Instead we use a weighting scheme to "assign" each contact point to multiple
+bodies. It is also possible to create a rigid flex, by assigning all vertices to the same body. This is a way to
+re-purpose the new flex collision machinery to implement rigid non-convex mesh collisions (unlike mesh geoms which are
+convexified for collision purposes).
+
+**Deformation model**.
+
+In order to preserve the shape of the flex (in a soft sense), we need to generate passive or constraint forces. Prior to
+MuJoCo 3.0 this would involve a large number of tendons plus constraints on tendons and joints. This is still possible
+here, but inefficient both in terms of modeling and in terms of simulation when the flex is large. Instead, the design
+philosophy is to use a single set of parameters and provide two modeling choices: a new (soft) equality constraint type
+that applies to all edges of a given flex, which permits large time steps, or a discretized continuum representation,
+where each element is in a constant stress state, which is equivalent to piecewise linear finite elements and achieves
+improved realism and accuracy. The edge-based model could be seen as a "lumped" stiffness model, where the correct
+coupling of deformation modes (e.g. shear and volumetric) is averaged in a single quantity. The continuum model enables
+instead to specify shear and volumetic stiffnesses separately using the `Poisson's ratio
+<https://en.wikipedia.org/wiki/Poisson%27s_ratio>`__ of the material. For more details, see the `Saint Venant-Kirchhoff
+<https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__ hyperelastic model. This
+functionality is currently based on first-party :ref:`engine plugins<exPlugin>` as of MuJoCo 3.0 but may be integrated
+into the engine in future releases.
+
+**Creation and visualization**.
+
+.. code-block:: xml
+
+   <option timestep=".001"/>
+
+   <extension>
+      <plugin plugin="mujoco.elasticity.solid"/>
+   </extension>
+
+   <worldbody>
+      <flexcomp type="grid" count="24 4 4" spacing=".1 .1 .1" pos=".1 0 1.5"
+                radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="7">
+         <contact condim="3" solref="0.01 1" solimp=".95 .99 .0001" selfcollide="none"/>
+         <edge damping="1"/>
+         <plugin plugin="mujoco.elasticity.solid">
+            <config key="poisson" value="0.2"/>
+            <!--Units are in Pa (SI)-->
+            <config key="young" value="5e4"/>
+         </plugin>
+      </flexcomp>
+   </worldbody>
+
+Using the :ref:`flexcomp<body-flexcomp>` element, we can create flexes from meshes, including tetrahedral meshes, and
+automatically generate all the bodies/vertices and connect them with suitable elements. We can also create grids and
+other topologies automatically. This machinery makes it easy to create very large flexes, involving thousands or even
+tens of thousands of bodies, elements and edges. Obviously such simulations will not be fast. Even for medium-sized
+flexes, pruning of collision pairs and essential. This is why we have developed elaborate methods for pruning
+self-collisions; see XML reference.
+
+In case of 3D flexes made of tetrahedra, it may be useful to examine how the flex is "triangulated" internally. We have
+a special visualization mode that peels off the outer layers. Below is an example with the Stanford Bunny. Note how it
+has smaller tetrahedra on the outside and larger ones on the inside. This mesh design makes sense, because we want the
+collision surface to be accurate, but on the inside we just need soft material properties - which require less spatial
+resolution. In order to convert a surface mesh to a tetrahedral mesh, we recommend open tools like the
+`fTetWild library <https://github.com/wildmeshing/fTetWild>`__.
+
+|bunny1| |bunny2|
+
+
 .. _CInclude:
 
 Including files
@@ -1347,7 +1531,7 @@ mocap bodies around:
 The key thing to understand about mocap bodies is that the simulator treats them as being fixed. We are causing them
 to move from one simulation time step to the next by updating their position and orientation directly, but as far as
 the physics model is concerned their position and orientation are constant. So what happens if we make contact with a
-regular dynamic body, as in the composite object examples provided with the MuJoCo 2.0 distribution (recall that in
+regular dynamic body, as in the composite object examples provided with the MuJoCo distribution (recall that in
 those example we have a capsule probe which is a mocap body that we move with the mouse). A contact between two
 regular bodies will experience penetration as well as relative velocity, while contact with a mocap body is missing
 the relative velocity component because the simulator does not know that the mocap body itself is moving. So the
@@ -1499,7 +1683,7 @@ the same yet the damping component of the force is integrated implicitly.
 Restitution
 ~~~~~~~~~~~
 
-MuJoCo 2.0 introduced a new mechanism for specifying :at:`solref` as explained in :ref:`Solver parameters <CSolver>`.
+Another mechanism exists for specifying :at:`solref`, as explained in :ref:`Solver parameters <CSolver>`.
 When both numbers are non-positive, they are interpreted as (-stiffness, -damping) and scaled by the constraint
 impedance. To achieve perfect restitution for contacts and other constraints, set stiffness to some reasonably large
 value and damping to zero. Below is an example of a sphere bouncing on a plane with restitution coefficient of 1, so
@@ -1520,12 +1704,6 @@ in a visible way, and the energy fluctuates around the initial value instead of
    </worldbody>
 
 
-.. |image0| image:: images/modeling/impedance.png
-   :width: 600px
-.. |image1| image:: images/modeling/musclemodel.png
-   :width: 650px
-.. |image2| image:: images/modeling/musclerange.png
-   :width: 400px
 .. |image3| image:: images/modeling/tendonwraps.png
    :width: 500px
 .. |image4| image:: images/modeling/particle.png
@@ -1562,3 +1740,9 @@ in a visible way, and the energy fluctuates around the initial value instead of
    :height: 250px
 .. |particle| image:: images/models/particle.gif
    :width: 270px
+.. |bunny1| image:: images/modeling/bunny1.png
+   :width: 300px
+.. |bunny2| image:: images/modeling/bunny2.png
+   :width: 300px
+.. |coil| image:: images/modeling/coil.png
+   :width: 300px

doc/overview.rst

@@ -98,7 +98,7 @@ Separation of model and data
    :ref:`mjModel` is constructed by the compiler. :ref:`mjData` is constructed at runtime, given
    :ref:`mjModel`. This separation makes it easy to simulate multiple models as well as multiple states and controls for
    each model, in turn facilitating :ref:`multi-threading <siMultithread>` for sampling and :ref:`finite
-   differences <saDerivative>`. The top-level API functions reflect this basic separation, and have
+   differences <mjd_transitionFD>`. The top-level API functions reflect this basic separation, and have
    the format:
 
    .. code:: C
@@ -585,8 +585,21 @@ Equality constraints can impose additional constraints beyond those already impo
 and the joints/DOFs defined in it. They can be used to create loop joints, or in general model mechanical coupling.
 The internal forces that enforce these constraints are computed together with all other constraint forces. The
 available equality constraint types are: connect two bodies at a point (creating a ball joint outside the kinematic
-tree); weld two bodies together; make two surfaces slide on each other; fix the position of a joint or tendon; couple
-the positions of two joints or two tendons via a cubic polynomial.
+tree); weld two bodies together; fix the position of a joint or tendon; couple the positions of two joints or two
+tendons via a cubic polynomial; constrain the edges of a flex (i.e. deformable mesh) to their initial lengths.
+
+
+Flex
+^^^^
+
+Flexes were added in MuJoCo 3.0. They represent deformable meshes that can be 1, 2 or 3 dimensional (thus their elements
+are capsules, triangles or tetrahedra). Unlike geoms which are static shapes attached rigidly to a single body, the
+elements of a flex are deformable: they are constructed by connecting multiple bodies, thus the body positions and
+orientations determine the shape of the flex elements at runtime. These deformable elements suport collisions and
+contact forces, as well as generate passive and constraint forces which softly preserve the shape of the deformable
+entity. Automation is provided to load a mesh from a file, construct bodies corresponding to the mesh vertices,
+construct flex elements corresponding to the mesh faces (or lines or tetrahedra, depending on dimensionality), and
+obtain a corresponding deformable mesh.
 
 Contact pair
 ^^^^^^^^^^^^
@@ -596,7 +609,9 @@ sources: automated proximity tests and other filters collectively called "dynami
 geom pairs provided in the model. The latter is a separate type of model element. Because a contact involves a
 combination of two geoms, the explicit specification allows the user to define contact parameters in ways that cannot
 be done with the dynamic mechanism. It is also useful for fine-tuning the contact model, in particular adding contact
-pairs that were removed by an aggressive filtering scheme.
+pairs that were removed by an aggressive filtering scheme. The contact machinery is now extended to flex elements,
+which can create contact interactions between more than two bodies. However such collisions are automated and cannot
+be finetuned using contact pairs.
 
 Contact exclude
 ^^^^^^^^^^^^^^^
@@ -929,15 +944,20 @@ can be obtained as:
    int qposadr = -1, qveladr = -1;
 
    // make sure we have a floating body: it has a single free joint
-   if( bodyid>=0 && m->body_jntnum[bodyid]==1 &&
-       m->jnt_type[m->body_jntadr[bodyid]]==mjJNT_FREE )
-      {
-         // extract the addresses from the joint specification
-         qposadr = m->jnt_qposadr[m->body_jntadr[bodyid]];
-         qveladr = m->jnt_dofadr[m->body_jntadr[bodyid]];
-      }
+   if( bodyid>=0 && m->body_jntnum[bodyid]==1 && m->jnt_type[m->body_jntadr[bodyid]]==mjJNT_FREE ) {
+     // extract the addresses from the joint specification
+     qposadr = m->jnt_qposadr[m->body_jntadr[bodyid]];
+     qveladr = m->jnt_dofadr[m->body_jntadr[bodyid]];
+   }
 
 Now if everything went well (i.e., "myfloatingbody" was indeed a floating body), qposadr and qveladr are the addresses
 in qpos and qvel where the data for our floating body/joint lives. The position data is 7 numbers (3D position followed
 by unit quaternion) while the velocity data is 6 numbers (3D linear velocity followed by 3D angular velocity). These
 numbers can now be set to the desired pose and velocity of the body.
+
+The semantics of free joints are as follows. The linear postions of free joints are in the global frame, as are
+linear velocities. The orientation of a free joint (the quaternion) is also in the global frame. However, the rotational
+velocities of a free joint are in the local body frame. This is not so much a design decision but rather correct
+use of the topology of quaternions. Angular velocities live in the quaternion tangent space, which is defined locally
+for a certain orientation, so frame-local angular velocities are a natural parameterization.
+Accelerations are defined in the same space as the corresponding velocities.

doc/programming/extension.rst

@@ -199,6 +199,25 @@ When :ref:`mjData` is being copied via :ref:`mj_copyData`, MuJoCo will copy over
 code is responsible for setting up the plugin data for the newly copied :ref:`mjData`. To facilitate this, MuJoCo calls
 the ``copy`` callback from :ref:`mjpPlugin` for each plugin instance present.
 
+.. _exActuatorAct:
+
+Actuator activations
+""""""""""""""""""""
+
+When writing stateful actuator plugins, there are two choices for where to save the actuator state. One option is using
+``plugin_state`` as described above, and the other is to use ``mjData.act`` by implementing the ``actuator_actdim`` and
+``actuator_act_dot`` callbacks on :ref:`mjpPlugin`.
+
+When using the latter option, the actuator plugin's state will be added to ``mjData.act``, and MuJoCo will
+automatically integrate ``mjData.act_dot`` values between timesteps. One advantage of this approach is that
+finite-differencing functions like :ref:`mjd_transitionFD` will work as they do for native actuators. The
+``mjpPlugin.advance`` callback will be called after ``act_dot`` is integrated, and actuator plugins may overwrite
+the ``act`` values at that point, if Euler integration isn't appropriate.
+
+Users may specify the :ref:`dyntype<actuator-plugin-dyntype>` attribute on actuator plugins, to introduce a filter or
+an integrator between user inputs and actuator activations. When they do, the activation variable introduced by
+``dyntype`` will be placed *after* the plugin's activation variables in the ``act`` array.
+
 .. _exRegistration:
 
 Registration
@@ -248,27 +267,35 @@ A future version of this section will include:
 * Things that developers need to keep in mind in order to ensure that plugins function correctly when :ref:`mjData` is
   copied, stepped, or reset.
 
-Currently, there are three directories of first-party plugins:
+There are several first-party plugin directories:
 
-* **elasticity:** The plugins in the `elasticity/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__
-  directory are passive forces based on continuum mechanics for 1-dimensional and 3-dimensional bodies. The
-  1D model is invariant under rotations and captures the large deformation of elastic cables, decoupling twisting and
-  bending strains. The 3D solid is a
-  `Saint Venant-Kirchhoff <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__
-  model discretized with piecewise linear finite elements, which is suitable for large deformations with small strains.
+* **actuator:** The plugins in the `actuator/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/actuator>`__
+  directory implement custom actuators, so far only a PID controller. See the
+  `README <https://github.com/google-deepmind/mujoco/blob/main/plugin/actuator/README.md>`__ for details.
+* **elasticity:** The plugins in the `elasticity/
+  <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__ directory are passive forces based on
+  continuum mechanics for 1-dimensional and 3-dimensional bodies. The 1D model is invariant under rotations and captures
+  the large deformation of elastic cables, decoupling twisting and bending strains. The 3D solid is a `Saint
+  Venant-Kirchhoff <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__ model
+  discretized with piecewise linear finite elements, which is suitable for large deformations with small strains. See
+  also :ref:`composite <CComposite>` and :ref:`deformable <CDeformable>` objects. For more information, please see the
+  `README <https://github.com/google-deepmind/mujoco/blob/main/plugin/elasticity/README.md>`__.
 * **sensor:** The plugins in the `sensor/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sensor>`__
   directory implement custom sensors. Currently the sole sensor plugin is the touch grid sensor, see the
   `README <https://github.com/google-deepmind/mujoco/blob/main/plugin/sensor/README.md>`__ for details.
-* **sdf:** The plugins in the `sdf/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sdf>`__
-  directory specify custom shapes in a mesh-free manner, by defining methods computing a signed distance field and its
-  gradient at query points. This shape then acts as a new geom type in the collision table at the top of
-  `engine_collision_driver.c <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__.
+* **sdf:** The plugins in the `sdf/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sdf>`__ directory
+  specify custom shapes in a mesh-free manner, by defining methods computing a signed distance field and its gradient at
+  query points. This shape then acts as a new geom type in the collision table at the top of `engine_collision_driver.c
+  <https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_collision_driver.c>`__. For more information
+  concerning the available SDFs and how to write your own implicit geometry, please see the `README
+  <https://github.com/google-deepmind/mujoco/blob/main/plugin/sdf/README.md>`__. The rest of this section will give more
+  detail concerning the collision algorithm and the plugin engine interface.
 
-  Collision points are found by minimizing the maximum of the two colliding SDFs via gradient descent.
-  Because SDFs are non-convex, multiple starting points are required in order to converge to multiple local minima.
-  The number of starting points is set using :ref:`sdf_initpoints<option-sdf_initpoints>`, and are
-  initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes.
-  The number of gradient descent iterations is set using :ref:`sdf_iterations<option-sdf_iterations>`.
+  Collision points are found by minimizing the function A + B + abs(max(A, B)), where A and B are the two colliding
+  SDFs, via gradient descent. Because SDFs are non-convex, multiple starting points are required in order to converge to
+  multiple local minima. The number of starting points is set using :ref:`sdf_initpoints<option-sdf_initpoints>`, and
+  are initialized using the Halton sequence inside the intersection of the axis-aligned bounding boxes. The number of
+  gradient descent iterations is set using :ref:`sdf_iterations<option-sdf_iterations>`.
 
   While *exact* SDFs---encoding the precise signed distance to the surface---are preferred, collisions are possible with
   any function whose value vanishes at the surface and grows monotonically away from it, with a negative sign in the
@@ -318,7 +345,6 @@ loading functions. The :ref:`mjpResourceProvider` struct stores three types of f
 .. _Uniform Resource Identifier: https://en.wikipedia.org/wiki/Uniform_Resource_Identifier
 
 Resource prefix
-
   Resources are identified by prefixes in their name. The chosen prefix should have a valid `Uniform Resource
   Identifier`_ (URI) scheme syntax. Resource names should also have a valid URI syntax, however this isn't enforced. A
   resource name with the syntax ``{prefix}:{filename}`` will match a provider using the scheme ``prefix``.  For

doc/programming/index.rst

@@ -35,7 +35,7 @@ Thread
    to process Tasks asynchronously. To enable use in MuJoCo, create a ThreadPool and assign it to the thread_pool field
    in mjData.
 UI framework
-   The UI framework (new in MuJoCo 2.0) is written in C. UI elements are rendered in OpenGL. It has its own event
+   The UI framework is written in C. UI elements are rendered in OpenGL. It has its own event
    mechanism and abstract hooks for keyboard and mouse input. The code samples use it with GLFW, but it can also be used
    with other window libraries.
 
@@ -165,7 +165,7 @@ links below, to make this documentation self-contained.
 `mjxmacro.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjxmacro.h>`__
    This file is optional and is not included by mujoco.h. It defines :ref:`X Macros <tyXMacro>` that can
    automate the mapping of mjModel and mjData into scripting languages, as well as other operations that require
-   accessing all fields of mjModel and mjData. See code sample :ref:`testxml.cc <saTestXML>`.
+   accessing all fields of mjModel and mjData.
 `mjexport.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjexport.h>`__
    Macros used for exporting public symbols from the MuJoCo library. This header should not be used directly by client
    code.
@@ -275,5 +275,6 @@ now lazily resolved at runtime after the switch to GLAD, the "nogl" libraries ar
 
     simulation
     visualization
+    ui
     samples
     extension

doc/programming/samples.rst

@@ -12,68 +12,66 @@ with the library.
 `testspeed <https://github.com/google-deepmind/mujoco/blob/main/sample/testspeed.cc>`_
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-This code sample tests the simulation speed for a given model. The command line arguments are the model file, the
-number of time steps to simulate, the number of parallel threads to use, and a flag to enable internal profiling (the
-last two are optional). When N threads are specified with N>1, the code allocates a single mjModel and per-thread
-mjData, and runs N identical simulations in parallel. The idea is to test performance with all cores active, similar
-to Reinforcement Learning scenarios where samples are collected in parallel. The optimal N usually equals the number
-of logical cores. By default the simulation starts from the model reference configuration qpos0 and qvel=0. However if
-a keyframe named "test" is present in the model, it is used as the initial state state.
+This code sample times the simulation of a given model. The timing is straightforward: the simulation of the passive
+dynamics (with optional control noise) is rolled-out for the specified number of steps, while collecting statistics
+about the number of contacts, scalar constraints, and CPU times from internal profiling. The results are then printed to
+the console. To simulate controlled dynamics instead of passive dynamics one can either install a control callback
+:ref:`mjcb_control`, or modify the code to set control signals explicitly, as explained in the :ref:`simulation loop
+<siSimulation>` section below. This command-line utility is run with
 
-The timing code is straightforward: the simulation of the passive dynamics is advanced for the specified number of
-steps, while collecting statistics about the number of contacts, scalar constraints, and CPU times from internal
-profiling. The results are then printed in the console. To simulate controlled dynamics instead of passive dynamics
-one can either install the control callback :ref:`mjcb_control`, or set control signals
-explicitly as explained in the :ref:`simulation loop <siSimulation>` section below.
+.. code-block:: Shell
 
-.. _saTestXML:
+   testspeed modelfile [nstep nthread ctrlnoise npoolthread]
 
-`testxml <https://github.com/google-deepmind/mujoco/blob/main/sample/testxml.cc>`_
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Where the command line arguments are
 
-This code sample tests the parser, compiler and XML writer. The testing code does the following:
+.. list-table::
+   :width: 95%
+   :align: left
+   :widths: 1 1 5
+   :header-rows: 1
 
--  Parse and compile a specified XML model in MJCF or URDF. This yields an mjModel structure ready for simulation;
--  Save the model as a temporary MJCF file, using a "canonical" subset of MJCF where a number of conversions have
-   already been performed by the compiler;
--  Parse and compile the temporary MJCF file. This yields a second mjModel structure ready for simulation;
--  Compare the two mjModel structures field by field, and print the field with the largest numerical difference. Since
-   MJCF is a text format, the real-valued numbers saved in it have lower precision than the double precision used
-   internally, thus we cannot expect the two models to be identical on the bit level. But we can expect the largest
-   difference to be on the order of 1e-6. A substantially larger difference indicates a bug in the parser, compiler or
-   XML writer - and should be reported.
+   * - Argument
+     - Default
+     - Meaning
+   * - ``modelfile``
+     - (required)
+     - path to model
+   * - ``nstep``
+     - 10000
+     - number of steps per rollout
+   * - ``nthread``
+     - 1
+     - number of threads running parallel rollouts
+   * - ``ctrlnoise``
+     - 0.01
+     - scale of pseudo-random noise injected into actuators
+   * - ``npoolthread``
+     - 1
+     - number of threads in engine-internal threadpool
 
-The code uses the :ref:`X Macros <tyXMacro>` described in the Reference chapter. This is a convenient way
-to apply the same operation to all fields in mjModel, without explicitly typing their names. The code sample
-:ref:`simulate.cc <saSimulate>` also uses X Macros to implement a watch, where the user can type the name of any mjData
-field which is resolved at runtime.
+**Notes:**
 
-.. _saCompile:
-
-`compile <https://github.com/google-deepmind/mujoco/blob/main/sample/compile.cc>`_
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This code sample evokes the built-in parser and compiler. It implements all possible model conversions from (MJCF, URDF,
-MJB) format to (MJCF, MJB, TXT) format. Models saved as MJCF use a canonical subset of our format as described in the
-:doc:`../modeling` chapter, and therefore MJCF-to-MJCF conversion will generally result in a different file.
-The TXT format is a human-readable road-map to the model. It cannot be loaded by MuJoCo, but can be a very useful aid
-during model development. It is in one-to-one correspondence with the compiled mjModel. Note also that one can use the
-function :ref:`mj_printData` to create a text file which is in one-to-one correspondence
-with mjData, although this is not done by the code sample.
-
-.. _saBasic:
-
-`basic <https://github.com/google-deepmind/mujoco/blob/main/sample/basic.cc>`_
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This code sample is a minimal interactive simulator. The model file must be provided as command-line argument. It
-opens an OpenGL window using the platform-independent GLFW library, and renders the simulation state at 60 fps while
-advancing the simulation in real-time. Press Backspace to reset the simulation. The mouse can be used to control the
-camera: left drag to rotate, right drag to translate in the vertical plane, shift right drag to translate in the
-horizontal plane, scroll or middle drag to zoom.
-
-The :ref:`Visualization` programming guide below explains how visualization works. This code sample is a minimal
-illustration of the concepts in that guide.
+- When ``nthread > 1`` is specified, the code allocates a single mjModel and per-thread mjData, and runs ``nthread``
+  identical simulations in parallel. This tests performance with all cores active, as in Reinforcement
+  Learning scenarios where samples are collected in parallel. The optimal ``nthread`` usually equals the number of
+  logical cores.
+- By default, the simulation starts from the model reference configuration with zero velocities. However, if a
+  keyframe named "test" is present in the model, it is used as the initial state.
+- The ``ctrlnoise`` argument prevents models from settling into a static state where, due to warmstarts, one can
+  measure artificially faster simulation.
+- When ``npoolthread > 1`` is specified, an engine-internal :ref:`mjThreadPool` is created with the specified number of
+  threads, to speed up simulation of large scenes. Note that while it is possible to to use both ``nthread`` and
+  ``npoolthread``, the scenarios for which one would want these different type of multithreading are usually mutually
+  exclusive.
+- For more repeatable performance statistics, run the tool with the ``performance``
+  `governor <https://www.kernel.org/doc/Documentation/cpu-freq/governors.txt>`__ on Linux, or the
+  ``High Performance`` power plan on Windows, to reduce noise from CPU scaling.
+- Many modern CPUs contain a mixture of "performance" and "efficiency" cores. Users should consider restricting the
+  process to only run on the same type of cores for more interpretable performance statistics. This can be done via the
+  `taskset <https://man7.org/linux/man-pages/man1/taskset.1.html>`__ command on Linux, or the
+  `start /affinity <https://learn.microsoft.com/en-us/windows-server/administration/windows-commands/start>`__
+  command on Windows (processor affinity cannot be specified through documented API means on macOS).
 
 .. _saSimulate:
 
@@ -83,10 +81,11 @@ illustration of the concepts in that guide.
 This code sample is a fully-featured interactive simulator. It opens an OpenGL window using the platform-independent
 GLFW library, and renders the simulation state in it. There is built-in help, simulation statistics, profiler, sensor
 data plots. The model file can be specified as a command-line argument, or loaded at runtime using drag-and-drop
-functionality. As of MuJoCo 2.0, this code sample uses the native UI to render various controls, and provides an
+functionality. This code sample uses the native UI to render various controls, and provides an
 illustration of how the new UI framework is intended to be used. Below is a screen-capture of ``simulate`` in action:
 
 ..  youtube:: 0ORsj_E17B0
+    :width: 95%
     :align: center
 
 Interaction is done with the mouse; built-in help with a summary of available commands is available by pressing the
@@ -115,6 +114,33 @@ power settings, the OS may reduce the CPU frequency; this is because :ref:`simul
 the time in order to slow down to realtime. This results in inaccurate timings. To avoid this problem, change the
 Windows power plan so that the minimum processor state is 100%.
 
+.. _saCompile:
+
+`compile <https://github.com/google-deepmind/mujoco/blob/main/sample/compile.cc>`_
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This code sample evokes the built-in parser and compiler. It implements all possible model conversions from (MJCF, URDF,
+MJB) format to (MJCF, MJB, TXT) format. Models saved as MJCF use a canonical subset of our format as described in the
+:doc:`../modeling` chapter, and therefore MJCF-to-MJCF conversion will generally result in a different file.
+The TXT format is a human-readable road-map to the model. It cannot be loaded by MuJoCo, but can be a very useful aid
+during model development. It is in one-to-one correspondence with the compiled mjModel. Note also that one can use the
+function :ref:`mj_printData` to create a text file which is in one-to-one correspondence
+with mjData, although this is not done by the code sample.
+
+.. _saBasic:
+
+`basic <https://github.com/google-deepmind/mujoco/blob/main/sample/basic.cc>`_
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This code sample is a minimal interactive simulator. The model file must be provided as command-line argument. It
+opens an OpenGL window using the platform-independent GLFW library, and renders the simulation state at 60 fps while
+advancing the simulation in real-time. Press Backspace to reset the simulation. The mouse can be used to control the
+camera: left drag to rotate, right drag to translate in the vertical plane, shift right drag to translate in the
+horizontal plane, scroll or middle drag to zoom.
+
+The :ref:`Visualization` programming guide below explains how visualization works. This code sample is a minimal
+illustration of the concepts in that guide.
+
 .. _saRecord:
 
 `record <https://github.com/google-deepmind/mujoco/blob/main/sample/record.cc>`_
@@ -152,69 +178,3 @@ proliferation of overlapping technologies, which differ not only between platfor
 case of Linux. The addition of a couple of extra functions (such as those provided by OSMesa for example) could have
 avoided a lot of confusion. EGL is a newer standard from Khronos which aims to do this, and it is gaining popularity.
 But we cannot yet assume that all users have it installed.
-
-.. _saDerivative:
-
-`derivative <https://github.com/google-deepmind/mujoco/blob/main/sample/derivative.cc>`_
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This code sample illustrates the numerical approximation of forward and inverse dynamics derivatives via finite
-differences. The process involves a number of epochs. In each epoch the simulation is advanced for a specified number
-of steps, derivatives are computed at the last state, and timing and accuracy statistics are collected. The averages
-over epochs are printed at the end.
-
-The code can be incorporated in user projects where derivatives are needed, and can also be used as a stand-alone tool
-for estimating CPU time and numerical accuracy. Accuracy is estimated in the function ``checkderiv()`` using several
-mathematical identities about the derivatives of inverse functions; the residuals being computed would be zero if the
-derivatives were exact. Note that these identities involve matrix multiplications which may affect the accuracy
-estimates. Timing tests are applied only to the parallel section, where the function ``worker()`` is executed in
-multiple threads using OpenMP. There are fewer threads than forward/inverse dynamics evaluations, thus each thread
-executes multiple evaluations. For a more general discussion of parallel processing in MuJoCo see
-:ref:`multi-threading <siMultithread>` below.
-
-Recall than for a differentiable function ``f(x)`` the derivative can be approximated as
-
-.. code-block:: Text
-
-     df/dx = (f(x+eps)-f(x))/eps
-
-where ``eps`` is a small number. One can also use the centered finite difference method, which is two times slower but
-more accurate. Here ``f`` is one of the functions
-
-.. code-block:: Text
-
-     forward dynamics:  qacc(qfrc_applied, qvel, qpos)
-     inverse dynamics:  qfrc_inverse(qacc, qvel, qpos)
-
-The code sample computes six Jacobian matrices, containing the derivative of each function with respect to its three
-arguments. The results are stored in the array ``deriv``. All six Jacobian matrices are square, with dimensionality
-equal to the number of degrees of freedom ``mjModel.nv``. When the model configuration includes quaternion joints,
-mjData.qpos has larger dimensionality than the other vectors, however the derivative is only defined in the tangent
-space to the configuration manifold. This is why, when differentiating with respect to the elements of ``mjData.qpos``,
-we do not directly add ``eps`` but instead use the function :ref:`mju_quatIntegrate` to perturb the quaternion in the
-tangent space, keeping it normalized. This technique should also be used in any other situation where quaternions need
-to be perturbed.
-
-There are some important subtleties in this code that improve speed as well as accuracy. To speed up the computation,
-we re-use intermediate results whenever possible. This relies on the skip mechanism described under :ref:`forward
-dynamics <siForward>` and :ref:`inverse dynamics <siInverse>` below. We first perturb force dimensions, keeping
-position and velocity fixed. In this way we avoid recomputing results that depend on position and velocity but not on
-force. Then we perturb velocity dimensions, and avoid recomputing results that depend on position but not on velocity
-or force. Finally we perturb position dimensions - which requires full computation because everything depends on
-position.
-
-Accuracy depends on the value of ``eps`` which is user-adjustable, as well as the shape of the function. In the case
-of forward dynamics however, the function evaluation involves an iterative constraint solver, and this must be handled
-with care. In general, the difference between ``f(x+eps)`` and ``f(x)`` is very small, thus any noise affecting the
-two function evaluations differently can make the resulting derivatives meaningless. Different warm-starts or
-different number of solver iterations can act as such noise here. Therefore we fix the warm-start ``mjData.qacc`` to a
-value pre-computed at the center point, using ``nwarmup`` extra major iterations to obtain a more accurate warm-start.
-We also fix the number of solver iterations to ``niter`` and set ``mjModel.opt.tolerance = 0``; this disables the early
-termination mechanism. Note that the original simulation options are restored in the serial code which advances the
-state.
-
-We emphasize that the above subtleties are not high-order corrections that can be incorporated later. In the presence
-of unilateral constraints, numerical derivatives are hard to compute and there is no shortcut around it; indeed they
-would not even be defined if it wasn't for our soft-constraint model. Making the constraints softer results in more
-accurate results. This effect is so strong that in some situations it makes sense to intentionally work with the wrong
-model, i.e., a model that is softer than desired, so as to obtain more accurate derivatives.

doc/programming/simulation.rst

@@ -428,9 +428,7 @@ However, MuJoCo is designed not only for simulation but also for more advanced a
 optimization, machine learning etc. In such settings one often needs to sample the dynamics at a cloud of nearby
 states, or approximate derivatives via finite differences - which is another form of sampling. If the samples are
 arranged on a grid, where only the position or only the velocity or only the control is different from the center
-point, then the above mechanism can improve performance by about a factor of 2. The code sample :ref:`derivative.cc
-<saDerivative>` illustrates this approach, and also shows how :ref:`multi-threading <siMultithread>` can be used for
-additional speedup.
+point, then the above mechanism can improve performance by about a factor of 2.
 
 .. _siInverse:
 
@@ -494,10 +492,7 @@ bodies, we may eventually implement within-step multi-threading, but for now thi
 Rather than speed up a single simulation, we prefer to use multi-threading to speed up sampling operations that are
 common in more advanced applications. Simulation is inherently serial over time (the output of one mj_step is the
 input to the next), while in sampling many calls to either forward or inverse dynamics can be executed in parallel
-since there are no dependencies among them, except perhaps for a common initial state. The code sample
-:ref:`derivative.cc <saDerivative>` illustrates one important example of sampling, namely the approximation of
-dynamics derivatives via finite differences. Here we will not repeat the material from that section, but will instead
-explain MuJoCo's general approach to parallel processing.
+since there are no dependencies among them, except perhaps for a common initial state.
 
 MuJoCo was designed for multi-threading from its beginning. Unlike most existing simulators where the notion of
 dynamical system state is difficult to map to the software state and is often distributed among multiple objects, in
@@ -515,7 +510,7 @@ management.
 
    // allocate per-thread mjData
    mjData* d[64];
-   for( int n=0; n<nthread; n++ )
+   for( int n=0; n < nthread; n++ )
        d[n] = mj_makeData(m);
 
    // ... serial code, perhaps using its own mjData* dmain
@@ -541,8 +536,8 @@ writes to its own mjData. Therefore no further synchronization among threads is
 The above template reflects a particular style of parallel processing. Instead of creating a large number of threads,
 one for each work item, and letting OpenMP distribute them among processors, we rely on manual scheduling. More
 precisely, we create as many threads as there are processors, and then within the ``worker`` function we distribute the
-work explicitly among threads (not shown here, but see :ref:`derivative.cc <saDerivative>` for an example). This
-approach is more efficient because the thread-specific mjData is large compared to the processor cache.
+work explicitly among threads. This approach is more efficient because the thread-specific mjData is large compared to
+the processor cache.
 
 We also use a shared mjModel for cache-efficiency. In some situations it may not be possible to use the same mjModel
 for all threads. One obvious reason is that mjModel may need to be modified within the thread function. Another reason
@@ -676,10 +671,7 @@ correct way to do it is the hard way:
    mju_copy(myqposqvel, d->qpos, m->nq + m->nv);
 
 The :ref:`X Macros <tyXMacro>` defined in the optional header file ``mjxmacro.h`` can be used to automate allocation of
-data structure that match mjModel and mjData, for example when writing a MuJoCo wrapper for a scripting language. In
-the code sample :ref:`testxml.cc <saTestXML>` we use these unusual macros to compare all data arrays from two instances
-of mjModel and find the one with the largest difference. Apparently X Macros were invented in the 1960's for assembly
-language, and remain a great idea.
+data structure that match mjModel and mjData, for example when writing a MuJoCo wrapper for a scripting language.
 
 .. _siStack:
 
@@ -841,13 +833,6 @@ forces are disabled) the underlying physical system is energy-conserving. In tha
 the total energy indicate inaccuracies in numerical integration. For such systems the Runge-Kutta integrator has much
 better performance than the default semi-implicit Euler integrator.
 
-Finally, the user can implement additional diagnostics as needed. Two examples were provided in the code samples
-``testxml.cc`` and ``derivative.cc``, where we computed model mismatches after save and load, and assessed the accuracy
-of the numerical derivatives respectively. Key to such diagnostics is to implement two different algorithms or
-simulation paths that compute the same quantity, and compare the results numerically. This type of sanity check is
-essential when dealing with complex dynamical systems where we do not really know what the numerical output should be;
-if we knew that, we would not be using a simulator in the first place.
-
 .. _siJacobian:
 
 Jacobians

doc/programming/ui.rst

@@ -0,0 +1,105 @@
+.. _UI:
+
+User Interface
+--------------
+
+MuJoCo has a native UI framework. Its use is illustrated in the :ref:`simulate.cc <saSimulate>` viewer. It is
+designed to be fast in terms of updating and rendering, easy to use both for the developer and for the user,
+cross-platform, and integrated with the native MuJoCo renderer. In order to achieve these design goals, we have omitted
+many features and customization options that are available in other UI frameworks, and instead focused on efficiency and
+automation.
+
+
+
+.. _uiDesign:
+
+Design overview
+~~~~~~~~~~~~~~~
+
+
+Native OpenGL rendering
+  We do not use any helper tools or libraries. Instead we provide C code for rendering all UI elements directly in
+  OpenGL. We support multiple UIs, each of which is a virtual rectangle that can be taller than the visible window. The
+  elements of each UI are rendered offscreen in auxiliary OpenGL buffers, via minimal updates, only when changes are
+  necessary. At each screen refresh we then copy the pixels from these auxiliary buffers to the window framebuffer, and
+  also implement vertical scroll bars when the window is smaller than the UI. This copy operation is done on the GPU and
+  is very fast.
+
+
+Platform abstraction
+  The software design has 3 layers: OpenGL rendering of UI elements working in conjunction with the MuJoCo renderer
+  (which is fully cross-platform); abstract functions for access to windows, keyboard and mouse, defined as pure virtual
+  functions in the ``PlatformUIAdapter`` class; and an implementation of those functions in the derived class
+  ``GlfwAdapter``. `GLFW <https://www.glfw.org/>`__ itself is cross-platform. Nevertheless, we have opted for this
+  layered design in order to separate generic from platform-specific functionality. If GLFW needs to be replaced with
+  another similar framework for some reason, only ``GlfwAdapter`` will need to be rewritten.
+
+
+Themes and appearance
+  Individual UI elements do not allow customization in terms of appearance or layout. Instead, we use themes for colors
+  and spacing, and arrange all UI elements automatically. Several built-in themes are provided and the user can design
+  custom themes, however the entire UI uses a single theme for all elements. Appearance is minimalist: mostly colored
+  rectangles with text. Bitmaps and other custom decorations are not supported. The UI element types are check boxes,
+  radio button groups, selection lists, sliders, text edit boxes, static text, buttons, separators. These elements are
+  grouped into sections which can be expanded and collapsed.
+
+
+Layout and rectangles
+  Each UI is one virtual rectangle, whose width is determined by the theme and whose height is determined by the
+  sections, items within each section, and also the expand/collapse state of each section. The sizes and auxiliary
+  buffers for these virtual rectangles are handled automatically when the UI is updated. Each UI has a visible rectangle
+  on the screen, and in addition there are other rectangles -- for 3D rendering, 2D figures, and possibly custom OpenGL
+  rendering. All these visible rectangles are saved (in :ref:`mjuiState`) and are used to determine where mouse events
+  should be directed. The rectangle layout is updated by a callback provided by the user.
+
+
+Static allocation and creation
+  Rather than allocating and deallocating a large number of objects corresponding to UI elements and linking them
+  together, we create a single C struct (type :ref:`mjUI`) with static allocation supporting some maximum number of
+  sections and elements; and then keep a record of how many are in use. UI creation is simplified by helper functions
+  whose input is a C struct (type :ref:`mjuiDef`) that is essentially a table where each row describes one UI element
+  (see below). This makes it possible to construct elaborate user interfaces with surprisingly little C code.
+  Programmatic UI creation is also possible, for example when populating a UI with sliders corresponding to MuJoCo model
+  joints.
+
+
+Minimal state
+  The UI is designed to be as stateless as possible, so as to simplify development. This has two aspects. First, instead
+  of replicating user data within the UI elements, we store pointers to user data. For example, we might create a UI
+  slider and set its data pointer to ``mjData* d->qpos+7``. This slider will visualize as well as control the 7th scalar
+  component of the qpos vector of a MuJoCo model. Thus, when the simulation is updated, we have to remember to update
+  the UI as well. And furthermore we have to disable UI editing when the simulation is being updated. But the advantage
+  is that the UI becomes easier to construct, and there is no danger of discrepancies between user data and the UI.
+  Second, the UI elements themselves are mostly stateless. Instead we keep track of a minimal set of global states, in
+  particular mouse and keyboard state, section expand/collapse, contents of the text box being edited if any.
+
+
+Automated enable and disable
+  While each UI item can be set in enabled or disabled state directly, we also provide automation as follows. Each UI
+  item can be assigned an integer category. Then a :ref:`mjfItemEnable` callback determines whether each category should
+  be enabled or disabled, based on some program-specific conditions. For example, sliders that can change the values of
+  MuJoCo model joints should be disabled when the simulation state is being updated.
+
+
+
+.. _uiAPI:
+
+Main API
+~~~~~~~~
+
+Click on the links below for detailed API reference of the main UI data structures and functions.
+
+
+**Main data structures:**
+
+- :ref:`mjUI<mjUI>`: An entire UI.
+- :ref:`mjuiState<mjuiState>`: Global UI state.
+- :ref:`mjuiDef<mjuiDef>`: One entry in the definition table used for UI construction.
+
+
+**Main functions:**
+
+- :ref:`mjui_update<mjui_update>`: Main UI update function.
+- :ref:`mjui_render<mjui_render>`: Renders the UI.
+- :ref:`mjui_event<mjui_event>`: Low-level event handler.
+- :ref:`mjui_add<mjui_add>`: Helper function used to construct a UI.

doc/python.rst

@@ -96,13 +96,13 @@ perturbations will not work unless the user explicitly synchronizes incoming eve
 The ``launch_passive`` function returns a handle which can be used to interact with the viewer. It has the following
 attributes:
 
-- ``scn``, ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvScene`, :ref:`mjvCamera`,
-  :ref:`mjvOption`, and :ref:`mjvPerturb` structs, respectively.
+- ``cam``, ``opt``, and ``pert`` properties: correspond to :ref:`mjvCamera`, :ref:`mjvOption`, and :ref:`mjvPerturb`
+  structs, respectively.
 
 - ``lock()``: provides a mutex lock for the viewer as a context manager. Since the viewer operates its own
   thread, user code must ensure that it is holding the viewer lock before modifying any physics or visualization
-  state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``scn``,
-  ``cam``, ``opt``, and ``pert`` properties of the viewer handle.
+  state. These include the ``mjModel`` and ``mjData`` instance passed to ``launch_passive``, and also the ``cam``,
+  ``opt``, and ``pert`` properties of the viewer handle.
 
 - ``sync()``: synchronizes state between ``mjModel``, ``mjData``, and GUI user inputs since the previous call to
   ``sync``. In order to allow user scripts to make arbitrary modifications to ``mjModel`` and ``mjData`` without
@@ -124,6 +124,44 @@ attributes:
 - ``is_running()``: returns ``True`` if the viewer window is running and ``False`` if it is closed.
   This method can be safely called without locking.
 
+- ``user_scn``: an :ref:`mjvScene` object that allows users to add change rendering flags and add custom
+  visualization geoms to the rendered scene. This is separate from the ``mjvScene`` that the viewer uses internally to
+  render the final scene, and is entirely under the user's control. User scripts can call e.g. :ref:`mjv_initGeom` or
+  :ref:`mjv_makeConnector` to add visualization geoms to ``user_scn``, and upon the next call to ``sync()``, the viewer
+  will incorporate these geoms to future rendered images. Similarly, user scripts can make changes to ``user_scn.flags``
+  which would be picked up at the next call to ``sync()``. The ``sync()`` call also copies changes to rendering flags
+  made via the GUI back into ``user_scn`` to preserve consistency. For example:
+
+  .. code-block:: python
+
+    with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
+
+      # Enable wireframe rendering of the entire scene.
+      viewer.user_scn.flags[mujoco.mjtRndFlag.mjRND_WIREFRAME] = 1
+      viewer.sync()
+
+      while viewer.is_running():
+        ...
+        # Step the physics.
+        mujoco.mj_step(m, d)
+
+        # Add a 3x3x3 grid of variously colored spheres to the middle of the scene.
+        viewer.user_scn.ngeom = 0
+        i = 0
+        for x, y, z in itertools.product(*((range(-1, 2),) * 3)):
+          mujoco.mjv_initGeom(
+              viewer.user_scn.geoms[i],
+              type=mujoco.mjtGeom.mjGEOM_SPHERE,
+              size=[0.02, 0, 0],
+              pos=0.1*np.array([x, y, z]),
+              mat=np.eye(3).flatten(),
+              rgba=0.5*np.array([x + 1, y + 1, z + 1, 2])
+          )
+          i += 1
+        viewer.user_scn.ngeom = i
+        viewer.sync()
+        ...
+
 The viewer handle can also be used as a context manager which calls ``close()`` automatically upon exit. A minimal
 example of a user script that uses ``launch_passive`` might look like the following. (Note that example is a simple
 illustrative example that does **not** necessarily keep the physics ticking at the correct wallclock rate.)
@@ -160,29 +198,33 @@ illustrative example that does **not** necessarily keep the physics ticking at t
       if time_until_next_step > 0:
         time.sleep(time_until_next_step)
 
-Optionally, ``viewer.launch_passive`` also accepts a callable as a keyword argument ``key_callback``, which gets called
-each time a keyboard event occurs in the viewer window. This allows user scripts to react to various key presses, e.g.,
-pause or resume the run loop when the spacebar is pressed.
+Optionally, ``viewer.launch_passive`` accepts the following keyword arguments.
 
-.. code-block:: python
+- ``key_callback``: A callable which gets called each time a keyboard event occurs in the viewer window. This allows
+  user scripts to react to various key presses, e.g., pause or resume the run loop when the spacebar is pressed.
 
-  paused = False
+  .. code-block:: python
 
-  def key_callback(keycode):
-    if chr(keycode) == ' ':
-      nonlocal paused
-      paused = not paused
+    paused = False
 
-  ...
+    def key_callback(keycode):
+      if chr(keycode) == ' ':
+        nonlocal paused
+        paused = not paused
 
-  with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
-    while viewer.is_running():
-      ...
-      if not paused:
-        mujoco.mj_step(m, d)
-        viewer.sync()
-      ...
+    ...
 
+    with mujoco.viewer.launch_passive(m, d, key_callback=key_callback) as viewer:
+      while viewer.is_running():
+        ...
+        if not paused:
+          mujoco.mj_step(m, d)
+          viewer.sync()
+        ...
+
+- ``show_left_ui`` and ``show_right_ui``: Boolean arguments indicating whether UI panels should be visible
+  or hidden when the viewer is launched. Note that regardless of the values specified, the user can still toggle the
+  visibility of these panels after launch by pressing Tab or Shift+Tab.
 
 .. _PyUsage:
 
@@ -430,28 +472,32 @@ Open-loop rollouts
 ==================
 
 We include a code sample showing how to add additional C/C++ functionality, exposed as a Python module via pybind11. The
-sample, implemented in ``rollout.cc`` and wrapped in ``rollout.py``, implements a common use case where tight loops
-implemented outside of Python are beneficial: rolling out a trajectory (i.e., calling ``mj_step()`` in a loop), given an
-intial state and sequence of controls, and returning subsequent states and sensor values. The canonical usage form is
+sample, implemented in `rollout.cc <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.cc>`__
+and wrapped in `rollout.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout.py>`__,
+implements a common use case where tight loops implemented outside of Python are beneficial: rolling out a trajectory
+(i.e., calling ``mj_step()`` in a loop), given an intial state and sequence of controls, and returning subsequent states
+and sensor values. The basic usage form is
 
 .. code-block:: python
 
-   state, sensordata = rollout.rollout(model, data, initial_state, ctrl)
+   state, sensordata = rollout.rollout(model, data, initial_state, control)
 
-``initial_state`` is a ``nstate x nqva`` array, with ``nstate`` initial states of length ``nqva``, where ``nqva =
-model.nq + model.nv + model.na`` is the size of the full MuJoCo mechanical state: positions (``data.qpos``), velocities
-(``data.qvel``) and actuator activations (``data.act``). ``ctrl`` is a ``nstate x nstep x nu`` array of control
-sequences.
+``initial_state`` is a ``nroll x nstate`` array, with ``nroll`` initial states of size ``nstate``, where
+``nstate = mj_stateSize(model, mjtState.mjSTATE_FULLPHYSICS)`` is the size of the
+:ref:`full physics state<geFullPhysics>`. ``control`` is a ``nroll x nstep x ncontrol`` array of controls. Controls are
+by default the ``mjModel.nu`` standard actuators, but any combination of :ref:`user input<geInput>` arrays can be
+specified by passing an optional ``control_spec`` bitflag.
+
+If a rollout diverges, the current state and sensor values are used to fill the remainder of the trajectory.
+Therefore, non-increasing time values can be used to detect diverged rollouts.
 
 The ``rollout`` function is designed to be completely stateless, so all inputs of the stepping pipeline are set and any
-values already present in the given ``MjData`` instance will have no effect on the output. In order to facilitate this,
-all inputs including ``time`` and ``qacc_warmstart`` are set to default values, as are auxillary controls
-(``qfrc_applied``, ``xfrc_applied`` and ``mocap_{pos,quat}``). These can also be optionally set by the user.
+values already present in the given ``MjData`` instance will have no effect on the output.
 
 Since the Global Interpreter Lock can be released, this function can be efficiently threaded using Python threads. See
 the ``test_threading`` function in
-`rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`_ for an example of
-threaded operation.
+`rollout_test.py <https://github.com/google-deepmind/mujoco/blob/main/python/mujoco/rollout_test.py>`__ for an example
+of threaded operation (and more generally for usage examples).
 
 .. _PyMjpy_migration:
 

doc/requirements.txt

@@ -1,4 +1,4 @@
-Sphinx==4.5.0
+Sphinx==5.3.0
 furo==2022.9.29
 sphinxcontrib-bibtex==2.6.1
 sphinxcontrib-katex==0.9.4
@@ -11,7 +11,7 @@ nbsphinx==0.9.1
 pandoc==1.1.0
 pygments==2.15.0
 jq==1.4.1
-Jinja2==2.11.3
+Jinja2~=3.0
 wheel
 # see https://github.com/aws/aws-sam-cli/issues/3661 regarding markupsafe
 markupsafe==2.0.1

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.1.2.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.1.2/lib/libmujoco.so.3.1.2`` and rename it as ``libmujoco.so``.
 
 Windows
 _______

include/mujoco/mjdata.h

@@ -85,6 +85,10 @@ typedef enum mjtTimer_ {     // internal timers
   mjTIMER_POS_MAKE,          // make constraints
   mjTIMER_POS_PROJECT,       // project constraints
 
+  // breakdown of mj_collision
+  mjTIMER_COL_BROAD,         // broadphase
+  mjTIMER_COL_NARROW,        // narrowphase
+
   mjNTIMER                   // number of timers
 } mjtTimer;
 
@@ -92,10 +96,10 @@ typedef enum mjtTimer_ {     // internal timers
 //---------------------------------- mjContact -----------------------------------------------------
 
 struct mjContact_ {                // result of collision detection functions
-  // contact parameters set by geom-specific collision detector
+  // contact parameters set by near-phase collision function
   mjtNum  dist;                    // distance between nearest points; neg: penetration
   mjtNum  pos[3];                  // position of contact point: midpoint between geoms
-  mjtNum  frame[9];                // normal is in [0-2]
+  mjtNum  frame[9];                // normal is in [0-2], points from geom[0] to geom[1]
 
   // contact parameters set by mj_collideGeoms
   mjtNum  includemargin;           // include if dist<includemargin=margin-gap
@@ -108,12 +112,16 @@ struct mjContact_ {                // result of collision detection functions
   mjtNum  mu;                      // friction of regularized cone, set by mj_makeConstraint
   mjtNum  H[36];                   // cone Hessian, set by mj_updateConstraint
 
-  // contact descriptors set by mj_collideGeoms
+  // contact descriptors set by mj_collideXXX
   int     dim;                     // contact space dimensionality: 1, 3, 4 or 6
-  int     geom1;                   // id of geom 1
-  int     geom2;                   // id of geom 2
+  int     geom1;                   // id of geom 1; deprecated, use geom[0]
+  int     geom2;                   // id of geom 2; deprecated, use geom[1]
+  int     geom[2];                 // geom ids; -1 for flex
+  int     flex[2];                 // flex ids; -1 for geom
+  int     elem[2];                 // element ids; -1 for geom or flex vertex
+  int     vert[2];                 // vertex ids;  -1 for geom or flex element
 
-  // flag set by mj_instantianteEquality
+  // flag set by mj_setContact or mj_instantiateContact
   int     exclude;                 // 0: include, 1: in gap, 2: fused, 3: no dofs
 
   // address computed by mj_instantiateContact
@@ -167,7 +175,7 @@ struct mjData_ {
 
   // memory utilization stats
   size_t  maxuse_stack;                      // maximum stack allocation in bytes
-  size_t  maxuse_threadstack[mjMAXTHREADS];  // maximum stack allocation per thread in bytes
+  size_t  maxuse_threadstack[mjMAXTHREAD];   // maximum stack allocation per thread in bytes
   size_t  maxuse_arena;                      // maximum arena allocation in bytes
   int     maxuse_con;                        // maximum number of contacts
   int     maxuse_efc;                        // maximum number of scalar constraints
@@ -183,12 +191,6 @@ struct mjData_ {
   int     solver_nnz[mjNISLAND];    // number of non-zeros in Hessian or efc_AR, per island
   mjtNum  solver_fwdinv[2];         // forward-inverse comparison: qfrc, efc
 
-  // collision statistics
-  int     nbodypair_broad;   // number of body pairs in collision according to the broad-phase
-  int     nbodypair_narrow;  // number of body pairs actually in collision in the narrow-phase
-  int     ngeompair_mid;     // number of geom pairs in collision according to the mid-phase
-  int     ngeompair_narrow;  // number of geom pairs actually in collision in the narrow-phase
-
   // variable sizes
   int     ne;                // number of equality constraints
   int     nf;                // number of friction constraints
@@ -262,17 +264,26 @@ struct mjData_ {
 
   // computed by mj_fwdPosition/mj_comPos
   mjtNum* subtree_com;       // center of mass of each subtree                   (nbody x 3)
-  mjtNum* cdof;              // com-based motion axis of each dof                (nv x 6)
+  mjtNum* cdof;              // com-based motion axis of each dof (rot:lin)      (nv x 6)
   mjtNum* cinert;            // com-based body inertia and mass                  (nbody x 10)
 
+  // computed by mj_fwdPosition/mj_flex
+  mjtNum* flexvert_xpos;     // Cartesian flex vertex positions                  (nflexvert x 3)
+  mjtNum* flexelem_aabb;     // flex element bounding boxes (center, size)       (nflexelem x 6)
+  int*    flexedge_J_rownnz; // number of non-zeros in Jacobian row              (nflexedge x 1)
+  int*    flexedge_J_rowadr; // row start address in colind array                (nflexedge x 1)
+  int*    flexedge_J_colind; // column indices in sparse Jacobian                (nflexedge x nv)
+  mjtNum* flexedge_J;        // flex edge Jacobian                               (nflexedge x nv)
+  mjtNum* flexedge_length;   // flex edge lengths                                (nflexedge x 1)
+
   // computed by mj_fwdPosition/mj_tendon
   int*    ten_wrapadr;       // start address of tendon's path                   (ntendon x 1)
   int*    ten_wrapnum;       // number of wrap points in path                    (ntendon x 1)
   int*    ten_J_rownnz;      // number of non-zeros in Jacobian row              (ntendon x 1)
   int*    ten_J_rowadr;      // row start address in colind array                (ntendon x 1)
   int*    ten_J_colind;      // column indices in sparse Jacobian                (ntendon x nv)
-  mjtNum* ten_length;        // tendon lengths                                   (ntendon x 1)
   mjtNum* ten_J;             // tendon Jacobian                                  (ntendon x nv)
+  mjtNum* ten_length;        // tendon lengths                                   (ntendon x 1)
   int*    wrap_obj;          // geom id; -1: site; -2: pulley                    (nwrap*2 x 1)
   mjtNum* wrap_xpos;         // Cartesian 3D points in all path                  (nwrap*2 x 3)
 
@@ -290,31 +301,37 @@ struct mjData_ {
   mjtNum* qLDiagSqrtInv;     // 1/sqrt(diag(D))                                  (nv x 1)
 
   // computed by mj_collisionTree
+  mjtNum*  bvh_aabb_dyn;     // global bounding box (center, size)               (nbvhdynamic x 6)
   mjtByte* bvh_active;       // volume has been added to collisions              (nbvh x 1)
 
   //-------------------- POSITION, VELOCITY dependent
 
   // computed by mj_fwdVelocity
+  mjtNum* flexedge_velocity; // flex edge velocities                             (nflexedge x 1)
   mjtNum* ten_velocity;      // tendon velocities                                (ntendon x 1)
   mjtNum* actuator_velocity; // actuator velocities                              (nu x 1)
 
   // computed by mj_fwdVelocity/mj_comVel
-  mjtNum* cvel;              // com-based velocity [3D rot; 3D tran]             (nbody x 6)
-  mjtNum* cdof_dot;          // time-derivative of cdof                          (nv x 6)
+  mjtNum* cvel;              // com-based velocity (rot:lin)                     (nbody x 6)
+  mjtNum* cdof_dot;          // time-derivative of cdof (rot:lin)                (nv x 6)
 
   // computed by mj_fwdVelocity/mj_rne (without acceleration)
   mjtNum* qfrc_bias;         // C(qpos,qvel)                                     (nv x 1)
 
   // computed by mj_fwdVelocity/mj_passive
-  mjtNum* qfrc_passive;      // passive force                                    (nv x 1)
+  mjtNum* qfrc_spring;       // passive spring force                             (nv x 1)
+  mjtNum* qfrc_damper;       // passive damper force                             (nv x 1)
+  mjtNum* qfrc_gravcomp;     // passive gravity compensation force               (nv x 1)
+  mjtNum* qfrc_fluid;        // passive fluid force                              (nv x 1)
+  mjtNum* qfrc_passive;      // total passive force                              (nv x 1)
 
   // computed by mj_sensorVel/mj_subtreeVel if needed
   mjtNum* subtree_linvel;    // linear velocity of subtree com                   (nbody x 3)
   mjtNum* subtree_angmom;    // angular momentum about subtree com               (nbody x 3)
 
   // computed by mj_Euler or mj_implicit
-  mjtNum*   qH;              // L'*D*L factorization of modified M               (nM x 1)
-  mjtNum*   qHDiagInv;       // 1/diag(D) of modified M                          (nv x 1)
+  mjtNum* qH;                // L'*D*L factorization of modified M               (nM x 1)
+  mjtNum* qHDiagInv;         // 1/diag(D) of modified M                          (nv x 1)
 
   // computed by mj_resetData
   int*    D_rownnz;          // non-zeros in each row                            (nv x 1)

include/mujoco/mjmacro.h

@@ -34,6 +34,9 @@
 #define mjDISABLED(x) (m->opt.disableflags & (x))
 #define mjENABLED(x)  (m->opt.enableflags & (x))
 
+// is actuator disabled
+#define mjACTUATORDISABLED(i) (m->opt.disableactuator & (1 << m->actuator_group[i]))
+
 // annotation for functions that accept printf-like variadic arguments
 #ifndef mjPRINTFLIKE
   #if defined(__GNUC__)

include/mujoco/mjmodel.h

@@ -109,9 +109,11 @@ typedef enum mjtGeom_ {           // type of geometric shape
   mjGEOM_ARROW1,                  // arrow without wedges
   mjGEOM_ARROW2,                  // arrow in both directions
   mjGEOM_LINE,                    // line
+  mjGEOM_LINEBOX,                 // box with line edges
+  mjGEOM_FLEX,                    // flex
   mjGEOM_SKIN,                    // skin
   mjGEOM_LABEL,                   // text label
-  mjGEOM_TRIANGLE,                // triangle connecting a frame
+  mjGEOM_TRIANGLE,                // triangle
 
   mjGEOM_NONE         = 1001      // missing geom type
 } mjtGeom;
@@ -166,6 +168,7 @@ typedef enum mjtEq_ {             // type of equality constraint
   mjEQ_WELD,                      // fix relative position and orientation of two bodies
   mjEQ_JOINT,                     // couple the values of two scalar joints with cubic
   mjEQ_TENDON,                    // couple the lengths of two tendons with cubic
+  mjEQ_FLEX,                      // fix all edge lengths of a flex
   mjEQ_DISTANCE                   // unsupported, will cause an error if used
 } mjtEq;
 
@@ -228,6 +231,7 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_SITE,                     // site
   mjOBJ_CAMERA,                   // camera
   mjOBJ_LIGHT,                    // light
+  mjOBJ_FLEX,                     // flex
   mjOBJ_MESH,                     // mesh
   mjOBJ_SKIN,                     // skin
   mjOBJ_HFIELD,                   // heightfield
@@ -243,7 +247,9 @@ typedef enum mjtObj_ {            // type of MujoCo object
   mjOBJ_TEXT,                     // text
   mjOBJ_TUPLE,                    // tuple
   mjOBJ_KEY,                      // keyframe
-  mjOBJ_PLUGIN                    // plugin instance
+  mjOBJ_PLUGIN,                   // plugin instance
+
+  mjNOBJECT                       // number of object types
 } mjtObj;
 
 
@@ -353,6 +359,15 @@ typedef enum mjtLRMode_ {         // mode for actuator length range computation
 } mjtLRMode;
 
 
+typedef enum mjtFlexSelf_ {       // mode for flex selfcollide
+  mjFLEXSELF_NONE   = 0,          // no self-collisions
+  mjFLEXSELF_NARROW,              // skip midphase, go directly to narrowphase
+  mjFLEXSELF_BVH,                 // use BVH in midphase (if midphase enabled)
+  mjFLEXSELF_SAP,                 // use SAP in midphase
+  mjFLEXSELF_AUTO                 // choose between BVH and SAP automatically
+} mjtFlexSelf;
+
+
 //---------------------------------- mjLROpt -------------------------------------------------------
 
 struct mjLROpt_ {                 // options for mj_setLengthRange()
@@ -408,6 +423,7 @@ struct mjOption_ {                // physics options
   mjtNum o_margin;                // margin
   mjtNum o_solref[mjNREF];        // solref
   mjtNum o_solimp[mjNIMP];        // solimp
+  mjtNum o_friction[5];           // friction
 
   // discrete settings
   int integrator;                 // integration mode (mjtIntegrator)
@@ -420,6 +436,7 @@ struct mjOption_ {                // physics options
   int mpr_iterations;             // maximum number of MPR solver iterations
   int disableflags;               // bit flags for disabling standard features
   int enableflags;                // bit flags for enabling optional features
+  int disableactuator;            // bit flags for disabling actuators by group id
 
   // sdf collision settings
   int sdf_initpoints;             // number of starting points for gradient descent
@@ -439,9 +456,10 @@ struct mjVisual_ {                // visualization options
     float linewidth;              // line width for wireframe and ray rendering
     float glow;                   // glow coefficient for selected body
     float realtime;               // initial real-time factor (1: real time)
-    int offwidth;                 // width of offscreen buffer
-    int offheight;                // height of offscreen buffer
-    int ellipsoidinertia;         // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   offwidth;               // width of offscreen buffer
+    int   offheight;              // height of offscreen buffer
+    int   ellipsoidinertia;       // geom for inertia visualization (0: box, 1: ellipsoid)
+    int   bvactive;               // visualize active bounding volumes (0: no, 1: yes)
   } global;
 
   struct {                        // rendering quality
@@ -492,6 +510,7 @@ struct mjVisual_ {                // visualization options
     float framewidth;             // bodyframe axis width
     float constraint;             // constraint width
     float slidercrank;            // slidercrank width
+    float frustum;                // frustum zfar plane
   } scale;
 
   struct {                        // color of decor elements
@@ -517,6 +536,9 @@ struct mjVisual_ {                // visualization options
     float constraint[4];          // constraint
     float slidercrank[4];         // slidercrank
     float crankbroken[4];         // used when crank must be stretched/broken
+    float frustum[4];             // camera frustum
+    float bv[4];                  // bounding volume
+    float bvactive[4];            // active bounding volume
   } rgba;
 };
 typedef struct mjVisual_ mjVisual;
@@ -545,12 +567,22 @@ struct mjModel_ {
   int nu;                         // number of actuators/controls = dim(ctrl)
   int na;                         // number of activation states = dim(act)
   int nbody;                      // number of bodies
-  int nbvh;                       // number of bounding volumes in all bodies
+  int nbvh;                       // number of total bounding volumes in all bodies
+  int nbvhstatic;                 // number of static bounding volumes (aabb stored in mjModel)
+  int nbvhdynamic;                // number of dynamic bounding volumes (aabb stored in mjData)
   int njnt;                       // number of joints
   int ngeom;                      // number of geoms
   int nsite;                      // number of sites
   int ncam;                       // number of cameras
   int nlight;                     // number of lights
+  int nflex;                      // number of flexes
+  int nflexvert;                  // number of vertices in all flexes
+  int nflexedge;                  // number of edges in all flexes
+  int nflexelem;                  // number of elements in all flexes
+  int nflexelemdata;              // number of element vertex ids in all flexes
+  int nflexshelldata;             // number of shell fragment vertex ids in all flexes
+  int nflexevpair;                // number of element-vertex pairs in all flexes
+  int nflextexcoord;              // number of vertices with texture coordinates
   int nmesh;                      // number of meshes
   int nmeshvert;                  // number of vertices in all meshes
   int nmeshnormal;                // number of normals in all meshes
@@ -638,7 +670,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // body is simple (has diagonal M)          (nbody x 1)
+  mjtByte*  body_simple;          // 1: diag M; 2: diag M, sliders only       (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)
@@ -649,16 +681,19 @@ struct mjModel_ {
   mjtNum*   body_inertia;         // diagonal inertia in ipos/iquat frame     (nbody x 3)
   mjtNum*   body_invweight0;      // mean inv inert in qpos0 (trn, rot)       (nbody x 2)
   mjtNum*   body_gravcomp;        // antigravity force, units of body weight  (nbody x 1)
+  mjtNum*   body_margin;          // MAX over all geom margins                (nbody x 1)
   mjtNum*   body_user;            // user data                                (nbody x nuser_body)
   int*      body_plugin;          // plugin instance id; -1: not in use       (nbody x 1)
+  int*      body_contype;         // OR over all geom contypes                (nbody x 1)
+  int*      body_conaffinity;     // OR over all geom conaffinities           (nbody x 1)
   int*      body_bvhadr;          // address of bvh root                      (nbody x 1)
   int*      body_bvhnum;          // number of bounding volumes               (nbody x 1)
 
   // bounding volume hierarchy
   int*      bvh_depth;            // depth in the bounding volume hierarchy   (nbvh x 1)
   int*      bvh_child;            // left and right children in tree          (nbvh x 2)
-  int*      bvh_geomid;           // geom id of the node; -1: non-leaf        (nbvh x 1)
-  mjtNum*   bvh_aabb;             // bounding box of node (center, size)      (nbvh x 6)
+  int*      bvh_nodeid;           // geom or elem id of node; -1: non-leaf    (nbvh x 1)
+  mjtNum*   bvh_aabb;             // local bounding box (center, size)        (nbvhstatic x 6)
 
   // joints
   int*      jnt_type;             // type of joint (mjtJoint)                 (njnt x 1)
@@ -767,6 +802,60 @@ struct mjModel_ {
   float*    light_diffuse;        // diffuse rgb (alpha=1)                    (nlight x 3)
   float*    light_specular;       // specular rgb (alpha=1)                   (nlight x 3)
 
+  // flexes: contact properties
+  int*      flex_contype;         // flex contact type                        (nflex x 1)
+  int*      flex_conaffinity;     // flex contact affinity                    (nflex x 1)
+  int*      flex_condim;          // contact dimensionality (1, 3, 4, 6)      (nflex x 1)
+  int*      flex_priority;        // flex contact priority                    (nflex x 1)
+  mjtNum*   flex_solmix;          // mix coef for solref/imp in contact pair  (nflex x 1)
+  mjtNum*   flex_solref;          // constraint solver reference: contact     (nflex x mjNREF)
+  mjtNum*   flex_solimp;          // constraint solver impedance: contact     (nflex x mjNIMP)
+  mjtNum*   flex_friction;        // friction for (slide, spin, roll)         (nflex x 3)
+  mjtNum*   flex_margin;          // detect contact if dist<margin            (nflex x 1)
+  mjtNum*   flex_gap;             // include in solver if dist<margin-gap     (nflex x 1)
+  mjtByte*  flex_internal;        // internal flex collision enabled          (nflex x 1)
+  int*      flex_selfcollide;     // self collision mode (mjtFlexSelf)        (nflex x 1)
+  int*      flex_activelayers;    // number of active element layers, 3D only (nflex x 1)
+
+  // flexes: other properties
+  int*      flex_dim;             // 1: lines, 2: triangles, 3: tetrahedra    (nflex x 1)
+  int*      flex_matid;           // material id for rendering                (nflex x 1)
+  int*      flex_group;           // group for visibility                     (nflex x 1)
+  int*      flex_vertadr;         // first vertex address                     (nflex x 1)
+  int*      flex_vertnum;         // number of vertices                       (nflex x 1)
+  int*      flex_edgeadr;         // first edge address                       (nflex x 1)
+  int*      flex_edgenum;         // number of edges                          (nflex x 1)
+  int*      flex_elemadr;         // first element address                    (nflex x 1)
+  int*      flex_elemnum;         // number of elements                       (nflex x 1)
+  int*      flex_elemdataadr;     // first element vertex id address          (nflex x 1)
+  int*      flex_shellnum;        // number of shells                         (nflex x 1)
+  int*      flex_shelldataadr;    // first shell data address                 (nflex x 1)
+  int*      flex_evpairadr;       // first evpair address                     (nflex x 1)
+  int*      flex_evpairnum;       // number of evpairs                        (nflex x 1)
+  int*      flex_texcoordadr;     // address in flex_texcoord; -1: none       (nflex x 1)
+  int*      flex_vertbodyid;      // vertex body ids                          (nflexvert x 1)
+  int*      flex_edge;            // edge vertex ids (2 per edge)             (nflexedge x 2)
+  int*      flex_elem;            // element vertex ids (dim+1 per elem)      (nflexelemdata x 1)
+  int*      flex_elemlayer;       // element distance from surface, 3D only   (nflexelem x 1)
+  int*      flex_shell;           // shell fragment vertex ids (dim per frag) (nflexshelldata x 1)
+  int*      flex_evpair;          // (element, vertex) collision pairs        (nflexevpair x 2)
+  mjtNum*   flex_vert;            // vertex positions in local body frames    (nflexvert x 3)
+  mjtNum*   flex_xvert0;          // Cartesian vertex positions in qpos0      (nflexvert x 3)
+  mjtNum*   flexedge_length0;     // edge lengths in qpos0                    (nflexedge x 1)
+  mjtNum*   flexedge_invweight0;  // edge inv. weight in qpos0                (nflexedge x 1)
+  mjtNum*   flex_radius;          // radius around primitive element          (nflex x 1)
+  mjtNum*   flex_edgestiffness;   // edge stiffness                           (nflex x 1)
+  mjtNum*   flex_edgedamping;     // edge damping                             (nflex x 1)
+  mjtByte*  flex_edgeequality;    // is edge equality constraint defined      (nflex x 1)
+  mjtByte*  flex_rigid;           // are all verices in the same body         (nflex x 1)
+  mjtByte*  flexedge_rigid;       // are both edge vertices in same body      (nflexedge x 1)
+  mjtByte*  flex_centered;        // are all vertex coordinates (0,0,0)       (nflex x 1)
+  mjtByte*  flex_flatskin;        // render flex skin with flat shading       (nflex x 1)
+  int*      flex_bvhadr;          // address of bvh root; -1: no bvh          (nflex x 1)
+  int*      flex_bvhnum;          // number of bounding volumes               (nflex x 1)
+  float*    flex_rgba;            // rgba when material is omitted            (nflex x 4)
+  float*    flex_texcoord;        // vertex texture coordinates               (nflextexcoord x 2)
+
   // meshes
   int*      mesh_vertadr;         // first vertex address                     (nmesh x 1)
   int*      mesh_vertnum;         // number of vertices                       (nmesh x 1)
@@ -812,20 +901,23 @@ struct mjModel_ {
   int*      skin_bonebodyid;      // body id of each bone                     (nskinbone x 1)
   int*      skin_bonevertid;      // mesh ids of vertices in each bone        (nskinbonevert x 1)
   float*    skin_bonevertweight;  // weights of vertices in each bone         (nskinbonevert x 1)
+  int*      skin_pathadr;         // address of asset path for skin; -1: none (nskin x 1)
 
   // height fields
-  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                  (nhfield x 4)
-  int*      hfield_nrow;          // number of rows in grid                   (nhfield x 1)
-  int*      hfield_ncol;          // number of columns in grid                (nhfield x 1)
-  int*      hfield_adr;           // address in hfield_data                   (nhfield x 1)
-  float*    hfield_data;          // elevation data                           (nhfielddata x 1)
+  mjtNum*   hfield_size;          // (x, y, z_top, z_bottom)                    (nhfield x 4)
+  int*      hfield_nrow;          // number of rows in grid                     (nhfield x 1)
+  int*      hfield_ncol;          // number of columns in grid                  (nhfield x 1)
+  int*      hfield_adr;           // address in hfield_data                     (nhfield x 1)
+  float*    hfield_data;          // elevation data                             (nhfielddata x 1)
+  int*      hfield_pathadr;       // address of asset path for hfield; -1: none (nhfield x 1)
 
   // textures
-  int*      tex_type;             // texture type (mjtTexture)                (ntex x 1)
-  int*      tex_height;           // number of rows in texture image          (ntex x 1)
-  int*      tex_width;            // number of columns in texture image       (ntex x 1)
-  int*      tex_adr;              // address in rgb                           (ntex x 1)
-  mjtByte*  tex_rgb;              // rgb (alpha = 1)                          (ntexdata x 1)
+  int*      tex_type;             // texture type (mjtTexture)                  (ntex x 1)
+  int*      tex_height;           // number of rows in texture image            (ntex x 1)
+  int*      tex_width;            // number of columns in texture image         (ntex x 1)
+  int*      tex_adr;              // address in rgb                             (ntex x 1)
+  mjtByte*  tex_rgb;              // rgb (alpha = 1)                            (ntexdata x 1)
+  int*      tex_pathadr;         // address of asset path for texture; -1: none (ntex x 1)
 
   // materials
   int*      mat_texid;            // texture id; -1: none                     (nmat x 1)
@@ -841,7 +933,7 @@ struct mjModel_ {
   int*      pair_dim;             // contact dimensionality                   (npair x 1)
   int*      pair_geom1;           // id of geom1                              (npair x 1)
   int*      pair_geom2;           // id of geom2                              (npair x 1)
-  int*      pair_signature;       // (body1+1) << 16 + body2+1                (npair x 1)
+  int*      pair_signature;       // body1 << 16 + body2                      (npair x 1)
   mjtNum*   pair_solref;          // solver reference: contact normal         (npair x mjNREF)
   mjtNum*   pair_solreffriction;  // solver reference: contact friction       (npair x mjNREF)
   mjtNum*   pair_solimp;          // solver impedance: contact                (npair x mjNIMP)
@@ -850,7 +942,7 @@ struct mjModel_ {
   mjtNum*   pair_friction;        // tangent1, 2, spin, roll1, 2              (npair x 5)
 
   // excluded body pairs for collision detection
-  int*      exclude_signature;    // (body1+1) << 16 + body2+1                (nexclude x 1)
+  int*      exclude_signature;    // body1 << 16 + body2                      (nexclude x 1)
 
   // equality constraints
   int*      eq_type;              // constraint type (mjtEq)                  (neq x 1)
@@ -970,6 +1062,7 @@ struct mjModel_ {
   int*      name_siteadr;         // site name pointers                       (nsite x 1)
   int*      name_camadr;          // camera name pointers                     (ncam x 1)
   int*      name_lightadr;        // light name pointers                      (nlight x 1)
+  int*      name_flexadr;         // flex name pointers                       (nflex x 1)
   int*      name_meshadr;         // mesh name pointers                       (nmesh x 1)
   int*      name_skinadr;         // skin name pointers                       (nskin x 1)
   int*      name_hfieldadr;       // hfield name pointers                     (nhfield x 1)

include/mujoco/mjplugin.h

@@ -17,6 +17,7 @@
 
 #include <mujoco/mjdata.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>
 #include <mujoco/mjvisualize.h>
 
 
@@ -106,6 +107,15 @@ struct mjpPlugin_ {
   // called by mjv_updateScene (optional)
   void (*visualize)(const mjModel*m, mjData* d, const mjvOption* opt, mjvScene* scn, int instance);
 
+  // methods specific to actuators (optional)
+
+  // dimension of the actuator state for the plugin (excluding state from actuator's dyntype)
+  int (*actuator_actdim)(const mjModel*m, int instance, int actuator_id);
+
+  // updates the actuator plugin's entries in act_dot
+  // called after native act_dot is computed and before the compute callback
+  void (*actuator_act_dot)(const mjModel* m, mjData* d, int instance);
+
   // methods specific to signed distance fields (optional)
 
   // signed distance from the surface
@@ -117,6 +127,9 @@ struct mjpPlugin_ {
   // called during compilation for marching cubes
   mjtNum (*sdf_staticdistance)(const mjtNum point[3], const mjtNum* attributes);
 
+  // convert attributes and provide defaults if not present
+  void (*sdf_attribute)(mjtNum attribute[], const char* name[], const char* value[]);
+
   // bounding box of implicit surface
   void (*sdf_aabb)(mjtNum aabb[6], const mjtNum* attributes);
 };

include/mujoco/mjrender.h

@@ -43,6 +43,10 @@ typedef enum mjtFramebuffer_ {    // OpenGL framebuffer option
   mjFB_OFFSCREEN                  // offscreen buffer
 } mjtFramebuffer;
 
+typedef enum mjtDepthMap_ {       // depth mapping for `mjr_readPixels`
+  mjDEPTH_ZERONEAR    = 0,        // standard depth map; 0: znear, 1: zfar
+  mjDEPTH_ZEROFAR     = 1         // reversed depth map; 1: znear, 0: zfar
+} mjtDepthMap;
 
 typedef enum mjtFontScale_ {      // font scale, used at context creation
   mjFONTSCALE_50      = 50,       // 50% scale, suitable for low-res rendering
@@ -155,6 +159,9 @@ struct mjrContext_ {              // custom OpenGL context
 
   // pixel output format
   int     readPixelFormat;        // default color pixel format for mjr_readPixels
+
+  // depth output format
+  int     readDepthMap;           // depth mapping: mjDEPTH_ZERONEAR or mjDEPTH_ZEROFAR
 };
 typedef struct mjrContext_ mjrContext;
 

include/mujoco/mjthread.h

@@ -15,7 +15,7 @@
 #ifndef MUJOCO_INCLUDE_MJTHREAD_H_
 #define MUJOCO_INCLUDE_MJTHREAD_H_
 
-#define mjMAXTHREADS 128       // maximum number of threads in a thread pool
+#define mjMAXTHREAD 128        // maximum number of threads in a thread pool
 
 typedef enum mjtTaskStatus_ {  // status values for mjTask
   mjTASK_NEW = 0,              // newly created

include/mujoco/mjui.h

@@ -18,7 +18,7 @@
 #include <mujoco/mjrender.h>
 
 #define mjMAXUISECT     10      // maximum number of sections
-#define mjMAXUIITEM     100     // maximum number of items per section
+#define mjMAXUIITEM     200     // maximum number of items per section
 #define mjMAXUITEXT     300     // maximum number of chars in edittext and other
 #define mjMAXUINAME     40      // maximum number of chars in name
 #define mjMAXUIMULTI    35      // maximum number of radio/select items in group

include/mujoco/mjvisualize.h

@@ -74,6 +74,7 @@ typedef enum mjtLabel_ {          // object labeling
   mjLABEL_TENDON,                 // tendon labels
   mjLABEL_ACTUATOR,               // actuator labels
   mjLABEL_CONSTRAINT,             // constraint labels
+  mjLABEL_FLEX,                   // flex labels
   mjLABEL_SKIN,                   // skin labels
   mjLABEL_SELECTION,              // selected object
   mjLABEL_SELPNT,                 // coordinates of selection point
@@ -124,7 +125,12 @@ typedef enum mjtVisFlag_ {        // flags enabling model element visualization
   mjVIS_SELECT,                   // selection point
   mjVIS_STATIC,                   // static bodies
   mjVIS_SKIN,                     // skin
-  mjVIS_MIDPHASE,                 // mid-phase bounding volume hierarchy
+  mjVIS_FLEXVERT,                 // flex vertices
+  mjVIS_FLEXEDGE,                 // flex edges
+  mjVIS_FLEXFACE,                 // flex element faces
+  mjVIS_FLEXSKIN,                 // flex smooth skin (disables the rest)
+  mjVIS_BODYBVH,                  // body bounding volume hierarchy
+  mjVIS_FLEXBVH,                  // flex bounding volume hierarchy
   mjVIS_MESHBVH,                  // mesh bounding volume hierarchy
   mjVIS_SDFITER,                  // iterations of SDF gradient descent
 
@@ -159,6 +165,7 @@ typedef enum mjtStereo_ {         // type of stereo rendering
 
 struct mjvPerturb_ {              // object selection and perturbation
   int      select;                // selected body id; non-positive: none
+  int      flexselect;            // selected flex id; negative: none
   int      skinselect;            // selected skin id; negative: none
   int      active;                // perturbation bitmask (mjtPertBit)
   int      active2;               // secondary perturbation bitmask (mjtPertBit)
@@ -219,7 +226,7 @@ struct mjvGeom_ {                 // abstract geom
   int      category;              // visual category
   int      texid;                 // texture id; -1: no texture
   int      texuniform;            // uniform cube mapping
-  int      texcoord;              // mesh geom has texture coordinates
+  int      texcoord;              // mesh or flex geom has texture coordinates
   int      segid;                 // segmentation id; -1: not shown
 
   // OpenGL info
@@ -270,9 +277,11 @@ struct mjvOption_ {                  // abstract visualization options
   mjtByte  jointgroup[mjNGROUP];     // joint visualization by group
   mjtByte  tendongroup[mjNGROUP];    // tendon visualization by group
   mjtByte  actuatorgroup[mjNGROUP];  // actuator visualization by group
+  mjtByte  flexgroup[mjNGROUP];      // flex visualization by group
   mjtByte  skingroup[mjNGROUP];      // skin visualization by group
   mjtByte  flags[mjNVISFLAG];        // visualization flags (indexed by mjtVisFlag)
-  int bvh_depth;                     // depth of the bounding volume hierarchy to be visualized
+  int      bvh_depth;                // depth of the bounding volume hierarchy to be visualized
+  int      flex_layer;               // element layer to be visualized for 3D flex
 };
 typedef struct mjvOption_ mjvOption;
 
@@ -286,13 +295,32 @@ struct mjvScene_ {                // abstract scene passed to OpenGL renderer
   mjvGeom* geoms;                 // buffer for geoms (ngeom)
   int*     geomorder;             // buffer for ordering geoms by distance to camera (ngeom)
 
+  // flex data
+  int      nflex;                 // number of flexes
+  int*     flexedgeadr;           // address of flex edges (nflex)
+  int*     flexedgenum;           // number of edges in flex (nflex)
+  int*     flexvertadr;           // address of flex vertices (nflex)
+  int*     flexvertnum;           // number of vertices in flex (nflex)
+  int*     flexfaceadr;           // address of flex faces (nflex)
+  int*     flexfacenum;           // number of flex faces allocated (nflex)
+  int*     flexfaceused;          // number of flex faces currently in use (nflex)
+  int*     flexedge;              // flex edge data (2*nflexedge)
+  float*   flexvert;              // flex vertices (3*nflexvert)
+  float*   flexface;              // flex faces vertices (9*sum(flexfacenum))
+  float*   flexnormal;            // flex face normals (9*sum(flexfacenum))
+  float*   flextexcoord;          // flex face texture coordinates (6*sum(flexfacenum))
+  mjtByte  flexvertopt;           // copy of mjVIS_FLEXVERT mjvOption flag
+  mjtByte  flexedgeopt;           // copy of mjVIS_FLEXEDGE mjvOption flag
+  mjtByte  flexfaceopt;           // copy of mjVIS_FLEXFACE mjvOption flag
+  mjtByte  flexskinopt;           // copy of mjVIS_FLEXSKIN mjvOption flag
+
   // skin data
   int      nskin;                 // number of skins
   int*     skinfacenum;           // number of faces in skin (nskin)
   int*     skinvertadr;           // address of skin vertices (nskin)
   int*     skinvertnum;           // number of vertices in skin (nskin)
-  float*   skinvert;              // skin vertex data (nskin)
-  float*   skinnormal;            // skin normal data (nskin)
+  float*   skinvert;              // skin vertex data (3*nskinvert)
+  float*   skinnormal;            // skin normal data (3*nskinvert)
 
   // OpenGL lights
   int      nlight;                // number of lights currently in buffer
@@ -375,7 +403,7 @@ struct mjvSceneState_ {
   int nbuffer;                     // size of the buffer in bytes
   void* buffer;                    // heap-allocated memory for all arrays in this struct
   int maxgeom;                     // maximum number of mjvGeom supported by this state object
-  mjvScene plugincache;            // scratch space for vis geoms inserted by plugins
+  mjvScene scratch;                // scratch space for vis geoms inserted by the user and plugins
 
   // fields in mjModel that are necessary to re-render a scene
   struct {
@@ -384,6 +412,7 @@ struct mjvSceneState_ {
     int na;
     int nbody;
     int nbvh;
+    int nbvhstatic;
     int njnt;
     int ngeom;
     int nsite;
@@ -391,6 +420,9 @@ struct mjvSceneState_ {
     int nlight;
     int nmesh;
     int nskin;
+    int nflex;
+    int nflexvert;
+    int nflextexcoord;
     int nskinvert;
     int nskinface;
     int nskinbone;
@@ -404,6 +436,7 @@ struct mjvSceneState_ {
     int nnames;
     int npaths;
     int nsensordata;
+    int narena;
 
     mjOption opt;
     mjVisual vis;
@@ -427,7 +460,7 @@ struct mjvSceneState_ {
 
     int* bvh_depth;
     int* bvh_child;
-    int* bvh_geomid;
+    int* bvh_nodeid;
     mjtNum* bvh_aabb;
 
     int* jnt_type;
@@ -455,6 +488,7 @@ struct mjvSceneState_ {
 
     mjtNum* cam_fovy;
     mjtNum* cam_ipd;
+    float* cam_intrinsic;
     float* cam_sensorsize;
 
     mjtByte* light_directional;
@@ -467,6 +501,28 @@ struct mjvSceneState_ {
     float* light_diffuse;
     float* light_specular;
 
+    mjtByte* flex_flatskin;
+    int* flex_dim;
+    int* flex_matid;
+    int* flex_group;
+    int* flex_vertadr;
+    int* flex_vertnum;
+    int* flex_elem;
+    int* flex_elemlayer;
+    int* flex_elemadr;
+    int* flex_elemnum;
+    int* flex_elemdataadr;
+    int* flex_shell;
+    int* flex_shellnum;
+    int* flex_shelldataadr;
+    int* flex_texcoordadr;
+    int* flex_bvhadr;
+    int* flex_bvhnum;
+    mjtNum* flex_radius;
+    float* flex_rgba;
+
+    int* hfield_pathadr;
+
     int* mesh_bvhadr;
     int* mesh_bvhnum;
     int* mesh_texcoordadr;
@@ -493,6 +549,9 @@ struct mjvSceneState_ {
     int* skin_bonebodyid;
     int* skin_bonevertid;
     float* skin_bonevertweight;
+    int* skin_pathadr;
+
+    int* tex_pathadr;
 
     int* mat_texid;
     mjtByte* mat_texuniform;
@@ -591,6 +650,7 @@ struct mjvSceneState_ {
     mjtNum* ten_length;
     mjtNum* wrap_xpos;
 
+    mjtNum* bvh_aabb_dyn;
     mjtByte* bvh_active;
     int* island_dofadr;
     int* island_dofind;
@@ -598,8 +658,11 @@ struct mjvSceneState_ {
     int* efc_island;
     int* tendon_efcadr;
 
+    mjtNum* flexvert_xpos;
+
     mjContact* contact;
     mjtNum* efc_force;
+    void* arena;
   } data;
 };
 typedef struct mjvSceneState_ mjvSceneState;

include/mujoco/mjxmacro.h

@@ -43,6 +43,7 @@
     X( int,     mpr_iterations    ) \
     X( int,     disableflags      ) \
     X( int,     enableflags       ) \
+    X( int,     disableactuator   ) \
     X( int,     sdf_initpoints    ) \
     X( int,     sdf_iterations    )
 
@@ -58,7 +59,8 @@
     X( wind,            3       )   \
     X( magnetic,        3       )   \
     X( o_solref,        mjNREF  )   \
-    X( o_solimp,        mjNIMP  )
+    X( o_solimp,        mjNIMP  )   \
+    X( o_friction,      5       )
 
 
 //-------------------------------- mjModel ---------------------------------------------------------
@@ -71,11 +73,21 @@
     XMJV( na )                 \
     XMJV( nbody )              \
     XMJV( nbvh )               \
+    XMJV( nbvhstatic )         \
+    X   ( nbvhdynamic )        \
     XMJV( njnt )               \
     XMJV( ngeom )              \
     XMJV( nsite )              \
     XMJV( ncam )               \
     XMJV( nlight )             \
+    XMJV( nflex )              \
+    XMJV( nflexvert )          \
+    X   ( nflexedge )          \
+    X   ( nflexelem )          \
+    X   ( nflexelemdata )      \
+    X   ( nflexshelldata )     \
+    X   ( nflexevpair )        \
+    XMJV( nflextexcoord )      \
     XMJV( nmesh )              \
     X   ( nmeshvert )          \
     X   ( nmeshnormal )        \
@@ -185,14 +197,17 @@
     XMJV( mjtNum,  body_inertia,          nbody,         3                    ) \
     X   ( mjtNum,  body_invweight0,       nbody,         2                    ) \
     X   ( mjtNum,  body_gravcomp,         nbody,         1                    ) \
+    X   ( mjtNum,  body_margin,           nbody,         1                    ) \
     X   ( mjtNum,  body_user,             nbody,         MJ_M(nuser_body)     ) \
     X   ( int,     body_plugin,           nbody,         1                    ) \
+    X   ( int,     body_contype,          nbody,         1                    ) \
+    X   ( int,     body_conaffinity,      nbody,         1                    ) \
     XMJV( int,     body_bvhadr,           nbody,         1                    ) \
     XMJV( int,     body_bvhnum,           nbody,         1                    ) \
     XMJV( int,     bvh_depth,             nbvh,          1                    ) \
     XMJV( int,     bvh_child,             nbvh,          2                    ) \
-    XMJV( int,     bvh_geomid,            nbvh,          1                    ) \
-    XMJV( mjtNum,  bvh_aabb,              nbvh,          6                    ) \
+    XMJV( int,     bvh_nodeid,            nbvh,          1                    ) \
+    XMJV( mjtNum,  bvh_aabb,              nbvhstatic,    6                    ) \
     XMJV( int,     jnt_type,              njnt,          1                    ) \
     X   ( int,     jnt_qposadr,           njnt,          1                    ) \
     X   ( int,     jnt_dofadr,            njnt,          1                    ) \
@@ -262,7 +277,7 @@
     X   ( int,     cam_targetbodyid,      ncam,          1                    ) \
     X   ( int,     cam_resolution,        ncam,          2                    ) \
     XMJV( float,   cam_sensorsize,        ncam,          2                    ) \
-    X   ( float,   cam_intrinsic,         ncam,          4                    ) \
+    XMJV( float,   cam_intrinsic,         ncam,          4                    ) \
     X   ( mjtNum,  cam_pos,               ncam,          3                    ) \
     X   ( mjtNum,  cam_quat,              ncam,          4                    ) \
     X   ( mjtNum,  cam_poscom0,           ncam,          3                    ) \
@@ -288,6 +303,56 @@
     XMJV( float,   light_ambient,         nlight,        3                    ) \
     XMJV( float,   light_diffuse,         nlight,        3                    ) \
     XMJV( float,   light_specular,        nlight,        3                    ) \
+    X   ( int,     flex_contype,          nflex,         1                    ) \
+    X   ( int,     flex_conaffinity,      nflex,         1                    ) \
+    X   ( int,     flex_condim,           nflex,         1                    ) \
+    X   ( int,     flex_priority,         nflex,         1                    ) \
+    X   ( mjtNum,  flex_solmix,           nflex,         1                    ) \
+    X   ( mjtNum,  flex_solref,           nflex,         mjNREF               ) \
+    X   ( mjtNum,  flex_solimp,           nflex,         mjNIMP               ) \
+    X   ( mjtNum,  flex_friction,         nflex,         3                    ) \
+    X   ( mjtNum,  flex_margin,           nflex,         1                    ) \
+    X   ( mjtNum,  flex_gap,              nflex,         1                    ) \
+    X   ( mjtByte, flex_internal,         nflex,         1                    ) \
+    X   ( int,     flex_selfcollide,      nflex,         1                    ) \
+    X   ( int,     flex_activelayers,     nflex,         1                    ) \
+    XMJV( int,     flex_dim,              nflex,         1                    ) \
+    XMJV( int,     flex_matid,            nflex,         1                    ) \
+    XMJV( int,     flex_group,            nflex,         1                    ) \
+    XMJV( int,     flex_vertadr,          nflex,         1                    ) \
+    XMJV( int,     flex_vertnum,          nflex,         1                    ) \
+    X   ( int,     flex_edgeadr,          nflex,         1                    ) \
+    X   ( int,     flex_edgenum,          nflex,         1                    ) \
+    XMJV( int,     flex_elemadr,          nflex,         1                    ) \
+    XMJV( int,     flex_elemnum,          nflex,         1                    ) \
+    XMJV( int,     flex_elemdataadr,      nflex,         1                    ) \
+    XMJV( int,     flex_shellnum,         nflex,         1                    ) \
+    XMJV( int,     flex_shelldataadr,     nflex,         1                    ) \
+    X   ( int,     flex_evpairadr,        nflex,         1                    ) \
+    X   ( int,     flex_evpairnum,        nflex,         1                    ) \
+    XMJV( int,     flex_texcoordadr,      nflex,         1                    ) \
+    X   ( int,     flex_vertbodyid,       nflexvert,     1                    ) \
+    X   ( int,     flex_edge,             nflexedge,     2                    ) \
+    XMJV( int,     flex_elem,             nflexelemdata, 1                    ) \
+    XMJV( int,     flex_elemlayer,        nflexelem,     1                    ) \
+    XMJV( int,     flex_shell,            nflexshelldata,1                    ) \
+    X   ( int,     flex_evpair,           nflexevpair,   2                    ) \
+    X   ( mjtNum,  flex_vert,             nflexvert,     3                    ) \
+    X   ( mjtNum,  flex_xvert0,           nflexvert,     3                    ) \
+    X   ( mjtNum,  flexedge_length0,      nflexedge,     1                    ) \
+    X   ( mjtNum,  flexedge_invweight0,   nflexedge,     1                    ) \
+    XMJV( mjtNum,  flex_radius,           nflex,         1                    ) \
+    X   ( mjtNum,  flex_edgestiffness,    nflex,         1                    ) \
+    X   ( mjtNum,  flex_edgedamping,      nflex,         1                    ) \
+    X   ( mjtByte, flex_edgeequality,     nflex,         1                    ) \
+    X   ( mjtByte, flex_rigid,            nflex,         1                    ) \
+    X   ( mjtByte, flexedge_rigid,        nflexedge,     1                    ) \
+    X   ( mjtByte, flex_centered,         nflex,         1                    ) \
+    XMJV( mjtByte, flex_flatskin,         nflex,         1                    ) \
+    XMJV( int,     flex_bvhadr,           nflex,         1                    ) \
+    XMJV( int,     flex_bvhnum,           nflex,         1                    ) \
+    XMJV( float,   flex_rgba,             nflex,         4                    ) \
+    X   ( float,   flex_texcoord,         nflextexcoord, 2                    ) \
     X   ( int,     mesh_vertadr,          nmesh,         1                    ) \
     X   ( int,     mesh_vertnum,          nmesh,         1                    ) \
     X   ( int,     mesh_normaladr,        nmesh,         1                    ) \
@@ -330,16 +395,19 @@
     XMJV( int,     skin_bonebodyid,       nskinbone,     1                    ) \
     XMJV( int,     skin_bonevertid,       nskinbonevert, 1                    ) \
     XMJV( float,   skin_bonevertweight,   nskinbonevert, 1                    ) \
+    XMJV( int,     skin_pathadr,          nskin,         1                    ) \
     X   ( mjtNum,  hfield_size,           nhfield,       4                    ) \
     X   ( int,     hfield_nrow,           nhfield,       1                    ) \
     X   ( int,     hfield_ncol,           nhfield,       1                    ) \
     X   ( int,     hfield_adr,            nhfield,       1                    ) \
     X   ( float,   hfield_data,           nhfielddata,   1                    ) \
+    XMJV( int,     hfield_pathadr,        nhfield,       1                    ) \
     X   ( int,     tex_type,              ntex,          1                    ) \
     X   ( int,     tex_height,            ntex,          1                    ) \
     X   ( int,     tex_width,             ntex,          1                    ) \
     X   ( int,     tex_adr,               ntex,          1                    ) \
     X   ( mjtByte, tex_rgb,               ntexdata,      1                    ) \
+    XMJV( int,     tex_pathadr,           ntex,          1                    ) \
     XMJV( int,     mat_texid,             nmat,          1                    ) \
     XMJV( mjtByte, mat_texuniform,        nmat,          1                    ) \
     XMJV( float,   mat_texrepeat,         nmat,          2                    ) \
@@ -456,6 +524,7 @@
     XMJV( int,     name_siteadr,          nsite,         1                    ) \
     XMJV( int,     name_camadr,           ncam,          1                    ) \
     XMJV( int,     name_lightadr,         nlight,        1                    ) \
+    X   ( int,     name_flexadr,          nflex,         1                    ) \
     X   ( int,     name_meshadr,          nmesh,         1                    ) \
     X   ( int,     name_skinadr,          nskin,         1                    ) \
     X   ( int,     name_hfieldadr,        nhfield,       1                    ) \
@@ -522,6 +591,13 @@
     XMJV( mjtNum,    subtree_com,       nbody,       3           ) \
     X   ( mjtNum,    cdof,              nv,          6           ) \
     X   ( mjtNum,    cinert,            nbody,       10          ) \
+    XMJV( mjtNum,    flexvert_xpos,     nflexvert,   3           ) \
+    X   ( mjtNum,    flexelem_aabb,     nflexelem,   6           ) \
+    X   ( int,       flexedge_J_rownnz, nflexedge,   1           ) \
+    X   ( int,       flexedge_J_rowadr, nflexedge,   1           ) \
+    X   ( int,       flexedge_J_colind, nflexedge,   MJ_M(nv)    ) \
+    X   ( mjtNum,    flexedge_J,        nflexedge,   MJ_M(nv)    ) \
+    X   ( mjtNum,    flexedge_length,   nflexedge,   1           ) \
     XMJV( int,       ten_wrapadr,       ntendon,     1           ) \
     XMJV( int,       ten_wrapnum,       ntendon,     1           ) \
     X   ( int,       ten_J_rownnz,      ntendon,     1           ) \
@@ -538,12 +614,18 @@
     X   ( mjtNum,    qLD,               nM,          1           ) \
     X   ( mjtNum,    qLDiagInv,         nv,          1           ) \
     X   ( mjtNum,    qLDiagSqrtInv,     nv,          1           ) \
+    XMJV( mjtNum,    bvh_aabb_dyn,      nbvhdynamic, 6           ) \
     XMJV( mjtByte,   bvh_active,        nbvh,        1           ) \
+    X   ( mjtNum,    flexedge_velocity, nflexedge,   1           ) \
     X   ( mjtNum,    ten_velocity,      ntendon,     1           ) \
     X   ( mjtNum,    actuator_velocity, nu,          1           ) \
     X   ( mjtNum,    cvel,              nbody,       6           ) \
     X   ( mjtNum,    cdof_dot,          nv,          6           ) \
     X   ( mjtNum,    qfrc_bias,         nv,          1           ) \
+    X   ( mjtNum,    qfrc_spring,       nv,          1           ) \
+    X   ( mjtNum,    qfrc_damper,       nv,          1           ) \
+    X   ( mjtNum,    qfrc_gravcomp,     nv,          1           ) \
+    X   ( mjtNum,    qfrc_fluid,        nv,          1           ) \
     X   ( mjtNum,    qfrc_passive,      nv,          1           ) \
     X   ( mjtNum,    subtree_linvel,    nbody,       3           ) \
     X   ( mjtNum,    subtree_angmom,    nbody,       3           ) \
@@ -644,10 +726,6 @@
     X( int,       maxuse_con         ) \
     X( int,       maxuse_efc         ) \
     X( int,       solver_nisland     ) \
-    X( int,       nbodypair_broad    ) \
-    X( int,       nbodypair_narrow   ) \
-    X( int,       ngeompair_mid      ) \
-    X( int,       ngeompair_narrow   ) \
     X( int,       ne                 ) \
     X( int,       nf                 ) \
     X( int,       nl                 ) \
@@ -661,7 +739,7 @@
 
 // vector fields of mjData
 #define MJDATA_VECTOR                                                \
-    X( size_t,         maxuse_threadstack, mjMAXTHREADS, 1         ) \
+    X( size_t,         maxuse_threadstack, mjMAXTHREAD,  1         ) \
     X( mjWarningStat,  warning,            mjNWARNING,   1         ) \
     X( mjTimerStat,    timer,              mjNTIMER,     1         ) \
     X( mjSolverStat,   solver,             mjNILSAND,    mjNSOLVER ) \

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 312
 
 // needed to define size_t, fabs and log10
 #include <stdlib.h>
@@ -264,6 +264,9 @@ MJAPI void mj_Euler(const mjModel* m, mjData* d);
 // Runge-Kutta explicit order-N integrator.
 MJAPI void mj_RungeKutta(const mjModel* m, mjData* d, int N);
 
+// Implicit-in-velocity integrators.
+MJAPI void mj_implicit(const mjModel* m, mjData* d);
+
 // Run position-dependent computations in inverse dynamics.
 MJAPI void mj_invPosition(const mjModel* m, mjData* d);
 
@@ -312,6 +315,9 @@ MJAPI void mj_comPos(const mjModel* m, mjData* d);
 // Compute camera and light positions and orientations.
 MJAPI void mj_camlight(const mjModel* m, mjData* d);
 
+// Compute flex-related quantities.
+MJAPI void mj_flex(const mjModel* m, mjData* d);
+
 // Compute tendon lengths, velocities and moment arms.
 MJAPI void mj_tendon(const mjModel* m, mjData* d);
 
@@ -333,10 +339,10 @@ MJAPI void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, i
 // Compute cvel, cdof_dot.
 MJAPI void mj_comVel(const mjModel* m, mjData* d);
 
-// Compute qfrc_passive from spring-dampers, viscosity and density.
+// Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.
 MJAPI void mj_passive(const mjModel* m, mjData* d);
 
-// subtree linear velocity and angular momentum
+// Sub-tree linear velocity and angular momentum: compute subtree_linvel, subtree_angmom.
 MJAPI void mj_subtreeVel(const mjModel* m, mjData* d);
 
 // RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term.
@@ -518,6 +524,12 @@ MJAPI mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
 MJAPI mjtNum mju_rayGeom(const mjtNum pos[3], const mjtNum mat[9], const mjtNum size[3],
                          const mjtNum pnt[3], const mjtNum vec[3], int geomtype);
 
+// Intersect ray with flex, return nearest distance or -1 if no intersection,
+// and also output nearest vertex id.
+MJAPI mjtNum mju_rayFlex(const mjModel* m, const mjData* d, int flex_layer, mjtByte flg_vert,
+                         mjtByte flg_edge, mjtByte flg_face, mjtByte flg_skin, int flexid,
+                         const mjtNum* pnt, const mjtNum* vec, int vertid[1]);
+
 // Intersect ray with skin, return nearest distance or -1 if no intersection,
 // and also output nearest vertex id.
 MJAPI mjtNum mju_raySkin(int nface, int nvert, const int* face, const float* vert,
@@ -593,10 +605,11 @@ MJAPI void mjv_applyPerturbForce(const mjModel* m, mjData* d, const mjvPerturb*
 // Return the average of two OpenGL cameras.
 MJAPI mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera* cam2);
 
-// Select geom or skin with mouse, return bodyid; -1: none selected.
+// Select geom, flex or skin with mouse, return bodyid; -1: none selected.
 MJAPI int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
                      mjtNum aspectratio, mjtNum relx, mjtNum rely,
-                     const mjvScene* scn, mjtNum selpnt[3], int geomid[1], int skinid[1]);
+                     const mjvScene* scn, mjtNum selpnt[3],
+                     int geomid[1], int flexid[1], int skinid[1]);
 
 
 //---------------------------------- Visualization -------------------------------------------------
@@ -1105,7 +1118,7 @@ MJAPI void mju_bandMulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
 // Address of diagonal element i in band-dense matrix representation.
 MJAPI int mju_bandDiag(int i, int ntotal, int nband, int ndense);
 
-// Eigenvalue decomposition of symmetric 3x3 matrix.
+// Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.
 MJAPI int mju_eig3(mjtNum eigval[3], mjtNum eigvec[9], mjtNum quat[4], const mjtNum mat[9]);
 
 // minimize 0.5*x'*H*x + x'*g  s.t. lower <= x <= upper, return rank or -1 if failed

introspect/enums.py

@@ -90,9 +90,11 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjGEOM_ARROW1', 101),
              ('mjGEOM_ARROW2', 102),
              ('mjGEOM_LINE', 103),
-             ('mjGEOM_SKIN', 104),
-             ('mjGEOM_LABEL', 105),
-             ('mjGEOM_TRIANGLE', 106),
+             ('mjGEOM_LINEBOX', 104),
+             ('mjGEOM_FLEX', 105),
+             ('mjGEOM_SKIN', 106),
+             ('mjGEOM_LABEL', 107),
+             ('mjGEOM_TRIANGLE', 108),
              ('mjGEOM_NONE', 1001),
          ]),
      )),
@@ -167,7 +169,8 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjEQ_WELD', 1),
              ('mjEQ_JOINT', 2),
              ('mjEQ_TENDON', 3),
-             ('mjEQ_DISTANCE', 4),
+             ('mjEQ_FLEX', 4),
+             ('mjEQ_DISTANCE', 5),
          ]),
      )),
     ('mjtWrap',
@@ -246,22 +249,24 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjOBJ_SITE', 6),
              ('mjOBJ_CAMERA', 7),
              ('mjOBJ_LIGHT', 8),
-             ('mjOBJ_MESH', 9),
-             ('mjOBJ_SKIN', 10),
-             ('mjOBJ_HFIELD', 11),
-             ('mjOBJ_TEXTURE', 12),
-             ('mjOBJ_MATERIAL', 13),
-             ('mjOBJ_PAIR', 14),
-             ('mjOBJ_EXCLUDE', 15),
-             ('mjOBJ_EQUALITY', 16),
-             ('mjOBJ_TENDON', 17),
-             ('mjOBJ_ACTUATOR', 18),
-             ('mjOBJ_SENSOR', 19),
-             ('mjOBJ_NUMERIC', 20),
-             ('mjOBJ_TEXT', 21),
-             ('mjOBJ_TUPLE', 22),
-             ('mjOBJ_KEY', 23),
-             ('mjOBJ_PLUGIN', 24),
+             ('mjOBJ_FLEX', 9),
+             ('mjOBJ_MESH', 10),
+             ('mjOBJ_SKIN', 11),
+             ('mjOBJ_HFIELD', 12),
+             ('mjOBJ_TEXTURE', 13),
+             ('mjOBJ_MATERIAL', 14),
+             ('mjOBJ_PAIR', 15),
+             ('mjOBJ_EXCLUDE', 16),
+             ('mjOBJ_EQUALITY', 17),
+             ('mjOBJ_TENDON', 18),
+             ('mjOBJ_ACTUATOR', 19),
+             ('mjOBJ_SENSOR', 20),
+             ('mjOBJ_NUMERIC', 21),
+             ('mjOBJ_TEXT', 22),
+             ('mjOBJ_TUPLE', 23),
+             ('mjOBJ_KEY', 24),
+             ('mjOBJ_PLUGIN', 25),
+             ('mjNOBJECT', 26),
          ]),
      )),
     ('mjtConstraint',
@@ -371,6 +376,18 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjLRMODE_ALL', 3),
          ]),
      )),
+    ('mjtFlexSelf',
+     EnumDecl(
+         name='mjtFlexSelf',
+         declname='enum mjtFlexSelf_',
+         values=dict([
+             ('mjFLEXSELF_NONE', 0),
+             ('mjFLEXSELF_NARROW', 1),
+             ('mjFLEXSELF_BVH', 2),
+             ('mjFLEXSELF_SAP', 3),
+             ('mjFLEXSELF_AUTO', 4),
+         ]),
+     )),
     ('mjtTaskStatus',
      EnumDecl(
          name='mjtTaskStatus',
@@ -440,7 +457,9 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjTIMER_POS_COLLISION', 10),
              ('mjTIMER_POS_MAKE', 11),
              ('mjTIMER_POS_PROJECT', 12),
-             ('mjNTIMER', 13),
+             ('mjTIMER_COL_BROAD', 13),
+             ('mjTIMER_COL_NARROW', 14),
+             ('mjNTIMER', 15),
          ]),
      )),
     ('mjtCatBit',
@@ -503,13 +522,14 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjLABEL_TENDON', 7),
              ('mjLABEL_ACTUATOR', 8),
              ('mjLABEL_CONSTRAINT', 9),
-             ('mjLABEL_SKIN', 10),
-             ('mjLABEL_SELECTION', 11),
-             ('mjLABEL_SELPNT', 12),
-             ('mjLABEL_CONTACTPOINT', 13),
-             ('mjLABEL_CONTACTFORCE', 14),
-             ('mjLABEL_ISLAND', 15),
-             ('mjNLABEL', 16),
+             ('mjLABEL_FLEX', 10),
+             ('mjLABEL_SKIN', 11),
+             ('mjLABEL_SELECTION', 12),
+             ('mjLABEL_SELPNT', 13),
+             ('mjLABEL_CONTACTPOINT', 14),
+             ('mjLABEL_CONTACTFORCE', 15),
+             ('mjLABEL_ISLAND', 16),
+             ('mjNLABEL', 17),
          ]),
      )),
     ('mjtFrame',
@@ -557,10 +577,15 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjVIS_SELECT', 21),
              ('mjVIS_STATIC', 22),
              ('mjVIS_SKIN', 23),
-             ('mjVIS_MIDPHASE', 24),
-             ('mjVIS_MESHBVH', 25),
-             ('mjVIS_SDFITER', 26),
-             ('mjNVISFLAG', 27),
+             ('mjVIS_FLEXVERT', 24),
+             ('mjVIS_FLEXEDGE', 25),
+             ('mjVIS_FLEXFACE', 26),
+             ('mjVIS_FLEXSKIN', 27),
+             ('mjVIS_BODYBVH', 28),
+             ('mjVIS_FLEXBVH', 29),
+             ('mjVIS_MESHBVH', 30),
+             ('mjVIS_SDFITER', 31),
+             ('mjNVISFLAG', 32),
          ]),
      )),
     ('mjtRndFlag',
@@ -626,6 +651,15 @@ ENUMS: Mapping[str, EnumDecl] = dict([
              ('mjFB_OFFSCREEN', 1),
          ]),
      )),
+    ('mjtDepthMap',
+     EnumDecl(
+         name='mjtDepthMap',
+         declname='enum mjtDepthMap_',
+         values=dict([
+             ('mjDEPTH_ZERONEAR', 0),
+             ('mjDEPTH_ZEROFAR', 1),
+         ]),
+     )),
     ('mjtFontScale',
      EnumDecl(
          name='mjtFontScale',

introspect/enums_test.py

@@ -61,7 +61,7 @@ class EnumsTest(absltest.TestCase):
     self.assertEqual(enum_decl.values['mjGEOM_ARROW'], 100)
     self.assertEqual(enum_decl.values['mjGEOM_ARROW1'], 101)
     self.assertEqual(enum_decl.values['mjGEOM_ARROW2'], 102)
-    self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 106)
+    self.assertEqual(enum_decl.values['mjGEOM_TRIANGLE'], 108)
     # Skip a few...
 
     self.assertEqual(enum_decl.values['mjGEOM_NONE'], 1001)

introspect/functions.py

@@ -1156,6 +1156,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Runge-Kutta explicit order-N integrator.',
      )),
+    ('mj_implicit',
+     FunctionDecl(
+         name='mj_implicit',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+         ),
+         doc='Implicit-in-velocity integrators.',
+     )),
     ('mj_invPosition',
      FunctionDecl(
          name='mj_invPosition',
@@ -1456,6 +1476,26 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Compute camera and light positions and orientations.',
      )),
+    ('mj_flex',
+     FunctionDecl(
+         name='mj_flex',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData'),
+                 ),
+             ),
+         ),
+         doc='Compute flex-related quantities.',
+     )),
     ('mj_tendon',
      FunctionDecl(
          name='mj_tendon',
@@ -1646,7 +1686,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='Compute qfrc_passive from spring-dampers, viscosity and density.',
+         doc='Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.',  # pylint: disable=line-too-long
      )),
     ('mj_subtreeVel',
      FunctionDecl(
@@ -1666,7 +1706,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='subtree linear velocity and angular momentum',
+         doc='Sub-tree linear velocity and angular momentum: compute subtree_linvel, subtree_angmom.',  # pylint: disable=line-too-long
      )),
     ('mj_rne',
      FunctionDecl(
@@ -3161,6 +3201,69 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Intersect ray with pure geom, return nearest distance or -1 if no intersection.',  # pylint: disable=line-too-long
      )),
+    ('mju_rayFlex',
+     FunctionDecl(
+         name='mju_rayFlex',
+         return_type=ValueType(name='mjtNum'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='m',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjModel', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='d',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjData', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='flex_layer',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='flg_vert',
+                 type=ValueType(name='mjtByte'),
+             ),
+             FunctionParameterDecl(
+                 name='flg_edge',
+                 type=ValueType(name='mjtByte'),
+             ),
+             FunctionParameterDecl(
+                 name='flg_face',
+                 type=ValueType(name='mjtByte'),
+             ),
+             FunctionParameterDecl(
+                 name='flg_skin',
+                 type=ValueType(name='mjtByte'),
+             ),
+             FunctionParameterDecl(
+                 name='flexid',
+                 type=ValueType(name='int'),
+             ),
+             FunctionParameterDecl(
+                 name='pnt',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='vec',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjtNum', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='vertid',
+                 type=ArrayType(
+                     inner_type=ValueType(name='int'),
+                     extents=(1,),
+                 ),
+             ),
+         ),
+         doc='Intersect ray with flex, return nearest distance or -1 if no intersection, and also output nearest vertex id.',  # pylint: disable=line-too-long
+     )),
     ('mju_raySkin',
      FunctionDecl(
          name='mju_raySkin',
@@ -3815,6 +3918,13 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                      extents=(1,),
                  ),
              ),
+             FunctionParameterDecl(
+                 name='flexid',
+                 type=ArrayType(
+                     inner_type=ValueType(name='int'),
+                     extents=(1,),
+                 ),
+             ),
              FunctionParameterDecl(
                  name='skinid',
                  type=ArrayType(
@@ -3823,7 +3933,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='Select geom or skin with mouse, return bodyid; -1: none selected.',  # pylint: disable=line-too-long
+         doc='Select geom, flex or skin with mouse, return bodyid; -1: none selected.',  # pylint: disable=line-too-long
      )),
     ('mjv_defaultOption',
      FunctionDecl(
@@ -7258,7 +7368,7 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
                  ),
              ),
          ),
-         doc='Eigenvalue decomposition of symmetric 3x3 matrix.',
+         doc="Eigenvalue decomposition of symmetric 3x3 matrix, mat = eigvec * diag(eigval) * eigvec'.",  # pylint: disable=line-too-long
      )),
     ('mju_boxQP',
      FunctionDecl(

mjx/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include mujoco/mjx/test_data *

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 introduces the MJX API
+and trains a reinforcement learning policy in a few minutes: [![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.

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

mjx/mujoco/mjx/__init__.py

@@ -0,0 +1,50 @@
+# 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.collision_driver import ncon
+from mujoco.mjx._src.constraint import count_constraints
+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 euler
+from mujoco.mjx._src.forward import forward
+from mujoco.mjx._src.forward import fwd_acceleration
+from mujoco.mjx._src.forward import fwd_actuation
+from mujoco.mjx._src.forward import fwd_position
+from mujoco.mjx._src.forward import fwd_velocity
+from mujoco.mjx._src.forward import rungekutta4
+from mujoco.mjx._src.forward import step
+from mujoco.mjx._src.io import get_data
+from mujoco.mjx._src.io import get_data_into
+from mujoco.mjx._src.io import make_data
+from mujoco.mjx._src.io import put_data
+from mujoco.mjx._src.io import put_model
+from mujoco.mjx._src.passive import passive
+from mujoco.mjx._src.ray import ray
+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 rne
+from mujoco.mjx._src.smooth import transmission
+from mujoco.mjx._src.solver import solve
+from mujoco.mjx._src.support import full_m
+from mujoco.mjx._src.support import is_sparse
+from mujoco.mjx._src.support import mul_m
+from mujoco.mjx._src.types import *

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.
+# ==============================================================================

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

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

mjx/mujoco/mjx/_src/collision_driver.py

@@ -0,0 +1,381 @@
+# 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
+
+
+# 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
+
+  # MuJoCo does not collide planes with other planes or hfields
+  if t1 == GeomType.PLANE and t2 == GeomType.PLANE:
+    return
+  if t1 == GeomType.PLANE and t2 == GeomType.HFIELD:
+    return
+
+  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,
+  )
+  return con
+
+
+def _max_contact_points(m: Union[Model, mujoco.MjModel]) -> 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: Union[Model, mujoco.MjModel]) -> int:
+  """Returns the number of contacts computed in MJX given a model."""
+  if m.opt.disableflags & DisableBit.CONTACT:
+    return 0
+
+  candidates = collision_candidates(m)
+  max_count = _max_contact_points(m)
+
+  count = 0
+  for k, v in candidates.items():
+    fn = get_collision_fn(k[0:2])
+    if fn is None:
+      continue
+    count += len(v) * fn.ncon  # pytype: disable=attribute-error
+
+  return min(max_count, count) if max_count > -1 else count
+
+
+def collision(m: Model, d: Data) -> Data:
+  """Collides geometries."""
+  if ncon(m) == 0:
+    return d.replace(contact=Contact.zero())
+
+  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:
+    raise RuntimeError('No contacts found.')
+
+  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)
+
+  return d.replace(contact=contact)

mjx/mujoco/mjx/_src/collision_driver_test.py

@@ -0,0 +1,514 @@
+# 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
+from mujoco.mjx._src import collision_driver
+from mujoco.mjx._src import test_util
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Contact
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import DisableBit
+from mujoco.mjx._src.types import Model
+# pylint: enable=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.put_model(m)
+  d = mujoco.MjData(m)
+  dx = mjx.put_data(m, 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 = """
+    <mujoco>
+      <worldbody>
+        <geom size="40 40 40" type="plane"/>
+        <body pos="0 0 0.4">
+          <joint type="free"/>
+          <geom size="0.5" type="sphere"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  _SPHERE_SPHERE = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom pos="0 0 0" size="0.2" type="sphere"/>
+        </body>
+        <body >
+          <joint type="free"/>
+          <geom pos="0 0.3 0" size="0.11" type="sphere"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  _SPHERE_CAP = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom pos="0 0.3 0" size="0.05" type="sphere"/>
+        </body>
+        <body>
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="0.0 -0.5 0.14 0.0 0.5 0.14" size="0.1" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  @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 = """
+    <mujoco>
+      <worldbody>
+        <body pos="0.52 0 0.52">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.05" type="sphere"/>
+        </body>
+        <body>
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.5 0.5 0.5" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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 = """
+    <mujoco>
+      <worldbody>
+        <geom size="40 40 40" type="plane"/>
+        <body pos="0 0 0.4">
+          <joint type="free"/>
+          <geom fromto="-1 0 0 1 0 0" size="0.5" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  _CAP_CAP = """
+    <mujoco model="two_capsules">
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom fromto="0.62235904  0.58846647 0.651046 1.5330081 0.33564585 0.977849"
+           size="0.05" type="capsule"/>
+        </body>
+        <body>
+          <joint type="free"/>
+          <geom fromto="0.5505271 0.60345304 0.476661 1.3900293 0.30709633 0.932082"
+           size="0.05" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  @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 = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom fromto="-0.5 0.1 0.25 0.5 0.1 0.25" size="0.1" type="capsule"/>
+        </body>
+        <body>
+          <joint type="free"/>
+          <geom fromto="-0.5 0.1 0.1 0.5 0.1 0.1" size="0.1" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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 = """
+    <mujoco>
+      <worldbody>
+        <body pos="0 0 0.54">
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
+        </body>
+        <body>
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.5 0.5 0.5" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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 = """
+    <mujoco>
+      <worldbody>
+        <body pos="0.5 0 0.55" euler="0 30 0">
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="-0.6 0 0 0.6 0 0" size="0.05" type="capsule"/>
+        </body>
+        <body>
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.5 0.5 0.5" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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)
+    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 = """
+    <mujoco>
+      <worldbody>
+        <geom size="40 40 40" type="plane"/>
+        <body pos="0 0 0.7" euler="45 0 0">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.5 0.5 0.5" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2)
+
+  _BOX_BOX = """
+    <mujoco>
+      <worldbody>
+        <body pos="0.0 1.0 0.2">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.2 0.2 0.2" type="box"/>
+        </body>
+        <body pos="0.1 1.0 0.495" euler="0.1 -0.1 0">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.1 0.1 0.1" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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)
+    np.testing.assert_array_less(c.dist, 0)
+    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 = """
+    <mujoco>
+      <worldbody>
+        <body pos="-1.0 -1.0 0.2">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.2 0.2 0.2" type="box"/>
+        </body>
+        <body pos="-1.0 -1.2 0.55" euler="0 45 30">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.1 0.1 0.1" type="box"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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)
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2)
+
+  _CONVEX_CONVEX = """
+    <mujoco>
+      <asset>
+        <mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.1 0.1 0.1" />
+        <mesh name="dodecahedron" file="meshes/dodecahedron.stl" scale="0.01 0.01 0.01" />
+      </asset>
+      <worldbody>
+        <body pos="0.0 2.0 0.096">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.2 0.2 0.2" type="mesh" mesh="tetrahedron"/>
+        </body>
+        <body pos="0.0 2.0 0.289" euler="0.1 -0.1 45">
+          <joint axis="1 0 0" type="free"/>
+          <geom size="0.1 0.1 0.1" type="mesh" mesh="dodecahedron"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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)
+    np.testing.assert_array_less(0, c.dist[1:])
+    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 = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom size="0.2"/>
+          <geom size="0.2"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  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 = """
+    <mujoco>
+      <worldbody>
+        <body>
+          <joint type="free"/>
+          <geom size="0.2"/>
+          <body pos="0.0 0.0 0.1">
+            <joint type="hinge"/>
+            <geom size="0.2"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_filter_parent_child(self):
+    """Tests that parent-child collisions get filtered."""
+    m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
+    mx = mjx.put_model(m)
+    d = mujoco.MjData(m)
+    dx = mjx.put_data(m, 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.put_model(m)
+    d = mujoco.MjData(m)
+    dx = mjx.put_data(m, 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 NconTest(parameterized.TestCase):
+  """Tests ncon."""
+
+  def test_ncon(self):
+    m = test_util.load_test_file('constraints.xml')
+    ncon = collision_driver.ncon(m)
+    self.assertEqual(ncon, 16)
+
+  def test_disable_contact(self):
+    m = test_util.load_test_file('constraints.xml')
+    m.opt.disableflags |= DisableBit.CONTACT
+    ncon = collision_driver.ncon(m)
+    self.assertEqual(ncon, 0)
+
+
+class TopKContactTest(absltest.TestCase):
+  """Tests top-k contacts."""
+
+  _CAPSULES = """
+    <mujoco>
+      <custom>
+        <numeric data="2" name="max_contact_points"/>
+      </custom>
+      <worldbody>
+        <body pos="0 0 0.54">
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
+        </body>
+        <body pos="0 0 0.54">
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
+        </body>
+        <body pos="0 0 0.54">
+          <joint axis="1 0 0" type="free"/>
+          <geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+  def test_top_k_contacts(self):
+    m = mujoco.MjModel.from_xml_string(self._CAPSULES)
+    mx_top_k = mjx.put_model(m)
+    mx_all = mx_top_k.replace(
+        nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([])
+    )
+    d = mujoco.MjData(m)
+    dx = mjx.put_data(m, 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.contact.dist.shape, (3,))
+    self.assertEqual(dx_top_k.contact.dist.shape, (2,))
+
+
+if __name__ == '__main__':
+  absltest.main()

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

mjx/mujoco/mjx/_src/constraint.py

@@ -0,0 +1,379 @@
+# 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 Optional, Tuple, Union
+
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco.mjx._src import collision_driver
+from mujoco.mjx._src import math
+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):
+  """Support data for creating constraint matrices."""
+  J: jax.Array
+  pos: jax.Array
+  pos_norm: jax.Array
+  invweight: jax.Array
+  solref: jax.Array
+  solimp: jax.Array
+  frictionloss: jax.Array
+
+
+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_equality_connect(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for connect equality constraints."""
+
+  ids = np.nonzero(m.eq_type == EqType.CONNECT)[0]
+
+  if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
+    return None
+
+  id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
+
+  @jax.vmap
+  def fn(data, id1, id2):
+    anchor1, anchor2 = data[0:3], data[3:6]
+    # find global points
+    pos1 = d.xmat[id1] @ anchor1 + d.xpos[id1]
+    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
+
+    return j, cpos, jp.repeat(math.norm(cpos), 3)
+
+  # concatenate to drop connect grouping dimension
+  j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
+  invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0]
+  invweight = jp.repeat(invweight, 3)
+  solref = jp.tile(m.eq_solref[ids], (3, 1))
+  solimp = jp.tile(m.eq_solimp[ids], (3, 1))
+  frictionloss = jp.zeros_like(pos_norm)
+
+  return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
+
+
+def _instantiate_equality_weld(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for weld equality constraints."""
+
+  ids = np.nonzero(m.eq_type == EqType.WELD)[0]
+
+  if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
+    return None
+
+  id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
+
+  @jax.vmap
+  def fn(data, id1, id2):
+    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 * torquescale))
+
+    return j, pos, jp.repeat(math.norm(pos), 6)
+
+  # concatenate to drop weld grouping dimension
+  j, pos, pos_norm = jax.tree_map(jp.concatenate, fn(data, id1, id2))
+  invweight = m.body_invweight0[id1] + m.body_invweight0[id2]
+  invweight = jp.repeat(invweight, 3)
+  solref = jp.tile(m.eq_solref[ids], (6, 1))
+  solimp = jp.tile(m.eq_solimp[ids], (6, 1))
+  frictionloss = jp.zeros_like(pos_norm)
+
+  return _Efc(j, pos, pos_norm, invweight, solref, solimp, frictionloss)
+
+
+def _instantiate_equality_joint(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for joint equality constraints."""
+
+  ids = np.nonzero(m.eq_type == EqType.JOINT)[0]
+
+  if (m.opt.disableflags & DisableBit.EQUALITY) or ids.size == 0:
+    return None
+
+  id1, id2, data = m.eq_obj1id[ids], m.eq_obj2id[ids], m.eq_data[ids]
+  dofadr1, dofadr2 = m.jnt_dofadr[id1], m.jnt_dofadr[id2]
+  qposadr1, qposadr2 = m.jnt_qposadr[id1], m.jnt_qposadr[id2]
+
+  @jax.vmap
+  def fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2):
+    pos1, pos2 = d.qpos[qposadr1], d.qpos[qposadr2]
+    ref1, ref2 = m.qpos0[qposadr1], m.qpos0[qposadr2]
+    pos2, ref2 = pos2 * (id2 > -1), ref2 * (id2 > -1)
+
+    dif = pos2 - ref2
+    dif_power = jp.power(dif, jp.arange(0, 5))
+
+    deriv = jp.dot(data[1:5], dif_power[:4] * jp.arange(1, 5))
+    j = jp.zeros((m.nv)).at[dofadr1].set(1.0).at[dofadr2].set(-deriv)
+    pos = pos1 - ref1 - jp.dot(data[:5], dif_power)
+    return j, pos
+
+  j, pos = fn(data, id2, dofadr1, dofadr2, qposadr1, qposadr2)
+  invweight = m.dof_invweight0[dofadr1] + m.dof_invweight0[dofadr2] * (id2 > -1)
+  solref, solimp = m.eq_solref[ids], m.eq_solimp[ids]
+  frictionloss = jp.zeros_like(pos)
+
+  return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
+
+
+def _instantiate_friction(m: Model, d: Data) -> Optional[_Efc]:
+  # TODO(robotics-team): implement _instantiate_friction
+  del m, d
+  return None
+
+
+def _instantiate_limit_ball(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for ball joint limits."""
+
+  ids = np.nonzero((m.jnt_type == JointType.BALL) & m.jnt_limited)[0]
+
+  if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
+    return None
+
+  jnt_range = m.jnt_range[ids]
+  jnt_margin = m.jnt_margin[ids]
+  qposadr = np.array([np.arange(q, q + 4) for q in m.jnt_qposadr[ids]])
+  dofadr = np.array([np.arange(d, d + 3) for d in m.jnt_dofadr[ids]])
+
+  @jax.vmap
+  def fn(jnt_range, jnt_margin, qposadr, dofadr):
+    axis, angle = math.quat_to_axis_angle(d.qpos[qposadr])
+    j = jp.zeros(m.nv).at[dofadr].set(-axis)
+    pos = jp.amax(jnt_range) - angle - jnt_margin
+    active = pos < 0
+    return j * active, pos * active
+
+  j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
+  invweight = m.dof_invweight0[m.jnt_dofadr[ids]]
+  solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
+  frictionloss = jp.zeros_like(pos)
+
+  return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
+
+
+def _instantiate_limit_slide_hinge(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for slide and hinge joint limits."""
+
+  slide_hinge = np.isin(m.jnt_type, (JointType.SLIDE, JointType.HINGE))
+  ids = np.nonzero(slide_hinge & m.jnt_limited)[0]
+
+  if (m.opt.disableflags & DisableBit.LIMIT) or ids.size == 0:
+    return None
+
+  jnt_range = m.jnt_range[ids]
+  jnt_margin = m.jnt_margin[ids]
+  qposadr = m.jnt_qposadr[ids]
+  dofadr = m.jnt_dofadr[ids]
+
+  @jax.vmap
+  def fn(jnt_range, jnt_margin, qposadr, dofadr):
+    dist_min = d.qpos[qposadr] - jnt_range[0]
+    dist_max = jnt_range[1] - d.qpos[qposadr]
+    j = jp.zeros(m.nv).at[dofadr].set((dist_min < dist_max) * 2 - 1)
+    pos = jp.minimum(dist_min, dist_max) - jnt_margin
+    active = pos < 0
+    return j * active, pos * active
+
+  j, pos = fn(jnt_range, jnt_margin, qposadr, dofadr)
+  invweight = m.dof_invweight0[dofadr]
+  solref, solimp = m.jnt_solref[ids], m.jnt_solimp[ids]
+  frictionloss = jp.zeros_like(pos)
+
+  return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
+
+
+def _instantiate_contact(m: Model, d: Data) -> Optional[_Efc]:
+  """Calculates constraint rows for contacts."""
+
+  if collision_driver.ncon(m) == 0:
+    return None
+
+  @jax.vmap
+  def fn(c: Contact):
+    dist = c.dist - c.includemargin
+    geom_bodyid = jp.array(m.geom_bodyid)
+    body1, body2 = geom_bodyid[c.geom1], geom_bodyid[c.geom2]
+    diff = support.jac_dif_pair(m, d, c.pos, body1, body2)
+    t = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0]
+
+    # rotate Jacobian differences to contact frame
+    diff_con = c.frame @ diff.T
+
+    # TODO(robotics-simulation): add support for other friction dimensions
+    # 4 pyramidal friction directions
+    js, invweights = [], []
+    for diff_tan, friction in zip(diff_con[1:], c.friction[:2]):
+      for f in (friction, -friction):
+        js.append(diff_con[0] + diff_tan * f)
+        invweights.append((t + f * f * t) * 2 * f * f)
+
+    active = dist < 0
+    j, invweight = jp.stack(js) * active, jp.stack(invweights)
+    pos = jp.repeat(dist, 4) * active
+    solref, solimp = jp.tile(c.solref, (4, 1)), jp.tile(c.solimp, (4, 1))
+
+    return j, invweight, pos, solref, solimp
+
+  res = fn(d.contact)
+  # remove contact grouping dimension:
+  j, invweight, pos, solref, solimp = jax.tree_map(jp.concatenate, res)
+  frictionloss = jp.zeros_like(pos)
+
+  return _Efc(j, pos, pos, invweight, solref, solimp, frictionloss)
+
+
+def count_constraints(
+    m: Union[Model, mujoco.MjModel]
+) -> 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_connect = (m.eq_type == EqType.CONNECT).sum()
+    ne_weld = (m.eq_type == EqType.WELD).sum()
+    ne_joint = (m.eq_type == EqType.JOINT).sum()
+    ne = ne_connect * 3 + ne_weld * 6 + ne_joint
+
+  nf = 0
+
+  if m.opt.disableflags & DisableBit.LIMIT:
+    nl = 0
+  else:
+    nl = int(m.jnt_limited.sum())
+
+  nc = collision_driver.ncon(m) * 4
+
+  return ne, nf, nl, nc
+
+
+def make_constraint(m: Model, d: Data) -> Data:
+  """Creates constraint jacobians and other supporting data."""
+
+  if m.opt.disableflags & DisableBit.CONSTRAINT:
+    efcs = ()
+  else:
+    efcs = tuple(efc for efc in (
+        _instantiate_equality_connect(m, d),
+        _instantiate_equality_weld(m, d),
+        _instantiate_equality_joint(m, d),
+        _instantiate_friction(m, d),
+        _instantiate_limit_ball(m, d),
+        _instantiate_limit_slide_hinge(m, d),
+        _instantiate_contact(m, d),
+    ) if efc is not None)
+
+  if not efcs:
+    z = jp.empty(0)
+    d = d.replace(efc_J=jp.empty((0, m.nv)))
+    d = d.replace(efc_D=z, efc_aref=z, efc_frictionloss=z)
+    return d
+
+  efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs)
+
+  @jax.vmap
+  def fn(efc):
+    k, b, imp = _kbi(m, efc.solref, efc.solimp, efc.pos_norm)
+    r = jp.maximum(efc.invweight * (1 - imp) / imp, mujoco.mjMINVAL)
+    aref = -b * (efc.J @ d.qvel) - k * imp * efc.pos
+    return aref, r
+
+  aref, r = fn(efc)
+  d = d.replace(efc_J=efc.J, efc_D=1 / r, efc_aref=aref)
+  d = d.replace(efc_frictionloss=efc.frictionloss)
+
+  return d

mjx/mujoco/mjx/_src/constraint_test.py

@@ -0,0 +1,107 @@
+# 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 jax import numpy as jp
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import constraint
+from mujoco.mjx._src import test_util
+import numpy as np
+
+
+# tolerance for difference between MuJoCo and MJX constraint calculations,
+# mostly due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10   # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ConstraintTest(absltest.TestCase):
+
+  def test_constraints(self):
+    """Test constraints."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = mjx.make_constraint(mx, dx)
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_J, dx.efc_J[nnz].reshape(-1), 'efc_J')
+    _assert_eq(d.efc_D, dx.efc_D[nnz], 'efc_D')
+    _assert_eq(d.efc_aref, dx.efc_aref[nnz], 'efc_aref')
+    _assert_eq(d.efc_frictionloss, dx.efc_frictionloss[nnz], 'efc_frictionloss')
+
+  def test_disable_refsafe(self):
+    m = test_util.load_test_file('constraints.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 | mjx.DisableBit.REFSAFE
+    mx = mjx.device_put(m)
+    k, *_ = constraint._kbi(mx, solimp, solref, pos)
+    self.assertEqual(k, 1 / (0.99**2 * timeconst**2))
+
+  def test_disable_constraint(self):
+    m = test_util.load_test_file('constraints.xml')
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONSTRAINT
+    ne, nf, nl, nc = mjx.count_constraints(m)
+    self.assertEqual(ne, 0)
+    self.assertEqual(nf, 0)
+    self.assertEqual(nl, 0)
+    self.assertEqual(nc, 0)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 0)
+
+  def test_disable_equality(self):
+    m = test_util.load_test_file('constraints.xml')
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EQUALITY
+    ne, nf, nl, nc = mjx.count_constraints(m)
+    self.assertEqual(ne, 0)
+    self.assertEqual(nf, 0)
+    self.assertEqual(nl, 2)
+    self.assertEqual(nc, 64)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 66)  # only joint range, contact
+
+  def test_disable_contact(self):
+    m = test_util.load_test_file('constraints.xml')
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.CONTACT
+    ne, nf, nl, nc = mjx.count_constraints(m)
+    self.assertEqual(ne, 10)
+    self.assertEqual(nf, 0)
+    self.assertEqual(nl, 2)
+    self.assertEqual(nc, 0)
+    dx = constraint.make_constraint(mjx.put_model(m), mjx.make_data(m))
+    self.assertEqual(dx.efc_J.shape[0], 12)  # only joint range, limit
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/dataclasses.py

@@ -0,0 +1,156 @@
+# 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 Any, 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),
+        any(dataclasses.is_dataclass(a) for a 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), val.dtype) 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[0], dtype=meta[1]))
+      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
+
+  @classmethod
+  def fields(cls) -> tuple[dataclasses.Field[Any], ...]:
+    return dataclasses.fields(cls)
+
+  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)}
+  )

mjx/mujoco/mjx/_src/device.py

@@ -0,0 +1,307 @@
+# 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
+import numpy as np
+
+_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,
+    ),
+    mujoco.mjtSolver: (
+        mujoco.MjModel.opt,
+        mujoco.MjOption.solver,
+    ),
+}
+
+_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.Data, 'site_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)),
+    (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.Data, 'site_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."""
+
+  # 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])
+    unsupported = [mj_type(t) for t in unsupported_typs]  # pylint: disable=too-many-function-args
+    if unsupported:
+      raise NotImplementedError(f'{unsupported} not implemented.')
+
+  # check condim
+  if any(dim != 3 for dim in m.geom_condim) or any(
+      dim != 3 for dim in m.pair_dim
+  ):
+    raise NotImplementedError('Only condim=3 is supported.')
+
+  if m.ntendon:
+    raise NotImplementedError('Tendons are not supported.')
+
+  # check collision geom types
+  candidate_set = collision_driver.collision_candidates(m)
+  for g1, g2, *_ in candidate_set:
+    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}.')
+
+  if not np.allclose(m.dof_frictionloss, 0):
+    raise NotImplementedError('dof_frictionloss is not implemented.')
+
+
+@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
+  """
+  warnings.warn(
+      'device_put is deprecated, use put_model and put_data instead',
+      category=DeprecationWarning,
+  )
+
+  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
+  """
+  warnings.warn(
+      'device_get_into is deprecated, use get_data instead',
+      category=DeprecationWarning,
+  )
+
+  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):
+      ncon = value.contact.dist.shape[0]
+      nefc = value.efc_J.shape[0]
+      mujoco._functions._realloc_con_efc(  # pylint: disable=protected-access
+          result, ncon=ncon, nefc=nefc
+      )
+      result.ncon = ncon
+      result.nefc = nefc
+      efc_start = nefc - ncon * 4
+      result.contact.efc_address[:] = np.arange(efc_start, nefc, 4)
+      result.contact.dim[:] = 3
+
+    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

mjx/mujoco/mjx/_src/device_test.py

@@ -0,0 +1,203 @@
+# 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('constraints.xml', 'pendula.xml')
+  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('constraints.xml', 'pendula.xml')
+  def testdevice_get(self, fname):
+    """Test getting MjData from a device."""
+    m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
+    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('constraints.xml', 'pendula.xml')
+  def testdevice_get_batched(self, fname):
+    """Test getting MjData from a device."""
+    m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
+    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(
+        '<mujoco><option solver="PGS"/><worldbody/></mujoco>'
+    )
+    with self.assertRaises(NotImplementedError):
+      mjx.device_put(m)
+
+  def test_integrator(self):
+    m = mujoco.MjModel.from_xml_string(
+        '<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
+    )
+    with self.assertRaises(NotImplementedError):
+      mjx.device_put(m)
+
+  def test_cone(self):
+    m = mujoco.MjModel.from_xml_string(
+        '<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
+    )
+    with self.assertRaises(NotImplementedError):
+      mjx.device_put(m)
+
+  def test_dyn(self):
+    m = test_util.load_test_file('pendula.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('pendula.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('pendula.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('constraints.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("""
+      <mujoco>
+        <worldbody>
+          <body>
+            <joint axis="1 0 0" type="free"/>
+            <geom size="0.2 0.2 0.2" type="box"/>
+          </body>
+          <body>
+            <joint axis="1 0 0" type="free"/>
+            <geom size="0.1 0.1" type="cylinder"/>
+          </body>
+        </worldbody>
+      </mujoco>
+    """)
+    with self.assertRaises(NotImplementedError):
+      mjx.device_put(m)
+
+  def test_tendon(self):
+    m = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <worldbody>
+        <body name="left_thigh" pos="0 0.1 -0.04">
+          <joint axis="0 1 0" name="left_hip_y" type="hinge"/>
+          <geom fromto="0 0 0 0 -0.01 -.34" name="left_thigh1" size="0.06" type="capsule"/>
+          <body name="left_shin" pos="0 -0.01 -0.403">
+            <joint axis="0 -1 0" name="left_knee" pos="0 0 .02" range="-160 -2" type="hinge"/>
+            <geom fromto="0 0 0 0 0 -.3" name="left_shin1" size="0.049" type="capsule"/>
+          </body>
+        </body>
+        </worldbody>
+        <tendon>
+          <fixed name="left_hipknee">
+            <joint coef="-1" joint="left_hip_y"/>
+            <joint coef="1" joint="left_knee"/>
+          </fixed>
+        </tendon>
+      </mujoco>
+    """)
+    with self.assertRaises(NotImplementedError):
+      mjx.device_put(m)
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/forward.py

@@ -0,0 +1,371 @@
+# 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
+# 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 fwd_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 = collision_driver.collision(m, d)
+  d = constraint.make_constraint(m, d)
+  d = smooth.transmission(m, d)
+  return d
+
+
+@named_scope
+def fwd_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 fwd_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 in (DynType.FILTER, DynType.FILTEREXACT):
+      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 fwd_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,))
+
+
+def _next_activation(m: Model, d: Data, act_dot: jax.Array) -> jax.Array:
+  """Returns the next act given the current act_dot, after clamping."""
+  act = d.act
+
+  if not m.na:
+    return act
+
+  actrange = jp.where(
+      m.actuator_actlimited[:, None],
+      m.actuator_actrange,
+      jp.array([-jp.inf, jp.inf]),
+  )
+
+  def fn(dyntype, dynprm, act, act_dot, actrange):
+    if dyntype == DynType.FILTEREXACT:
+      tau = jp.clip(dynprm[0], a_min=mujoco.mjMINVAL)
+      act = act + act_dot * tau * (1 - jp.exp(-m.opt.timestep / tau))
+    else:
+      act = act + act_dot * m.opt.timestep
+    act = jp.clip(act, actrange[0], actrange[1])
+    return act
+
+  args = (m.actuator_dyntype, m.actuator_dynprm, act, act_dot, actrange)
+  act = scan.flat(m, fn, 'uuaau', 'a', *args, group_by='u')
+
+  return act.reshape(m.na)
+
+
+@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 = _next_activation(m, d, act_dot)
+
+  # 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
+    dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
+    dh = smooth.factor_m(m, dh)
+    qfrc = d.qfrc_smooth + d.qfrc_constraint
+    qacc = smooth.solve_m(m, dh, qfrc)
+  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 = fwd_position(m, d)
+  d = fwd_velocity(m, d)
+  d = fwd_actuation(m, d)
+  d = fwd_acceleration(m, d)
+
+  if d.efc_J.size == 0:
+    d = d.replace(qacc=d.qacc_smooth)
+    return d
+
+  d = named_scope(solver.solve)(m, d)
+
+  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

mjx/mujoco/mjx/_src/forward_test.py

@@ -0,0 +1,178 @@
+# 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
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+
+# tolerance for difference between MuJoCo and MJX forward calculations - mostly
+# due to float precision
+_TOLERANCE = 1e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class ForwardTest(absltest.TestCase):
+
+  def test_forward(self):
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    # apply some control and xfrc input
+    d.ctrl = np.array([-18, 0.59, 0.47])
+    d.xfrc_applied[0, 2] = 0.1  # torque
+    d.xfrc_applied[1, 4] = 0.3  # linear force
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
+    mujoco.mj_forward(m, d)
+
+    mx = mjx.put_model(m)
+
+    # fwd_actuation
+    dx = jax.jit(mjx.fwd_actuation)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'act_dot')
+    _assert_attr_eq(d, dx, 'qfrc_actuator')
+
+    # fwd_accleration (fwd_position and fwd_velocity already tested elsewhere)
+    dx = jax.jit(mjx.fwd_acceleration)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_smooth')
+    _assert_attr_eq(d, dx, 'qacc_smooth')
+
+    # euler
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
+    mujoco.mj_Euler(m, d)
+    _assert_attr_eq(d, dx, 'act')
+    _assert_attr_eq(d, dx, 'qpos')
+    _assert_attr_eq(d, dx, 'time')
+
+  def test_step(self):
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    # apply some control and xfrc input
+    d.ctrl = np.array([-18, 0.59, 0.47])
+    d.xfrc_applied[0, 2] = 0.1  # torque
+    d.xfrc_applied[1, 4] = 0.3  # linear force
+    mujoco.mj_step(m, d, 100)  # get some dynamics going
+
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
+    mujoco.mj_step(m, d)
+    _assert_attr_eq(d, dx, 'act')
+    _assert_attr_eq(d, dx, 'time')
+    _assert_attr_eq(d, dx, 'qvel')
+    _assert_attr_eq(d, dx, 'qpos')
+
+  def test_rk4(self):
+    m = mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <option integrator="RK4">
+            <flag constraint="disable"/>
+          </option>
+          <worldbody>
+            <geom type="plane" size="1 1 .01" pos="0 0 -1"/>
+            <body pos="0.15 0 0">
+              <joint type="hinge" axis="0 1 0"/>
+              <geom type="capsule" size="0.02" fromto="0 0 0 .1 0 0"/>
+              <body pos="0.1 0 0">
+                <joint type="slide" axis="1 0 0" stiffness="200"/>
+                <geom type="capsule" size="0.015" fromto="-.1 0 0 .1 0 0"/>
+              </body>
+            </body>
+          </worldbody>
+        </mujoco>
+        """)
+
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.array([0.2, -0.1])
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d))
+    mujoco.mj_RungeKutta(m, d, 4)
+
+    _assert_attr_eq(d, dx, 'qvel')
+    _assert_attr_eq(d, dx, 'qpos')
+    _assert_attr_eq(d, dx, 'act')
+    _assert_attr_eq(d, dx, 'time')
+
+  def test_disable_eulerdamp(self):
+    m = test_util.load_test_file('pendula.xml')
+    self.assertTrue((m.dof_damping > 0).any())
+    m.opt.disableflags = m.opt.disableflags | mjx.DisableBit.EULERDAMP
+
+    d = mujoco.MjData(m)
+    d.qvel[:] = 1.0
+    d.qacc[:] = 1.0
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
+
+    np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep)
+
+
+class ActuatorTest(absltest.TestCase):
+  _DYN_XML = """
+    <mujoco>
+      <compiler autolimits="true"/>
+      <worldbody>
+        <body name="box">
+          <joint name="slide1" type="slide" axis="1 0 0" />
+          <joint name="slide2" type="slide" axis="0 1 0" />
+          <joint name="slide3" type="slide" axis="0 0 1" />
+          <joint name="slide4" type="slide" axis="1 1 0" />
+          <geom type="box" size=".05 .05 .05" mass="1"/>
+        </body>
+      </worldbody>
+      <actuator>
+        <general joint="slide1" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide2" dyntype="integrator" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide3" dyntype="filter" dynprm="0.1" gainprm="1.1" />
+        <general joint="slide4" dyntype="filterexact" dynprm="0.1" gainprm="1.1" />
+      </actuator>
+    </mujoco>
+  """
+
+  def test_dyntype(self):
+    m = mujoco.MjModel.from_xml_string(self._DYN_XML)
+    d = mujoco.MjData(m)
+    d.ctrl = np.array([1.5, 1.5, 1.5, 1.5])
+    d.act = np.array([0.5, 0.5, 0.5])
+
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_fwdActuation(m, d)
+    dx = jax.jit(mjx.fwd_actuation)(mx, dx)
+    _assert_attr_eq(d, dx, 'act_dot')
+
+    mujoco.mj_Euler(m, d)
+    dx = jax.jit(mjx.euler)(mx, dx)
+    _assert_attr_eq(d, dx, 'act')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/io.py

@@ -0,0 +1,412 @@
+# 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."""
+
+import copy
+from typing import List, Union
+
+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 mesh
+from mujoco.mjx._src import support
+from mujoco.mjx._src import types
+import numpy as np
+import scipy
+
+
+def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
+  """Puts mujoco.MjOption onto a device, resulting in mjx.Option."""
+  if o.integrator not in set(types.IntegratorType):
+    raise NotImplementedError(f'{mujoco.mjtIntegrator(o.integrator)}')
+
+  if o.cone not in set(types.ConeType):
+    raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
+
+  if o.jacobian not in set(types.JacobianType):
+    raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}')
+
+  if o.solver not in set(types.SolverType):
+    raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
+
+  for i in range(mujoco.mjtEnableBit.mjNENABLE):
+    if o.enableflags & 2**i:
+      raise NotImplementedError(f'{mujoco.mjtEnableBit(2 ** i)}')
+
+  static_fields = {
+      f.name: copy.copy(getattr(o, f.name))
+      for f in types.Option.fields()
+      if f.type in (int, bytes, np.ndarray)
+  }
+  static_fields['integrator'] = types.IntegratorType(o.integrator)
+  static_fields['cone'] = types.ConeType(o.cone)
+  static_fields['jacobian'] = types.JacobianType(o.jacobian)
+  static_fields['solver'] = types.SolverType(o.solver)
+  static_fields['disableflags'] = types.DisableBit(o.disableflags)
+
+  device_fields = {
+      f.name: copy.copy(getattr(o, f.name))
+      for f in types.Option.fields()
+      if f.type is jax.Array
+  }
+  device_fields = jax.device_put(device_fields, device=device)
+
+  has_fluid_params = o.density > 0 or o.viscosity > 0 or o.wind.any()
+
+  return types.Option(
+      has_fluid_params=has_fluid_params,
+      **static_fields,
+      **device_fields,
+  )
+
+
+def _put_statistic(s: mujoco.MjStatistic, device=None) -> types.Statistic:
+  """Puts mujoco.MjStatistic onto a device, resulting in mjx.Statistic."""
+  return types.Statistic(
+      meaninertia=jax.device_put(s.meaninertia, device=device)
+  )
+
+
+def put_model(m: mujoco.MjModel, device=None) -> types.Model:
+  """Puts mujoco.MjModel onto a device, resulting in mjx.Model."""
+
+  if m.ntendon:
+    raise NotImplementedError('tendons are not supported')
+
+  if (m.geom_condim != 3).any() or (m.pair_dim != 3).any():
+    raise NotImplementedError('only condim=3 is supported')
+
+  # check collision geom types
+  for g1, g2, *_ in collision_driver.collision_candidates(m):
+    if collision_driver.get_collision_fn((g1, g2)) is None:
+      g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2)
+      raise NotImplementedError(f'({g1}, {g2}) has no collision function')
+
+  for enum_field, enum_type, mj_type in (
+      (m.actuator_biastype, types.BiasType, mujoco.mjtBias),
+      (m.actuator_dyntype, types.DynType, mujoco.mjtDyn),
+      (m.actuator_gaintype, types.GainType, mujoco.mjtGain),
+      (m.actuator_trntype, types.TrnType, mujoco.mjtTrn),
+      (m.eq_type, types.EqType, mujoco.mjtEq),
+  ):
+    missing = set(enum_field) - set(enum_type)
+    if missing:
+      raise NotImplementedError(
+          f'{[mj_type(m) for m in missing]} not supported'
+      )
+
+  if not np.allclose(m.dof_frictionloss, 0):
+    raise NotImplementedError('dof_frictionloss is not implemented.')
+
+  opt = _put_option(m.opt, device=device)
+  stat = _put_statistic(m.stat, device=device)
+
+  static_fields = {
+      f.name: getattr(m, f.name)
+      for f in types.Model.fields()
+      if f.type in (int, bytes, np.ndarray)
+  }
+  static_fields['geom_rgba'] = static_fields['geom_rgba'].reshape((-1, 4))
+  static_fields['mat_rgba'] = static_fields['mat_rgba'].reshape((-1, 4))
+
+  device_fields = {
+      f.name: copy.copy(getattr(m, f.name))  # copy because device_put is async
+      for f in types.Model.fields()
+      if f.type is jax.Array
+  }
+  device_fields.update(mesh.get(m))
+  device_fields = jax.device_put(device_fields, device=device)
+
+  return types.Model(
+      opt=opt,
+      stat=stat,
+      **static_fields,
+      **device_fields,
+  )
+
+
+def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
+  """Allocate and initialize Data."""
+
+  ncon = collision_driver.ncon(m)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  zero_0 = jp.zeros(0, dtype=jp.float32)
+  zero_nv = jp.zeros(m.nv, dtype=jp.float32)
+  zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
+  zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=jp.float32)
+  zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32)
+  zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32)
+  zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32)
+  zero_nbody_3_3 = jp.zeros((m.nbody, 3, 3), dtype=jp.float32)
+  zero_nefc = jp.zeros(nefc, dtype=jp.float32)
+  zero_na = jp.zeros(m.na, dtype=jp.float32)
+  zero_nu = jp.zeros(m.nu, dtype=jp.float32)
+  zero_njnt_3 = jp.zeros((m.njnt, 3), dtype=jp.float32)
+  zero_nm = jp.zeros(m.nM, dtype=jp.float32)
+
+  # create first d to get num contacts and nc
+  d = types.Data(
+      solver_niter=jp.array(0, dtype=jp.int32),
+      time=jp.array(0.0),
+      qpos=jp.array(m.qpos0),
+      qvel=zero_nv,
+      act=zero_na,
+      qacc_warmstart=zero_nv,
+      ctrl=zero_nu,
+      qfrc_applied=zero_nv,
+      xfrc_applied=zero_nbody_6,
+      eq_active=jp.zeros(m.neq, dtype=jp.int32),
+      qacc=zero_nv,
+      act_dot=zero_na,
+      xpos=zero_nbody_3,
+      xquat=jp.zeros((m.nbody, 4), dtype=jp.float32),
+      xmat=zero_nbody_3_3,
+      xipos=zero_nbody_3,
+      ximat=zero_nbody_3_3,
+      xanchor=zero_njnt_3,
+      xaxis=zero_njnt_3,
+      geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32),
+      geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32),
+      site_xpos=jp.zeros((m.nsite, 3), dtype=jp.float32),
+      site_xmat=jp.zeros((m.nsite, 3, 3), dtype=jp.float32),
+      subtree_com=zero_nbody_3,
+      cdof=zero_nv_6,
+      cinert=zero_nbody_10,
+      actuator_length=zero_nu,
+      actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
+      crb=zero_nbody_10,
+      qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
+      qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
+      qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
+      contact=types.Contact.zero(ncon),
+      efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
+      efc_frictionloss=zero_nefc,
+      efc_D=zero_nefc,
+      actuator_velocity=zero_nu,
+      cvel=zero_nbody_6,
+      cdof_dot=zero_nv_6,
+      qfrc_bias=zero_nv,
+      qfrc_passive=zero_nv,
+      efc_aref=zero_nefc,
+      qfrc_actuator=zero_nv,
+      qfrc_smooth=zero_nv,
+      qacc_smooth=zero_nv,
+      qfrc_constraint=zero_nv,
+      qfrc_inverse=zero_nv,
+      efc_force=zero_nefc,
+  )
+
+  return d
+
+
+def _get_contact(
+    c: mujoco._structs._MjContactList,
+    cx: types.Contact,
+    efc_start: int,
+):
+  """Converts mjx.Contact to mujoco._structs._MjContactList."""
+  con_id = np.nonzero(cx.dist <= 0)[0]
+  for field in types.Contact.fields():
+    value = getattr(cx, field.name)[con_id]
+    if field.name == 'frame':
+      value = value.reshape((-1, 9))
+    getattr(c, field.name)[:] = value
+
+  ncon = cx.dist.shape[0]
+  c.efc_address[:] = np.arange(efc_start, efc_start + ncon * 4, 4)[con_id]
+
+
+def get_data(
+    m: mujoco.MjModel, d: types.Data
+) -> Union[mujoco.MjData, List[mujoco.MjData]]:
+  """Gets mjx.Data from a device, resulting in mujoco.MjData or List[MjData]."""
+  batched = len(d.qpos.shape) > 1
+  batch_size = d.qpos.shape[0] if batched else 1
+
+  if batched:
+    result = [mujoco.MjData(m) for _ in range(batch_size)]
+  else:
+    result = mujoco.MjData(m)
+
+  get_data_into(result, m, d)
+
+  return result
+
+
+def get_data_into(
+    result: Union[mujoco.MjData, List[mujoco.MjData]],
+    m: mujoco.MjModel,
+    d: types.Data,
+):
+  """Gets mjx.Data from a device into an existing mujoco.MjData or list."""
+  batched = isinstance(result, list)
+  if batched and len(d.qpos.shape) < 2:
+    raise ValueError('dst is a list, but d is not batched.')
+  if not batched and len(d.qpos.shape) >= 2:
+    raise ValueError('dst is a an MjData, but d is batched.')
+
+  d = jax.device_get(d)
+
+  batch_size = d.qpos.shape[0] if batched else 1
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  efc_type = np.array([
+      mujoco.mjtConstraint.mjCNSTR_EQUALITY,
+      mujoco.mjtConstraint.mjCNSTR_FRICTION_DOF,
+      mujoco.mjtConstraint.mjCNSTR_LIMIT_JOINT,
+      mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
+  ]).repeat([ne, nf, nl, nc])
+
+  dof_i, dof_j = [], []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      dof_i.append(i)
+      dof_j.append(j)
+      j = m.dof_parentid[j]
+
+  for i in range(batch_size):
+    d_i = jax.tree_map(lambda x, i=i: x[i], d) if batched else d
+    result_i = result[i] if batched else result
+    ncon = (d_i.contact.dist <= 0).sum()
+    efc_active = (d_i.efc_J != 0).any(axis=1)
+    efc_con = efc_type == mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL
+    nefc, nc = efc_active.sum(), (efc_active & efc_con).sum()
+    result_i.nnzJ = nefc * m.nv
+    mujoco._functions._realloc_con_efc(result_i, ncon=ncon, nefc=nefc)  # pylint: disable=protected-access
+    result_i.efc_J_rownnz[:] = np.repeat(m.nv, nefc)
+    result_i.efc_J_rowadr[:] = np.arange(0, nefc * m.nv, m.nv)
+    result_i.efc_J_colind[:] = np.tile(np.arange(m.nv), nefc)
+
+    for field in types.Data.fields():
+      if field.name == 'contact':
+        _get_contact(result_i.contact, d_i.contact, nefc - nc)
+        continue
+
+      value = getattr(d_i, field.name)
+
+      if field.name in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+        value = value.reshape((-1, 9))
+
+      if field.name in ('efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+        value = value[efc_active]
+
+      if field.name == 'efc_J':
+        value = value[efc_active].reshape(-1)
+
+      if field.name == 'qM' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLD' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLDiagInv' and not support.is_sparse(m):
+        value = np.ones(m.nv)
+
+      if value.shape:
+        getattr(result_i, field.name)[:] = value
+      else:
+        setattr(result_i, field.name, value)
+
+    result_i.efc_type[:] = efc_type[efc_active]
+
+
+def _put_contact(
+    c: mujoco._structs._MjContactList, ncon: int, device=None
+) -> types.Contact:
+  """Puts mujoco.structs._MjContactList onto a device, resulting in mjx.Contact."""
+  fields = {
+      f.name: copy.copy(getattr(c, f.name)) for f in types.Contact.fields()
+  }
+  fields['frame'] = fields['frame'].reshape((-1, 3, 3))
+  pad_size = ncon - c.dist.shape[0]
+  pad_fn = lambda x: np.concatenate(
+      (x, np.zeros((pad_size,) + x.shape[1:], dtype=x.dtype))
+  )
+  fields = jax.tree_map(pad_fn, fields)
+  fields['dist'][-pad_size:] = np.inf
+  fields = jax.device_put(fields, device=device)
+
+  return types.Contact(**fields)
+
+
+def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
+  """Puts mujoco.MjData onto a device, resulting in mjx.Data."""
+  ncon = collision_driver.ncon(m)
+  ne, nf, nl, nc = constraint.count_constraints(m)
+  nefc = ne + nf + nl + nc
+
+  for d_val, val, name in (
+      (d.ncon, ncon, 'ncon'),
+      (d.ne, ne, 'ne'),
+      (d.nf, nf, 'nf'),
+      (d.nl, nl, 'nl'),
+      (d.nefc, nefc, 'nefc'),
+  ):
+    if d_val > val:
+      raise ValueError(f'd.{name} too high, d.{name} = {d_val}, model = {val}')
+
+  fields = {
+      f.name: copy.copy(getattr(d, f.name))  # copy because device_put is async
+      for f in types.Data.fields()
+      if f.type is jax.Array
+  }
+
+  for fname in ('xmat', 'ximat', 'geom_xmat', 'site_xmat'):
+    fields[fname] = fields[fname].reshape((-1, 3, 3))
+
+  # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
+  # efc_J is also optionally column-sparse (typically for large nv).  MJX is
+  # neither: it contains zeros for inactive constraints, and efc_J is always
+  # (nefc, nv).  this may change in the future.
+  if mujoco.mj_isSparse(m):
+    nr = d.efc_J_rownnz.shape[0]
+    efc_j = np.zeros((nr, m.nv))
+    for i in range(nr):
+      rowadr = d.efc_J_rowadr[i]
+      for j in range(d.efc_J_rownnz[i]):
+        efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
+    fields['efc_J'] = efc_j
+  else:
+    fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
+
+  for fname in ('efc_J', 'efc_frictionloss', 'efc_D', 'efc_aref', 'efc_force'):
+    value = np.zeros((nefc, m.nv)) if fname == 'efc_J' else np.zeros(nefc)
+    for i in range(4):
+      value_beg = sum([ne, nf, nl][:i])
+      d_beg = sum([d.ne, d.nf, d.nl][:i])
+      size = [d.ne, d.nf, d.nl, d.nefc - d.nl - d.nf - d.ne][i]
+      value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
+    fields[fname] = value
+
+  # convert qM and qLD if jacobian is dense
+  if not support.is_sparse(m):
+    fields['qM'] = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, fields['qM'], d.qM)
+    # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
+    try:
+      fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM'])
+    except scipy.linalg.LinAlgError:
+      # this happens when qM is empty or unstable simulation
+      fields['qLD'] = np.zeros((m.nv, m.nv))
+    fields['qLDiagInv'] = np.zeros(0)
+
+  fields = jax.device_put(fields, device=device)
+  fields['contact'] = _put_contact(d.contact, ncon, device=device)
+
+  return types.Data(**fields)

mjx/mujoco/mjx/_src/io_test.py

@@ -0,0 +1,438 @@
+# 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 jax import numpy as jp
+import mujoco
+from mujoco import mjx
+import numpy as np
+
+
+_MULTIPLE_CONVEX_OBJECTS = """
+  <mujoco>
+    <option timestep="0.001" jacobian="dense"/>
+    <default>
+      <geom solref=".006 1"/>
+    </default>
+    <asset>
+      <mesh name="box" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 1 -1 -1 1" scale="1 1 .1"/>
+      <mesh name="boxoid" vertex="-1 -1 -1 1 -1 -1 1 1 -1 1 1 1 1 -1 1 -1 1 -1 -1 1 .5 -1 -1 2" scale=".3 .2 .1"/>
+      <mesh name="pentaprism" vertex="1 0 0 0.309 0.951 0 -0.809 0.588 0 -0.809 -0.588 0 0.309 -0.951 0 1 0 1 0.309 0.951 1 -0.809 0.588 1 -0.809 -0.588 1 0.309 -0.951 1" scale=".2 .2 .1"/>
+    </asset>
+    <worldbody>
+      <geom type="plane" pos="0 0 -.5" size="3 3 .01"/>
+      <geom type="mesh" mesh="box" pos="0 0 -.15" euler="3 7 30"/>
+      <body pos="-.3 -.3 .3">
+        <freejoint/>
+        <geom type="mesh" mesh="boxoid" rgba=".8 0 0 1" euler="3 5 -130"/>
+      </body>
+      <body pos=".3 .3 0.3">
+        <freejoint/>
+        <geom type="box" euler="3 5 -80" size=".3 .2 .1" rgba="0 .8 0 1"/>
+      </body>
+      <body pos=".3 .3 .6">
+        <freejoint/>
+        <geom type="mesh" mesh="pentaprism" rgba="0 0 .8 1"/>
+      </body>
+      <body pos=".6 -.3 .3">
+        <joint name="joint" axis="1 0 0" type="hinge" range="-45 45"/>
+        <geom type="capsule" size=".2 .05" rgba=".6 0 .6 1"/>
+      </body>
+    </worldbody>
+    <actuator>
+      <motor joint="joint"/>
+    </actuator>
+  </mujoco>
+"""
+
+_MULTIPLE_CONSTRAINTS = """
+  <mujoco>
+    <worldbody>
+      <geom type="plane" size="3 3 .01"/>
+      <body name="cap1" pos="-.3 -.3 .2">
+        <freejoint/>
+        <geom type="capsule" size=".2 .05"/>
+        <body name="cap2" pos=".6 -.3 .3">
+          <joint axis="0 1 0" type="hinge" range="-45 45"/>
+          <joint axis="1 0 0" type="hinge" range="-0.001 0.001"/>
+          <geom type="capsule" size=".2 .05"/>
+          <site pos="-0.214 -0.078 0" quat="0.664 0.664 -0.242 -0.242"/>
+        </body>
+      </body>
+    </worldbody>
+    <equality>
+      <connect body1="cap2" anchor="0 0 1"/>
+    </equality>
+  </mujoco>
+"""
+
+
+class ModelIOTest(parameterized.TestCase):
+  """IO tests for mjx.Model."""
+
+  def test_put_model(self):
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
+    mx = mjx.put_model(m)
+    self.assertEqual(mx.nq, m.nq)
+    self.assertEqual(mx.nv, m.nv)
+    self.assertEqual(mx.nu, m.nu)
+    self.assertEqual(mx.na, m.na)
+    self.assertEqual(mx.nbody, m.nbody)
+    self.assertEqual(mx.njnt, m.njnt)
+    self.assertEqual(mx.ngeom, m.ngeom)
+    self.assertEqual(mx.nmesh, m.nmesh)
+    self.assertEqual(mx.npair, m.npair)
+    self.assertEqual(mx.nexclude, m.nexclude)
+    self.assertEqual(mx.neq, m.neq)
+    self.assertEqual(mx.nnumeric, m.nnumeric)
+    self.assertEqual(mx.nM, m.nM)
+    self.assertAlmostEqual(mx.opt.timestep, m.opt.timestep)
+
+    np.testing.assert_allclose(mx.body_parentid, m.body_parentid)
+    np.testing.assert_allclose(mx.geom_type, m.geom_type)
+    np.testing.assert_allclose(mx.geom_bodyid, m.geom_bodyid)
+    np.testing.assert_almost_equal(mx.geom_solref, m.geom_solref)
+    np.testing.assert_almost_equal(mx.geom_pos, m.geom_pos)
+    self.assertLen(mx.geom_convex_face, 6)
+    self.assertLen(mx.geom_convex_vert, 6)
+    self.assertLen(mx.geom_convex_edge, 6)
+    self.assertLen(mx.geom_convex_facenormal, 6)
+
+    np.testing.assert_allclose(mx.jnt_type, m.jnt_type)
+    np.testing.assert_allclose(mx.jnt_dofadr, m.jnt_dofadr)
+    np.testing.assert_allclose(mx.jnt_bodyid, m.jnt_bodyid)
+    np.testing.assert_allclose(mx.jnt_limited, m.jnt_limited)
+    np.testing.assert_almost_equal(mx.jnt_axis, m.jnt_axis)
+
+    np.testing.assert_allclose(mx.actuator_trntype, m.actuator_trntype)
+    np.testing.assert_allclose(mx.actuator_dyntype, m.actuator_dyntype)
+    np.testing.assert_allclose(mx.actuator_gaintype, m.actuator_gaintype)
+    np.testing.assert_allclose(mx.actuator_biastype, m.actuator_biastype)
+    np.testing.assert_allclose(mx.actuator_trnid, m.actuator_trnid)
+
+  def test_fluid_params(self):
+    """Test that has_fluid_params is set when fluid params are present."""
+    m = mjx.put_model(
+        mujoco.MjModel.from_xml_string(
+            '<mujoco><option viscosity="3.0"/><worldbody/></mujoco>'
+        )
+    )
+    self.assertTrue(m.opt.has_fluid_params)
+
+  def test_implicit_not_implemented(self):
+    """Test that MJX guards against models with unimplemented features."""
+
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_cone_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_pgs_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(
+          mujoco.MjModel.from_xml_string(
+              '<mujoco><option solver="PGS"/><worldbody/></mujoco>'
+          )
+      )
+
+  def test_tendon_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint name="left_hip" type="hinge"/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+          <tendon>
+            <fixed>
+              <joint coef="1" joint="left_hip"/>
+            </fixed>
+          </tendon>
+        </mujoco>"""))
+
+  def test_condim_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05" condim="1"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+  def test_cylinder_not_implemented(self):
+    with self.assertRaises(NotImplementedError):
+      mjx.put_model(mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <body>
+              <freejoint/>
+              <geom type="cylinder" size="0.05 0.05"/>
+            </body>
+            <body>
+              <freejoint/>
+              <geom size="0.05"/>
+            </body>
+          </worldbody>
+        </mujoco>"""))
+
+
+class DataIOTest(parameterized.TestCase):
+  """IO tests for mjx.Data."""
+
+  def test_make_data(self):
+    """Test that make_data returns the correct shapes."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONVEX_OBJECTS)
+    d = mjx.make_data(m)
+
+    nq = 22
+    nbody = 5
+    ncon = 46
+    nv = 19
+    nefc = 185
+
+    self.assertEqual(d.qpos.shape, (nq,))
+    self.assertEqual(d.qvel.shape, (nv,))
+    self.assertEqual(d.act.shape, (0,))
+    self.assertEqual(d.qacc_warmstart.shape, (nv,))
+    self.assertEqual(d.ctrl.shape, (1,))
+    self.assertEqual(d.qfrc_applied.shape, (nv,))
+    self.assertEqual(d.xfrc_applied.shape, (nbody, 6))
+    self.assertEqual(d.eq_active.shape, (0,))
+    self.assertEqual(d.qacc.shape, (nv,))
+    self.assertEqual(d.act_dot.shape, (0,))
+    self.assertEqual(d.xpos.shape, (nbody, 3))
+    self.assertEqual(d.xquat.shape, (nbody, 4))
+    self.assertEqual(d.xmat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xipos.shape, (nbody, 3))
+    self.assertEqual(d.ximat.shape, (nbody, 3, 3))
+    self.assertEqual(d.xanchor.shape, (4, 3))
+    self.assertEqual(d.xaxis.shape, (4, 3))
+    self.assertEqual(d.geom_xpos.shape, (6, 3))
+    self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
+    self.assertEqual(d.subtree_com.shape, (nbody, 3))
+    self.assertEqual(d.cdof.shape, (nv, 6))
+    self.assertEqual(d.cinert.shape, (nbody, 10))
+    self.assertEqual(d.crb.shape, (nbody, 10))
+    self.assertEqual(d.actuator_length.shape, (1,))
+    self.assertEqual(d.actuator_moment.shape, (1, nv))
+    self.assertEqual(d.qM.shape, (nv, nv))
+    self.assertEqual(d.qLD.shape, (nv, nv))
+    self.assertEqual(d.qLDiagInv.shape, (0,))
+    self.assertEqual(d.contact.dist.shape, (ncon,))
+    self.assertEqual(d.contact.pos.shape, (ncon, 3))
+    self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
+    self.assertEqual(d.contact.solref.shape, (ncon, 2))
+    self.assertEqual(d.contact.solimp.shape, (ncon, 5))
+    self.assertEqual(d.contact.geom1.shape, (ncon,))
+    self.assertEqual(d.contact.geom2.shape, (ncon,))
+    self.assertEqual(d.efc_J.shape, (nefc, nv))
+    self.assertEqual(d.efc_frictionloss.shape, (nefc,))
+    self.assertEqual(d.efc_D.shape, (nefc,))
+    self.assertEqual(d.actuator_velocity.shape, (1,))
+    self.assertEqual(d.cvel.shape, (nbody, 6))
+    self.assertEqual(d.cdof_dot.shape, (nv, 6))
+    self.assertEqual(d.qfrc_bias.shape, (nv,))
+    self.assertEqual(d.qfrc_passive.shape, (nv,))
+    self.assertEqual(d.efc_aref.shape, (nefc,))
+    self.assertEqual(d.qfrc_actuator.shape, (nv,))
+    self.assertEqual(d.qfrc_smooth.shape, (nv,))
+    self.assertEqual(d.qacc_smooth.shape, (nv,))
+    self.assertEqual(d.qfrc_constraint.shape, (nv,))
+    self.assertEqual(d.qfrc_inverse.shape, (nv,))
+    self.assertEqual(d.efc_force.shape, (nefc,))
+
+  def test_put_data(self):
+    """Test that put_data puts the correct data for dense and sparse."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+
+    # check a few fields
+    np.testing.assert_allclose(dx.qpos, d.qpos)
+    np.testing.assert_allclose(dx.xpos, d.xpos)
+    np.testing.assert_allclose(dx.cvel, d.cvel)
+    np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
+
+    # check that qM is transformed properly
+    qm = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
+
+    # 4 contacts, 2 for each capsule against the plane
+    self.assertEqual(dx.contact.dist.shape, (4,))
+    self.assertEqual(d.ncon, 1)  # however only 1 contact in this step
+    np.testing.assert_allclose(dx.contact.dist[0], d.contact.dist[0])
+    self.assertTrue(np.isinf(dx.contact.dist[1:]).all())
+    self.assertEqual(dx.contact.frame.shape, (4, 3, 3))
+    np.testing.assert_allclose(
+        dx.contact.frame[0].reshape(9), d.contact.frame[0]
+    )
+    np.testing.assert_allclose(dx.contact.frame[1:], 0)
+
+    # xmat, ximat, geom_xmat are all shape transformed
+    self.assertEqual(dx.xmat.shape, (3, 3, 3))
+    self.assertEqual(dx.ximat.shape, (3, 3, 3))
+    self.assertEqual(dx.geom_xmat.shape, (3, 3, 3))
+    self.assertEqual(dx.site_xmat.shape, (1, 3, 3))
+    np.testing.assert_allclose(dx.xmat.reshape((3, 9)), d.xmat)
+    np.testing.assert_allclose(dx.ximat.reshape((3, 9)), d.ximat)
+    np.testing.assert_allclose(dx.geom_xmat.reshape((3, 9)), d.geom_xmat)
+    np.testing.assert_allclose(dx.site_xmat.reshape((1, 9)), d.site_xmat)
+
+    # efc_ are also shape transformed and padded
+    self.assertEqual(dx.efc_J.shape, (21, 8))  # nefc, nv
+    d_efc_j = d.efc_J.reshape((-1, 8))
+    np.testing.assert_allclose(dx.efc_J[:3], d_efc_j[:3])  # connect eq
+    np.testing.assert_allclose(dx.efc_J[3], d_efc_j[3])  # one active limit
+    np.testing.assert_allclose(dx.efc_J[4], 0)  # one inactive limit
+    np.testing.assert_allclose(dx.efc_J[5:9], d_efc_j[4:8])  # contact
+    np.testing.assert_allclose(dx.efc_J[9:], 0)  # no contact
+
+    # check another efc_ too
+    self.assertEqual(dx.efc_aref.shape, (21,))  # nefc
+    np.testing.assert_allclose(dx.efc_aref[:3], d.efc_aref[:3])
+    np.testing.assert_allclose(dx.efc_aref[3], d.efc_aref[3])
+    np.testing.assert_allclose(dx.efc_aref[4], 0)
+    np.testing.assert_allclose(dx.efc_aref[5:9], d.efc_aref[4:8])
+    np.testing.assert_allclose(dx.efc_aref[9:], 0)
+
+    # check sparse transform is correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_sparse = mjx.put_data(m, d)
+    np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
+
+    # check sparse mass matrices are correct
+    np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
+
+    # check dense mass matrices are correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_from_dense = mjx.put_data(m, d)
+    qm = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
+
+  def test_get_data(self):
+    """Test that get_data makes correct MjData."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2: mujoco.MjData = mjx.get_data(m, dx)
+
+    # check a few fields
+    np.testing.assert_allclose(d_2.qpos, d.qpos)
+    np.testing.assert_allclose(d_2.xpos, d.xpos)
+    np.testing.assert_allclose(d_2.cvel, d.cvel)
+    np.testing.assert_allclose(d_2.cdof_dot, d.cdof_dot)
+    np.testing.assert_allclose(d_2.qM, d.qM)
+
+    # only 1 contact active
+    self.assertEqual(d_2.contact.dist.shape, (1,))
+    self.assertEqual(d_2.ncon, 1)
+    np.testing.assert_allclose(d_2.contact.dist, d.contact.dist)
+    self.assertEqual(d_2.contact.frame.shape, (1, 9))
+    np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
+
+    # xmat, ximat, geom_xmat, site_xmat are all shape transformed
+    self.assertEqual(d_2.xmat.shape, (3, 9))
+    self.assertEqual(d_2.ximat.shape, (3, 9))
+    self.assertEqual(d_2.geom_xmat.shape, (3, 9))
+    self.assertEqual(d_2.site_xmat.shape, (1, 9))
+    np.testing.assert_allclose(d_2.xmat, d.xmat)
+    np.testing.assert_allclose(d_2.ximat, d.ximat)
+    np.testing.assert_allclose(d_2.geom_xmat, d.geom_xmat)
+    np.testing.assert_allclose(d_2.site_xmat, d.site_xmat)
+
+    # efc_* are also shape transformed and filtered
+    self.assertEqual(d_2.efc_J.shape, (64,))  # nefc * nv
+    np.testing.assert_allclose(d_2.efc_J, d.efc_J)
+    self.assertEqual(d_2.efc_aref.shape, (8,))  # nefc
+    np.testing.assert_allclose(d_2.efc_aref, d.efc_aref)
+
+    # efc_address is created on demand
+    np.testing.assert_allclose(d_2.contact.efc_address, d.contact.efc_address)
+
+  def test_get_data_batched(self):
+    """Test that get_data makes correct List[MjData] for batched Data."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    # second data in batch has contact dist > 0, disables contact
+    dx_b = jax.tree_map(lambda x: jp.stack((x, x + 0.05)), dx)
+    ds = mjx.get_data(m, dx_b)
+    self.assertLen(ds, 2)
+    np.testing.assert_allclose(ds[0].qpos, d.qpos)
+    np.testing.assert_allclose(ds[1].qpos, d.qpos + 0.05, atol=1e-8)
+    self.assertEqual(ds[0].ncon, 1)
+    self.assertEqual(ds[1].ncon, 0)
+
+  def test_get_data_into(self):
+    """Test that get_data_into correctly populates an MjData."""
+
+    m = mujoco.MjModel.from_xml_string(_MULTIPLE_CONSTRAINTS)
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx = mjx.put_data(m, d)
+    d_2 = mujoco.MjData(m)
+    mjx.get_data_into(d_2, m, dx)
+
+    # check a few fields
+    np.testing.assert_allclose(d_2.qpos, d.qpos)
+    np.testing.assert_allclose(d_2.xpos, d.xpos)
+    np.testing.assert_allclose(d_2.qM, d.qM)
+
+    # only 1 contact active
+    self.assertEqual(d_2.contact.dist.shape, (1,))
+    self.assertEqual(d_2.ncon, 1)
+    np.testing.assert_allclose(d_2.contact.dist, d.contact.dist)
+    self.assertEqual(d_2.contact.frame.shape, (1, 9))
+    np.testing.assert_allclose(d_2.contact.frame, d.contact.frame)
+
+if __name__ == '__main__':
+  absltest.main()

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

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()

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

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()

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.replace(qfrc_passive=jp.zeros(m.nv))
+
+  # 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

mjx/mujoco/mjx/_src/passive_test.py

@@ -0,0 +1,79 @@
+# 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."""
+
+from absl.testing import absltest
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX passive calculations - mostly
+# due to float precision
+_TOLERANCE = 1e-7
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class PassiveTest(absltest.TestCase):
+
+  def test_passive(self):
+    m = test_util.load_test_file('pendula.xml')
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.ctrl = np.array([0.1, -0.1, 0.2, 0.3, -0.4])
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
+
+    # test with fluid forces
+    m.opt.density = 0.01
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
+
+    m.opt.viscosity = 0.02
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
+
+    m.opt.wind = np.array([0.03, 0.04, 0.05])
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_passive')
+
+    # test disable passive
+    mx = mx.tree_replace({'opt.disableflags': mjx.DisableBit.PASSIVE})
+    dx = jax.jit(mjx.passive)(mx, mjx.put_data(m, d))
+    np.testing.assert_allclose(dx.qfrc_passive, 0)
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/ray.py

@@ -0,0 +1,279 @@
+# 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 for ray interesection testing."""
+
+from typing import Sequence, Tuple
+
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco.mjx._src import math
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import GeomType
+from mujoco.mjx._src.types import Model
+# pylint: enable=g-importing-member
+import numpy as np
+
+
+def _ray_quad(
+    a: jax.Array, b: jax.Array, c: jax.Array
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns two solutions for quadratic: a*x^2 + 2*b*x + c = 0."""
+  det = b * b - a * c
+  det_2 = jp.sqrt(det)
+
+  x0, x1 = (-b - det_2) / a, (-b + det_2) / a
+  x0 = jp.where((det < mujoco.mjMINVAL) | (x0 < 0), jp.inf, x0)
+  x1 = jp.where((det < mujoco.mjMINVAL) | (x1 < 0), jp.inf, x1)
+
+  return x0, x1
+
+
+def _ray_plane(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a plane."""
+  x = -pnt[2] / vec[2]
+
+  valid = vec[2] <= -mujoco.mjMINVAL  # z-vec pointing towards front face
+  valid &= x >= 0
+  # only within rendered rectangle
+  p = pnt[0:2] + x * vec[0:2]
+  valid &= jp.all((size[0:2] <= 0) | (jp.abs(p) <= size[0:2]))
+
+  return jp.where(valid, x, jp.inf)
+
+
+def _ray_sphere(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a sphere."""
+  x0, x1 = _ray_quad(vec @ vec, vec @ pnt, pnt @ pnt - size[0] * size[0])
+  x = jp.where(jp.isinf(x0), x1, x0)
+
+  return x
+
+
+def _ray_capsule(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a capsule."""
+
+  # cylinder round side: (x*lvec+lpnt)'*(x*lvec+lpnt) = size[0]*size[0]
+  a = vec[0:2] @ vec[0:2]
+  b = vec[0:2] @ pnt[0:2]
+  c = pnt[0:2] @ pnt[0:2] - size[0] * size[0]
+
+  # solve a*x^2 + 2*b*x + c = 0
+  x0, x1 = _ray_quad(a, b, c)
+  x = jp.where(jp.isinf(x0), x1, x0)
+
+  # make sure round solution is between flat sides
+  x = jp.where(jp.abs(pnt[2] + x * vec[2]) <= size[1], x, jp.inf)
+
+  # top cap
+  dif = pnt - jp.array([0, 0, size[1]])
+  x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
+  # accept only top half of sphere
+  x = jp.where((pnt[2] + x0 * vec[2] >= size[1]) & (x0 < x), x0, x)
+  x = jp.where((pnt[2] + x1 * vec[2] >= size[1]) & (x1 < x), x1, x)
+
+  # bottom cap
+  dif = pnt + jp.array([0, 0, size[1]])
+  x0, x1 = _ray_quad(vec @ vec, vec @ dif, dif @ dif - size[0] * size[0])
+
+  # accept only bottom half of sphere
+  x = jp.where((pnt[2] + x0 * vec[2] <= -size[1]) & (x0 < x), x0, x)
+  x = jp.where((pnt[2] + x1 * vec[2] <= -size[1]) & (x1 < x), x1, x)
+
+  return x
+
+
+def _ray_box(
+    size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a box."""
+
+  iface = jp.array([(1, 2), (0, 2), (0, 1), (1, 2), (0, 2), (0, 1)])
+
+  # side +1, -1
+  # solution of pnt[i] + x * vec[i] = side * size[i]
+  x = jp.concatenate([(size - pnt) / vec, (-size - pnt) / vec])
+
+  # intersection with face
+  p0 = pnt[iface[:, 0]] + x * vec[iface[:, 0]]
+  p1 = pnt[iface[:, 1]] + x * vec[iface[:, 1]]
+  valid = jp.abs(p0) <= size[iface[:, 0]]
+  valid &= jp.abs(p1) <= size[iface[:, 1]]
+
+  return jp.min(jp.where(valid, x, jp.inf))
+
+
+def _ray_triangle(
+    vert: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+    b0: jax.Array,
+    b1: jax.Array,
+) -> jax.Array:
+  """Returns the distance at which a ray intersects with a triangle."""
+  # project difference vectors in ray normal plane
+  planar = jp.dot(jp.array([b0, b1]), (vert - pnt).T)
+
+  # determine if origin is inside planar projection of triangle
+  # A = (p0-p2, p1-p2), b = -p2, solve A*t = b
+  A = jp.array(  # pylint: disable=invalid-name
+      [planar[:, 0] - planar[:, 2], planar[:, 1] - planar[:, 2]]
+  ).T.flatten()
+  b = -planar[:, 2]
+  det = A[0] * A[3] - A[1] * A[2]
+  valid = jp.abs(det) >= mujoco.mjMINVAL
+
+  t0 = (A[3] * b[0] - A[1] * b[1]) / det
+  t1 = (-A[2] * b[0] + A[0] * b[1]) / det
+  valid &= (t0 >= 0) & (t1 >= 0) & (t0 + t1 <= 1)
+
+  # intersect ray with plane of triangle
+  nrm = jp.cross(vert[0] - vert[2], vert[1] - vert[2])
+  denom = jp.dot(vec, nrm)
+  valid &= jp.abs(denom) >= mujoco.mjMINVAL
+
+  dist = jp.where(valid, -jp.dot(pnt - vert[2], nrm) / denom, jp.inf)
+
+  return dist
+
+
+def _ray_mesh(
+    m: Model,
+    geom_id: np.ndarray,
+    unused_size: jax.Array,
+    pnt: jax.Array,
+    vec: jax.Array,
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns the best distance and geom_id for ray mesh intersections."""
+  data_id = m.geom_dataid[geom_id]
+
+  ray_basis = lambda x: math.orthogonals(math.normalize(x))
+  b0, b1 = jax.vmap(ray_basis)(vec)
+
+  faceadr = np.append(m.mesh_faceadr, m.nmeshface)
+  vertadr = np.append(m.mesh_vertadr, m.nmeshvert)
+
+  dists = []
+  for i, id_ in enumerate(data_id):
+    face = m.mesh_face[faceadr[id_] : faceadr[id_ + 1]]
+    vert = m.mesh_vert[vertadr[id_] : vertadr[id_ + 1]]
+    dist = jax.vmap(_ray_triangle, in_axes=(0, None, None, None, None))(
+        vert[face], pnt[i], vec[i], b0[i], b1[i]
+    )
+    dists.append(dist)
+
+  # map the triangle id to data id
+  tri_id = np.append(0, (faceadr[data_id + 1] - faceadr[data_id]).cumsum())
+  tri_data_id = np.zeros(tri_id[-1], dtype=np.int32)
+  tri_data_id[tri_id[:-1]] = 1
+  tri_data_id = tri_data_id.cumsum() - 1
+
+  dists = jp.concatenate(dists)
+  min_id = jp.argmin(dists)
+  # Grab the best distance amongst all meshes, bypassing the argmin in `ray`.
+  # This avoids having to compute the best distance per mesh.
+  dist = dists[min_id, None]
+  id_ = jp.array(geom_id)[jp.array(tri_data_id)[min_id], None]
+
+  return dist, id_
+
+
+_RAY_FUNC = {
+    GeomType.PLANE: _ray_plane,
+    GeomType.SPHERE: _ray_sphere,
+    GeomType.CAPSULE: _ray_capsule,
+    GeomType.BOX: _ray_box,
+    GeomType.MESH: _ray_mesh,
+}
+
+
+def ray(
+    m: Model,
+    d: Data,
+    pnt: jax.Array,
+    vec: jax.Array,
+    geomgroup: Sequence[int] = (),
+    flg_static: bool = True,
+    bodyexclude: int = -1,
+) -> Tuple[jax.Array, jax.Array]:
+  """Returns the geom id and distance at which a ray intersects with a geom.
+
+  Args:
+    m: MJX model
+    d: MJX data
+    pnt: ray origin point (3,)
+    vec: ray direction    (3,)
+    geomgroup: group inclusion/exclusion mask, or empty to ignore
+    flg_static: if True, allows rays to intersect with static geoms
+    bodyexclude: ignore geoms on specified body id
+
+  Returns:
+    dist: distance from ray origin to geom surface (or -1.0 for no intersection)
+    id: id of intersected geom (or -1 for no intersection)
+  """
+
+  dists, ids = [], []
+  geom_filter = m.geom_bodyid != bodyexclude
+  geom_filter &= (m.geom_matid != -1) | (m.geom_rgba[:, 3] != 0)
+  geom_filter &= (m.geom_matid == -1) | (m.mat_rgba[m.geom_matid, 3] != 0)
+  geom_filter &= flg_static | (m.body_weldid[m.geom_bodyid] != 0)
+  if geomgroup:
+    geomgroup = np.array(geomgroup, dtype=bool)
+    geom_filter &= geomgroup[np.clip(m.geom_group, 0, mujoco.mjNGROUP)]
+
+  # map ray to local geom frames
+  geom_pnts = jax.vmap(lambda x, y: x.T @ (pnt - y))(d.geom_xmat, d.geom_xpos)
+  geom_vecs = jax.vmap(lambda x: x.T @ vec)(d.geom_xmat)
+
+  for geom_type, fn in _RAY_FUNC.items():
+    id_, = np.nonzero(geom_filter & (m.geom_type == geom_type))
+
+    if id_.size == 0:
+      continue
+
+    args = m.geom_size[id_], geom_pnts[id_], geom_vecs[id_]
+
+    if geom_type == GeomType.MESH:
+      dist, id_ = fn(m, id_, *args)
+    else:
+      dist = jax.vmap(fn)(*args)
+
+    dists, ids = dists + [dist], ids + [id_]
+
+  if not ids:
+    return jp.array(-1), jp.array(-1.0)
+
+  dists = jp.concatenate(dists)
+  ids = jp.concatenate(ids)
+  min_id = jp.argmin(dists)
+  dist = jp.where(jp.isinf(dists[min_id]), -1, dists[min_id])
+  id_ = jp.where(jp.isinf(dists[min_id]), -1, ids[min_id])
+
+  return dist, id_

mjx/mujoco/mjx/_src/ray_test.py

@@ -0,0 +1,257 @@
+# 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 ray functions."""
+
+from absl.testing import absltest
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX ray calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+class RayTest(absltest.TestCase):
+
+  def test_ray_nothing(self):
+    """Tests that MJX ray returns -1 when nothing is hit."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    pnt, vec = jp.array([12.146, 1.865, 3.895]), jp.array([0, 0, -1.0])
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_plane(self):
+    """Tests MJX ray<>plane matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at a slight angle
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 0, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking on wrong side of plane
+    pnt = jp.array([0, 0, -0.5])
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_sphere(self):
+    """Tests MJX ray<>sphere matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at sphere at a slight angle
+    pnt, vec = jp.array([0, 0, 1.6]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 1, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_capsule(self):
+    """Tests MJX ray<>capsule matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at capsule at a slight angle
+    pnt, vec = jp.array([0.5, 1, 1.6]), jp.array([0, 0.05, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking up at capsule from below
+    pnt, vec = jp.array([-0.5, 1, 0.05]), jp.array([0, 0.05, 1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking at cylinder of capsule from the side
+    pnt, vec = jp.array([0, 1, 0.75]), jp.array([1, 0, 0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 2, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_box(self):
+    """Tests MJX ray<>box matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # looking down at box at a slight angle
+    pnt, vec = jp.array([1, 0, 1.6]), jp.array([0, 0.05, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 3, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # looking up at box from below
+    pnt, vec = jp.array([1, 0, 0.05]), jp.array([0, 0.05, 1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 3, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+  def test_ray_mesh(self):
+    """Tests MJX ray<>mesh matches MuJoCo."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    # look at the tetrahedron
+    pnt, vec = jp.array([2.0, 2.0, 2.0]), -jp.array([
+        1.0,
+        1.0,
+        1.0,
+    ])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 4, 'geom_id')
+
+    pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid)
+    _assert_eq(geomid, 4, 'geom_id')
+    _assert_eq(dist, mj_dist, 'dist-tetrahedron')
+
+    # look at the dodecahedron
+    pnt, vec = jp.array([4.0, 2.0, 2.0]), -jp.array([
+        2.0,
+        1.0,
+        1.0,
+    ])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, 5, 'geom_id')
+
+    pnt, vec, geomid = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, geomid)
+    _assert_eq(geomid, 5, 'geom_id')
+    _assert_eq(dist, mj_dist, 'dist-dodecahedron')
+
+  def test_ray_geomgroup(self):
+    """Tests ray geomgroup filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnums=(4,))
+
+    # hits plane with geom_group[0] = 1
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    geomgroup = (1, 0, 0, 0, 0, 0)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup)
+    _assert_eq(geomid, 0, 'geom_id')
+    pnt, vec, unused = np.array(pnt), np.array(vec), np.zeros(1, dtype=np.int32)
+    mj_dist = mujoco.mj_ray(m, d, pnt, vec, None, 1, -1, unused)
+    _assert_eq(dist, mj_dist, 'dist')
+
+    # nothing hit with geom_group[0] = 0
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    geomgroup = (0, 0, 0, 0, 0, 0)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, geomgroup)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_flg_static(self):
+    """Tests ray flg_static filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnames=('flg_static',))
+
+    # nothing hit with flg_static = False
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, flg_static=False)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_bodyexclude(self):
+    """Tests ray bodyexclude filter."""
+    m = test_util.load_test_file('ray.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+    ray_fn = jax.jit(mjx.ray, static_argnames=('bodyexclude',))
+
+    # nothing hit with bodyexclude = 0 (world body)
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = ray_fn(mx, dx, pnt, vec, bodyexclude=0)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+  def test_ray_invisible(self):
+    """Tests ray doesn't hit transparent geoms."""
+    m = test_util.load_test_file('ray.xml')
+    # nothing hit with transparent geoms:
+    m.geom_rgba = 0
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx, dx = mjx.put_model(m), mjx.put_data(m, d)
+
+    pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0])
+    vec /= jp.linalg.norm(vec)
+    dist, geomid = jax.jit(mjx.ray)(mx, dx, pnt, vec)
+    _assert_eq(geomid, -1, 'geom_id')
+    _assert_eq(dist, -1, 'dist')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/scan.py

@@ -0,0 +1,489 @@
+# 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 not x.shape[0]:
+      return x
+    elif (
+        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}')
+
+  # split out numpy and jax args
+  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)]
+
+  # remove empty args that we should not vmap over
+  args = jax.tree_map(lambda a: a if a.shape[0] else None, 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,
+      's': m.nsite,
+  }
+  for idx, (arg, typ) in enumerate(zip(args, in_types)):
+    if len(arg) != size[typ]:
+      raise IndexError(
+          (
+              f'f argument "{idx}" with type "{typ}" has length "{len(arg)}"'
+              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],
+        m.actuator_trnid[ids_u, 1] == -1,  # key by refsite being present
+    )
+
+  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, 0]
+            if m.actuator_trntype[i] == TrnType.JOINT
+            else -1
+        ),
+        's': (
+            m.actuator_trnid[i]
+            if m.actuator_trntype[i] == TrnType.SITE
+            else np.array([-1, -1])
+        ),
+    }
+    v, q = np.array([-1]), np.array([-1])
+    if m.actuator_trntype[i] == TrnType.JOINT:
+      # v/q are associated with the joint transmissions only
+      v = np.nonzero(m.dof_jntid == typ_ids['j'])[0]
+      q = np.nonzero(_q_jointid(m) == typ_ids['j'])[0]
+
+    typ_ids.update({'v': v, 'q': q})
+
+    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)
+
+  # group together bodies that will be processed together.  grouping key:
+  #   1) the tree depth: parent bodies are processed first, so that they are
+  #      available as carry input to child bodies (or reverse if reverse=True)
+  #   2) the types of arguments passed to f, both carry and *args:
+  #      * for 'b' arguments, there is no extra grouping
+  #      * for 'j' arguments, we group by joint type
+  #      * for 'q' arguments, we group by q width
+  #      * for 'v' arguments, we group by dof width
+  depths = np.zeros(m.nbody, dtype=np.int32)
+
+  # map key => body id
+  key_body_ids = {}
+  for body_id in range(m.nbody):
+    parent_id = -1
+    if body_id > 0:
+      parent_id = m.body_parentid[body_id]
+      depths[body_id] = 1 + depths[parent_id]
+
+    # create grouping key: depth, carry, args
+    key = (depths[body_id],)
+
+    for i, t in enumerate(out_types + in_types):
+      id_ = parent_id if i < len(out_types) else body_id
+      if t == 'b':
+        continue
+      elif t == 'j':
+        key += (tuple(m.jnt_type[np.nonzero(m.jnt_bodyid == id_)[0]]))
+      elif t == 'v':
+        key += (len(np.nonzero(m.dof_bodyid == id_)[0]),)
+      elif t == 'q':
+        key += (len(np.nonzero(_q_bodyid(m) == id_)[0]),)
+
+    body_ids = key_body_ids.get(key, np.array([], dtype=np.int32))
+    key_body_ids[key] = np.append(body_ids, body_id)
+
+  # find parent keys of each key.  a key may have multiple parents if the
+  # carry output keys of distinct parents are the same. e.g.:
+  # - depth 0 body 1 (slide joint)
+  # -- depth 1 body 1 (hinge joint)
+  # - depth 0 body 2 (ball joint)
+  # -- depth 1 body 2 (hinge joint)
+  # given a scan with 'j' in the in_types, we would group depth 0 bodies
+  # separately but we may group depth 1 bodies together
+  key_parents = {}
+
+  for key, body_ids in key_body_ids.items():
+    body_ids = body_ids[body_ids != 0]  # ignore worldbody, has no parent
+    if body_ids.size == 0:
+      continue
+    # find any key which has a body id that is a parent of these body_ids
+    pids = m.body_parentid[body_ids]
+    parents = {k for k, v in key_body_ids.items() if np.isin(v, pids).any()}
+    key_parents[key] = list(sorted(parents))
+
+  # key => take indices
+  key_in_take, key_y_take = {}, {}
+  for key, body_ids in key_body_ids.items():
+    for i, typ in enumerate(in_types + out_types):
+      if typ == 'b':
+        ids = body_ids
+      elif typ == 'j':
+        ids = np.stack([np.nonzero(m.jnt_bodyid == b)[0] for b in body_ids])
+      elif typ == 'v':
+        ids = np.stack([np.nonzero(m.dof_bodyid == b)[0] for b in body_ids])
+      elif typ == 'q':
+        ids = np.stack([np.nonzero(_q_bodyid(m) == b)[0] for b in body_ids])
+      else:
+        raise ValueError(f'Unknown in_type: {typ}')
+      if i < len(in_types):
+        key_in_take.setdefault(key, []).append(ids)
+      else:
+        key_y_take.setdefault(key, []).append(np.hstack(ids))
+
+  # use this grouping to take the right data subsets and call vmap(f)
+  keys = sorted(key_body_ids, reverse=reverse)
+  key_y = {}
+  for key in keys:
+    carry = None
+
+    if reverse:
+      child_keys = [k for k, v in key_parents.items() if key in v]
+
+      for child_key in child_keys:
+        y = key_y[child_key]
+        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)
+    elif key in key_parents:
+      ys = [key_y[p] for p in key_parents[key]]
+      y = jax.tree_map(lambda *x: jp.concatenate(x), *ys)
+      body_ids = np.concatenate([key_body_ids[p] for p in key_parents[key]])
+      parent_ids = m.body_parentid[key_body_ids[key]]
+      take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i)
+      carry = jax.tree_map(take_fn, y)
+
+    f_args = [_take(arg, ids) for arg, ids in zip(args, key_in_take[key])]
+    key_y[key] = _nvmap(f, carry, *f_args)
+
+  # slice None results from the final output
+  keys = [k for k in keys if key_y[k] is not None]
+
+  # concatenate ys, drop grouping dimensions, put back in order
+  y = []
+  for i, typ in enumerate(out_types):
+    y_typ = [key_y[key] for key in keys]
+    if len(out_types) > 1:
+      y_typ = [y_[i] for y_ in y_typ]
+    if typ != 'b':
+      y_typ = jax.tree_map(jp.concatenate, y_typ)
+    y_typ = jax.tree_map(lambda *x: jp.concatenate(x), *y_typ)
+    y_take = np.argsort(np.concatenate([key_y_take[key][i] for key in keys]))
+    _check_output(y_typ, y_take, typ, i)
+    y.append(_take(y_typ, y_take))
+
+  y = y[0] if len(out_types) == 1 else y
+
+  return y

mjx/mujoco/mjx/_src/scan_test.py

@@ -0,0 +1,228 @@
+# 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 = """
+      <mujoco>
+        <compiler inertiafromgeom="true"/>
+        <worldbody>
+          <body>
+            <joint type="free"/>
+            <geom size=".15" mass="1" type="sphere"/>
+            <body>
+              <joint axis="1 0 0" pos="1 0 0" type="ball"/>
+              <geom size=".15" mass="2" type="sphere"/>
+            </body>
+            <body>
+              <joint axis="1 0 0" pos="2 0 0" type="hinge"/>
+              <joint axis="0 1 0" pos="3 0 0" type="slide"/>
+              <geom size=".15" mass="3" type="sphere"/>
+            </body>
+          </body>
+        </worldbody>
+      </mujoco>
+    """
+
+  def test_flat_empty(self):
+    """Test scanning over just world body."""
+    m = mujoco.MjModel.from_xml_string("""
+      <mujoco model="world_body">
+        <worldbody/>
+      </mujoco>
+    """)
+    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 = """
+    <mujoco>
+      <option timestep="0.02"/>
+      <compiler autolimits="true"/>
+      <default>
+        <geom contype="0" conaffinity="0"/>
+      </default>
+      <worldbody>
+        <body pos="0 0 -0.5">
+          <joint type="free" name="joint0" range="-37.81 86.15"/>
+          <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
+          <body pos="0 0 -0.5">
+            <joint axis="1 0 0" type="hinge" name="joint1" range="-32.88 5.15" actuatorfrcrange="-0.48 0.72"/>
+            <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
+            <body pos="0 0 -0.5">
+              <joint axis="0 1 0" type="slide" name="joint2" range="-83.31 54.77"/>
+              <joint axis="1 0 0" type="slide" name="joint3" range="-87.84 60.58" actuatorfrcrange="-0.62 0.28"/>
+              <joint axis="0 1 0" type="hinge" name="joint4" range="-49.96 83.03" actuatorfrcrange="-0.01 0.64"/>
+              <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
+              <body pos="0 0 -0.5">
+                <joint axis="1 0 0" type="slide" name="joint5" range="-1.53 87.07" actuatorfrcrange="-0.06 0.26"/>
+                <joint axis="0 1 0" type="slide" name="joint6" range="-63.57 9.95"/>
+                <joint axis="1 0 0" type="hinge" name="joint7" range="-22.50 41.20" actuatorfrcrange="-0.41 0.84"/>
+                <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
+                <body pos="0 0 -0.5">
+                  <joint axis="0 1 0" type="hinge" name="joint8" range="-34.67 72.56"/>
+                  <joint axis="1 0 0" type="hinge" name="joint9" range="-1.81 16.22" actuatorfrcrange="-0.41 0.14"/>
+                  <joint axis="1 0 0" type="ball" name="joint10" range="0.00 6.02" actuatorfrcrange="-0.74 0.12"/>
+                  <geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </worldbody>
+      <actuator>
+        <position joint="joint6" gear="1"/>
+        <intvelocity joint="joint1" kp="2000" actrange="-0.7 2.3"/>
+        <motor joint="joint3" gear="42"/>
+        <velocity joint="joint2" kv="123"/>
+        <position joint="joint0" ctrlrange="-0.9472 0.9472"/>
+        <intvelocity joint="joint7" kp="2000" actrange="-0.7 2.3"/>
+        <general joint="joint5" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
+        <general joint="joint4" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
+      </actuator>
+    </mujoco>
+  """
+
+  def test_scan_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'
+    )
+
+    actuator_trnid = m.actuator_trnid[:, 0]
+    np.testing.assert_array_equal(gear, m.actuator_gear)
+    np.testing.assert_array_equal(jnt_typ, m.jnt_type[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 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 actuator_trnid]
+    )
+    np.testing.assert_array_equal(qadr, expected_qadr)
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/smooth.py

@@ -0,0 +1,494 @@
+# 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
+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
+from mujoco.mjx._src.types import TrnType
+# pylint: enable=g-importing-member
+import numpy as np
+
+
+def kinematics(m: Model, d: Data) -> Data:
+  """Converts position/velocity from generalized coordinates to maximal."""
+
+  def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat):
+    # calculate joint anchors, axes, body pos and quat in global frame
+    # also normalize qpos while we're at it
+
+    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)
+  d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos)
+  d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat)
+
+  if m.ngeom:
+    geom_xpos, geom_xmat = local_to_global(
+        xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat
+    )
+    d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat)
+
+  if m.nsite:
+    site_xpos, site_xmat = local_to_global(
+        xpos[m.site_bodyid], xquat[m.site_bodyid], m.site_pos, m.site_quat
+    )
+    d = d.replace(site_xpos=site_xpos, site_xmat=site_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 < mujoco.mjMINVAL, (3, 1)).T
+  # take maximum to avoid NaN in gradient of jp.where
+  subtree_com = jax.vmap(jp.divide)(pos, jp.maximum(mass, mujoco.mjMINVAL))
+  subtree_com = jp.where(cond, d.xipos, subtree_com)
+  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 * inert) @ ximat.T + h @ h.T * mass
+    # cinert is triu(inert), mass * off, mass
+    inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
+    return jp.concatenate([inert, off * mass, mass[None]])
+
+  root_com = subtree_com[m.body_rootid]
+  offset = d.xipos - root_com
+  cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
+  d = d.replace(cinert=cinert)
+
+  # 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)
+
+  crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
+  crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
+  qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
+  d = d.replace(qM=qm)
+
+  return d
+
+
+def factor_m(m: Model, d: Data) -> Data:
+  """Gets factorizaton of inertia-like matrix M, assumed spd."""
+
+  if not support.is_sparse(m):
+    qh, _ = jax.scipy.linalg.cho_factor(d.qM)
+    d = d.replace(qLD=qh)
+    return d
+
+  # 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 = d.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 ."""
+
+  if not support.is_sparse(m):
+    return jax.scipy.linalg.cho_solve((d.qLD, False), x)
+
+  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 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 _site_dof_mask(m: Model) -> np.ndarray:
+  """Creates a dof mask for site transmissions."""
+  mask = np.ones((m.nu, m.nv))
+  for i in np.nonzero(m.actuator_trnid[:, 1] != -1)[0]:
+    id_, refid = m.actuator_trnid[i]
+    # intialize last dof address for each body
+    b0 = m.body_weldid[m.site_bodyid[id_]]
+    b1 = m.body_weldid[m.site_bodyid[refid]]
+    dofadr0 = m.body_dofadr[b0] + m.body_dofnum[b0] - 1
+    dofadr1 = m.body_dofadr[b1] + m.body_dofnum[b1] - 1
+
+    # find common ancestral dof, if any
+    while dofadr0 != dofadr1:
+      if dofadr0 < dofadr1:
+        dofadr1 = m.dof_parentid[dofadr1]
+      else:
+        dofadr0 = m.dof_parentid[dofadr0]
+      if dofadr0 == -1 or dofadr1 == -1:
+        break
+
+    # if common ancestral dof was found, clear the columns of its parental chain
+    da = dofadr0 if dofadr0 == dofadr1 else -1
+    while da >= 0:
+      mask[i, da] = 0.0
+      da = m.dof_parentid[da]
+
+  return mask
+
+
+def transmission(m: Model, d: Data) -> Data:
+  """Computes actuator/transmission lengths and moments."""
+  # TODO: consider combining transmission calculation into fwd_actuation.
+  if not m.nu:
+    return d
+
+  def fn(
+      trntype,
+      trnid,
+      gear,
+      jnt_typ,
+      m_j,
+      qpos,
+      has_refsite,
+      site_dof_mask,
+      site_xpos,
+      site_xmat,
+      site_quat,
+  ):
+    if trntype == TrnType.JOINT:
+      if jnt_typ == JointType.FREE:
+        length = jp.zeros(1)
+        moment = gear
+        m_j = m_j + jp.arange(6)
+      elif jnt_typ == JointType.BALL:
+        axis, angle = math.quat_to_axis_angle(qpos)
+        length = jp.dot(axis * angle, gear[:3])[None]
+        moment = gear[:3]
+        m_j = m_j + jp.arange(3)
+      elif jnt_typ in (JointType.SLIDE, JointType.HINGE):
+        length = qpos * gear[0]
+        moment = gear[:1]
+        m_j = m_j[None]
+      else:
+        raise RuntimeError(f'unrecognized joint type: {JointType(jnt_typ)}')
+
+      moment = jp.zeros((m.nv,)).at[m_j].set(moment)
+    elif trntype == TrnType.SITE:
+      length = jp.zeros(1)
+      id_, refid = jp.array(m.site_bodyid)[trnid]
+      jacp, jacr = support.jac(m, d, site_xpos[0], id_)
+      frame_xmat = site_xmat[0]
+      if has_refsite:
+        vecp = site_xmat[1].T @ (site_xpos[0] - site_xpos[1])
+        vecr = math.quat_sub(site_quat[0], site_quat[1])
+        length += jp.dot(jp.concatenate([vecp, vecr]), gear)
+        jacrefp, jacrefr = support.jac(m, d, site_xpos[1], refid)
+        jacp, jacr = jacp - jacrefp, jacr - jacrefr
+        frame_xmat = site_xmat[1]
+
+      jac = jp.concatenate((jacp, jacr), axis=1) * site_dof_mask[:, None]
+      wrench = jp.concatenate((frame_xmat @ gear[:3], frame_xmat @ gear[3:]))
+      moment = jac @ wrench
+    else:
+      raise RuntimeError(f'unrecognized trntype: {TrnType(trntype)}')
+
+    return length, moment
+
+  # pre-compute values for site transmissions
+  has_refsite = m.actuator_trnid[:, 1] != -1
+  site_dof_mask = _site_dof_mask(m)
+  site_quat = jax.vmap(math.quat_mul)(m.site_quat, d.xquat[m.site_bodyid])
+
+  length, moment = scan.flat(
+      m,
+      fn,
+      'uuujjquusss',
+      'uu',
+      m.actuator_trntype,
+      jp.array(m.actuator_trnid),
+      m.actuator_gear,
+      m.jnt_type,
+      jp.array(m.jnt_dofadr),
+      d.qpos,
+      has_refsite,
+      jp.array(site_dof_mask),
+      d.site_xpos,
+      d.site_xmat,
+      site_quat,
+      group_by='u',
+  )
+  length = length.reshape((m.nu,))
+  moment = moment.reshape((m.nu, m.nv))
+
+  d = d.replace(actuator_length=length, actuator_moment=moment)
+  return d

mjx/mujoco/mjx/_src/smooth_test.py

@@ -0,0 +1,161 @@
+# 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
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+# tolerance for difference between MuJoCo and MJX smooth calculations - mostly
+# due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name):
+  tol = _TOLERANCE * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class SmoothTest(absltest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # although we already have generous padding of thresholds, it doesn't hurt
+    # to also fix the seed to reduce test flakiness
+    np.random.seed(0)
+
+  def test_smooth(self):
+    """Tests MJX smooth functions match MuJoCo smooth functions."""
+
+    m = test_util.load_test_file('pendula.xml')
+    # # force MJX sparse for testing:
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+
+    # kinematics
+    dx = jax.jit(mjx.kinematics)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'xanchor')
+    _assert_attr_eq(d, dx, 'xaxis')
+    _assert_attr_eq(d, dx, 'xpos')
+    _assert_attr_eq(d, dx, 'xquat')
+    _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat')
+    _assert_attr_eq(d, dx, 'xipos')
+    _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat')
+    _assert_attr_eq(d, dx, 'geom_xpos')
+    _assert_eq(d.geom_xmat.reshape((-1, 3, 3)), dx.geom_xmat, 'geom_xmat')
+    _assert_attr_eq(d, dx, 'site_xpos')
+    _assert_eq(d.site_xmat.reshape((-1, 3, 3)), dx.site_xmat, 'site_xmat')
+    # com_pos
+    dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'subtree_com')
+    _assert_attr_eq(d, dx, 'cinert')
+    _assert_attr_eq(d, dx, 'cdof')
+    # crb
+    dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'crb')
+    _assert_attr_eq(d, dx, 'qM')
+    # factor_m
+    dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qLD')
+    _assert_attr_eq(d, dx, 'qLDiagInv')
+    # com_vel
+    dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'cvel')
+    _assert_attr_eq(d, dx, 'cdof_dot')
+    # rne
+    dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'qfrc_bias')
+    # transmission
+    dx = jax.jit(mjx.transmission)(mx, mjx.put_data(m, d))
+    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_moment')
+
+  def test_disable_gravity(self):
+    m = mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <option>
+            <flag gravity="disable"/>
+          </option>
+          <worldbody>
+            <body>
+              <joint type="free"/>
+              <geom size="0.1"/>
+            </body>
+          </worldbody>
+        </mujoco>
+        """)
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = jax.jit(mjx.rne)(mx, dx)
+    np.testing.assert_allclose(dx.qfrc_bias, 0)
+
+  def test_site_transmission(self):
+    m = mujoco.MjModel.from_xml_string("""
+        <mujoco>
+        <compiler autolimits="true"/>
+        <worldbody>
+          <body>
+            <joint type="free"/>
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+            <site name="site1"/>
+            <body>
+              <joint type="hinge"/>
+              <geom size="0.1" mass="1"/>
+              <site name="site2" pos="0.1 0.2 0.3"/>
+            </body>
+          </body>
+          <body pos="1 0 0">
+            <joint name="slide" type="hinge"/>
+            <geom type="box" size=".05 .05 .05" mass="1"/>
+          </body>
+        </worldbody>
+        <actuator>
+          <position site="site1" kv="0.1" gear="1 2 3 0 0 0"/>
+          <position site="site1" kv="0.2" gear="0 0 0 1 2 3"/>
+          <position site="site2" kv="0.3" gear="0 3 0 0 0 1"/>
+          <position joint="slide" kv="0.05" />
+          <position site="site2" refsite="site1" gear="1 2 3 0.5 0.4 0.6"/>
+        </actuator>
+        </mujoco>
+      """)
+    d = mujoco.MjData(m)
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_transmission(m, d)
+    dx = jax.jit(mjx.transmission)(mx, dx)
+    _assert_attr_eq(d, dx, 'actuator_length')
+    _assert_attr_eq(d, dx, 'actuator_moment')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/solver.py

@@ -0,0 +1,372 @@
+# 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.
+# ==============================================================================
+"""Constraint solvers."""
+
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco.mjx._src import constraint
+from mujoco.mjx._src import math
+from mujoco.mjx._src import smooth
+from mujoco.mjx._src import support
+# 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
+from mujoco.mjx._src.types import SolverType
+# pylint: enable=g-importing-member
+
+
+class _Context(PyTreeNode):
+  """Data updated during each 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,)
+    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 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
+  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) -> '_Context':
+    jaref = d.efc_J @ d.qacc - d.efc_aref
+    # TODO(robotics-team): determine nv at which sparse mul is faster
+    ma = support.mul_m(m, d, d.qacc)
+    nv_0 = jp.zeros(m.nv)
+    ctx = _Context(
+        qacc=d.qacc,
+        qfrc_constraint=d.qfrc_constraint,
+        Jaref=jaref,
+        efc_force=d.efc_force,
+        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 = _update_constraint(m, d, ctx)
+    if grad:
+      ctx = _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,
+      m: Model,
+      ctx: _Context,
+      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 friction constraints
+    ne, nf, *_ = constraint.count_constraints(m)
+    active = ((ctx.Jaref + alpha * jv) < 0).at[:ne + nf].set(True)
+    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] + (quad_total[2] == 0) * mujoco.mjMINVAL
+    return _LSPoint(alpha=alpha, cost=cost, deriv_0=deriv_0, deriv_1=deriv_1)
+
+
+class _LSContext(PyTreeNode):
+  """Data updated during each 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 _update_constraint(m: Model, d: Data, ctx: _Context) -> _Context:
+  """Updates constraint force and resulting cost given latst solver iteration.
+
+  Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c
+
+  Args:
+    m: model defining constraints
+    d: data which contains latest qacc and smooth terms
+    ctx: current solver context
+
+  Returns:
+    context with new constraint force and costs
+  """
+  # TODO(robotics-team): add friction constraints
+
+  # only count active constraints
+  ne, nf, *_ = constraint.count_constraints(m)
+  active = (ctx.Jaref < 0).at[:ne + nf].set(True)
+
+  efc_force = d.efc_D * -ctx.Jaref * active
+  qfrc_constraint = d.efc_J.T @ efc_force
+  gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth)
+  cost = 0.5 * jp.sum(d.efc_D * ctx.Jaref * ctx.Jaref * active) + gauss
+
+  ctx = ctx.replace(
+      qfrc_constraint=qfrc_constraint,
+      gauss=gauss,
+      cost=cost,
+      prev_cost=ctx.cost,
+      efc_force=efc_force,
+  )
+
+  return ctx
+
+
+def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
+  """Updates grad and M / grad given latest solver iteration.
+
+  Corresponds to CGupdateGradient in mujoco/src/engine/engine_solver.c
+
+  Args:
+    m: model defining constraints
+    d: data which contains latest smooth terms
+    ctx: current solver context
+
+  Returns:
+    context with new grad and M / grad
+  Raises:
+    NotImplementedError: for unsupported solver type
+  """
+
+  grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint
+
+  if m.opt.solver == SolverType.CG:
+    mgrad = smooth.solve_m(m, d, grad)
+  elif m.opt.solver == SolverType.NEWTON:
+    ne, nf, *_ = constraint.count_constraints(m)
+    active = (ctx.Jaref < 0).at[:ne + nf].set(True)
+    h = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T)
+    dh = d.replace(qM=h)
+    dh = smooth.factor_m(m, dh)
+    mgrad = smooth.solve_m(m, dh, grad)
+  else:
+    raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
+
+  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 _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
+  """Performs a zoom linesearch to find optimal search step size.
+
+  Args:
+    m: model defining search options and other needed terms
+    d: data with inertia matrix and other needed terms
+    ctx: current solver 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 = support.mul_m(m, d, 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 a: _LSPoint.create(m, ctx, a, 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 solve(m: Model, d: Data) -> Data:
+  """Finds forces that satisfy constraints using conjugate gradient descent."""
+
+  def cond(ctx: _Context) -> 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: _Context) -> _Context:
+    ctx = _linesearch(m, d, ctx)
+    prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad  # pylint: disable=invalid-name
+    ctx = _update_constraint(m, d, ctx)
+    ctx = _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 = _Context.create(m, d.replace(qacc=d.qacc_warmstart), grad=False)
+    smth = _Context.create(m, d.replace(qacc=d.qacc_smooth), grad=False)
+    qacc = jp.where(warm.cost < smth.cost, d.qacc_warmstart, d.qacc_smooth)
+  d = d.replace(qacc=qacc)
+
+  ctx = _Context.create(m, d)
+  if m.opt.iterations == 1:
+    ctx = body(ctx)
+  else:
+    ctx = jax.lax.while_loop(cond, body, ctx)
+
+  d = d.replace(
+      qacc_warmstart=ctx.qacc,
+      qacc=ctx.qacc,
+      qfrc_constraint=ctx.qfrc_constraint,
+      efc_force=ctx.efc_force,
+  )
+
+  return d

mjx/mujoco/mjx/_src/solver_test.py

@@ -0,0 +1,103 @@
+# 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
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+import numpy as np
+
+
+# tolerance for difference between MuJoCo and MJX constraint calculations,
+# mostly due to float precision
+_TOLERANCE = 5e-5
+
+
+def _assert_eq(a, b, name, tol=_TOLERANCE):
+  tol = tol * 10  # avoid test noise
+  err_msg = f'mismatch: {name}'
+  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
+
+
+def _assert_attr_eq(a, b, attr):
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+
+
+class SolverTest(absltest.TestCase):
+
+  def test_newton(self):
+    """Test newton solver."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+
+  def test_cg(self):
+    """Test CG solver."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+
+  def test_no_warmstart(self):
+    """Test no warmstart."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    nnz = dx.efc_J.any(axis=1)
+    # without warmstart, the solution is not as close
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2)
+
+  def test_dense(self):
+    """Test solver works with dense mass matrices."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/_src/support.py

@@ -0,0 +1,188 @@
+# 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 Optional, Tuple, Union
+
+import jax
+from jax import numpy as jp
+import mujoco
+from mujoco.mjx._src import scan
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import JacobianType
+from mujoco.mjx._src.types import Model
+# pylint: enable=g-importing-member
+
+
+def is_sparse(m: Union[mujoco.MjModel, Model]) -> bool:
+  """Return True if this model should create sparse mass matrices.
+
+  Args:
+    m: a MuJoCo or MJX model
+
+  Returns:
+    True if provided model should create sparse mass matrices
+
+  Modern TPUs have specialized hardware for rapidly operating over sparse
+  matrices, whereas GPUs tend to be faster with dense matrices as long as they
+  fit onto the device.  As such, the default behavior in MJX (via
+  ``JacobianType.AUTO``) is sparse if ``nv`` is >= 60 or MJX detects a TPU as
+  the default backend, otherwise dense.
+  """
+  # AUTO is a rough heuristic - you may see better performance for your workload
+  # and compute by explicitly setting jacobian to dense or sparse
+  if m.opt.jacobian == JacobianType.AUTO:
+    return m.nv >= 60 or jax.default_backend() == 'tpu'
+  return m.opt.jacobian == JacobianType.SPARSE
+
+
+def make_m(
+    m: Model, a: jax.Array, b: jax.Array, d: Optional[jax.Array] = None
+) -> jax.Array:
+  """Computes M = a @ b.T + diag(d)."""
+
+  ij = []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      ij.append((i, j))
+      j = m.dof_parentid[j]
+
+  i, j = (jp.array(x) for x in zip(*ij))
+
+  if not is_sparse(m):
+    qm = a @ b.T
+    if d is not None:
+      qm += jp.diag(d)
+    mask = jp.zeros((m.nv, m.nv), dtype=bool).at[(i, j)].set(True)
+    qm = qm * mask
+    qm = qm + jp.tril(qm, -1).T
+    return qm
+
+  a_i = jp.take(a, i, axis=0)
+  b_j = jp.take(b, j, axis=0)
+  qm = jax.vmap(jp.dot)(a_i, b_j)
+
+  # add diagonal
+  if d is not None:
+    qm = qm.at[m.dof_Madr].add(d)
+
+  return qm
+
+
+def full_m(m: Model, d: Data) -> jax.Array:
+  """Reconstitute dense mass matrix from qM."""
+
+  if not is_sparse(m):
+    return d.qM
+
+  ij = []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      ij.append((i, j))
+      j = m.dof_parentid[j]
+
+  i, j = (jp.array(x) for x in zip(*ij))
+
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
+
+  # also set upper triangular
+  mat = mat + jp.tril(mat, -1).T
+
+  return mat
+
+
+def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
+  """Multiply vector by inertia matrix."""
+
+  if not is_sparse(m):
+    return d.qM @ vec
+
+  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 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)

mjx/mujoco/mjx/_src/support_test.py

@@ -0,0 +1,121 @@
+# 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):
+
+  def test_mul_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # first test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    vec = np.random.random(m.nv)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    mj_vec = np.zeros(m.nv)
+    mujoco.mj_mulM(m, d, mj_vec, vec)
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+    # also check dense
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+  def test_full_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # for the model to be sparse to exercise MJX full_M
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_full_m = jax.jit(support.full_m)(mx, dx)
+    mj_full_m = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, mj_full_m, d.qM)
+    np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5)
+
+  @parameterized.parameters('constraints.xml', 'pendula.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.put_model(m)
+    dx = mjx.put_data(m, 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('pendula.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, 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()

mjx/mujoco/mjx/_src/test_util.py

@@ -0,0 +1,369 @@
+# 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, Tuple
+from xml.etree import ElementTree as ET
+
+from etils import epath
+import mujoco
+import numpy as np
+
+_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity']
+_DYN_TYPES = ['none', 'integrator', 'filter', 'filterexact']
+_DYN_PRMS = ['0.189', '2.1']
+_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 = ['12', '1', '0', '0.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 = ['2.1 0.0 3.3 0 2.3 0', '5.0 3.1 0 2.3 0.0 1.1']
+
+
+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 | None = None,
+    site: str | None = None,
+    refsite: str | None = None,
+) -> Dict[str, str]:
+  """Returns attributes for an actuator."""
+  if joint:
+    attr = {'joint': joint}
+  elif site:
+    attr = {'site': site}
+  else:
+    raise ValueError('must provide a joint or site name')
+
+  if refsite:
+    attr['refsite'] = refsite
+
+  attr['gear'] = np.random.choice(_GEARS)
+
+  # set actuator type
+  if actuator_type == 'position':
+    attr['kp'] = np.random.choice(_KP_POS)
+    attr['kv'] = np.random.choice(_KV_VEL)
+  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)
+
+  # set dyntype
+  if actuator_type == 'general':
+    attr['dyntype'] = np.random.choice(_DYN_TYPES)
+    if attr['dyntype'] != 'none':
+      attr['dynprm'] = np.random.choice(_DYN_PRMS)
+
+  # ctrlrange
+  if p(50) and actuator_type != 'intvelocity':
+    lb, ub = -np.random.uniform(), np.random.uniform()
+    attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}'
+
+  # forcerange
+  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}'})
+    ET.SubElement(child, 'site', {'name': f'site{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)
+
+  bodies = list(mjcf.iter('body'))
+  n_bodies = len(bodies)
+
+  # 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 = []
+
+    # joint transmission
+    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))
+
+    # site transmission
+    for i in range(np.random.randint(0, n_bodies)):
+      actuator_type = np.random.choice(_ACTUATOR_TYPES)
+      attr = _make_actuator(actuator_type, site=f'site{i}')
+      actuators.append((actuator_type, attr))
+
+    # site transmission with refsite
+    for i in range(np.random.randint(0, n_bodies)):
+      j = np.random.randint(0, n_bodies)
+      actuator_type = np.random.choice(_ACTUATOR_TYPES)
+      attr = _make_actuator(actuator_type, site=f'site{i}', refsite=f'site{j}')
+      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
+  body_names = [b.get('name') for b in 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

mjx/mujoco/mjx/_src/types.py

@@ -0,0 +1,699 @@
+# 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
+from mujoco.mjx._src.dataclasses import PyTreeNode  # pylint: disable=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 JacobianType(enum.IntEnum):
+  """Type of constraint Jacobian.
+
+  Attributes:
+    DENSE: dense
+    SPARSE: sparse
+    AUTO: sparse if nv>60 and device is TPU, dense otherwise
+  """
+  DENSE = mujoco.mjtJacobian.mjJAC_DENSE
+  SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
+  AUTO = mujoco.mjtJacobian.mjJAC_AUTO
+
+
+class SolverType(enum.IntEnum):
+  """Constraint solver algorithm.
+
+  Attributes:
+    CG: Conjugate gradient (primal)
+  """
+  # unsupported: PGS
+  CG = mujoco.mjtSolver.mjSOL_CG
+  NEWTON = mujoco.mjtSolver.mjSOL_NEWTON
+
+
+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
+  JOINT = mujoco.mjtEq.mjEQ_JOINT
+  # unsupported: TENDON, DISTANCE
+
+
+class TrnType(enum.IntEnum):
+  """Type of actuator transmission.
+
+  Attributes:
+    JOINT: force on joint
+    SITE: force on site
+  """
+  JOINT = mujoco.mjtTrn.mjTRN_JOINT
+  SITE = mujoco.mjtTrn.mjTRN_SITE
+  # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, BODY
+
+
+class DynType(enum.IntEnum):
+  """Type of actuator dynamics.
+
+  Attributes:
+    NONE: no internal dynamics; ctrl specifies force
+    INTEGRATOR: integrator: da/dt = u
+    FILTER: linear filter: da/dt = (u-a) / tau
+    FILTEREXACT: linear filter: da/dt = (u-a) / tau, with exact integration
+  """
+  NONE = mujoco.mjtDyn.mjDYN_NONE
+  INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR
+  FILTER = mujoco.mjtDyn.mjDYN_FILTER
+  FILTEREXACT = mujoco.mjtDyn.mjDYN_FILTEREXACT
+  # unsupported: 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
+    jacobian:         matrix layout for mass matrices (dense or sparse)
+                      (note that this is different from MuJoCo, where jacobian
+                      specifies whether efc_J and its accompanying matrices
+                      are dense or sparse.
+    solver:           solver algorithm
+    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
+  jacobian: JacobianType
+  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
+    nsite: number of sites
+    nmesh: number of meshes
+    nmeshvert: number of vertices in all meshes
+    nmeshface: number of triangular faces in all meshes
+    nmat: number of materials
+    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_dataid: id of geom's mesh/hfield; -1: none           (ngeom,)
+    geom_group: group for visibility                          (ngeom,)
+    geom_matid: material id for rendering                     (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<margin-gap         (ngeom,)
+    geom_gap: include in solver if dist<margin-gap            (ngeom,)
+    geom_rgba: rgba when material is omitted                  (ngeom, 4)
+    site_bodyid: id of site's body                            (nsite,)
+    site_pos: local position offset rel. to body              (nsite, 3)
+    site_quat: local orientation offset rel. to body          (nsite, 4)
+    mat_rgba: rgba                                            (nmat, 4)
+    mesh_vertadr: first vertex address                        (nmesh x 1)
+    mesh_faceadr: first face address                          (nmesh x 1)
+    mesh_vert: vertex positions for all meshes                (nmeshvert, 3)
+    mesh_face: vertex face data                               (nmeshface, 3)
+    geom_convex_face: vertex face data, MJX only              (ngeom,)
+    geom_convex_vert: vertex data, MJX only                   (ngeom,)
+    geom_convex_edge: unique edge data, MJX only              (ngeom,)
+    geom_convex_facenormal: normal face data, MJX only        (ngeom,)
+    pair_dim: contact dimensionality                          (npair,)
+    pair_geom1: id of geom1                                   (npair,)
+    pair_geom2: id of geom2                                   (npair,)
+    pair_solref: solver reference: contact normal             (npair, mjNREF)
+    pair_solreffriction: solver reference: contact friction   (npair, mjNREF)
+    pair_solimp: solver impedance: contact                    (npair, mjNIMP)
+    pair_margin: include in solver if dist<margin-gap         (npair,)
+    pair_gap: include in solver if dist<margin-gap            (npair,)
+    pair_friction: tangent1, 2, spin, roll1, 2                (npair, 5)
+    exclude_signature: (body1+1) << 16 + body2+1              (nexclude,)
+    eq_type: constraint type (mjtEq)                          (neq,)
+    eq_obj1id: id of object 1                                 (neq,)
+    eq_obj2id: id of object 2                                 (neq,)
+    eq_active0: initial enable/disable constraint state       (neq,)
+    eq_solref: constraint solver reference                    (neq, mjNREF)
+    eq_solimp: constraint solver impedance                    (neq, mjNIMP)
+    eq_data: numeric data for constraint                      (neq, mjNEQDATA)
+    actuator_trntype: transmission type (mjtTrn)              (nu,)
+    actuator_dyntype: dynamics type (mjtDyn)                  (nu,)
+    actuator_gaintype: gain type (mjtGain)                    (nu,)
+    actuator_biastype: bias type (mjtBias)                    (nu,)
+    actuator_trnid: transmission id: joint, tendon, site      (nu, 2)
+    actuator_actadr: first activation address; -1: stateless  (nu,)
+    actuator_actnum: number of activation variables           (nu,)
+    actuator_ctrllimited: is control limited                  (nu,)
+    actuator_forcelimited: is force limited                   (nu,)
+    actuator_actlimited: is activation limited                (nu,)
+    actuator_dynprm: dynamics parameters                      (nu, mjNDYN)
+    actuator_gainprm: gain parameters                         (nu, mjNGAIN)
+    actuator_biasprm: bias parameters                         (nu, mjNBIAS)
+    actuator_ctrlrange: range of controls                     (nu, 2)
+    actuator_forcerange: range of forces                      (nu, 2)
+    actuator_actrange: range of activations                   (nu, 2)
+    actuator_gear: scale length and transmitted force         (nu, 6)
+    numeric_adr: address of field in numeric_data             (nnumeric,)
+    numeric_data: array of all numeric fields                 (nnumericdata,)
+    name_numericadr: numeric name pointers                    (nnumeric,)
+    names: names of all objects, 0-terminated                 (nnames,)
+  """
+  nq: int
+  nv: int
+  nu: int
+  na: int
+  nbody: int
+  njnt: int
+  ngeom: int
+  nsite: int
+  nmesh: int
+  nmeshvert: int
+  nmeshface: int
+  nmat: int
+  npair: int
+  nexclude: int
+  neq: int
+  nnumeric: int
+  nM: int  # pylint:disable=invalid-name
+  opt: Option
+  stat: Statistic
+  qpos0: jax.Array
+  qpos_spring: jax.Array
+  body_parentid: np.ndarray
+  body_rootid: np.ndarray
+  body_weldid: np.ndarray
+  body_jntnum: np.ndarray
+  body_jntadr: np.ndarray
+  body_dofnum: np.ndarray
+  body_dofadr: np.ndarray
+  body_geomnum: np.ndarray
+  body_geomadr: np.ndarray
+  body_pos: jax.Array
+  body_quat: jax.Array
+  body_ipos: jax.Array
+  body_iquat: jax.Array
+  body_mass: jax.Array
+  body_subtreemass: jax.Array
+  body_inertia: jax.Array
+  body_invweight0: jax.Array
+  jnt_type: np.ndarray
+  jnt_qposadr: np.ndarray
+  jnt_dofadr: np.ndarray
+  jnt_bodyid: np.ndarray
+  jnt_limited: np.ndarray
+  jnt_actfrclimited: np.ndarray
+  jnt_solref: jax.Array
+  jnt_solimp: jax.Array
+  jnt_pos: jax.Array
+  jnt_axis: jax.Array
+  jnt_stiffness: jax.Array
+  jnt_range: jax.Array
+  jnt_actfrcrange: jax.Array
+  jnt_margin: jax.Array
+  dof_bodyid: np.ndarray
+  dof_jntid: np.ndarray
+  dof_parentid: np.ndarray
+  dof_Madr: np.ndarray  # pylint:disable=invalid-name
+  dof_solref: jax.Array
+  dof_solimp: jax.Array
+  dof_frictionloss: jax.Array
+  dof_armature: jax.Array
+  dof_damping: jax.Array
+  dof_invweight0: jax.Array
+  dof_M0: jax.Array  # pylint:disable=invalid-name
+  geom_type: np.ndarray
+  geom_contype: np.ndarray
+  geom_conaffinity: np.ndarray
+  geom_condim: np.ndarray
+  geom_bodyid: np.ndarray
+  geom_dataid: np.ndarray
+  geom_group: np.ndarray
+  geom_matid: np.ndarray
+  geom_priority: np.ndarray
+  geom_solmix: jax.Array
+  geom_solref: jax.Array
+  geom_solimp: jax.Array
+  geom_size: jax.Array
+  geom_pos: jax.Array
+  geom_quat: jax.Array
+  geom_friction: jax.Array
+  geom_margin: jax.Array
+  geom_gap: jax.Array
+  geom_rgba: np.ndarray
+  site_bodyid: np.ndarray
+  site_pos: jax.Array
+  site_quat: jax.Array
+  mesh_vertadr: np.ndarray
+  mesh_faceadr: np.ndarray
+  mesh_vert: np.ndarray
+  mesh_face: np.ndarray
+  mat_rgba: np.ndarray
+  pair_dim: np.ndarray
+  pair_geom1: np.ndarray
+  pair_geom2: np.ndarray
+  geom_convex_face: Sequence[jax.Array]
+  geom_convex_vert: Sequence[jax.Array]
+  geom_convex_edge: Sequence[jax.Array]
+  geom_convex_facenormal: Sequence[jax.Array]
+  pair_solref: jax.Array
+  pair_solreffriction: jax.Array
+  pair_solimp: jax.Array
+  pair_margin: jax.Array
+  pair_gap: jax.Array
+  pair_friction: jax.Array
+  exclude_signature: np.ndarray
+  eq_type: np.ndarray
+  eq_obj1id: np.ndarray
+  eq_obj2id: np.ndarray
+  eq_active0: np.ndarray
+  eq_solref: jax.Array
+  eq_solimp: jax.Array
+  eq_data: jax.Array
+  actuator_trntype: np.ndarray
+  actuator_dyntype: np.ndarray
+  actuator_gaintype: np.ndarray
+  actuator_biastype: np.ndarray
+  actuator_trnid: np.ndarray
+  actuator_actadr: np.ndarray
+  actuator_actnum: np.ndarray
+  actuator_ctrllimited: np.ndarray
+  actuator_forcelimited: np.ndarray
+  actuator_actlimited: np.ndarray
+  actuator_dynprm: jax.Array
+  actuator_gainprm: jax.Array
+  actuator_biasprm: jax.Array
+  actuator_ctrlrange: jax.Array
+  actuator_forcerange: jax.Array
+  actuator_actrange: jax.Array
+  actuator_gear: jax.Array
+  numeric_adr: np.ndarray
+  numeric_data: np.ndarray
+  name_numericadr: np.ndarray
+  names: bytes
+
+
+class Contact(PyTreeNode):
+  """Result of collision detection functions.
+
+  Attributes:
+    dist: distance between nearest points; neg: penetration
+    pos: position of contact point: midpoint between geoms            (3,)
+    frame: normal is in [0-2]                                         (9,)
+    includemargin: include if dist<includemargin=margin-gap           (1,)
+    friction: tangent1, 2, spin, roll1, 2                             (5,)
+    solref: constraint solver reference, normal direction             (mjNREF,)
+    solreffriction: constraint solver reference, friction directions  (mjNREF,)
+    solimp: constraint solver impedance                               (mjNIMP,)
+    geom1: id of geom 1
+    geom2: id of geom 2
+  """
+  dist: jax.Array
+  pos: jax.Array
+  frame: jax.Array
+  includemargin: jax.Array
+  friction: jax.Array
+  solref: jax.Array
+  solreffriction: jax.Array
+  solimp: jax.Array
+  # unsupported: mu, H, dim
+  geom1: jax.Array
+  geom2: jax.Array
+  # unsupported: efc_address, exclude
+
+  @classmethod
+  def zero(cls, ncon: int = 0) -> 'Contact':
+    """Returns a contact filled with zeros."""
+    return Contact(
+        dist=jp.zeros(ncon),
+        pos=jp.zeros((ncon, 3,)),
+        frame=jp.zeros((ncon, 3, 3)),
+        includemargin=jp.zeros(ncon),
+        friction=jp.zeros((ncon, 5)),
+        solref=jp.zeros((ncon, mujoco.mjNREF)),
+        solreffriction=jp.zeros((ncon, mujoco.mjNREF)),
+        solimp=jp.zeros((ncon, mujoco.mjNIMP,)),
+        geom1=jp.zeros(ncon, dtype=jp.int32),
+        geom2=jp.zeros(ncon, dtype=jp.int32),
+    )
+
+
+class Data(PyTreeNode):
+  """Dynamic state that updates each step.
+
+  Attributes:
+    solver_niter: number of solver iterations, per island         (mjNISLAND,)
+    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)
+    site_xpos: Cartesian site position                            (nsite, 3)
+    site_xmat: Cartesian site orientation                         (nsite, 9)
+    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                                  if sparse: (nM,)
+                                                       if dense:  (nv, nv)
+    qLD: L'*D*L (or Cholesky) factorization of M.      if sparse: (nM,)
+                                                       if dense:  (nv, nv)
+    qLDiagInv: 1/diag(D)                               if sparse: (nv,)
+                                                       if dense:  (0,)
+    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,)
+    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
+  # 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
+  site_xpos: jax.Array
+  site_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  # pylint:disable=invalid-name
+  qLD: jax.Array  # pylint:disable=invalid-name
+  qLDiagInv: jax.Array  # pylint:disable=invalid-name
+  contact: Contact
+  efc_J: jax.Array  # pylint:disable=invalid-name
+  efc_frictionloss: jax.Array
+  efc_D: jax.Array  # pylint:disable=invalid-name
+  # 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:
+  qfrc_actuator: jax.Array
+  qfrc_smooth: jax.Array
+  qacc_smooth: jax.Array
+  qfrc_constraint: jax.Array
+  qfrc_inverse: jax.Array
+  efc_force: jax.Array

mjx/mujoco/mjx/integration_test/collision_driver_test.py

@@ -0,0 +1,90 @@
+# 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 absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+# pylint: disable=g-importing-member
+from mujoco.mjx._src.types import Contact
+# pylint: enable=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)
+
+
+class CollisionDriverIntegrationTest(parameterized.TestCase):
+
+  @parameterized.parameters(list(range(256)))
+  def test_collision_driver(self, seed):
+    enable_contact = False if seed == 0 else True
+    mjcf = test_util.create_mjcf(
+        seed,
+        body_pos=(0.0, 0.0, 0.14),
+        disable_actuation_pct=100,
+        root_always_free=True,
+        min_trees=1,
+        max_trees=5,
+        max_tree_depth=1,
+        enable_contact=enable_contact,
+    )
+
+    m = mujoco.MjModel.from_xml_string(mjcf)
+    mx = mjx.put_model(m)
+    d = mujoco.MjData(m)
+    dx = mjx.put_data(m, 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)
+
+    if not d.contact.geom1.shape[0]:
+      self.assertTrue((dx.contact.dist > 0).all())
+      return  # no contacts to test
+
+    # re-order MJX contacts to match MJ order
+    idx_mj = list(zip(d.contact.geom1, d.contact.geom2))
+    idx_mjx = list(zip(dx.contact.geom1, dx.contact.geom2))
+    idx_mjx = [tuple(np.array(i)) for i in idx_mjx]
+    self.assertSequenceEqual(set(idx_mjx), set(idx_mj))
+    idx = sorted(range(len(idx_mj)), key=lambda x: idx_mj.index(idx_mjx[x]))
+
+    mjx_contact = jax.tree_map(
+        lambda x: x.take(np.array(idx), axis=0), dx.contact
+    )
+    for field in dataclasses.fields(Contact):
+      _assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/integration_test/forward_test.py

@@ -0,0 +1,70 @@
+# 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
+import jax
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+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 ActuationIntegrationTest(parameterized.TestCase):
+
+  @parameterized.parameters(list(range(30)))
+  def test_actuation(self, seed):
+    mjcf = test_util.create_mjcf(
+        seed,
+        max_trees=1,
+        max_tree_depth=5,
+        body_pos=(0.0, 0.0, -0.5),
+        geom_pos=(0.0, 0.5, 0.0),
+        disable_actuation_pct=10,
+        add_actuators=True,
+        max_contact_excludes=2,
+        enable_contact=False,
+    )
+    m = mujoco.MjModel.from_xml_string(mjcf)
+    actuation_jit_fn = jax.jit(mjx.fwd_actuation)
+
+    # init
+    d = mujoco.MjData(m)
+    d.ctrl = np.random.normal(scale=10, size=m.nu)
+    d.act = np.random.normal(scale=10, size=m.na)
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_fwdPosition(m, d)
+    mujoco.mj_fwdVelocity(m, d)
+
+    # put on device
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_fwdActuation(m, d)
+    dx = actuation_jit_fn(mx, dx)
+
+    _assert_attr_eq(d, dx, 'qfrc_actuator', seed, f'actuator{seed}')
+    _assert_attr_eq(d, dx, 'act_dot', seed, f'actuator{seed}')
+
+
+if __name__ == '__main__':
+  absltest.main()

mjx/mujoco/mjx/integration_test/smooth_test.py

@@ -0,0 +1,78 @@
+# 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
+import mujoco
+from mujoco import mjx
+from mujoco.mjx._src import test_util
+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 TransmissionIntegrationTest(parameterized.TestCase):
+
+  @parameterized.parameters(list(range(30)))
+  def test_transmission(self, seed):
+    """Tests mjx.transmission matches mj_transmission."""
+    mjcf = test_util.create_mjcf(
+        seed,
+        max_trees=1,
+        max_tree_depth=5,
+        body_pos=(0.0, 0.0, -0.5),
+        geom_pos=(0.0, 0.5, 0.0),
+        disable_actuation_pct=10,
+        add_actuators=True,
+        max_contact_excludes=2,
+        enable_contact=False,
+    )
+    m = mujoco.MjModel.from_xml_string(mjcf)
+    transmission_jit_fn = jax.jit(mjx.transmission)
+
+    # init
+    d = mujoco.MjData(m)
+    d.ctrl = np.random.normal(scale=10, size=m.nu)
+    d.act = np.random.normal(scale=10, size=m.na)
+    d.qpos = np.random.normal(m.nq)
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_forward(m, d)
+
+    # put on device
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_transmission(m, d)
+    dx = transmission_jit_fn(mx, dx)
+
+    _assert_attr_eq(d, dx, 'actuator_length', seed, f'transmission{seed}')
+    _assert_attr_eq(
+        d, dx, 'actuator_moment', seed, f'transmission{seed}', atol=1e-4
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()