diff --git a/CMakeLists.txt b/CMakeLists.txt index bc6c5126..8a76576a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -28,7 +28,7 @@ set(MSVC_INCREMENTAL_DEFAULT ON) project( mujoco - VERSION 2.3.8 + VERSION 3.0.1 DESCRIPTION "MuJoCo Physics Simulator" HOMEPAGE_URL "https://mujoco.org" ) diff --git a/cmake/MujocoDependencies.cmake b/cmake/MujocoDependencies.cmake index 48c82034..4e3e2c8d 100644 --- a/cmake/MujocoDependencies.cmake +++ b/cmake/MujocoDependencies.cmake @@ -59,7 +59,7 @@ set(MUJOCO_DEP_VERSION_benchmark ) 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 diff --git a/cmake/sdflib-optional-dependencies.patch b/cmake/sdflib-optional-dependencies.patch deleted file mode 100644 index 46f44931..00000000 --- a/cmake/sdflib-optional-dependencies.patch +++ /dev/null @@ -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 - - #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& 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& values, glm::vec3 fracPart, float nodeSize, std::array& 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; - - inline static float interpolateValue(const std::array& values, glm::vec3 fracPart) -@@ -433,6 +428,16 @@ struct TriCubicInterpolation - - return sum; - } -+#else -+ inline static float interpolateValue(const std::array& 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& 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 -+#endif -+ - #include "utils/Mesh.h" - #include "utils/TriangleUtils.h" - #include "OctreeSdfUtils.h" -@@ -1008,6 +1011,7 @@ struct VHQueries - } - }; - -+#ifdef ENOKI_AVAILABLE - template - 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 - #include - #include -+#ifdef ASSIMP_AVAILABLE - #include - #include - #include -+#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 - 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 mVertices; - std::vector 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 -+#include - #include - #include - #include -@@ -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 - #include - -@@ -46,8 +48,11 @@ OctreeSdf::OctreeSdf(const Mesh& mesh, BoundingBox box, - break; - case OctreeSdf::InitAlgorithm::CONTINUITY: - //initOctreeWithContinuity>(mesh, startDepth, depth, terminationThreshold, terminationRule); -- // initOctreeWithContinuity>(mesh, startDepth, depth, terminationThreshold, terminationRule); -+#ifdef OPENMP_AVAILABLE - initOctreeWithContinuityNoDelay>(mesh, startDepth, depth, terminationThreshold, terminationRule, numThreads); -+#else -+ initOctreeWithContinuity>(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& 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 - #include -+#ifdef OPENMP_AVAILABLE - #include -+#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 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 -+} diff --git a/dist/mujoco.rc b/dist/mujoco.rc index 6804238a..eecb3c24 100644 --- a/dist/mujoco.rc +++ b/dist/mujoco.rc @@ -1,6 +1,6 @@ 1 VERSIONINFO -FILEVERSION 2,3,8,0 -PRODUCTVERSION 2,3,8,0 +FILEVERSION 3,0,1,0 +PRODUCTVERSION 3,0,1,0 FILEOS 0x4 FILETYPE 0x1 { @@ -9,9 +9,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "2.3.8" + VALUE "ProductVersion", "3.0.1" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "2.3.8" + VALUE "FileVersion", "3.0.1" VALUE "InternalName", "mujoco.dll" VALUE "OriginalFilename", "mujoco.dll" VALUE "CompanyName", "Google DeepMind" diff --git a/dist/simulate.rc b/dist/simulate.rc index a7a8ed35..bcf38e79 100644 --- a/dist/simulate.rc +++ b/dist/simulate.rc @@ -1,8 +1,8 @@ MUJOCO ICON "mujoco.ico" 1 VERSIONINFO -FILEVERSION 2,3,8,0 -PRODUCTVERSION 2,3,8,0 +FILEVERSION 3,0,1,0 +PRODUCTVERSION 3,0,1,0 FILEOS 0x4 FILETYPE 0x1 { @@ -11,9 +11,9 @@ FILETYPE 0x1 BLOCK "040904b0" { VALUE "ProductName", "MuJoCo" - VALUE "ProductVersion", "2.3.8" + VALUE "ProductVersion", "3.0.1" VALUE "FileDescription", "MuJoCo" - VALUE "FileVersion", "2.3.8" + VALUE "FileVersion", "3.0.1" VALUE "InternalName", "simulate.exe" VALUE "OriginalFilename", "simulate.exe" VALUE "CompanyName", "Google DeepMind" diff --git a/doc/APIreference/APIglobals.rst b/doc/APIreference/APIglobals.rst index 8a6502d1..467bca5e 100644 --- a/doc/APIreference/APIglobals.rst +++ b/doc/APIreference/APIglobals.rst @@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr - Maximum number of UI rectangles. Defined in `mjui.h `_. * - ``mjVERSION_HEADER`` - - 238 + - 301 - The version of the MuJoCo headers; changes with every release. This is an integer equal to 100x the software version, so 210 corresponds to version 2.1. Defined in mujoco.h. The API function :ref:`mj_version` returns a number with the same meaning but for the compiled library. diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index a97481a7..afb6e339 100644 --- a/doc/XMLreference.rst +++ b/doc/XMLreference.rst @@ -1321,9 +1321,9 @@ also known as terrain map, is a 2D matrix of elevation data. The data can be spe .. _asset-hfield-content_type: -:at:`content_type`: :at-val: `string, optional` +:at:`content_type`: :at-val:`string, optional` If the file attribute is specified, then this sets the - `Media Type `_ (formerly known as MIME types) of the + `Media Type `__ (formerly known as MIME types) of the file to be loaded. Any filename extensions will be overloaded. Currently ``image/png`` and ``image/vnd.mujoco.hfield`` are supported. @@ -1576,173 +1576,23 @@ Associate this mesh with an :ref:`engine plugin`. Either :at:`plugin` Instance name, used for explicit plugin instantiation. - -.. _deformable-skin: .. _asset-skin: :el-prefix:`asset/` |-| **skin** (*) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Skinned meshes (or skins) were added in MuJoCo 2.0. These are deformable meshes whose vertex positions and normals are -computed each time the model is rendered. MuJoCo skins are only used for visualization and do not affect the physics -in any way. In particular, collisions involve the geoms of the bodies to which the skin is attached, and not the skin -itself. Unlike regular meshes which are referenced from geoms and participate in collisions, the skin is not -referenced from anywhere else in the model. It is a stand-alone asset that is used by renderer and not by the -simulator. - -The skin has vertex positions and normals updated at runtime, and triangle faces and optional texture coordinates -which are predefined. It also has "bones" used for updating. Bones are regular MuJoCo bodies referenced with the -:el:`bone` subelement. Each bone has a list of vertex indices and corresponding real-valued weights which specify how -much the bone position and orientation influence the corresponding vertex. The vertex has local coordinates with -respect to every bone that influences it. The local coordinates are computed by the model compiler, given global -vertex coordinates and global bind poses for each body. The bind poses do not have to correspond to the model -reference configuration qpos0. Note that the vertex positions and bone bind poses provided in the skin definition are -always global, even if the model itself is defined in local coordinates. - -At runtime the local coordinates of each vertex with respect to each bone that influences it are converted to global -coordinates, and averaged in proportion to the corresponding weights to obtain a single set of 3D coordinates for each -vertex. Normals then are computed automatically given the resulting global vertex positions and face information. -Finally, the skin can be inflated by applying an offset to each vertex position along its (computed) normal. -Skins are one-sided for rendering purposes; this is because back-face culling is needed to avoid shading and aliasing -artifacts. When the skin is a closed 3D shape this does not matter because the back sides cannot be seen. But if the -skin is a 2D object, we have to specify both sides and offset them slightly to avoid artifacts. Note that the -composite objects introduced in MuJoCo 2.0 generate skins automatically. So one can save an XML model with a composite -object, and obtain an elaborate example of how a skin is specified in the XML. - -Similar to meshes, skins can be specified directly in the XML via attributes documented later, or loaded from a binary -SKN file which is in a custom format. The specification of skins is more complex than meshes because of the bone -subelements. The file format starts with a header of 4 integers: nvertex, ntexcoord, nface, nbone. The first three are -the same as in meshes, and specify the total number of vertices, texture coordinate pairs, and triangle faces in the -skin. ntexcoord can be zero or equal to nvertex. nbone specifies the number of MuJoCo bodies that will be used as -bones in the skin. The header is followed by the vertex, texcoord and face data, followed by a specification for each -bone. The bone specification contains the name of the corresponding model body, 3D bind position, 4D bind quaternion, -number of vertices influenced by the bone, and the vertex index array and weight array. Body names are represented as -fixed-length character arrays and are expected to be 0-terminated. Characters after the first 0 are ignored. The -contents of the SKN file are: - -.. code:: Text - - (int32) nvertex - (int32) ntexcoord - (int32) nface - (int32) nbone - (float) vertex_positions[3*nvertex] - (float) vertex_texcoords[2*ntexcoord] - (int32) face_vertex_indices[3*nface] - for each bone: - (char) body_name[40] - (float) bind_position[3] - (float) bind_quaternion[4] - (int32) vertex_count - (int32) vertex_index[vertex_count] - (float) vertex_weight[vertex_count] - -Similar to the other custom binary formats used in MuJoCo, the file size in bytes is strictly enforced by the model -compiler. The skin file format has subelements so the overall file size formula is difficult to write down, but should -be clear from the above specification. - -.. _deformable-skin-name: .. _asset-skin-name: - -:at:`name`: :at-val:`string, optional` - Name of the skin. - -.. _deformable-skin-file: .. _asset-skin-file: - -:at:`file`: :at-val:`string, optional` - The SKN file from which the skin will be loaded. The path is determined as described in the meshdir attribute of - :ref:`compiler `. If the file is omitted, the skin specification must be provided in the XML using the - attributes below. - -.. _deformable-skin-vertex: .. _asset-skin-vertex: - -:at:`vertex`: :at-val:`real(3*nvert), optional` - Vertex 3D positions, in the global bind pose where the skin is defined. - -.. _deformable-skin-texcoord: .. _asset-skin-texcoord: - -:at:`texcoord`: :at-val:`real(2*nvert), optional` - Vertex 2D texture coordinates, between 0 and 1. Note that skin and geom texturing are somewhat different. Geoms can - use automated texture coordinate generation while skins cannot. This is because skin data are computed directly in - global coordinates. So if the material references a texture, one should specify explicit texture coordinates for the - skin using this attribute. Otherwise the texture will appear to be stationary in the world while the skin moves - around (creating an interesting effect but probably not as intended). - -.. _deformable-skin-face: .. _asset-skin-face: - -:at:`face`: :at-val:`int(3*nface), optional` - Trinagular skin faces. Each face is a triple of vertex indices, which are integers between zero and nvert-1. - -.. _deformable-skin-inflate: .. _asset-skin-inflate: - -:at:`inflate`: :at-val:`real, "0"` - If this number is not zero, the position of vertex during updating will be offset along the vertex normal, but the - distance specified in this attribute. This is particularly useful for skins representing flexible 2D shapes. - -.. _deformable-skin-material: .. _asset-skin-material: - -:at:`material`: :at-val:`string, optional` - If specified, this attribute applies a material to the skin. - -.. _deformable-skin-rgba: .. _asset-skin-rgba: - -:at:`rgba`: :at-val:`real(4), "0.5 0.5 0.5 1"` - Instead of creating material assets and referencing them, this attribute can be used to set color and transparency - only. This is not as flexible as the material mechanism, but is more convenient and is often sufficient. If the value - of this attribute is different from the internal default, it takes precedence over the material. - -.. _deformable-skin-group: .. _asset-skin-group: -:at:`group`: :at-val:`int, "0"` - Integer group to which the skin belongs. This attribute can be used for custom tags. It is also used by the - visualizer to enable and disable the rendering of entire groups of skins. - - -.. _skin-bone: - -:el-prefix:`skin/` |-| **bone** (*) -''''''''''''''''''''''''''''''''''' - -This element defines a bone of the skin. The bone is a regular MuJoCo body which is referenced by name here. - - -.. _skin-bone-body: - -:at:`body`: :at-val:`string, required` - Name of the body corresponding to this bone. - -.. _skin-bone-bindpos: - -:at:`bindpos`: :at-val:`real(3), required` - Global body position corresponding to the bind pose. - -.. _skin-bone-bindquat: - -:at:`bindquat`: :at-val:`real(4), required` - Global body orientation corresponding to the bind pose. - -.. _skin-bone-vertid: - -:at:`vertid`: :at-val:`int(nvert), required` - Integer indices of the vertices influenced by this bone. The vertex index corresponds to the order of the vertex in - the skin mesh. The number of vertex indices specified here (nvert) must equal the number of vertex weights specified - with the next attribute. The same vertex may be influenced by multiple bones, and each vertex must be influenced by - at least one bone. - -.. _skin-bone-vertweight: - -:at:`vertweight`: :at-val:`real(nvert), required` - Weights for the vertices influenced by this bone, in the same order as the vertex indices. Negative weights are - allowed (which is needed for cubic interpolation for example) however the sum of all bone weights for a given vertex - must be positive. +:ref:`Skins` have been moved under the new grouping element :ref:`deformable`. They can +still be specified here but this functionality is now deprecated and will be removed in the future. .. _asset-material: @@ -1750,7 +1600,7 @@ This element defines a bone of the skin. The bone is a regular MuJoCo body which :el-prefix:`asset/` |-| **material** (*) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -This element creates a material asset. It can be referenced from :ref:`skins `, :ref:`geoms `, +This element creates a material asset. It can be referenced from :ref:`skins `, :ref:`geoms `, :ref:`sites ` and :ref:`tendons ` to set their appearance. Note that all these elements also have a local rgba attribute, which is more convenient when only colors need to be adjusted, because it does not require creating materials and referencing them. Materials are useful for adjusting appearance properties beyond color. However @@ -1920,19 +1770,17 @@ adjust it properly through the XML. :at:`o_margin`: :at-val:`real, "0"` This attribute replaces the margin parameter of all active contact pairs when :ref:`Contact override ` is - enabled. Otherwise MuJoCo uses the element-specific margin attribute of :ref:`geom ` or :ref:`pair - ` depending on how the contact pair was generated. See also :ref:`Collision` in the Computation - chapter. The related gap parameter does not have a global override. + enabled. Otherwise MuJoCo uses the element-specific margin attribute of :ref:`geom` or + :ref:`pair` depending on how the contact pair was generated. See also :ref:`Collision` in the + Computation chapter. The related gap parameter does not have a global override. .. _option-o_solref: - .. _option-o_solimp: - .. _option-o_friction: :at:`o_solref`, :at:`o_solimp`, :at:`o_friction` - These attributes replace the solref, solimp, and friction parameters of all active contact pairs when contact - override is enabled. See :ref:`CSolver` for details. + These attributes replace the solref, solimp and friction parameters of all active contact pairs when contact override is + enabled. See :ref:`CSolver` for details. .. _option-integrator: @@ -3714,156 +3562,344 @@ Associate this composite with an :ref:`engine plugin`. Either :at:`plu Instance name, used for explicit plugin instantiation. + .. _body-flexcomp: -:el-prefix:`body/` |-| **flex** (*) -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +:el-prefix:`body/` |-| **flexcomp** (*) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Similar to :el:`composite`, this element (new in MuJoCo 3.0) is not a model element, but rather a macro which expands +into multiple model elements representing a deformable entity. In particular this macro creates one +:ref:`flex` element, a number of bodies that are children of the body in which the :el:`flexcomp` is +defined, and optionally one :ref:`flex equality` which constrains all flex edges to their initial length. +A number of attributes are specified here and then passed through to the automatically-constructed flex. The primary +role of :el:`flexcomp` is to automate the creation of a (possibly large) collection of moving bodies with corresponding +joints, and connect them with stretchable flex elements. See :ref:`flex` and :ref:`deformable +objects` documentation for specifics on how flexes work. Here we only describe the automated construction +process. + +An important distinction between :el:`flex` and :el:`flexcomp` is that the flex references bodies and specifies vertex +coordinates in the frames of those bodies, while the flexcomp defines *points*. Each flexcomp point corresponds to one +body and one vertex in the underlying flex. If the flexcomp point is *pinned*, the corresponding flex body is the parent +body of the flexcomp, while the corresponding flex vertex coordinates equal the flexcomp point coordinates. If the +flexcomp point is not pinned, a new child body is created at the coordinates of the flexcomp point (within the flexcomp +parent body), and then the coordinates of the flex vertex within that new body are (0,0,0). The mechanism for +:ref:`pinning` flexcomp points is explained below. + +Composite objects (available prior to MuJoCo 3.0) needed bodies with geoms for collisions, and sites for connecting +tendons which generated shape-preserving forces. In contrast, flexes generate their own collisions and shape-preserving +forces (as well as rendering), thus the bodies created here are much simpler: no geoms, sites or tendons are needed. +Most of the bodies created here have 3 orthogonal slider joints, corresponding to freely moving point masses. In some +cases we generate radial slider joints, allowing only expansion and contraction. Since no geoms are generated, the +bodies need to have explicit inertial parameters. + +Below is a simple example of a flexcomp, modeling a (somewhat flexible) double pendulum with one end pinned to the +world: + +.. code-block:: xml + + + + + + + + + +This flexcomp has 3 points, however the first point is pinned to the world (i.e. the parent of the flexcomp) and so only +two bodies are automatically created, namely FL_1 and FL_2. Here is what this flexcomp generates after loading and +saving the XML: + +.. code-block:: xml + + + + + + + + + + + + + + + + + + + + + + + -.. _body-flexcomp-name: .. _body-flexcomp-class: -.. _body-flexcomp-type: +:at:`class`: :at-val:`string, optional` + Defaults class for setting unspecified attributes. + +.. _body-flexcomp-name: + +:at:`name`: :at-val:`string, required` + The name of the flex element being generated automatically. This name is used as a prefix for all bodies that are + automatically generated here, and is also referenced by the corresponding flex equality constraint (if applicable). .. _body-flexcomp-dim: +:at:`dim`: :at-val:`int(1), "2"` + Dimensionality of the flex object. This value must be 1, 2 or 3. The flex elements are capsules in 1D, triangles with + radius in 2D, and tetrahedra with radius in 3D. Certain flexcomp types imply a dimensionality, in which case the + value specified here is ignored. + +.. _body-flexcomp-type: + +:at:`type`: :at-val:`[grid, box, cylinder, ellipsoid, mesh, gmsh, direct], "grid"` + This attribute determines the type of :el:`flexcomp` object. The remaining attributes and sub-elements are then + interpreted according to the type. Default settings are also adjusted depending on the type. Different types + correspond to different methods for specifying the flexcomp points and the stretchable elements that connect them. + They fall in three categories: direct specification entered in the XML, direct specification loaded from file, and + automated generation from higher-level specification. + + **grid** generates a rectangular grid of points in 1D, 2D or 3D as specified by :at:`dim`. The number of points in + each dimension is determined by :at:`count` while the grid spacing in each dimension is determined by :at:`spacing`. + Make sure the spacing is sufficiently large relative to :at:`radius` to avoid permanent contacts. In 2D and 3D the + grid is automatically triangulated, and corresponding flex elements are created (triangles or tetrahedra). In 1D the + elements are capsules connecting consecutive pairs of points. + + **box** generates a 3D box object, however flex bodies are only generated on the outer shell. Each flex body has a + radial slider joint allowing it to move in and out from the center of the box. The parent body would normally be a + floating body. The box surface is triangulated, and each flex element is a tetrahedron connecting the center of the + box with one triangle face. :at:`count` and :at:`spacing` determine the count and spacing of the flex bodies, similar + to the **grid** type in 3D. Note that the resulting flex has the same topology as the box generated by + :el:`composite`. + + **cylinder** is the same as **box**, except the points are projected on the surface of a cylinder. + + **ellipsoid** is the same as **box**, except the points are projected on the surface of an ellipsoid. + + **mesh** loads the flexcomp points and elements (i.e. triangles) from a mesh file, in the same file formats as mesh + assets. A mesh asset is not actually added to the model. Instead the vertex and face data from the mesh file are used + to populate the point and element data of the flexcomp. :at:`dim` is automatically set to 2. Recall that a mesh asset + in MuJoCo can be used as a rigid geom attached to a single body. In contrast, the flex generated here corresponds to + a soft mesh with the same initial shape, where each vertex is a separate moving body (unless pinned). + + **gmsh** is similar to mesh, but it loads a `GMSH file `__ + in format 4.1 (ascii or binary). The file extension can be anything; the parser recognizes the format by examining + the file header. This is a very rich file format, allowing all kinds of elements with different dimensionality and + topology. MuJoCo only supports GMSH element types 1, 2, 4 which happen to correspond to our 1D, 2D and 3D flexes. + Only the Nodes and Elements sections of the GMHS file are processed, and used to populate the point and element data + of the flexcomp. The parser will generate an error if the GMSH file contains meshes that are not supported by MuJoCo. + :at:`dim` is automatically set to the dimensionality specified in the GMSH file. Presently this is the only mechanism + to load a large tetrahedral mesh in MuJoCo and generate a corresponding soft entity. If such a mesh is available in a + different file format, use the freely available `GMSH software `__ to convert it to GMSH 4.1. + + **direct** allows the user to specify the point and element data of the flexcomp directly in the XML. Note that + flexcomp will still generate moving bodies automatically, as well as automate other settings; so it still provides + convenience compared to specifing the corresponding flex directly. + .. _body-flexcomp-count: +:at:`count`: :at-val:`int(3), "10 10 10"` + The number of automatically generated points in each dimension. This and the next attribute only apply to types grid, + box, cylinder, ellipsoid. + .. _body-flexcomp-spacing: -.. _body-flexcomp-radius: - -.. _body-flexcomp-rigid: - -.. _body-flexcomp-mass: - -.. _body-flexcomp-inertiabox: - -.. _body-flexcomp-scale: - -.. _body-flexcomp-file: +:at:`spacing`: :at-val:`real(3), "0.02 0.02 0.02"` + The spacing between the automatically generated points in each dimension. The spacing should be sufficiently large + compared to the radius, to avoid permanent contacts. .. _body-flexcomp-point: +:at:`point`: :at-val:`real(3*npoint), optional` + + The 3D coordinates of the points. This attribute is only used with type **direct**. All other flexcomp types generate + their own points. The points are used to construct bodies and vertices as explained earlier. + .. _body-flexcomp-element: -.. _body-flexcomp-material: +:at:`element`: :at-val:`int((dim+1)*npoint), optional` -.. _body-flexcomp-rgba: + The zero-based point ids forming each flex elements. This attribute is only used with type **direct**. All other + flexcomp types generate their own elements. This data is passed through to the automatically-generated flex. .. _body-flexcomp-texcoord: +:at:`texcoord`: :at-val:`real(2*npoint), optional` + + Texture coordinates of each point, passed through to the automatically-generated flex. Note that flexcomp does not + generate texture coordinates automatically, except for 2D grids. For all other types, the user can specify explicit + texture coordinates here, even if the points themselves were generated automatically. This requires understanding of + the layout of the automatically-generated points and how they correspond to the texture referenced by the material. + +.. _body-flexcomp-mass: + +:at:`mass`: :at-val:`real(1), "1"` + The mass of each automatically-generated body equals this value divided by the number of points. Note that pinning + some points does not affect the mass of the other bodies. + +.. _body-flexcomp-inertiabox: + +:at:`inertiabox`: :at-val:`real(1), "0.005"` + Even though the automatically-generated bodies have the physics of point masses, with slider joints, MuJoCo still + requires each body to have rotational inertia. The inertias generated here are diagonal, and are computed such that + the corresponding equivalent-inertia boxes have sides equal to this value. + +.. _body-flexcomp-file: + +:at:`file`: :at-val:`string, optional` + The name of the file from which a **mesh** or a **gmsh** is loaded. For mesh, the file extentsion is used to + determine the file format. Supported formats are the same as in :ref:`mesh assets`. For gmsh, the file is + expected to be in GMSH format 4.1, ascii or binary. + +.. _body-flexcomp-rigid: + +:at:`rigid`: :at-val:`[true, false], "false"` + If this is true, all points correspond to vertices within the parent body, and no new bodies are created. This is + equivalent to pinning all points. Note that if all points are indeed pinned, the model compiler will detect that the + flex is rigid (which behaves is a non-convex mesh in collision detection). + .. _body-flexcomp-pos: +:at:`pos`: :at-val:`real(3), "0 0 0"` + This 3D vector translates all points relative to the frame of the parent body. + .. _body-flexcomp-quat: +:at:`quat`: :at-val:`real(4), "1 0 0 0"` + This is a quaternion rotation of all points around the :at:`pos` vector specified above. Together these two vectors + define a pose transformation, used to position and orient the points as needed. + .. _body-flexcomp-axisangle: - .. _body-flexcomp-xyaxes: - .. _body-flexcomp-zaxis: - .. _body-flexcomp-euler: -.. _body-flexcomp-selfcollide: +:at:`axisangle`, :at:`xyaxes`, :at:`zaxis`, :at:`euler` + Alternative specification of rotation, that can be used instead of :at:`quat`. + +.. _body-flexcomp-scale: + +:at:`scale`: :at-val:`real(3), "1 1 1"` + Scaling of all point coordinates, for types that specify coordinates explicitly. Scaling is applied after the pose + transformation. + +.. _body-flexcomp-radius: +.. _body-flexcomp-material: +.. _body-flexcomp-rgba: +.. _body-flexcomp-group: .. _body-flexcomp-flatskin: -.. _flex-edge: -.. _flexcomp-edge: +:at:`radius`, :at:`material`, :at:`rgba`, :at:`group`, :at:`flatskin` -.. _flex-edge-equality: -.. _flexcomp-edge-equality: +These attributes are directly passed through to the automatically-generated :ref:`flex` object and have +the same meaning. -.. _flex-edge-solref: -.. _flexcomp-edge-solref: - -.. _flex-edge-solimp: -.. _flexcomp-edge-solimp: - -.. _flex-edge-stiffness: -.. _flexcomp-edge-stiffness: - -.. _flex-edge-damping: -.. _flexcomp-edge-damping: - -.. _flex-contact: .. _flexcomp-contact: -.. _flex-contact-contype: +:el-prefix:`flexcomp/` |-| **contact** (*) +'''''''''''''''''''''''''''''''''''''''''' + +.. _flexcomp-contact-internal: +.. _flexcomp-contact-selfcollide: +.. _flexcomp-contact-activelayers: .. _flexcomp-contact-contype: - -.. _flex-contact-conaffinity: .. _flexcomp-contact-conaffinity: - -.. _flex-contact-condim: .. _flexcomp-contact-condim: - -.. _flex-contact-priority: .. _flexcomp-contact-priority: - -.. _flex-contact-friction: .. _flexcomp-contact-friction: - -.. _flex-contact-solmix: .. _flexcomp-contact-solmix: - -.. _flex-contact-solref: .. _flexcomp-contact-solref: - -.. _flex-contact-solimp: .. _flexcomp-contact-solimp: - -.. _flex-contact-margin: .. _flexcomp-contact-margin: - -.. _flex-contact-gap: .. _flexcomp-contact-gap: +:at:`internal`, :at:`selfcollide`, :at:`activelayers`, :at:`contype`, :at:`conaffinity`, :at:`condim`, :at:`priority`, +:at:`friction`, :at:`solmix`, :at:`solimp`, :at:`margin`, :at:`gap` + +Same as in :ref:`flex/contact`. All attributes are passed through to the automatically-generated flex. + +.. _flexcomp-edge: + +:el-prefix:`flexcomp/` |-| **edge** (*) +''''''''''''''''''''''''''''''''''''''' + +Each flex element has one edge in 1D (coinciding with the capsule element), three edges in 2D, and six edges in 3D. The +edges are generated automatically when the flex element is compiled, and the user cannot specify them directly. This +element is used to adjust the properties of all edges in the flex. + +.. _flexcomp-edge-equality: + +:at:`equality`: :at-val:`[true, false], "false"` + When enabled, an equality constraint of :ref:`type flex` is added to the model, referencing the + automatically-generated flex by name. + +.. _flexcomp-edge-solref: +.. _flexcomp-edge-solimp: + +:at:`solref`, :at:`solimp` + The standard constraint parameters, passed through to the automatically generated equality constraint. + +.. _flexcomp-edge-stiffness: +.. _flexcomp-edge-damping: + +:at:`stiffness`, :at:`damping` + Edge stiffness and damping, passed through to the automatically generated flex. + + .. _flexcomp-pin: +:el-prefix:`flexcomp/` |-| **pin** (*) +'''''''''''''''''''''''''''''''''''''' + +Each point is either pinned or not pinned. The effect of pinning was explained earlier. This element is used to specify +which points are pinned. Note that each attribute below can be used to specify multiple pins, and in addition to that, +the :el:`pin` element itself can be repeated for user convenience. The effects are cumulative; pinning the same point +multiple times is allowed. + .. _flexcomp-pin-id: +:at:`id`: :at-val:`int(n), required` + Zero-based ids of points to pin. When the points are automatically-generaged, the user needs to understand their + layout in order to decide which points to pin. This can be done by first creating a flexcomp without any pins, + loading it in the simulator, and showing the body labels. + .. _flexcomp-pin-range: +:at:`range`: :at-val:`int(2*n), required` + Ranges of points to pin. Each range is specified by two integers. + .. _flexcomp-pin-grid: +:at:`grid`: :at-val:`int(dim*n), required` + Grid coordinates of points to pin. This can only be used with type grid. + .. _flexcomp-pin-gridrange: +:at:`gridrange`: :at-val:`int(2*dim*n), required` + Ranges of grid coordinates of points to pin. Each range is specified by (dim) integers for the minimum of the range + followed by (dim) integers for the maximum of the range. This can only be used with type grid. -.. _deformable: +.. _flexcomp-plugin: -.. _deformable-flex: +:el-prefix:`flexcomp/` |-| **plugin** (?) +''''''''''''''''''''''''''''''''''''''''' -.. _deformable-flex-name: +Associate this flexcomp with an :ref:`engine plugin`. Either :at:`plugin` or :at:`instance` are required. -.. _deformable-flex-group: +.. _flexcomp-plugin-plugin: -.. _deformable-flex-material: +:at:`plugin`: :at-val:`string, optional` + Plugin identifier, used for implicit plugin instantiation. -.. _deformable-flex-radius: - -.. _deformable-flex-rgba: - -.. _deformable-flex-texcoord: - -.. _deformable-flex-flatskin: - -.. _deformable-flex-selfcollide: - -.. _deformable-flex-dim: - -.. _deformable-flex-body: - -.. _deformable-flex-vertex: - -.. _deformable-flex-element: - - - -**deformable** (*) -~~~~~~~~~~~~~~~~~~ +.. _flexcomp-plugin-instance: +:at:`instance`: :at-val:`string, optional` + Instance name, used for explicit plugin instantiation. .. _contact: @@ -3984,6 +4020,342 @@ the :ref:`pair ` element above are checked for collisions. The name of the second body in the pair. +.. _deformable: + +**deformable** (*) +~~~~~~~~~~~~~~~~~~ + +This is a grouping element and does not have any attributes. It groups elements that specify deformable objects, namely flexes and skins. + + +.. _deformable-flex: + +:el-prefix:`deformable/` |-| **flex** (*) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Flexible objects (or flexes) were added in MuJoCo 3.0. These are collections of massless stretchable geometric elements +(capsules, triangles or tetrahedra) connecting vertices that are defined within different moving body frames. These +stretchable elements support collisions and contact forces, which are then distributed to all the interconnected bodies. +Flexes also generate passive and constraint forces as needed to simulate deformable entities with the desired material +properties. The modeling of flexes is automated and simplified by the :ref:`flexcomp` element. In most +cases, the user will specify a :el:`flexcomp` which will then automatically construct the corresponding low-level +:el:`flex`. See :ref:`deformable objects` for additional information. + +.. _deformable-flex-name: + +:at:`name`: :at-val:`string, optional` + Name of the flex. + +.. _deformable-flex-dim: + +:at:`dim`: :at-val:`int, "2"` + Dimensionality of the flex. Allowed values are 1, 2 and 3. In 1D the elements are capsules, in 2D the elements are + triangles with radius, in 3D the elements are tetrahedra with (optional) radius. + +.. _deformable-flex-radius: + +:at:`radius`: :at-val:`real, "0.005"` + Radius of all flex elements. It can be zero in 3D, but must be positive in 1D and 2D. The radius affects both + collision detection and rendering. In 1D and 2D it is needed to make the elements volumetric. + +.. _deformable-flex-body: + +:at:`body`: :at-val:`string(nvert or 1), required` + An array of MuJoCo body names (separated by white space) to which each vertex belongs. The number of body names + should either equal the number of vertices (nvert), or be a single body. If a single body is specified, all vertices + are defined within that body - in which case the flex becomes a rigid body. The latter functionality effectively + creates a general non-convex mesh (unlike mesh geoms which are convexified for collision detection purposes). + +.. _deformable-flex-vertex: + +:at:`vertex`: :at-val:`real(3*nvert), optional` + The local coordinates of the vertices within the corresponding body frames. If this attribute is omitted, all + coordinates are (0,0,0) or in other words, the vertices coincide with the centers of the body frames. + +.. _deformable-flex-texcoord: + +:at:`texcoord`: :at-val:`real(2*nvert), optional` + Texture coordinates for each vertex. If omitted, texture mapping for this flex is disabled, even if a texture is + specified in the material. + +.. _deformable-flex-element: + +:at:`element`: :at-val:`int((dim+1)*nelem), required` + For each element of the flex, this lists the zero-based indices of the vertices forming that flex element. We need + two vertices to specify a capsule, three vertices to specify a triangle, and four vertices to specify a tetrahedron - + which is why the number of indices equals (dim+1) times the number of elements. In 2D, the vertices should be listed + in counter-clockwise order. In 1D and 3D the order is irrelevant; in 3D the model compiler will rearrange the + vertices as needed. Repeated vertex indices within a flex element are not allowed. The topology of the flex is not + enforced; it could corespond to a continuous soft body, or a collection of disconnected stretchable elements, or + anything in-between. + +.. _deformable-flex-flatskin: + +:at:`flatskin`: :at-val:`[true, false], "false"` + This attribute determines whether 2D and 3D flexes that are rendered in flexskin mode will use smooth or flat + shading. The default smooth shading is suitable in most cases, however if the object is intended to have visible + sharp edges (such as a cube) then flat shading is more natural. + +.. _deformable-flex-material: + +:at:`material`: :at-val:`string, optional` + If specified, this attribute applies a :ref:`material` to the flex. Note that textures specified in + the material will be applied only if the flex has explicit texture coordinates. + +.. _deformable-flex-rgba: + +:at:`rgba`: :at-val:`real(4), "0.5 0.5 0.5 1"` + Instead of creating material assets and referencing them, this attribute can be used to set color and transparency + only. This is not as flexible as the material mechanism, but is more convenient and is often sufficient. If the value + of this attribute is different from the internal default, it takes precedence over the material. + +.. _deformable-flex-group: + +:at:`group`: :at-val:`int, "0"` + Integer group to which the flex belongs. This attribute can be used for custom tags. It is also used by the + visualizer to enable and disable the rendering of entire groups of flexes. + + +.. _flex-edge: + +:el-prefix:`flex/` |-| **edge** (?) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This element adjusts the passive or constraint properties of all edges of the flex. A flex edge can have a damping +passive force and an :ref:`equality constraint` associated with it, resulting in edge constraint forces. +In the latter case, passive forces are usually unnecessary. For a 1D flex, an edge can also have a passive stiffness, +while ``Solid`` or ``Membrane`` first-party plugins can be used for the 2D and 3D case, respectively. which would +generally make edge constraints unnecessary. However these are modeling choices left to the user. MuJoCo allows all +these mechanisms to be combined as desired. + +.. _flex-edge-stiffness: + +:at:`stiffness`: :at-val:`real(1), "0"` + Stiffness of all edges. Only for 1D flex. For 2D and 3D, plugins must be used. + +.. _flex-edge-damping: + +:at:`damping`: :at-val:`real(1), "0"` + Damping of all edges. + +.. _flex-contact: + +:el-prefix:`flex/` |-| **contact** (?) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This element adjusts the contact properties of the flex. It is mostly identical to geom contact properties, with some +extensions specific to flexes. + +.. _flex-contact-internal: + +:at:`internal`: :at-val:`[true, false], "true"` + Enables or disables internal collisions which prevent flex self-penetration and element inversion. Note that flex + elements that have shared vertices cannot collide (or else there will be permanent contacts). In 1D and 2D, internal + collision checks rely on predefined vertex-element pairs, where the vertex is treated as a sphere with the same + radius as the flex. These spheres correspond to non-shared vertices of neighboring elements on the periphery of the + flex. The pre-defined vertex-element pairs are generated by the model compiler automatically. In 3D, internal + collision checks are performed within each tetraheron: each vertex is collided with the plane corresponding to the + opposing triangle face (again using the flex radius). The resulting contacts are always created with condim 1, gap 0, + margin 0. + +.. _flex-contact-selfcollide: + +:at:`selfcollide`: :at-val:`[none, narrow, bvh, sap, auto], "auto"` + This determines the strategy for midphase collision pruning of element pairs belonging to the same flex. **none** + means flex elements cannot collide with each other. **narrow** means narrow phase only (i.e. all pairs are checked). + This is a diagnostic tool and is never a good idea in practice. **bvh** and **sap** refer to bounding volume + hierarchies and sweep-and-prune (which are two different strategies for midphase collision pruning). **auto** selects + **sap** in 1D and 2D, and **bvh** in 3D. Which strategy performs better depends on the specifics of the model. The + automatic setting is just a simple rule which we have found to perform well in general. + +.. _flex-contact-activelayers: + +:at:`activelayers`: :at-val:`int(1), "1"` + This only has an effect for 3D flexes. Each tetrahedron is labeled by the model compiler with an integer + corresponding to (graph) distance to the outside surface of the flex. Thus outside-facing elements are in layer 0, + their neighbors are in layer 1, etc. This attribute specifies how many layers will be allowed to participate in + collisions. The default setting 1 means that only one layer (i.e. layer 0) can collide, with itself and with the rest + of the world. This is usually sufficient, however if the outer layer is composed of small tetrahedra, another body + can "pierce" it and get stuck. In that case the value should be increased. + + +.. _flex-contact-contype: +.. _flex-contact-conaffinity: +.. _flex-contact-condim: +.. _flex-contact-priority: +.. _flex-contact-friction: +.. _flex-contact-solmix: +.. _flex-contact-solref: +.. _flex-contact-solimp: +.. _flex-contact-margin: +.. _flex-contact-gap: + +.. |deformable/flex/contact attrib list| replace:: + :at:`contype`, :at:`conaffinity`, :at:`condim`, :at:`priority`, :at:`friction`, + :at:`solmix`, :at:`solref`, :at:`solimp`, :at:`margin`, :at:`gap` + +|deformable/flex/contact attrib list| + Same meaning as regular :ref:`geom ` attributes. + + +.. _deformable-skin: + +:el-prefix:`deformable/` |-| **skin** (*) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +These are deformable meshes whose vertex positions and normals are computed each time the model is rendered. MuJoCo +skins are only used for visualization and do not affect the physics in any way. In particular, collisions involve the +geoms of the bodies to which the skin is attached, and not the skin itself. Unlike regular meshes which are referenced +from geoms and participate in collisions, the skin is not referenced from anywhere else in the model. It is a +stand-alone element that is used by renderer and not by the simulator. + +The skin has vertex positions and normals updated at runtime, and triangle faces and optional texture coordinates +which are predefined. It also has "bones" used for updating. Bones are regular MuJoCo bodies referenced with the +:el:`bone` subelement. Each bone has a list of vertex indices and corresponding real-valued weights which specify how +much the bone position and orientation influence the corresponding vertex. The vertex has local coordinates with +respect to every bone that influences it. The local coordinates are computed by the model compiler, given global +vertex coordinates and global bind poses for each body. The bind poses do not have to correspond to the model +reference configuration qpos0. Note that the vertex positions and bone bind poses provided in the skin definition are +always global, even if the model itself is defined in local coordinates. + +At runtime the local coordinates of each vertex with respect to each bone that influences it are converted to global +coordinates, and averaged in proportion to the corresponding weights to obtain a single set of 3D coordinates for each +vertex. Normals then are computed automatically given the resulting global vertex positions and face information. +Finally, the skin can be inflated by applying an offset to each vertex position along its (computed) normal. +Skins are one-sided for rendering purposes; this is because back-face culling is needed to avoid shading and aliasing +artifacts. When the skin is a closed 3D shape this does not matter because the back sides cannot be seen. But if the +skin is a 2D object, we have to specify both sides and offset them slightly to avoid artifacts. Note that the +composite objects introduced in MuJoCo 2.0 generate skins automatically. So one can save an XML model with a composite +object, and obtain an elaborate example of how a skin is specified in the XML. + +Similar to meshes, skins can be specified directly in the XML via attributes documented later, or loaded from a binary +SKN file which is in a custom format. The specification of skins is more complex than meshes because of the bone +subelements. The file format starts with a header of 4 integers: nvertex, ntexcoord, nface, nbone. The first three are +the same as in meshes, and specify the total number of vertices, texture coordinate pairs, and triangle faces in the +skin. ntexcoord can be zero or equal to nvertex. nbone specifies the number of MuJoCo bodies that will be used as +bones in the skin. The header is followed by the vertex, texcoord and face data, followed by a specification for each +bone. The bone specification contains the name of the corresponding model body, 3D bind position, 4D bind quaternion, +number of vertices influenced by the bone, and the vertex index array and weight array. Body names are represented as +fixed-length character arrays and are expected to be 0-terminated. Characters after the first 0 are ignored. The +contents of the SKN file are: + +.. code:: Text + + (int32) nvertex + (int32) ntexcoord + (int32) nface + (int32) nbone + (float) vertex_positions[3*nvertex] + (float) vertex_texcoords[2*ntexcoord] + (int32) face_vertex_indices[3*nface] + for each bone: + (char) body_name[40] + (float) bind_position[3] + (float) bind_quaternion[4] + (int32) vertex_count + (int32) vertex_index[vertex_count] + (float) vertex_weight[vertex_count] + +Similar to the other custom binary formats used in MuJoCo, the file size in bytes is strictly enforced by the model +compiler. The skin file format has subelements so the overall file size formula is difficult to write down, but should +be clear from the above specification. + +.. _deformable-skin-name: + +:at:`name`: :at-val:`string, optional` + Name of the skin. + +.. _deformable-skin-file: + +:at:`file`: :at-val:`string, optional` + The SKN file from which the skin will be loaded. The path is determined as described in the meshdir attribute of + :ref:`compiler `. If the file is omitted, the skin specification must be provided in the XML using the + attributes below. + +.. _deformable-skin-vertex: + +:at:`vertex`: :at-val:`real(3*nvert), optional` + Vertex 3D positions, in the global bind pose where the skin is defined. + +.. _deformable-skin-texcoord: + +:at:`texcoord`: :at-val:`real(2*nvert), optional` + Vertex 2D texture coordinates, between 0 and 1. Note that skin and geom texturing are somewhat different. Geoms can + use automated texture coordinate generation while skins cannot. This is because skin data are computed directly in + global coordinates. So if the material references a texture, one should specify explicit texture coordinates for the + skin using this attribute. Otherwise the texture will appear to be stationary in the world while the skin moves + around (creating an interesting effect but probably not as intended). + +.. _deformable-skin-face: + +:at:`face`: :at-val:`int(3*nface), optional` + Trinagular skin faces. Each face is a triple of vertex indices, which are integers between zero and nvert-1. + +.. _deformable-skin-inflate: + +:at:`inflate`: :at-val:`real, "0"` + If this number is not zero, the position of vertex during updating will be offset along the vertex normal, but the + distance specified in this attribute. This is particularly useful for skins representing flexible 2D shapes. + +.. _deformable-skin-material: + +:at:`material`: :at-val:`string, optional` + If specified, this attribute applies a material to the skin. + +.. _deformable-skin-rgba: + +:at:`rgba`: :at-val:`real(4), "0.5 0.5 0.5 1"` + Instead of creating material assets and referencing them, this attribute can be used to set color and transparency + only. This is not as flexible as the material mechanism, but is more convenient and is often sufficient. If the value + of this attribute is different from the internal default, it takes precedence over the material. + +.. _deformable-skin-group: + +:at:`group`: :at-val:`int, "0"` + Integer group to which the skin belongs. This attribute can be used for custom tags. It is also used by the + visualizer to enable and disable the rendering of entire groups of skins. + + +.. _skin-bone: + +:el-prefix:`skin/` |-| **bone** (*) +''''''''''''''''''''''''''''''''''' + +This element defines a bone of the skin. The bone is a regular MuJoCo body which is referenced by name here. + + +.. _skin-bone-body: + +:at:`body`: :at-val:`string, required` + Name of the body corresponding to this bone. + +.. _skin-bone-bindpos: + +:at:`bindpos`: :at-val:`real(3), required` + Global body position corresponding to the bind pose. + +.. _skin-bone-bindquat: + +:at:`bindquat`: :at-val:`real(4), required` + Global body orientation corresponding to the bind pose. + +.. _skin-bone-vertid: + +:at:`vertid`: :at-val:`int(nvert), required` + Integer indices of the vertices influenced by this bone. The vertex index corresponds to the order of the vertex in + the skin mesh. The number of vertex indices specified here (nvert) must equal the number of vertex weights specified + with the next attribute. The same vertex may be influenced by multiple bones, and each vertex must be influenced by + at least one bone. + +.. _skin-bone-vertweight: + +:at:`vertweight`: :at-val:`real(nvert), required` + Weights for the vertices influenced by this bone, in the same order as the vertex indices. Negative weights are + allowed (which is needed for cubic interpolation for example) however the sum of all bone weights for a given vertex + must be positive. + + + .. _equality: **equality** (*) @@ -4083,8 +4455,7 @@ of the other body, without any joint elements in the child body. :at:`body2`: :at-val:`string, optional` Name of the second body. If this attribute is omitted, the second body is the world body. Welding a body to the world - and changing the corresponding component of :ref:`mjData.eq_active` at runtime can be used to fix the body - temporarily. + and changing the corresponding component of mjModel.eq_active at runtime can be used to fix the body temporarily. .. _equality-weld-relpose: @@ -4193,19 +4564,26 @@ This element constrains the length of one tendon to be a quartic polynomial of a .. _equality-flex: :el-prefix:`equality/` |-| **flex** (*) -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +This element constrains the lengths of all edges of a specified flex to their respective lengths in the initial model +configuration. In this way the edges are used to maintain the shape of the deformable entity. Note that all other +equality constraint types add a fixed number of scalar constraints, while this element adds as many scalar constraints +as there are edges in the specified flex. .. _equality-flex-name: - .. _equality-flex-class: +.. _equality-flex-active: +.. _equality-flex-solref: +.. _equality-flex-solimp: + +:at:`name`, :at:`class`, :at:`active`, :at:`solref`, :at:`solimp` + Same as in :ref:`connect ` element. .. _equality-flex-flex: -.. _equality-flex-active: - -.. _equality-flex-solref: - -.. _equality-flex-solimp: +:at:`flex`: :at-val:`string, required` + Name of the flex whose edges are being constrained. .. _equality-distance: @@ -7346,44 +7724,6 @@ tendon, slidersite, cranksite. All :ref:`adhesion ` attributes are available here except: name, class, body. -.. _default-flex: - -.. _default-flex-contype: - -.. _default-flex-conaffinity: - -.. _default-flex-condim: - -.. _default-flex-priority: - -.. _default-flex-material: - -.. _default-flex-friction: - -.. _default-flex-solmix: - -.. _default-flex-solref: - -.. _default-flex-solimp: - -.. _default-flex-margin: - -.. _default-flex-gap: - -.. _default-flex-stiffness: - -.. _default-flex-damping: - -.. _default-flex-radius: - -.. _default-flex-rgba: - -.. _default-flex-dim: - -:el-prefix:`default/` |-| **flex** (?) -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - .. _custom: **custom** (*) diff --git a/doc/XMLschema.rst b/doc/XMLschema.rst index 0458347c..149ebc74 100644 --- a/doc/XMLschema.rst +++ b/doc/XMLschema.rst @@ -485,19 +485,19 @@ | :ref:`flexcomp | \* | :class: mjcf-attributes | | ` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`name` | :ref:`class` | :ref:`type` | :ref:`dim` | | +| | | | :ref:`name` | :ref:`class` | :ref:`type` | :ref:`group` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`flatskin` | :ref:`count` | :ref:`spacing` | :ref:`radius` | | +| | | | :ref:`dim` | :ref:`count` | :ref:`spacing` | :ref:`radius` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`rigid` | :ref:`mass` | :ref:`inertiabox` | :ref:`scale` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`file` | :ref:`point` | :ref:`element` | :ref:`texcoord` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`material` | :ref:`rgba` | :ref:`selfcollide` | :ref:`flatskin` | | +| | | | :ref:`material` | :ref:`rgba` | :ref:`flatskin` | :ref:`pos` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`pos` | :ref:`quat` | :ref:`axisangle` | :ref:`xyaxes` | | +| | | | :ref:`quat` | :ref:`axisangle` | :ref:`xyaxes` | :ref:`zaxis` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`zaxis` | :ref:`euler` | | | | +| | | | :ref:`euler` | | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flexcomp |br| |_2| |L| | | .. table:: | @@ -517,7 +517,9 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`friction` | :ref:`solmix` | :ref:`solref` | :ref:`solimp` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`margin` | :ref:`gap` | | | | +| | | | :ref:`margin` | :ref:`gap` | :ref:`internal` | :ref:`selfcollide` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`activelayers` | | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flexcomp |br| |_2| |L| | | .. table:: | @@ -527,6 +529,20 @@ | | | | :ref:`id` | :ref:`range` | :ref:`grid` | :ref:`gridrange` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| |_2| flexcomp |br| |_2| |L| | | .. table:: | +| :ref:`plugin | \* | :class: mjcf-attributes | +| ` | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`plugin` | :ref:`instance` | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | ++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ +| |_3| plugin |br| |_3| |L| | | .. table:: | +| :ref:`config | \* | :class: mjcf-attributes | +| ` | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`key` | :ref:`value` | | | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | ++------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | mujoco |br| |L| | | *no attributes* | | :ref:`deformable` | | | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ @@ -536,9 +552,9 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`name` | :ref:`group` | :ref:`dim` | :ref:`radius` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`material` | :ref:`rgba` | :ref:`flatskin` | :ref:`selfcollide` | | +| | | | :ref:`material` | :ref:`rgba` | :ref:`flatskin` | :ref:`body` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`body` | :ref:`vertex` | :ref:`element` | :ref:`texcoord` | | +| | | | :ref:`vertex` | :ref:`element` | :ref:`texcoord` | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flex |br| |_2| |L| | | .. table:: | @@ -549,7 +565,9 @@ | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | | | | | :ref:`friction` | :ref:`solmix` | :ref:`solref` | :ref:`solimp` | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | -| | | | :ref:`margin` | :ref:`gap` | | | | +| | | | :ref:`margin` | :ref:`gap` | :ref:`internal` | :ref:`selfcollide` | | +| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +| | | | :ref:`activelayers` | | | | | | | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ | +------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | |_2| flex |br| |_2| |L| | | .. table:: | diff --git a/doc/changelog.rst b/doc/changelog.rst index 1c5cd20f..151c6ebf 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -5,18 +5,27 @@ Changelog Upcoming version (not yet released) ----------------------------------- +Bug fixes +^^^^^^^^^ + +1. Fix in simulate: correct handling of "Pause update", "Fullscreen" and "VSync" buttons. + +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`. Island discovery can be activated using a new :ref:`enable flag`. - If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island, - see video. + - MJX is designed to work with on-device reinforcement learning algorithms. This Colab notebook demonstrates using + MJX along with reinforcement learning to train humanoid and quadruped robots to locomote: |colab| + - The MJX API is compatible with MuJoCo but is missing some features in this release. See the outline of + :ref:`MJX feature parity ` for more details. + +.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg + :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb .. youtube:: QewlEqIZi1o :align: right @@ -28,7 +37,37 @@ 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` for more details. -3. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine +.. youtube:: ra2bTiZHGlw + :align: right + :width: 240px + +3. Added new low-level model element called ``flex``, used to define deformable objects. These + `simplicial complexes `__ can be of dimension 1, 2 + or 3, corresponding to stretchable lines, triangles or tetrahedra. Two new MJCF elements are used + to define flexes. The top-level :ref:`deformable` section contains the low-level flex definition. + The :ref:`flexcomp` element, similar to :ref:`composite` is a convenience macro for + creating deformables, and supports the GMSH tetrahedral file format. + + - Added `shell `__ passive force plugin, + computing bending forces 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` and :ref:`composite`, + and both are modifiable by the first-party + `elasticity plugins `__. 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`. Island discovery can be activated using a new :ref:`enable flag`. + If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island, + see video. Island discovery is currently disabled for models that have deformable objects (see 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` and the @@ -40,16 +79,8 @@ New features Engine-internal threading is a work in progress and currently only available in first-party code via the :ref:`testspeed` utility, exposed with the ``npoolthread`` flag. -.. youtube:: ra2bTiZHGlw - :align: right - :width: 240px - -4. Added capability to initialize :ref:`composite` particles with arbitrary positions. - -5. Added `shell `__ passive force plugin: - - - Collisions use spheres located at mesh vertices. - - Stretching as tendon constraints and bending using a constant precomputed Hessian (cotangent operator). +6. Added capability to initialize :ref:`composite` particles from OBJ files. Fixes :github:issue:`642` + and :github:issue:`674`. General ^^^^^^^ @@ -57,32 +88,32 @@ General .. admonition:: Breaking API changes :class: attention - 6. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``. + 7. Removed the macros ``mjMARKSTACK`` and ``mjFREESTACK``. **Migration:** These macros have been replaced by new functions :ref:`mj_markStack` and :ref:`mj_freeStack`. These functions manage the :ref:`mjData stack` in a fully encapsulated way (i.e., without introducing a local variable at the call site). - 7. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an + 8. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an arbitrary number of bytes and has an additional argument for specifying the alignment of the returned pointer. **Migration:** The functionality for allocating ``mjtNum`` arrays is now available via :ref:`mj_stackAllocNum`. - 8. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``, + 9. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``, and ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s. - 9. Changed :ref:`mjData.solver`, the array used to collect solver diagnostic information. - This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix. - Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. - If the solver does not use islands, only row 0 is filled. + 10. Changed :ref:`mjData.solver`, the array used to collect solver diagnostic information. + This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix. + Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island. + If the solver does not use islands, only row 0 is filled. - - The new constant :ref:`mjNISLAND` was set to 20. - - :ref:`mjNSOLVER` was reduced from 1000 to 200. - - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. - - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer - vectors of length ``mjNISLAND``. + - The new constant :ref:`mjNISLAND` was set to 20. + - :ref:`mjNSOLVER` was reduced from 1000 to 200. + - Added :ref:`mjData.solver_nisland`: the number of islands for which the solver ran. + - Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer + vectors of length ``mjNISLAND``. - 10. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute. + 11. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute. **Migration:** @@ -93,40 +124,39 @@ General :ref:`conaffinity` attributes in the model and then setting them globally to ``0`` using |br| `` ``. - 11. Removed the :at:`rope` and :at:`cloth` composite objects. + 12. Removed the :at:`rope` and :at:`cloth` composite objects. **Migration:** Users should use the :at:`cable` and :at:`shell` elasticity plugins. - 12. Added :ref:`mjData.eq_active` user input variable, for enabling/disabling the state of equality + 13. Added :ref:`mjData.eq_active` user input variable, for enabling/disabling the state of equality constraints. Renamed ``mjModel.eq_active`` to :ref:`mjModel.eq_active0`, which now has the semantic of - "initial value of ``mjData.eq_active``". - Fixes `#876 `__ + "initial value of ``mjData.eq_active``". Fixes :github:issue:`876`. **Migration:** Replace uses of ``mjModel.eq_active`` with ``mjData.eq_active``. - 13. Changed the default of :ref:`autolimits` from "false" to "true". This is a minor breaking + 14. Changed the default of :ref:`autolimits` from "false" to "true". This is a minor breaking change. The potential breakage applies to models which have elements with "range" defined and "limited" not set. Such models cannot be loaded since version 2.2.2 (July 2022). -14. Added a new :ref:`dyntype`, ``filterexact``, which updates first-order filter states with +15. Added a new :ref:`dyntype`, ``filterexact``, which updates first-order filter states with the exact formula rather than with Euler integration. -15. Added an actuator attribute, :ref:`actearly`, which uses semi-implicit integration for +16. Added an actuator attribute, :ref:`actearly`, which uses semi-implicit integration for actuator forces: using the next step's actuator state to compute the current actuator forces. -16. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to +17. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to :ref:`actuatorfrcrange` and :ref:`actuatorfrclimited`, respectively. -17. Added the flag :ref:`eulerdamp`, which disables implicit integration of joint damping in the +18. Added the flag :ref:`eulerdamp`, which disables implicit integration of joint damping in the Euler integrator. See the :ref:`Numerical Integration` section for more details. -18. Added the flag :ref:`invdiscrete`, which enables discrete-time inverse dynamics for all +19. Added the flag :ref:`invdiscrete`, which enables discrete-time inverse dynamics for all :ref:`integrators` other than ``RK4``. See the flag documentation for more details. -19. Added :ref:`ls_iterations` and :ref:`ls_tolerance` options for adjusting +20. Added :ref:`ls_iterations` and :ref:`ls_tolerance` options for adjusting linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning. -20. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to - mesh assets. Fixes `#409 `__ . -21. Added camera :ref:`resolution` attribute and :ref:`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` attribute and :ref:`camprojection` sensor. If camera resolution is set to positive values, the camera projection sensor will report the location of a target site, projected onto the camera image, in pixel coordinates. -22. Added :ref:`camera` calibration attributes: +23. Added :ref:`camera` calibration attributes: - The new attributes are :ref:`resolution`, :ref:`focal`, :ref:`focalpixel`, :ref:`principal`, @@ -135,23 +165,21 @@ General attributes are specified. See the following `example model `__. - Note that these attributes only take effect for offline rendering and do not affect interactive visualisation. -23. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values +24. Implemented reversed Z rendering for better depth precision. An enum :ref:`mjtDepthMap` was added with values ``mjDEPTH_ZERONEAR`` and ``mjDEPTH_ZEROFAR``, which can be used to set the new ``readDepthMap`` attribute in :ref:`mjrContext` to control how the depth returned by :ref:`mjr_readPixels` is mapped from ``znear`` to ``zfar``. - `Contribution `__ by - `Levi Burner `__. -24. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the + Contribution :github:pull:`978` by `Levi Burner `__. +25. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the `WriteReadCompare `__ test. -25. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`. +26. Deleted the code sample ``derivative``. Functionality provided by :ref:`mjd_transitionFD`. Python bindings ^^^^^^^^^^^^^^^ -26. Fixed `#870 `__ where calling ``update_scene`` with an invalid - camera name used the default camera. -27. Added ``user_scn`` to the :ref:`passive viewer` handle, which allows users to add custom - visualization geoms (`#1023 `__). -28. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch`` +27. Fixed :github:issue:`870` where calling ``update_scene`` with an invalid camera name used the default camera. +28. Added ``user_scn`` to the :ref:`passive viewer` handle, which allows users to add custom + visualization geoms (:github:issue:`1023`). +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 @@ -161,21 +189,31 @@ Simulate :align: right :width: 240px -29. Added **state history** mechanism to :ref:`simulate` and the managed +30. Added **state history** mechanism to :ref:`simulate` and the managed :ref:`Python viewer`. State history can be viewed by scrubbing the History slider and (more precisely) with the left and right arrow keys. See screen capture: -30. The ``LOADING...`` label is now shown correctly. - `Contribution `__ by +31. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by `Levi Burner `__. +Documentation +^^^^^^^^^^^^^ + +.. youtube:: nljr0X79vI0 + :align: right + :width: 240px + +32. Added :doc:`detailed documentation ` of fluid force modeling, and an illustrative example model + showing `tumbling cards `__ using the + ellipsoid-based fluid model. + Bug fixes ^^^^^^^^^ -31. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of +33. Fixed a bug that was causing :ref:`geom margin` to be ignored during the construction of midphase collision trees. -32. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints. +34. Fixed a bug that was generating incorrect values in ``efc_diagApprox`` for weld equality constraints. Version 2.3.7 (July 20, 2023) @@ -204,13 +242,10 @@ Python bindings 7. The :ref:`passive viewer` handle now exposes ``update_hfield``, ``update_mesh``, and ``update_texture`` methods to allow users to update renderable assets. - (`#812 `_, - `#958 `_, - `#965 `_) -#. Allow a custom keyboard event callback to be specified in the :ref:`passive viewer`. - (`#766 `_) -#. Fix GLFW crash when Python exits while the passive viewer is running. - (`#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` + (:github:issue:`766`). +#. Fix GLFW crash when Python exits while the passive viewer is running (:github:issue:`790`). Models ^^^^^^ diff --git a/doc/computation/fluid.rst b/doc/computation/fluid.rst index 0497b921..070cb7b3 100644 --- a/doc/computation/fluid.rst +++ b/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`, 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`, 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 @@ -178,12 +198,12 @@ We present the following result. .. admonition:: Lemma :class: note - Given an ellipsoid with semi-axes :math:`(d_x, d_y, d_z)` aligned with the coordinate axes :math:`(x, y, z)`, and a + 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_z^4 d_x^4 u_y^2 + d_x^4 d_y^4 u_z^2}{d_y^2 d_z^2 u_x^2 + d_z^2 d_x^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 @@ -202,10 +222,10 @@ We present the following result. **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:`(d_x, d_y, d_z)` be the semi-axis lengths of the ellipsoid. Without loss of generality, it is sufficient to + :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/ d_x^2 \right., \left. 1 \middle/ d_y^2 \right., \left. 1 \middle/ d_z^2 \right. \right)\mathclose{}` + :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 @@ -266,9 +286,9 @@ We present the following result. .. math:: \begin{align*} - Q'_{xx} &= \frac{1}{d_x^2} R_{xx}^2 + \frac{1}{d_y^2} R_{yx}^2 + \frac{1}{d_z^2} R_{zx}^2 , \\ - Q'_{yy} &= \frac{1}{d_x^2} R_{xy}^2 + \frac{1}{d_y^2} R_{yy}^2 + \frac{1}{d_z^2} R_{zy}^2 , \\ - Q'_{xy} &= \frac{1}{d_x^2} R_{xx} R_{xy} + \frac{1}{d_y^2} R_{yx} R_{yy} + \frac{1}{d_z^2} R_{zx} R_{zy} , + 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 @@ -277,7 +297,7 @@ We present the following result. 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 d_x d_y d_z}{\sqrt{d_x^2 n_x^2 + d_y^2 n_y^2 + d_z^2 n_z^2}}, + = \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}}`. @@ -293,9 +313,9 @@ We present the following result. 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}(d_x, d_y, d_z)`. Furthermore, if :math:`\mathbf{\tilde{u}}` is a vector + 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}}=(d_x \tilde{u}_x, d_y \tilde{u}_y, d_z \tilde{u}_z)` is tangent to the + :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 @@ -303,15 +323,15 @@ We present the following result. 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} = (d_x \tilde{v}_x, d_y \tilde{v}_y, d_z \tilde {v}_z)` and - :math:`\mathbf{w} = (d_x \tilde{w}_x, d_y \tilde{w}_y, d_z \tilde {w}_z)` respectively, and they remain orthogonal + 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} - = (d_y d_z \tilde{u}_x, d_z d_x \tilde{u}_y, d_x d_y \tilde{u}_z) - = \left( \frac{d_y d_z}{d_x} u_x, \frac{d_z d_x}{d_y} u_y, \frac{d_x d_y}{d_z} u_z \right). + = (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 @@ -359,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 @@ -375,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 ~~~~~~~~~~~~ @@ -411,7 +431,7 @@ bagheri2016`. See screen capture of the 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} ~ + @@ -424,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: @@ -435,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 ~~~~~~~~~~~~ @@ -487,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. @@ -512,13 +531,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` above and the corresponding unit vector :math:`\hat{\mathbf{n}}_{s, \mathbf{v}}`. @@ -539,25 +558,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). diff --git a/doc/conf.py b/doc/conf.py index 53c6d9a5..ca6252bc 100644 --- a/doc/conf.py +++ b/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'] diff --git a/doc/images/mjx/mujoco_vs_mjx_large_scene.png b/doc/images/mjx/mujoco_vs_mjx_large_scene.png new file mode 100644 index 00000000..d9ef4b54 Binary files /dev/null and b/doc/images/mjx/mujoco_vs_mjx_large_scene.png differ diff --git a/doc/images/modeling/bunny1.png b/doc/images/modeling/bunny1.png new file mode 100644 index 00000000..d1b6bb3f Binary files /dev/null and b/doc/images/modeling/bunny1.png differ diff --git a/doc/images/modeling/bunny2.png b/doc/images/modeling/bunny2.png new file mode 100644 index 00000000..51b9810b Binary files /dev/null and b/doc/images/modeling/bunny2.png differ diff --git a/doc/images/modeling/coil.png b/doc/images/modeling/coil.png new file mode 100644 index 00000000..9fd9f233 Binary files /dev/null and b/doc/images/modeling/coil.png differ diff --git a/doc/images/modeling/flexelem.png b/doc/images/modeling/flexelem.png new file mode 100644 index 00000000..39df7463 Binary files /dev/null and b/doc/images/modeling/flexelem.png differ diff --git a/doc/index.rst b/doc/index.rst index e4225fd3..0aba3e4c 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -14,6 +14,11 @@ programming/index.rst APIreference/index.rst python + MJX unity models changelog + + +.. sidebar-links:: + :github: diff --git a/doc/mjx.rst b/doc/mjx.rst new file mode 100644 index 00000000..9189afe7 --- /dev/null +++ b/doc/mjx.rst @@ -0,0 +1,312 @@ +========== +MuJoCo XLA +========== + +Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the +`mjx `__ directory. MJX allows MuJoCo to run on compute +hardware supported by the `XLA `__ compiler via the +`JAX `__ framework. MJX runs on a +`all platforms supported by JAX `__: Nvidia +and AMD GPUs, Apple Silicon, and `Google Cloud TPUs `__. + +The MJX API is consistent with the main simulation functions in the MuJoCo API, although it is currently missing some +features. While the :ref:`API documentation ` is applicable to both libraries, we indicate features +unsupported by MJX in the :ref:`notes ` below. + +MJX is distributed as a separate package called ``mujoco-mjx`` on `PyPI `__. +Although it depends on the main ``mujoco`` package for model compilation and visualization, it is a re-implementation of +MuJoCo that uses the same algorithms as the MuJoCo implementation. However, in order to properly leverage JAX, MJX +deliberately diverges from the MuJoCo API in a few places, see below. + +MJX is a successor to the `generalized physics pipeline `__ +in Google's `Brax `__ physics and reinforcement learning library. MJX was built +by core contributors to both MuJoCo and Brax, who will together continue to support both Brax (for its reinforcement +learning algorithms and included environments) and MJX (for its physics algorithms). A future version of Brax will +depend on the ``mujoco-mjx`` package, and Brax's existing +`generalized pipeline `__ will be deprecated. This change +will be largely transparent to users of Brax. + +.. _MjxNotebook: + +Tutorial notebook +================= + +The following IPython notebook demonstrates the use of MJX along with reinforcement learning to train humanoid and +quadruped robots to locomote: |colab|. + +.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg + :target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb + +.. _MjxInstallation: + +Installation +============ + +The recommended way to install this package is via `PyPI `__: + +.. code-block:: shell + + pip install mujoco-mjx + +A copy of the MuJoCo library is provided as part of this package's depdendencies and does **not** need to be downloaded +or installed separately. + +.. _MjxUsage: + +Basic usage +=========== + +Once installed, the package can be imported via ``from mujoco import mjx``. Structs, functions, and enums are available +directly from the top-level ``mjx`` module. + +.. _MjxStructs: + +Structs +------- + +Before running MJX functions on an accelerator device, structs must be copied onto the device via the ``mjx.device_put`` +function. Placing an :ref:`mjModel` on device yields an ``mjx.Model``. Placing an :ref:`mjData` on device yields +an ``mjx.Data``: + +.. code-block:: python + + model = mujoco.MjModel.from_xml_string("...") + data = mujoco.MjData(model) + mjx_model = mjx.device_put(model) + mjx_data = mjx.device_put(data) + +These MJX variants mirror their MuJoCo counterparts but have three key differences: + +#. Fields in ``mjx.Model`` and ``mjx.Data`` are JAX arrays copied onto device, instead of numpy arrays. +#. Some fields are missing from ``mjx.Model`` and ``mjx.Data`` for features that are + :ref:`unsupported ` in MJX. +#. Arrays in ``mjx.Model`` and ``mjx.Data`` support adding batch dimensions. Batch dimensions are a natural way to + express domain randomization (in the case of ``mjx.Model``) or high-throughput simulation for reinforcement learning + (in the case of ``mjx.Data``). + + +Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An ``mjx.Data`` may be created by calling +``mjx.make_data``, which mirrors the :ref:`mj_makeData` function in MuJoCo: + +.. code-block:: python + + model = mujoco.MjModel.from_xml_string("...") + mjx_model = mjx.device_put(model) + mjx_data = mjx.make_data(model) + +Using ``mx.make_data`` may be preferable when constructing batched ``mjx.Data`` structures inside of a ``vmap``. + +.. _MjxFunctions: + +Functions +--------- + +MuJoCo functions are exposed as MJX functions of the same name, but following +`PEP 8 `__-compliant names. Most of the :ref:`main simulation ` and +some of the :ref:`sub-components ` for forward simulation are available from the top-level ``mjx`` module. + +MJX functions are not `JIT compiled `__ by default -- we +leave it to the user to JIT MJX functions, or JIT their own functions that reference MJX functions. See the +:ref:`minimal example ` below. + +.. _MjxEnums: + +Enums and constants +------------------- + +MJX enums are available as ``mjx.EnumType.ENUM_VALUE``, for example ``mjx.JointType.FREE``. Enums for unsupported MJX +features are omitted from the MJX enum declaration. MJX declares no constants but references MuJoCo constants directly. + +.. _MjxExample: + +Minimal example +--------------- + +.. code-block:: python + + # Throw a ball at 100 different velocities. + + import jax + import mujoco + from mujoco import mjx + + XML=r""" + + + + + + + + + """ + + model = mujoco.MjModel.from_xml_string(XML) + mjx_model = mjx.device_put(model) + + @jax.vmap + def batched_step(vel): + mjx_data = mjx.make_data(mjx_model) + qvel = mjx_data.qvel.at[0].set(vel) + mjx_data = mjx_data.replace(qvel=qvel) + pos = mjx.step(mjx_model, mjx_data).qpos[0] + return pos + + vel = jax.numpy.arange(0.0, 1.0, 0.01) + pos = jax.jit(batched_step)(vel) + print(pos) + +.. _MjxFeatureParity: + +Feature Parity +============== + +MJX supports most of the main simulation features of MuJoCo, with a few exceptions. MJX will raise an exception if +asked to copy to device an :ref:`mjModel` with field values referencing unsupported features. + +The following features are **fully supported** in MJX: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - Dynamics + - :ref:`Forward ` + * - :ref:`Joint ` + - ``FREE``, ``BALL``, ``SLIDE``, ``HINGE`` + * - :ref:`Transmission ` + - ``TRN_JOINT`` + * - :ref:`Actuation ` + - ``DYN_NONE``, ``DYN_INTEGRATOR``, ``DYN_FILTER``, ``GAIN_FIXED``, ``GAIN_AFFINE``, ``BIAS_NONE``, + ``BIAS_AFFINE`` + * - :ref:`Geom ` + - ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` + * - :ref:`Constraint ` + - ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL`` + * - :ref:`Integrator ` + - ``EULER``, ``RK4`` + * - :ref:`Cone ` + - ``PYRAMIDAL`` + * - :ref:`Condim ` + - 3 + * - :ref:`Solver ` + - ``CG`` + * - Fluid Model + - :ref:`flInertia` + +The following features are **in development** and coming soon: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - Dynamics + - :ref:`Inverse ` + * - :ref:`Transmission ` + - ``TRN_TENDON`` + * - :ref:`Geom ` + - ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF`` + * - :ref:`Integrator ` + - ``IMPLICIT``, ``IMPLICITFAST`` + * - :ref:`Cone ` + - ``ELLIPTIC`` + * - :ref:`Condim ` + - 1, 4, 6 + * - :ref:`Solver ` + - ``NEWTON`` + * - Fluid Model + - :ref:`flEllipsoid` + * - :ref:`Tendons ` + - :ref:`Spatial `, :ref:`Fixed ` + +The following features are **unsupported**: + +.. list-table:: + :width: 90% + :align: left + :widths: 1 5 + :header-rows: 1 + + * - Category + - Feature + * - :ref:`Transmission ` + - ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``, ``MUSCLE`` + * - :ref:`Solver ` + - ``PGS`` + * - :ref:`Callbacks ` + - ``mjDYN_USER``, ``mjGAIN_USER``, ``mjBIAS_USER``, ``mjSENS_USER`` + +.. _MjxSharpBits: + +🔪 MJX - The Sharp Bits 🔪 +========================== + +GPUs and TPUs have unique performance tradeoffs that MJX is subject to. MJX specializes in simulating big batches of +parallel identical physics scenes using algorithms that can be efficiently vectorized on +`SIMD hardware `__. This specialization is useful +for machine learning workloads such as `reinforcement learning `__ +that require massive data throughput. + +There are certain workflows that MJX is ill-suited for: + +Single scene simulation + Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been + carefully optimized for CPU. MJX works best when simulating thousands or tens of thousands of scenes in parallel. + +Large, complex scenes with many contacts + Accelerators exhibit poor performance for + `branching code `__. + Branching is used in broad-phase collision detection, when identifying potential collisions between large numbers of + bodies in a scene. MJX ships with a simple branchless broad-phase algorithm (see performance tuning) but it is not as + powerful as the one in MuJoCo. + + To see how this affects simulation, let us consider a physics scene with increasing numbers of physics bodies. We + simulate a scene with a variable number of humanoids (from 1 to 10) and then compare MJX's performance on an Nvidia + A100 GPU to MuJoCo on a 12-core workstation: + + .. figure:: images/mjx/mujoco_vs_mjx_large_scene.png + :width: 658px + :align: center + + Notice that as we increase the number of humanoids (which increases the number of potential contacts in a scene), MJX + performance degrades more rapidly than MuJoCo. At the limit, for such a large scene, MuJoCo performance nearly + matches MJX. + +Scenes with collisions between meshes with many vertices + MJX supports mesh geometries and can determine if two meshes are colliding using branchless versions of + `mesh collision algorithms `__. + These algorithms work well for smaller meshes (with hundreds of vertices) but suffer with large meshes. With careful + tuning, MJX can simulate scenes with mesh collisions well -- see the MJX + `shadow hand `__ + config for an example. + +.. _MjxPerformance: + +Performance tuning +================== + +For MJX to perform well, some configuration parameters should be adjusted from their default MuJoCo values: + +:ref:`option` element + For now, solver must be set to ``CG`` (but Newton is on its way!). The ``iterations`` and ``ls_iterations`` + attributes---which control solver and linesearch iterations, respectively---should be brought down to just low enough + that the simulation remains stable. Accurate solver forces are not so important in reinforcement learning in which + domain randomization is often used to add noise to physics for sim2real. + +:ref:`contact-pair` element + Consider explicitly marking geoms for collision detection to reduce the number of contacts that MJX must consider + during each step. Enabling only an explicit list of valid contacts can have a dramatic effect on simulation + performance in MJX. Doing this well often requires an understanding of the task -- for example, the + `OpenAI Gym Humanoid `__ task resets when + the humanoid starts to fall, so full contact with the floor is not needed. + +:ref:`option-flag` element + Disabling ``eulerdamp`` can help performance and is often not needed for stability. diff --git a/doc/modeling.rst b/doc/modeling.rst index 20718400..8a4d5d0e 100644 --- a/doc/modeling.rst +++ b/doc/modeling.rst @@ -1069,6 +1069,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 +` 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 ` section, which outlines how to use +:ref:`flexcomp` to create such an object. It is easy to port models create with composite particles to +flex, see the folder `elasticity/ `__ for +several examples. + **1D grid**. |image6| |image7| @@ -1111,59 +1119,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 - - - - - - - + + + -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. + + + + + + + + + + + + + +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`, 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 ` for extensible strings in a tensile loading +scenario (e.g. a stretched rubber band). **Cloth**. -|image12| |image13| - -.. code-block:: xml - - - - - - - - - - -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 ` for simulating thin +elastic structures. **Box**. @@ -1222,6 +1222,94 @@ 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 ` 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` described earlier was actually one such element, but it is merely used for +visualization. We now have a related element :ref:`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`. 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` which automates the creation of the low-level flex, +similar to how :ref:`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: + +|flexelem| + +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 +`__ of the material. For more details, see the `Saint Venant-Kirchhoff +`__ hyperelastic model. This +functionality is currently based on first-party :ref:`engine plugins` as of MuJoCo 3.0 but may be integrated +into the engine in future releases. + +**Creation and visualization**. + +.. code-block:: xml + + + + + + + + + + + + + + + + + +Using the :ref:`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. + +|bunny1| |bunny2| + + .. _CInclude: Including files @@ -1562,3 +1650,11 @@ in a visible way, and the energy fluctuates around the initial value instead of :height: 250px .. |particle| image:: images/models/particle.gif :width: 270px +.. |flexelem| image:: images/modeling/flexelem.png + :width: 400px +.. |bunny1| image:: images/modeling/bunny1.png + :width: 300px +.. |bunny2| image:: images/modeling/bunny2.png + :width: 300px +.. |coil| image:: images/modeling/coil.png + :width: 300px diff --git a/doc/overview.rst b/doc/overview.rst index bb89671e..da10f6e4 100644 --- a/doc/overview.rst +++ b/doc/overview.rst @@ -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 ^^^^^^^^^^^^^^^ diff --git a/doc/programming/extension.rst b/doc/programming/extension.rst index 9b9dba15..5d8167e3 100644 --- a/doc/programming/extension.rst +++ b/doc/programming/extension.rst @@ -256,6 +256,7 @@ Currently, there are three directories of first-party plugins: bending strains. The 3D solid is a `Saint Venant-Kirchhoff `__ model discretized with piecewise linear finite elements, which is suitable for large deformations with small strains. + See also :ref:`composite ` and :ref:`deformable ` objects. * **sensor:** The plugins in the `sensor/ `__ directory implement custom sensors. Currently the sole sensor plugin is the touch grid sensor, see the `README `__ for details. diff --git a/doc/unity.rst b/doc/unity.rst index 9fe13ef9..13ddec5a 100644 --- a/doc/unity.rst +++ b/doc/unity.rst @@ -30,14 +30,14 @@ _____ The MuJoCo app needs to be run at least once before the native library can be used, in order to register the library as a trusted binary. Then, copy the dynamic library file from -``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.2.3.8.dylib`` (it can be +``/Applications/MuJoCo.app/Contents/Frameworks/mujoco.framework/Versions/Current/libmujoco.3.0.1.dylib`` (it can be found by browsing the contents of ``MuJoCo.app``) and rename it as ``mujoco.dylib``. Linux _____ Expand the ``tar.gz`` archive to ``~/.mujoco``. Then copy the dynamic library from -``~/.mujoco/mujoco-2.3.8/lib/libmujoco.so.2.3.8`` and rename it as ``libmujoco.so``. +``~/.mujoco/mujoco-3.0.1/lib/libmujoco.so.3.0.1`` and rename it as ``libmujoco.so``. Windows _______ diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 61d63325..3cd4ff6b 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -24,7 +24,7 @@ extern "C" { #endif // header version; should match the library version as returned by mj_version() -#define mjVERSION_HEADER 238 +#define mjVERSION_HEADER 301 // needed to define size_t, fabs and log10 #include diff --git a/mjx/MANIFEST.in b/mjx/MANIFEST.in new file mode 100644 index 00000000..b963ce6e --- /dev/null +++ b/mjx/MANIFEST.in @@ -0,0 +1,2 @@ +recursive-include mujoco/mjx/test_data * +recursive-include mujoco/mjx/benchmark *.obj *.stl *.xml diff --git a/mjx/README.md b/mjx/README.md new file mode 100644 index 00000000..4bb727bd --- /dev/null +++ b/mjx/README.md @@ -0,0 +1,54 @@ +# MuJoCo XLA (MJX) + +[![PyPI Python Version][pypi-versions-badge]][pypi] +[![PyPI version][pypi-badge]][pypi] + +[pypi-versions-badge]: https://img.shields.io/pypi/pyversions/mujoco-mjx +[pypi-badge]: https://badge.fury.io/py/mujoco-mjx.svg +[pypi]: https://pypi.org/project/mujoco-mjx/ + +This package is a re-implementation of the +[MuJoCo physics engine](https://github.com/google-deepmind/mujoco) in +[JAX](https://github.com/google/jax). This library is developed and maintained +by Google DeepMind, and is kept up-to-date with the latest developments in +MuJoCo itself. + +The `mujoco-mjx` package is API-compatible with MuJoCo, but is missing some +features found in MuJoCo. See our +[documentation](https://mujoco.readthedocs.io/en/stable/mjx.html) for more +details concerning feature parity. + +## Installation + +The recommended way to install this package is via [PyPI](https://pypi.org/project/mujoco-mjx/): + +```sh +pip install mujoco-mjx +``` + +## Usage + +Once installed, the package can be imported via `from mujoco import mjx`. Please +consult our [documentation](https://mujoco.readthedocs.io/en/stable/mjx.html) +for further detail on the package's API. + +We recommend going through the tutorial notebook which 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. diff --git a/mjx/cuda_requirements.txt b/mjx/cuda_requirements.txt new file mode 100644 index 00000000..3898472f --- /dev/null +++ b/mjx/cuda_requirements.txt @@ -0,0 +1,17 @@ +-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html + +jax[cuda12_local]==0.4.13; python_version == '3.8' \ + --hash=sha256:03bfe6749dfe647f16f15f6616638adae6c4a7ca7167c75c21961ecfd3a3baaa +jax[cuda12_local]==0.4.18; python_version >= '3.9' \ + --hash=sha256:c3ab72ea2f1c5d8ccf2561e79f6562fb2964629f3e55b3ac1c11c48b64c20336 +jaxlib==0.4.13+cuda12.cudnn89; python_version == '3.8' \ + --hash=sha256:31372f41dc28ecb11a5cc5573ae632e8ee9cfb63788edc8ad8a77e8d3279f569 +jaxlib==0.4.18+cuda12.cudnn89; python_version >= '3.9' \ + --hash=sha256:14f74ff081882ea091c121e355051b35932e39cb7ff7242b88a87f3690f3ca90 \ + --hash=sha256:7c87dc2d68257b02e83c04be88a3c447373ee7077d65f43545bcbda5bfe2231d \ + --hash=sha256:4d16e9c7592e1aaca0b3d28d2c8beba415a2721bb7001f2947728247951a250d \ + --hash=sha256:759c08c69f4a5b1e6b39c3e4eff908a04ce3b2b483bb594ed624407c7d12d110 \ + --hash=sha256:a7a04dbe1851cd50d07691282116aee49a2f0be7838e55b76d7ada86db06be62 \ + --hash=sha256:2bf842db3d58c8c6c52fbc8ed3fabefd7b91a21746cd59d3eaf3522eea229b53 \ + --hash=sha256:35d265ef9bb3835a14580cbaa9402060f117e46056f80e0996405fff3964667a \ + --hash=sha256:0e4352f24d629e912965e6435e140c1b06086243a098651f2d01b75f3738b51c diff --git a/mjx/mujoco/mjx/__init__.py b/mjx/mujoco/mjx/__init__.py new file mode 100644 index 00000000..3ec6312a --- /dev/null +++ b/mjx/mujoco/mjx/__init__.py @@ -0,0 +1,34 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Public API for MJX.""" + +# pylint:disable=g-importing-member +from mujoco.mjx._src.collision_driver import collision +from mujoco.mjx._src.constraint import make_constraint +from mujoco.mjx._src.device import device_get_into +from mujoco.mjx._src.device import device_put +from mujoco.mjx._src.forward import forward +from mujoco.mjx._src.forward import step +from mujoco.mjx._src.io import make_data +from mujoco.mjx._src.passive import passive +from mujoco.mjx._src.smooth import com_pos +from mujoco.mjx._src.smooth import com_vel +from mujoco.mjx._src.smooth import crb +from mujoco.mjx._src.smooth import factor_m +from mujoco.mjx._src.smooth import kinematics +from mujoco.mjx._src.smooth import mul_m +from mujoco.mjx._src.smooth import rne +from mujoco.mjx._src.smooth import transmission +from mujoco.mjx._src.types import * diff --git a/mjx/mujoco/mjx/_src/__init__.py b/mjx/mujoco/mjx/_src/__init__.py new file mode 100644 index 00000000..4ad51d69 --- /dev/null +++ b/mjx/mujoco/mjx/_src/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/mjx/mujoco/mjx/_src/collision_base.py b/mjx/mujoco/mjx/_src/collision_base.py new file mode 100644 index 00000000..c843dceb --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_base.py @@ -0,0 +1,64 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collision base.""" + +import dataclasses +from typing import Dict, List, Optional, Tuple + +import jax +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import GeomType +# pylint: enable=g-importing-member + +Contact = Tuple[jax.Array, jax.Array, jax.Array] + + +@dataclasses.dataclass +class Candidate: + geom1: int + geom2: int + ipair: int + geomp: int # priority geom + dim: int + + +CandidateSet = Dict[ + Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]], + List[Candidate], +] + + +class GeomInfo(PyTreeNode): + """Collision info for a geom.""" + + pos: jax.Array + mat: jax.Array + size: jax.Array + face: Optional[jax.Array] = None + vert: Optional[jax.Array] = None + edge: Optional[jax.Array] = None + facenorm: Optional[jax.Array] = None + + +class SolverParams(PyTreeNode): + """Contact solver params.""" + + friction: jax.Array + solref: jax.Array + solreffriction: jax.Array + solimp: jax.Array + margin: jax.Array + gap: jax.Array diff --git a/mjx/mujoco/mjx/_src/collision_convex.py b/mjx/mujoco/mjx/_src/collision_convex.py new file mode 100644 index 00000000..62d357f9 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_convex.py @@ -0,0 +1,724 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Convex collisions.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import math +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Contact +from mujoco.mjx._src.collision_base import GeomInfo +# pylint: enable=g-importing-member + + +def _closest_segment_point_plane( + a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Gets the closest point between a line segment and a plane. + + Args: + a: first line segment point + b: second line segment point + p0: point on plane + plane_normal: plane normal + + Returns: + closest point between the line segment and the plane + """ + # Parametrize a line segment as S(t) = a + t * (b - a), plug it into the plane + # equation dot(n, S(t)) - d = 0, then solve for t to get the line-plane + # intersection. We then clip t to be in [0, 1] to be on the line segment. + n = plane_normal + d = jp.sum(p0 * n) # shortest distance from origin to plane + denom = jp.sum(n * (b - a)) + t = (d - jp.sum(n * a)) / (denom + 1e-6 * (denom == 0.0)) + t = jp.clip(t, 0, 1) + segment_point = a + t * (b - a) + + return segment_point + + +def _closest_triangle_point( + p0: jax.Array, p1: jax.Array, p2: jax.Array, pt: jax.Array +) -> jax.Array: + """Gets the closest point between a triangle and a point in space. + + Args: + p0: triangle point + p1: triangle point + p2: triangle point + pt: point to test + + Returns: + closest point on the triangle w.r.t point pt + """ + # Parametrize the triangle s.t. a point inside the triangle is + # Q = p0 + u * e0 + v * e1, when 0 <= u <= 1, 0 <= v <= 1, and + # 0 <= u + v <= 1. Let e0 = (p1 - p0) and e1 = (p2 - p0). + # We analytically minimize the distance between the point pt and Q. + e0 = p1 - p0 + e1 = p2 - p0 + a = e0.dot(e0) + b = e0.dot(e1) + c = e1.dot(e1) + d = pt - p0 + # The determinant is 0 only if the angle between e1 and e0 is 0 + # (i.e. the triangle has overlapping lines). + det = a * c - b * b + u = (c * e0.dot(d) - b * e1.dot(d)) / det + v = (-b * e0.dot(d) + a * e1.dot(d)) / det + inside = (0 <= u) & (u <= 1) & (0 <= v) & (v <= 1) & (u + v <= 1) + closest_p = p0 + u * e0 + v * e1 + d0 = (closest_p - pt).dot(closest_p - pt) + + # If the closest point is outside the triangle, it must be on an edge, so we + # check each triangle edge for a closest point to the point pt. + closest_p1, d1 = math.closest_segment_point_and_dist(p0, p1, pt) + closest_p = jp.where((d0 < d1) & inside, closest_p, closest_p1) + min_d = jp.where((d0 < d1) & inside, d0, d1) + + closest_p2, d2 = math.closest_segment_point_and_dist(p1, p2, pt) + closest_p = jp.where(d2 < min_d, closest_p2, closest_p) + min_d = jp.minimum(min_d, d2) + + closest_p3, d3 = math.closest_segment_point_and_dist(p2, p0, pt) + closest_p = jp.where(d3 < min_d, closest_p3, closest_p) + + return closest_p + + +def _closest_segment_triangle_points( + a: jax.Array, + b: jax.Array, + p0: jax.Array, + p1: jax.Array, + p2: jax.Array, + triangle_normal: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Gets the closest points between a line segment and triangle. + + Args: + a: first line segment point + b: second line segment point + p0: triangle point + p1: triangle point + p2: triangle point + triangle_normal: normal of triangle + + Returns: + closest point on the triangle w.r.t the line segment + """ + # The closest triangle point is either on the edge or within the triangle. + # First check triangle edges for the closest point. + # TODO(robotics-simulation): consider vmapping over closest point functions + seg_pt1, tri_pt1 = math.closest_segment_to_segment_points(a, b, p0, p1) + d1 = (seg_pt1 - tri_pt1).dot(seg_pt1 - tri_pt1) + seg_pt2, tri_pt2 = math.closest_segment_to_segment_points(a, b, p1, p2) + d2 = (seg_pt2 - tri_pt2).dot(seg_pt2 - tri_pt2) + seg_pt3, tri_pt3 = math.closest_segment_to_segment_points(a, b, p0, p2) + d3 = (seg_pt3 - tri_pt3).dot(seg_pt3 - tri_pt3) + + # Next, handle the case where the closest triangle point is inside the + # triangle. Either the line segment intersects the triangle or a segment + # endpoint is closest to a point inside the triangle. + seg_pt4 = _closest_segment_point_plane(a, b, p0, triangle_normal) + tri_pt4 = _closest_triangle_point(p0, p1, p2, seg_pt4) + d4 = (seg_pt4 - tri_pt4).dot(seg_pt4 - tri_pt4) + + # Get the point with minimum distance from the line segment point to the + # triangle point. + distance = jp.array([[d1, d2, d3, d4]]) + min_dist = jp.amin(distance) + mask = (distance == min_dist).T + seg_pt = jp.array([seg_pt1, seg_pt2, seg_pt3, seg_pt4]) * mask + tri_pt = jp.array([tri_pt1, tri_pt2, tri_pt3, tri_pt4]) * mask + seg_pt = jp.sum(seg_pt, axis=0) / jp.sum(mask) + tri_pt = jp.sum(tri_pt, axis=0) / jp.sum(mask) + + return seg_pt, tri_pt + + +def _manifold_points( + poly: jax.Array, poly_mask: jax.Array, poly_norm: jax.Array +) -> jax.Array: + """Chooses four points on the polygon with approximately maximal area.""" + dist_mask = jp.where(poly_mask, 0.0, -1e6) + a_idx = jp.argmax(dist_mask) + a = poly[a_idx] + # choose point b furthest from a + b_idx = (((a - poly) ** 2).sum(axis=1) + dist_mask).argmax() + b = poly[b_idx] + # choose point c furthest along the axis orthogonal to (a-b) + ab = jp.cross(poly_norm, a - b) + ap = a - poly + c_idx = (jp.abs(ap.dot(ab)) + dist_mask).argmax() + c = poly[c_idx] + # choose point d furthest from the other two triangle edges + ac = jp.cross(poly_norm, a - c) + bc = jp.cross(poly_norm, b - c) + bp = b - poly + dist_bp = jp.abs(bp.dot(bc)) + dist_mask + dist_ap = jp.abs(ap.dot(ac)) + dist_mask + d_idx = jp.concatenate([dist_bp, dist_ap]).argmax() % poly.shape[0] + return jp.array([a_idx, b_idx, c_idx, d_idx]) + + +def _project_pt_onto_plane( + pt: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Projects a point onto a plane along the plane normal.""" + dist = (pt - plane_pt).dot(plane_normal) + return pt - dist * plane_normal + + +def _project_poly_onto_plane( + poly: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array +) -> jax.Array: + """Projects a polygon onto a plane using the plane normal.""" + return jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])( + poly, plane_pt, math.normalize(plane_normal) + ) + + +def _project_poly_onto_poly_plane( + poly1: jax.Array, norm1: jax.Array, poly2: jax.Array, norm2: jax.Array +) -> jax.Array: + """Projects poly1 onto the poly2 plane along poly1's normal.""" + d = poly2[0].dot(norm2) + denom = norm1.dot(norm2) + t = (d - poly1.dot(norm2)) / (denom + 1e-6 * (denom == 0.0)) + new_poly = poly1 + t.reshape(-1, 1) * norm1 + return new_poly + + +def _point_in_front_of_plane( + plane_pt: jax.Array, plane_normal: jax.Array, pt: jax.Array +) -> jax.Array: + """Checks if a point is strictly in front of a plane.""" + return (pt - plane_pt).dot(plane_normal) > 1e-6 + + +def _clip_edge_to_planes( + edge_p0: jax.Array, + edge_p1: jax.Array, + plane_pts: jax.Array, + plane_normals: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Clips an edge against side planes. + + We return two clipped points, and a mask to include the new edge or not. + + Args: + edge_p0: the first point on the edge + edge_p1: the second point on the edge + plane_pts: side plane points + plane_normals: side plane normals + + Returns: + new_ps: new edge points that are clipped against side planes + mask: a boolean mask, True if an edge point is a valid clipped point and + False otherwise + """ + p0, p1 = edge_p0, edge_p1 + p0_in_front = jax.vmap(jp.dot)(p0 - plane_pts, plane_normals) > 1e-6 + p1_in_front = jax.vmap(jp.dot)(p1 - plane_pts, plane_normals) > 1e-6 + + # Get candidate clipped points along line segment (p0, p1) by clipping against + # all clipping planes. + candidate_clipped_ps = jax.vmap( + _closest_segment_point_plane, in_axes=[None, None, 0, 0] + )(p0, p1, plane_pts, plane_normals) + + def clip_edge_point(p0, p1, p0_in_front, clipped_ps): + @jax.vmap + def choose_edge_point(in_front, clipped_p): + return jp.where(in_front, clipped_p, p0) + + # Pick the clipped point if p0 is in front of the clipping plane. Otherwise + # keep p0 as the edge point. + new_edge_ps = choose_edge_point(p0_in_front, clipped_ps) + + # Pick the clipped point that is most along the edge direction. + # This degenerates to picking the original point p0 if p0 is *not* in front + # of any clipping planes. + dists = jp.dot(new_edge_ps - p0, p1 - p0) + new_edge_p = new_edge_ps[jp.argmax(dists)] + return new_edge_p + + # Clip each edge point. + new_p0 = clip_edge_point(p0, p1, p0_in_front, candidate_clipped_ps) + new_p1 = clip_edge_point(p1, p0, p1_in_front, candidate_clipped_ps) + clipped_pts = jp.array([new_p0, new_p1]) + + # Keep the original points if both points are in front of any of the clipping + # planes, rather than creating a new clipped edge. If the entire subject edge + # is in front of any clipping plane, we need to grab an edge from the clipping + # polygon instead. + both_in_front = p0_in_front & p1_in_front + mask = ~jp.any(both_in_front) + new_ps = jp.where(mask, clipped_pts, jp.array([p0, p1])) + # Mask out crossing clipped edge points. + mask = jp.where((p0 - p1).dot(new_ps[0] - new_ps[1]) < 0, False, mask) + return new_ps, jp.array([mask, mask]) + + +def _clip( + clipping_poly: jax.Array, + subject_poly: jax.Array, + clipping_normal: jax.Array, + subject_normal: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Clips a subject polygon against a clipping polygon. + + A parallelized clipping algorithm for convex polygons. The result is a set of + vertices on the clipped subject polygon in the subject polygon plane. + + Args: + clipping_poly: the polygon that we use to clip the subject polygon against + subject_poly: the polygon that gets clipped + clipping_normal: normal of the clipping polygon + subject_normal: normal of the subject polygon + + Returns: + clipped_pts: points on the clipped polygon + mask: True if a point is in the clipping polygon, False otherwise + """ + # Get clipping edge points, edge planes, and edge normals. + clipping_p0 = jp.roll(clipping_poly, 1, axis=0) + clipping_plane_pts = clipping_p0 + clipping_p1 = clipping_poly + clipping_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])( + clipping_p1 - clipping_p0, + clipping_normal, + ) + + # Get subject edge points, edge planes, and edge normals. + subject_edge_p0 = jp.roll(subject_poly, 1, axis=0) + subject_plane_pts = subject_edge_p0 + subject_edge_p1 = subject_poly + subject_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])( + subject_edge_p1 - subject_edge_p0, + subject_normal, + ) + + # Clip all edges of the subject poly against clipping side planes. + clipped_edges0, masks0 = jax.vmap( + _clip_edge_to_planes, in_axes=[0, 0, None, None] + )( + subject_edge_p0, + subject_edge_p1, + clipping_plane_pts, + clipping_plane_normals, + ) + + # Project the clipping poly onto the subject plane. + clipping_p0_s = _project_poly_onto_poly_plane( + clipping_p0, clipping_normal, subject_poly, subject_normal + ) + clipping_p1_s = _project_poly_onto_poly_plane( + clipping_p1, clipping_normal, subject_poly, subject_normal + ) + + # Clip all edges of the clipping poly against subject planes. + clipped_edges1, masks1 = jax.vmap( + _clip_edge_to_planes, in_axes=[0, 0, None, None] + )(clipping_p0_s, clipping_p1_s, subject_plane_pts, subject_plane_normals) + + # Merge the points and reshape. + clipped_edges = jp.concatenate([clipped_edges0, clipped_edges1]) + masks = jp.concatenate([masks0, masks1]) + clipped_points = clipped_edges.reshape((-1, 3)) + mask = masks.reshape(-1) + + return clipped_points, mask + + +def _create_contact_manifold( + clipping_poly: jax.Array, + subject_poly: jax.Array, + clipping_norm: jax.Array, + subject_norm: jax.Array, + sep_axis: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Creates a contact manifold between two convex polygons. + + The polygon faces are expected to have a counter clockwise winding order so + that clipping plane normals point away from the polygon center. + + Args: + clipping_poly: the reference polygon to clip the contact against. + subject_poly: the subject polygon to clip contacts onto. + clipping_norm: the clipping polygon normal. + subject_norm: the subject polygon normal. + sep_axis: the separating axis + + Returns: + tuple of dist, pos, and normal + """ + # Clip the subject (incident) face onto the clipping (reference) face. + # The incident points are clipped points on the subject polygon. + poly_incident, mask = _clip( + clipping_poly, subject_poly, clipping_norm, subject_norm + ) + # The reference points are clipped points on the clipping polygon. + poly_ref = _project_poly_onto_plane( + poly_incident, clipping_poly[0], clipping_norm + ) + behind_clipping_plane = _point_in_front_of_plane( + clipping_poly[0], -clipping_norm, poly_incident + ) + mask = mask & behind_clipping_plane + + # Choose four contact points. + best = _manifold_points(poly_ref, mask, clipping_norm) + contact_pts = jp.take(poly_ref, best, axis=0) + mask_pts = jp.take(mask, best, axis=0) + penetration_dir = jp.take(poly_incident, best, axis=0) - contact_pts + penetration = penetration_dir.dot(-clipping_norm) + + dist = jp.where(mask_pts, -penetration, jp.ones_like(penetration)) + pos = contact_pts + normal = -jp.stack([sep_axis] * 4, 0) + return dist, pos, normal + + +def _sat_hull_hull( + faces_a: jax.Array, + faces_b: jax.Array, + vertices_a: jax.Array, + vertices_b: jax.Array, + normals_a: jax.Array, + normals_b: jax.Array, + unique_edges_a: jax.Array, + unique_edges_b: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Runs the Separating Axis Test for a pair of hulls. + + Given two convex hulls, the Separating Axis Test finds a separating axis + between all edge pairs and face pairs. Edge pairs create a single contact + point and face pairs create a contact manifold (up to four contact points). + We return both the edge and face contacts. Valid contacts can be checked with + dist < 0. Resulting edge contacts should be preferred over face contacts. + + Args: + faces_a: An ndarray of hull A's polygon faces. + faces_b: An ndarray of hull B's polygon faces. + vertices_a: Vertices for hull A. + vertices_b: Vertices for hull B. + normals_a: Normal vectors for hull A's polygon faces. + normals_b: Normal vectors for hull B's polygon faces. + unique_edges_a: Unique edges for hull A. + unique_edges_b: Unique edges for hull B. + + Returns: + tuple of dist, pos, and normal + """ + # get the separating axes + edge_dir_a = unique_edges_a[:, 0] - unique_edges_a[:, 1] + edge_dir_b = unique_edges_b[:, 0] - unique_edges_b[:, 1] + edge_dir_a_r = jp.tile(edge_dir_a, reps=(unique_edges_b.shape[0], 1)) + edge_dir_b_r = jp.repeat(edge_dir_b, repeats=unique_edges_a.shape[0], axis=0) + edge_edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r) + edge_edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))( + edge_edge_axes + ) + + axes = jp.concatenate([normals_a, normals_b, edge_edge_axes]) + + # for each separating axis, get the support + @jax.vmap + def get_support(axis): + support_a = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_a) + support_b = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_b) + dist1 = support_a.max() - support_b.min() + dist2 = support_b.max() - support_a.min() + sign = jp.where(dist1 > dist2, -1, 1) + dist = jp.minimum(dist1, dist2) + dist = jp.where(~jp.all(axis == 0.0), dist, 1e6) # degenerate axis + return dist, sign + + support, sign = get_support(axes) + + # choose the best separating axis + best_idx = jp.argmin(support) + best_sign = sign[best_idx] + best_axis = axes[best_idx] + is_edge_contact = best_idx >= (normals_a.shape[0] + normals_b.shape[0]) + + # get the (reference) face most aligned with the separating axis + dist_a = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_a) + dist_b = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_b) + a_max = dist_a.argmax() + b_max = dist_b.argmax() + a_min = dist_a.argmin() + b_min = dist_b.argmin() + + ref_face = jp.where(best_sign > 0, faces_a[a_max], faces_b[b_max]) + ref_face_norm = jp.where(best_sign > 0, normals_a[a_max], normals_b[b_max]) + incident_face = jp.where(best_sign > 0, faces_b[b_min], faces_a[a_min]) + incident_face_norm = jp.where( + best_sign > 0, normals_b[b_min], normals_a[a_min] + ) + + dist, pos, normal = _create_contact_manifold( + ref_face, + incident_face, + ref_face_norm, + incident_face_norm, + -best_sign * best_axis, + ) + + # For edge contacts, we use the clipped face point, mainly for performance + # reasons. For small penetration, the clipped face point is roughly the edge + # contact point. + idx = dist.argmin() + dist = jp.where( + is_edge_contact, + jp.array([dist[idx], 1, 1, 1]), + dist, + ) + pos = jp.where(is_edge_contact, jp.tile(pos[idx], (4, 1)), pos) + + return dist, pos, normal + + +def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contacts between a plane and a convex object.""" + vert = convex.vert + + # get points in the convex frame + plane_pos = convex.mat.T @ (plane.pos - convex.pos) + n = convex.mat.T @ plane.mat[:, 2] + support = (plane_pos - vert) @ n + idx = _manifold_points(vert, support > 0, n) + pos = vert[idx] + + # convert to world frame + pos = convex.pos + pos @ convex.mat.T + n = plane.mat[:, 2] + + frame = jp.stack([math.make_frame(n)] * 4, axis=0) + unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1 + dist = jp.where(unique, -support[idx], 1) + return dist, pos, frame + + +def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contact between a sphere and a convex object.""" + faces = jp.take(convex.vert, convex.face, axis=0) + normals = convex.facenorm + + # Put sphere in convex frame. + sphere_pos = convex.mat.T @ (sphere.pos - convex.pos) + + # Get support from face normals. + @jax.vmap + def get_support(faces, normal): + pos = sphere_pos - normal * sphere.size[0] + return jp.dot(pos - faces[0], normal) + + support = get_support(faces, normals) + + # Pick the face with minimal penetration as long as it has support. + support = jp.where(support >= 0, -1e12, support) + best_idx = support.argmax() + face = faces[best_idx] + normal = normals[best_idx] + + # Get closest point between the polygon face and the sphere center point. + # Project the sphere center point onto poly plane. If it's inside polygon + # edge normals, then we're done. + pt = _project_pt_onto_plane(sphere_pos, face[0], normal) + edge_p0 = jp.roll(face, 1, axis=0) + edge_p1 = face + edge_normals = jax.vmap(jp.cross, in_axes=[0, None])( + edge_p1 - edge_p0, + normal, + ) + edge_dist = jax.vmap( + lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm) + )(edge_p0, edge_normals) + inside = jp.all(edge_dist <= 0) # lte to handle degenerate edges + + # If the point is outside edge normals, project onto the closest edge plane + # that the point is in front of. + degenerate_edge = jp.all(edge_normals == 0, axis=1) + behind = edge_dist < 0.0 + edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist) + idx = edge_dist.argmin() + edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt) + + pt = jp.where(inside, pt, edge_pt) + + # Get the normal, dist, and contact position. + n, d = math.normalize_with_norm(pt - sphere_pos) + spt = sphere_pos + n * sphere.size[0] + dist = d - sphere.size[0] + pos = (pt + spt) * 0.5 + + # Go back to world frame. + n = convex.mat @ n + pos = convex.mat @ pos + convex.pos + + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact: + """Calculates contacts between a capsule and a convex object.""" + # Get convex transformed normals, faces, and vertices. + faces = jp.take(convex.vert, convex.face, axis=0) + normals = convex.facenorm + + # Put capsule in convex frame. + cap_pos = convex.mat.T @ (cap.pos - convex.pos) + axis, length = cap.mat[:, 2], cap.size[1] + axis = convex.mat.T @ axis + seg = axis * length + cap_pts = jp.array([ + cap_pos - seg, + cap_pos + seg, + ]) + + # Get support from face normals. + @jax.vmap + def get_support(face, normal): + pts = cap_pts - normal * cap.size[0] + sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts) + return sup.min() + + support = get_support(faces, normals) + has_support = jp.all(support < 0) + + # Pick the face with minimal penetration as long as it has support. + support = jp.where(support >= 0, -1e12, support) + best_idx = support.argmax() + face = faces[best_idx] + normal = normals[best_idx] + + # Clip the edge against side planes and create two contact points against the + # face. + edge_p0 = jp.roll(face, 1, axis=0) + edge_p1 = face + edge_normals = jax.vmap(jp.cross, in_axes=[0, None])( + edge_p1 - edge_p0, + normal, + ) + cap_pts_clipped, mask = _clip_edge_to_planes( + cap_pts[0], cap_pts[1], edge_p0, edge_normals + ) + cap_pts_clipped = cap_pts_clipped - normal * cap.size[0] + face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])( + cap_pts_clipped, face[0], normal + ) + # Create variables for the face contact. + pos = (cap_pts_clipped + face_pts) * 0.5 + norm = jp.stack([normal] * 2, 0) + penetration = jp.where( + mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1 + ) + + # Get a potential edge contact. + edge_closest, cap_closest = jax.vmap( + math.closest_segment_to_segment_points, in_axes=[0, 0, None, None] + )(edge_p0, edge_p1, cap_pts[0], cap_pts[1]) + e_idx = ((edge_closest - cap_closest) ** 2).sum(axis=1).argmin() + cap_closest_pt, edge_closest_pt = cap_closest[e_idx], edge_closest[e_idx] + edge_axis = cap_closest_pt - edge_closest_pt + edge_axis, edge_dist = math.normalize_with_norm(edge_axis) + edge_pos = ( + edge_closest_pt + (cap_closest_pt - edge_axis * cap.size[0]) + ) * 0.5 + edge_norm = edge_axis + edge_penetration = cap.size[0] - edge_dist + has_edge_contact = edge_penetration > 0 + + # Get the contact info. + pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos) + n = -jp.where(has_edge_contact, norm.at[0].set(edge_norm), norm) + + # Go back to world frame. + pos = convex.pos + pos @ convex.mat.T + n = n @ convex.mat.T + + dist = -jp.where( + has_edge_contact, penetration.at[0].set(edge_penetration), penetration + ) + frame = jax.vmap(math.make_frame)(n) + return dist, pos, frame + + +def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact: + """Calculates contacts between two convex objects.""" + if c1.face is None or c2.face is None or c1.vert is None or c2.vert is None: + raise AssertionError('Mesh info missing.') + # pad face vertices so that we can broadcast between geom1 and geom2 + s1, s2 = c1.face.shape[-1], c2.face.shape[-1] + if s1 < s2: + face = jp.pad(c1.face, ((0, 0), (0, s2 - s1)), 'edge') + c1 = c1.replace(face=face) + elif s2 < s1: + face = jp.pad(c2.face, ((0, 0), (0, s1 - s2)), 'edge') + c2 = c2.replace(face=face) + + # ensure that the first object has fewer verts + swapped = c1.vert.shape[0] > c2.vert.shape[0] + if swapped: + c1, c2 = c2, c1 + + faces1 = jp.take(c1.vert, c1.face, axis=0) + faces2 = jp.take(c2.vert, c2.face, axis=0) + + to_local_pos = c2.mat.T @ (c1.pos - c2.pos) + to_local_mat = c2.mat.T @ c1.mat + + faces1 = to_local_pos + faces1 @ to_local_mat.T + normals1 = c1.facenorm @ to_local_mat.T + normals2 = c2.facenorm + + vertices1 = to_local_pos + c1.vert @ to_local_mat.T + vertices2 = c2.vert + + unique_edges1 = jp.take(vertices1, c1.edge, axis=0) + unique_edges2 = jp.take(vertices2, c2.edge, axis=0) + + dist, pos, normal = _sat_hull_hull( + faces1, + faces2, + vertices1, + vertices2, + normals1, + normals2, + unique_edges1, + unique_edges2, + ) + + # Go back to world frame. + pos = c2.pos + pos @ c2.mat.T + normal = normal @ c2.mat.T + normal = -normal if swapped else normal + + frame = jax.vmap(math.make_frame)(normal) + return dist, pos, frame + + +# store ncon as function attributes +plane_convex.ncon = 4 +sphere_convex.ncon = 1 +capsule_convex.ncon = 2 +convex_convex.ncon = 4 diff --git a/mjx/mujoco/mjx/_src/collision_driver.py b/mjx/mujoco/mjx/_src/collision_driver.py new file mode 100644 index 00000000..b5f58130 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_driver.py @@ -0,0 +1,373 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collide geometries.""" + +from typing import Callable, Dict, Optional, Sequence, Tuple, Union + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_base +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Candidate +from mujoco.mjx._src.collision_base import CandidateSet +from mujoco.mjx._src.collision_base import GeomInfo +from mujoco.mjx._src.collision_base import SolverParams +from mujoco.mjx._src.collision_convex import capsule_convex +from mujoco.mjx._src.collision_convex import convex_convex +from mujoco.mjx._src.collision_convex import plane_convex +from mujoco.mjx._src.collision_convex import sphere_convex +from mujoco.mjx._src.collision_primitive import capsule_capsule +from mujoco.mjx._src.collision_primitive import plane_capsule +from mujoco.mjx._src.collision_primitive import plane_sphere +from mujoco.mjx._src.collision_primitive import sphere_capsule +from mujoco.mjx._src.collision_primitive import sphere_sphere +from mujoco.mjx._src.types import Contact +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +# pair-wise collision functions +_COLLISION_FUNC = { + (GeomType.PLANE, GeomType.SPHERE): plane_sphere, + (GeomType.PLANE, GeomType.CAPSULE): plane_capsule, + (GeomType.PLANE, GeomType.BOX): plane_convex, + (GeomType.PLANE, GeomType.MESH): plane_convex, + (GeomType.SPHERE, GeomType.SPHERE): sphere_sphere, + (GeomType.SPHERE, GeomType.CAPSULE): sphere_capsule, + (GeomType.SPHERE, GeomType.BOX): sphere_convex, + (GeomType.SPHERE, GeomType.MESH): sphere_convex, + (GeomType.CAPSULE, GeomType.CAPSULE): capsule_capsule, + (GeomType.CAPSULE, GeomType.BOX): capsule_convex, + (GeomType.CAPSULE, GeomType.MESH): capsule_convex, + (GeomType.BOX, GeomType.BOX): convex_convex, + (GeomType.BOX, GeomType.MESH): convex_convex, + (GeomType.MESH, GeomType.MESH): convex_convex, +} + + +def get_collision_fn( + key: Tuple[Union[GeomType, mujoco.mjtGeom], Union[GeomType, mujoco.mjtGeom]] +) -> Optional[Callable[[GeomInfo, GeomInfo], collision_base.Contact]]: + """Returns a collision function given a pair of geom types.""" + return _COLLISION_FUNC.get(key, None) + + +def _add_candidate( + result: CandidateSet, + m: Union[Model, mujoco.MjModel], + g1: int, + g2: int, + ipair: int = -1, +): + """Adds a candidate to test for collision.""" + t1, t2 = m.geom_type[g1], m.geom_type[g2] + if t1 > t2: + t1, t2, g1, g2 = t2, t1, g2, g1 + + def mesh_key(i): + convex_data = [[None] * m.ngeom] * 3 + if isinstance(m, Model): + convex_data = [m.geom_convex_face, m.geom_convex_vert, m.geom_convex_edge] + key = tuple((-1,) if v[i] is None else v[i].shape for v in convex_data) + return key + + k1, k2 = mesh_key(g1), mesh_key(g2) + + candidates = {(c.geom1, c.geom2) for c in result.get((t1, t2, k1, k2), [])} + if (g1, g2) in candidates: + return + + if ipair > -1: + candidate = Candidate(g1, g2, ipair, -1, m.pair_dim[ipair]) + elif m.geom_priority[g1] != m.geom_priority[g2]: + gp = g1 if m.geom_priority[g1] > m.geom_priority[g2] else g2 + candidate = Candidate(g1, g2, -1, gp, m.geom_condim[gp]) + else: + dim = max(m.geom_condim[g1], m.geom_condim[g2]) + candidate = Candidate(g1, g2, -1, -1, dim) + + result.setdefault((t1, t2, k1, k2), []).append(candidate) + + +def _pair_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params for pair geoms.""" + ipair = jp.array([c.ipair for c in candidates]) + friction = jp.clip(m.pair_friction[ipair], a_min=mujoco.mjMINMU) + solref = m.pair_solref[ipair] + solreffriction = m.pair_solreffriction[ipair] + solimp = m.pair_solimp[ipair] + margin = m.pair_margin[ipair] + gap = m.pair_gap[ipair] + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _priority_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params from priority geoms.""" + geomp = jp.array([c.geomp for c in candidates]) + friction = m.geom_friction[geomp][:, jp.array([0, 0, 1, 2, 2])] + solref = m.geom_solref[geomp] + solreffriction = jp.zeros(geomp.shape + (mujoco.mjNREF,)) + solimp = m.geom_solimp[geomp] + g = jp.array([(c.geom1, c.geom2) for c in candidates]) + margin = jp.amax(m.geom_margin[g.T], axis=0) + gap = jp.amax(m.geom_gap[g.T], axis=0) + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _dynamic_params( + m: Model, + candidates: Sequence[Candidate], +) -> SolverParams: + """Gets solver params for dynamic geoms.""" + g1 = jp.array([c.geom1 for c in candidates]) + g2 = jp.array([c.geom2 for c in candidates]) + + friction = jp.maximum(m.geom_friction[g1], m.geom_friction[g2]) + # copy friction terms for the full geom pair + friction = friction[:, jp.array([0, 0, 1, 2, 2])] + + minval = jp.array(mujoco.mjMINVAL) + solmix1, solmix2 = m.geom_solmix[g1], m.geom_solmix[g2] + mix = solmix1 / (solmix1 + solmix2) + mix = jp.where((solmix1 < minval) & (solmix2 < minval), 0.5, mix) + mix = jp.where((solmix1 < minval) & (solmix2 >= minval), 0.0, mix) + mix_fn = jax.vmap(lambda a, b, m: m * a + (1 - m) * b) + + solref1, solref2 = m.geom_solref[g1], m.geom_solref[g2] + solref = jp.minimum(solref1, solref2) + s_mix = mix_fn(solref1, solref2, mix) + solref = jp.where((solref1[0] > 0) & (solref2[0] > 0), s_mix, solref) + solreffriction = jp.zeros(g1.shape + (mujoco.mjNREF,)) + solimp = mix_fn(m.geom_solimp[g1], m.geom_solimp[g2], mix) + margin = jp.maximum(m.geom_margin[g1], m.geom_margin[g2]) + gap = jp.maximum(m.geom_gap[g1], m.geom_gap[g2]) + + return SolverParams(friction, solref, solreffriction, solimp, margin, gap) + + +def _pair_info( + m: Model, d: Data, geom1: Sequence[int], geom2: Sequence[int] +) -> Tuple[GeomInfo, GeomInfo, Sequence[Dict[str, Optional[int]]]]: + """Returns geom pair info for calculating collision.""" + g1, g2 = jp.array(geom1), jp.array(geom2) + info1 = GeomInfo( + d.geom_xpos[g1], + d.geom_xmat[g1], + m.geom_size[g1], + ) + info2 = GeomInfo( + d.geom_xpos[g2], + d.geom_xmat[g2], + m.geom_size[g2], + ) + in_axes1 = in_axes2 = jax.tree_map(lambda x: 0, info1) + if m.geom_convex_face[geom1[0]] is not None: + info1 = info1.replace( + face=jp.stack([m.geom_convex_face[i] for i in geom1]), + vert=jp.stack([m.geom_convex_vert[i] for i in geom1]), + edge=jp.stack([m.geom_convex_edge[i] for i in geom1]), + facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom1]), + ) + in_axes1 = in_axes1.replace(face=0, vert=0, edge=0, facenorm=0) + if m.geom_convex_face[geom2[0]] is not None: + info2 = info2.replace( + face=jp.stack([m.geom_convex_face[i] for i in geom2]), + vert=jp.stack([m.geom_convex_vert[i] for i in geom2]), + edge=jp.stack([m.geom_convex_edge[i] for i in geom2]), + facenorm=jp.stack([m.geom_convex_facenormal[i] for i in geom2]), + ) + in_axes2 = in_axes2.replace(face=0, vert=0, edge=0, facenorm=0) + return info1, info2, [in_axes1, in_axes2] + + +def _body_pair_filter( + m: Union[Model, mujoco.MjModel], b1: int, b2: int +) -> bool: + """Filters body pairs for collision.""" + dsbl_filterparent = m.opt.disableflags & DisableBit.FILTERPARENT + weld1 = m.body_weldid[b1] + weld2 = m.body_weldid[b2] + parent_weld1 = m.body_weldid[m.body_parentid[weld1]] + parent_weld2 = m.body_weldid[m.body_parentid[weld2]] + + if weld1 == weld2: + # filter out self-collisions + return True + + if ( + not dsbl_filterparent + and weld1 != 0 + and weld2 != 0 + and (weld1 == parent_weld2 or weld2 == parent_weld1) + ): + # filter out parent-child collisions + return True + + return False + + +def _collide_geoms( + m: Model, + d: Data, + geom_types: Tuple[GeomType, GeomType], + candidates: Sequence[Candidate], +) -> Contact: + """Collides a geom pair.""" + fn = get_collision_fn(geom_types) + if not fn: + return Contact.zero() + + # group sol params by different candidate types + typ_cands = {} + for c in candidates: + typ = (c.ipair > -1, c.geomp > -1) + typ_cands.setdefault(typ, []).append(c) + + geom1, geom2, params = [], [], [] + for (pair, priority), candidates in typ_cands.items(): + geom1.extend([c.geom1 for c in candidates]) + geom2.extend([c.geom2 for c in candidates]) + if pair: + params.append(_pair_params(m, candidates)) + elif priority: + params.append(_priority_params(m, candidates)) + else: + params.append(_dynamic_params(m, candidates)) + + # call contact function + g1, g2, in_axes = _pair_info(m, d, geom1, geom2) + res = jax.vmap(fn, in_axes=in_axes)(g1, g2) + dist, pos, frame = jax.tree_map(jp.concatenate, res) + + params = jax.tree_map(lambda *x: jp.concatenate(x), *params) + geom1, geom2 = jp.array(geom1), jp.array(geom2) + # repeat params by the number of contacts per geom pair + n_repeat = dist.shape[-1] // geom1.shape[0] + geom1, geom2, params = jax.tree_map( + lambda x: jp.repeat(x, n_repeat, axis=0), + (geom1, geom2, params), + ) + + con = Contact( + dist=dist, + pos=pos, + frame=frame, + includemargin=params.margin - params.gap, + friction=params.friction, + solref=params.solref, + solreffriction=params.solreffriction, + solimp=params.solimp, + geom1=geom1, + geom2=geom2, + dim=np.array([]), + efc_address=np.array([]), + ) + return con + + +def _max_contact_points(m: Model) -> int: + """Returns the maximum number of contact points when set as a numeric.""" + for i in range(m.nnumeric): + name = m.names[m.name_numericadr[i] :].decode('utf-8').split('\x00', 1)[0] + if name == 'max_contact_points': + return int(m.numeric_data[m.numeric_adr[i]]) + + return -1 + + +def collision_candidates(m: Union[Model, mujoco.MjModel]) -> CandidateSet: + """Returns candidates for collision checking.""" + candidate_set = {} + + for ipair in range(m.npair): + g1, g2 = m.pair_geom1[ipair], m.pair_geom2[ipair] + _add_candidate(candidate_set, m, g1, g2, ipair) + + body_pairs = [] + exclude_signature = set(m.exclude_signature) + for b1 in range(m.nbody): + for b2 in range(b1, m.nbody): + signature = (b1 << 16) + (b2) + if signature in exclude_signature: + continue + if _body_pair_filter(m, b1, b2): + continue + body_pairs.append((b1, b2)) + + for b1, b2 in body_pairs: + start1 = m.body_geomadr[b1] + end1 = m.body_geomadr[b1] + m.body_geomnum[b1] + for g1 in range(start1, end1): + start2 = m.body_geomadr[b2] + end2 = m.body_geomadr[b2] + m.body_geomnum[b2] + for g2 in range(start2, end2): + mask = m.geom_contype[g1] & m.geom_conaffinity[g2] + mask |= m.geom_contype[g2] & m.geom_conaffinity[g1] + if mask != 0: + _add_candidate(candidate_set, m, g1, g2) + + return candidate_set + + +def ncon(m: Model) -> int: + """Returns the number of contacts computed in MJX given a model.""" + candidates = collision_candidates(m) + max_count = _max_contact_points(m) + count = sum([ + len(v) * get_collision_fn(k[0:2]).ncon for k, v in candidates.items() # pytype: disable=attribute-error + ]) + return min(max_count, count) if max_count > -1 else count + + +def collision(m: Model, d: Data) -> Data: + """Collides geometries.""" + candidate_set = collision_candidates(m) + + contacts = [] + for key, candidates in candidate_set.items(): + geom_types = key[0:2] + contacts.append(_collide_geoms(m, d, geom_types, candidates)) + + if not contacts: + return d.replace(contact=Contact.zero(), ncon=0) + + contact = jax.tree_map(lambda *x: jp.concatenate(x), *contacts) + + max_contact_points = _max_contact_points(m) + if max_contact_points > -1 and contact.dist.shape[0] > max_contact_points: + # get top-k contacts + _, idx = jax.lax.top_k(-contact.dist, k=max_contact_points) + contact = jax.tree_map(lambda x, idx=idx: jp.take(x, idx, axis=0), contact) + + ncon_ = contact.dist.shape[0] + ns = d.ne + d.nf + d.nl + # TODO(robotics-simulation): add support for other friction dimensions + contact = contact.replace(efc_address=np.arange(ns, ns + d.ncon * 4, 4)) + contact = contact.replace(dim=3 * np.ones(ncon_, dtype=np.int32)) + + return d.replace(contact=contact, ncon=ncon_) diff --git a/mjx/mujoco/mjx/_src/collision_driver_test.py b/mjx/mujoco/mjx/_src/collision_driver_test.py new file mode 100644 index 00000000..71b7b553 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_driver_test.py @@ -0,0 +1,499 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests the collision driver.""" + +import dataclasses +from typing import Dict, Optional, Tuple + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import jax.numpy as jp +import mujoco +from mujoco import mjx +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Contact +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model +# pylint: emable=g-importing-member +import numpy as np + + +def _assert_attr_eq(mjx_d, mj_d, attr, name, atol): + if attr == 'efc_address': + # we do not test efc_address since it gets set in constraint logic + return + err_msg = f'mismatch: {attr} in run: {name}' + mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr) + if attr == 'frame': + mj_d = mj_d.reshape((-1, 3, 3)) + if mjx_d.shape != mj_d.shape: + raise AssertionError(f'{attr} shape mismatch: {mjx_d.shape}, {mj_d.shape}') + np.testing.assert_allclose(mjx_d, mj_d, err_msg=err_msg, atol=atol) + + +def _collide( + mjcf: str, assets: Optional[Dict[str, str]] = None +) -> Tuple[mujoco.MjModel, mujoco.MjData, Model, Data]: + m = mujoco.MjModel.from_xml_string(mjcf, assets or {}) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + return d, dx + + +class SphereCollisionTest(parameterized.TestCase): + _SPHERE_PLANE = """ + + + + + + + + + + """ + + _SPHERE_SPHERE = """ + + + + + + + + + + + + + """ + + _SPHERE_CAP = """ + + + + + + + + + + + + + """ + + @parameterized.parameters( + ('sphere_plane', _SPHERE_PLANE), + ('sphere_sphere', _SPHERE_SPHERE), + ('sphere_cap', _SPHERE_CAP), + ) + def test_sphere(self, name, mjcf): + d, dx = _collide(mjcf) + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5) + + _SPHERE_CONVEX = """ + + + + + + + + + + + + + """ + + def test_sphere_convex(self): + d, dx = _collide(self._SPHERE_CONVEX) + + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, 'sphere_convex', 1e-4) + + +class CapsuleCollisionTest(parameterized.TestCase): + _CAP_PLANE = """ + + + + + + + + + + """ + + _CAP_CAP = """ + + + + + + + + + + + + + """ + + @parameterized.parameters( + ('capsule_plane', _CAP_PLANE), + ('capsule_capsule', _CAP_CAP), + ) + def test_capsule(self, name, mjcf): + d, dx = _collide(mjcf) + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4) + + _PARALLEL_CAP = """ + + + + + + + + + + + + + """ + + def test_parallel_capsules(self): + """Tests that two parallel capsules are colliding at the midpoint.""" + _, dx = _collide(self._PARALLEL_CAP) + + np.testing.assert_allclose(dx.contact.dist, -0.05) + np.testing.assert_allclose( + dx.contact.pos[0], + np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]), + atol=1e-5, + ) + np.testing.assert_allclose( + dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5 + ) + + _CAP_BOX = """ + + + + + + + + + + + + + """ + + def test_capsule_convex(self): + """Tests a capsule-convex collision for a face contact.""" + d, dx = _collide(self._CAP_BOX) + + for field in dataclasses.fields(Contact): + _assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4) + + _CAP_EDGE_BOX = """ + + + + + + + + + + + + + """ + + def test_capsule_convex_edge(self): + """Tests a capsule-convex collision for an edge contact.""" + d, dx = _collide(self._CAP_EDGE_BOX) + + c = dx.contact + self.assertEqual(c.pos.shape[0], 2) + self.assertGreater(c.dist[1], 0) + # extract the contact point with penetration + c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = c.replace(dim=c.dim[np.array([0])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'capsule_convex_edge', 1e-4) + + +class ConvexTest(absltest.TestCase): + """Tests the convex contact functions.""" + + _BOX_PLANE = """ + + + + + + + + + + """ + + def test_box_plane(self): + """Tests box collision with a plane.""" + d, dx = _collide(self._BOX_PLANE) + + np.testing.assert_array_less(dx.contact.dist[:2], 0) + np.testing.assert_array_less(-dx.contact.dist[2:], 0) + # extract the contact points with penetration + c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact) + c = c.replace(dim=c.dim[np.array([0, 1])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'box_plane', 1e-2) + + _BOX_BOX = """ + + + + + + + + + + + + + """ + + def test_box_box(self): + """Tests a face contact for a box-box collision.""" + d, dx = _collide(self._BOX_BOX) + c = dx.contact + + self.assertEqual(c.pos.shape[0], 4) + 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 = """ + + + + + + + + + + + + + """ + + def test_box_box_edge(self): + """Tests an edge contact for a box-box collision.""" + d, dx = _collide(self._BOX_BOX_EDGE) + + # Only one contact point. + np.testing.assert_array_less(dx.contact.dist[:1], 0) + np.testing.assert_array_less(-dx.contact.dist[1:], 0) + # extract the contact point with penetration + c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact) + c = c.replace(dim=c.dim[np.array([0])]) + for field in dataclasses.fields(Contact): + _assert_attr_eq(c, d.contact, field.name, 'box_box_edge', 1e-2) + + _CONVEX_CONVEX = """ + + + + + + + + + + + + + + + + + """ + + def test_convex_convex(self): + """Tests generic convex-convex collision.""" + directory = epath.resource_path('mujoco.mjx') + assets = { + 'meshes/tetrahedron.stl': ( + directory / 'test_data' / 'meshes/tetrahedron.stl' + ).read_bytes(), + 'meshes/dodecahedron.stl': ( + directory / 'test_data' / 'meshes/dodecahedron.stl' + ).read_bytes(), + } + _, dx = _collide(self._CONVEX_CONVEX, assets=assets) + c = dx.contact + + # Only one contact point for an edge contact. + self.assertLess(c.dist[0], 0) + 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 = """ + + + + + + + + + + """ + + def test_filter_self_collision(self): + """Tests that self collisions get filtered.""" + d, dx = _collide(self._SELF_COLLISION) + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 0) + + _PARENT_CHILD = """ + + + + + + + + + + + + + """ + + def test_filter_parent_child(self): + """Tests that parent-child collisions get filtered.""" + m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 0) + + def test_disable_filter_parent_child(self): + """Tests that filterparent flag disables parent-child filtering.""" + m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD) + m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_FILTERPARENT + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx, dx) + dx = collision_jit_fn(mx, dx) + + # one collision between parent-child spheres + self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0]) + self.assertEqual(dx.contact.pos.shape[0], 1) + + +class TopKContactTest(absltest.TestCase): + """Tests top-k contacts.""" + + _CAPSULES = """ + + + + + + + + + + + + + + + + + + + + """ + + def test_top_k_contacts(self): + m = mujoco.MjModel.from_xml_string(self._CAPSULES) + mx_top_k = mjx.device_put(m) + mx_all = mx_top_k.replace( + nnumeric=0, name_numericadr=np.array([]), numeric_data=np.array([]) + ) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + collision_jit_fn = jax.jit(mjx.collision) + kinematics_jit_fn = jax.jit(mjx.kinematics) + dx = kinematics_jit_fn(mx_all, dx) + + dx_all = collision_jit_fn(mx_all, dx) + dx_top_k = collision_jit_fn(mx_top_k, dx) + + self.assertEqual(dx_all.ncon, 3) + self.assertEqual(dx_top_k.ncon, 2) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/collision_primitive.py b/mjx/mujoco/mjx/_src/collision_primitive.py new file mode 100644 index 00000000..8c2dd529 --- /dev/null +++ b/mjx/mujoco/mjx/_src/collision_primitive.py @@ -0,0 +1,126 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Collision primitives.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import math +# pylint: disable=g-importing-member +from mujoco.mjx._src.collision_base import Contact +from mujoco.mjx._src.collision_base import GeomInfo +# pylint: enable=g-importing-member + + +def _plane_sphere( + plane_normal: jax.Array, + plane_pos: jax.Array, + sphere_pos: jax.Array, + radius: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Returns the penetration and contact point between a plane and sphere.""" + cdist = jp.dot(sphere_pos - plane_pos, plane_normal) + dist = cdist - radius + pos = sphere_pos - plane_normal * (radius + 0.5 * dist) + return dist, pos + + +def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact: + """Calculates contact between a plane and a sphere.""" + n = plane.mat[:, 2] + dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact: + """Calculates two contacts between a capsule and a plane.""" + n, axis = plane.mat[:, 2], cap.mat[:, 2] + # align contact frames with capsule axis + b, b_norm = math.normalize_with_norm(axis - n * jp.dot(n, axis)) + y, z = jp.array([0.0, 1.0, 0.0]), jp.array([0.0, 0.0, 1.0]) + b = jp.where(b_norm < 0.5, jp.where((-0.5 < n[1]) & (n[1] < 0.5), y, z), b) + frame = jp.array([[n, b, jp.cross(n, b)]]) + segment = axis * cap.size[1] + contacts = [] + for offset in [segment, -segment]: + dist, pos = _plane_sphere(n, plane.pos, cap.pos + offset, cap.size[0]) + dist = jp.expand_dims(dist, axis=0) + pos = jp.expand_dims(pos, axis=0) + contacts.append((dist, pos, frame)) + return jax.tree_map(lambda *x: jp.concatenate(x), *contacts) + + +def _sphere_sphere( + pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array +) -> Contact: + """Returns the penetration, contact point, and normal between two spheres.""" + n, dist = math.normalize_with_norm(pos2 - pos1) + n = jp.where(dist == 0.0, jp.array([1.0, 0.0, 0.0]), n) + dist = dist - (radius1 + radius2) + pos = pos1 + n * (radius1 + dist * 0.5) + return dist, pos, n + + +def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Contact: + """Calculates contact between two spheres.""" + dist, pos, n = _sphere_sphere(s1.pos, s1.size[0], s2.pos, s2.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact: + """Calculates one contact between a sphere and a capsule.""" + axis, length = cap.mat[:, 2], cap.size[1] + segment = axis * length + pt = math.closest_segment_point( + cap.pos - segment, cap.pos + segment, sphere.pos + ) + dist, pos, n = _sphere_sphere(sphere.pos, sphere.size[0], pt, cap.size[0]) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + + +def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact: + """Calculates one contact between two capsules.""" + axis1, length1, axis2, length2 = ( + cap1.mat[:, 2], + cap1.size[1], + cap2.mat[:, 2], + cap2.size[1], + ) + seg1, seg2 = axis1 * length1, axis2 * length2 + pt1, pt2 = math.closest_segment_to_segment_points( + cap1.pos - seg1, + cap1.pos + seg1, + cap2.pos - seg2, + cap2.pos + seg2, + ) + radius1, radius2 = cap1.size[0], cap2.size[0] + dist, pos, n = _sphere_sphere(pt1, radius1, pt2, radius2) + return jax.tree_map( + lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n)) + ) + +# store ncon as function attributes +plane_sphere.ncon = 1 +plane_capsule.ncon = 2 +sphere_sphere.ncon = 1 +sphere_capsule.ncon = 1 +capsule_capsule.ncon = 1 diff --git a/mjx/mujoco/mjx/_src/constraint.py b/mjx/mujoco/mjx/_src/constraint.py new file mode 100644 index 00000000..bd0bbaa6 --- /dev/null +++ b/mjx/mujoco/mjx/_src/constraint.py @@ -0,0 +1,340 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Core non-smooth constraint functions.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math +from mujoco.mjx._src import scan +from mujoco.mjx._src import support +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import Contact +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import EqType +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +class _Efc(PyTreeNode): + J: jax.Array + R: jax.Array + aref: jax.Array + frictionloss: jax.Array + + @classmethod + def zero(cls, m: Model) -> '_Efc': + z = jp.empty((0,)) + return _Efc(J=jp.empty((0, m.nv)), R=z, aref=z, frictionloss=z) + + +def _kbi( + m: Model, + solref: jax.Array, + solimp: jax.Array, + pos: jax.Array, +) -> Tuple[jax.Array, jax.Array, jax.Array]: + """Calculates stiffness, damping, and impedance of a constraint.""" + timeconst, dampratio = solref + + if not m.opt.disableflags & DisableBit.REFSAFE: + timeconst = jp.maximum(timeconst, 2 * m.opt.timestep) * (timeconst > 0) + + dmin, dmax, width, mid, power = solimp + + dmin = jp.clip(dmin, mujoco.mjMINIMP, mujoco.mjMAXIMP) + dmax = jp.clip(dmax, mujoco.mjMINIMP, mujoco.mjMAXIMP) + width = jp.maximum(0, width) + mid = jp.clip(mid, mujoco.mjMINIMP, mujoco.mjMAXIMP) + power = jp.maximum(1, power) + + # See https://mujoco.readthedocs.io/en/latest/modeling.html#solver-parameters + k = 1 / (dmax * dmax * timeconst * timeconst * dampratio * dampratio) + b = 2 / (dmax * timeconst) + # TODO(robotics-simulation): check various solparam settings in model gen test + k = jp.where(dampratio <= 0, -dampratio / (dmax * dmax), k) + b = jp.where(timeconst <= 0, -timeconst / dmax, b) + + imp_x = jp.abs(pos) / width + imp_a = (1.0 / jp.power(mid, power - 1)) * jp.power(imp_x, power) + imp_b = 1 - (1.0 / jp.power(1 - mid, power - 1)) * jp.power(1 - imp_x, power) + imp_y = jp.where(imp_x < mid, imp_a, imp_b) + imp = dmin + imp_y * (dmax - dmin) + imp = jp.clip(imp, dmin, dmax) + imp = jp.where(imp_x > 1.0, dmax, imp) + + return k, b, imp # corresponds to K, B, I of efc_KBIP + + +def _instantiate_connect(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporting data for connect equality constraints.""" + + if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0: + return _Efc.zero(m) + + connect_id = np.nonzero(m.eq_type == EqType.CONNECT)[0] + + if connect_id.size == 0: + return _Efc.zero(m) + + body1id, body2id = m.eq_obj1id[connect_id], m.eq_obj2id[connect_id] + data = m.eq_data[connect_id] + solref, solimp = m.eq_solref[connect_id], m.eq_solimp[connect_id] + + def fn(data, id1, id2, solref, solimp): + anchor1, anchor2 = data[0:3], data[3:6] + # find global points + pos1 = d.xmat[id1] @ anchor1 + d.xpos[id1] + pos2 = d.xmat[id2] @ anchor2 + d.xpos[id2] + + # compute position error + cpos = pos1 - pos2 + + # compute Jacobian difference (opposite of contact: 0 - 1) + jacp1, _ = support.jac(m, d, pos1, id1) + jacp2, _ = support.jac(m, d, pos2, id2) + j = (jacp1 - jacp2).T + + # impedance, inverse constraint mass, reference acceleration + k, b, imp = _kbi(m, solref, solimp, math.norm(cpos)) + invweight = m.body_invweight0[id1, 0] + m.body_invweight0[id2, 0] + r = jp.maximum(invweight * (1 - imp) / imp, mujoco.mjMINVAL).repeat(3) + aref = -b * (j @ d.qvel) - k * imp * cpos + + return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp) + + return jax.tree_map(jp.concatenate, efcs) + + +def _instantiate_weld(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporting data for connect weld constraints.""" + + if (m.opt.disableflags & DisableBit.EQUALITY) or m.neq == 0: + return _Efc.zero(m) + + weld_id = np.nonzero(m.eq_type == EqType.WELD)[0] + + if weld_id.size == 0: + return _Efc.zero(m) + + body1id, body2id = m.eq_obj1id[weld_id], m.eq_obj2id[weld_id] + data = m.eq_data[weld_id] + solref, solimp = m.eq_solref[weld_id], m.eq_solimp[weld_id] + + def fn(data, id1, id2, solref, solimp): + anchor1, anchor2 = data[0:3], data[3:6] + relpose, torquescale = data[6:10], data[10] + + # find global points + pos1 = d.xmat[id1] @ anchor2 + d.xpos[id1] + pos2 = d.xmat[id2] @ anchor1 + d.xpos[id2] + + # compute position error + cpos = pos1 - pos2 + + # compute Jacobian difference (opposite of contact: 0 - 1) + jacp1, jacr1 = support.jac(m, d, pos1, id1) + jacp2, jacr2 = support.jac(m, d, pos2, id2) + jacdifp = jacp1 - jacp2 + jacdifr = (jacr1 - jacr2) * torquescale + + # compute orientation error: neg(q1) * q0 * relpose (axis components only) + quat = math.quat_mul(d.xquat[id1], relpose) + quat1 = math.quat_inv(d.xquat[id2]) + crot = math.quat_mul(quat1, quat)[1:] # copy axis components + + # correct rotation Jacobian: 0.5 * neg(q1) * (jac0-jac1) * q0 * relpose + jac_fn = lambda j: math.quat_mul(math.quat_mul_axis(quat1, j), quat)[1:] + jacdifr = 0.5 * jax.vmap(jac_fn)(jacdifr) + + j = jp.concatenate((jacdifp.T, jacdifr.T)) + pos = jp.concatenate((cpos, crot)) + + # impedance, inverse constraint mass, reference acceleration + k, b, imp = _kbi(m, solref, solimp, math.norm(pos.at[3:].mul(torquescale))) + invweight = m.body_invweight0[id1] + m.body_invweight0[id2] + r = jp.maximum(invweight * (1 - imp) / imp, mujoco.mjMINVAL).repeat(3) + aref = -b * (j @ d.qvel) - k * imp * pos + + return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + efcs = jax.vmap(fn)(data, body1id, body2id, solref, solimp) + + return jax.tree_map(jp.concatenate, efcs) + + +def _instantiate_friction(m: Model, d: Data) -> _Efc: + # TODO(robotics-team): implement _instantiate_friction + del d + return _Efc.zero(m) + + +def _instantiate_limit(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporting data for joint limits.""" + + if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any(): + return _Efc.zero(m) + + def fn(jnt_typs, jnt_range, solref, solimp, margin, qpos, dofs, invweight0): + js, rs, arefs = [], [], [] + qpos_i, dof_i = 0, 0 + + for i in range(len(jnt_typs)): + jnt_typ = JointType(jnt_typs[i]) + + if jnt_typ == JointType.FREE: + return None # omit constraint rows for free joints + elif jnt_typ == JointType.BALL: + axis, angle = math.quat_to_axis_angle(qpos[qpos_i : qpos_i + 4]) + dist = jp.amax(jnt_range[i]) - angle + j = jp.sum( + jax.vmap(jp.multiply)(dofs[dof_i : dof_i + 3], -axis), axis=0 + ) + elif jnt_typ in (JointType.HINGE, JointType.SLIDE): + dist_min = qpos[qpos_i] - jnt_range[i, 0] + dist_max = jnt_range[i, 1] - qpos[qpos_i] + dist = jp.minimum(dist_min, dist_max) + j = dofs[dof_i] * ((dist_min < dist_max) * 2 - 1) + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + dist = dist - margin[i] + k, b, imp = _kbi(m, solref[i], solimp[i], dist) + r = jp.maximum(invweight0[dof_i] * (1 - imp) / imp, mujoco.mjMINVAL) + aref = -b * (j @ d.qvel) - k * imp * dist + j, aref = j * (dist < 0), aref * (dist < 0) + js, rs, arefs = js + [j], rs + [r], arefs + [aref] + dof_i, qpos_i = dof_i + jnt_typ.dof_width(), qpos_i + jnt_typ.qpos_width() + + return jp.stack(js), jp.stack(rs), jp.stack(arefs) + + j, r, aref = scan.flat( + m, + fn, + 'jjjjjqvv', + 'jjj', + m.jnt_type, + m.jnt_range, + m.jnt_solref, + m.jnt_solimp, + m.jnt_margin, + d.qpos, + jp.eye(m.nv), + m.dof_invweight0, + ) + + return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + +def _instantiate_contact(m: Model, d: Data) -> _Efc: + """Returns jacobians and supporitng data for contacts.""" + + if (m.opt.disableflags & DisableBit.CONTACT) or d.ncon == 0: + return _Efc.zero(m) + + def fn(contact: Contact): + dist = contact.dist - contact.includemargin + k, b, imp = _kbi(m, contact.solref, contact.solimp, dist) + + geom_bodyid = jp.array(m.geom_bodyid) + body1, body2 = geom_bodyid[contact.geom1], geom_bodyid[contact.geom2] + diff = support.jac_dif_pair(m, d, contact.pos, body1, body2) + t = m.body_invweight0[body1, 0] + m.body_invweight0[body2, 0] + + # rotate Jacobian differences to contact frame + diff_con = contact.frame @ diff.T + + # TODO(robotics-simulation): add support for other friction dimensions + # 4 pyramidal friction directions + js, rs = [], [] + for diff_tan, friction in zip(diff_con[1:], contact.friction[:2]): + for f in (friction, -friction): + js.append(diff_con[0] + diff_tan * f) + rs.append((t + f * f * t) * 2 * f * f * (1 - imp) / imp) + + j, r = jp.stack(js), jp.stack(rs) + r = jp.maximum(r, mujoco.mjMINVAL) + aref = -b * (j @ d.qvel) - k * imp * dist + mask_fn = jax.vmap(lambda x, mask=(dist < 0): x * mask) + j, aref = jax.tree_map(mask_fn, (j, aref)) + + return _Efc(J=j, R=r, aref=aref, frictionloss=jp.zeros_like(r)) + + return jax.tree_map(jp.concatenate, jax.vmap(fn)(d.contact)) + + +def count_constraints(m: Model, d: Data) -> Tuple[int, int, int, int]: + """Returns equality, friction, limit, and contact constraint counts.""" + if m.opt.disableflags & DisableBit.CONSTRAINT: + return 0, 0, 0, 0 + + if m.opt.disableflags & DisableBit.EQUALITY: + ne = 0 + else: + ne_weld = (m.eq_type == EqType.WELD).sum() + ne_connect = (m.eq_type == EqType.CONNECT).sum() + ne = ne_weld * 6 + ne_connect * 3 + + nf = 0 + + if (m.opt.disableflags & DisableBit.LIMIT) or not m.jnt_limited.any(): + nl = 0 + else: + nl = (m.jnt_type != JointType.FREE).sum() + + if (m.opt.disableflags & DisableBit.CONTACT): + nc = 0 + else: + nc = d.ncon * 4 + + return ne, nf, nl, nc + + +def make_constraint(m: Model, d: Data) -> Data: + """Creates constraint jacobians and other supporting data.""" + + ns = sum(count_constraints(m, d)[:-1]) + # TODO(robotics-simulation): make device_put set nefc/efc_address instead + d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + d.ncon * 4, 4)}) + + if m.opt.disableflags & DisableBit.CONSTRAINT: + efc = _Efc.zero(m) + else: + efcs = ( + _instantiate_connect(m, d), + _instantiate_weld(m, d), + _instantiate_friction(m, d), + _instantiate_limit(m, d), + _instantiate_contact(m, d), + ) + efc = jax.tree_map(lambda *x: jp.concatenate(x), *efcs) + + d = d.replace( + efc_J=efc.J, + efc_D=1 / efc.R, + efc_aref=efc.aref, + efc_frictionloss=efc.frictionloss, + nefc=efc.aref.shape[0], + ) + + return d diff --git a/mjx/mujoco/mjx/_src/constraint_test.py b/mjx/mujoco/mjx/_src/constraint_test.py new file mode 100644 index 00000000..ce289df7 --- /dev/null +++ b/mjx/mujoco/mjx/_src/constraint_test.py @@ -0,0 +1,144 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for constraint functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import constraint +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(a, b, name, step, fname, atol=1e-3, rtol=1e-3): + err_msg = f'mismatch: {name} at step {step} in {fname}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +class ConstraintTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(test_util.TEST_FILES)) + def testconstraints(self, seed, fname): + """Test constraints.""" + np.random.seed(seed) + + # exclude convex.xml since convex contacts are not exactly equivalent + if fname == 'convex.xml': + return + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + for i in range(100): + dx = dx.replace(qpos=jax.device_put(d.qpos), qvel=jax.device_put(d.qvel)) + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + nnz_filter = dx.efc_J.any(axis=1) + + mj_efc_j = d.efc_J.reshape((-1, m.nv)) + mjx_efc_j = dx.efc_J[nnz_filter] + _assert_eq(mj_efc_j, mjx_efc_j, 'efc_J', i, fname) + + mjx_efc_d = dx.efc_D[nnz_filter] + _assert_eq(d.efc_D, mjx_efc_d, 'efc_D', i, fname) + + mjx_efc_aref = dx.efc_aref[nnz_filter] + _assert_eq(d.efc_aref, mjx_efc_aref, 'efc_aref', i, fname) + + mjx_efc_frictionloss = dx.efc_frictionloss[nnz_filter] + _assert_eq( + d.efc_frictionloss, + mjx_efc_frictionloss, + 'efc_frictionloss', + i, + fname, + ) + + def test_disable_refsafe(self): + m = test_util.load_test_file('ant.xml') + + timeconst = m.opt.timestep / 4.0 # timeconst < 2 * timestep + solimp = jp.array([timeconst, 1.0]) + solref = jp.array([0.8, 0.99, 0.001, 0.2, 2]) + pos = jp.ones(3) + + m.opt.disableflags = m.opt.disableflags | DisableBit.REFSAFE + mx = mjx.device_put(m) + k, *_ = constraint._kbi(mx, solimp, solref, pos) + self.assertEqual(k, 1 / (0.99**2 * timeconst**2)) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.REFSAFE + mx = mjx.device_put(m) + k, *_ = constraint._kbi(mx, solimp, solref, pos) + self.assertEqual(k, 1 / (0.99**2 * (2 * m.opt.timestep) ** 2)) + + def test_disableconstraint(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONSTRAINT + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = constraint.make_constraint(mx, dx) + self.assertGreater(dx.efc_J.shape[0], 1) + + m.opt.disableflags = m.opt.disableflags | DisableBit.CONSTRAINT + mx = mjx.device_put(m) + dx = constraint.make_constraint(mx, dx) + self.assertEqual(dx.efc_J.shape[0], 0) + + def test_disable_equality(self): + m = test_util.load_test_file('weld.xml') + d = mujoco.MjData(m) + + m.opt.disableflags = m.opt.disableflags | DisableBit.EQUALITY + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = constraint.make_constraint(mx, dx) + self.assertEqual(dx.efc_J.shape[0], 0) + + def test_disable_contact(self): + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + d.qpos[2] = 0.0 + mujoco.mj_forward(m, d) + + m.opt.disableflags = m.opt.disableflags & ~DisableBit.CONTACT + mx, dx = mjx.device_put(m), mjx.device_put(d) + dx = dx.tree_replace( + {'contact.frame': dx.contact.frame.reshape((-1, 3, 3))} + ) + efc = constraint._instantiate_contact(mx, dx) + self.assertIsNotNone(efc) + + m.opt.disableflags = m.opt.disableflags | DisableBit.CONTACT + mx, dx = mjx.device_put(m), mjx.device_put(d) + efc = constraint._instantiate_contact(mx, dx) + self.assertEqual(efc.J.shape[0], 0) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/dataclasses.py b/mjx/mujoco/mjx/_src/dataclasses.py new file mode 100644 index 00000000..0936bec6 --- /dev/null +++ b/mjx/mujoco/mjx/_src/dataclasses.py @@ -0,0 +1,151 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Wrapper that automatically registers dataclass as a Jax PyTree.""" + +import copy +import dataclasses + +import typing +from typing import Dict, Optional, Sequence, TypeVar +import jax +import numpy as np + +_T = TypeVar('_T') + + +def dataclass(clz: _T) -> _T: + """Wraps a dataclass with metadata for which fields are pytrees. + + This is based off flax.struct.dataclass, but instead of using field + descriptors to specify which fields are pytrees, we follow a simple rule: + a leaf field is a pytree node if and only if it's a jax.Array + + Args: + clz: the class to register as a dataclass + + Returns: + the resulting dataclass, registered with Jax + """ + data_clz = dataclasses.dataclass(frozen=True)(clz) + meta_fields, data_fields = [], [] + for field in dataclasses.fields(data_clz): + if any(( + field.type is jax.Array, + dataclasses.is_dataclass(field.type), + jax.Array in typing.get_args(field.type), + )): + data_fields.append(field) + else: + meta_fields.append(field) + + def replace(self, **updates): + """"Returns a new object replacing the specified fields with new values.""" + return dataclasses.replace(self, **updates) + + data_clz.replace = replace + + def iterate_clz_with_keys(x): + # numpy arrays are not hashable, so convert them to tuples for jit cache + to_tup = lambda x: tuple(x) if len(x.shape) == 1 else tuple(map(to_tup, x)) + + def to_meta(field, obj): + val = getattr(obj, field.name) + return to_tup(val) if isinstance(val, np.ndarray) else val + + def to_data(field, obj): + return (jax.tree_util.GetAttrKey(field.name), getattr(obj, field.name)) + + data = tuple(to_data(f, x) for f in data_fields) + meta = tuple(to_meta(f, x) for f in meta_fields) + return data, meta + + def clz_from_iterable(meta, data): + + def from_meta(field, meta): + if field.type is np.ndarray: + return (field.name, np.array(meta)) + else: + return (field.name, meta) + + from_data = lambda field, meta: (field.name, meta) + + meta_args = tuple(from_meta(f, m) for f, m in zip(meta_fields, meta)) + data_args = tuple(from_data(f, m) for f, m in zip(data_fields, data)) + + return data_clz(**dict(meta_args + data_args)) + + jax.tree_util.register_pytree_with_keys( + data_clz, iterate_clz_with_keys, clz_from_iterable + ) + + return data_clz + + +TNode = TypeVar('TNode', bound='PyTreeNode') + + +class PyTreeNode: + """Base class for dataclasses that should act like a JAX pytree node. + + This base class additionally avoids type checking errors when using PyType. + """ + + def __init_subclass__(cls): + dataclass(cls) + + def __init__(self, *args, **kwargs): + # stub for pytype + raise NotImplementedError + + def replace(self: TNode, **overrides) -> TNode: + # stub for pytype + raise NotImplementedError + + def tree_replace( + self, params: Dict[str, Optional[jax.typing.ArrayLike]] + ) -> 'PyTreeNode': + new = self + for k, v in params.items(): + new = _tree_replace(new, k.split('.'), v) + return new + + +def _tree_replace( + base: PyTreeNode, + attr: Sequence[str], + val: Optional[jax.typing.ArrayLike], +) -> PyTreeNode: + """Sets attributes in a struct.dataclass with values.""" + if not attr: + return base + + # special case for List attribute + if len(attr) > 1 and isinstance(getattr(base, attr[0]), list): + lst = copy.deepcopy(getattr(base, attr[0])) + + for i, g in enumerate(lst): + if not hasattr(g, attr[1]): + continue + v = val if not hasattr(val, '__iter__') else val[i] + lst[i] = _tree_replace(g, attr[1:], v) + + return base.replace(**{attr[0]: lst}) + + if len(attr) == 1: + return base.replace(**{attr[0]: val}) + + return base.replace( + **{attr[0]: _tree_replace(getattr(base, attr[0]), attr[1:], val)} + ) diff --git a/mjx/mujoco/mjx/_src/device.py b/mjx/mujoco/mjx/_src/device.py new file mode 100644 index 00000000..3c3c213f --- /dev/null +++ b/mjx/mujoco/mjx/_src/device.py @@ -0,0 +1,282 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Get and put mujoco data on/off device.""" + +import copy +import dataclasses +from typing import Any, Dict, Iterable, List, Union, overload +import warnings + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import mesh +from mujoco.mjx._src import types + +_MJ_TYPE_ATTR = { + mujoco.mjtBias: (mujoco.MjModel.actuator_biastype,), + mujoco.mjtDyn: (mujoco.MjModel.actuator_dyntype,), + mujoco.mjtEq: (mujoco.MjModel.eq_type,), + mujoco.mjtGain: (mujoco.MjModel.actuator_gaintype,), + mujoco.mjtTrn: (mujoco.MjModel.actuator_trntype,), + mujoco.mjtCone: ( + mujoco.MjModel.opt, + mujoco.MjOption.cone, + ), + mujoco.mjtIntegrator: ( + mujoco.MjModel.opt, + mujoco.MjOption.integrator, + ), +} + +_TYPE_MAP = { + mujoco._structs._MjContactList: types.Contact, # pylint: disable=protected-access + mujoco.MjData: types.Data, + mujoco.MjModel: types.Model, + mujoco.MjOption: types.Option, + mujoco.MjStatistic: types.Statistic, + mujoco.mjtBias: types.BiasType, + mujoco.mjtCone: types.ConeType, + mujoco.mjtDisableBit: types.DisableBit, + mujoco.mjtDyn: types.DynType, + mujoco.mjtEq: types.EqType, + mujoco.mjtGain: types.GainType, + mujoco.mjtIntegrator: types.IntegratorType, + mujoco.mjtSolver: types.SolverType, + mujoco.mjtTrn: types.TrnType, +} + +_TRANSFORMS = { + (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-1] + (3, 3)), + (types.Model, 'actuator_trnid'): lambda x: x[:, 0], + (types.Contact, 'frame'): ( + lambda x: x.reshape(x.shape[:-1] + (3, 3)) # pylint: disable=g-long-lambda + if x is not None and x.shape[0] else jp.zeros((0, 3, 3)) + ), +} + +_INVERSE_TRANSFORMS = { + (types.Data, 'ximat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Data, 'xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Data, 'geom_xmat'): lambda x: x.reshape(x.shape[:-2] + (9,)), + (types.Contact, 'frame'): ( + lambda x: x.reshape(x.shape[:-2] + (9,)) # pylint: disable=g-long-lambda + if x is not None and x.shape[0] else jp.zeros((0, 9)) + ), +} + +_DERIVED = mesh.DERIVED.union( + # efc_J is dense in MJX, sparse in MJ. ignore for now. + {(types.Data, 'efc_J'), (types.Option, 'has_fluid_params')} +) + + +def _model_derived(value: mujoco.MjModel) -> Dict[str, Any]: + return {k: jax.device_put(v) for k, v in mesh.get(value).items()} + + +def _data_derived(value: mujoco.MjData) -> Dict[str, Any]: + return {'efc_J': jax.device_put(value.efc_J)} + + +def _option_derived(value: types.Option) -> Dict[str, Any]: + has_fluid = ( + value.density > 0 or value.viscosity > 0 or (value.wind != 0.0).any() + ) + return {'has_fluid_params': has_fluid} + + +def _validate(m: mujoco.MjModel): + """Validates that an mjModel is compatible with MJX.""" + if m.opt.solver not in set(types.SolverType): + name = mujoco.mjtSolver(m.opt.solver).name + warnings.warn(f'Solver {name} is not supported, reverting to CG.') + m.opt.solver = mujoco.mjtSolver.mjSOL_CG.value + + # check enum types + for mj_type, attrs in _MJ_TYPE_ATTR.items(): + val = m + for attr in attrs: + val = attr.fget(val) # pytype: disable=attribute-error + + typs = set(val) if isinstance(val, Iterable) else {val} + unsupported_typs = typs - set(_TYPE_MAP[mj_type]) + if unsupported_typs: + raise NotImplementedError(f'{unsupported_typs} not implemented.') + + # check condim + if any(dim != 3 for dim in m.geom_condim) or any( + dim != 3 for dim in m.pair_dim + ): + raise NotImplementedError('Only condim=3 is supported.') + + # check collision geom types + candidate_set = collision_driver.collision_candidates(m) + for g1, g2, *_ in candidate_set: + g1, g2 = mujoco.mjtGeom(g1), mujoco.mjtGeom(g2) + if g1 == mujoco.mjtGeom.mjGEOM_PLANE and g2 in ( + mujoco.mjtGeom.mjGEOM_PLANE, + mujoco.mjtGeom.mjGEOM_HFIELD, + ): + # MuJoCo does not collide planes with other planes or hfields + continue + if collision_driver.get_collision_fn((g1, g2)) is None: + raise NotImplementedError(f'({g1}, {g2}) collisions not implemented.') + + # TODO(erikfrey): warn for high solver iterations, nefc, etc. + + # mjNDISABLE is not a DisableBit flag, so must be explicitly ignored + disablebit_members = set(mujoco.mjtDisableBit.__members__.values()) - { + mujoco.mjtDisableBit.mjNDISABLE} + unsupported_disable = disablebit_members - { + mujoco.mjtDisableBit(t.value) for t in types.DisableBit + } + for f in unsupported_disable: + if f & m.opt.disableflags: + warnings.warn(f'Ignoring disable flag {f.name}.') + + # mjNENABLE is not an EnableBit flag, so must be explicitly ignored + unsupported_enable = set(mujoco.mjtEnableBit.__members__.values()) - { + mujoco.mjtEnableBit.mjNENABLE + } + for f in unsupported_enable: + if f & m.opt.enableflags: + warnings.warn(f'Ignoring enable flag {f.name}.') + + +@overload +def device_put(value: mujoco.MjData) -> types.Data: + ... + + +@overload +def device_put(value: mujoco.MjModel) -> types.Model: + ... + + +def device_put(value): + """Places mujoco data onto a device. + + Args: + value: a mujoco struct to transfer + + Returns: + on-device MJX struct reflecting the input value + """ + clz = _TYPE_MAP.get(type(value)) + if clz is None: + raise NotImplementedError(f'{type(value)} is not supported for device_put.') + + if isinstance(value, mujoco.MjModel): + _validate(value) # type: ignore + + init_kwargs = {} + for f in dataclasses.fields(clz): # type: ignore + if (clz, f.name) in _DERIVED: + continue + + field_value = getattr(value, f.name) + if (clz, f.name) in _TRANSFORMS: + field_value = _TRANSFORMS[(clz, f.name)](field_value) + + if f.type is jax.Array: + field_value = jax.device_put(field_value) + elif type(field_value) in _TYPE_MAP.keys(): + field_value = device_put(field_value) + + init_kwargs[f.name] = copy.copy(field_value) + + derived_kwargs = {} + if isinstance(value, mujoco.MjModel): + derived_kwargs = _model_derived(value) + elif isinstance(value, mujoco.MjData): + derived_kwargs = _data_derived(value) + elif isinstance(value, mujoco.MjOption): + derived_kwargs = _option_derived(value) + + return clz(**init_kwargs, **derived_kwargs) # type: ignore + + +@overload +def device_get_into( + result: Union[mujoco.MjData, List[mujoco.MjData]], value: types.Data +): + ... + + +def device_get_into(result, value): + """Transfers data off device into a mujoco MjData. + + Data on device often has a batch dimension which adds (N,) to the beginning + of each array shape where N = batch size. + + If result is a single MjData, arrays are copied over with the batch dimension + intact. If result is a list, the list must be length N and will be populated + with distinct MjData structs where the batch dimension is stripped. + + Args: + result: struct (or list of structs) to transfer into + value: device value to transfer + + Raises: + RuntimeError: if result length doesn't match data batch size + """ + + value = jax.device_get(value) + + if isinstance(result, list): + array_shapes = [s.shape for s in jax.tree_util.tree_flatten(value)[0]] + + if any(len(s) < 1 or s[0] != array_shapes[0][0] for s in array_shapes): + raise ValueError('unrecognizable batch dimension in value') + + batch_size = array_shapes[0][0] + + if len(result) != batch_size: + raise ValueError( + f"result length ({len(result)}) doesn't match value batch size" + f' ({batch_size})' + ) + + for i in range(batch_size): + value_i = jax.tree_map(lambda x, i=i: x[i], value) + device_get_into(result[i], value_i) + + else: + if isinstance(result, mujoco.MjData): + mujoco._functions._realloc_con_efc( # pylint: disable=protected-access + result, ncon=value.ncon, nefc=value.nefc + ) + + for f in dataclasses.fields(value): # type: ignore + if (type(value), f.name) in _DERIVED: + continue + + field_value = getattr(value, f.name) + + if (type(value), f.name) in _INVERSE_TRANSFORMS: + field_value = _INVERSE_TRANSFORMS[(type(value), f.name)](field_value) + + if type(field_value) in _TYPE_MAP.values(): + device_get_into(getattr(result, f.name), field_value) + continue + + try: + setattr(result, f.name, field_value) + except AttributeError: + getattr(result, f.name)[:] = field_value diff --git a/mjx/mujoco/mjx/_src/device_test.py b/mjx/mujoco/mjx/_src/device_test.py new file mode 100644 index 00000000..bdc4e2d4 --- /dev/null +++ b/mjx/mujoco/mjx/_src/device_test.py @@ -0,0 +1,184 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for moving mujoco structs on and off device.""" + +import dataclasses + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import device +from mujoco.mjx._src import test_util +from mujoco.mjx._src import types +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(testcase, a, b, attr=None, name=None): + if (type(a), attr) in device._DERIVED: + return + + if attr: + a, b = getattr(a, attr), getattr(b, attr) + + if isinstance(a, PyTreeNode): + for field in dataclasses.fields(a): + _assert_eq(testcase, a, b, field.name, type(a).__name__) + return + + typ = {'Model': types.Model, 'Data': types.Data, + 'Contact': types.Contact}.get(name) + if (typ, attr) in device._TRANSFORMS: + b = device._TRANSFORMS[(typ, attr)](b) + + err_msg = f'mismatch: {attr} in {name}' + if not hasattr(b, 'shape') or not b.shape: + testcase.assertEqual(a, b, err_msg) + return + + a, b = np.array(a), np.array(b) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=1e-8) + + +class DeviceTest(parameterized.TestCase): + + @parameterized.parameters(test_util.TEST_FILES) + def testdevice_put(self, fname): + """Test putting MjData and MjModel on device.""" + m = test_util.load_test_file(fname) + # advance state to ensure non-zero fields + d = mujoco.MjData(m) + for _ in range(10): + mujoco.mj_step(m, d) + + _assert_eq(self, mjx.device_put(d), d) + _assert_eq(self, mjx.device_put(m), m) + + @parameterized.parameters(test_util.TEST_FILES) + def testdevice_get(self, fname): + """Test getting MjData from a device.""" + m = test_util.load_test_file(fname) + mx = device.device_put(m) + dx = mjx.make_data(mx) + d = mujoco.MjData(m) + device.device_get_into(d, dx) + _assert_eq(self, dx, d) + + @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + def testdevice_get_batched(self, fname): + """Test getting MjData from a device.""" + m = test_util.load_test_file(fname) + mx = device.device_put(m) + batch_size = 32 + + # create mjx_data and batch it + dx = mjx.make_data(mx) + dx = jax.tree_map( + lambda x: jp.repeat(x, batch_size).reshape((batch_size,) + x.shape), + dx, + ) + ds = [mujoco.MjData(m) for _ in range(batch_size - 1)] + + with self.assertRaises(ValueError): + device.device_get_into(ds, dx) + + ds = [mujoco.MjData(m) for _ in range(batch_size)] + device.device_get_into(ds, dx) + dx = jax.device_get(dx) # faster indexing for testing + for i in range(batch_size): + _assert_eq(self, jax.tree_map(lambda x, i=i: x[i], dx), ds[i]) + + +class ValidateInputTest(absltest.TestCase): + + def test_solver(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertWarns(UserWarning): + mx = mjx.device_put(m) + self.assertEqual(mx.opt.solver, mujoco.mjtSolver.mjSOL_CG) + + def test_integrator(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_cone(self): + m = mujoco.MjModel.from_xml_string( + '' + ) + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_trn(self): + m = test_util.load_test_file('ant.xml') + m.actuator_trntype[0] = mujoco.mjtTrn.mjTRN_SITE + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_dyn(self): + m = test_util.load_test_file('ant.xml') + m.actuator_dyntype[0] = mujoco.mjtDyn.mjDYN_MUSCLE + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_gain(self): + m = test_util.load_test_file('ant.xml') + m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_bias(self): + m = test_util.load_test_file('ant.xml') + m.actuator_gaintype[0] = mujoco.mjtGain.mjGAIN_MUSCLE + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_condim(self): + m = test_util.load_test_file('ant.xml') + for i in [1, 4, 6]: + m.geom_condim[0] = i + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + def test_geoms(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + + + + """) + with self.assertRaises(NotImplementedError): + _ = mjx.device_put(m) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py new file mode 100644 index 00000000..42faeca2 --- /dev/null +++ b/mjx/mujoco/mjx/_src/forward.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Forward step functions.""" + +import functools +from typing import Optional, Sequence + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import constraint +from mujoco.mjx._src import math +from mujoco.mjx._src import passive +from mujoco.mjx._src import scan +from mujoco.mjx._src import smooth +from mujoco.mjx._src import solver +from mujoco.mjx._src import support +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import BiasType +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import DynType +from mujoco.mjx._src.types import GainType +from mujoco.mjx._src.types import IntegratorType +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import SolverType +# pylint: enable=g-importing-member +import numpy as np + +# RK4 tableau +_RK4_A = np.array([ + [0.5, 0.0, 0.0], + [0.0, 0.5, 0.0], + [0.0, 0.0, 1.0], +]) +_RK4_B = np.array([1.0 / 6.0, 1.0 / 3.0, 1.0 / 3.0, 1.0 / 6.0]) + + +def named_scope(fn, name: str = ''): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with jax.named_scope(name or getattr(fn, '__name__')): + res = fn(*args, **kwargs) + return res + + return wrapper + + +@named_scope +def _position(m: Model, d: Data) -> Data: + """Position-dependent computations.""" + # TODO(robotics-simulation): tendon + d = smooth.kinematics(m, d) + d = smooth.com_pos(m, d) + d = smooth.crb(m, d) + d = smooth.factor_m(m, d, d.qM) + d = collision_driver.collision(m, d) + d = constraint.make_constraint(m, d) + d = smooth.transmission(m, d) + return d + + +@named_scope +def _velocity(m: Model, d: Data) -> Data: + """Velocity-dependent computations.""" + d = d.replace(actuator_velocity=d.actuator_moment @ d.qvel) + d = smooth.com_vel(m, d) + d = passive.passive(m, d) + d = smooth.rne(m, d) + return d + + +@named_scope +def _actuation(m: Model, d: Data) -> Data: + """Actuation-dependent computations.""" + if not m.nu or m.opt.disableflags & DisableBit.ACTUATION: + return d.replace( + act_dot=jp.zeros((m.na,)), + qfrc_actuator=jp.zeros((m.nv,)), + ) + + ctrl = d.ctrl + if not m.opt.disableflags & DisableBit.CLAMPCTRL: + ctrlrange = jp.where( + m.actuator_ctrllimited[:, None], + m.actuator_ctrlrange, + jp.array([-jp.inf, jp.inf]), + ) + ctrl = jp.clip(ctrl, ctrlrange[:, 0], ctrlrange[:, 1]) + + # act_dot for stateful actuators + def get_act_dot(dyn_typ, dyn_prm, ctrl, act): + if dyn_typ == DynType.NONE: + act_dot = jp.array(0.0) + elif dyn_typ == DynType.INTEGRATOR: + act_dot = ctrl + elif dyn_typ == DynType.FILTER: + act_dot = (ctrl - act) / jp.clip(dyn_prm[0], mujoco.mjMINVAL) + else: + raise NotImplementedError(f'dyntype {dyn_typ.name} not implemented.') + return act_dot + + act_dot = jp.zeros((m.na,)) + if m.na: + act_dot = scan.flat( + m, + get_act_dot, + 'uuua', + 'a', + m.actuator_dyntype, + m.actuator_dynprm, + ctrl, + d.act, + group_by='u', + ) + + ctrl_act = ctrl + if m.na: + act_last_dim = d.act[m.actuator_actadr + m.actuator_actnum - 1] + ctrl_act = jp.where(m.actuator_actadr == -1, ctrl, act_last_dim) + + def get_force(*args): + gain_t, gain_p, bias_t, bias_p, len_, vel, ctrl_act = args + + typ, prm = GainType(gain_t), gain_p + if typ == GainType.FIXED: + gain = prm[0] + elif typ == GainType.AFFINE: + gain = prm[0] + prm[1] * len_ + prm[2] * vel + else: + raise RuntimeError(f'unrecognized gaintype {typ.name}.') + + typ, prm = BiasType(bias_t), bias_p + bias = jp.array(0.0) + if typ == BiasType.AFFINE: + bias = prm[0] + prm[1] * len_ + prm[2] * vel + + return gain * ctrl_act + bias + + force = scan.flat( + m, + get_force, + 'uuuuuuu', + 'u', + m.actuator_gaintype, + m.actuator_gainprm, + m.actuator_biastype, + m.actuator_biasprm, + d.actuator_length, + d.actuator_velocity, + ctrl_act, + group_by='u', + ) + forcerange = jp.where( + m.actuator_forcelimited[:, None], + m.actuator_forcerange, + jp.array([-jp.inf, jp.inf]), + ) + force = jp.clip(force, forcerange[:, 0], forcerange[:, 1]) + + qfrc_actuator = d.actuator_moment.T @ force + + # clamp qfrc_actuator + actfrcrange = jp.where( + m.jnt_actfrclimited[:, None], + m.jnt_actfrcrange, + jp.array([-jp.inf, jp.inf]), + ) + ids = sum( + ([i] * JointType(j).dof_width() for i, j in enumerate(m.jnt_type)), [] + ) + actfrcrange = jp.take(actfrcrange, jp.array(ids), axis=0) + qfrc_actuator = jp.clip(qfrc_actuator, actfrcrange[:, 0], actfrcrange[:, 1]) + + d = d.replace(act_dot=act_dot, qfrc_actuator=qfrc_actuator) + return d + + +@named_scope +def _acceleration(m: Model, d: Data) -> Data: + """Add up all non-constraint forces, compute qacc_smooth.""" + qfrc_applied = d.qfrc_applied + support.xfrc_accumulate(m, d) + qfrc_smooth = d.qfrc_passive - d.qfrc_bias + d.qfrc_actuator + qfrc_applied + qacc_smooth = smooth.solve_m(m, d, qfrc_smooth) + d = d.replace(qfrc_smooth=qfrc_smooth, qacc_smooth=qacc_smooth) + return d + + +@named_scope +def _integrate_pos( + jnt_typs: Sequence[str], qpos: jax.Array, qvel: jax.Array, dt: jax.Array +) -> jax.Array: + """Integrate position given velocity.""" + qs, qi, vi = [], 0, 0 + + for jnt_typ in jnt_typs: + if jnt_typ == JointType.FREE: + pos = qpos[qi : qi + 3] + dt * qvel[vi : vi + 3] + quat = math.quat_integrate( + qpos[qi + 3 : qi + 7], qvel[vi + 3 : vi + 6], dt + ) + qs.append(jp.concatenate([pos, quat])) + qi, vi = qi + 7, vi + 6 + elif jnt_typ == JointType.BALL: + quat = math.quat_integrate(qpos[qi : qi + 4], qvel[vi : vi + 3], dt) + qs.append(quat) + qi, vi = qi + 4, vi + 3 + elif jnt_typ in (JointType.HINGE, JointType.SLIDE): + pos = qpos[qi] + dt * qvel[vi] + qs.append(pos[None]) + qi, vi = qi + 1, vi + 1 + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + return jp.concatenate(qs) if qs else jp.empty((0,)) + + +@named_scope +def _advance( + m: Model, + d: Data, + act_dot: jax.Array, + qacc: jax.Array, + qvel: Optional[jax.Array] = None, +) -> Data: + """Advance state and time given activation derivatives and acceleration.""" + act = d.act + if m.na: + act = d.act + act_dot * m.opt.timestep + actrange = jp.where( + m.actuator_actlimited[:, None], + m.actuator_actrange, + jp.array([-jp.inf, jp.inf]), + ) + fn = lambda act, actrange: jp.clip(act, actrange[0], actrange[1]) + act = scan.flat(m, fn, 'au', 'a', act, actrange, group_by='u') + + # advance velocities + d = d.replace(qvel=d.qvel + qacc * m.opt.timestep) + + # advance positions with qvel if given, d.qvel otherwise (semi-implicit) + qvel = d.qvel if qvel is None else qvel + integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep) + qpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d.qpos, qvel) + + # advance time + time = d.time + m.opt.timestep + + return d.replace(act=act, qpos=qpos, time=time) + + +@named_scope +def _euler(m: Model, d: Data) -> Data: + """Euler integrator, semi-implicit in velocity.""" + # integrate damping implicitly + qacc = d.qacc + if not m.opt.disableflags & DisableBit.EULERDAMP: + # TODO(robotics-simulation): can this be done with a smaller perf hit + mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping) + dh = smooth.factor_m(m, d, mh) + qfrc = d.qfrc_smooth + d.qfrc_constraint + qacc = smooth.solve_m(m, dh, qfrc) + return _advance(m, d, d.act_dot, qacc) + + +@named_scope +def _rungekutta4(m: Model, d: Data) -> Data: + """Runge-Kutta explicit order 4 integrator.""" + d_t0 = d + # pylint: disable=invalid-name + A, B = _RK4_A, _RK4_B + C = jp.tril(A).sum(axis=0) # C(i) = sum_j A(i,j) + T = d.time + C * m.opt.timestep + # pylint: enable=invalid-name + + kqvel = d.qvel # intermediate RK solution + # RK solutions sum + qvel, qacc, act_dot = jax.tree_map( + lambda k: B[0] * k, (kqvel, d.qacc, d.act_dot) + ) + integrate_fn = lambda *args: _integrate_pos(*args, dt=m.opt.timestep) + + def f(carry, x): + qvel, qacc, act_dot, kqvel, d = carry + a, b, t = x # tableau numbers + dqvel, dqacc, dact_dot = jax.tree_map( + lambda k: a * k, (kqvel, d.qacc, d.act_dot) + ) + # get intermediate RK solutions + kqpos = scan.flat(m, integrate_fn, 'jqv', 'q', m.jnt_type, d_t0.qpos, dqvel) + kact = d_t0.act + dact_dot * m.opt.timestep + kqvel = d_t0.qvel + dqacc * m.opt.timestep + d = d.replace(qpos=kqpos, qvel=kqvel, act=kact, time=t) + d = forward(m, d) + + qvel += b * kqvel + qacc += b * d.qacc + act_dot += b * d.act_dot + + return (qvel, qacc, act_dot, kqvel, d), None + + abt = jp.vstack([jp.diag(A), B[1:4], T]).T + out, _ = jax.lax.scan(f, (qvel, qacc, act_dot, kqvel, d), abt, unroll=3) + qvel, qacc, act_dot, *_ = out + + d = _advance(m, d_t0, act_dot, qacc, qvel) + return d + + +@named_scope +def forward(m: Model, d: Data) -> Data: + """Forward dynamics.""" + d = _position(m, d) + d = _velocity(m, d) + d = _actuation(m, d) + d = _acceleration(m, d) + + if d.efc_J.size == 0: + d = d.replace(qacc=d.qacc_smooth) + return d + + if m.opt.solver == SolverType.CG: + d = named_scope(solver.cg_solve)(m, d) + else: + raise NotImplementedError(f'solver {m.opt.solver} not implemented.') + + return d + + +@named_scope +def step(m: Model, d: Data) -> Data: + """Advance simulation.""" + d = forward(m, d) + + if m.opt.integrator == IntegratorType.EULER: + d = _euler(m, d) + elif m.opt.integrator == IntegratorType.RK4: + d = _rungekutta4(m, d) + else: + raise NotImplementedError(f'integrator {m.opt.integrator} not implemented.') + + return d diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py new file mode 100644 index 00000000..0857db2f --- /dev/null +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -0,0 +1,125 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for forward functions.""" + +import itertools + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import forward +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-3, rtol=1e-3): + err_msg = f'mismatch: {attr} at step {step} in {fname}' + a, b = getattr(a, attr), getattr(b, attr) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +class ForwardTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(test_util.TEST_FILES)) + def test_forward(self, seed, fname): + """Test mujoco mj forward function matches mujoco_mjx forward function.""" + if fname in ('weld.xml',): + return + + np.random.seed(seed) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) * 0.05 + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + + _assert_attr_eq(d, dx, 'qfrc_smooth', i, fname) + _assert_attr_eq(d, dx, 'qacc_smooth', i, fname) + + @parameterized.parameters(itertools.product(test_util.TEST_FILES, (0, 1))) + def test_step(self, fname, integrator_type): + """Test mujoco mj step matches mujoco_mjx step.""" + if fname in ( + 'mixed_joint_pendulum.xml', + 'ball_pendulum.xml', + 'convex.xml', + 'humanoid.xml', + 'triple_pendulum.xml', # TODO(b/301485081) + 'weld.xml', + ): + # skip models with big constraint violations at step 0 or too slow to run + return + + np.random.seed(integrator_type) + m = test_util.load_test_file(fname) + step_jit_fn = jax.jit(forward.step) + + m.opt.integrator = integrator_type + int_typ = 'euler' if integrator_type == 0 else 'rk4' + test_name = f'{fname} - {int_typ}' + steps = 100 if int_typ == 'euler' else 30 + dt = m.opt.timestep + m.opt.timestep = dt if int_typ == 'euler' else dt * 3 + + mx = mjx.device_put(m) + d = mujoco.MjData(m) + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.normal(m.nv) * 0.05 + for i in range(steps): + # in order to avoid re-jitting, reuse the same mj_data shape + qpos, qvel = d.qpos, d.qvel + d = mujoco.MjData(m) + d.qpos, d.qvel = qpos, qvel + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + dx = step_jit_fn(mx, dx) + + _assert_attr_eq(d, dx, 'qpos', i, test_name, atol=1e-2) + _assert_attr_eq(d, dx, 'qvel', i, test_name, atol=1e-2) + _assert_attr_eq(d, dx, 'act', i, test_name) + _assert_attr_eq(d, dx, 'time', i, test_name) + + def test_disable_eulerdamp(self): + m = test_util.load_test_file('ant.xml') + m.opt.disableflags = m.opt.disableflags | DisableBit.EULERDAMP + + d = mujoco.MjData(m) + mx = mjx.device_put(m) + self.assertTrue((mx.dof_damping > 0).any()) + dx = mjx.device_put(d) + dx = jax.jit(forward.forward)(mx, dx) + + dx = dx.replace(qvel=jp.ones_like(dx.qvel), qacc=jp.ones_like(dx.qacc)) + dx = jax.jit(forward._euler)(mx, dx) + np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py new file mode 100644 index 00000000..67e50a7e --- /dev/null +++ b/mjx/mujoco/mjx/_src/io.py @@ -0,0 +1,105 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Functions to initialize, load, or save data.""" + +from jax import numpy as jp +from mujoco.mjx._src import collision_driver +from mujoco.mjx._src import constraint +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Contact +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np + + +def make_data(m: Model) -> Data: + """Allocate and initialize Data.""" + + # create first d to get num contacts and nc + d = Data( + solver_niter=jp.array(0, dtype=jp.int32), + ne=0, + nf=0, + nl=0, + nefc=0, + ncon=0, + time=jp.zeros((), dtype=jp.float32), + qpos=m.qpos0, + qvel=jp.zeros(m.nv, dtype=jp.float32), + act=jp.zeros(m.na, dtype=jp.float32), + qacc_warmstart=jp.zeros(m.nv, dtype=jp.float32), + ctrl=jp.zeros(m.nu, dtype=jp.float32), + qfrc_applied=jp.zeros(m.nv, dtype=jp.float32), + xfrc_applied=jp.zeros((m.nbody, 6), dtype=jp.float32), + eq_active=jp.zeros(m.neq, dtype=jp.int32), + qacc=jp.zeros(m.nv, dtype=jp.float32), + act_dot=jp.zeros(m.na, dtype=jp.float32), + xpos=jp.zeros((m.nbody, 3), dtype=jp.float32), + xquat=jp.zeros((m.nbody, 4), dtype=jp.float32), + xmat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), + xipos=jp.zeros((m.nbody, 3), dtype=jp.float32), + ximat=jp.zeros((m.nbody, 3, 3), dtype=jp.float32), + xanchor=jp.zeros((m.njnt, 3), dtype=jp.float32), + xaxis=jp.zeros((m.njnt, 3), dtype=jp.float32), + geom_xpos=jp.zeros((m.ngeom, 3), dtype=jp.float32), + geom_xmat=jp.zeros((m.ngeom, 3, 3), dtype=jp.float32), + subtree_com=jp.zeros((m.nbody, 3), dtype=jp.float32), + cdof=jp.zeros((m.nv, 6), dtype=jp.float32), + cinert=jp.zeros((m.nbody, 10), dtype=jp.float32), + actuator_length=jp.zeros(m.nu, dtype=jp.float32), + actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32), + crb=jp.zeros((m.nbody, 10), dtype=jp.float32), + qM=jp.zeros(m.nM, dtype=jp.float32), + qLD=jp.zeros(m.nM, dtype=jp.float32), + qLDiagInv=jp.zeros(m.nv, dtype=jp.float32), + qLDiagSqrtInv=jp.zeros(m.nv, dtype=jp.float32), + contact=Contact.zero(), + efc_J=jp.zeros((), dtype=jp.float32), + efc_frictionloss=jp.zeros((), dtype=jp.float32), + efc_D=jp.zeros((), dtype=jp.float32), + actuator_velocity=jp.zeros(m.nu, dtype=jp.float32), + cvel=jp.zeros((m.nbody, 6), dtype=jp.float32), + cdof_dot=jp.zeros((m.nv, 6), dtype=jp.float32), + qfrc_bias=jp.zeros(m.nv, dtype=jp.float32), + qfrc_passive=jp.zeros(m.nv, dtype=jp.float32), + efc_aref=jp.zeros((), dtype=jp.float32), + actuator_force=jp.zeros(m.nu, dtype=jp.float32), + qfrc_actuator=jp.zeros(m.nv, dtype=jp.float32), + qfrc_smooth=jp.zeros(m.nv, dtype=jp.float32), + qacc_smooth=jp.zeros(m.nv, dtype=jp.float32), + qfrc_constraint=jp.zeros(m.nv, dtype=jp.float32), + qfrc_inverse=jp.zeros(m.nv, dtype=jp.float32), + efc_force=jp.zeros((), dtype=jp.float32), + ) + + # get contact data with correct shapes + ncon = collision_driver.ncon(m) + d = d.replace(contact=Contact.zero((ncon,)), ncon=ncon) + d = d.tree_replace({'contact.dim': 3 * np.ones(ncon)}) + + ne, nf, nl, nc = constraint.count_constraints(m, d) + d = d.replace(ne=ne, nf=nf, nl=nl, nefc=ne + nf + nl + nc) + ns = ne + nf + nl + d = d.tree_replace({'contact.efc_address': np.arange(ns, ns + ncon * 4, 4)}) + d = d.replace( + efc_J=jp.zeros((d.nefc, m.nv), dtype=jp.float32), + efc_frictionloss=jp.zeros(d.nefc, dtype=jp.float32), + efc_D=jp.zeros(d.nefc, dtype=jp.float32), + efc_aref=jp.zeros(d.nefc, dtype=jp.float32), + efc_force=jp.zeros(d.nefc, dtype=jp.float32), + ) + + return d diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py new file mode 100644 index 00000000..269509a8 --- /dev/null +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -0,0 +1,42 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for io functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from mujoco import mjx +from mujoco.mjx._src import test_util + + +class IoTest(parameterized.TestCase): + + @parameterized.parameters(test_util.TEST_FILES) + def test_make_data(self, fname): + """Test that data created by make_data matches data returned by step.""" + + m = test_util.load_test_file(fname) + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + dx_step = mjx.step(mx, dx) + + _, dx_treedef = jax.tree_util.tree_flatten(dx) + _, dx_step_treedef = jax.tree_util.tree_flatten(dx_step) + + self.assertEqual(dx_treedef, dx_step_treedef) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/math.py b/mjx/mujoco/mjx/_src/math.py new file mode 100644 index 00000000..e96e6ee0 --- /dev/null +++ b/mjx/mujoco/mjx/_src/math.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Some useful math functions.""" + +from typing import Optional, Tuple, Union + +import jax +from jax import numpy as jp + + +def norm( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> jax.Array: + """Calculates a linalg.norm(x) that's safe for gradients at x=0. + + Avoids a poorly defined gradient for jnp.linal.norm(0) see + https://github.com/google/jax/issues/3058 for details + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + Norm of the array x. + """ + + is_zero = jp.allclose(x, 0.0) + # temporarily swap x with ones if is_zero, then swap back + x = jp.where(is_zero, jp.ones_like(x), x) + n = jp.linalg.norm(x, axis=axis) + n = jp.where(is_zero, 0.0, n) + return n + + +def normalize_with_norm( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> Tuple[jax.Array, jax.Array]: + """Normalizes an array. + + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + A tuple of (normalized array x, the norm). + """ + n = norm(x, axis=axis) + x = x / (n + 1e-6 * (n == 0.0)) + return x, n + + +def normalize( + x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None +) -> jax.Array: + """Normalizes an array. + + Args: + x: A jnp.array + axis: The axis along which to compute the norm + + Returns: + normalized array x + """ + return normalize_with_norm(x, axis=axis)[0] + + +def rotate(vec: jax.Array, quat: jax.Array) -> jax.Array: + """Rotates a vector vec by a unit quaternion quat. + + Args: + vec: (3,) a vector + quat: (4,) a quaternion + + Returns: + ndarray(3) containing vec rotated by quat. + """ + if len(vec.shape) != 1: + raise ValueError('vec must have no batch dimensions.') + s, u = quat[0], quat[1:] + r = 2 * (jp.dot(u, vec) * u) + (s * s - jp.dot(u, u)) * vec + r = r + 2 * s * jp.cross(u, vec) + return r + + +def quat_inv(q: jp.ndarray) -> jp.ndarray: + """Calculates the inverse of quaternion q. + + Args: + q: (4,) quaternion [w, x, y, z] + + Returns: + The inverse of q, where qmult(q, inv_quat(q)) = [1, 0, 0, 0]. + """ + return q * jp.array([1, -1, -1, -1]) + + +def quat_sub(u: jax.Array, v: jax.Array) -> jax.Array: + """Subtracts two quaternions (u - v) as a 3D velocity.""" + q = quat_mul(quat_inv(v), u) + axis, angle = quat_to_axis_angle(q) + return axis * angle + + +def quat_mul(u: jax.Array, v: jax.Array) -> jax.Array: + """Multiplies two quaternions. + + Args: + u: (4,) quaternion (w,x,y,z) + v: (4,) quaternion (w,x,y,z) + + Returns: + A quaternion u * v. + """ + return jp.array([ + u[0] * v[0] - u[1] * v[1] - u[2] * v[2] - u[3] * v[3], + u[0] * v[1] + u[1] * v[0] + u[2] * v[3] - u[3] * v[2], + u[0] * v[2] - u[1] * v[3] + u[2] * v[0] + u[3] * v[1], + u[0] * v[3] + u[1] * v[2] - u[2] * v[1] + u[3] * v[0], + ]) + + +def quat_mul_axis(q: jax.Array, axis: jax.Array) -> jax.Array: + """Multiplies a quaternion and an axis. + + Args: + q: (4,) quaternion (w,x,y,z) + axis: (3,) axis (x,y,z) + + Returns: + A quaternion q * axis + """ + return jp.array([ + -q[1] * axis[0] - q[2] * axis[1] - q[3] * axis[2], + q[0] * axis[0] + q[2] * axis[2] - q[3] * axis[1], + q[0] * axis[1] + q[3] * axis[0] - q[1] * axis[2], + q[0] * axis[2] + q[1] * axis[1] - q[2] * axis[0], + ]) + + +# TODO(erikfrey): benchmark this against brax's quat_to_3x3 +def quat_to_mat(q: jax.Array) -> jax.Array: + """Converts a quaternion into a 9-dimensional rotation matrix.""" + q = jp.outer(q, q) + + return jp.array([ + [ + q[0, 0] + q[1, 1] - q[2, 2] - q[3, 3], + 2 * (q[1, 2] - q[0, 3]), + 2 * (q[1, 3] + q[0, 2]), + ], + [ + 2 * (q[1, 2] + q[0, 3]), + q[0, 0] - q[1, 1] + q[2, 2] - q[3, 3], + 2 * (q[2, 3] - q[0, 1]), + ], + [ + 2 * (q[1, 3] - q[0, 2]), + 2 * (q[2, 3] + q[0, 1]), + q[0, 0] - q[1, 1] - q[2, 2] + q[3, 3], + ], + ]) + + +def quat_to_axis_angle(q: jax.Array) -> Tuple[jax.Array, jax.Array]: + """Converts a quaternion into axis and angle.""" + axis, sin_a_2 = normalize_with_norm(q[1:]) + angle = 2 * jp.arctan2(sin_a_2, q[0]) + angle = jp.where(angle > jp.pi, angle - 2 * jp.pi, angle) + + return axis, angle + + +def axis_angle_to_quat(axis: jax.Array, angle: jax.Array) -> jax.Array: + """Provides a quaternion that describes rotating around axis by angle. + + Args: + axis: (3,) axis (x,y,z) + angle: () float angle to rotate by + + Returns: + A quaternion that rotates around axis by angle + """ + s, c = jp.sin(angle * 0.5), jp.cos(angle * 0.5) + return jp.insert(axis * s, 0, c) + + +def quat_integrate(q: jax.Array, v: jax.Array, dt: jax.Array) -> jax.Array: + """Integrates a quaternion given angular velocity and dt.""" + v, norm_ = normalize_with_norm(v) + angle = dt * norm_ + q_res = axis_angle_to_quat(v, angle) + q_res = quat_mul(q, q_res) + return normalize(q_res) + + +def inert_mul(i: jax.Array, v: jax.Array) -> jax.Array: + """Multiply inertia by motion, producing force. + + Args: + i: (10,) inertia (inertia matrix, position, mass) + v: (6,) spatial motion + + Returns: + resultant force + """ + tri_id = jp.array([[0, 3, 4], [3, 1, 5], [4, 5, 2]]) # cinert inr order + inr, pos, mass = i[tri_id], i[6:9], i[9] + ang = jp.dot(inr, v[:3]) + jp.cross(pos, v[3:]) + vel = mass * v[3:] - jp.cross(pos, v[:3]) + return jp.concatenate((ang, vel)) + + +def transform_motion(vel: jax.Array, offset: jax.Array, rotmat: jax.Array): + """Transform spatial motion. + + Args: + vel: (6,) spatial motion (3 angular, 3 linear) + offset: (3,) translation + rotmat: (3, 3) rotation + + Returns: + 6d spatial velocity + """ + # TODO(robotics-simulation): are quaternions faster here + ang, vel = vel[:3], vel[3:] + vel = rotmat.T @ (vel - jp.cross(offset, ang)) + ang = rotmat.T @ ang + return jp.concatenate([ang, vel]) + + +def motion_cross(u, v): + """Cross product of two motions. + + Args: + u: (6,) spatial motion + v: (6,) spatial motion + + Returns: + resultant spatial motion + """ + ang = jp.cross(u[:3], v[:3]) + vel = jp.cross(u[3:], v[:3]) + jp.cross(u[:3], v[3:]) + return jp.concatenate((ang, vel)) + + +def motion_cross_force(v, f): + """Cross product of a motion and force. + + Args: + v: (6,) spatial motion + f: (6,) force + + Returns: + resultant force + """ + ang = jp.cross(v[:3], f[:3]) + jp.cross(v[3:], f[3:]) + vel = jp.cross(v[:3], f[3:]) + return jp.concatenate((ang, vel)) + + +def orthogonals(a: jax.Array) -> Tuple[jax.Array, jax.Array]: + """Returns orthogonal vectors `b` and `c`, given a vector `a`.""" + y, z = jp.array([0, 1, 0]), jp.array([0, 0, 1]) + b = jp.where((-0.5 < a[1]) & (a[1] < 0.5), y, z) + b = b - a * a.dot(b) + # normalize b. however if a is a zero vector, zero b as well. + b = normalize(b) * jp.any(a) + return b, jp.cross(a, b) + + +def make_frame(a: jax.Array) -> jax.Array: + """Makes a right-handed 3D frame given a direction.""" + a = normalize(a) + b, c = orthogonals(a) + return jp.array([a, b, c]) + + +# Geometry. + + +def closest_segment_point( + a: jax.Array, b: jax.Array, pt: jax.Array +) -> jax.Array: + """Returns the closest point on the a-b line segment to a point pt.""" + ab = b - a + t = jp.dot(pt - a, ab) / (jp.dot(ab, ab) + 1e-6) + return a + jp.clip(t, 0.0, 1.0) * ab + + +def closest_segment_point_and_dist( + a: jax.Array, b: jax.Array, pt: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Returns closest point on the line segment and the distance squared.""" + closest = closest_segment_point(a, b, pt) + dist = (pt - closest).dot(pt - closest) + return closest, dist + + +def closest_segment_to_segment_points( + a0: jax.Array, a1: jax.Array, b0: jax.Array, b1: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Returns closest points between two line segments.""" + # Gets the closest segment points by first finding the closest points + # between two lines. Points are then clipped to be on the line segments + # and edge cases with clipping are handled. + dir_a, len_a = normalize_with_norm(a1 - a0) + dir_b, len_b = normalize_with_norm(b1 - b0) + + # Segment mid-points. + half_len_a = len_a * 0.5 + half_len_b = len_b * 0.5 + a_mid = a0 + dir_a * half_len_a + b_mid = b0 + dir_b * half_len_b + + # Translation between two segment mid-points. + trans = a_mid - b_mid + + # Parametrize points on each line as follows: + # point_on_a = a_mid + t_a * dir_a + # point_on_b = b_mid + t_b * dir_b + # and analytically minimize the distance between the two points. + dira_dot_dirb = dir_a.dot(dir_b) + dira_dot_trans = dir_a.dot(trans) + dirb_dot_trans = dir_b.dot(trans) + denom = 1 - dira_dot_dirb * dira_dot_dirb + + orig_t_a = (-dira_dot_trans + dira_dot_dirb * dirb_dot_trans) / (denom + 1e-6) + orig_t_b = dirb_dot_trans + orig_t_a * dira_dot_dirb + t_a = jp.clip(orig_t_a, -half_len_a, half_len_a) + t_b = jp.clip(orig_t_b, -half_len_b, half_len_b) + + best_a = a_mid + dir_a * t_a + best_b = b_mid + dir_b * t_b + + # Resolve edge cases where both closest points are clipped to the segment + # endpoints by recalculating the closest segment points for the current + # clipped points, and then picking the pair of points with smallest + # distance. An example of this edge case is when lines intersect but line + # segments don't. + new_a, d1 = closest_segment_point_and_dist(a0, a1, best_b) + new_b, d2 = closest_segment_point_and_dist(b0, b1, best_a) + best_a = jp.where(d1 < d2, new_a, best_a) + best_b = jp.where(d1 < d2, best_b, new_b) + + return best_a, best_b diff --git a/mjx/mujoco/mjx/_src/math_test.py b/mjx/mujoco/mjx/_src/math_test.py new file mode 100644 index 00000000..c2678321 --- /dev/null +++ b/mjx/mujoco/mjx/_src/math_test.py @@ -0,0 +1,202 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for math.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax.numpy as jp +from mujoco.mjx._src import math +import numpy as np + + +def _get_rand_point(seed=None): + if seed is not None: + np.random.seed(seed) + verts = np.random.randn(1, 3) + return verts[0, :] + + +def _get_rand_line_segment(seed=None): + if seed is not None: + np.random.seed(seed) + verts = np.random.randn(2, 3) + return verts[0, :], verts[1, :] + + +def _get_rand_unit(seed: int): + np.random.seed(seed) + theta = np.random.random(1) * 2 * np.pi + a = (np.random.random(1) - 0.5) * 2.0 + phi = np.arccos(a) + x = np.sin(phi) * np.cos(theta) + y = np.sin(phi) * np.sin(theta) + z = np.cos(phi) + return jp.array([x, y, z]).squeeze() + + +class OrthoganalsTest(parameterized.TestCase): + """Tests the orthogonals function.""" + + @parameterized.parameters(range(30)) + def test_orthogonals(self, i): + a = _get_rand_unit(i) + b, c = math.orthogonals(a) + np.testing.assert_almost_equal(jp.linalg.norm(a), 1) + np.testing.assert_almost_equal(jp.linalg.norm(b), 1) + np.testing.assert_almost_equal(jp.linalg.norm(c), 1) + self.assertAlmostEqual(np.abs(a.dot(b)), 0, 6) + self.assertAlmostEqual(np.abs(b.dot(c)), 0, 6) + self.assertAlmostEqual(np.abs(a.dot(c)), 0, 6) + + +def _minimize(fn, sample_fn, lb, ub, tol, max_iter=20, seed=42): + """Minimize a function using the cross-entropy method.""" + assert lb.shape == ub.shape, "bounds need to have the same shape" + np.random.seed(seed) + + i, n = 0, 1_000 + mu = (ub + lb) * 0.5 + sigma = (ub - lb) * 0.5 + size = lb.shape[0] + val, prev_val = fn(mu), None + + while prev_val is None or np.abs(val - prev_val) > tol: + params = sample_fn(mu, sigma, n, size, lb, ub) + vals = np.array([fn(p) for p in params]) + if val < vals.min(): # early exit + return mu + idx = vals.argsort() + best_idx = idx[: int(n * 0.05)] + mu = params[best_idx].mean(axis=0) + sigma = params[best_idx].std(axis=0) + 1e-10 + + prev_val = val + val = fn(mu) + + i += 1 + if i == max_iter: + break + + return mu + + +def _closest_segment_to_segment_points(a0, a1, b0, b1): + dir_a = a1 - a0 + len_a = np.sqrt(dir_a.dot(dir_a)) + half_len_a = len_a / 2 + dir_a = dir_a / len_a + + dir_b = b1 - b0 + len_b = np.sqrt(dir_b.dot(dir_b)) + half_len_b = len_b / 2 + dir_b = dir_b / len_b + + a_mid = a0 + dir_a * half_len_a + b_mid = b0 + dir_b * half_len_b + + # Parametrize both line segments. + def fn(t): + best_a = a_mid + dir_a * t[0] + best_b = b_mid + dir_b * t[1] + return (best_a - best_b).dot(best_a - best_b) + + def sample_fn(mu, sigma, n, size, lb, ub): + params = np.random.normal(mu, sigma, size=(n, size)) + params = np.clip(params, lb, ub) + return params + + lb = np.array([-half_len_a, -half_len_b]) + ub = np.array([half_len_a, half_len_b]) + ta, tb = _minimize(fn, sample_fn, lb, ub, tol=1e-4) + best_a = a_mid + dir_a * ta + best_b = b_mid + dir_b * tb + return best_a, best_b + + +class ClosestSegmentSegmentPointsTest(parameterized.TestCase): + """Tests for closest segment-to-segment points.""" + + def test_closest_segments_points(self): + a0 = jp.array([0.73432405, 0.12372768, 0.20272314]) + a1 = jp.array([1.10600128, 0.88555209, 0.65209485]) + b0 = jp.array([0.85599262, 0.61736299, 0.9843583]) + b1 = jp.array([1.84270939, 0.92891793, 1.36343326]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [1.09063, 0.85404, 0.63351], 5) + self.assertSequenceAlmostEqual(best_b, [0.99596, 0.66156, 1.03813], 5) + + def test_intersecting_segments(self): + """Tests segments that intersect.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([-1.0, 0.0, 0.0]), jp.array([1.0, 0.0, 0.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5) + + def test_intersecting_lines(self): + """Tests that intersecting lines get clipped.""" + a0, a1 = jp.array([0.2, 0.2, 0.0]), jp.array([1.0, 1.0, 0.0]) + b0, b1 = jp.array([0.2, 0.4, 0.0]), jp.array([1.0, 2.0, 0.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.3, 0.3, 0.0], 2) + self.assertSequenceAlmostEqual(best_b, [0.2, 0.4, 0.0], 2) + + def test_parallel_segments(self): + """Tests that parallel segments have closest points at the midpoint.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([1.0, 0.0, -1.0]), jp.array([1.0, 0.0, 1.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.0], 5) + + def test_parallel_offset_segments(self): + """Tests that offset parallel segments are close at segment endpoints.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([1.0, 0.0, 1.0]), jp.array([1.0, 0.0, 3.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 1.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 1.0], 5) + + def test_zero_length_segments(self): + """Test that zero length segments don't return NaNs.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, -1.0]) + b0, b1 = jp.array([1.0, 0.0, 0.1]), jp.array([1.0, 0.0, 0.1]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, -1.0], 5) + self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.1], 5) + + def test_overlapping_segments(self): + """Tests that perfectly overlapping segments intersect at the midpoints.""" + a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + b0, b1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0]) + best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1) + self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5) + self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5) + + params = list(zip(np.repeat(np.arange(10), 10), np.tile(np.arange(10), 10))) + + @parameterized.parameters(*params) + def test_closest_segment_to_segment_points(self, i, j): + a0, a1 = _get_rand_line_segment(i) + b0, b1 = _get_rand_line_segment(j) + expected = _closest_segment_to_segment_points(a0, a1, b0, b1) + ans = math.closest_segment_to_segment_points(a0, a1, b0, b1) + expected_dist = (expected[0] - expected[1]).dot(expected[0] - expected[1]) + test_dist = (ans[0] - ans[1]).dot(ans[0] - ans[1]) + self.assertAlmostEqual(expected_dist, test_dist, 4) + + +if __name__ == "__main__": + absltest.main() diff --git a/mjx/mujoco/mjx/_src/mesh.py b/mjx/mujoco/mjx/_src/mesh.py new file mode 100644 index 00000000..eaa5d3bf --- /dev/null +++ b/mjx/mujoco/mjx/_src/mesh.py @@ -0,0 +1,232 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Mesh processing.""" + +import itertools +from typing import Dict, Optional, Sequence, Tuple + +import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import GeomType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member +import numpy as np +from scipy import spatial +import trimesh + + +_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1))) +# pyformat: disable +# Rectangular box faces using a counter-clockwise winding order convention. +_BOX_FACES = [ + 0, 4, 5, 1, # left + 0, 2, 6, 4, # bottom + 6, 7, 5, 4, # front + 2, 3, 7, 6, # right + 1, 5, 7, 3, # top + 0, 1, 3, 2, # back +] +# pyformat: enable +_MAX_HULL_FACE_VERTICES = 20 +_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {} +_DERIVED_ARGS = [ + 'geom_convex_face', + 'geom_convex_vert', + 'geom_convex_edge', + 'geom_convex_facenormal', +] +DERIVED = {(Model, d) for d in _DERIVED_ARGS} + + +def _box(size: np.ndarray): + """Creates a mesh for a box with rectangular faces.""" + box_corners = np.array(_BOX_CORNERS) + vert = box_corners * size.reshape(-1, 3) + face = np.array([_BOX_FACES]).reshape(-1, 4) + return vert, face + + +def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray: + """Calculates face normals given vertices and face indexes.""" + assert len(vert.shape) == 2 and len(face.shape) == 2, ( + f'vert and face should have dim of 2, got {len(vert.shape)} and ' + f'{len(face.shape)}' + ) + face_vert = vert[face, :] + # use CCW winding order convention + edge0 = face_vert[:, 1, :] - face_vert[:, 0, :] + edge1 = face_vert[:, -1, :] - face_vert[:, 0, :] + face_norm = np.cross(edge0, edge1) + face_norm = face_norm / np.linalg.norm(face_norm, axis=1).reshape((-1, 1)) + return face_norm + + +def _get_unique_edges(vert: np.ndarray, face: np.ndarray) -> np.ndarray: + """Returns unique edges. + + Args: + vert: (n_vert, 3) vertices + face: (n_face, n_vert) face index array + + Returns: + edges: tuples of vertex indexes for each edge + """ + r_face = np.roll(face, 1, axis=1) + edges = np.concatenate(np.array([face, r_face]).T) + + # do a first pass to remove duplicates + edges.sort(axis=1) + edges = np.unique(edges, axis=0) + edges = edges[edges[:, 0] != edges[:, 1]] # get rid of edges from padded face + + # get normalized edge directions + edge_vert = vert.take(edges, axis=0) + edge_dir = edge_vert[:, 0] - edge_vert[:, 1] + norms = np.sqrt(np.sum(edge_dir**2, axis=1)) + edge_dir = edge_dir / norms.reshape((-1, 1)) + + # get the first unique edge for all pairwise comparisons + diff1 = edge_dir[:, None, :] - edge_dir[None, :, :] + diff2 = edge_dir[:, None, :] + edge_dir[None, :, :] + matches = (np.linalg.norm(diff1, axis=-1) < 1e-6) | ( + np.linalg.norm(diff2, axis=-1) < 1e-6 + ) + matches = np.tril(matches).sum(axis=-1) + unique_edge_idx = np.where(matches == 1)[0] + + return edges[unique_edge_idx] + + +def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray: + """Calculates the convex hull for a set of points on a plane.""" + # project points onto the closest axis plane + best_axis = np.abs(np.eye(3).dot(normal)).argmax() + axis = np.eye(3)[best_axis] + d = points.dot(axis).reshape((-1, 1)) + axis_points = points - d * axis + axis_points = axis_points[:, list({0, 1, 2} - {best_axis})] + + # get the polygon face, and make the points ccw wrt the face normal + c = spatial.ConvexHull(axis_points) + order_ = np.where(axis.dot(normal) > 0, 1, -1) + order_ *= np.where(best_axis == 1, -1, 1) + hull_point_idx = c.vertices[::order_] + assert (axis_points - c.points).sum() == 0 + + return hull_point_idx + + +def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray: + """Merges coplanar facets.""" + if not tm.facets: + return tm.faces.copy() # no facets + if not tm.faces.shape[0]: + raise ValueError('Mesh has no faces.') + + # Get faces. + face_idx = set(range(tm.faces.shape[0])) - set(np.concatenate(tm.facets)) + face_idx = np.array(list(face_idx)) + faces = tm.faces[face_idx] if face_idx.shape[0] > 0 else np.array([]) + + # Get facets. + facets = [] + for i, facet in enumerate(tm.facets): + point_idx = np.unique(tm.faces[facet]) + points = tm.vertices[point_idx] + normal = tm.facets_normal[i] + + # convert triangulated facet to a polygon + hull_point_idx = _convex_hull_2d(points, normal) + face = point_idx[hull_point_idx] + + # resize faces that exceed max polygon vertices + every = face.shape[0] // _MAX_HULL_FACE_VERTICES + 1 + face = face[::every] + facets.append(face) + + # Pad facets so that they can be stacked. + max_len = max(f.shape[0] for f in facets) if facets else faces.shape[1] + assert max_len <= _MAX_HULL_FACE_VERTICES + for i, f in enumerate(facets): + if f.shape[0] < max_len: + f = np.pad(f, (0, max_len - f.shape[0]), 'edge') + facets[i] = f + + if not faces.shape[0]: + assert facets + return np.array(facets) # no faces, return facets + + # Merge faces and facets. + faces = np.pad(faces, ((0, 0), (0, max_len - faces.shape[1])), 'edge') + return np.concatenate([faces, facets]) + + +def _get_faces_verts( + m: mujoco.MjModel, +) -> Tuple[Sequence[np.ndarray], Sequence[np.ndarray]]: + """Extracts mesh faces and vertices from MjModel.""" + verts, faces = [], [] + for i in range(m.nmesh): + last = (i + 1) >= m.nmesh + face_start = m.mesh_faceadr[i] + face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0] + face = m.mesh_face[face_start:face_end] + faces.append(face) + + vert_start = m.mesh_vertadr[i] + vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0] + vert = m.mesh_vert[vert_start:vert_end] + verts.append(vert) + return verts, faces + + +def _geom_mesh_kwargs( + vert: np.ndarray, face: np.ndarray +) -> Dict[str, np.ndarray]: + """Generates convex mesh attributes for mjx.Model.""" + tm = trimesh.Trimesh(vertices=vert, faces=face) + tm_convex = trimesh.convex.convex_hull(tm) + vert = np.array(tm_convex.vertices) + face = _merge_coplanar(tm_convex) + return { + 'geom_convex_face': face, + 'geom_convex_vert': vert, + 'geom_convex_edge': _get_unique_edges(vert, face), + 'geom_convex_facenormal': _get_face_norm(vert, face), + } + + +def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]: + """Derives geom mesh attributes for mjx.Model from MjModel.""" + kwargs = {k: [] for k in _DERIVED_ARGS} + verts, faces = _get_faces_verts(m) + for geomid in range(m.ngeom): + dataid = m.geom_dataid[geomid] + typ = m.geom_type[geomid] + if typ == GeomType.BOX: + vert, face = _box(m.geom_size[geomid]) + elif dataid >= 0: + vert, face = verts[dataid], faces[dataid] + else: + kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS} + continue + + key = (hash(vert.data.tobytes()), hash(face.data.tobytes())) + if key not in _CONVEX_CACHE: + _CONVEX_CACHE[key] = _geom_mesh_kwargs(vert, face) + + kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS} + + return kwargs diff --git a/mjx/mujoco/mjx/_src/mesh_test.py b/mjx/mujoco/mjx/_src/mesh_test.py new file mode 100644 index 00000000..7be8197a --- /dev/null +++ b/mjx/mujoco/mjx/_src/mesh_test.py @@ -0,0 +1,119 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for mesh.py.""" + +from absl.testing import absltest +from mujoco.mjx._src import mesh +import numpy as np + + +class GeomMeshKwargsTest(absltest.TestCase): + + def test_pyramid(self): + """Tests that a triangulated pyramid converts to merged coplanar faces.""" + vert = np.array([ + [-0.025, 0.05, 0.05], + [-0.025, -0.05, -0.05], + [-0.025, -0.05, 0.05], + [-0.025, 0.05, -0.05], + [0.075, 0.0, 0.0], + ]) + face = np.array( + [[0, 1, 2], [0, 3, 1], [0, 4, 3], [0, 2, 4], [2, 1, 4], [1, 3, 4]] + ) + h = mesh._geom_mesh_kwargs(vert, face) + + # get index of vertices in h['geom_convex_vert'] for vertices in vert + dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile( + h['geom_convex_vert'], (vert.shape[0], 1) + ) + dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1)) + vidx = np.argmin(dist, axis=0) + + # check verts + np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx]) + + # check face vertices + map_ = {v: k for k, v in enumerate(vidx)} + h_face = np.vectorize(map_.get)(h['geom_convex_face']) + face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()]) + expected_face_verts = sorted([ + (0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)]) + self.assertSequenceEqual( + face_verts, + expected_face_verts, + ) + + # check edges + unique_edge = np.vectorize(map_.get)(h['geom_convex_edge']) + unique_edge = np.array(sorted(unique_edge.tolist())) + np.testing.assert_array_equal( + unique_edge, + np.array([[0, 2], [0, 3], [0, 4], [1, 4], [2, 4], [3, 4]]), + ) + + # face normals + self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3)) + + +class ConvexHull2DTest(absltest.TestCase): + + def test_convex_hull_2d_axis1(self): + """Tests for the correct winding order of a polgyon with +y normal.""" + pts = np.array([ + [-0.04634297, -0.06652775, 0.05853534], + [-0.01877651, -0.08309858, -0.05236476], + [0.02362804, -0.08010745, 0.05499557], + [0.04066505, -0.09034877, -0.01354446], + [-0.07255043, -0.06837638, -0.00781699], + ]) + normal = np.array([-0.18467607, -0.97768016, 0.10018111]) + idx = mesh._convex_hull_2d(pts, normal) + expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0]) + expected /= np.linalg.norm(expected) + np.testing.assert_array_almost_equal(normal, expected) + + def test_convex_hull_2d_axis2(self): + """Tests for the correct winding order for a polgyon with +z normal.""" + pts = np.array([ + [0.08607829, -0.03881998, -0.03291714], + [-0.01877651, -0.08309858, -0.05236476], + [0.05470364, 0.00027677, -0.08371042], + [-0.01010019, -0.02708892, -0.0957297], + [0.04066505, -0.09034877, -0.01354446], + ]) + normal = np.array([0.3839915, -0.60171936, -0.70034587]) + idx = mesh._convex_hull_2d(pts, normal) + expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0]) + expected /= np.linalg.norm(expected) + np.testing.assert_array_almost_equal(normal, expected) + + +class UniqueEdgesTest(absltest.TestCase): + + def test_tetrahedron_edges(self): + """Tests unique edges for a tetrahedron.""" + vert = np.array( + [[-0.1, 0.0, -0.1], [0.0, 0.1, 0.1], [0.1, 0.0, -0.1], [0.0, -0.1, 0.1]] + ) + face = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 1], [2, 1, 3]]) + idx = mesh._get_unique_edges(vert, face) + np.testing.assert_array_equal( + idx, np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]]) + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/passive.py b/mjx/mujoco/mjx/_src/passive.py new file mode 100644 index 00000000..9b0b41f0 --- /dev/null +++ b/mjx/mujoco/mjx/_src/passive.py @@ -0,0 +1,141 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Passive forces.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import math +from mujoco.mjx._src import scan +from mujoco.mjx._src import support +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def _inertia_box_fluid_model( + m: Model, + inertia: jax.Array, + mass: jax.Array, + root_com: jax.Array, + xipos: jax.Array, + ximat: jax.Array, + cvel: jax.Array, +) -> Tuple[jax.Array, jax.Array]: + """Fluid forces based on inertia-box approximation.""" + box = jp.repeat(inertia[None, :], 3, axis=0) + box *= jp.ones((3, 3)) - 2 * jp.eye(3) + box = 6.0 * jp.clip(jp.sum(box, axis=-1), a_min=1e-12) + box = jp.sqrt(box / jp.maximum(mass, 1e-12)) * (mass > 0.0) + + # transform to local coordinate frame + offset = xipos - root_com + lvel = math.transform_motion(cvel, offset, ximat) + lwind = ximat.T @ m.opt.wind + lvel = lvel.at[3:].add(-lwind) + + # set viscous force and torque + diam = jp.mean(box, axis=-1) + lfrc_ang = lvel[:3] * -jp.pi * diam**3 * m.opt.viscosity + lfrc_vel = lvel[3:] * -3.0 * jp.pi * diam * m.opt.viscosity + + # add lift and drag force and torque + scale_vel = jp.array([box[1] * box[2], box[0] * box[2], box[0] * box[1]]) + scale_ang = jp.array([ + box[0] * (box[1] ** 4 + box[2] ** 4), + box[1] * (box[0] ** 4 + box[2] ** 4), + box[2] * (box[0] ** 4 + box[1] ** 4), + ]) + lfrc_vel -= 0.5 * m.opt.density * scale_vel * jp.abs(lvel[3:]) * lvel[3:] + lfrc_ang -= ( + 1.0 * m.opt.density * scale_ang * jp.abs(lvel[:3]) * lvel[:3] / 64.0 + ) + + # rotate to global orientation: lfrc -> bfrc + force, torque = ximat @ lfrc_vel, ximat @ lfrc_ang + + return force, torque + + +def passive(m: Model, d: Data) -> Data: + """Adds all passive forces.""" + if m.opt.disableflags & DisableBit.PASSIVE: + return d + + # joint-level springs + def fn(jnt_typs, stiffness, qpos_spring, qpos): + qpos_i = 0 + qfrcs = [] + for i in range(len(jnt_typs)): + jnt_typ = JointType(jnt_typs[i]) + q = qpos[qpos_i : qpos_i + jnt_typ.qpos_width()] + qs = qpos_spring[qpos_i : qpos_i + jnt_typ.qpos_width()] + qfrc = jp.zeros(jnt_typ.dof_width()) + if jnt_typ == JointType.FREE: + qfrc = qfrc.at[:3].set(-stiffness[i] * (q[:3] - qs[:3])) + qfrc = qfrc.at[3:6].set(-stiffness[i] * math.quat_sub(q[3:7], qs[3:7])) + elif jnt_typ == JointType.BALL: + qfrc = -stiffness[i] * math.quat_sub(q, qs) + elif jnt_typ in ( + JointType.SLIDE, + JointType.HINGE, + ): + qfrc = -stiffness[i] * (q - qs) + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + qfrcs.append(qfrc) + qpos_i += jnt_typ.qpos_width() + return jp.concatenate(qfrcs) + + qfrc_passive = scan.flat( + m, + fn, + 'jjqq', + 'v', + m.jnt_type, + m.jnt_stiffness, + m.qpos_spring, + d.qpos, + ) + + # dof-level dampers + qfrc_passive -= m.dof_damping * d.qvel + + # TODO(robotics-simulation): body-level gravity compensation + + # body-level viscosity, lift and drag + if m.opt.has_fluid_params: + force, torque = jax.vmap( + _inertia_box_fluid_model, in_axes=(None, 0, 0, 0, 0, 0, 0) + )( + m, + m.body_inertia, + m.body_mass, + d.subtree_com[jp.array(m.body_rootid)], + d.xipos, + d.ximat, + d.cvel, + ) + qfrc_target = jax.vmap(support.apply_ft, in_axes=(None, None, 0, 0, 0, 0))( + m, d, force, torque, d.xipos, jp.arange(m.nbody) + ) + qfrc_passive += jp.sum(qfrc_target, axis=0) + + d = d.replace(qfrc_passive=qfrc_passive) + return d diff --git a/mjx/mujoco/mjx/_src/passive_test.py b/mjx/mujoco/mjx/_src/passive_test.py new file mode 100644 index 00000000..49264014 --- /dev/null +++ b/mjx/mujoco/mjx/_src/passive_test.py @@ -0,0 +1,114 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests passive forces.""" + +import itertools + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import jax.numpy as jp +import mujoco +from mujoco import mjx +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5): + err_msg = f'mismatch: {attr} at step {step} in {fname}' + a, b = getattr(a, attr), getattr(b, attr) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +class PassiveTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(('ant.xml', 'mixed_joint_pendulum.xml'))) + def test_stiffness_damping(self, seed, fname): + """Tests stiffness and damping on Ant.""" + np.random.seed(seed) + path = epath.resource_path('mujoco.mjx') / 'test_data' + path /= fname + m = mujoco.MjModel.from_xml_string(path.read_text()) + + # set stiffness/damping + m.jnt_stiffness = np.random.uniform(size=m.njnt) + m.dof_damping = np.random.uniform(size=m.nv) + d = mujoco.MjData(m) + d.qvel = np.random.random(m.nv) # random kick + + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + passive_jit_fn = jax.jit(mjx.passive) + + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + dx = passive_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) + + @parameterized.parameters( + itertools.product(range(3), ('triple_pendulum.xml',)) + ) + def test_fluid(self, seed, fname): + np.random.seed(seed) + path = epath.resource_path('mujoco.mjx') / 'test_data' + path /= fname + m = mujoco.MjModel.from_xml_string(path.read_text()) + + # set density/viscosity/wind + m.opt.density = np.random.uniform() + m.opt.viscosity = np.random.uniform() + m.opt.wind = np.random.uniform() + + passive_jit_fn = jax.jit(mjx.passive) + + mx = mjx.device_put(m) + d = mujoco.MjData(m) + d.qvel = np.random.random(m.nv) # random kick + + for i in range(100): + mujoco.mj_step(m, d) + dx = mjx.device_put(d) + mujoco.mj_passive(m, d) + dx = passive_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'qfrc_passive', i, fname) + + def test_disable_passive(self): + m = mujoco.MjModel.from_xml_string(""" + + + + + + + + + + """) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + dx = dx.replace(qvel=jp.ones(mx.nv)) + + passive_jit_fn = jax.jit(mjx.passive) + dx = passive_jit_fn(mx, dx) + np.testing.assert_equal(dx.qfrc_passive, np.zeros(mx.nv)) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/scan.py b/mjx/mujoco/mjx/_src/scan.py new file mode 100644 index 00000000..2975a32f --- /dev/null +++ b/mjx/mujoco/mjx/_src/scan.py @@ -0,0 +1,432 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Scan across data ordered by body joint types and kinematic tree order.""" + +from typing import Any, Callable, TypeVar + +import jax +from jax import numpy as jp +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +from mujoco.mjx._src.types import TrnType +# pylint: enable=g-importing-member +import numpy as np + + +Y = TypeVar('Y') + + +# TODO(erikfrey): re-check if this really helps perf +def _take(obj: Y, idx: np.ndarray) -> Y: + """Takes idxs on any pytree given to it. + + XLA executes x[jp.array([1, 2, 3])] slower than x[1:4], so we detect when + take indices are contiguous, and convert them to slices. + + Args: + obj: an input pytree + idx: indices to take + + Returns: + obj pytree with leaves taken by idxs + """ + + if isinstance(obj, np.ndarray): + return obj[idx] + + def take(x): + # TODO(erikfrey): if this helps perf, add support for striding too + if ( + len(idx.shape) == 1 + and idx.size > 0 + and (idx == np.arange(idx[0], idx[0] + idx.size)).all() + and (idx > 0).all() + ): + x = x[idx[0] : idx[-1] + 1] + else: + x = x.take(jp.array(idx), axis=0, mode='wrap') + return x + + return jax.tree_map(take, obj) + + +def _q_bodyid(m: Model) -> np.ndarray: + """Returns the bodyid for each qpos adress.""" + q_bodyids = [np.array([], dtype=np.int32)] + for jnt_type, jnt_bodyid in zip(m.jnt_type, m.jnt_bodyid): + width = {JointType.FREE: 7, JointType.BALL: 4}.get(jnt_type, 1) + q_bodyids.append(np.repeat(jnt_bodyid, width)) + return np.concatenate(q_bodyids) + + +def _q_jointid(m: Model) -> np.ndarray: + """Returns the jointid for each qpos adress.""" + q_jointid = [np.array([], dtype=np.int32)] + for i, jnt_type in enumerate(m.jnt_type): + width = {JointType.FREE: 7, JointType.BALL: 4}.get(jnt_type, 1) + q_jointid.append(np.repeat(i, width)) + return np.concatenate(q_jointid) + + +def _index(haystack: np.ndarray, needle: np.ndarray) -> np.ndarray: + """Returns indexes in haystack for elements in needle.""" + idx = np.argsort(haystack) + sorted_haystack = haystack[idx] + sorted_idx = np.searchsorted(sorted_haystack, needle) + idx = np.take(idx, sorted_idx, mode='clip') + idx[haystack[idx] != needle] = -1 + + return idx + + +def _nvmap(f: Callable[..., Y], *args) -> Y: + """A vmap that accepts numpy arrays. + + Numpy arrays are statically vmapped, and the elements are passed to f as + static arguments. The implication is that all the elements of numpy array + arguments must be the same. + + Args: + f: function to be mapped over + *args: args to be mapped along, passed to f + + Returns: + the result of vmapping f over args + + Raises: + RuntimeError: if numpy arg elements do not match + """ + for arg in args: + if isinstance(arg, np.ndarray) and not np.all(arg == arg[0]): + raise RuntimeError(f'numpy arg elements do not match: {arg}') + + np_args = [a[0] if isinstance(a, np.ndarray) else None for a in args] + args = [a if n is None else None for n, a in zip(np_args, args)] + in_axes = [None if a is None else 0 for a in args] + + def outer_f(*args, np_args=np_args): + args = [a if n is None else n for n, a in zip(args, np_args)] + return f(*args) + + return jax.vmap(outer_f, in_axes=in_axes)(*args) + + +def _check_input(m: Model, args: Any, in_types: str) -> None: + """Checks that scan input has the right shape.""" + size = {'b': m.nbody, 'j': m.njnt, 'q': m.nq, 'v': m.nv, 'u': m.nu, 'a': m.na} + for idx, (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], + ) + + def type_ids_j(m, i): + return { + 'b': i, + 'j': np.nonzero(m.jnt_bodyid == i)[0], + 'v': np.nonzero(m.dof_bodyid == i)[0], + 'q': np.nonzero(_q_bodyid(m) == i)[0], + } + + def type_ids_u(m, i): + typ_ids = { + 'u': i, + 'a': m.actuator_actadr[i], + 'j': ( + m.actuator_trnid[i] + if m.actuator_trntype[i] == TrnType.JOINT + else np.array(-1) + ), + } + # v/q associated with joint transmissions + typ_ids.update({ + 'v': np.nonzero(m.dof_jntid == typ_ids['j'])[0], + 'q': np.nonzero(_q_jointid(m) == typ_ids['j'])[0], + }) + return typ_ids + + # build up a grouping of type take-ids in body/actuator order + key_typ_ids, order = {}, [] + all_types = set(in_types + out_types) + n_items = {'j': m.nbody, 'u': m.nu}[group_by] + for i in np.arange(n_items, dtype=np.int32): + typ_ids = type_ids_j(m, i) if group_by == 'j' else type_ids_u(m, i) + + # create grouping key + key = ( + key_j(typ_ids['j']) + if group_by == 'j' + else key_u(typ_ids['u'], typ_ids['j']) + ) + order.append((key, typ_ids)) + + # add ids per type to the corresponding group + for t in all_types: + out = key_typ_ids.setdefault(key, {}) + val = np.expand_dims(typ_ids[t], axis=0) + out[t] = np.concatenate((out[t], val)) if t in out else val + + key_typ_ids = list(sorted(key_typ_ids.items())) + + # use this grouping to take the right data subsets and call vmap(f) + ys = [] + for _, typ_ids in key_typ_ids: + # only execute f if we would actually take something from the result + if any(typ_ids[v].size > 0 for v in out_types): + f_args = [_take(arg, typ_ids[typ]) for arg, typ in zip(args, in_types)] + y = _nvmap(f, *f_args) + ys.append(y) + else: + ys.append(None) + + # remove None results from the final output + key_typ_ids = [v for y, v in zip(ys, key_typ_ids) if y is not None] + ys = [y for y in ys if y is not None] + ys_keys = set([k for k, *_ in key_typ_ids]) + order = [o for k, o in order if k in ys_keys] + + # get the original input order + order = [[o[t] for o in order] for t in all_types] + order = [ + np.concatenate(o) if isinstance(o[0], np.ndarray) else np.array(o) + for o in order + ] + order = dict(zip(all_types, order)) + + # concatenate back to a single tree and drop the grouping dimension + f_ret_is_seq = isinstance(ys[0], (list, tuple)) + ys = ys if f_ret_is_seq else [[y] for y in ys] + flat_ = {'j': 'b', 'u': 'uaj'}[group_by] + ys = [ + [v if typ in flat_ else jp.concatenate(v) for v, typ in zip(y, out_types)] + for y in ys + ] + ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys) + + # put concatenated results back in order + reordered_ys = [] + for i, (y, typ) in enumerate(zip(ys, out_types)): + _check_output(y, order[typ], typ, i) + ids = np.concatenate([np.hstack(v[typ]) for _, v in key_typ_ids]) + input_order = order[typ][np.where(order[typ] != -1)] + reordered_ys.append(_take(y, _index(ids, input_order))) + y = reordered_ys if f_ret_is_seq else reordered_ys[0] + + return y + + +def body_tree( + m: Model, + f: Callable[..., Y], + in_types: str, + out_types: str, + *args, + reverse: bool = False, +) -> Y: + r"""Scan ``f`` across bodies in tree order, carrying results up/down the tree. + + This function groups bodies according to level and attached joints, then calls + vmap(f) on them.\ + + Args: + m: an mjx mjmodel + f: a function to be scanned with the following type signature: + def f(y, *args) -> y + where + ``y`` is the carry value and return value + ``*args`` are input arguments with types matching ``in_types`` + in_types: string specifying the type of each input arg: + 'b': split according to bodies + 'j': split according to joint types + 'q': split according to generalized coordinates (len(qpos)) + 'v': split according to degrees of freedom (len(qvel)) + out_types: string specifying the types the output dimension matches + *args: the input arguments corresponding to ``in_types`` + reverse: if True, scans up the body tree from leaves to root, otherwise + root to leaves + + Returns: + The stacked outputs of ``f`` matching the model's body order. + + Raises: + IndexError: if function output shape does not match out_types shape + """ + _check_input(m, args, in_types) + depth_fn = lambda i, p=m.body_parentid: int(i > 0) and 1 + depth_fn(p[i]) + typ_body_id = { + 'j': m.jnt_bodyid, + 'v': m.dof_bodyid, + 'q': _q_bodyid(m), + } + key_parents = {} + + # build up groupings of bodies and type ids using (level, (jnt_type,)) keys + key_typ_ids, key_body_ids = {}, {} + for body_id in np.arange(m.nbody, dtype=np.int32): + depth = depth_fn(body_id) + + # create grouping key + if any(t in 'jqv' for t in in_types + out_types): + jnts = np.nonzero(typ_body_id['j'] == body_id)[0] + jnts_p = np.nonzero(typ_body_id['j'] == m.body_parentid[body_id])[0] + key = depth, tuple(m.jnt_type[jnts]) + parent_key = depth - 1, tuple(m.jnt_type[jnts_p]) + else: + key, parent_key = (depth, ()), (depth - 1, ()) + + key_parents[key] = parent_key + body_ids = key_body_ids.get(key, np.array([], dtype=np.int32)) + key_body_ids[key] = np.append(body_ids, body_id) + + # add ids per type + for t in set(in_types + out_types): + out = key_typ_ids.setdefault(key, {}) + id_ = body_id if t == 'b' else np.nonzero(typ_body_id[t] == body_id)[0] + id_ = np.expand_dims(id_, axis=0) + out[t] = np.concatenate((out[t], id_)) if t in out else id_ + + key_typ_ids = list(sorted(key_typ_ids.items(), reverse=reverse)) + + # use this grouping to take the right data subsets and call vmap(f) + key_y = {} + for key, typ_ids in key_typ_ids: + carry = None + + if reverse: + child_keys = [k for k, v in key_parents.items() if v == key] + + for child_key in child_keys: + y = key_y[child_key] + body_ids = key_body_ids[key] + parent_ids = m.body_parentid[key_body_ids[child_key]] + id_map = _index(body_ids, parent_ids) + + def index_sum(x, i=id_map, s=body_ids.size): + return jax.ops.segment_sum(x, i, s) + + y = jax.tree_map(index_sum, y) + carry = y if carry is None else jax.tree_map(jp.add, carry, y) + else: + parent_key = key_parents[key] + y = key_y.get(parent_key) + + if y is not None: + body_ids = key_body_ids[parent_key] + parent_ids = m.body_parentid[key_body_ids[key]] + take_fn = lambda x, i=_index(body_ids, parent_ids): _take(x, i) + carry = jax.tree_map(take_fn, y) + + f_args = [_take(arg, typ_ids[typ]) for arg, typ in zip(args, in_types)] + key_y[key] = _nvmap(f, carry, *f_args) + + # slice None results from the final output + key_typ_ids = [(k, v) for k, v in key_typ_ids if key_y[k] is not None] + + # concatenate back to a single tree and drop the grouping dimension + ys = [key_y[key] for key, _ in key_typ_ids] + f_ret_is_seq = isinstance(ys[0], (list, tuple)) + ys = ys if f_ret_is_seq else [[y] for y in ys] + ys = [ + [v if typ == 'b' else jp.concatenate(v) for v, typ in zip(y, out_types)] + for y in ys + ] + ys = jax.tree_map(lambda *x: jp.concatenate(x), *ys) + + # put concatenated results back into body order + reordered_ys = [] + for i, (y, typ) in enumerate(zip(ys, out_types)): + ids = np.concatenate([np.hstack(v[typ]) for _, v in key_typ_ids]) + take_ids = _index(ids, np.sort(ids)) + _check_output(y, take_ids, typ, i) + reordered_ys.append(_take(y, take_ids)) + y = reordered_ys if f_ret_is_seq else reordered_ys[0] + + return y diff --git a/mjx/mujoco/mjx/_src/scan_test.py b/mjx/mujoco/mjx/_src/scan_test.py new file mode 100644 index 00000000..9b845de1 --- /dev/null +++ b/mjx/mujoco/mjx/_src/scan_test.py @@ -0,0 +1,227 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for scan functions.""" + +from absl.testing import absltest +from jax import numpy as jp +import mujoco +from mujoco import mjx +# pylint: disable=g-importing-member +from mujoco.mjx._src import scan +from mujoco.mjx._src.types import JointType +# pylint: enable=g-importing-member +import numpy as np + + +class ScanTest(absltest.TestCase): + + _MULTI_DOF_XML = """ + + + + + + + + + + + + + + + + + + + """ + + def test_flat_empty(self): + """Test scanning over just world body.""" + m = mujoco.MjModel.from_xml_string(""" + + + + """) + m = mjx.device_put(m) + + def fn(body_id): + return body_id + 1 + + b_in = jp.array([1]) + b_expect = jp.array([2]) + b_out = scan.flat(m, fn, 'b', 'b', b_in) + + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + def test_flat_joints(self): + """Tests scanning over bodies with joints of different types.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) + m = mjx.device_put(m) + + # we will test two functions: + # 1) j_fn receives jnt_types as a jp array + # 2) s_fn receives jnt_types as a static np array and can switch on it + j_fn = lambda jnt_pos, val: val + jp.sum(jnt_pos) + s_fn = lambda jnt_types, val: val + sum(jnt_types) + + b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]]) + b_expect = jp.array([[0, 0], [1, 1], [3, 3], [8, 8]]) + b_out = scan.flat(m, j_fn, 'jb', 'b', m.jnt_pos, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.flat(m, s_fn, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # None should be omitted from the results + def no_free(jnt_types, val): + if tuple(jnt_types) == (JointType.FREE,): + return None + return val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) + b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # we should not call functions for which we know we will discard the results + def no_world(jnt_types, val): + if jnt_types.size == 0: + self.fail('world has no dofs, should not be called') + return val + sum(jnt_types) + v_in = jp.ones((m.nv, 1)) + scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in) + + def test_body_tree(self): + """Tests tree scanning over bodies with different joint counts.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML) + m = mjx.device_put(m) + + # we will test two functions: + # 1) j_fn receives jnt_pos which is a jp array + # 2) s_fn receives jnt_types which is a static np array + def j_fn(carry, jnt_pos, val): + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + jp.sum(jnt_pos) + + def s_fn(carry, jnt_types, val): + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + sum(jnt_types) + + b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]]) + b_expect = jp.array([[0, 0], [1, 1], [4, 4], [9, 9]]) + + b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # and reverse too: + b_expect = jp.array([[12, 12], [12, 12], [3, 3], [8, 8]]) + b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in, reverse=True) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in, reverse=True) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + # None should be omitted from the results + def no_free(carry, jnt_types, val): + if tuple(jnt_types) == (JointType.FREE,): + return None + carry = jp.zeros_like(val) if carry is None else carry + return carry + val + sum(jnt_types) + b_expect = jp.array([[0, 0], [3, 3], [8, 8]]) + b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in) + np.testing.assert_equal(np.array(b_out), np.array(b_expect)) + + _MULTI_ACT_XML = """ + + + """ + + def testscan_actuators(self): + """Tests scanning over actuators.""" + m = mujoco.MjModel.from_xml_string(self._MULTI_ACT_XML) + m = mjx.device_put(m) + + fn = lambda *args: args + args = ( + m.actuator_gear, + m.jnt_type, + jp.arange(m.nq), + jp.arange(m.nv), + jp.array([1.4, 1.1]), + ) + gear, jnt_typ, qadr, vadr, act = scan.flat( + m, fn, 'ujqva', 'ujqva', *args, group_by='u' + ) + + np.testing.assert_array_equal(gear, m.actuator_gear) + np.testing.assert_array_equal(jnt_typ, m.jnt_type[m.actuator_trnid]) + np.testing.assert_array_equal(act, jp.array([1.4, 1.1])) + expected_vadr = np.concatenate( + [np.nonzero(m.dof_jntid == trnid)[0] for trnid in m.actuator_trnid] + ) + np.testing.assert_array_equal(vadr, expected_vadr) + expected_qadr = np.concatenate( + [np.nonzero(scan._q_jointid(m) == i)[0] for i in m.actuator_trnid] + ) + np.testing.assert_array_equal(qadr, expected_qadr) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py new file mode 100644 index 00000000..2b421f2f --- /dev/null +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -0,0 +1,467 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Core smooth dynamics functions.""" + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math +from mujoco.mjx._src import scan +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import JointType +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def kinematics(m: Model, d: Data) -> Data: + """Converts position/velocity from generalized coordinates to maximal.""" + + def fn(carry, jnt_typs, jnt_pos, jnt_axis, qpos, qpos0, pos, quat): + # calculate joint anchors, axes, body pos and quat in global frame + # also normalize qpos while we're at it + + if carry is not None: + _, _, _, parent_pos, parent_quat, _ = carry + pos = parent_pos + math.rotate(pos, parent_quat) + quat = math.quat_mul(parent_quat, quat) + + anchors, axes = [], [] + + qpos_i = 0 + for i, jnt_typ in enumerate(jnt_typs): + if jnt_typ == JointType.FREE: + anchor, axis = qpos[qpos_i : qpos_i + 3], jp.array([0.0, 0.0, 1.0]) + else: + anchor = math.rotate(jnt_pos[i], quat) + pos + axis = math.rotate(jnt_axis[i], quat) + anchors, axes = anchors + [anchor], axes + [axis] + + if jnt_typ == JointType.FREE: + pos = qpos[qpos_i : qpos_i + 3] + quat = math.normalize(qpos[qpos_i + 3 : qpos_i + 7]) + qpos = qpos.at[qpos_i + 3 : qpos_i + 7].set(quat) + qpos_i += 7 + elif jnt_typ == JointType.BALL: + qloc = math.normalize(qpos[qpos_i : qpos_i + 4]) + qpos = qpos.at[qpos_i : qpos_i + 4].set(qloc) + quat = math.quat_mul(quat, qloc) + pos = anchor - math.rotate(jnt_pos[i], quat) # off-center rotation + qpos_i += 4 + elif jnt_typ == JointType.HINGE: + angle = qpos[qpos_i] - qpos0[qpos_i] + qloc = math.axis_angle_to_quat(jnt_axis[i], angle) + quat = math.quat_mul(quat, qloc) + pos = anchor - math.rotate(jnt_pos[i], quat) # off-center rotation + qpos_i += 1 + elif jnt_typ == JointType.SLIDE: + pos += axis * (qpos[qpos_i] - qpos0[qpos_i]) + qpos_i += 1 + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + anchor = jp.stack(anchors) if anchors else jp.empty((0, 3)) + axis = jp.stack(axes) if axes else jp.empty((0, 3)) + mat = math.quat_to_mat(quat) + + return qpos, anchor, axis, pos, quat, mat + + qpos, xanchor, xaxis, xpos, xquat, xmat = scan.body_tree( + m, + fn, + 'jjjqqbb', + 'qjjbbb', + m.jnt_type, + m.jnt_pos, + m.jnt_axis, + d.qpos, + m.qpos0, + m.body_pos, + m.body_quat, + ) + + @jax.vmap + def local_to_global(pos1, quat1, pos2, quat2): + pos = pos1 + math.rotate(pos2, quat1) + mat = math.quat_to_mat(math.quat_mul(quat1, quat2)) + return pos, mat + + # TODO(erikfrey): confirm that quats are more performant for mjx than mats + xipos, ximat = local_to_global(xpos, xquat, m.body_ipos, m.body_iquat) + geom_xpos, geom_xmat = local_to_global( + xpos[m.geom_bodyid], xquat[m.geom_bodyid], m.geom_pos, m.geom_quat + ) + + d = d.replace(qpos=qpos, xanchor=xanchor, xaxis=xaxis, xpos=xpos) + d = d.replace(xquat=xquat, xmat=xmat, xipos=xipos, ximat=ximat) + d = d.replace(geom_xpos=geom_xpos, geom_xmat=geom_xmat) + + return d + + +def com_pos(m: Model, d: Data) -> Data: + """Maps inertias and motion dofs to global frame centered at subtree-CoM.""" + + # calculate center of mass of each subtree + def subtree_sum(carry, xipos, body_mass): + pos, mass = xipos * body_mass, body_mass + if carry is not None: + subtree_pos, subtree_mass = carry + pos, mass = pos + subtree_pos, mass + subtree_mass + return pos, mass + + pos, mass = scan.body_tree( + m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True + ) + cond = jp.tile(mass < jp.array(mujoco.mjMINVAL), (3, 1)).T + subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass)) + d = d.replace(subtree_com=subtree_com) + + # map inertias to frame centered at subtree_com + @jax.vmap + def inert_com(inert, ximat, off, mass): + h = jp.cross(off, -jp.eye(3)) + inert = ximat @ jp.diag(inert) @ ximat.T + h @ h.T * mass + # cinert is triu(inert), mass * off, mass + inert = inert[(jp.array([0, 1, 2, 0, 0, 1]), jp.array([0, 1, 2, 1, 2, 2]))] + return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)]) + + root_com = subtree_com[jp.array(m.body_rootid)] + offset = d.xipos - root_com + cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass) + d = d.replace(cinert=cinert) + + # map motion dofs to global frame centered at subtree_com + def cdof_fn(jnt_typs, root_com, xmat, xanchor, xaxis): + cdofs = [] + + dof_com_fn = lambda a, o: jp.concatenate([a, jp.cross(a, o)]) + + for i, jnt_typ in enumerate(jnt_typs): + offset = root_com - xanchor[i] + if jnt_typ == JointType.FREE: + cdofs.append(jp.eye(3, 6, 3)) # free translation + cdofs.append(jax.vmap(dof_com_fn, in_axes=(0, None))(xmat.T, offset)) + elif jnt_typ == JointType.BALL: + cdofs.append(jax.vmap(dof_com_fn, in_axes=(0, None))(xmat.T, offset)) + elif jnt_typ == JointType.HINGE: + cdof = dof_com_fn(xaxis[i], offset) + cdofs.append(jp.expand_dims(cdof, 0)) + elif jnt_typ == JointType.SLIDE: + cdof = jp.concatenate((jp.zeros((3,)), xaxis[i])) + cdofs.append(jp.expand_dims(cdof, 0)) + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + + cdof = jp.concatenate(cdofs) if cdofs else jp.empty((0, 6)) + + return cdof + + cdof = scan.flat( + m, + cdof_fn, + 'jbbjj', + 'v', + m.jnt_type, + root_com, + d.xmat, + d.xanchor, + d.xaxis, + ) + d = d.replace(cdof=cdof) + + return d + + +def crb(m: Model, d: Data) -> Data: + """Runs composite rigid body inertia algorithm.""" + + def crb_fn(crb_child, crb_body): + if crb_child is not None: + crb_body += crb_child + return crb_body + + crb_body = scan.body_tree(m, crb_fn, 'b', 'b', d.cinert, reverse=True) + crb_body = crb_body.at[0].set(0.0) + d = d.replace(crb=crb_body) + + # TODO(erikfrey): do centralized take fn? + crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0) + crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof) + + dof_i, dof_j, diag = [], [], [] + for i in range(m.nv): + diag.append(len(dof_i)) + j = i + while j > -1: + dof_i, dof_j = dof_i + [i], dof_j + [j] + j = m.dof_parentid[j] + + crb_codf_i = jp.take(crb_cdof, jp.array(dof_i), axis=0) + cdof_j = jp.take(d.cdof, jp.array(dof_j), axis=0) + qm = jax.vmap(jp.dot)(crb_codf_i, cdof_j) + + # add armature to diagonal + qm = qm.at[jp.array(diag)].add(m.dof_armature) + + d = d.replace(qM=qm) + + return d + + +def factor_m( + m: Model, + d: Data, + qM: jax.Array, # pylint:disable=invalid-name +) -> Data: + """Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd.""" + + # build up indices for where we will do backwards updates over qLD + # TODO(erikfrey): do fewer updates by combining non-overlapping ranges + dof_madr = jp.array(m.dof_Madr) + updates = {} + madr_ds = [] + for i in range(m.nv): + madr_d = madr_ij = m.dof_Madr[i] + j = i + while True: + madr_ds.append(madr_d) + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + madr_j_range = tuple(m.dof_Madr[j : j + 2]) + updates.setdefault(madr_j_range, []).append((madr_d, madr_ij)) + + qld = qM + + for (out_beg, out_end), vals in sorted(updates.items(), reverse=True): + madr_d, madr_ij = jp.array(vals).T + + @jax.vmap + def off_diag_fn(madr_d, madr_ij, qld=qld, width=out_end - out_beg): + qld_row = jax.lax.dynamic_slice(qld, (madr_ij,), (width,)) + return -(qld_row[0] / qld[madr_d]) * qld_row + + qld_update = jp.sum(off_diag_fn(madr_d, madr_ij), axis=0) + qld = qld.at[out_beg:out_end].add(qld_update) + # TODO(erikfrey): determine if this minimum value guarding is necessary: + # qld = qld.at[dof_madr].set(jp.maximum(qld[dof_madr], _MJ_MINVAL)) + + qld_diag = qld[dof_madr] + qld = (qld / qld[jp.array(madr_ds)]).at[dof_madr].set(qld_diag) + + d = d.replace(qLD=qld, qLDiagInv=1 / qld_diag) + + return d + + +def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array: + """Computes sparse backsubstitution: x = inv(L'*D*L)*y .""" + + updates_i, updates_j = {}, {} + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + updates_i.setdefault(i, []).append((madr_ij, j)) + updates_j.setdefault(j, []).append((madr_ij, i)) + + # x <- inv(L') * x + for j, vals in sorted(updates_j.items(), reverse=True): + madr_ij, i = jp.array(vals).T + x = x.at[j].add(-jp.sum(d.qLD[madr_ij] * x[i])) + + # x <- inv(D) * x + x = x * d.qLDiagInv + + # x <- inv(L) * x + for i, vals in sorted(updates_i.items()): + madr_ij, j = jp.array(vals).T + x = x.at[i].add(-jp.sum(d.qLD[madr_ij] * x[j])) + + return x + + +def dense_m(m: Model, d: Data) -> jax.Array: + """Reconstitute dense mass matrix from qM.""" + + is_, js, madr_ijs = [], [], [] + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij] + + i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs)) + + mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM[madr_ij]) + + # diagonal, upper triangular, lower triangular + mat = jp.diag(d.qM[jp.array(m.dof_Madr)]) + mat + mat.T + + return mat + + +def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array: + """Multiply vector by inertia matrix.""" + + diag_mul = d.qM[jp.array(m.dof_Madr)] * vec + + is_, js, madr_ijs = [], [], [] + for i in range(m.nv): + madr_ij, j = m.dof_Madr[i], i + + while True: + madr_ij, j = madr_ij + 1, m.dof_parentid[j] + if j == -1: + break + is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij] + + i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs)) + + out = diag_mul.at[i].add(d.qM[madr_ij] * vec[j]) + out = out.at[j].add(d.qM[madr_ij] * vec[i]) + + return out + + +def com_vel(m: Model, d: Data) -> Data: + """Computes cvel, cdof_dot.""" + + # forward scan down tree: accumulate link center of mass velocity + def fn(parent, jnt_typs, cdof, qvel): + cvel = jp.zeros((6,)) if parent is None else parent[0] + + cross_fn = jax.vmap(math.motion_cross, in_axes=(None, 0)) + cdof_x_qvel = jax.vmap(jp.multiply)(cdof, qvel) + + dof_beg = 0 + cdof_dots = [] + for jnt_typ in jnt_typs: + dof_end = dof_beg + JointType(jnt_typ).dof_width() + if jnt_typ == JointType.FREE: + cvel += jp.sum(cdof_x_qvel[:3], axis=0) + cdof_ang_dot = cross_fn(cvel, cdof[3:]) + cvel += jp.sum(cdof_x_qvel[3:], axis=0) + cdof_dots.append(jp.concatenate((jp.zeros((3, 6)), cdof_ang_dot))) + else: + cdof_dots.append(cross_fn(cvel, cdof[dof_beg:dof_end])) + cvel += jp.sum(cdof_x_qvel[dof_beg:dof_end], axis=0) + dof_beg = dof_end + + cdof_dot = jp.concatenate(cdof_dots) if cdof_dots else jp.empty((0, 6)) + return cvel, cdof_dot + + cvel, cdof_dot = scan.body_tree( + m, + fn, + 'jvv', + 'bv', + m.jnt_type, + d.cdof, + d.qvel, + ) + + d = d.replace(cvel=cvel, cdof_dot=cdof_dot) + + return d + + +def rne(m: Model, d: Data) -> Data: + """Computes inverse dynamics using the recursive Newton-Euler algorithm.""" + # forward scan over tree: accumulate link center of mass acceleration + def cacc_fn(cacc, cdof_dot, qvel): + if cacc is None: + if m.opt.disableflags & DisableBit.GRAVITY: + cacc = jp.zeros((6,)) + else: + cacc = jp.concatenate((jp.zeros((3,)), -m.opt.gravity)) + + cacc += jp.sum(jax.vmap(jp.multiply)(cdof_dot, qvel), axis=0) + + return cacc + + cacc = scan.body_tree(m, cacc_fn, 'vv', 'b', d.cdof_dot, d.qvel) + + def frc(cinert, cacc, cvel): + frc = math.inert_mul(cinert, cacc) + frc += math.motion_cross_force(cvel, math.inert_mul(cinert, cvel)) + + return frc + + loc_cfrc = jax.vmap(frc)(d.cinert, cacc, d.cvel) + + # backward scan up tree: accumulate body forces + def cfrc_fn(cfrc_child, cfrc): + if cfrc_child is not None: + cfrc += cfrc_child + return cfrc + + cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True) + qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)]) + + d = d.replace(qfrc_bias=qfrc_bias) + + return d + + +def transmission(m: Model, d: Data) -> Data: + """Computes actuator/transmission lengths and moments.""" + if not m.nu: + return d + + def fn(gear, jnt_typ, m_i, m_j, qpos): + # handles joint transmissions only + if jnt_typ == JointType.FREE: + length = jp.zeros(1) + moment = gear + m_i = jp.repeat(m_i, 6) + m_j = m_j + jp.arange(6) + elif jnt_typ == JointType.BALL: + axis, _ = math.quat_to_axis_angle(qpos) + length = jp.dot(axis, gear[:3])[None] + moment = gear[:3] + m_i = jp.repeat(m_i, 3) + m_j = m_j + jp.arange(3) + elif jnt_typ in (JointType.SLIDE, JointType.HINGE): + length = qpos * gear[0] + moment = gear[:1] + m_i, m_j = m_i[None], m_j[None] + else: + raise RuntimeError(f'unrecognized joint type: {jnt_typ}') + return length, moment, m_i, m_j + + length, m_val, m_i, m_j = scan.flat( + m, + fn, + 'ujujq', + 'uvvv', + m.actuator_gear, + m.jnt_type, + jp.arange(m.nu), + jp.array(m.jnt_dofadr), + d.qpos, + group_by='u', + ) + moment = jp.zeros((m.nu, m.nv)).at[m_i, m_j].set(m_val) + length = length.reshape((m.nu,)) + d = d.replace(actuator_length=length, actuator_moment=moment) + return d diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py new file mode 100644 index 00000000..4cdfc6ce --- /dev/null +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -0,0 +1,166 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for smooth dynamics functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import test_util +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import DisableBit +# pylint: enable=g-importing-member +import numpy as np + + +def _assert_eq(a, b, name, step, fname, atol=1e-5, rtol=1e-5): + err_msg = f'mismatch: {name} at step {step} in {fname}' + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-5, rtol=1e-5): + err_msg = f'mismatch: {attr} at step {step} in {fname}' + a, b = getattr(a, attr), getattr(b, attr) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +class SmoothTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(test_util.TEST_FILES)) + def test_smooth(self, seed, fname): + """Tests mujoco mj smooth functions match mujoco_mjx smooth functions.""" + if fname in ('convex.xml', 'weld.xml'): + return + + np.random.seed(seed) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + + kinematics_jit_fn = jax.jit(mjx.kinematics) + com_pos_jit_fn = jax.jit(mjx.com_pos) + crb_jit_fn = jax.jit(mjx.crb) + factor_m_fn = jax.jit(mjx.factor_m) + com_vel_jit_fn = jax.jit(mjx.com_vel) + rne_jit_fn = jax.jit(mjx.rne) + mul_m_jit_fn = jax.jit(mjx.mul_m) + transmission_jit_fn = jax.jit(mjx.transmission) + + mx = mjx.device_put(m) + dx = mjx.make_data(mx) + + # give the system a little kick to ensure we have non-identity rotations + d.qvel = np.random.random(m.nv) + for i in range(100): + qpos, qvel = d.qpos.copy(), d.qvel.copy() + mujoco.mj_step(m, d) + + # kinematics + dx = kinematics_jit_fn(mx, dx.replace(qpos=qpos, qvel=qvel)) + _assert_attr_eq(d, dx, 'xanchor', i, fname) + _assert_attr_eq(d, dx, 'xaxis', i, fname) + _assert_attr_eq(d, dx, 'xpos', i, fname) + _assert_attr_eq(d, dx, 'xquat', i, fname) + _assert_eq(d.xmat.reshape((-1, 3, 3)), dx.xmat, 'xmat', i, fname) + _assert_attr_eq(d, dx, 'xipos', i, fname) + _assert_eq(d.ximat.reshape((-1, 3, 3)), dx.ximat, 'ximat', i, fname) + _assert_attr_eq(d, dx, 'geom_xpos', i, fname) + _assert_eq( + d.geom_xmat.reshape((-1, 3, 3)), + dx.geom_xmat, + 'geom_xmat', + i, + fname, + ) + + # com_pos + dx = com_pos_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'subtree_com', i, fname) + _assert_attr_eq(d, dx, 'cinert', i, fname) + _assert_attr_eq(d, dx, 'cdof', i, fname) + + # crb + dx = crb_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'crb', i, fname) + _assert_attr_eq(d, dx, 'qM', i, fname) + + # factor_m + dx = factor_m_fn(mx, dx, dx.qM) + _assert_attr_eq(d, dx, 'qLD', i, fname, atol=1e-3) + _assert_attr_eq(d, dx, 'qLDiagInv', i, fname, atol=1e-3, rtol=1e-4) + + # com_vel + dx = com_vel_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'cvel', i, fname) + _assert_attr_eq(d, dx, 'cdof_dot', i, fname) + + # rne + dx = rne_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'qfrc_bias', i, fname, atol=1e-4) + + # mul_m (auxilliary function, not part of smooth step) + vec = np.random.random(m.nv) + mjx_vec = mul_m_jit_fn(mx, dx, jp.array(vec)) + mj_vec = np.zeros(m.nv) + mujoco.mj_mulM(m, d, mj_vec, vec) + _assert_eq(mj_vec, mjx_vec, 'mul_m', i, fname, atol=1e-4) + + # transmission + dx = transmission_jit_fn(mx, dx) + _assert_attr_eq(d, dx, 'actuator_length', i, fname) + _assert_attr_eq(d, dx, 'actuator_moment', i, fname) + + +class DisableGravityTest(absltest.TestCase): + + def test_disabled(self): + m = mujoco.MjModel.from_xml_string(""" + + + """) + mx = mjx.device_put(m) + d = mujoco.MjData(m) + dx = mjx.device_put(d) + + # test with gravity + step_jit_fn = jax.jit(mjx.step) + dx = step_jit_fn(mx, dx) + np.testing.assert_array_almost_equal( + dx.qpos, np.array([0.0, 0.0, -9.81e-4, 1.0, 0.0, 0.0, 0.0]), decimal=7 + ) + + # test with gravity disabled + mx = mx.tree_replace( + {'opt.disableflags': mx.opt.disableflags | DisableBit.GRAVITY} + ) + dx = mjx.device_put(d) + step_jit_fn = jax.jit(mjx.step) + dx = step_jit_fn(mx, dx) + np.testing.assert_equal( + dx.qpos, np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py new file mode 100644 index 00000000..0aefd027 --- /dev/null +++ b/mjx/mujoco/mjx/_src/solver.py @@ -0,0 +1,356 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""CG and Newton solvers.""" + +from typing import Optional + +import jax +from jax import numpy as jp +import mujoco +from mujoco.mjx._src import math +from mujoco.mjx._src import smooth +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import DisableBit +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +class _CGContext(PyTreeNode): + """Data updated during each cg solver iteration. + + Attributes: + qacc: acceleration (from Data) (nv,) + qfrc_constraint: constraint force (from Data) (nv,) + Jaref: Jac*qacc - aref (nefc,) + efc_force: constraint force in constraint space (nefc,) + M: dense mass matrix, populated for nv < 100 (nv, nv) + Ma: M*qacc (nv,) + grad: gradient of master cost (nv,) + Mgrad: M / grad (nv,) + search: linesearch vector (nv,) + gauss: gauss Cost + cost: constraint + Gauss cost + prev_cost: cost from previous cg iter + solver_niter: number of solver iterations + """ + + qacc: jax.Array + qfrc_constraint: jax.Array + Jaref: jax.Array # pylint: disable=invalid-name + efc_force: jax.Array + M: Optional[jax.Array] + Ma: jax.Array # pylint: disable=invalid-name + grad: jax.Array + Mgrad: jax.Array # pylint: disable=invalid-name + search: jax.Array + gauss: jax.Array + cost: jax.Array + prev_cost: jax.Array + solver_niter: jax.Array + + @classmethod + def create(cls, m: Model, d: Data, grad: bool = True) -> '_CGContext': + jaref = d.efc_J @ d.qacc - d.efc_aref + # TODO(robotics-team): determine nv at which sparse mul is faster + M = smooth.dense_m(m, d) if m.nv < 100 else None # pylint: disable=invalid-name + ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc + nv_0 = jp.zeros((m.nv,)) + ctx = _CGContext( + qacc=d.qacc, + qfrc_constraint=d.qfrc_constraint, + Jaref=jaref, + efc_force=-jaref * d.efc_D, + M=M, + Ma=ma, + grad=nv_0, + Mgrad=nv_0, + search=nv_0, + gauss=0.0, + cost=jp.inf, + prev_cost=0.0, + solver_niter=0, + ) + ctx = _cg_update_constraint(m, d, ctx) + if grad: + ctx = _cg_update_gradient(m, d, ctx) + ctx = ctx.replace(search=-ctx.Mgrad) # start with preconditioned gradient + + return ctx + + +class _LSPoint(PyTreeNode): + """Line search evaluation point. + + Attributes: + alpha: step size that reduces f(x + alpha * p) given search direction p + cost: line search cost + deriv_0: first derivative of quadratic + deriv_1: second derivative of quadratic + """ + + alpha: jax.Array + cost: jax.Array + deriv_0: jax.Array + deriv_1: jax.Array + + @classmethod + def create( + cls, + ctx: _CGContext, + alpha: jax.Array, + jv: jax.Array, + quad: jax.Array, + quad_gauss: jax.Array, + ) -> '_LSPoint': + """Creates a linesearch point with first and second derivatives.""" + # roughly corresponds to CGEval in mujoco/src/engine/engine_solver.c + + # TODO(robotics-team): change this to support equality, friction constraints + active = (ctx.Jaref + alpha * jv) < 0 + quad = jax.vmap(jp.multiply)(quad, active) # only active + quad_total = quad_gauss + jp.sum(quad, axis=0) + + cost = alpha * alpha * quad_total[2] + alpha * quad_total[1] + quad_total[0] + deriv_0 = 2 * alpha * quad_total[2] + quad_total[1] + deriv_1 = 2 * quad_total[2] + return _LSPoint(alpha=alpha, cost=cost, deriv_0=deriv_0, deriv_1=deriv_1) + + +class _LSContext(PyTreeNode): + """Data updated during each cg line search iteration. + + Attributes: + lo: low point bounding the line search interval + hi: high point bounding the line search interval + swap: True if low or hi was swapped in the line search iteration + ls_iter: number of linesearch iterations + """ + + lo: _LSPoint + hi: _LSPoint + swap: jax.Array + ls_iter: jax.Array + + +def _while_loop_scan(cond_fun, body_fun, init_val, max_iter): + """Scan-based implementation (jit ok, reverse-mode autodiff ok).""" + def _iter(val): + next_val = body_fun(val) + next_cond = cond_fun(next_val) + return next_val, next_cond + + def _fun(tup, it): + val, cond = tup + # When cond is met, we start doing no-ops. + return jax.lax.cond(cond, _iter, lambda x: (x, False), val), it + + init = (init_val, cond_fun(init_val)) + return jax.lax.scan(_fun, init, None, length=max_iter)[0][0] + + +def _cg_update_constraint(m: Model, d: Data, ctx: _CGContext) -> _CGContext: + """Updates constraint force and resulting cost given latst CG iteration. + + Corresponds to CGupdateConstraint in mujoco/src/engine/engine_solver.c + + Args: + m: model defining constraints + d: data which contains latest qacc and smooth terms + ctx: current CG context + + Returns: + context with new constraint force and costs + """ + del m + + # TODO(robotics-team): add equality, friction constraints + # also consider moving to _constraint.py to match mujoco layout + + jaref = ctx.Jaref * (ctx.Jaref < 0) # non-negative constraints + efc_force = -jaref * d.efc_D + qfrc_constraint = d.efc_J.T @ efc_force + gauss = 0.5 * jp.dot(ctx.Ma - d.qfrc_smooth, ctx.qacc - d.qacc_smooth) + cost = 0.5 * jp.sum(jaref * jaref * d.efc_D) + gauss + + ctx = ctx.replace( + qfrc_constraint=qfrc_constraint, + gauss=gauss, + cost=cost, + prev_cost=ctx.cost, + efc_force=efc_force, + ) + + return ctx + + +def _cg_update_gradient(m: Model, d: Data, ctx: _CGContext) -> _CGContext: + """Updates grad and M / grad given latest CG iteration. + + Corresponds to CGupdateGradient in mujoco/src/engine/engine_solver.c + + Args: + m: model defining constraints + d: data which contains latest smooth terms + ctx: current CG contet + + Returns: + context with new grad and M / grad + """ + + grad = ctx.Ma - d.qfrc_smooth - ctx.qfrc_constraint + mgrad = smooth.solve_m(m, d, grad) + + ctx = ctx.replace(grad=grad, Mgrad=mgrad) + + return ctx + + +def _rescale(m: Model, value: jax.Array) -> jax.Array: + return value / (m.stat.meaninertia * max(1, m.nv)) + + +def _cg_search(m: Model, d: Data, ctx: _CGContext) -> _CGContext: + """Performs a zoom linesearch to find optimal search step size. + + Args: + m: model defining search options and other needed terms + d: data with inertia matrix and other needed terms + ctx: current CG context + + Returns: + updated context with new qacc, Ma, Jaref + """ + smag = math.norm(ctx.search) * m.stat.meaninertia * max(1, m.nv) + gtol = m.opt.tolerance * m.opt.ls_tolerance * smag + + # compute Mv, Jv + mv = smooth.mul_m(m, d, ctx.search) if ctx.M is None else ctx.M @ ctx.search + jv = d.efc_J @ ctx.search + + # prepare quadratics + quad_gauss = jp.stack(( + ctx.gauss, + jp.dot(ctx.search, ctx.Ma) - jp.dot(ctx.search, d.qfrc_smooth), + 0.5 * jp.dot(ctx.search, mv), + )) + quad = jp.stack((0.5 * ctx.Jaref * ctx.Jaref, jv * ctx.Jaref, 0.5 * jv * jv)) + quad = (quad * d.efc_D).T + + point_fn = lambda alpha: _LSPoint.create(ctx, alpha, jv, quad, quad_gauss) + + def cond(ctx: _LSContext) -> jax.Array: + done = ctx.ls_iter >= m.opt.ls_iterations + done |= ~ctx.swap # if we did not adjust the interval + done |= (ctx.lo.deriv_0 < 0) & (ctx.lo.deriv_0 > -gtol) + done |= (ctx.hi.deriv_0 > 0) & (ctx.hi.deriv_0 < gtol) + + return ~done + + def body(ctx: _LSContext) -> _LSContext: + # always compute new bracket boundaries and a midpoint + lo, hi = ctx.lo, ctx.hi + lo_next = point_fn(lo.alpha - lo.deriv_0 / lo.deriv_1) + hi_next = point_fn(hi.alpha - hi.deriv_0 / hi.deriv_1) + mid = point_fn(0.5 * (lo.alpha + hi.alpha)) + + # we swap lo/hi if: + # 1) they are not correctly at a bracket boundary (e.g. lo.deriv_0 > 0), OR + # 2) if moving to next or mid narrows the bracket + swap_lo_next = (lo.deriv_0 > 0) | (lo.deriv_0 < lo_next.deriv_0) + lo = jax.tree_map(lambda x, y: jp.where(swap_lo_next, y, x), lo, lo_next) + swap_lo_mid = (mid.deriv_0 < 0) & (lo.deriv_0 < mid.deriv_0) + lo = jax.tree_map(lambda x, y: jp.where(swap_lo_mid, y, x), lo, mid) + + swap_hi_next = (hi.deriv_0 < 0) | (hi.deriv_0 > hi_next.deriv_0) + hi = jax.tree_map(lambda x, y: jp.where(swap_hi_next, y, x), hi, hi_next) + swap_hi_mid = (mid.deriv_0 > 0) & (hi.deriv_0 > mid.deriv_0) + hi = jax.tree_map(lambda x, y: jp.where(swap_hi_mid, y, x), hi, mid) + + swap = swap_lo_next | swap_lo_mid | swap_hi_next | swap_hi_mid + + ctx = ctx.replace(lo=lo, hi=hi, swap=swap, ls_iter=ctx.ls_iter + 1) + + return ctx + + # initialize interval + p0 = point_fn(jp.array(0.0)) + lo = point_fn(p0.alpha - p0.deriv_0 / p0.deriv_1) + lesser_fn = lambda x, y: jp.where(lo.deriv_0 < p0.deriv_0, x, y) + hi = jax.tree_map(lesser_fn, p0, lo) + lo = jax.tree_map(lesser_fn, lo, p0) + ls_ctx = _LSContext(lo=lo, hi=hi, swap=jp.array(True), ls_iter=0) + ls_ctx = _while_loop_scan(cond, body, ls_ctx, m.opt.ls_iterations) + + # move to new solution if improved + lo, hi = ls_ctx.lo, ls_ctx.hi + improved = (lo.cost < p0.cost) | (hi.cost < p0.cost) + alpha = jp.where(lo.cost < hi.cost, lo.alpha, hi.alpha) + qacc = ctx.qacc + improved * ctx.search * alpha + ma = ctx.Ma + improved * mv * alpha + jaref = ctx.Jaref + improved * jv * alpha + + ctx = ctx.replace(qacc=qacc, Ma=ma, Jaref=jaref) + + return ctx + + +def cg_solve(m: Model, d: Data) -> Data: + """Finds forces that satisfy constraints using conjugate gradient descent.""" + + def cond(ctx: _CGContext) -> jax.Array: + improvement = _rescale(m, ctx.prev_cost - ctx.cost) + gradient = _rescale(m, math.norm(ctx.grad)) + + done = ctx.solver_niter >= m.opt.iterations + done |= improvement < m.opt.tolerance + done |= gradient < m.opt.tolerance + + return ~done + + def body(ctx: _CGContext) -> _CGContext: + ctx = _cg_search(m, d, ctx) + prev_grad, prev_Mgrad = ctx.grad, ctx.Mgrad # pylint: disable=invalid-name + ctx = _cg_update_constraint(m, d, ctx) + ctx = _cg_update_gradient(m, d, ctx) + + # polak-ribiere: + beta = jp.dot(ctx.grad, ctx.Mgrad - prev_Mgrad) + beta = beta / jp.maximum(mujoco.mjMINVAL, jp.dot(prev_grad, prev_Mgrad)) + beta = jp.maximum(0, beta) + search = -ctx.Mgrad + beta * ctx.search + ctx = ctx.replace(search=search, solver_niter=ctx.solver_niter + 1) + + return ctx + + # warmstart: + qacc = d.qacc_smooth + if not m.opt.disableflags & DisableBit.WARMSTART: + warm = _CGContext.create(m, d.replace(qacc=d.qacc_warmstart), grad=False) + smth = _CGContext.create(m, d.replace(qacc=d.qacc_smooth), grad=False) + qacc = jp.where(warm.cost < smth.cost, d.qacc_warmstart, d.qacc_smooth) + d = d.replace(qacc=qacc) + + ctx = jax.lax.while_loop(cond, body, _CGContext.create(m, d)) + + d = d.replace( + qacc_warmstart=ctx.qacc, + qacc=ctx.qacc, + qfrc_constraint=ctx.qfrc_constraint, + efc_force=ctx.efc_force, + ) + + return d diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py new file mode 100644 index 00000000..6a4f7792 --- /dev/null +++ b/mjx/mujoco/mjx/_src/solver_test.py @@ -0,0 +1,130 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for forward functions.""" + +from absl.testing import absltest +from absl.testing import parameterized +from etils import epath +import jax +import mujoco +from mujoco import mjx +import numpy as np + + +def _assert_attr_eq(a, b, attr, step, fname, atol=1e-2, rtol=1e-2): + err_msg = f'mismatch: {attr} at step {step} in {fname}' + a, b = getattr(a, attr), getattr(b, attr) + np.testing.assert_allclose(a, b, err_msg=err_msg, atol=atol, rtol=rtol) + + +class Solver64Test(parameterized.TestCase): + """Tests solvers at 64 bit precision.""" + + def setUp(self): + super().setUp() + jax.config.update('jax_enable_x64', True) + + def tearDown(self): + super().tearDown() + jax.config.update('jax_enable_x64', False) + + @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) + def test_cg(self, seed, fname): + """Test mjx cg solver matches mujoco cg solver at 64 bit precision.""" + f = epath.resource_path('mujoco.mjx') / 'test_data' / fname + m = mujoco.MjModel.from_xml_string(f.read_text()) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + + jax.config.update('jax_enable_x64', True) + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + np.random.seed(seed) + d.qvel = 0.01 * np.random.random(m.nv) + + for i in range(100): + # in order to avoid re-jitting, reuse the same mj_data shape + save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth + d = mujoco.MjData(m) + d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + # at 64 bits the solutions returned by the two solvers are quite close + self.assertLessEqual(dx.solver_niter[0], d.solver_niter[0]) + _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname) + _assert_attr_eq(d, dx, 'qacc', i, fname) + + +class SolverTest(parameterized.TestCase): + + @parameterized.parameters(enumerate(('ant.xml', 'humanoid.xml'))) + def test_cg(self, seed, fname): + """Test mjx cg solver is close to mj at 32 bit precision. + + Args: + seed: int + fname: file to test + + At lower float resolution there's wiggle room in valid forces that satisfy + constraints. So instead let's mainly validate that mjx is finding solutions + with as good cost as mujoco, even if the resulting forces/accelerations + are not quite the same. + """ + f = epath.resource_path('mujoco.mjx') / 'test_data' / fname + m = mujoco.MjModel.from_xml_string(f.read_text()) + d = mujoco.MjData(m) + mx = mjx.device_put(m) + + forward_jit_fn = jax.jit(mjx.forward) + + # give the system a little kick to ensure we have non-identity rotations + np.random.seed(seed) + d.qvel = 0.01 * np.random.random(m.nv) + + for i in range(100): + # in order to avoid re-jitting, reuse the same mj_data shape + save = d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth + d = mujoco.MjData(m) + d.qpos, d.qvel, d.time, d.qacc_warmstart, d.qacc_smooth = save + dx = mjx.device_put(d) + + mujoco.mj_step(m, d) + dx = forward_jit_fn(mx, dx) + + def cost(qacc): + jaref = np.zeros(d.nefc) + mujoco.mj_mulJacVec(m, d, jaref, qacc) + jaref -= d.efc_aref + cost = np.array([0.0]) + mujoco.mj_constraintUpdate(m, d, jaref, cost, 0) + return cost[0] + + cost_mj, cost_mjx = cost(d.qacc), cost(dx.qacc) + + self.assertLessEqual( + cost_mjx, + cost_mj * 1.01, + msg=f'mismatch: {fname} at step {i}, cost too high', + ) + _assert_attr_eq(d, dx, 'qfrc_constraint', i, fname, atol=1e-1, rtol=1e-1) + _assert_attr_eq(d, dx, 'qacc', i, fname, atol=1e-1, rtol=1e-1) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py new file mode 100644 index 00000000..cfbdfadf --- /dev/null +++ b/mjx/mujoco/mjx/_src/support.py @@ -0,0 +1,81 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Engine support functions.""" + +from typing import Tuple + +import jax +from jax import numpy as jp +from mujoco.mjx._src import scan +# pylint: disable=g-importing-member +from mujoco.mjx._src.types import Data +from mujoco.mjx._src.types import Model +# pylint: enable=g-importing-member + + +def jac( + m: Model, d: Data, point: jax.Array, body_id: jax.Array +) -> Tuple[jax.Array, jax.Array]: + """Compute pair of (NV, 3) Jacobians of global point attached to body.""" + fn = lambda carry, b: b if carry is None else b + carry + mask = (jp.arange(m.nbody) == body_id) * 1 + mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True) + mask = mask[jp.array(m.dof_bodyid)] > 0 + + offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]] + jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof) + jacp = jax.vmap(jp.multiply)(jacp, mask) + jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask) + + return jacp, jacr + + +def jac_dif_pair( + m: Model, + d: Data, + pos: jax.Array, + body_1: jax.Array, + body_2: jax.Array, +) -> jax.Array: + """Compute Jacobian difference for two body points.""" + jacp2, _ = jac(m, d, pos, body_2) + jacp1, _ = jac(m, d, pos, body_1) + return jacp2 - jacp1 + + +def apply_ft( + m: Model, + d: Data, + force: jax.Array, + torque: jax.Array, + point: jax.Array, + body_id: jax.Array, +) -> jax.Array: + """Apply Cartesian force and torque.""" + jacp, jacr = jac(m, d, point, body_id) + return jacp @ force + jacr @ torque + + +def xfrc_accumulate(m: Model, d: Data) -> jax.Array: + """Accumulate xfrc_applied into a qfrc.""" + qfrc = jax.vmap(apply_ft, in_axes=(None, None, 0, 0, 0, 0))( + m, + d, + d.xfrc_applied[:, :3], + d.xfrc_applied[:, 3:], + d.xipos, + jp.arange(m.nbody), + ) + return jp.sum(qfrc, axis=0) diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py new file mode 100644 index 00000000..fe88fc84 --- /dev/null +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -0,0 +1,81 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for support.""" + +from absl.testing import absltest +from absl.testing import parameterized +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx +from mujoco.mjx._src import support +from mujoco.mjx._src import test_util +import numpy as np + + +class SupportTest(parameterized.TestCase): + + @parameterized.parameters(set(test_util.TEST_FILES) - {'convex.xml'}) + def test_jac(self, fname): + np.random.seed(0) + + m = test_util.load_test_file(fname) + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + mx = mjx.device_put(m) + dx = mjx.device_put(d) + point = np.random.randn(3) + body = np.random.choice(m.nbody) + jacp, jacr = jax.jit(support.jac)(mx, dx, point, body) + + jacp_expected, jacr_expected = np.zeros((3, m.nv)), np.zeros((3, m.nv)) + mujoco.mj_jac(m, d, jacp_expected, jacr_expected, point, body) + np.testing.assert_almost_equal(jacp, jacp_expected.T, 6) + np.testing.assert_almost_equal(jacr, jacr_expected.T, 6) + + def test_xfrc_accumulate(self): + """Tests that xfrc_accumulate ouput matches mj_xfrcAccumulate.""" + np.random.seed(0) + + m = test_util.load_test_file('ant.xml') + d = mujoco.MjData(m) + mujoco.mj_step(m, d) + mx = mjx.device_put(m) + dx = mjx.device_put(d) + self.assertFalse((dx.xipos == 0.0).all()) + + xfrc = np.random.rand(*dx.xfrc_applied.shape) + + d.xfrc_applied[:] = xfrc + dx = dx.replace(xfrc_applied=jp.array(xfrc)) + + qfrc = jax.jit(support.xfrc_accumulate)(mx, dx) + qfrc_expected = np.zeros(m.nv) + for i in range(1, m.nbody): + mujoco.mj_applyFT( + m, + d, + d.xfrc_applied[i, :3], + d.xfrc_applied[i, 3:], + d.xipos[i], + i, + qfrc_expected, + ) + + np.testing.assert_almost_equal(qfrc, qfrc_expected, 6) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/test_util.py b/mjx/mujoco/mjx/_src/test_util.py new file mode 100644 index 00000000..5ac50816 --- /dev/null +++ b/mjx/mujoco/mjx/_src/test_util.py @@ -0,0 +1,340 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Utilities for testing.""" + +import sys +from typing import Dict, List, Tuple +from xml.etree import ElementTree as ET + +from etils import epath +import mujoco +import numpy as np + +TEST_FILES: List[str] = [ + 'ant.xml', + 'ball_pendulum.xml', + 'cherry_pendulum.xml', + 'convex.xml', + 'humanoid.xml', + 'mixed_joint_pendulum.xml', + 'single_pendulum.xml', + 'slide_pendulum.xml', + 'triple_pendulum.xml', + 'triple_pendulum_free.xml', + 'weld.xml', +] + +_ACTUATOR_TYPES = ['motor', 'velocity', 'position', 'general', 'intvelocity'] +_JOINT_TYPES = ['free', 'hinge', 'slide', 'ball'] +_JOINT_AXES = ['1 0 0', '0 1 0', '0 0 1'] +_FRICTIONS = ['1.2 0.003 0.0002', '0.2 0.0001 0.0005'] +_KP_POS = ['1', '2'] +_KP_INTVEL = ['10000', '2000'] +_KV_VEL = ['123', '1'] +_PAIR_FRICTIONS = ['1.2 0.9 0.003 0.0002 0.0001'] +_SOLREFS = ['0.04 1.01', '0.05 1.02', '0.03 1.1', '0.015 1.0'] +_SOLIMPS = [ + '0.75 0.94 0.002 0.2 2', + '0.8 0.99 0.001 0.3 6', + '0.6 0.9 0.003 0.1 1', +] +_DIMS = ['3'] +_MARGINS = ['0.0', '0.01', '0.02'] +_GAPS = ['0.0', '0.005'] +_GEARS = ['20', '50', '100'] + + +def p(pct: int) -> bool: + assert 0 <= pct <= 100 + return np.random.uniform(low=0, high=100) < pct + + +def _make_joint(joint_type: str, name: str) -> Dict[str, str]: + """Returns attributes for a joint.""" + joint_attr = {'type': joint_type, 'name': name} + + if joint_type not in ('free', 'ball'): + joint_attr['axis'] = np.random.choice(_JOINT_AXES) + lb, ub = -np.random.uniform() * 90, np.random.uniform() * 90 + joint_attr['range'] = f'{lb:.2f} {ub:.2f}' + elif joint_type == 'ball': + joint_attr['axis'] = '1 0 0' + ub = np.random.uniform() * 90 + joint_attr['range'] = f'0.0 {ub:.2f}' + + if p(50) and joint_type != 'free': + lb, ub = -np.random.uniform(), np.random.uniform() + joint_attr['actuatorfrcrange'] = f'{lb:.2f} {ub:.2f}' + + if joint_type not in ('free',): + joint_attr['damping'] = '{:.2f}'.format(np.random.uniform() * 20) + joint_attr['stiffness'] = '{:.2f}'.format(np.random.uniform() * 20) + + return joint_attr + + +def _geom_solparams( + pair: bool = False, enable_contact: bool = True +) -> Dict[str, str]: + """Returns geom solver parameters.""" + params = { + 'contype': np.random.choice(['0', '1']) if enable_contact else '0', + 'conaffinity': np.random.choice(['0', '1']) if enable_contact else '0', + 'priority': np.random.choice(['-1', '2']), + 'solmix': np.random.choice(['0.0', '1.6']), + 'friction': np.random.choice(_FRICTIONS), + 'condim': np.random.choice(_DIMS), + } + pair_params = { + 'solreffriction': np.random.choice(_SOLREFS), + 'friction': np.random.choice(_PAIR_FRICTIONS), + 'condim': np.random.choice(_DIMS), + } + params = pair_params if pair else params + params.update({ + 'solimp': np.random.choice(_SOLIMPS), + 'solref': np.random.choice(_SOLREFS), + 'margin': np.random.choice(_MARGINS), + 'gap': np.random.choice(_GAPS), + }) + + return params + + +def _make_geom( + pos: str, size: float, name: str, enable_contact: bool = True +) -> Dict[str, str]: + """Returns attributes for a sphere geom.""" + attr = { + 'pos': pos, + 'type': 'sphere', + 'name': name, + 'size': f'{size:.2f}', + 'mass': '1', + } + attr.update(_geom_solparams(pair=False, enable_contact=enable_contact)) + + return attr + + +def _make_actuator(actuator_type: str, joint: str) -> Dict[str, str]: + """Returns attributes for an actuator.""" + attr = {'joint': joint} + if actuator_type == 'motor': + attr['gear'] = np.random.choice(_GEARS) + elif actuator_type == 'position': + attr['kp'] = np.random.choice(_KP_POS) + elif actuator_type == 'general': + attr['biastype'] = 'affine' + attr['gainprm'] = '35 0 0' + attr['biasprm'] = '0 -35 -0.65' + elif actuator_type == 'intvelocity': + attr['kp'] = np.random.choice(_KP_INTVEL) + lb, ub = -np.random.uniform(), np.random.uniform() + attr['actrange'] = f'{lb:.2f} {ub:.2f}' + elif actuator_type == 'velocity': + attr['kv'] = np.random.choice(_KV_VEL) + + if p(50) and actuator_type != 'intvelocity': + lb, ub = -np.random.uniform(), np.random.uniform() + attr['ctrlrange'] = f'{lb:.2f} {ub:.2f}' + + if p(50): + lb, ub = -np.random.uniform(), np.random.uniform() + attr['forcerange'] = f'{lb*10:.2f} {ub*10:.2f}' + + return attr + + +def create_mjcf( + seed: int, + min_trees: int = 1, + max_trees: int = 1, + max_tree_depth: int = 5, + body_pos: Tuple[float, float, float] = (0.0, 0.0, -0.5), + geom_pos: Tuple[float, float, float] = (0.0, 0.0, 0.0), + max_stacked_joints=4, + max_geoms_per_body=2, + max_contact_excludes=1, + max_contact_pairs=4, + disable_actuation_pct: int = 0, + add_actuators: bool = False, + root_always_free: bool = False, + enable_contact: bool = True, +) -> str: + """Creates a random MJCF for testing. + + Args: + seed: seed for rng + min_trees: minimum number of kinematic trees to generate + max_trees: maximum number of kinematic trees to generate + max_tree_depth: the maximum tree depth + body_pos: the default body position relative to the parent + geom_pos: the default geom position in the body frame + max_stacked_joints: maximum number of joints to stack for each body + max_geoms_per_body: maximum number of geoms per body + max_contact_excludes: maximum number of bodies to exlude from contact + max_contact_pairs: maximum number of explicit geom contact pairs in the xml + disable_actuation_pct: the percentage of time to disable actuation via the + disable flag + add_actuators: whether to add actuators + root_always_free: if True, the root body of each kinematic tree has a free + joint with the world + enable_contact: if False, disables all contacts via contype/conaffinity + + Returns: + an XML string for the MuJoCo config + Raises: + AssertionError when args are not in the correct ranges + """ + np.random.seed(seed) + + assert min_trees <= max_trees + assert max_tree_depth >= 1 + assert 0 <= disable_actuation_pct <= 100 + assert max_stacked_joints >= 1 + assert max_geoms_per_body >= 1 + assert max_contact_excludes >= 1 + assert max_contact_pairs >= 1 + + mjcf = ET.Element('mujoco') + opt = ET.SubElement(mjcf, 'option', {'timestep': '0.005', 'solver': 'CG'}) + world = ET.SubElement(mjcf, 'worldbody') + ET.SubElement(mjcf, 'compiler', {'autolimits': 'true'}) + + # disable flags + if p(disable_actuation_pct): + ET.SubElement(opt, 'flag', {'actuation': 'disable'}) + + ET.SubElement( + world, + 'geom', + { + 'name': 'plane', + 'type': 'plane', + 'contype': '1' if enable_contact else '0', + 'conaffinity': '1' if enable_contact else '0', + 'size': '40 40 40', + }, + ) + + # kinematic trees + tree_depth = np.random.randint(1, max_tree_depth + 1) + + def make_tree(body: ET.Element, depth: int) -> None: + if depth >= tree_depth: + return + + z_pos = np.random.uniform(low=-1, high=1) * 0.01 # small jitter + pos = f'{body_pos[0]:.3f} {body_pos[1]:.3f} {body_pos[2] + z_pos:.3f}' + n_bodies = len(list(mjcf.iter('body'))) + child = ET.SubElement(body, 'body', {'pos': pos, 'name': f'body{n_bodies}'}) + + n_joints = len(list(mjcf.iter('joint'))) + for nj in range(np.random.randint(1, max_stacked_joints + 1)): + joint_type = np.random.choice(_JOINT_TYPES) + if nj == 0 and depth == 0 and root_always_free: + joint_type = 'free' + + # free joint only allowed at top level + while joint_type == 'free' and (depth > 0 or nj > 0): + joint_type = np.random.choice(_JOINT_TYPES) + + joint_attr = _make_joint(joint_type, name=f'joint{n_joints + nj}') + ET.SubElement(child, 'joint', joint_attr) + + prev_joints = child.findall('joint') + had_ball_or_free = any( + [j.get('type') in ('ball', 'free') for j in prev_joints] + ) + if had_ball_or_free: + break # do not stack more joints + + n_geoms = len(list(mjcf.iter('geom'))) + for _ in range(np.random.randint(1, max_geoms_per_body + 1)): + pos = ('{:.2f} ' * 3).format(*geom_pos).strip() + size = 0.2 + np.random.uniform(low=-1, high=1) * 0.02 + geom_attr = _make_geom( + pos, size, name=f'geom{n_geoms}', enable_contact=enable_contact + ) + ET.SubElement(child, 'geom', geom_attr) + n_geoms += 1 + + make_tree(child, depth + 1) + + num_trees = np.random.randint(min_trees, max_trees + 1) + for _ in range(num_trees): + make_tree(world, 0) + + # actuators + if add_actuators: + actuator = ET.SubElement(mjcf, 'actuator') + n_joints = len(list(mjcf.iter('joint'))) + nu = np.random.randint(1, n_joints + 1) + actuators = [] + for i in range(nu): + actuator_type = np.random.choice(_ACTUATOR_TYPES) + attr = _make_actuator(actuator_type, joint=f'joint{i}') + actuators.append((actuator_type, attr)) + + np.random.shuffle(actuators) + for typ, attr in actuators: + ET.SubElement(actuator, typ, attr) + + # contact pairs + contact = ET.SubElement(mjcf, 'contact') + geoms = list(mjcf.iter('geom')) + geom_names = [geom.get('name') for geom in geoms] + n_geoms = len(geoms) + pairs = set() + for _ in range(min(max_contact_pairs, n_geoms * (n_geoms - 1) // 2)): + if p(80): + continue + + geom1, geom2 = np.random.choice(geom_names, replace=False, size=2) + if geom1 > geom2: + geom1, geom2 = geom2, geom1 + + if (geom1, geom2) in pairs: + continue + + pairs.add((geom1, geom2)) + attr = {'geom1': geom1, 'geom2': geom2} + attr.update(_geom_solparams(pair=True)) + ET.SubElement(contact, 'pair', attr) + + # exclude contacts + bodies = list(mjcf.iter('body')) + body_names = [b.get('name') for b in bodies] + n_bodies = len(bodies) + for _ in range(min(max_contact_excludes, (n_bodies * (n_bodies - 1) // 2))): + if p(50): + continue + + body1, body2 = np.random.choice(body_names, replace=False, size=2) + ET.SubElement(contact, 'exclude', {'body1': body1, 'body2': body2}) + + # ElementTree.indent is not available before Python 3.9 + if sys.version_info.minor >= 9: + ET.indent(mjcf) + + return ET.tostring(mjcf).decode('utf-8') + + +def load_test_file(name: str) -> mujoco.MjModel: + """Loads a mujoco.MjModel based on the file name.""" + path = epath.resource_path('mujoco.mjx') / 'test_data' / name + m = mujoco.MjModel.from_xml_path(path.as_posix()) + return m diff --git a/mjx/mujoco/mjx/_src/test_util_test.py b/mjx/mujoco/mjx/_src/test_util_test.py new file mode 100644 index 00000000..e8d32dab --- /dev/null +++ b/mjx/mujoco/mjx/_src/test_util_test.py @@ -0,0 +1,38 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the test_util.""" + +from absl.testing import absltest +from etils import epath +from mujoco.mjx._src import test_util + + +class TestUtilTest(absltest.TestCase): + + def test_files_in_test_data_match(self): + directory = epath.resource_path('mujoco.mjx') / 'test_data' + files = set([f.name for f in directory.glob('*.xml')]) + self.assertSetEqual( + files, + set(test_util.TEST_FILES), + msg=( + '`_test_util.TEST_FILES` must match the files in the ' + 'test_data/*.xml directory' + ), + ) + + +if __name__ == '__main__': + absltest.main() diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py new file mode 100644 index 00000000..bd3f9fe1 --- /dev/null +++ b/mjx/mujoco/mjx/_src/types.py @@ -0,0 +1,660 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Base types used in MJX.""" + +import enum +from typing import Sequence + +import jax +import jax.numpy as jp +import mujoco +# pylint: disable=g-importing-member +from mujoco.mjx._src.dataclasses import PyTreeNode +# pylint: enable=g-importing-member +import numpy as np + + +class DisableBit(enum.IntFlag): + """Disable default feature bitflags. + + Attributes: + CONSTRAINT: entire constraint solver + EQUALITY: equality constraints + FRICTIONLOSS: joint and tendon frictionloss constraints + LIMIT: joint and tendon limit constraints + CONTACT: contact constraints + PASSIVE: passive forces + GRAVITY: gravitational forces + CLAMPCTRL: clamp control to specified range + WARMSTART: warmstart constraint solver + ACTUATION: apply actuation forces + REFSAFE: integrator safety: make ref[0]>=2*timestep + """ + CONSTRAINT = mujoco.mjtDisableBit.mjDSBL_CONSTRAINT + EQUALITY = mujoco.mjtDisableBit.mjDSBL_EQUALITY + LIMIT = mujoco.mjtDisableBit.mjDSBL_LIMIT + CONTACT = mujoco.mjtDisableBit.mjDSBL_CONTACT + PASSIVE = mujoco.mjtDisableBit.mjDSBL_PASSIVE + GRAVITY = mujoco.mjtDisableBit.mjDSBL_GRAVITY + CLAMPCTRL = mujoco.mjtDisableBit.mjDSBL_CLAMPCTRL + WARMSTART = mujoco.mjtDisableBit.mjDSBL_WARMSTART + ACTUATION = mujoco.mjtDisableBit.mjDSBL_ACTUATION + REFSAFE = mujoco.mjtDisableBit.mjDSBL_REFSAFE + EULERDAMP = mujoco.mjtDisableBit.mjDSBL_EULERDAMP + FILTERPARENT = mujoco.mjtDisableBit.mjDSBL_FILTERPARENT + # unsupported: FRICTIONLOSS, SENSOR, MIDPHASE + + +class JointType(enum.IntEnum): + """Type of degree of freedom. + + Attributes: + FREE: global position and orientation (quat) (7,) + BALL: orientation (quat) relative to parent (4,) + SLIDE: sliding distance along body-fixed axis (1,) + HINGE: rotation angle (rad) around body-fixed axis (1,) + """ + FREE = mujoco.mjtJoint.mjJNT_FREE + BALL = mujoco.mjtJoint.mjJNT_BALL + SLIDE = mujoco.mjtJoint.mjJNT_SLIDE + HINGE = mujoco.mjtJoint.mjJNT_HINGE + + def dof_width(self) -> int: + return {0: 6, 1: 3, 2: 1, 3: 1}[self.value] + + def qpos_width(self) -> int: + return {0: 7, 1: 4, 2: 1, 3: 1}[self.value] + + +class IntegratorType(enum.IntEnum): + """Integrator mode. + + Attributes: + EULER: semi-implicit Euler + RK4: 4th-order Runge Kutta + """ + EULER = mujoco.mjtIntegrator.mjINT_EULER + RK4 = mujoco.mjtIntegrator.mjINT_RK4 + # unsupported: IMPLICIT, IMPLICITFAST + + +class GeomType(enum.IntEnum): + """Type of geometry. + + Attributes: + PLANE: plane + HFIELD: height field + SPHERE: sphere + CAPSULE: capsule + ELLIPSOID: ellipsoid + CYLINDER: cylinder + BOX: box + MESH: mesh + """ + + PLANE = mujoco.mjtGeom.mjGEOM_PLANE + HFIELD = mujoco.mjtGeom.mjGEOM_HFIELD + SPHERE = mujoco.mjtGeom.mjGEOM_SPHERE + CAPSULE = mujoco.mjtGeom.mjGEOM_CAPSULE + ELLIPSOID = mujoco.mjtGeom.mjGEOM_ELLIPSOID + CYLINDER = mujoco.mjtGeom.mjGEOM_CYLINDER + BOX = mujoco.mjtGeom.mjGEOM_BOX + MESH = mujoco.mjtGeom.mjGEOM_MESH + # unsupported: NGEOMTYPES, ARROW*, LINE, SKIN, LABEL, NONE + + +class ConeType(enum.IntEnum): + """Type of friction cone. + + Attributes: + PYRAMIDAL: pyramidal + """ + PYRAMIDAL = mujoco.mjtCone.mjCONE_PYRAMIDAL + # unsupported: ELLIPTIC + + +class SolverType(enum.IntEnum): + """Constraint solver algorithm. + + Attributes: + CG: Conjugate gradient (primal) + """ + # unsupported: PGS, NEWTON + CG = mujoco.mjtSolver.mjSOL_CG + + +class EqType(enum.IntEnum): + """Type of equality constraint. + + Attributes: + CONNECT: connect two bodies at a point (ball joint) + WELD: fix relative position and orientation of two bodies + JOINT: couple the values of two scalar joints with cubic + """ + CONNECT = mujoco.mjtEq.mjEQ_CONNECT + WELD = mujoco.mjtEq.mjEQ_WELD + # unsupported: JOINT, TENDON, DISTANCE + + +class TrnType(enum.IntEnum): + """Type of actuator transmission. + + Attributes: + JOINT: force on joint + """ + JOINT = mujoco.mjtTrn.mjTRN_JOINT + # unsupported: JOINTINPARENT, SLIDERCRANK, TENDON, SITE, BODY + + +class DynType(enum.IntEnum): + """Type of actuator dynamics. + + Attributes: + NONE: no internal dynamics; ctrl specifies force + INTEGRATOR: integrator: da/dt = u + """ + NONE = mujoco.mjtDyn.mjDYN_NONE + INTEGRATOR = mujoco.mjtDyn.mjDYN_INTEGRATOR + FILTER = mujoco.mjtDyn.mjDYN_FILTER + # unsupported: FILTEREXACT, MUSCLE, USER + + +class GainType(enum.IntEnum): + """Type of actuator gain. + + Attributes: + FIXED: fixed gain + AFFINE: const + kp*length + kv*velocity + """ + FIXED = mujoco.mjtGain.mjGAIN_FIXED + AFFINE = mujoco.mjtGain.mjGAIN_AFFINE + # unsupported: MUSCLE, USER + + +class BiasType(enum.IntEnum): + """Type of actuator bias. + + Attributes: + NONE: no bias + AFFINE: const + kp*length + kv*velocity + """ + NONE = mujoco.mjtBias.mjBIAS_NONE + AFFINE = mujoco.mjtBias.mjBIAS_AFFINE + # unsupported: MUSCLE, USER + + +class Option(PyTreeNode): + """Physics options. + + Attributes: + timestep: timestep + tolerance: main solver tolerance + ls_tolerance: CG/Newton linesearch tolerance + gravity: gravitational acceleration (3,) + wind: wind (for lift, drag and viscosity) + density: density of medium + viscosity: viscosity of medium + has_fluid_params: automatically set by mjx if wind/density/viscosity are + nonzero. Not used by mj + integrator: integration mode + cone: type of friction cone + solver: solver algorithm + integrator: integration mode + iterations: number of main solver iterations + ls_iterations: maximum number of CG/Newton linesearch iterations + disableflags: bit flags for disabling standard features + """ + timestep: jax.Array + tolerance: jax.Array + ls_tolerance: jax.Array + # unsupported: apirate, impratio, noslip_tolerance, mpr_tolerance + gravity: jax.Array + wind: jax.Array + density: jax.Array + viscosity: jax.Array + has_fluid_params: bool + # unsupported: magnetic, o_margin, o_solref, o_solimp + integrator: IntegratorType + cone: ConeType + # unsupported: jacobian + solver: SolverType + iterations: int + ls_iterations: int + # unsupported: noslip_iterations, mpr_iterations + disableflags: DisableBit + # unsupported: enableflags + + +class Statistic(PyTreeNode): + """Model statistics (in qpos0). + + Attributes: + meaninertia: mean diagonal inertia + """ + meaninertia: jax.Array + # unsupported: meanmass, meansize, extent, center + + +class Model(PyTreeNode): + """Static model of the scene that remains unchanged with each physics step. + + Attributes: + nq: number of generalized coordinates = dim(qpos) + nv: number of degrees of freedom = dim(qvel) + nu: number of actuators/controls = dim(ctrl) + na: number of activation states = dim(act) + nbody: number of bodies + njnt: number of joints + ngeom: number of geoms + nmesh: number of meshes + npair: number of predefined geom pairs + nexclude: number of excluded geom pairs + neq: number of equality constraints + nnumeric: number of numeric custom fields + nM: number of non-zeros in sparse inertia matrix + opt: physics options + stat: model statistics + qpos0: qpos values at default pose (nq,) + qpos_spring: reference pose for springs (nq,) + body_parentid: id of body's parent (nbody,) + body_rootid: id of root above body (nbody,) + body_weldid: id of body that this body is welded to (nbody,) + body_jntnum: number of joints for this body (nbody,) + body_jntadr: start addr of joints; -1: no joints (nbody,) + body_dofnum: number of motion degrees of freedom (nbody,) + body_dofadr: start addr of dofs; -1: no dofs (nbody,) + body_geomnum: number of geoms (nbody,) + body_geomadr: start addr of geoms; -1: no geoms (nbody,) + body_pos: position offset rel. to parent body (nbody, 3) + body_quat: orientation offset rel. to parent body (nbody, 4) + body_ipos: local position of center of mass (nbody, 3) + body_iquat: local orientation of inertia ellipsoid (nbody, 4) + body_mass: mass (nbody,) + body_subtreemass: mass of subtree starting at this body (nbody,) + body_inertia: diagonal inertia in ipos/iquat frame (nbody, 3) + body_invweight0: mean inv inert in qpos0 (trn, rot) (nbody, 2) + jnt_type: type of joint (mjtJoint) (njnt,) + jnt_qposadr: start addr in 'qpos' for joint's data (njnt,) + jnt_dofadr: start addr in 'qvel' for joint's data (njnt,) + jnt_bodyid: id of joint's body (njnt,) + jnt_group: group for visibility (njnt,) + jnt_limited: does joint have limits (njnt,) + jnt_solref: constraint solver reference: limit (njnt, mjNREF) + jnt_solimp: constraint solver impedance: limit (njnt, mjNIMP) + jnt_pos: local anchor position (njnt, 3) + jnt_axis: local joint axis (njnt, 3) + jnt_stiffness: stiffness coefficient (njnt,) + jnt_range: joint limits (njnt, 2) + jnt_actfrcrange: range of total actuator force (njnt, 2) + jnt_margin: min distance for limit detection (njnt,) + dof_bodyid: id of dof's body (nv,) + dof_jntid: id of dof's joint (nv,) + dof_parentid: id of dof's parent; -1: none (nv,) + dof_Madr: dof address in M-diagonal (nv,) + dof_solref: constraint solver reference:frictionloss (nv, mjNREF) + dof_solimp: constraint solver impedance:frictionloss (nv, mjNIMP) + dof_frictionloss: dof friction loss (nv,) + dof_armature: dof armature inertia/mass (nv,) + dof_damping: damping coefficient (nv,) + dof_invweight0: diag. inverse inertia in qpos0 (nv,) + dof_M0: diag. inertia in qpos0 (nv,) + geom_type: geometric type (mjtGeom) (ngeom,) + geom_contype: geom contact type (ngeom,) + geom_conaffinity: geom contact affinity (ngeom,) + geom_condim: contact dimensionality (1, 3, 4, 6) (ngeom,) + geom_bodyid: id of geom's body (ngeom,) + geom_priority: geom contact priority (ngeom,) + geom_solmix: mixing coef for solref/imp in geom pair (ngeom,) + geom_solref: constraint solver reference: contact (ngeom, mjNREF) + geom_solimp: constraint solver impedance: contact (ngeom, mjNIMP) + geom_size: geom-specific size parameters (ngeom, 3) + geom_pos: local position offset rel. to body (ngeom, 3) + geom_quat: local orientation offset rel. to body (ngeom, 4) + geom_friction: friction for (slide, spin, roll) (ngeom, 3) + geom_margin: include in solver if dist 'Contact': + """Returns a contact filled with zeros.""" + return Contact( + dist=jp.zeros(shape), + pos=jp.zeros(shape + (3,)), + frame=jp.zeros(shape + (3, 3)), + includemargin=jp.zeros(shape), + friction=jp.zeros(shape + (5,)), + solref=jp.zeros(shape + (mujoco.mjNREF,)), + solreffriction=jp.zeros(shape + (mujoco.mjNREF,)), + solimp=jp.zeros(shape + (mujoco.mjNIMP,)), + dim=np.zeros(shape, dtype=np.int32), + geom1=jp.zeros(shape, dtype=jp.int32), + geom2=jp.zeros(shape, dtype=jp.int32), + efc_address=np.zeros(shape, dtype=np.int32), + ) + + +class Data(PyTreeNode): + """Dynamic state that updates each step. + + Attributes: + solver_niter: number of solver iterations, per island (mjNISLAND,) + ne: number of equality constraints + nf: number of friction constraints + nl: number of limit constraints + nefc: number of constraints + ncon: nubmer of contacts + time: simulation time + qpos: position (nq,) + qvel: velocity (nv,) + act: actuator activation (na,) + qacc_warmstart: acceleration used for warmstart (nv,) + ctrl: control (nu,) + qfrc_applied: applied generalized force (nv,) + xfrc_applied: applied Cartesian force/torque (nbody, 6) + eq_active: enable/disable constraints (neq,) + qacc: acceleration (nv,) + act_dot: time-derivative of actuator activation (na,) + xpos: Cartesian position of body frame (nbody, 3) + xquat: Cartesian orientation of body frame (nbody, 4) + xmat: Cartesian orientation of body frame (nbody, 3, 3) + xipos: Cartesian position of body com (nbody, 3) + ximat: Cartesian orientation of body inertia (nbody, 3, 3) + xanchor: Cartesian position of joint anchor (njnt, 3) + xaxis: Cartesian joint axis (njnt, 3) + geom_xpos: Cartesian geom position (ngeom, 3) + geom_xmat: Cartesian geom orientation (ngeom, 3, 3) + subtree_com: center of mass of each subtree (nbody, 3) + cdof: com-based motion axis of each dof (nv, 6) + cinert: com-based body inertia and mass (nbody, 10) + actuator_length: actuator lengths (nu,) + actuator_moment: actuator moments (nu, nv) + crb: com-based composite inertia and mass (nbody, 10) + qM: total inertia (sparse) (nM,) + qLD: L'*D*L factorization of M (sparse) (nM,) + qLDiagInv: 1/diag(D) (nv,) + qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,) + contact: list of all detected contacts (ncon,) + efc_J: constraint Jacobian (nefc, nv) + efc_frictionloss: frictionloss (friction) (nefc,) + efc_D: constraint mass (nefc,) + actuator_velocity: actuator velocities (nu,) + cvel: com-based velocity [3D rot; 3D tran] (nbody, 6) + cdof_dot: time-derivative of cdof (nv, 6) + qfrc_bias: C(qpos,qvel) (nv,) + qfrc_passive: passive force (nv,) + efc_aref: reference pseudo-acceleration (nefc,) + actuator_force: actuator force in actuation space (nu,) + qfrc_actuator: actuator force (nv,) + qfrc_smooth: net unconstrained force (nv,) + qacc_smooth: unconstrained acceleration (nv,) + qfrc_constraint: constraint force (nv,) + qfrc_inverse: net external force; should equal: (nv,) + qfrc_applied + J'*xfrc_applied + qfrc_actuator + efc_force: constraint force in constraint space (nefc,) + """ + # solver statistics: + solver_niter: jax.Array + # sizes (variable in MJ, constant in MJX) + ne: int + nf: int + nl: int + nefc: int + ncon: int + # global properties: + time: jax.Array + # state: + qpos: jax.Array + qvel: jax.Array + act: jax.Array + qacc_warmstart: jax.Array + # control: + ctrl: jax.Array + qfrc_applied: jax.Array + xfrc_applied: jax.Array + eq_active: jax.Array + # dynamics: + qacc: jax.Array + act_dot: jax.Array + # position dependent: + xpos: jax.Array + xquat: jax.Array + xmat: jax.Array + xipos: jax.Array + ximat: jax.Array + xanchor: jax.Array + xaxis: jax.Array + geom_xpos: jax.Array + geom_xmat: jax.Array + subtree_com: jax.Array + cdof: jax.Array + cinert: jax.Array + crb: jax.Array + actuator_length: jax.Array + actuator_moment: jax.Array + qM: jax.Array + qLD: jax.Array + qLDiagInv: jax.Array + qLDiagSqrtInv: jax.Array + contact: Contact + efc_J: jax.Array + efc_frictionloss: jax.Array + efc_D: jax.Array + # position, velocity dependent: + actuator_velocity: jax.Array + cvel: jax.Array + cdof_dot: jax.Array + qfrc_bias: jax.Array + qfrc_passive: jax.Array + efc_aref: jax.Array + # position, velcoity, control & acceleration dependent: + actuator_force: jax.Array + qfrc_actuator: jax.Array + qfrc_smooth: jax.Array + qacc_smooth: jax.Array + qfrc_constraint: jax.Array + qfrc_inverse: jax.Array + efc_force: jax.Array diff --git a/mjx/mujoco/mjx/benchmark/__init__.py b/mjx/mujoco/mjx/benchmark/__init__.py new file mode 100644 index 00000000..4ad51d69 --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== diff --git a/mjx/mujoco/mjx/benchmark/benchmark.py b/mjx/mujoco/mjx/benchmark/benchmark.py new file mode 100644 index 00000000..b0a59af0 --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/benchmark.py @@ -0,0 +1,116 @@ +# Copyright 2023 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Run benchmarks on various devices.""" + +import sys +import time + +from absl import flags +from etils import epath +import google_benchmark as benchmark +import jax +from jax import numpy as jp +import mujoco +from mujoco import mjx + +FLAGS = flags.FLAGS + +_PATHS = { + 'humanoid': 'benchmark/model/humanoid/humanoid.xml', + 'barkour': 'benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml', + 'shadow_hand': 'benchmark/model/shadow_hand/scene_right.xml', +} + + +_BATCH_SIZE = { + ('humanoid', 'TPU v5 lite'): 1024, + ('barkour', 'TPU v5 lite'): 1024, + ('shadow_hand', 'TPU v5 lite'): 1024, + ('humanoid', 'Tesla V100-SXM2-16GB'): 8192, + ('barkour', 'Tesla V100-SXM2-16GB'): 8192, + ('shadow_hand', 'Tesla V100-SXM2-16GB'): 4096, + ('humanoid', 'cpu'): 64, + ('barkour', 'cpu'): 64, + ('shadow_hand', 'cpu'): 64, +} + +flags.DEFINE_string('model', 'humanoid', 'Model to benchmark') +flags.DEFINE_string('device', 'cpu', 'Device benchmark is running on') + + +def _measure_fn(state, init_fn, step_fn, batch_size: int = 1024) -> float: + """Reports jit time and op time for a function.""" + + step_count = 100 if FLAGS.device == 'cpu' else 1000 + + @jax.jit + def run_batch(seed: jp.ndarray): + rngs = jax.random.split(jax.random.PRNGKey(seed), batch_size) + init_state = jax.vmap(init_fn)(rngs) + + @jax.vmap + def run(state): + def step(state, _): + state = step_fn(state) + return state, () + + return jax.lax.scan(step, state, (), length=step_count) + + return run(init_state) + + # run once to jit + beg = time.perf_counter() + jax.tree_util.tree_map(lambda x: x.block_until_ready(), run_batch(0)) + first_t = time.perf_counter() - beg + + times = [] + while state: + beg = time.perf_counter() + batch = run_batch(jp.array(len(times))) + jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch) + times.append(time.perf_counter() - beg) + + op_time = jp.mean(jp.array(times)) + batch_sps = batch_size * step_count / op_time + + state.counters['jit_time'] = first_t - op_time + state.counters['batch_sps'] = batch_sps + + +@benchmark.option.unit(benchmark.kSecond) +def _run(state: benchmark.State): + """Benchmark a model.""" + + f = epath.resource_path('mujoco.mjx') / _PATHS[FLAGS.model] + m = mujoco.MjModel.from_xml_path(f.as_posix()) + m = mjx.device_put(m) + + def init(rng): + d = mjx.make_data(m) + qvel = 0.01 * jax.random.normal(rng, shape=(m.nv,)) + d = d.replace(qvel=qvel) + return d + + def step(d): + return mjx.step(m, d) + + batch_size = _BATCH_SIZE[(FLAGS.model, jax.devices()[0].device_kind)] + _measure_fn(state, init, step, batch_size=batch_size) + + +if __name__ == '__main__': + FLAGS(sys.argv) + benchmark.register(_run, name=FLAGS.model + '_' + FLAGS.device) + benchmark.main() diff --git a/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md b/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md new file mode 100644 index 00000000..b7cfdcef --- /dev/null +++ b/mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md @@ -0,0 +1,35 @@ +# Google Barkour v0 Joystick Policy + +## Overview + +This folder contains a training script for a flat-terrain joystick policy for the [Barkour v0 Quadruped](https://ai.googleblog.com/2023/05/barkour-benchmarking-animal-level.html) which demonstrates sim2real transfer. + +`barkour_joystick.py` contains the environment definition, while the [colab](https://colab.research.google.com/github/google/brax/blob/main/experimental/barkour_v0/barkour_v0_joystick.ipynb) shows how to train the policy. + +

+ +

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