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

Merge branch 'usd-integration' of github.com:awesome-aj0123/mujoco into usd-integration

Browse Source
This commit is contained in:
Abhishek Joshi
2023-10-19 12:25:52 -05:00
parent f094f964ed 57b14f178b
commit 35097a6afa
185 changed files with 118687 additions and 4357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+1 -1
View File
@@ -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"
)
cmake/MujocoDependencies.cmake
+4 -3
View File
@@ -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
cmake/sdflib-optional-dependencies.patch
-585
View File
@@ -1,585 +0,0 @@
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 20551cf..0d2a364 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -70,15 +70,25 @@ add_custom_target(copyShaders ALL SOURCES ${SHADER_FILES})
# Add dependencies
add_subdirectory(libs)
+if(SDFLIB_USE_ENOKI)
+ target_link_libraries(${PROJECT_NAME} PUBLIC enoki)
+ target_link_libraries(${PROJECT_NAME} PUBLIC fcpw)
+ target_compile_definitions(${PROJECT_NAME} PUBLIC -DENOKI_AVAILABLE)
+endif()
+
+if(SDFLIB_USE_ASSIMP)
+ target_link_libraries(${PROJECT_NAME} PUBLIC assimp)
+ target_compile_definitions(${PROJECT_NAME} PUBLIC -DASSIMP_AVAILABLE)
+endif()
+
+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
+ target_link_libraries(${PROJECT_NAME} PUBLIC args)
+ target_link_libraries(${PROJECT_NAME} PUBLIC stb_image)
+endif()
+
target_link_libraries(${PROJECT_NAME} PUBLIC glm)
-target_link_libraries(${PROJECT_NAME} PUBLIC assimp)
-target_link_libraries(${PROJECT_NAME} PUBLIC args)
target_link_libraries(${PROJECT_NAME} PUBLIC spdlog)
target_link_libraries(${PROJECT_NAME} PUBLIC cereal)
-target_link_libraries(${PROJECT_NAME} PUBLIC enoki)
-target_link_libraries(${PROJECT_NAME} PUBLIC eigen)
-target_link_libraries(${PROJECT_NAME} PUBLIC fcpw)
-target_link_libraries(${PROJECT_NAME} PUBLIC stb_image)
target_link_libraries(${PROJECT_NAME} PUBLIC icg)
if(CMAKE_CXX_COMPILER_ID MATCHES GNU)
@@ -86,17 +96,22 @@ if(CMAKE_CXX_COMPILER_ID MATCHES GNU)
endif()
# Add openMP
-if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
- message("Enabling openmp llvm extension")
- set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp:llvm")
-else()
- find_package(OpenMP)
- if(NOT OpenMP_CXX_FOUND)
- message(FATAL_ERROR "OpenMP not found")
- endif()
- message("OpenMP version ${OpenMP_CXX_VERSION}")
- target_link_libraries(${PROJECT_NAME} PUBLIC OpenMP::OpenMP_CXX)
-endif()
+if(SDFLIB_USE_OPENMP)
+ if (CMAKE_CXX_COMPILER_ID STREQUAL "MSVC")
+ message("Enabling openmp llvm extension")
+ set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /openmp:llvm")
+ target_compile_definitions(${PROJECT_NAME} PUBLIC -DOPENMP_AVAILABLE)
+ else()
+ find_package(OpenMP)
+ if(OpenMP_CXX_FOUND)
+ message("OpenMP version ${OpenMP_CXX_VERSION}")
+ target_link_libraries(${PROJECT_NAME} PUBLIC OpenMP::OpenMP_CXX)
+ target_compile_definitions(${PROJECT_NAME} PUBLIC -DOPENMP_AVAILABLE)
+ else()
+ message("Disabling openmp")
+ endif()
+ endif()
+endif()
# Add executable
if (NOT UNIX)
@@ -160,6 +175,7 @@ if(SDFLIB_BUILD_DEBUG_APPS)
add_executable(GJKtest src/tools/GJKtest/main.cpp)
target_link_libraries(GJKtest PUBLIC ${PROJECT_NAME})
+ target_link_libraries(${PROJECT_NAME} PUBLIC eigen)
add_executable(CalculateInterpolationParameters src/tools/CalculateInterpolationParameters/main.cpp)
target_link_libraries(CalculateInterpolationParameters PUBLIC ${PROJECT_NAME})
diff --git a/include/SdfLib/ExactOctreeSdf.h b/include/SdfLib/ExactOctreeSdf.h
index 79ad82d..06dd7ac 100644
--- a/include/SdfLib/ExactOctreeSdf.h
+++ b/include/SdfLib/ExactOctreeSdf.h
@@ -214,6 +214,8 @@ private:
};
}
+#ifdef OPENMP_AVAILABLE
#include "ExactOctreeSdfDepthFirst.h"
+#endif
#endif
\ No newline at end of file
diff --git a/include/SdfLib/InterpolationMethods.h b/include/SdfLib/InterpolationMethods.h
index 077dfb2..f707d5b 100644
--- a/include/SdfLib/InterpolationMethods.h
+++ b/include/SdfLib/InterpolationMethods.h
@@ -4,7 +4,10 @@
#include <array>
#include "utils/TriangleUtils.h"
+
+#ifdef ENOKI_AVAILABLE
#include "enoki/array.h"
+#endif
namespace sdflib
{
@@ -236,15 +239,6 @@ struct TriLinearInterpolation
// outCoeff[63] = 8 * inValues[0][0] + 4 * inValues[0][1] * nodeSize + 4 * inValues[0][2] * nodeSize + 4 * inValues[0][3] * nodeSize + -8 * inValues[1][0] + 4 * inValues[1][1] * nodeSize + -4 * inValues[1][2] * nodeSize + -4 * inValues[1][3] * nodeSize + -8 * inValues[2][0] + -4 * inValues[2][1] * nodeSize + 4 * inValues[2][2] * nodeSize + -4 * inValues[2][3] * nodeSize + 8 * inValues[3][0] + -4 * inValues[3][1] * nodeSize + -4 * inValues[3][2] * nodeSize + 4 * inValues[3][3] * nodeSize + -8 * inValues[4][0] + -4 * inValues[4][1] * nodeSize + -4 * inValues[4][2] * nodeSize + 4 * inValues[4][3] * nodeSize + 8 * inValues[5][0] + -4 * inValues[5][1] * nodeSize + 4 * inValues[5][2] * nodeSize + -4 * inValues[5][3] * nodeSize + 8 * inValues[6][0] + 4 * inValues[6][1] * nodeSize + -4 * inValues[6][2] * nodeSize + -4 * inValues[6][3] * nodeSize + -8 * inValues[7][0] + 4 * inValues[7][1] * nodeSize + 4 * inValues[7][2] * nodeSize + 4 * inValues[7][3] * nodeSize + 0.0f;
// }
-// inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart)
-// {
-// return 0.0f
-// + values[0] + values[1] * fracPart[0] + values[2] * fracPart[0] * fracPart[0] + values[3] * fracPart[0] * fracPart[0] * fracPart[0] + values[4] * fracPart[1] + values[5] * fracPart[0] * fracPart[1] + values[6] * fracPart[0] * fracPart[0] * fracPart[1] + values[7] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] + values[8] * fracPart[1] * fracPart[1] + values[9] * fracPart[0] * fracPart[1] * fracPart[1] + values[10] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[11] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[12] * fracPart[1] * fracPart[1] * fracPart[1] + values[13] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[14] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[15] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1]
-// + values[16] * fracPart[2] + values[17] * fracPart[0] * fracPart[2] + values[18] * fracPart[0] * fracPart[0] * fracPart[2] + values[19] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] + values[20] * fracPart[1] * fracPart[2] + values[21] * fracPart[0] * fracPart[1] * fracPart[2] + values[22] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[23] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[24] * fracPart[1] * fracPart[1] * fracPart[2] + values[25] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[26] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[27] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[28] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[29] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[30] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[31] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2]
-// + values[32] * fracPart[2] * fracPart[2] + values[33] * fracPart[0] * fracPart[2] * fracPart[2] + values[34] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[35] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[36] * fracPart[1] * fracPart[2] * fracPart[2] + values[37] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[38] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[39] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[40] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[41] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[42] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[43] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[44] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[45] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[46] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[47] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2]
-// + values[48] * fracPart[2] * fracPart[2] * fracPart[2] + values[49] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[50] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[51] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[52] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[53] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[54] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[55] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[56] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[57] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[58] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[59] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[60] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[61] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[62] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[63] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2];
-// }
-
// inline static void interpolateVertexValues(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart, float nodeSize, std::array<float, VALUES_PER_VERTEX>& outValues)
// {
// outValues[0] = 0.0f
@@ -383,6 +377,7 @@ struct TriCubicInterpolation
outCoeff[63] = 8 * inValues[0][0] + 4 * inValues[0][1] + 4 * inValues[0][2] + 4 * inValues[0][3] + 2 * inValues[0][4] + 2 * inValues[0][5] + 2 * inValues[0][6] + 1 * inValues[0][7] + -8 * inValues[1][0] + 4 * inValues[1][1] + -4 * inValues[1][2] + -4 * inValues[1][3] + 2 * inValues[1][4] + 2 * inValues[1][5] + -2 * inValues[1][6] + 1 * inValues[1][7] + -8 * inValues[2][0] + -4 * inValues[2][1] + 4 * inValues[2][2] + -4 * inValues[2][3] + 2 * inValues[2][4] + -2 * inValues[2][5] + 2 * inValues[2][6] + 1 * inValues[2][7] + 8 * inValues[3][0] + -4 * inValues[3][1] + -4 * inValues[3][2] + 4 * inValues[3][3] + 2 * inValues[3][4] + -2 * inValues[3][5] + -2 * inValues[3][6] + 1 * inValues[3][7] + -8 * inValues[4][0] + -4 * inValues[4][1] + -4 * inValues[4][2] + 4 * inValues[4][3] + -2 * inValues[4][4] + 2 * inValues[4][5] + 2 * inValues[4][6] + 1 * inValues[4][7] + 8 * inValues[5][0] + -4 * inValues[5][1] + 4 * inValues[5][2] + -4 * inValues[5][3] + -2 * inValues[5][4] + 2 * inValues[5][5] + -2 * inValues[5][6] + 1 * inValues[5][7] + 8 * inValues[6][0] + 4 * inValues[6][1] + -4 * inValues[6][2] + -4 * inValues[6][3] + -2 * inValues[6][4] + -2 * inValues[6][5] + 2 * inValues[6][6] + 1 * inValues[6][7] + -8 * inValues[7][0] + 4 * inValues[7][1] + 4 * inValues[7][2] + 4 * inValues[7][3] + -2 * inValues[7][4] + -2 * inValues[7][5] + -2 * inValues[7][6] + 1 * inValues[7][7];
}
+#ifdef ENOKI_AVAILABLE
using vec4 = enoki::Array<float, 4>;
inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart)
@@ -433,6 +428,16 @@ struct TriCubicInterpolation
return sum;
}
+#else
+ inline static float interpolateValue(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart)
+ {
+ return 0.0f
+ + values[0] + values[1] * fracPart[0] + values[2] * fracPart[0] * fracPart[0] + values[3] * fracPart[0] * fracPart[0] * fracPart[0] + values[4] * fracPart[1] + values[5] * fracPart[0] * fracPart[1] + values[6] * fracPart[0] * fracPart[0] * fracPart[1] + values[7] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] + values[8] * fracPart[1] * fracPart[1] + values[9] * fracPart[0] * fracPart[1] * fracPart[1] + values[10] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[11] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] + values[12] * fracPart[1] * fracPart[1] * fracPart[1] + values[13] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[14] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] + values[15] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1]
+ + values[16] * fracPart[2] + values[17] * fracPart[0] * fracPart[2] + values[18] * fracPart[0] * fracPart[0] * fracPart[2] + values[19] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] + values[20] * fracPart[1] * fracPart[2] + values[21] * fracPart[0] * fracPart[1] * fracPart[2] + values[22] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[23] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] + values[24] * fracPart[1] * fracPart[1] * fracPart[2] + values[25] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[26] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[27] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] + values[28] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[29] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[30] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] + values[31] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2]
+ + values[32] * fracPart[2] * fracPart[2] + values[33] * fracPart[0] * fracPart[2] * fracPart[2] + values[34] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[35] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] + values[36] * fracPart[1] * fracPart[2] * fracPart[2] + values[37] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[38] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[39] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] + values[40] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[41] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[42] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[43] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[44] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[45] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[46] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] + values[47] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2]
+ + values[48] * fracPart[2] * fracPart[2] * fracPart[2] + values[49] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[50] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[51] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[2] * fracPart[2] * fracPart[2] + values[52] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[53] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[54] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[55] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[56] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[57] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[58] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[59] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[60] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[61] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[62] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2] + values[63] * fracPart[0] * fracPart[0] * fracPart[0] * fracPart[1] * fracPart[1] * fracPart[1] * fracPart[2] * fracPart[2] * fracPart[2];
+ }
+#endif
inline static glm::vec3 interpolateGradient(const std::array<float, NUM_COEFFICIENTS>& values, glm::vec3 fracPart)
{
@@ -493,4 +498,4 @@ struct TriCubicInterpolation
};
}
-#endif
\ No newline at end of file
+#endif
diff --git a/include/SdfLib/TrianglesInfluence.h b/include/SdfLib/TrianglesInfluence.h
index 3f3d33f..fc2ca52 100644
--- a/include/SdfLib/TrianglesInfluence.h
+++ b/include/SdfLib/TrianglesInfluence.h
@@ -1,7 +1,10 @@
#ifndef TRIANGLES_INFLUENCE_H
#define TRIANGLES_INFLUENCE_H
+#ifdef ENOKI_AVAILABLE
#include <fcpw/fcpw.h>
+#endif
+
#include "utils/Mesh.h"
#include "utils/TriangleUtils.h"
#include "OctreeSdfUtils.h"
@@ -1008,6 +1011,7 @@ struct VHQueries
}
};
+#ifdef ENOKI_AVAILABLE
template<typename T>
struct FCPWQueries
{
@@ -1118,6 +1122,8 @@ struct FCPWQueries
{
}
};
+#endif
+
}
-#endif
\ No newline at end of file
+#endif
diff --git a/include/SdfLib/utils/Mesh.h b/include/SdfLib/utils/Mesh.h
index 28d5486..7d21e44 100644
--- a/include/SdfLib/utils/Mesh.h
+++ b/include/SdfLib/utils/Mesh.h
@@ -4,9 +4,11 @@
#include <string>
#include <vector>
#include <glm/glm.hpp>
+#ifdef ASSIMP_AVAILABLE
#include <assimp/Importer.hpp>
#include <assimp/scene.h>
#include <assimp/postprocess.h>
+#endif
#include "SdfLib/utils/UsefullSerializations.h"
namespace sdflib
@@ -43,6 +45,23 @@ struct BoundingBox
return glm::length(glm::max(q,glm::vec3(0.0f))) + glm::min(glm::max(q.x, glm::max(q.y,q.z)),0.0f);
}
+ float getDistance(glm::vec3 point, glm::vec3& outGradient) const
+ {
+ glm::vec3 a = glm::abs(point) - getSize();
+ int k = a[0] > a[1] ? 0 : 1;
+ int l = a[2] > a[k] ? 2 : k;
+ if (a[l] < 0) {
+ outGradient[l] = point[l] / glm::abs(point[l]);
+ } else {
+ glm::vec3 b = glm::max(a, glm::vec3(0.0f));
+ float c = glm::length(b);
+ outGradient[0] = a[0] > 0 ? b[0] / c * point[0] / glm::abs(point[0]) : 0;
+ outGradient[1] = a[1] > 0 ? b[1] / c * point[1] / glm::abs(point[1]) : 0;
+ outGradient[2] = a[2] > 0 ? b[2] / c * point[2] / glm::abs(point[2]) : 0;
+ }
+ return getDistance(point);
+ }
+
template<class Archive>
void serialize(Archive & archive)
{
@@ -54,8 +73,10 @@ class Mesh
{
public:
Mesh() {}
+#ifdef ASSIMP_AVAILABLE
Mesh(std::string filePath);
Mesh(const aiMesh* mesh);
+#endif
Mesh(glm::vec3* vertices, uint32_t numVertices,
uint32_t* indices, uint32_t numIndices);
@@ -74,7 +95,9 @@ public:
void computeNormals();
void applyTransform(glm::mat4 trans);
private:
+#ifdef ASSIMP_AVAILABLE
void initMesh(const aiMesh* mesh);
+#endif
std::vector<glm::vec3> mVertices;
std::vector<uint32_t> mIndices;
@@ -83,4 +106,4 @@ private:
};
}
-#endif
\ No newline at end of file
+#endif
diff --git a/include/SdfLib/utils/TriangleUtils.h b/include/SdfLib/utils/TriangleUtils.h
index 9f930ed..6ee2304 100644
--- a/include/SdfLib/utils/TriangleUtils.h
+++ b/include/SdfLib/utils/TriangleUtils.h
@@ -2,6 +2,7 @@
#define TRIANGLE_UTILS_H
#include <glm/glm.hpp>
+#include <algorithm>
#include <vector>
#include <array>
#include <map>
@@ -401,4 +402,4 @@ namespace TriangleUtils
}
}
-#endif
\ No newline at end of file
+#endif
diff --git a/libs/CMakeLists.txt b/libs/CMakeLists.txt
index b48bf39..ea22b9a 100644
--- a/libs/CMakeLists.txt
+++ b/libs/CMakeLists.txt
@@ -14,37 +14,41 @@ if(NOT glm_lib_POPULATED)
endif()
# assimp
-FetchContent_Declare(assimp_lib
- GIT_REPOSITORY https://github.com/assimp/assimp.git
- GIT_TAG 9519a62dd20799c5493c638d1ef5a6f484e5faf1 # 5.2.5
-)
-
-if(NOT assimp_lib)
- FetchContent_Populate(assimp_lib)
-
- set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
- set(BUILD_SHARED_LIBS OFF)
- set(ASSIMP_BUILD_ASSIMP_TOOLS OFF)
- set(ASSIMP_BUILD_TESTS OFF)
- set(ASSIMP_INSTALL OFF)
- set(ASSIMP_INJECT_DEBUG_POSTFIX OFF)
- set(ASSIMP_BUILD_ASSIMP_VIEW OFF)
+if(SDF_USE_ASSIMP)
+ FetchContent_Declare(assimp_lib
+ GIT_REPOSITORY https://github.com/assimp/assimp.git
+ GIT_TAG 9519a62dd20799c5493c638d1ef5a6f484e5faf1 # 5.2.5
+ )
- add_subdirectory(${assimp_lib_SOURCE_DIR} ${assimp_lib_BINARY_DIR})
+ if(NOT assimp_lib)
+ FetchContent_Populate(assimp_lib)
+
+ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
+ set(BUILD_SHARED_LIBS OFF)
+ set(ASSIMP_BUILD_ASSIMP_TOOLS OFF)
+ set(ASSIMP_BUILD_TESTS OFF)
+ set(ASSIMP_INSTALL OFF)
+ set(ASSIMP_INJECT_DEBUG_POSTFIX OFF)
+ set(ASSIMP_BUILD_ASSIMP_VIEW OFF)
+
+ add_subdirectory(${assimp_lib_SOURCE_DIR} ${assimp_lib_BINARY_DIR})
+ endif()
endif()
# args
-FetchContent_Declare(args_lib
- GIT_REPOSITORY https://github.com/Taywee/args.git
- GIT_TAG a48e1f880813b367d2354963a58dedbf2b708584 # 6.3.0
-)
-
-FetchContent_GetProperties(args_lib)
-if(NOT args_lib_POPULATED)
- FetchContent_Populate(args_lib)
- add_library(args INTERFACE)
- target_include_directories(args INTERFACE ${args_lib_SOURCE_DIR})
-endif()
+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
+ FetchContent_Declare(args_lib
+ GIT_REPOSITORY https://github.com/Taywee/args.git
+ GIT_TAG a48e1f880813b367d2354963a58dedbf2b708584 # 6.3.0
+ )
+
+ FetchContent_GetProperties(args_lib)
+ if(NOT args_lib_POPULATED)
+ FetchContent_Populate(args_lib)
+ add_library(args INTERFACE)
+ target_include_directories(args INTERFACE ${args_lib_SOURCE_DIR})
+ endif()
+ endif()
# spdlog
FetchContent_Declare(spdlog_lib
@@ -76,47 +80,53 @@ if(NOT cereal_lib_POPULATED)
endif()
# Enoki
-FetchContent_Declare(enoki_lib
- GIT_REPOSITORY https://github.com/mitsuba-renderer/enoki.git
- GIT_TAG 2a18afa
-)
-FetchContent_GetProperties(enoki_lib)
-if(NOT enoki_lib_POPULATED)
- FetchContent_Populate(enoki_lib)
- add_library(enoki INTERFACE)
- add_subdirectory(${enoki_lib_SOURCE_DIR} ${enoki_lib_BINARY_DIR})
- target_include_directories(enoki INTERFACE ${enoki_lib_SOURCE_DIR}/include)
-endif()
-
-# eigen
-FetchContent_Declare(eigen_lib
-GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
-GIT_TAG 46126273552afe13692929523d34006f54c19719 # 3.4
-)
+if(SDFLIB_USE_ENOKI)
+ FetchContent_Declare(enoki_lib
+ GIT_REPOSITORY https://github.com/mitsuba-renderer/enoki.git
+ GIT_TAG 2a18afa
+ )
+ FetchContent_GetProperties(enoki_lib)
+ if(NOT enoki_lib_POPULATED)
+ FetchContent_Populate(enoki_lib)
+ add_library(enoki INTERFACE)
+ add_subdirectory(${enoki_lib_SOURCE_DIR} ${enoki_lib_BINARY_DIR})
+ target_include_directories(enoki INTERFACE ${enoki_lib_SOURCE_DIR}/include)
+ endif()
-FetchContent_GetProperties(eigen_lib)
-if(NOT eigen_lib_POPULATED)
- FetchContent_Populate(eigen_lib)
- add_library(eigen INTERFACE)
- target_include_directories(eigen INTERFACE ${eigen_lib_SOURCE_DIR})
+ # FCPW
+ FetchContent_Declare(fcpw_lib
+ GIT_REPOSITORY https://github.com/rohan-sawhney/fcpw.git
+ GIT_TAG dd65ec2
+ )
+
+ FetchContent_GetProperties(fcpw_lib)
+ if(NOT fcpw_lib_POPULATED)
+ FetchContent_Populate(fcpw_lib)
+ add_subdirectory(${fcpw_lib_SOURCE_DIR} ${fcpw_lib_BINARY_DIR})
+ target_include_directories(fcpw INTERFACE ${fcpw_lib_SOURCE_DIR})
+ endif()
endif()
-# FCPW
-FetchContent_Declare(fcpw_lib
- GIT_REPOSITORY https://github.com/rohan-sawhney/fcpw.git
- GIT_TAG dd65ec2
-)
-
-FetchContent_GetProperties(fcpw_lib)
-if(NOT fcpw_lib_POPULATED)
- FetchContent_Populate(fcpw_lib)
- add_subdirectory(${fcpw_lib_SOURCE_DIR} ${fcpw_lib_BINARY_DIR})
- target_include_directories(fcpw INTERFACE ${fcpw_lib_SOURCE_DIR})
-endif()
+# eigen
+if(SDFLIB_BUILD_DEBUG_APPS)
+ FetchContent_Declare(eigen_lib
+ GIT_REPOSITORY https://gitlab.com/libeigen/eigen.git
+ GIT_TAG 46126273552afe13692929523d34006f54c19719 # 3.4
+ )
+
+ FetchContent_GetProperties(eigen_lib)
+ if(NOT eigen_lib_POPULATED)
+ FetchContent_Populate(eigen_lib)
+ add_library(eigen INTERFACE)
+ target_include_directories(eigen INTERFACE ${eigen_lib_SOURCE_DIR})
+ endif()
+ endif()
# stb
-add_library(stb_image INTERFACE)
-target_include_directories(stb_image INTERFACE stb)
+if(SDFLIB_BUILD_APPS OR SDFLIB_BUILD_DEBUG_APPS)
+ add_library(stb_image INTERFACE)
+ target_include_directories(stb_image INTERFACE stb)
+ endif()
# icg
add_library(icg INTERFACE)
diff --git a/src/sdf/OctreeSdf.cpp b/src/sdf/OctreeSdf.cpp
index ef8ed4d..0e1eb97 100644
--- a/src/sdf/OctreeSdf.cpp
+++ b/src/sdf/OctreeSdf.cpp
@@ -6,7 +6,9 @@
#include "SdfLib/InterpolationMethods.h"
#include "sdf/OctreeSdfDepthFirst.h"
#include "sdf/OctreeSdfBreadthFirst.h"
+#ifdef OPENMP_AVAILABLE
#include "sdf/OctreeSdfBreadthFirstNoDelay.h"
+#endif
#include <array>
#include <stack>
@@ -46,8 +48,11 @@ OctreeSdf::OctreeSdf(const Mesh& mesh, BoundingBox box,
break;
case OctreeSdf::InitAlgorithm::CONTINUITY:
//initOctreeWithContinuity<PerNodeRegionTrianglesInfluence<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
- // initOctreeWithContinuity<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
+#ifdef OPENMP_AVAILABLE
initOctreeWithContinuityNoDelay<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule, numThreads);
+#else
+ initOctreeWithContinuity<VHQueries<InterpolationMethod>>(mesh, startDepth, depth, terminationThreshold, terminationRule);
+#endif
break;
// case OctreeSdf::InitAlgorithm::GPU_IMPLEMENTATION:
// Timer time;
@@ -78,7 +83,7 @@ float OctreeSdf::getDistance(glm::vec3 sample) const
startArrayPos.y < 0 || startArrayPos.y >= mStartGridSize ||
startArrayPos.z < 0 || startArrayPos.z >= mStartGridSize)
{
- return mBox.getDistance(sample) + glm::sqrt(3.0f) * mBox.getSize().x;
+ return mBox.getDistance(sample) + mMinBorderValue;
}
const OctreeNode* currentNode = &mOctreeData[startArrayPos.z * mStartGridXY + startArrayPos.y * mStartGridSize + startArrayPos.x];
@@ -108,7 +113,7 @@ float OctreeSdf::getDistance(glm::vec3 sample, glm::vec3& outGradient) const
startArrayPos.y < 0 || startArrayPos.y >= mStartGridSize ||
startArrayPos.z < 0 || startArrayPos.z >= mStartGridSize)
{
- return mBox.getDistance(sample) + mMinBorderValue;
+ return mBox.getDistance(sample, outGradient) + mMinBorderValue;
}
const OctreeNode* currentNode = &mOctreeData[startArrayPos.z * mStartGridXY + startArrayPos.y * mStartGridSize + startArrayPos.x];
@@ -253,4 +258,4 @@ void OctreeSdf::getDepthDensity(std::vector<float>& depthsDensity)
size *= 0.125f;
}
}
-}
\ No newline at end of file
+}
diff --git a/src/sdf/OctreeSdfDepthFirst.h b/src/sdf/OctreeSdfDepthFirst.h
index 53ee4b2..196d191 100644
--- a/src/sdf/OctreeSdfDepthFirst.h
+++ b/src/sdf/OctreeSdfDepthFirst.h
@@ -7,7 +7,10 @@
#include "SdfLib/OctreeSdfUtils.h"
#include <array>
#include <stack>
+#ifdef OPENMP_AVAILABLE
#include <omp.h>
+#endif
+
namespace sdflib
{
@@ -381,7 +384,9 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
};
const uint32_t voxlesPerAxis = 1 << startDepth;
+#ifdef OPENMP_AVAILABLE
if(numThreads < 2)
+#endif
{
// Create the grid
mOctreeData.resize(voxlesPerAxis * voxlesPerAxis * voxlesPerAxis);
@@ -402,7 +407,8 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
mValueRange = mainThread.valueRange;
}
- else
+#ifdef OPENMP_AVAILABLE
+ else
{
std::vector<ThreadContext> threadsContext(numThreads, mainThread);
@@ -511,6 +517,7 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
}
#endif
}
+#endif
#ifdef SDFLIB_PRINT_STATISTICS
SPDLOG_INFO("Used an octree of max depth {}", maxDepth);
@@ -544,4 +551,4 @@ void OctreeSdf::initOctree(const Mesh& mesh, uint32_t startDepth, uint32_t maxDe
}
}
-#endif
\ No newline at end of file
+#endif
diff --git a/src/utils/Mesh.cpp b/src/utils/Mesh.cpp
index b407d38..6fff5fe 100644
--- a/src/utils/Mesh.cpp
+++ b/src/utils/Mesh.cpp
@@ -5,6 +5,7 @@
namespace sdflib
{
+#ifdef ASSIMP_AVAILABLE
Mesh::Mesh(std::string filePath)
{
Assimp::Importer import;
@@ -28,6 +29,7 @@ Mesh::Mesh(const aiMesh* mesh)
{
initMesh(mesh);
}
+#endif
Mesh::Mesh(glm::vec3* vertices, uint32_t numVertices,
uint32_t* indices, uint32_t numIndices)
@@ -39,7 +41,7 @@ Mesh::Mesh(glm::vec3* vertices, uint32_t numVertices,
std::memcpy(mIndices.data(), indices, sizeof(uint32_t) * numIndices);
}
-
+#ifdef ASSIMP_AVAILABLE
void Mesh::initMesh(const aiMesh* mesh)
{
if(!(mesh->mPrimitiveTypes & aiPrimitiveType_TRIANGLE))
@@ -83,6 +85,7 @@ void Mesh::initMesh(const aiMesh* mesh)
computeNormals();
}
}
+#endif
void Mesh::computeBoundingBox()
{
@@ -134,4 +137,4 @@ void Mesh::applyTransform(glm::mat4 trans)
computeBoundingBox();
}
-}
\ No newline at end of file
+}
dist/mujoco.rc
Vendored
+4 -4
View File
@@ -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"
dist/simulate.rc
Vendored
+4 -4
View File
@@ -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"
doc/APIreference/APIglobals.rst
+1 -1
View File
@@ -522,7 +522,7 @@ shown in the table below. Their names are in the format ``mjKEY_XXX``. They corr
- Maximum number of UI rectangles.
Defined in `mjui.h <https://github.com/google-deepmind/mujoco/blob/main/include/mujoco/mjui.h>`_.
* - ``mjVERSION_HEADER``
- 238
- 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.
doc/XMLreference.rst
+633 -293
View File
File diff suppressed because it is too large Load Diff
doc/XMLschema.rst
+27 -9
View File
@@ -485,19 +485,19 @@
| :ref:`flexcomp | \* | :class: mjcf-attributes |
| <body-flexcomp>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<body-flexcomp-name>` | :ref:`class<body-flexcomp-class>` | :ref:`type<body-flexcomp-type>` | :ref:`dim<body-flexcomp-dim>` | |
| | | | :ref:`name<body-flexcomp-name>` | :ref:`class<body-flexcomp-class>` | :ref:`type<body-flexcomp-type>` | :ref:`group<body-flexcomp-group>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`flatskin<body-flexcomp-flatskin>` | :ref:`count<body-flexcomp-count>` | :ref:`spacing<body-flexcomp-spacing>` | :ref:`radius<body-flexcomp-radius>` | |
| | | | :ref:`dim<body-flexcomp-dim>` | :ref:`count<body-flexcomp-count>` | :ref:`spacing<body-flexcomp-spacing>` | :ref:`radius<body-flexcomp-radius>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`rigid<body-flexcomp-rigid>` | :ref:`mass<body-flexcomp-mass>` | :ref:`inertiabox<body-flexcomp-inertiabox>` | :ref:`scale<body-flexcomp-scale>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`file<body-flexcomp-file>` | :ref:`point<body-flexcomp-point>` | :ref:`element<body-flexcomp-element>` | :ref:`texcoord<body-flexcomp-texcoord>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`material<body-flexcomp-material>` | :ref:`rgba<body-flexcomp-rgba>` | :ref:`selfcollide<body-flexcomp-selfcollide>` | :ref:`flatskin<body-flexcomp-flatskin>` | |
| | | | :ref:`material<body-flexcomp-material>` | :ref:`rgba<body-flexcomp-rgba>` | :ref:`flatskin<body-flexcomp-flatskin>` | :ref:`pos<body-flexcomp-pos>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`pos<body-flexcomp-pos>` | :ref:`quat<body-flexcomp-quat>` | :ref:`axisangle<body-flexcomp-axisangle>` | :ref:`xyaxes<body-flexcomp-xyaxes>` | |
| | | | :ref:`quat<body-flexcomp-quat>` | :ref:`axisangle<body-flexcomp-axisangle>` | :ref:`xyaxes<body-flexcomp-xyaxes>` | :ref:`zaxis<body-flexcomp-zaxis>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`zaxis<body-flexcomp-zaxis>` | :ref:`euler<body-flexcomp-euler>` | | | |
| | | | :ref:`euler<body-flexcomp-euler>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
@@ -517,7 +517,9 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`friction<flexcomp-contact-friction>` | :ref:`solmix<flexcomp-contact-solmix>` | :ref:`solref<flexcomp-contact-solref>` | :ref:`solimp<flexcomp-contact-solimp>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`margin<flexcomp-contact-margin>` | :ref:`gap<flexcomp-contact-gap>` | | | |
| | | | :ref:`margin<flexcomp-contact-margin>` | :ref:`gap<flexcomp-contact-gap>` | :ref:`internal<flexcomp-contact-internal>` | :ref:`selfcollide<flexcomp-contact-selfcollide>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`activelayers<flexcomp-contact-activelayers>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
@@ -527,6 +529,20 @@
| | | | :ref:`id<flexcomp-pin-id>` | :ref:`range<flexcomp-pin-range>` | :ref:`grid<flexcomp-pin-grid>` | :ref:`gridrange<flexcomp-pin-gridrange>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flexcomp |br| |_2| |L| | | .. table:: |
| :ref:`plugin | \* | :class: mjcf-attributes |
| <flexcomp-plugin>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`plugin<flexcomp-plugin-plugin>` | :ref:`instance<flexcomp-plugin-instance>` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_3| plugin |br| |_3| |L| | | .. table:: |
| :ref:`config | \* | :class: mjcf-attributes |
| <plugin-config>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`key<plugin-config-key>` | :ref:`value<plugin-config-value>` | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| mujoco |br| |L| | | *no attributes* |
| :ref:`deformable<deformable>` | | |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
@@ -536,9 +552,9 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`name<deformable-flex-name>` | :ref:`group<deformable-flex-group>` | :ref:`dim<deformable-flex-dim>` | :ref:`radius<deformable-flex-radius>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`material<deformable-flex-material>` | :ref:`rgba<deformable-flex-rgba>` | :ref:`flatskin<deformable-flex-flatskin>` | :ref:`selfcollide<deformable-flex-selfcollide>` | |
| | | | :ref:`material<deformable-flex-material>` | :ref:`rgba<deformable-flex-rgba>` | :ref:`flatskin<deformable-flex-flatskin>` | :ref:`body<deformable-flex-body>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`body<deformable-flex-body>` | :ref:`vertex<deformable-flex-vertex>` | :ref:`element<deformable-flex-element>` | :ref:`texcoord<deformable-flex-texcoord>` | |
| | | | :ref:`vertex<deformable-flex-vertex>` | :ref:`element<deformable-flex-element>` | :ref:`texcoord<deformable-flex-texcoord>` | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flex |br| |_2| |L| | | .. table:: |
@@ -549,7 +565,9 @@
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`friction<flex-contact-friction>` | :ref:`solmix<flex-contact-solmix>` | :ref:`solref<flex-contact-solref>` | :ref:`solimp<flex-contact-solimp>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`margin<flex-contact-margin>` | :ref:`gap<flex-contact-gap>` | | | |
| | | | :ref:`margin<flex-contact-margin>` | :ref:`gap<flex-contact-gap>` | :ref:`internal<flex-contact-internal>` | :ref:`selfcollide<flex-contact-selfcollide>` | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
| | | | :ref:`activelayers<flex-contact-activelayers>` | | | | |
| | | +-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+-----------------------------------------------------------------+ |
+------------------------------------+----+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| |_2| flex |br| |_2| |L| | | .. table:: |
doc/changelog.rst
+104 -69
View File
@@ -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<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
see video.
- MJX is designed to work with on-device reinforcement learning algorithms. This Colab notebook demonstrates using
MJX along with reinforcement learning to train humanoid and quadruped robots to locomote: |colab|
- The MJX API is compatible with MuJoCo but is missing some features in this release. See the outline of
:ref:`MJX feature parity <MjxFeatureParity>` for more details.
.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb
.. youtube:: QewlEqIZi1o
:align: right
@@ -28,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<exWriting>` 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 <https://en.wikipedia.org/wiki/Simplicial_complex>`__ can be of dimension 1, 2
or 3, corresponding to stretchable lines, triangles or tetrahedra. Two new MJCF elements are used
to define flexes. The top-level :ref:`deformable<deformable>` section contains the low-level flex definition.
The :ref:`flexcomp<body-flexcomp>` element, similar to :ref:`composite<body-composite>` is a convenience macro for
creating deformables, and supports the GMSH tetrahedral file format.
- Added `shell <https://github.com/deepmind/mujoco/blob/main/plugin/elasticity/shell.cc>`__ passive force plugin,
computing bending forces using a constant precomputed Hessian (cotangent operator).
**Note**: This feature is still under development and subject to change. In particular, deformable object
functionality is currently available both via :ref:`deformable<CDeformable>` and :ref:`composite<CComposite>`,
and both are modifiable by the first-party
`elasticity plugins <https://github.com/google-deepmind/mujoco/tree/main/plugin/elasticity>`__. We expect some of
this functionality to be unified in the future.
.. youtube:: Vc1tq0fFvQA
:align: right
:width: 240px
4. Added constraint island discovery with :ref:`mj_island`. Constraint islands are disjoint sets of constraints
and degrees-of-freedom that do not interact. The only solver which currently supports islands is
:ref:`CG<option-solver>`. Island discovery can be activated using a new :ref:`enable flag<option-flag-island>`.
If island discovery is enabled, geoms, contacts and tendons will be colored according to the corresponding island,
see video. Island discovery is currently disabled for models that have deformable objects (see previous item).
5. Added :ref:`mjThreadPool` and :ref:`mjTask` which allow for multi-threaded operations within the MuJoCo engine
pipeline. If engine-internal threading is enabled, the following operations will be multi-threaded:
- Island constraint resolution, if island discovery is :ref:`enabled<option-flag-island>` and the
@@ -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<saTestspeed>` utility, exposed with the ``npoolthread`` flag.
.. youtube:: ra2bTiZHGlw
:align: right
:width: 240px
4. Added capability to initialize :ref:`composite<body-composite>` particles with arbitrary positions.
5. Added `shell <https://github.com/deepmind/mujoco/blob/main/plugin/elasticity/shell.cc>`__ 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<body-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<siStack>` in a fully encapsulated way (i.e.,
without introducing a local variable at the call site).
7. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an
8. Renamed ``mj_stackAlloc`` to :ref:`mj_stackAllocNum`. The new function :ref:`mj_stackAllocByte` allocates an
arbitrary number of bytes and has an additional argument for specifying the alignment of the returned pointer.
**Migration:** The functionality for allocating ``mjtNum`` arrays is now available via :ref:`mj_stackAllocNum`.
8. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``,
9. Renamed the ``nstack`` field in :ref:`mjModel` and :ref:`mjData` to ``narena``. Changed ``narena``, ``pstack``,
and ``maxuse_stack`` to count number of bytes rather than number of :ref:`mjtNum` |-| s.
9. Changed :ref:`mjData.solver<mjData>`, the array used to collect solver diagnostic information.
This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix.
Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island.
If the solver does not use islands, only row 0 is filled.
10. Changed :ref:`mjData.solver<mjData>`, the array used to collect solver diagnostic information.
This array of :ref:`mjSolverStat` structs is now of length ``mjNISLAND * mjNSOLVER``, interpreted as as a matrix.
Each row of length ``mjNSOLVER`` contains separate solver statistics for each constraint island.
If the solver does not use islands, only row 0 is filled.
- The new constant :ref:`mjNISLAND<glNumeric>` was set to 20.
- :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
- Added :ref:`mjData.solver_nisland<mjData>`: the number of islands for which the solver ran.
- Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer
vectors of length ``mjNISLAND``.
- The new constant :ref:`mjNISLAND<glNumeric>` was set to 20.
- :ref:`mjNSOLVER<glNumeric>` was reduced from 1000 to 200.
- Added :ref:`mjData.solver_nisland<mjData>`: the number of islands for which the solver ran.
- Renamed ``mjData.solver_iter`` to ``solver_niter``. Both this member and ``mjData.solver_nnz`` are now integer
vectors of length ``mjNISLAND``.
10. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute.
11. Removed ``mjOption.collision`` and the associated ``option/collision`` attribute.
**Migration:**
@@ -93,40 +124,39 @@ General
:ref:`conaffinity<body-geom-conaffinity>` attributes in the model and then setting them globally to ``0`` using
|br| ``<default> <geom contype="0" conaffinity="0"/> </default>``.
11. Removed the :at:`rope` and :at:`cloth` composite objects.
12. Removed the :at:`rope` and :at:`cloth` composite objects.
**Migration:** Users should use the :at:`cable` and :at:`shell` elasticity plugins.
12. Added :ref:`mjData.eq_active<mjData>` user input variable, for enabling/disabling the state of equality
13. Added :ref:`mjData.eq_active<mjData>` user input variable, for enabling/disabling the state of equality
constraints. Renamed ``mjModel.eq_active`` to :ref:`mjModel.eq_active0<mjModel>`, which now has the semantic of
"initial value of ``mjData.eq_active``".
Fixes `#876 <https://github.com/google-deepmind/mujoco/discussions/876>`__
"initial value of ``mjData.eq_active``". Fixes :github:issue:`876`.
**Migration:** Replace uses of ``mjModel.eq_active`` with ``mjData.eq_active``.
13. Changed the default of :ref:`autolimits<compiler-autolimits>` from "false" to "true". This is a minor breaking
14. Changed the default of :ref:`autolimits<compiler-autolimits>` from "false" to "true". This is a minor breaking
change. The potential breakage applies to models which have elements with "range" defined and "limited" not set.
Such models cannot be loaded since version 2.2.2 (July 2022).
14. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
15. Added a new :ref:`dyntype<actuator-general-dyntype>`, ``filterexact``, which updates first-order filter states with
the exact formula rather than with Euler integration.
15. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
16. Added an actuator attribute, :ref:`actearly<actuator-general-actearly>`, which uses semi-implicit integration for
actuator forces: using the next step's actuator state to compute the current actuator forces.
16. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
17. Renamed ``actuatorforcerange`` and ``actuatorforcelimited``, introduced in the previous version to
:ref:`actuatorfrcrange<body-joint-actuatorfrcrange>` and
:ref:`actuatorfrclimited<body-joint-actuatorfrclimited>`, respectively.
17. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
18. Added the flag :ref:`eulerdamp<option-flag-eulerdamp>`, which disables implicit integration of joint damping in the
Euler integrator. See the :ref:`Numerical Integration<geIntegration>` section for more details.
18. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
19. Added the flag :ref:`invdiscrete<option-flag-invdiscrete>`, which enables discrete-time inverse dynamics for all
:ref:`integrators<option-integrator>` other than ``RK4``. See the flag documentation for more details.
19. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
20. Added :ref:`ls_iterations<option-ls_iterations>` and :ref:`ls_tolerance<option-ls_tolerance>` options for adjusting
linesearch stopping criteria in CG and Newton solvers. These can be useful for performance tuning.
20. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
mesh assets. Fixes `#409 <https://github.com/google-deepmind/mujoco/issues/409>`__ .
21. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
21. Added ``mesh_pos`` and ``mesh_quat`` fields to :ref:`mjModel` to store the normalizing transformation applied to
mesh assets. Fixes :github:issue:`409`.
22. Added camera :ref:`resolution<body-camera-resolution>` attribute and :ref:`camprojection<sensor-camprojection>`
sensor. If camera resolution is set to positive values, the camera projection sensor will report the location of a
target site, projected onto the camera image, in pixel coordinates.
22. Added :ref:`camera<body-camera>` calibration attributes:
23. Added :ref:`camera<body-camera>` calibration attributes:
- The new attributes are :ref:`resolution<body-camera-resolution>`, :ref:`focal<body-camera-focal>`,
:ref:`focalpixel<body-camera-focalpixel>`, :ref:`principal<body-camera-principal>`,
@@ -135,23 +165,21 @@ General
attributes are specified. See the following
`example model <https://github.com/deepmind/mujoco/blob/main/test/engine/testdata/vis_visualize/frustum.xml>`__.
- Note that these attributes only take effect for offline rendering and do not affect interactive visualisation.
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 <https://github.com/google-deepmind/mujoco/pull/978>`__ by
`Levi Burner <https://github.com/aftersomemath>`__.
24. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the
Contribution :github:pull:`978` by `Levi Burner <https://github.com/aftersomemath>`__.
25. Deleted the code sample ``testxml``. The functionality provided by this utility is implemented in the
`WriteReadCompare <https://github.com/google-deepmind/mujoco/blob/main/test/xml/xml_native_writer_test.cc>`__ test.
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 <https://github.com/google-deepmind/mujoco/issues/870>`__ where calling ``update_scene`` with an invalid
camera name used the default camera.
27. Added ``user_scn`` to the :ref:`passive viewer<PyViewerPassive>` handle, which allows users to add custom
visualization geoms (`#1023 <https://github.com/google-deepmind/mujoco/issues/870>`__).
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<PyViewerPassive>` handle, which allows users to add custom
visualization geoms (:github:issue:`1023`).
29. Added optional boolean keyword arguments ``show_left_ui`` and ``show_right_ui`` to the functions ``viewer.launch``
and ``viewer.launch_passive``, which allow users to launch a viewer with UI panels hidden.
Simulate
@@ -161,21 +189,31 @@ Simulate
:align: right
:width: 240px
29. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
30. Added **state history** mechanism to :ref:`simulate<saSimulate>` and the managed
:ref:`Python viewer<PyViewerManaged>`. State history can be viewed by scrubbing the History slider and (more
precisely) with the left and right arrow keys. See screen capture:
30. The ``LOADING...`` label is now shown correctly.
`Contribution <https://github.com/google-deepmind/mujoco/pull/1070>`__ by
31. The ``LOADING...`` label is now shown correctly. Contribution :github:pull:`1070` by
`Levi Burner <https://github.com/aftersomemath>`__.
Documentation
^^^^^^^^^^^^^
.. youtube:: nljr0X79vI0
:align: right
:width: 240px
32. Added :doc:`detailed documentation <computation/fluid>` of fluid force modeling, and an illustrative example model
showing `tumbling cards <https://github.com/google-deepmind/mujoco/blob/main/model/cards/cards.xml>`__ using the
ellipsoid-based fluid model.
Bug fixes
^^^^^^^^^
31. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
33. Fixed a bug that was causing :ref:`geom margin<body-geom-margin>` to be ignored during the construction of
midphase collision trees.
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<PyViewerPassive>` handle now exposes ``update_hfield``, ``update_mesh``, and
``update_texture`` methods to allow users to update renderable assets.
(`#812 <https://github.com/google-deepmind/mujoco/issues/812>`_,
`#958 <https://github.com/google-deepmind/mujoco/issues/958>`_,
`#965 <https://github.com/google-deepmind/mujoco/issues/965>`_)
#. Allow a custom keyboard event callback to be specified in the :ref:`passive viewer<PyViewerPassive>`.
(`#766 <https://github.com/google-deepmind/mujoco/issues/766>`_)
#. Fix GLFW crash when Python exits while the passive viewer is running.
(`#790 <https://github.com/google-deepmind/mujoco/issues/790>`_)
(Issues :github:issue:`812`, :github:issue:`958`, :github:issue:`965`).
#. Allow a custom keyboard event callback to be specified in the :ref:`passive viewer<PyViewerPassive>`
(:github:issue:`766`).
#. Fix GLFW crash when Python exits while the passive viewer is running (:github:issue:`790`).
Models
^^^^^^
doc/computation/fluid.rst
+91 -72
View File
@@ -26,36 +26,55 @@ to positive values. These parameters correspond to the density :math:`\rho` and
Inertia model
-------------
In this model, the shape of each body, for fluid dynamics purposes, is assumed to be the *equivalent inertia box*,
which can also be visualized. Each forward-facing (relative to the linear velocity) face of the box experiences force
along its normal direction. All faces also experience torque due to the angular velocity; this torque is obtained by
integrating the force resulting from the rotation over the surface area. In this sub-section, let :math:`v` and
:math:`\omega` denote the linear and angular body velocity in the body local frame (aligned with the equivalent
inertia box), and :math:`s` the 3D vector of box sizes. When the contributions from all faces are added, the resulting
force and torque applied to the body by a fluid of density :math:`\rho`, in local body coordinates, have the
:math:`i`-th component
In this model the shape of each body, for fluid dynamics purposes, is assumed to be the *equivalent inertia box*,
which can also be visualized. For a body with mass :math:`\mathcal{M}` and inertia matrix :math:`\mathcal{I}`, the
half-dimensions (i.e. half-width, half-depth and half-height) of the equivalent inertia box are
.. math::
\begin{align*}
r_x = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{yy} + \mathcal{I}_{zz} - \mathcal{I}_{xx} \right)} \\
r_y = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{zz} + \mathcal{I}_{xx} - \mathcal{I}_{yy} \right)} \\
r_z = \sqrt{\frac{3}{2 \mathcal{M}} \left(\mathcal{I}_{xx} + \mathcal{I}_{yy} - \mathcal{I}_{zz} \right)}
\end{align*}
Let :math:`\mathbf{v}` and :math:`\boldsymbol{\omega}` denote the linear and angular body velocity of the body in
the body-local frame (aligned with the equivalent inertia box). The force :math:`\mathbf{f}_{\text{inertia}}` and
torque :math:`\mathbf{g}_{\text{inertia}}` exerted by the fluid onto the solid are the sum of the terms
.. math::
\begin{align*}
\mathbf{f}_{\text{inertia}} &= \mathbf{f}_D + \mathbf{f}_V \\
\mathbf{g}_{\text{inertia}} &= \mathbf{g}_D + \mathbf{g}_V
\end{align*}
Here subscripts :math:`D` and :math:`V` denote quadratic Drag and Viscous resistance.
The quadratic drag terms depend on the density :math:`\rho` of the fluid, scale quadratically with the velocity
of the body, and are a valid approximation of the fluid forces at high Reynolds numbers.
The torque is obtained by integrating the force resulting from the rotation over the surface area.
The :math:`i`-th component of the force and torque can be written as
.. math::
\begin{aligned}
\text{density force} : \quad &- {1 \over 2} \rho s_j s_k |v_i| v_i \\
\text{density torque} : \quad &- {1 \over 64} \rho s_i \left(s_j^4 + s_k^4 \right) |\omega_i| \omega_i \\
f_{D, i} = \quad &- 2 \rho r_j r_k |v_i| v_i \\
g_{D, i} = \quad &- {1 \over 2} \rho r_i \left(r_j^4 + r_k^4 \right) |\omega_i| \omega_i \\
\end{aligned}
This model implicitly assumes high Reynolds numbers, with lift-to-drag ratio equal to the tangent of the angle of
attack. One can also specify a non-zero :ref:`wind<option-wind>`, which is a 3D vector subtracted from the body linear
velocity in the fluid dynamics computation.
Each body also experiences a force and a torque proportional to the viscosity :math:`\beta` and opposite to its linear and
angular velocity. Note that viscosity can be used independent of density, to make the simulation more damped. We use the
formulas for a sphere at low Reynolds numbers, with diameter :math:`d` equal to the average of the equivalent inertia
box sizes. The resulting 3D force and torque in local body coordinates are
The viscous resistance terms depend on the fluid viscosity :math:`\beta`, scale linearly with the body velocity, and
approximate the fluid forces at low Reynolds numbers. Note that viscosity can be used independent of density to make
the simulation more damped. We use the formulas for the equivalent sphere with radius
:math:`r_{eq} = (r_x + r_y + r_z) / 3` at low Reynolds numbers. The resulting 3D force and torque in local
body coordinates are
.. math::
\begin{aligned}
\text{viscosity force} : \quad &- 3 \beta \pi d v \\
\text{viscosity torque} : \quad &- \beta \pi d^3 \omega \\
f_{V, i} = \quad &- 6 \beta \pi r_{eq} v_i \\
g_{V, i} = \quad &- 8 \beta \pi r_{eq}^3 \omega_i \\
\end{aligned}
One can also affect these forces by specifing a non-zero :ref:`wind<option-wind>`, which is a 3D vector subtracted
from the body linear velocity in the fluid dynamics computation.
.. _flEllipsoid:
Ellipsoid model
@@ -116,35 +135,36 @@ also disables the inertia-based model for the parent body. The
- 1.0
Elements of the model are a generalization of :cite:t:`andersen2005b` to 3 dimensions.
The force :math:`\mathbf{f}_{\text{fluid}\rightarrow \text{solid}}` and torque
:math:`\mathbf{g}_{\text{fluid} \rightarrow \text{solid}}` exerted by the fluid onto the solid are
The force :math:`\mathbf{f}_{\text{ellipsoid}}` and torque
:math:`\mathbf{g}_{\text{ellipsoid}}` exerted by the fluid onto the solid are
the sum of of the terms
.. math::
\begin{align*}
\mathbf{f}_{\text{fluid} \rightarrow \text{solid}} &= \mathbf{f}_A + \mathbf{f}_D + \mathbf{f}_M + \mathbf{f}_K \\
\mathbf{g}_{\text{fluid} \rightarrow \text{solid}} &= \mathbf{g}_A + \mathbf{g}_D
\mathbf{f}_{\text{ellipsoid}} &= \mathbf{f}_A + \mathbf{f}_D + \mathbf{f}_M + \mathbf{f}_K + \mathbf{f}_V \\
\mathbf{g}_{\text{ellipsoid}} &= \mathbf{g}_A + \mathbf{g}_D + \mathbf{g}_V
\end{align*}
Where subscripts :math:`A`, :math:`D`, :math:`M` and :math:`K`, denote Added mass, viscous Drag, Magnus lift and
Kutta lift, respectively. The :math:`D`, :math:`M` and :math:`K` terms are scaled by the respective
:math:`C_D`, :math:`C_M` and :math:`C_K` coefficients above, while the added mass term cannot be scaled.
Where subscripts :math:`A`, :math:`D`, :math:`M`, :math:`K` and :math:`V` denote Added mass, viscous Drag, Magnus lift,
Kutta lift and Viscous resistance, respectively. The :math:`D`, :math:`M` and :math:`K` terms are scaled by the respective
:math:`C_D`, :math:`C_M` and :math:`C_K` coefficients above, the viscous resistance scales with the fluid viscosity
:math:`\beta`, while the added mass term cannot be scaled.
Notation
~~~~~~~~
We describe the motion of the object in an inviscid, incompressible quiescent fluid of density :math:`\rho`. The
arbitrarily-shaped object is described in the model as the equivalent ellipsoid of semi-axes
:math:`\mathbf{d} = \{d_x, d_y, d_z\}`.
:math:`\mathbf{r} = \{r_x, r_y, r_z\}`.
The problem is described in a reference frame aligned with the sides of the ellipsoid and moving with it. The
body has velocity :math:`\mathbf{v} = \{v_x, v_y, v_z\}` and angular velocity
:math:`\boldsymbol{\omega} = \{\omega_x, \omega_y, \omega_z\}`. We will also use
.. math::
\begin{align*}
d_\text{max} &= \max(d_x, d_y, d_z) \\
d_\text{min} &= \min(d_x, d_y, d_z) \\
d_\text{mid} &= d_x + d_y + d_z - d_\text{max} - d_\text{min}
r_\text{max} &= \max(r_x, r_y, r_z) \\
r_\text{min} &= \min(r_x, r_y, r_z) \\
r_\text{mid} &= r_x + r_y + r_z - r_\text{max} - r_\text{min}
\end{align*}
The Reynolds number is the ratio between inertial and viscous forces within a flow and is defined as :math:`Re=u~l/\beta`, where
@@ -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<flProjection>` 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).
doc/conf.py
+6
View File
@@ -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']
doc/images/mjx/mujoco_vs_mjx_large_scene.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 26 KiB

doc/images/modeling/bunny1.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 488 KiB

doc/images/modeling/bunny2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

doc/images/modeling/coil.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 576 KiB

doc/images/modeling/flexelem.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 91 KiB

doc/index.rst
+5
View File
@@ -14,6 +14,11 @@
programming/index.rst
APIreference/index.rst
python
MJX <mjx>
unity
models
changelog
.. sidebar-links::
:github:
doc/mjx.rst
+312
View File
@@ -0,0 +1,312 @@
==========
MuJoCo XLA
==========
Starting with version 3.0.0, MuJoCo includes MuJoCo XLA (MJX) under the
`mjx <https://github.com/google-deepmind/mujoco/tree/main/mjx>`__ directory. MJX allows MuJoCo to run on compute
hardware supported by the `XLA <https://www.tensorflow.org/xla>`__ compiler via the
`JAX <https://github.com/google/jax#readme>`__ framework. MJX runs on a
`all platforms supported by JAX <https://jax.readthedocs.io/en/latest/installation.html#supported-platforms>`__: Nvidia
and AMD GPUs, Apple Silicon, and `Google Cloud TPUs <https://cloud.google.com/tpu>`__.
The MJX API is consistent with the main simulation functions in the MuJoCo API, although it is currently missing some
features. While the :ref:`API documentation <Mainsimulation>` is applicable to both libraries, we indicate features
unsupported by MJX in the :ref:`notes <MjxFeatureParity>` below.
MJX is distributed as a separate package called ``mujoco-mjx`` on `PyPI <https://pypi.org/project/mujoco-mjx>`__.
Although it depends on the main ``mujoco`` package for model compilation and visualization, it is a re-implementation of
MuJoCo that uses the same algorithms as the MuJoCo implementation. However, in order to properly leverage JAX, MJX
deliberately diverges from the MuJoCo API in a few places, see below.
MJX is a successor to the `generalized physics pipeline <https://github.com/google/brax/tree/main/brax/generalized>`__
in Google's `Brax <https://github.com/google/brax>`__ physics and reinforcement learning library. MJX was built
by core contributors to both MuJoCo and Brax, who will together continue to support both Brax (for its reinforcement
learning algorithms and included environments) and MJX (for its physics algorithms). A future version of Brax will
depend on the ``mujoco-mjx`` package, and Brax's existing
`generalized pipeline <https://github.com/google/brax/tree/main/brax/generalized>`__ will be deprecated. This change
will be largely transparent to users of Brax.
.. _MjxNotebook:
Tutorial notebook
=================
The following IPython notebook demonstrates the use of MJX along with reinforcement learning to train humanoid and
quadruped robots to locomote: |colab|.
.. |colab| image:: https://colab.research.google.com/assets/colab-badge.svg
:target: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb
.. _MjxInstallation:
Installation
============
The recommended way to install this package is via `PyPI <https://pypi.org/project/mujoco-mjx/>`__:
.. code-block:: shell
pip install mujoco-mjx
A copy of the MuJoCo library is provided as part of this package's depdendencies and does **not** need to be downloaded
or installed separately.
.. _MjxUsage:
Basic usage
===========
Once installed, the package can be imported via ``from mujoco import mjx``. Structs, functions, and enums are available
directly from the top-level ``mjx`` module.
.. _MjxStructs:
Structs
-------
Before running MJX functions on an accelerator device, structs must be copied onto the device via the ``mjx.device_put``
function. Placing an :ref:`mjModel` on device yields an ``mjx.Model``. Placing an :ref:`mjData` on device yields
an ``mjx.Data``:
.. code-block:: python
model = mujoco.MjModel.from_xml_string("...")
data = mujoco.MjData(model)
mjx_model = mjx.device_put(model)
mjx_data = mjx.device_put(data)
These MJX variants mirror their MuJoCo counterparts but have three key differences:
#. Fields in ``mjx.Model`` and ``mjx.Data`` are JAX arrays copied onto device, instead of numpy arrays.
#. Some fields are missing from ``mjx.Model`` and ``mjx.Data`` for features that are
:ref:`unsupported <mjxFeatureParity>` in MJX.
#. Arrays in ``mjx.Model`` and ``mjx.Data`` support adding batch dimensions. Batch dimensions are a natural way to
express domain randomization (in the case of ``mjx.Model``) or high-throughput simulation for reinforcement learning
(in the case of ``mjx.Data``).
Neither ``mjx.Model`` nor ``mjx.Data`` are meant to be constructed manually. An ``mjx.Data`` may be created by calling
``mjx.make_data``, which mirrors the :ref:`mj_makeData` function in MuJoCo:
.. code-block:: python
model = mujoco.MjModel.from_xml_string("...")
mjx_model = mjx.device_put(model)
mjx_data = mjx.make_data(model)
Using ``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 <https://peps.python.org/pep-0008/>`__-compliant names. Most of the :ref:`main simulation <Mainsimulation>` and
some of the :ref:`sub-components <Subcomponents>` for forward simulation are available from the top-level ``mjx`` module.
MJX functions are not `JIT compiled <https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html>`__ by default -- we
leave it to the user to JIT MJX functions, or JIT their own functions that reference MJX functions. See the
:ref:`minimal example <MjxExample>` below.
.. _MjxEnums:
Enums and constants
-------------------
MJX enums are available as ``mjx.EnumType.ENUM_VALUE``, for example ``mjx.JointType.FREE``. Enums for unsupported MJX
features are omitted from the MJX enum declaration. MJX declares no constants but references MuJoCo constants directly.
.. _MjxExample:
Minimal example
---------------
.. code-block:: python
# Throw a ball at 100 different velocities.
import jax
import mujoco
from mujoco import mjx
XML=r"""
<mujoco>
<worldbody>
<body>
<freejoint/>
<geom size=".15" mass="1" type="sphere"/>
</body>
</worldbody>
</mujoco>
"""
model = mujoco.MjModel.from_xml_string(XML)
mjx_model = mjx.device_put(model)
@jax.vmap
def batched_step(vel):
mjx_data = mjx.make_data(mjx_model)
qvel = mjx_data.qvel.at[0].set(vel)
mjx_data = mjx_data.replace(qvel=qvel)
pos = mjx.step(mjx_model, mjx_data).qpos[0]
return pos
vel = jax.numpy.arange(0.0, 1.0, 0.01)
pos = jax.jit(batched_step)(vel)
print(pos)
.. _MjxFeatureParity:
Feature Parity
==============
MJX supports most of the main simulation features of MuJoCo, with a few exceptions. MJX will raise an exception if
asked to copy to device an :ref:`mjModel` with field values referencing unsupported features.
The following features are **fully supported** in MJX:
.. list-table::
:width: 90%
:align: left
:widths: 1 5
:header-rows: 1
* - Category
- Feature
* - Dynamics
- :ref:`Forward <mj_forward>`
* - :ref:`Joint <mjtJoint>`
- ``FREE``, ``BALL``, ``SLIDE``, ``HINGE``
* - :ref:`Transmission <mjtTrn>`
- ``TRN_JOINT``
* - :ref:`Actuation <geactuation>`
- ``DYN_NONE``, ``DYN_INTEGRATOR``, ``DYN_FILTER``, ``GAIN_FIXED``, ``GAIN_AFFINE``, ``BIAS_NONE``,
``BIAS_AFFINE``
* - :ref:`Geom <mjtGeom>`
- ``PLANE``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH``
* - :ref:`Constraint <mjtConstraint>`
- ``EQUALITY``, ``FRICTION_DOF``, ``LIMIT_JOINT``, ``CONTACT_PYRAMIDAL``
* - :ref:`Integrator <mjtIntegrator>`
- ``EULER``, ``RK4``
* - :ref:`Cone <mjtCone>`
- ``PYRAMIDAL``
* - :ref:`Condim <coContact>`
- 3
* - :ref:`Solver <mjtSolver>`
- ``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 <mj_inverse>`
* - :ref:`Transmission <mjtTrn>`
- ``TRN_TENDON``
* - :ref:`Geom <mjtGeom>`
- ``HFIELD``, ``ELLIPSOID``, ``CYLINDER``, ``SDF``
* - :ref:`Integrator <mjtIntegrator>`
- ``IMPLICIT``, ``IMPLICITFAST``
* - :ref:`Cone <mjtCone>`
- ``ELLIPTIC``
* - :ref:`Condim <coContact>`
- 1, 4, 6
* - :ref:`Solver <mjtSolver>`
- ``NEWTON``
* - Fluid Model
- :ref:`flEllipsoid`
* - :ref:`Tendons <tendon>`
- :ref:`Spatial <tendon-spatial>`, :ref:`Fixed <tendon-fixed>`
The following features are **unsupported**:
.. list-table::
:width: 90%
:align: left
:widths: 1 5
:header-rows: 1
* - Category
- Feature
* - :ref:`Transmission <mjtTrn>`
- ``TRN_JOINTINPARENT``, ``TRN_SLIDERCRANK``, ``TRN_SITE``, ``TRN_BODY``, ``MUSCLE``
* - :ref:`Solver <mjtSolver>`
- ``PGS``
* - :ref:`Callbacks <glphysics>`
- ``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 <https://en.wikipedia.org/wiki/Single_instruction,_multiple_data>`__. This specialization is useful
for machine learning workloads such as `reinforcement learning <https://en.wikipedia.org/wiki/Reinforcement_learning>`__
that require massive data throughput.
There are certain workflows that MJX is ill-suited for:
Single scene simulation
Simulating a single scene (1 instance of :ref:`mjData`), MJX can be **10x** slower than MuJoCo, which has been
carefully optimized for CPU. MJX works best when simulating thousands or tens of thousands of scenes in parallel.
Large, complex scenes with many contacts
Accelerators exhibit poor performance for
`branching code <https://aschrein.github.io/jekyll/update/2019/06/13/whatsup-with-my-branches-on-gpu.html#tldr>`__.
Branching is used in broad-phase collision detection, when identifying potential collisions between large numbers of
bodies in a scene. MJX ships with a simple branchless broad-phase algorithm (see performance tuning) but it is not as
powerful as the one in MuJoCo.
To see how this affects simulation, let us consider a physics scene with increasing numbers of 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 <https://ubm-twvideo01.s3.amazonaws.com/o1/vault/gdc2013/slides/822403Gregorius_Dirk_TheSeparatingAxisTest.pdf>`__.
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 <https://github.com/google-deepmind/mujoco/tree/main/mjx/mujoco/mjx/benchmark/model/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 <https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid_v4.py>`__ task resets when
the humanoid starts to fall, so full contact with the floor is not needed.
:ref:`option-flag` element
Disabling ``eulerdamp`` can help performance and is often not needed for stability.
doc/modeling.rst
+141 -45
View File
@@ -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
<CDeformable>` section. However, collisions are limited to the particle themselves and not to the whole boundary of the
skin that encloses them. This makes contacts very fast but does not guarantee that all penetrations can be avoided. For
a more complete treatment, see again the :ref:`deformable <CDeformable>` section, which outlines how to use
:ref:`flexcomp<body-flexcomp>` to create such an object. It is easy to port models create with composite particles to
flex, see the folder `elasticity/ <https://github.com/google-deepmind/mujoco/tree/main/model/plugin/elasticity>`__ for
several examples.
**1D grid**.
|image6| |image7|
@@ -1111,59 +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
<body name="B10" pos="0 0 1">
<freejoint/>
<composite type="rope" count="21 1 1" spacing="0.04" offset="0 0 2">
<joint kind="main" damping="0.005"/>
<geom type="capsule" size=".01 .015" rgba=".8 .2 .1 1"/>
</composite>
</body>
<extension>
<plugin plugin="mujoco.elasticity.cable"/>
</extension>
The remaining composite object types create kinematic trees of element bodies, and the parent body becomes the root of
the tree. This is why :el:`composite` appears inside a moving body, and not inside the world body as in particle and
grid objects. If it appeared inside the world body, the root of the composite object would not move. Unlike grids and
particles, the orientation of the element bodies here can change. The kinematic tree is constructed using (mostly)
hinge joints. In the case of rope and loop objects illustrated here, the tree is a chain. Note the naming of the
parent body. This name must correspond to one of the automatically-generated names of the element bodies. This
mechanism is used to specify where the composite object should attach to the parent. Compared to 1D grids, the rope
and loop are less jittery and can use capsule and ellipsoid geoms in addition to spheres (thus filling the gaps for
collision detection). However this comes at a price. Because we have long kinematic chains, the resulting differential
equations become stiff and can no longer be integrated at large timesteps. The examples we provide illustrate
comfortable timesteps where the models are stable. The rope can be easily tied into a knot using mouse perturbations,
as shown in the left plot. Using a larger number of smaller elements makes knots and other manipulations even easier.
The loop is similar to a rope but the first and last element bodies are connected with an equality constraint.
<worldbody>
<composite prefix="actuated" type="cable" curve="cos(s) sin(s) s" count="41 1 1"
size="0.25 .1 4" offset="0.25 0 .05" initial="none">
<plugin plugin="mujoco.elasticity.cable">
<!--Units are in Pa (SI)-->
<config key="twist" value="5e8"/>
<config key="bend" value="15e8"/>
<config key="vmax" value="0"/>
</plugin>
<joint kind="main" damping="0.15" armature="0.01"/>
<geom type="capsule" size=".005" rgba=".8 .2 .1 1"/>
</composite>
</worldbody>
The cable simulates an inextensible elastic 1D object having twist and bending stiffness. It is discretized using a
sequence of capsules or boxes. Its stiffness and inertia properties are computed directly from the given parameters and
the shape of the cross section, which allows for anisotropic behaviors, which can be found in e.g. belts or computer
cables. It is a single kinematic tree, so it is exactly inextensible without the use of additional constraints, enabling
the use of large time steps. The elastic model is geometrically exact and based on computing the Bishop or twist-free
frame of the centerline, i.e., the line passing through the center of the cross section. The orientations of the geoms
are expressed with respect to this frame and then decomposed into twist and bending components, hence different
stiffnesses can be set independently. Moreover, it is possible to specify if the stress-free configuration is flat or
curve, such as in the case of coil springs. The cable requires using a first-party :ref:`engine plugin<exPlugin>`, which
may be integrated directly into the engine in the future.
**Rope and loop**.
The rope and loop are deprecated. It is recommended to use the cable for simulating inextensible elastic rods that are
bent and twisted and 1D flex :ref:`deformable objects <CDeformable>` for extensible strings in a tensile loading
scenario (e.g. a stretched rubber band).
**Cloth**.
|image12| |image13|
.. code-block:: xml
<body name="B3_5" pos="0 0 1">
<freejoint/>
<composite type="cloth" count="9 9 1" spacing="0.05" flatinertia="0.01">
<joint kind="main" damping="0.001"/>
<skin material="matcarpet" texcoord="true" inflate="0.005" subgrid="2"/>
<geom type="capsule" size="0.015 0.01" rgba=".8 .2 .1 1"/>
</composite>
</body>
The cloth type is an alternative to a 2D grid, and has somewhat different properties. Similar to rope vs. 1D grid, the
cloth is less jittery than a 2D grid and can also fill collision holes better. This is done by using capsules or
ellipsoids, and arranging them in the pattern shown on the right. The geom capsules are shown in red, the kinematic
tree in thick blue, the equality-constrained tendons holding the cloth together in thin gray, and the joints in cyan.
The element body corresponding to the parent body has a floating joint rendered as a cube, while the rest of the tree
is constructed using pairs of hinge joints that form universal joints. Note the naming of the parent body: similar to
rope, it must coincide with one of the automatically-generated element body names in the composite object. Explicit
pinning is not possible. However if the parent is a static body, the cloth is essentially pinned but only at one
point. Similar to rope, the cloth object involves long kinematic chains that require relatively small timesteps and
some damping for stable integration. The parameters can be found in the XML model files in the software distribution.
The cloth is deprecated. It is recommended to use 2D flex :ref:`deformable objects <CDeformable>` for simulating thin
elastic structures.
**Box**.
@@ -1222,6 +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 <CComposite>` described earlier were intended to emulate soft bodies in what is effectively
a rigid-body simulator. This was possible because MuJoCo constraints are soft, but nevertheless it was limited in
functionality and modeling power. In MuJoCo 3.0 we have introduced true deformable objects involving new model elements.
The :ref:`skin<deformable-skin>` described earlier was actually one such element, but it is merely used for
visualization. We now have a related element :ref:`flex<deformable-flex>` which generates contact forces, constraint
forces and passive forces as needed to model a wide range of deformable entities. Both skins and flexes are now defined
within a new grouping element in the XML called :ref:`deformable<deformable>`. A flex is a low-level element that
specifies everything needed at runtime, but is difficult to design at modeling time. To aid with modeling, we have
further introduced the element :ref:`flexcomp<body-flexcomp>` which automates the creation of the low-level flex,
similar to how :ref:`composite<body-composite>` automates the creation of (collections of) MuJoCo objects needed to
emulate a soft body. Flexes may eventually supersede composites, but for now both are useful for somewhat different
purposes.
A flex is a collection of MuJoCo bodies that are connected with massless stretchable elements. These elements can be
capsules (1D flex), triangles (2D flex), or tetrahedra (3D flex). In all cases we allow a radius, which makes the
elements smooth and also volumetric in 1D and 2D. The primitive elements are illustrated below:
|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
<https://en.wikipedia.org/wiki/Poisson%27s_ratio>`__ of the material. For more details, see the `Saint Venant-Kirchhoff
<https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__ hyperelastic model. This
functionality is currently based on first-party :ref:`engine plugins<exPlugin>` as of MuJoCo 3.0 but may be integrated
into the engine in future releases.
**Creation and visualization**.
.. code-block:: xml
<extension>
<plugin plugin="mujoco.elasticity.solid"/>
</extension>
<worldbody>
<flexcomp type="grid" count="24 4 4" spacing=".1 .1 .1" pos=".1 0 1.5"
radius=".0" rgba="0 .7 .7 1" name="softbody" dim="3" mass="7">
<contact condim="3" solref="0.01 1" solimp=".95 .99 .0001" selfcollide="none"/>
<edge damping="1"/>
<plugin plugin="mujoco.elasticity.solid">
<config key="poisson" value="0.2"/>
<!--Units are in Pa (SI)-->
<config key="young" value="5e4"/>
</plugin>
</flexcomp>
</worldbody>
Using the :ref:`flexcomp<body-flexcomp>` element, we can create flexes from meshes, including tetrahedral meshes, and
automatically generate all the bodies/vertices and connect them with suitable elements. We can also create grids and
other topologies automatically. This machinery makes it easy to create very large flexes, involving thousands or even
tens of thousands of bodies, elements and edges. Obviously such simulations will not be fast. Even for medium-sized
flexes, pruning of collision pairs and essential. This is why we have developed elaborate methods for pruning
self-collisions; see XML reference.
In case of 3D flexes made of tetrahedra, it may be useful to examine how the flex is "triangulated" internally. We have
a special visualization mode that peels off the outer layers. Below is an example with the Stanford Bunny. Note how it
has smaller tetrahedra on the outside and larger ones on the inside. This mesh design makes sense, because we want the
collision surface to be accurate, but on the inside we just need soft material properties - which require less spatial
resolution.
|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
doc/overview.rst
+18 -3
View File
@@ -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
^^^^^^^^^^^^^^^
doc/programming/extension.rst
+1
View File
@@ -256,6 +256,7 @@ Currently, there are three directories of first-party plugins:
bending strains. The 3D solid is a
`Saint Venant-Kirchhoff <https://en.wikipedia.org/wiki/Hyperelastic_material#Saint_Venant%E2%80%93Kirchhoff_model>`__
model discretized with piecewise linear finite elements, which is suitable for large deformations with small strains.
See also :ref:`composite <CComposite>` and :ref:`deformable <CDeformable>` objects.
* **sensor:** The plugins in the `sensor/ <https://github.com/google-deepmind/mujoco/tree/main/plugin/sensor>`__
directory implement custom sensors. Currently the sole sensor plugin is the touch grid sensor, see the
`README <https://github.com/google-deepmind/mujoco/blob/main/plugin/sensor/README.md>`__ for details.
doc/unity.rst
+2 -2
View File
@@ -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
_______
include/mujoco/mujoco.h
+1 -1
View File
@@ -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 <stdlib.h>
mjx/MANIFEST.in
+2
View File
@@ -0,0 +1,2 @@
recursive-include mujoco/mjx/test_data *
recursive-include mujoco/mjx/benchmark *.obj *.stl *.xml
mjx/README.md
+54
View File
@@ -0,0 +1,54 @@
# MuJoCo XLA (MJX)
[![PyPI Python Version][pypi-versions-badge]][pypi]
[![PyPI version][pypi-badge]][pypi]
[pypi-versions-badge]: https://img.shields.io/pypi/pyversions/mujoco-mjx
[pypi-badge]: https://badge.fury.io/py/mujoco-mjx.svg
[pypi]: https://pypi.org/project/mujoco-mjx/
This package is a re-implementation of the
[MuJoCo physics engine](https://github.com/google-deepmind/mujoco) in
[JAX](https://github.com/google/jax). This library is developed and maintained
by Google DeepMind, and is kept up-to-date with the latest developments in
MuJoCo itself.
The `mujoco-mjx` package is API-compatible with MuJoCo, but is missing some
features found in MuJoCo. See our
[documentation](https://mujoco.readthedocs.io/en/stable/mjx.html) for more
details concerning feature parity.
## Installation
The recommended way to install this package is via [PyPI](https://pypi.org/project/mujoco-mjx/):
```sh
pip install mujoco-mjx
```
## Usage
Once installed, the package can be imported via `from mujoco import mjx`. Please
consult our [documentation](https://mujoco.readthedocs.io/en/stable/mjx.html)
for further detail on the package's API.
We recommend going through the tutorial notebook which introduces the MJX API
and trains a reinforcement learning policy in a few minutes: [![Open In
Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb)
## Versioning
The `major.minor.micro` portion of the version number matches the version of
MuJoCo that this library provides. Optionally, if we release updates to MJX that
target the same version of MuJoCo, a `.postN` suffix is added, for example
`3.0.1.post2` represents the second update to MJX for MuJoCo 3.0.1.
## License and Disclaimer
Copyright 2023 DeepMind Technologies Limited
MuJoCo and its libraries are licensed under the Apache License,
Version 2.0. You may obtain a copy of the License at
https://www.apache.org/licenses/LICENSE-2.0.
This is not an officially supported Google product.
mjx/cuda_requirements.txt
+17
View File
@@ -0,0 +1,17 @@
-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
jax[cuda12_local]==0.4.13; python_version == '3.8' \
--hash=sha256:03bfe6749dfe647f16f15f6616638adae6c4a7ca7167c75c21961ecfd3a3baaa
jax[cuda12_local]==0.4.18; python_version >= '3.9' \
--hash=sha256:c3ab72ea2f1c5d8ccf2561e79f6562fb2964629f3e55b3ac1c11c48b64c20336
jaxlib==0.4.13+cuda12.cudnn89; python_version == '3.8' \
--hash=sha256:31372f41dc28ecb11a5cc5573ae632e8ee9cfb63788edc8ad8a77e8d3279f569
jaxlib==0.4.18+cuda12.cudnn89; python_version >= '3.9' \
--hash=sha256:14f74ff081882ea091c121e355051b35932e39cb7ff7242b88a87f3690f3ca90 \
--hash=sha256:7c87dc2d68257b02e83c04be88a3c447373ee7077d65f43545bcbda5bfe2231d \
--hash=sha256:4d16e9c7592e1aaca0b3d28d2c8beba415a2721bb7001f2947728247951a250d \
--hash=sha256:759c08c69f4a5b1e6b39c3e4eff908a04ce3b2b483bb594ed624407c7d12d110 \
--hash=sha256:a7a04dbe1851cd50d07691282116aee49a2f0be7838e55b76d7ada86db06be62 \
--hash=sha256:2bf842db3d58c8c6c52fbc8ed3fabefd7b91a21746cd59d3eaf3522eea229b53 \
--hash=sha256:35d265ef9bb3835a14580cbaa9402060f117e46056f80e0996405fff3964667a \
--hash=sha256:0e4352f24d629e912965e6435e140c1b06086243a098651f2d01b75f3738b51c
mjx/mujoco/mjx/__init__.py
+34
View File
@@ -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 *
mjx/mujoco/mjx/_src/__init__.py
+14
View File
@@ -0,0 +1,14 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
mjx/mujoco/mjx/_src/collision_base.py
+64
View File
@@ -0,0 +1,64 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Collision base."""
import dataclasses
from typing import Dict, List, Optional, Tuple
import jax
# pylint: disable=g-importing-member
from mujoco.mjx._src.dataclasses import PyTreeNode
from mujoco.mjx._src.types import GeomType
# pylint: enable=g-importing-member
Contact = Tuple[jax.Array, jax.Array, jax.Array]
@dataclasses.dataclass
class Candidate:
geom1: int
geom2: int
ipair: int
geomp: int # priority geom
dim: int
CandidateSet = Dict[
Tuple[GeomType, GeomType, Tuple[int, ...], Tuple[int, ...]],
List[Candidate],
]
class GeomInfo(PyTreeNode):
"""Collision info for a geom."""
pos: jax.Array
mat: jax.Array
size: jax.Array
face: Optional[jax.Array] = None
vert: Optional[jax.Array] = None
edge: Optional[jax.Array] = None
facenorm: Optional[jax.Array] = None
class SolverParams(PyTreeNode):
"""Contact solver params."""
friction: jax.Array
solref: jax.Array
solreffriction: jax.Array
solimp: jax.Array
margin: jax.Array
gap: jax.Array
mjx/mujoco/mjx/_src/collision_convex.py
+724
View File
@@ -0,0 +1,724 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Convex collisions."""
from typing import Tuple
import jax
from jax import numpy as jp
from mujoco.mjx._src import math
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Contact
from mujoco.mjx._src.collision_base import GeomInfo
# pylint: enable=g-importing-member
def _closest_segment_point_plane(
a: jax.Array, b: jax.Array, p0: jax.Array, plane_normal: jax.Array
) -> jax.Array:
"""Gets the closest point between a line segment and a plane.
Args:
a: first line segment point
b: second line segment point
p0: point on plane
plane_normal: plane normal
Returns:
closest point between the line segment and the plane
"""
# Parametrize a line segment as S(t) = a + t * (b - a), plug it into the plane
# equation dot(n, S(t)) - d = 0, then solve for t to get the line-plane
# intersection. We then clip t to be in [0, 1] to be on the line segment.
n = plane_normal
d = jp.sum(p0 * n) # shortest distance from origin to plane
denom = jp.sum(n * (b - a))
t = (d - jp.sum(n * a)) / (denom + 1e-6 * (denom == 0.0))
t = jp.clip(t, 0, 1)
segment_point = a + t * (b - a)
return segment_point
def _closest_triangle_point(
p0: jax.Array, p1: jax.Array, p2: jax.Array, pt: jax.Array
) -> jax.Array:
"""Gets the closest point between a triangle and a point in space.
Args:
p0: triangle point
p1: triangle point
p2: triangle point
pt: point to test
Returns:
closest point on the triangle w.r.t point pt
"""
# Parametrize the triangle s.t. a point inside the triangle is
# Q = p0 + u * e0 + v * e1, when 0 <= u <= 1, 0 <= v <= 1, and
# 0 <= u + v <= 1. Let e0 = (p1 - p0) and e1 = (p2 - p0).
# We analytically minimize the distance between the point pt and Q.
e0 = p1 - p0
e1 = p2 - p0
a = e0.dot(e0)
b = e0.dot(e1)
c = e1.dot(e1)
d = pt - p0
# The determinant is 0 only if the angle between e1 and e0 is 0
# (i.e. the triangle has overlapping lines).
det = a * c - b * b
u = (c * e0.dot(d) - b * e1.dot(d)) / det
v = (-b * e0.dot(d) + a * e1.dot(d)) / det
inside = (0 <= u) & (u <= 1) & (0 <= v) & (v <= 1) & (u + v <= 1)
closest_p = p0 + u * e0 + v * e1
d0 = (closest_p - pt).dot(closest_p - pt)
# If the closest point is outside the triangle, it must be on an edge, so we
# check each triangle edge for a closest point to the point pt.
closest_p1, d1 = math.closest_segment_point_and_dist(p0, p1, pt)
closest_p = jp.where((d0 < d1) & inside, closest_p, closest_p1)
min_d = jp.where((d0 < d1) & inside, d0, d1)
closest_p2, d2 = math.closest_segment_point_and_dist(p1, p2, pt)
closest_p = jp.where(d2 < min_d, closest_p2, closest_p)
min_d = jp.minimum(min_d, d2)
closest_p3, d3 = math.closest_segment_point_and_dist(p2, p0, pt)
closest_p = jp.where(d3 < min_d, closest_p3, closest_p)
return closest_p
def _closest_segment_triangle_points(
a: jax.Array,
b: jax.Array,
p0: jax.Array,
p1: jax.Array,
p2: jax.Array,
triangle_normal: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
"""Gets the closest points between a line segment and triangle.
Args:
a: first line segment point
b: second line segment point
p0: triangle point
p1: triangle point
p2: triangle point
triangle_normal: normal of triangle
Returns:
closest point on the triangle w.r.t the line segment
"""
# The closest triangle point is either on the edge or within the triangle.
# First check triangle edges for the closest point.
# TODO(robotics-simulation): consider vmapping over closest point functions
seg_pt1, tri_pt1 = math.closest_segment_to_segment_points(a, b, p0, p1)
d1 = (seg_pt1 - tri_pt1).dot(seg_pt1 - tri_pt1)
seg_pt2, tri_pt2 = math.closest_segment_to_segment_points(a, b, p1, p2)
d2 = (seg_pt2 - tri_pt2).dot(seg_pt2 - tri_pt2)
seg_pt3, tri_pt3 = math.closest_segment_to_segment_points(a, b, p0, p2)
d3 = (seg_pt3 - tri_pt3).dot(seg_pt3 - tri_pt3)
# Next, handle the case where the closest triangle point is inside the
# triangle. Either the line segment intersects the triangle or a segment
# endpoint is closest to a point inside the triangle.
seg_pt4 = _closest_segment_point_plane(a, b, p0, triangle_normal)
tri_pt4 = _closest_triangle_point(p0, p1, p2, seg_pt4)
d4 = (seg_pt4 - tri_pt4).dot(seg_pt4 - tri_pt4)
# Get the point with minimum distance from the line segment point to the
# triangle point.
distance = jp.array([[d1, d2, d3, d4]])
min_dist = jp.amin(distance)
mask = (distance == min_dist).T
seg_pt = jp.array([seg_pt1, seg_pt2, seg_pt3, seg_pt4]) * mask
tri_pt = jp.array([tri_pt1, tri_pt2, tri_pt3, tri_pt4]) * mask
seg_pt = jp.sum(seg_pt, axis=0) / jp.sum(mask)
tri_pt = jp.sum(tri_pt, axis=0) / jp.sum(mask)
return seg_pt, tri_pt
def _manifold_points(
poly: jax.Array, poly_mask: jax.Array, poly_norm: jax.Array
) -> jax.Array:
"""Chooses four points on the polygon with approximately maximal area."""
dist_mask = jp.where(poly_mask, 0.0, -1e6)
a_idx = jp.argmax(dist_mask)
a = poly[a_idx]
# choose point b furthest from a
b_idx = (((a - poly) ** 2).sum(axis=1) + dist_mask).argmax()
b = poly[b_idx]
# choose point c furthest along the axis orthogonal to (a-b)
ab = jp.cross(poly_norm, a - b)
ap = a - poly
c_idx = (jp.abs(ap.dot(ab)) + dist_mask).argmax()
c = poly[c_idx]
# choose point d furthest from the other two triangle edges
ac = jp.cross(poly_norm, a - c)
bc = jp.cross(poly_norm, b - c)
bp = b - poly
dist_bp = jp.abs(bp.dot(bc)) + dist_mask
dist_ap = jp.abs(ap.dot(ac)) + dist_mask
d_idx = jp.concatenate([dist_bp, dist_ap]).argmax() % poly.shape[0]
return jp.array([a_idx, b_idx, c_idx, d_idx])
def _project_pt_onto_plane(
pt: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array
) -> jax.Array:
"""Projects a point onto a plane along the plane normal."""
dist = (pt - plane_pt).dot(plane_normal)
return pt - dist * plane_normal
def _project_poly_onto_plane(
poly: jax.Array, plane_pt: jax.Array, plane_normal: jax.Array
) -> jax.Array:
"""Projects a polygon onto a plane using the plane normal."""
return jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
poly, plane_pt, math.normalize(plane_normal)
)
def _project_poly_onto_poly_plane(
poly1: jax.Array, norm1: jax.Array, poly2: jax.Array, norm2: jax.Array
) -> jax.Array:
"""Projects poly1 onto the poly2 plane along poly1's normal."""
d = poly2[0].dot(norm2)
denom = norm1.dot(norm2)
t = (d - poly1.dot(norm2)) / (denom + 1e-6 * (denom == 0.0))
new_poly = poly1 + t.reshape(-1, 1) * norm1
return new_poly
def _point_in_front_of_plane(
plane_pt: jax.Array, plane_normal: jax.Array, pt: jax.Array
) -> jax.Array:
"""Checks if a point is strictly in front of a plane."""
return (pt - plane_pt).dot(plane_normal) > 1e-6
def _clip_edge_to_planes(
edge_p0: jax.Array,
edge_p1: jax.Array,
plane_pts: jax.Array,
plane_normals: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
"""Clips an edge against side planes.
We return two clipped points, and a mask to include the new edge or not.
Args:
edge_p0: the first point on the edge
edge_p1: the second point on the edge
plane_pts: side plane points
plane_normals: side plane normals
Returns:
new_ps: new edge points that are clipped against side planes
mask: a boolean mask, True if an edge point is a valid clipped point and
False otherwise
"""
p0, p1 = edge_p0, edge_p1
p0_in_front = jax.vmap(jp.dot)(p0 - plane_pts, plane_normals) > 1e-6
p1_in_front = jax.vmap(jp.dot)(p1 - plane_pts, plane_normals) > 1e-6
# Get candidate clipped points along line segment (p0, p1) by clipping against
# all clipping planes.
candidate_clipped_ps = jax.vmap(
_closest_segment_point_plane, in_axes=[None, None, 0, 0]
)(p0, p1, plane_pts, plane_normals)
def clip_edge_point(p0, p1, p0_in_front, clipped_ps):
@jax.vmap
def choose_edge_point(in_front, clipped_p):
return jp.where(in_front, clipped_p, p0)
# Pick the clipped point if p0 is in front of the clipping plane. Otherwise
# keep p0 as the edge point.
new_edge_ps = choose_edge_point(p0_in_front, clipped_ps)
# Pick the clipped point that is most along the edge direction.
# This degenerates to picking the original point p0 if p0 is *not* in front
# of any clipping planes.
dists = jp.dot(new_edge_ps - p0, p1 - p0)
new_edge_p = new_edge_ps[jp.argmax(dists)]
return new_edge_p
# Clip each edge point.
new_p0 = clip_edge_point(p0, p1, p0_in_front, candidate_clipped_ps)
new_p1 = clip_edge_point(p1, p0, p1_in_front, candidate_clipped_ps)
clipped_pts = jp.array([new_p0, new_p1])
# Keep the original points if both points are in front of any of the clipping
# planes, rather than creating a new clipped edge. If the entire subject edge
# is in front of any clipping plane, we need to grab an edge from the clipping
# polygon instead.
both_in_front = p0_in_front & p1_in_front
mask = ~jp.any(both_in_front)
new_ps = jp.where(mask, clipped_pts, jp.array([p0, p1]))
# Mask out crossing clipped edge points.
mask = jp.where((p0 - p1).dot(new_ps[0] - new_ps[1]) < 0, False, mask)
return new_ps, jp.array([mask, mask])
def _clip(
clipping_poly: jax.Array,
subject_poly: jax.Array,
clipping_normal: jax.Array,
subject_normal: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
"""Clips a subject polygon against a clipping polygon.
A parallelized clipping algorithm for convex polygons. The result is a set of
vertices on the clipped subject polygon in the subject polygon plane.
Args:
clipping_poly: the polygon that we use to clip the subject polygon against
subject_poly: the polygon that gets clipped
clipping_normal: normal of the clipping polygon
subject_normal: normal of the subject polygon
Returns:
clipped_pts: points on the clipped polygon
mask: True if a point is in the clipping polygon, False otherwise
"""
# Get clipping edge points, edge planes, and edge normals.
clipping_p0 = jp.roll(clipping_poly, 1, axis=0)
clipping_plane_pts = clipping_p0
clipping_p1 = clipping_poly
clipping_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])(
clipping_p1 - clipping_p0,
clipping_normal,
)
# Get subject edge points, edge planes, and edge normals.
subject_edge_p0 = jp.roll(subject_poly, 1, axis=0)
subject_plane_pts = subject_edge_p0
subject_edge_p1 = subject_poly
subject_plane_normals = jax.vmap(jp.cross, in_axes=[0, None])(
subject_edge_p1 - subject_edge_p0,
subject_normal,
)
# Clip all edges of the subject poly against clipping side planes.
clipped_edges0, masks0 = jax.vmap(
_clip_edge_to_planes, in_axes=[0, 0, None, None]
)(
subject_edge_p0,
subject_edge_p1,
clipping_plane_pts,
clipping_plane_normals,
)
# Project the clipping poly onto the subject plane.
clipping_p0_s = _project_poly_onto_poly_plane(
clipping_p0, clipping_normal, subject_poly, subject_normal
)
clipping_p1_s = _project_poly_onto_poly_plane(
clipping_p1, clipping_normal, subject_poly, subject_normal
)
# Clip all edges of the clipping poly against subject planes.
clipped_edges1, masks1 = jax.vmap(
_clip_edge_to_planes, in_axes=[0, 0, None, None]
)(clipping_p0_s, clipping_p1_s, subject_plane_pts, subject_plane_normals)
# Merge the points and reshape.
clipped_edges = jp.concatenate([clipped_edges0, clipped_edges1])
masks = jp.concatenate([masks0, masks1])
clipped_points = clipped_edges.reshape((-1, 3))
mask = masks.reshape(-1)
return clipped_points, mask
def _create_contact_manifold(
clipping_poly: jax.Array,
subject_poly: jax.Array,
clipping_norm: jax.Array,
subject_norm: jax.Array,
sep_axis: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
"""Creates a contact manifold between two convex polygons.
The polygon faces are expected to have a counter clockwise winding order so
that clipping plane normals point away from the polygon center.
Args:
clipping_poly: the reference polygon to clip the contact against.
subject_poly: the subject polygon to clip contacts onto.
clipping_norm: the clipping polygon normal.
subject_norm: the subject polygon normal.
sep_axis: the separating axis
Returns:
tuple of dist, pos, and normal
"""
# Clip the subject (incident) face onto the clipping (reference) face.
# The incident points are clipped points on the subject polygon.
poly_incident, mask = _clip(
clipping_poly, subject_poly, clipping_norm, subject_norm
)
# The reference points are clipped points on the clipping polygon.
poly_ref = _project_poly_onto_plane(
poly_incident, clipping_poly[0], clipping_norm
)
behind_clipping_plane = _point_in_front_of_plane(
clipping_poly[0], -clipping_norm, poly_incident
)
mask = mask & behind_clipping_plane
# Choose four contact points.
best = _manifold_points(poly_ref, mask, clipping_norm)
contact_pts = jp.take(poly_ref, best, axis=0)
mask_pts = jp.take(mask, best, axis=0)
penetration_dir = jp.take(poly_incident, best, axis=0) - contact_pts
penetration = penetration_dir.dot(-clipping_norm)
dist = jp.where(mask_pts, -penetration, jp.ones_like(penetration))
pos = contact_pts
normal = -jp.stack([sep_axis] * 4, 0)
return dist, pos, normal
def _sat_hull_hull(
faces_a: jax.Array,
faces_b: jax.Array,
vertices_a: jax.Array,
vertices_b: jax.Array,
normals_a: jax.Array,
normals_b: jax.Array,
unique_edges_a: jax.Array,
unique_edges_b: jax.Array,
) -> Tuple[jax.Array, jax.Array, jax.Array]:
"""Runs the Separating Axis Test for a pair of hulls.
Given two convex hulls, the Separating Axis Test finds a separating axis
between all edge pairs and face pairs. Edge pairs create a single contact
point and face pairs create a contact manifold (up to four contact points).
We return both the edge and face contacts. Valid contacts can be checked with
dist < 0. Resulting edge contacts should be preferred over face contacts.
Args:
faces_a: An ndarray of hull A's polygon faces.
faces_b: An ndarray of hull B's polygon faces.
vertices_a: Vertices for hull A.
vertices_b: Vertices for hull B.
normals_a: Normal vectors for hull A's polygon faces.
normals_b: Normal vectors for hull B's polygon faces.
unique_edges_a: Unique edges for hull A.
unique_edges_b: Unique edges for hull B.
Returns:
tuple of dist, pos, and normal
"""
# get the separating axes
edge_dir_a = unique_edges_a[:, 0] - unique_edges_a[:, 1]
edge_dir_b = unique_edges_b[:, 0] - unique_edges_b[:, 1]
edge_dir_a_r = jp.tile(edge_dir_a, reps=(unique_edges_b.shape[0], 1))
edge_dir_b_r = jp.repeat(edge_dir_b, repeats=unique_edges_a.shape[0], axis=0)
edge_edge_axes = jax.vmap(jp.cross)(edge_dir_a_r, edge_dir_b_r)
edge_edge_axes = jax.vmap(lambda x: math.normalize(x, axis=0))(
edge_edge_axes
)
axes = jp.concatenate([normals_a, normals_b, edge_edge_axes])
# for each separating axis, get the support
@jax.vmap
def get_support(axis):
support_a = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_a)
support_b = jax.vmap(jp.dot, in_axes=[None, 0])(axis, vertices_b)
dist1 = support_a.max() - support_b.min()
dist2 = support_b.max() - support_a.min()
sign = jp.where(dist1 > dist2, -1, 1)
dist = jp.minimum(dist1, dist2)
dist = jp.where(~jp.all(axis == 0.0), dist, 1e6) # degenerate axis
return dist, sign
support, sign = get_support(axes)
# choose the best separating axis
best_idx = jp.argmin(support)
best_sign = sign[best_idx]
best_axis = axes[best_idx]
is_edge_contact = best_idx >= (normals_a.shape[0] + normals_b.shape[0])
# get the (reference) face most aligned with the separating axis
dist_a = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_a)
dist_b = jax.vmap(jp.dot, in_axes=[None, 0])(best_axis, normals_b)
a_max = dist_a.argmax()
b_max = dist_b.argmax()
a_min = dist_a.argmin()
b_min = dist_b.argmin()
ref_face = jp.where(best_sign > 0, faces_a[a_max], faces_b[b_max])
ref_face_norm = jp.where(best_sign > 0, normals_a[a_max], normals_b[b_max])
incident_face = jp.where(best_sign > 0, faces_b[b_min], faces_a[a_min])
incident_face_norm = jp.where(
best_sign > 0, normals_b[b_min], normals_a[a_min]
)
dist, pos, normal = _create_contact_manifold(
ref_face,
incident_face,
ref_face_norm,
incident_face_norm,
-best_sign * best_axis,
)
# For edge contacts, we use the clipped face point, mainly for performance
# reasons. For small penetration, the clipped face point is roughly the edge
# contact point.
idx = dist.argmin()
dist = jp.where(
is_edge_contact,
jp.array([dist[idx], 1, 1, 1]),
dist,
)
pos = jp.where(is_edge_contact, jp.tile(pos[idx], (4, 1)), pos)
return dist, pos, normal
def plane_convex(plane: GeomInfo, convex: GeomInfo) -> Contact:
"""Calculates contacts between a plane and a convex object."""
vert = convex.vert
# get points in the convex frame
plane_pos = convex.mat.T @ (plane.pos - convex.pos)
n = convex.mat.T @ plane.mat[:, 2]
support = (plane_pos - vert) @ n
idx = _manifold_points(vert, support > 0, n)
pos = vert[idx]
# convert to world frame
pos = convex.pos + pos @ convex.mat.T
n = plane.mat[:, 2]
frame = jp.stack([math.make_frame(n)] * 4, axis=0)
unique = jp.tril(idx == idx[:, None]).sum(axis=1) == 1
dist = jp.where(unique, -support[idx], 1)
return dist, pos, frame
def sphere_convex(sphere: GeomInfo, convex: GeomInfo) -> Contact:
"""Calculates contact between a sphere and a convex object."""
faces = jp.take(convex.vert, convex.face, axis=0)
normals = convex.facenorm
# Put sphere in convex frame.
sphere_pos = convex.mat.T @ (sphere.pos - convex.pos)
# Get support from face normals.
@jax.vmap
def get_support(faces, normal):
pos = sphere_pos - normal * sphere.size[0]
return jp.dot(pos - faces[0], normal)
support = get_support(faces, normals)
# Pick the face with minimal penetration as long as it has support.
support = jp.where(support >= 0, -1e12, support)
best_idx = support.argmax()
face = faces[best_idx]
normal = normals[best_idx]
# Get closest point between the polygon face and the sphere center point.
# Project the sphere center point onto poly plane. If it's inside polygon
# edge normals, then we're done.
pt = _project_pt_onto_plane(sphere_pos, face[0], normal)
edge_p0 = jp.roll(face, 1, axis=0)
edge_p1 = face
edge_normals = jax.vmap(jp.cross, in_axes=[0, None])(
edge_p1 - edge_p0,
normal,
)
edge_dist = jax.vmap(
lambda plane_pt, plane_norm: (pt - plane_pt).dot(plane_norm)
)(edge_p0, edge_normals)
inside = jp.all(edge_dist <= 0) # lte to handle degenerate edges
# If the point is outside edge normals, project onto the closest edge plane
# that the point is in front of.
degenerate_edge = jp.all(edge_normals == 0, axis=1)
behind = edge_dist < 0.0
edge_dist = jp.where(degenerate_edge | behind, 1e12, edge_dist)
idx = edge_dist.argmin()
edge_pt = math.closest_segment_point(edge_p0[idx], edge_p1[idx], pt)
pt = jp.where(inside, pt, edge_pt)
# Get the normal, dist, and contact position.
n, d = math.normalize_with_norm(pt - sphere_pos)
spt = sphere_pos + n * sphere.size[0]
dist = d - sphere.size[0]
pos = (pt + spt) * 0.5
# Go back to world frame.
n = convex.mat @ n
pos = convex.mat @ pos + convex.pos
return jax.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
def capsule_convex(cap: GeomInfo, convex: GeomInfo) -> Contact:
"""Calculates contacts between a capsule and a convex object."""
# Get convex transformed normals, faces, and vertices.
faces = jp.take(convex.vert, convex.face, axis=0)
normals = convex.facenorm
# Put capsule in convex frame.
cap_pos = convex.mat.T @ (cap.pos - convex.pos)
axis, length = cap.mat[:, 2], cap.size[1]
axis = convex.mat.T @ axis
seg = axis * length
cap_pts = jp.array([
cap_pos - seg,
cap_pos + seg,
])
# Get support from face normals.
@jax.vmap
def get_support(face, normal):
pts = cap_pts - normal * cap.size[0]
sup = jax.vmap(lambda x: jp.dot(x - face[0], normal))(pts)
return sup.min()
support = get_support(faces, normals)
has_support = jp.all(support < 0)
# Pick the face with minimal penetration as long as it has support.
support = jp.where(support >= 0, -1e12, support)
best_idx = support.argmax()
face = faces[best_idx]
normal = normals[best_idx]
# Clip the edge against side planes and create two contact points against the
# face.
edge_p0 = jp.roll(face, 1, axis=0)
edge_p1 = face
edge_normals = jax.vmap(jp.cross, in_axes=[0, None])(
edge_p1 - edge_p0,
normal,
)
cap_pts_clipped, mask = _clip_edge_to_planes(
cap_pts[0], cap_pts[1], edge_p0, edge_normals
)
cap_pts_clipped = cap_pts_clipped - normal * cap.size[0]
face_pts = jax.vmap(_project_pt_onto_plane, in_axes=[0, None, None])(
cap_pts_clipped, face[0], normal
)
# Create variables for the face contact.
pos = (cap_pts_clipped + face_pts) * 0.5
norm = jp.stack([normal] * 2, 0)
penetration = jp.where(
mask & has_support, jp.dot(face_pts - cap_pts_clipped, normal), -1
)
# Get a potential edge contact.
edge_closest, cap_closest = jax.vmap(
math.closest_segment_to_segment_points, in_axes=[0, 0, None, None]
)(edge_p0, edge_p1, cap_pts[0], cap_pts[1])
e_idx = ((edge_closest - cap_closest) ** 2).sum(axis=1).argmin()
cap_closest_pt, edge_closest_pt = cap_closest[e_idx], edge_closest[e_idx]
edge_axis = cap_closest_pt - edge_closest_pt
edge_axis, edge_dist = math.normalize_with_norm(edge_axis)
edge_pos = (
edge_closest_pt + (cap_closest_pt - edge_axis * cap.size[0])
) * 0.5
edge_norm = edge_axis
edge_penetration = cap.size[0] - edge_dist
has_edge_contact = edge_penetration > 0
# Get the contact info.
pos = jp.where(has_edge_contact, pos.at[0].set(edge_pos), pos)
n = -jp.where(has_edge_contact, norm.at[0].set(edge_norm), norm)
# Go back to world frame.
pos = convex.pos + pos @ convex.mat.T
n = n @ convex.mat.T
dist = -jp.where(
has_edge_contact, penetration.at[0].set(edge_penetration), penetration
)
frame = jax.vmap(math.make_frame)(n)
return dist, pos, frame
def convex_convex(c1: GeomInfo, c2: GeomInfo) -> Contact:
"""Calculates contacts between two convex objects."""
if c1.face is None or c2.face is None or c1.vert is None or c2.vert is None:
raise AssertionError('Mesh info missing.')
# pad face vertices so that we can broadcast between geom1 and geom2
s1, s2 = c1.face.shape[-1], c2.face.shape[-1]
if s1 < s2:
face = jp.pad(c1.face, ((0, 0), (0, s2 - s1)), 'edge')
c1 = c1.replace(face=face)
elif s2 < s1:
face = jp.pad(c2.face, ((0, 0), (0, s1 - s2)), 'edge')
c2 = c2.replace(face=face)
# ensure that the first object has fewer verts
swapped = c1.vert.shape[0] > c2.vert.shape[0]
if swapped:
c1, c2 = c2, c1
faces1 = jp.take(c1.vert, c1.face, axis=0)
faces2 = jp.take(c2.vert, c2.face, axis=0)
to_local_pos = c2.mat.T @ (c1.pos - c2.pos)
to_local_mat = c2.mat.T @ c1.mat
faces1 = to_local_pos + faces1 @ to_local_mat.T
normals1 = c1.facenorm @ to_local_mat.T
normals2 = c2.facenorm
vertices1 = to_local_pos + c1.vert @ to_local_mat.T
vertices2 = c2.vert
unique_edges1 = jp.take(vertices1, c1.edge, axis=0)
unique_edges2 = jp.take(vertices2, c2.edge, axis=0)
dist, pos, normal = _sat_hull_hull(
faces1,
faces2,
vertices1,
vertices2,
normals1,
normals2,
unique_edges1,
unique_edges2,
)
# Go back to world frame.
pos = c2.pos + pos @ c2.mat.T
normal = normal @ c2.mat.T
normal = -normal if swapped else normal
frame = jax.vmap(math.make_frame)(normal)
return dist, pos, frame
# store ncon as function attributes
plane_convex.ncon = 4
sphere_convex.ncon = 1
capsule_convex.ncon = 2
convex_convex.ncon = 4
mjx/mujoco/mjx/_src/collision_driver.py
+373
View File
@@ -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_)
mjx/mujoco/mjx/_src/collision_driver_test.py
+499
View File
@@ -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 = """
<mujoco>
<worldbody>
<geom size="40 40 40" type="plane"/>
<body pos="0 0 0.4">
<joint type="free"/>
<geom size="0.5" type="sphere"/>
</body>
</worldbody>
</mujoco>
"""
_SPHERE_SPHERE = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom pos="0 0 0" size="0.2" type="sphere"/>
</body>
<body >
<joint type="free"/>
<geom pos="0 0.3 0" size="0.11" type="sphere"/>
</body>
</worldbody>
</mujoco>
"""
_SPHERE_CAP = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom pos="0 0.3 0" size="0.05" type="sphere"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom fromto="0.0 -0.5 0.14 0.0 0.5 0.14" size="0.1" type="capsule"/>
</body>
</worldbody>
</mujoco>
"""
@parameterized.parameters(
('sphere_plane', _SPHERE_PLANE),
('sphere_sphere', _SPHERE_SPHERE),
('sphere_cap', _SPHERE_CAP),
)
def test_sphere(self, name, mjcf):
d, dx = _collide(mjcf)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-5)
_SPHERE_CONVEX = """
<mujoco>
<worldbody>
<body pos="0.52 0 0.52">
<joint axis="1 0 0" type="free"/>
<geom size="0.05" type="sphere"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.5 0.5 0.5" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_sphere_convex(self):
d, dx = _collide(self._SPHERE_CONVEX)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'sphere_convex', 1e-4)
class CapsuleCollisionTest(parameterized.TestCase):
_CAP_PLANE = """
<mujoco>
<worldbody>
<geom size="40 40 40" type="plane"/>
<body pos="0 0 0.4">
<joint type="free"/>
<geom fromto="-1 0 0 1 0 0" size="0.5" type="capsule"/>
</body>
</worldbody>
</mujoco>
"""
_CAP_CAP = """
<mujoco model="two_capsules">
<worldbody>
<body>
<joint type="free"/>
<geom fromto="0.62235904 0.58846647 0.651046 1.5330081 0.33564585 0.977849"
size="0.05" type="capsule"/>
</body>
<body>
<joint type="free"/>
<geom fromto="0.5505271 0.60345304 0.476661 1.3900293 0.30709633 0.932082"
size="0.05" type="capsule"/>
</body>
</worldbody>
</mujoco>
"""
@parameterized.parameters(
('capsule_plane', _CAP_PLANE),
('capsule_capsule', _CAP_CAP),
)
def test_capsule(self, name, mjcf):
d, dx = _collide(mjcf)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, name, 1e-4)
_PARALLEL_CAP = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom fromto="-0.5 0.1 0.25 0.5 0.1 0.25" size="0.1" type="capsule"/>
</body>
<body>
<joint type="free"/>
<geom fromto="-0.5 0.1 0.1 0.5 0.1 0.1" size="0.1" type="capsule"/>
</body>
</worldbody>
</mujoco>
"""
def test_parallel_capsules(self):
"""Tests that two parallel capsules are colliding at the midpoint."""
_, dx = _collide(self._PARALLEL_CAP)
np.testing.assert_allclose(dx.contact.dist, -0.05)
np.testing.assert_allclose(
dx.contact.pos[0],
np.array([0.0, 0.1, (0.15 + 0.2) / 2.0]),
atol=1e-5,
)
np.testing.assert_allclose(
dx.contact.frame[0, 0, :], np.array([0, 0.0, -1.0]), atol=1e-5
)
_CAP_BOX = """
<mujoco>
<worldbody>
<body pos="0 0 0.54">
<joint axis="1 0 0" type="free"/>
<geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.5 0.5 0.5" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_capsule_convex(self):
"""Tests a capsule-convex collision for a face contact."""
d, dx = _collide(self._CAP_BOX)
for field in dataclasses.fields(Contact):
_assert_attr_eq(dx.contact, d.contact, field.name, 'capsule_convex', 1e-4)
_CAP_EDGE_BOX = """
<mujoco>
<worldbody>
<body pos="0.5 0 0.55" euler="0 30 0">
<joint axis="1 0 0" type="free"/>
<geom fromto="-0.6 0 0 0.6 0 0" size="0.05" type="capsule"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.5 0.5 0.5" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_capsule_convex_edge(self):
"""Tests a capsule-convex collision for an edge contact."""
d, dx = _collide(self._CAP_EDGE_BOX)
c = dx.contact
self.assertEqual(c.pos.shape[0], 2)
self.assertGreater(c.dist[1], 0)
# extract the contact point with penetration
c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
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 = """
<mujoco>
<worldbody>
<geom size="40 40 40" type="plane"/>
<body pos="0 0 0.7" euler="45 0 0">
<joint axis="1 0 0" type="free"/>
<geom size="0.5 0.5 0.5" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_box_plane(self):
"""Tests box collision with a plane."""
d, dx = _collide(self._BOX_PLANE)
np.testing.assert_array_less(dx.contact.dist[:2], 0)
np.testing.assert_array_less(-dx.contact.dist[2:], 0)
# extract the contact points with penetration
c = jax.tree_map(lambda x: jp.take(x, jp.array([0, 1]), axis=0), dx.contact)
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 = """
<mujoco>
<worldbody>
<body pos="0.0 1.0 0.2">
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="box"/>
</body>
<body pos="0.1 1.0 0.495" euler="0.1 -0.1 45">
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1 0.1" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_box_box(self):
"""Tests a face contact for a box-box collision."""
d, dx = _collide(self._BOX_BOX)
c = dx.contact
self.assertEqual(c.pos.shape[0], 4)
np.testing.assert_array_less(c.dist, 0)
np.testing.assert_array_almost_equal(c.pos[:, 2], np.array([0.39] * 4), 2)
np.testing.assert_array_almost_equal(
c.frame[:, 0, :], np.array([[0.0, 0.0, 1.0]] * 4)
)
np.testing.assert_array_almost_equal(
c.frame.reshape((-1, 9)), d.contact.frame[:4, :]
)
_BOX_BOX_EDGE = """
<mujoco>
<worldbody>
<body pos="-1.0 -1.0 0.2">
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="box"/>
</body>
<body pos="-1.0 -1.2 0.55" euler="0 45 30">
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1 0.1" type="box"/>
</body>
</worldbody>
</mujoco>
"""
def test_box_box_edge(self):
"""Tests an edge contact for a box-box collision."""
d, dx = _collide(self._BOX_BOX_EDGE)
# Only one contact point.
np.testing.assert_array_less(dx.contact.dist[:1], 0)
np.testing.assert_array_less(-dx.contact.dist[1:], 0)
# extract the contact point with penetration
c = jax.tree_map(lambda x: jp.take(x, 0, axis=0)[None], dx.contact)
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 = """
<mujoco>
<asset>
<mesh name="tetrahedron" file="meshes/tetrahedron.stl" scale="0.1 0.1 0.1" />
<mesh name="dodecahedron" file="meshes/dodecahedron.stl" scale="0.01 0.01 0.01" />
</asset>
<worldbody>
<body pos="0.0 2.0 0.096">
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="mesh" mesh="tetrahedron"/>
</body>
<body pos="0.0 2.0 0.289" euler="0.1 -0.1 45">
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1 0.1" type="mesh" mesh="dodecahedron"/>
</body>
</worldbody>
</mujoco>
"""
def test_convex_convex(self):
"""Tests generic convex-convex collision."""
directory = epath.resource_path('mujoco.mjx')
assets = {
'meshes/tetrahedron.stl': (
directory / 'test_data' / 'meshes/tetrahedron.stl'
).read_bytes(),
'meshes/dodecahedron.stl': (
directory / 'test_data' / 'meshes/dodecahedron.stl'
).read_bytes(),
}
_, dx = _collide(self._CONVEX_CONVEX, assets=assets)
c = dx.contact
# Only one contact point for an edge contact.
self.assertLess(c.dist[0], 0)
np.testing.assert_array_less(0, c.dist[1:])
np.testing.assert_array_almost_equal(c.frame[0, 0], np.array([0, 0, 1]))
class BodyPairFilterTest(absltest.TestCase):
"""Tests that certain body pairs get filtered."""
_SELF_COLLISION = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom size="0.2"/>
<geom size="0.2"/>
</body>
</worldbody>
</mujoco>
"""
def test_filter_self_collision(self):
"""Tests that self collisions get filtered."""
d, dx = _collide(self._SELF_COLLISION)
self.assertEqual(dx.contact.pos.shape[0], d.contact.pos.shape[0])
self.assertEqual(dx.contact.pos.shape[0], 0)
_PARENT_CHILD = """
<mujoco>
<worldbody>
<body>
<joint type="free"/>
<geom size="0.2"/>
<body pos="0.0 0.0 0.1">
<joint type="hinge"/>
<geom size="0.2"/>
</body>
</body>
</worldbody>
</mujoco>
"""
def test_filter_parent_child(self):
"""Tests that parent-child collisions get filtered."""
m = mujoco.MjModel.from_xml_string(self._PARENT_CHILD)
mx = mjx.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 = """
<mujoco>
<custom>
<numeric data="2" name="max_contact_points"/>
</custom>
<worldbody>
<body pos="0 0 0.54">
<joint axis="1 0 0" type="free"/>
<geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
</body>
<body pos="0 0 0.54">
<joint axis="1 0 0" type="free"/>
<geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
</body>
<body pos="0 0 0.54">
<joint axis="1 0 0" type="free"/>
<geom fromto="-0.4 0 0 0.4 0 0" size="0.05" type="capsule"/>
</body>
</worldbody>
</mujoco>
"""
def test_top_k_contacts(self):
m = mujoco.MjModel.from_xml_string(self._CAPSULES)
mx_top_k = mjx.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()
mjx/mujoco/mjx/_src/collision_primitive.py
+126
View File
@@ -0,0 +1,126 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Collision primitives."""
from typing import Tuple
import jax
from jax import numpy as jp
from mujoco.mjx._src import math
# pylint: disable=g-importing-member
from mujoco.mjx._src.collision_base import Contact
from mujoco.mjx._src.collision_base import GeomInfo
# pylint: enable=g-importing-member
def _plane_sphere(
plane_normal: jax.Array,
plane_pos: jax.Array,
sphere_pos: jax.Array,
radius: jax.Array,
) -> Tuple[jax.Array, jax.Array]:
"""Returns the penetration and contact point between a plane and sphere."""
cdist = jp.dot(sphere_pos - plane_pos, plane_normal)
dist = cdist - radius
pos = sphere_pos - plane_normal * (radius + 0.5 * dist)
return dist, pos
def plane_sphere(plane: GeomInfo, sphere: GeomInfo) -> Contact:
"""Calculates contact between a plane and a sphere."""
n = plane.mat[:, 2]
dist, pos = _plane_sphere(n, plane.pos, sphere.pos, sphere.size[0])
return jax.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
def plane_capsule(plane: GeomInfo, cap: GeomInfo) -> Contact:
"""Calculates two contacts between a capsule and a plane."""
n, axis = plane.mat[:, 2], cap.mat[:, 2]
# align contact frames with capsule axis
b, b_norm = math.normalize_with_norm(axis - n * jp.dot(n, axis))
y, z = jp.array([0.0, 1.0, 0.0]), jp.array([0.0, 0.0, 1.0])
b = jp.where(b_norm < 0.5, jp.where((-0.5 < n[1]) & (n[1] < 0.5), y, z), b)
frame = jp.array([[n, b, jp.cross(n, b)]])
segment = axis * cap.size[1]
contacts = []
for offset in [segment, -segment]:
dist, pos = _plane_sphere(n, plane.pos, cap.pos + offset, cap.size[0])
dist = jp.expand_dims(dist, axis=0)
pos = jp.expand_dims(pos, axis=0)
contacts.append((dist, pos, frame))
return jax.tree_map(lambda *x: jp.concatenate(x), *contacts)
def _sphere_sphere(
pos1: jax.Array, radius1: jax.Array, pos2: jax.Array, radius2: jax.Array
) -> Contact:
"""Returns the penetration, contact point, and normal between two spheres."""
n, dist = math.normalize_with_norm(pos2 - pos1)
n = jp.where(dist == 0.0, jp.array([1.0, 0.0, 0.0]), n)
dist = dist - (radius1 + radius2)
pos = pos1 + n * (radius1 + dist * 0.5)
return dist, pos, n
def sphere_sphere(s1: GeomInfo, s2: GeomInfo) -> Contact:
"""Calculates contact between two spheres."""
dist, pos, n = _sphere_sphere(s1.pos, s1.size[0], s2.pos, s2.size[0])
return jax.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
def sphere_capsule(sphere: GeomInfo, cap: GeomInfo) -> Contact:
"""Calculates one contact between a sphere and a capsule."""
axis, length = cap.mat[:, 2], cap.size[1]
segment = axis * length
pt = math.closest_segment_point(
cap.pos - segment, cap.pos + segment, sphere.pos
)
dist, pos, n = _sphere_sphere(sphere.pos, sphere.size[0], pt, cap.size[0])
return jax.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
def capsule_capsule(cap1: GeomInfo, cap2: GeomInfo) -> Contact:
"""Calculates one contact between two capsules."""
axis1, length1, axis2, length2 = (
cap1.mat[:, 2],
cap1.size[1],
cap2.mat[:, 2],
cap2.size[1],
)
seg1, seg2 = axis1 * length1, axis2 * length2
pt1, pt2 = math.closest_segment_to_segment_points(
cap1.pos - seg1,
cap1.pos + seg1,
cap2.pos - seg2,
cap2.pos + seg2,
)
radius1, radius2 = cap1.size[0], cap2.size[0]
dist, pos, n = _sphere_sphere(pt1, radius1, pt2, radius2)
return jax.tree_map(
lambda x: jp.expand_dims(x, axis=0), (dist, pos, math.make_frame(n))
)
# store ncon as function attributes
plane_sphere.ncon = 1
plane_capsule.ncon = 2
sphere_sphere.ncon = 1
sphere_capsule.ncon = 1
capsule_capsule.ncon = 1
mjx/mujoco/mjx/_src/constraint.py
+340
View File
@@ -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
mjx/mujoco/mjx/_src/constraint_test.py
+144
View File
@@ -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()
mjx/mujoco/mjx/_src/dataclasses.py
+151
View File
@@ -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)}
)
mjx/mujoco/mjx/_src/device.py
+282
View File
@@ -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
mjx/mujoco/mjx/_src/device_test.py
+184
View File
@@ -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(
'<mujoco><option solver="Newton"/><worldbody/></mujoco>'
)
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(
'<mujoco><option integrator="implicit"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
def test_cone(self):
m = mujoco.MjModel.from_xml_string(
'<mujoco><option cone="elliptic"/><worldbody/></mujoco>'
)
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
def test_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("""
<mujoco>
<worldbody>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.2 0.2 0.2" type="box"/>
</body>
<body>
<joint axis="1 0 0" type="free"/>
<geom size="0.1 0.1" type="cylinder"/>
</body>
</worldbody>
</mujoco>
""")
with self.assertRaises(NotImplementedError):
_ = mjx.device_put(m)
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/forward.py
+356
View File
@@ -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
mjx/mujoco/mjx/_src/forward_test.py
+125
View File
@@ -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()
mjx/mujoco/mjx/_src/io.py
+105
View File
@@ -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
mjx/mujoco/mjx/_src/io_test.py
+42
View File
@@ -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()
mjx/mujoco/mjx/_src/math.py
+356
View File
@@ -0,0 +1,356 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Some useful math functions."""
from typing import Optional, Tuple, Union
import jax
from jax import numpy as jp
def norm(
x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None
) -> jax.Array:
"""Calculates a linalg.norm(x) that's safe for gradients at x=0.
Avoids a poorly defined gradient for jnp.linal.norm(0) see
https://github.com/google/jax/issues/3058 for details
Args:
x: A jnp.array
axis: The axis along which to compute the norm
Returns:
Norm of the array x.
"""
is_zero = jp.allclose(x, 0.0)
# temporarily swap x with ones if is_zero, then swap back
x = jp.where(is_zero, jp.ones_like(x), x)
n = jp.linalg.norm(x, axis=axis)
n = jp.where(is_zero, 0.0, n)
return n
def normalize_with_norm(
x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None
) -> Tuple[jax.Array, jax.Array]:
"""Normalizes an array.
Args:
x: A jnp.array
axis: The axis along which to compute the norm
Returns:
A tuple of (normalized array x, the norm).
"""
n = norm(x, axis=axis)
x = x / (n + 1e-6 * (n == 0.0))
return x, n
def normalize(
x: jax.Array, axis: Optional[Union[Tuple[int, ...], int]] = None
) -> jax.Array:
"""Normalizes an array.
Args:
x: A jnp.array
axis: The axis along which to compute the norm
Returns:
normalized array x
"""
return normalize_with_norm(x, axis=axis)[0]
def rotate(vec: jax.Array, quat: jax.Array) -> jax.Array:
"""Rotates a vector vec by a unit quaternion quat.
Args:
vec: (3,) a vector
quat: (4,) a quaternion
Returns:
ndarray(3) containing vec rotated by quat.
"""
if len(vec.shape) != 1:
raise ValueError('vec must have no batch dimensions.')
s, u = quat[0], quat[1:]
r = 2 * (jp.dot(u, vec) * u) + (s * s - jp.dot(u, u)) * vec
r = r + 2 * s * jp.cross(u, vec)
return r
def quat_inv(q: jp.ndarray) -> jp.ndarray:
"""Calculates the inverse of quaternion q.
Args:
q: (4,) quaternion [w, x, y, z]
Returns:
The inverse of q, where qmult(q, inv_quat(q)) = [1, 0, 0, 0].
"""
return q * jp.array([1, -1, -1, -1])
def quat_sub(u: jax.Array, v: jax.Array) -> jax.Array:
"""Subtracts two quaternions (u - v) as a 3D velocity."""
q = quat_mul(quat_inv(v), u)
axis, angle = quat_to_axis_angle(q)
return axis * angle
def quat_mul(u: jax.Array, v: jax.Array) -> jax.Array:
"""Multiplies two quaternions.
Args:
u: (4,) quaternion (w,x,y,z)
v: (4,) quaternion (w,x,y,z)
Returns:
A quaternion u * v.
"""
return jp.array([
u[0] * v[0] - u[1] * v[1] - u[2] * v[2] - u[3] * v[3],
u[0] * v[1] + u[1] * v[0] + u[2] * v[3] - u[3] * v[2],
u[0] * v[2] - u[1] * v[3] + u[2] * v[0] + u[3] * v[1],
u[0] * v[3] + u[1] * v[2] - u[2] * v[1] + u[3] * v[0],
])
def quat_mul_axis(q: jax.Array, axis: jax.Array) -> jax.Array:
"""Multiplies a quaternion and an axis.
Args:
q: (4,) quaternion (w,x,y,z)
axis: (3,) axis (x,y,z)
Returns:
A quaternion q * axis
"""
return jp.array([
-q[1] * axis[0] - q[2] * axis[1] - q[3] * axis[2],
q[0] * axis[0] + q[2] * axis[2] - q[3] * axis[1],
q[0] * axis[1] + q[3] * axis[0] - q[1] * axis[2],
q[0] * axis[2] + q[1] * axis[1] - q[2] * axis[0],
])
# TODO(erikfrey): benchmark this against brax's quat_to_3x3
def quat_to_mat(q: jax.Array) -> jax.Array:
"""Converts a quaternion into a 9-dimensional rotation matrix."""
q = jp.outer(q, q)
return jp.array([
[
q[0, 0] + q[1, 1] - q[2, 2] - q[3, 3],
2 * (q[1, 2] - q[0, 3]),
2 * (q[1, 3] + q[0, 2]),
],
[
2 * (q[1, 2] + q[0, 3]),
q[0, 0] - q[1, 1] + q[2, 2] - q[3, 3],
2 * (q[2, 3] - q[0, 1]),
],
[
2 * (q[1, 3] - q[0, 2]),
2 * (q[2, 3] + q[0, 1]),
q[0, 0] - q[1, 1] - q[2, 2] + q[3, 3],
],
])
def quat_to_axis_angle(q: jax.Array) -> Tuple[jax.Array, jax.Array]:
"""Converts a quaternion into axis and angle."""
axis, sin_a_2 = normalize_with_norm(q[1:])
angle = 2 * jp.arctan2(sin_a_2, q[0])
angle = jp.where(angle > jp.pi, angle - 2 * jp.pi, angle)
return axis, angle
def axis_angle_to_quat(axis: jax.Array, angle: jax.Array) -> jax.Array:
"""Provides a quaternion that describes rotating around axis by angle.
Args:
axis: (3,) axis (x,y,z)
angle: () float angle to rotate by
Returns:
A quaternion that rotates around axis by angle
"""
s, c = jp.sin(angle * 0.5), jp.cos(angle * 0.5)
return jp.insert(axis * s, 0, c)
def quat_integrate(q: jax.Array, v: jax.Array, dt: jax.Array) -> jax.Array:
"""Integrates a quaternion given angular velocity and dt."""
v, norm_ = normalize_with_norm(v)
angle = dt * norm_
q_res = axis_angle_to_quat(v, angle)
q_res = quat_mul(q, q_res)
return normalize(q_res)
def inert_mul(i: jax.Array, v: jax.Array) -> jax.Array:
"""Multiply inertia by motion, producing force.
Args:
i: (10,) inertia (inertia matrix, position, mass)
v: (6,) spatial motion
Returns:
resultant force
"""
tri_id = jp.array([[0, 3, 4], [3, 1, 5], [4, 5, 2]]) # cinert inr order
inr, pos, mass = i[tri_id], i[6:9], i[9]
ang = jp.dot(inr, v[:3]) + jp.cross(pos, v[3:])
vel = mass * v[3:] - jp.cross(pos, v[:3])
return jp.concatenate((ang, vel))
def transform_motion(vel: jax.Array, offset: jax.Array, rotmat: jax.Array):
"""Transform spatial motion.
Args:
vel: (6,) spatial motion (3 angular, 3 linear)
offset: (3,) translation
rotmat: (3, 3) rotation
Returns:
6d spatial velocity
"""
# TODO(robotics-simulation): are quaternions faster here
ang, vel = vel[:3], vel[3:]
vel = rotmat.T @ (vel - jp.cross(offset, ang))
ang = rotmat.T @ ang
return jp.concatenate([ang, vel])
def motion_cross(u, v):
"""Cross product of two motions.
Args:
u: (6,) spatial motion
v: (6,) spatial motion
Returns:
resultant spatial motion
"""
ang = jp.cross(u[:3], v[:3])
vel = jp.cross(u[3:], v[:3]) + jp.cross(u[:3], v[3:])
return jp.concatenate((ang, vel))
def motion_cross_force(v, f):
"""Cross product of a motion and force.
Args:
v: (6,) spatial motion
f: (6,) force
Returns:
resultant force
"""
ang = jp.cross(v[:3], f[:3]) + jp.cross(v[3:], f[3:])
vel = jp.cross(v[:3], f[3:])
return jp.concatenate((ang, vel))
def orthogonals(a: jax.Array) -> Tuple[jax.Array, jax.Array]:
"""Returns orthogonal vectors `b` and `c`, given a vector `a`."""
y, z = jp.array([0, 1, 0]), jp.array([0, 0, 1])
b = jp.where((-0.5 < a[1]) & (a[1] < 0.5), y, z)
b = b - a * a.dot(b)
# normalize b. however if a is a zero vector, zero b as well.
b = normalize(b) * jp.any(a)
return b, jp.cross(a, b)
def make_frame(a: jax.Array) -> jax.Array:
"""Makes a right-handed 3D frame given a direction."""
a = normalize(a)
b, c = orthogonals(a)
return jp.array([a, b, c])
# Geometry.
def closest_segment_point(
a: jax.Array, b: jax.Array, pt: jax.Array
) -> jax.Array:
"""Returns the closest point on the a-b line segment to a point pt."""
ab = b - a
t = jp.dot(pt - a, ab) / (jp.dot(ab, ab) + 1e-6)
return a + jp.clip(t, 0.0, 1.0) * ab
def closest_segment_point_and_dist(
a: jax.Array, b: jax.Array, pt: jax.Array
) -> Tuple[jax.Array, jax.Array]:
"""Returns closest point on the line segment and the distance squared."""
closest = closest_segment_point(a, b, pt)
dist = (pt - closest).dot(pt - closest)
return closest, dist
def closest_segment_to_segment_points(
a0: jax.Array, a1: jax.Array, b0: jax.Array, b1: jax.Array
) -> Tuple[jax.Array, jax.Array]:
"""Returns closest points between two line segments."""
# Gets the closest segment points by first finding the closest points
# between two lines. Points are then clipped to be on the line segments
# and edge cases with clipping are handled.
dir_a, len_a = normalize_with_norm(a1 - a0)
dir_b, len_b = normalize_with_norm(b1 - b0)
# Segment mid-points.
half_len_a = len_a * 0.5
half_len_b = len_b * 0.5
a_mid = a0 + dir_a * half_len_a
b_mid = b0 + dir_b * half_len_b
# Translation between two segment mid-points.
trans = a_mid - b_mid
# Parametrize points on each line as follows:
# point_on_a = a_mid + t_a * dir_a
# point_on_b = b_mid + t_b * dir_b
# and analytically minimize the distance between the two points.
dira_dot_dirb = dir_a.dot(dir_b)
dira_dot_trans = dir_a.dot(trans)
dirb_dot_trans = dir_b.dot(trans)
denom = 1 - dira_dot_dirb * dira_dot_dirb
orig_t_a = (-dira_dot_trans + dira_dot_dirb * dirb_dot_trans) / (denom + 1e-6)
orig_t_b = dirb_dot_trans + orig_t_a * dira_dot_dirb
t_a = jp.clip(orig_t_a, -half_len_a, half_len_a)
t_b = jp.clip(orig_t_b, -half_len_b, half_len_b)
best_a = a_mid + dir_a * t_a
best_b = b_mid + dir_b * t_b
# Resolve edge cases where both closest points are clipped to the segment
# endpoints by recalculating the closest segment points for the current
# clipped points, and then picking the pair of points with smallest
# distance. An example of this edge case is when lines intersect but line
# segments don't.
new_a, d1 = closest_segment_point_and_dist(a0, a1, best_b)
new_b, d2 = closest_segment_point_and_dist(b0, b1, best_a)
best_a = jp.where(d1 < d2, new_a, best_a)
best_b = jp.where(d1 < d2, best_b, new_b)
return best_a, best_b
mjx/mujoco/mjx/_src/math_test.py
+202
View File
@@ -0,0 +1,202 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for math."""
from absl.testing import absltest
from absl.testing import parameterized
import jax.numpy as jp
from mujoco.mjx._src import math
import numpy as np
def _get_rand_point(seed=None):
if seed is not None:
np.random.seed(seed)
verts = np.random.randn(1, 3)
return verts[0, :]
def _get_rand_line_segment(seed=None):
if seed is not None:
np.random.seed(seed)
verts = np.random.randn(2, 3)
return verts[0, :], verts[1, :]
def _get_rand_unit(seed: int):
np.random.seed(seed)
theta = np.random.random(1) * 2 * np.pi
a = (np.random.random(1) - 0.5) * 2.0
phi = np.arccos(a)
x = np.sin(phi) * np.cos(theta)
y = np.sin(phi) * np.sin(theta)
z = np.cos(phi)
return jp.array([x, y, z]).squeeze()
class OrthoganalsTest(parameterized.TestCase):
"""Tests the orthogonals function."""
@parameterized.parameters(range(30))
def test_orthogonals(self, i):
a = _get_rand_unit(i)
b, c = math.orthogonals(a)
np.testing.assert_almost_equal(jp.linalg.norm(a), 1)
np.testing.assert_almost_equal(jp.linalg.norm(b), 1)
np.testing.assert_almost_equal(jp.linalg.norm(c), 1)
self.assertAlmostEqual(np.abs(a.dot(b)), 0, 6)
self.assertAlmostEqual(np.abs(b.dot(c)), 0, 6)
self.assertAlmostEqual(np.abs(a.dot(c)), 0, 6)
def _minimize(fn, sample_fn, lb, ub, tol, max_iter=20, seed=42):
"""Minimize a function using the cross-entropy method."""
assert lb.shape == ub.shape, "bounds need to have the same shape"
np.random.seed(seed)
i, n = 0, 1_000
mu = (ub + lb) * 0.5
sigma = (ub - lb) * 0.5
size = lb.shape[0]
val, prev_val = fn(mu), None
while prev_val is None or np.abs(val - prev_val) > tol:
params = sample_fn(mu, sigma, n, size, lb, ub)
vals = np.array([fn(p) for p in params])
if val < vals.min(): # early exit
return mu
idx = vals.argsort()
best_idx = idx[: int(n * 0.05)]
mu = params[best_idx].mean(axis=0)
sigma = params[best_idx].std(axis=0) + 1e-10
prev_val = val
val = fn(mu)
i += 1
if i == max_iter:
break
return mu
def _closest_segment_to_segment_points(a0, a1, b0, b1):
dir_a = a1 - a0
len_a = np.sqrt(dir_a.dot(dir_a))
half_len_a = len_a / 2
dir_a = dir_a / len_a
dir_b = b1 - b0
len_b = np.sqrt(dir_b.dot(dir_b))
half_len_b = len_b / 2
dir_b = dir_b / len_b
a_mid = a0 + dir_a * half_len_a
b_mid = b0 + dir_b * half_len_b
# Parametrize both line segments.
def fn(t):
best_a = a_mid + dir_a * t[0]
best_b = b_mid + dir_b * t[1]
return (best_a - best_b).dot(best_a - best_b)
def sample_fn(mu, sigma, n, size, lb, ub):
params = np.random.normal(mu, sigma, size=(n, size))
params = np.clip(params, lb, ub)
return params
lb = np.array([-half_len_a, -half_len_b])
ub = np.array([half_len_a, half_len_b])
ta, tb = _minimize(fn, sample_fn, lb, ub, tol=1e-4)
best_a = a_mid + dir_a * ta
best_b = b_mid + dir_b * tb
return best_a, best_b
class ClosestSegmentSegmentPointsTest(parameterized.TestCase):
"""Tests for closest segment-to-segment points."""
def test_closest_segments_points(self):
a0 = jp.array([0.73432405, 0.12372768, 0.20272314])
a1 = jp.array([1.10600128, 0.88555209, 0.65209485])
b0 = jp.array([0.85599262, 0.61736299, 0.9843583])
b1 = jp.array([1.84270939, 0.92891793, 1.36343326])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [1.09063, 0.85404, 0.63351], 5)
self.assertSequenceAlmostEqual(best_b, [0.99596, 0.66156, 1.03813], 5)
def test_intersecting_segments(self):
"""Tests segments that intersect."""
a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0])
b0, b1 = jp.array([-1.0, 0.0, 0.0]), jp.array([1.0, 0.0, 0.0])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5)
self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5)
def test_intersecting_lines(self):
"""Tests that intersecting lines get clipped."""
a0, a1 = jp.array([0.2, 0.2, 0.0]), jp.array([1.0, 1.0, 0.0])
b0, b1 = jp.array([0.2, 0.4, 0.0]), jp.array([1.0, 2.0, 0.0])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.3, 0.3, 0.0], 2)
self.assertSequenceAlmostEqual(best_b, [0.2, 0.4, 0.0], 2)
def test_parallel_segments(self):
"""Tests that parallel segments have closest points at the midpoint."""
a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0])
b0, b1 = jp.array([1.0, 0.0, -1.0]), jp.array([1.0, 0.0, 1.0])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5)
self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.0], 5)
def test_parallel_offset_segments(self):
"""Tests that offset parallel segments are close at segment endpoints."""
a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0])
b0, b1 = jp.array([1.0, 0.0, 1.0]), jp.array([1.0, 0.0, 3.0])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 1.0], 5)
self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 1.0], 5)
def test_zero_length_segments(self):
"""Test that zero length segments don't return NaNs."""
a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, -1.0])
b0, b1 = jp.array([1.0, 0.0, 0.1]), jp.array([1.0, 0.0, 0.1])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, -1.0], 5)
self.assertSequenceAlmostEqual(best_b, [1.0, 0.0, 0.1], 5)
def test_overlapping_segments(self):
"""Tests that perfectly overlapping segments intersect at the midpoints."""
a0, a1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0])
b0, b1 = jp.array([0.0, 0.0, -1.0]), jp.array([0.0, 0.0, 1.0])
best_a, best_b = math.closest_segment_to_segment_points(a0, a1, b0, b1)
self.assertSequenceAlmostEqual(best_a, [0.0, 0.0, 0.0], 5)
self.assertSequenceAlmostEqual(best_b, [0.0, 0.0, 0.0], 5)
params = list(zip(np.repeat(np.arange(10), 10), np.tile(np.arange(10), 10)))
@parameterized.parameters(*params)
def test_closest_segment_to_segment_points(self, i, j):
a0, a1 = _get_rand_line_segment(i)
b0, b1 = _get_rand_line_segment(j)
expected = _closest_segment_to_segment_points(a0, a1, b0, b1)
ans = math.closest_segment_to_segment_points(a0, a1, b0, b1)
expected_dist = (expected[0] - expected[1]).dot(expected[0] - expected[1])
test_dist = (ans[0] - ans[1]).dot(ans[0] - ans[1])
self.assertAlmostEqual(expected_dist, test_dist, 4)
if __name__ == "__main__":
absltest.main()
mjx/mujoco/mjx/_src/mesh.py
+232
View File
@@ -0,0 +1,232 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Mesh processing."""
import itertools
from typing import Dict, Optional, Sequence, Tuple
import mujoco
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import GeomType
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
import numpy as np
from scipy import spatial
import trimesh
_BOX_CORNERS = list(itertools.product((-1, 1), (-1, 1), (-1, 1)))
# pyformat: disable
# Rectangular box faces using a counter-clockwise winding order convention.
_BOX_FACES = [
0, 4, 5, 1, # left
0, 2, 6, 4, # bottom
6, 7, 5, 4, # front
2, 3, 7, 6, # right
1, 5, 7, 3, # top
0, 1, 3, 2, # back
]
# pyformat: enable
_MAX_HULL_FACE_VERTICES = 20
_CONVEX_CACHE: Dict[Tuple[int, int], Dict[str, np.ndarray]] = {}
_DERIVED_ARGS = [
'geom_convex_face',
'geom_convex_vert',
'geom_convex_edge',
'geom_convex_facenormal',
]
DERIVED = {(Model, d) for d in _DERIVED_ARGS}
def _box(size: np.ndarray):
"""Creates a mesh for a box with rectangular faces."""
box_corners = np.array(_BOX_CORNERS)
vert = box_corners * size.reshape(-1, 3)
face = np.array([_BOX_FACES]).reshape(-1, 4)
return vert, face
def _get_face_norm(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
"""Calculates face normals given vertices and face indexes."""
assert len(vert.shape) == 2 and len(face.shape) == 2, (
f'vert and face should have dim of 2, got {len(vert.shape)} and '
f'{len(face.shape)}'
)
face_vert = vert[face, :]
# use CCW winding order convention
edge0 = face_vert[:, 1, :] - face_vert[:, 0, :]
edge1 = face_vert[:, -1, :] - face_vert[:, 0, :]
face_norm = np.cross(edge0, edge1)
face_norm = face_norm / np.linalg.norm(face_norm, axis=1).reshape((-1, 1))
return face_norm
def _get_unique_edges(vert: np.ndarray, face: np.ndarray) -> np.ndarray:
"""Returns unique edges.
Args:
vert: (n_vert, 3) vertices
face: (n_face, n_vert) face index array
Returns:
edges: tuples of vertex indexes for each edge
"""
r_face = np.roll(face, 1, axis=1)
edges = np.concatenate(np.array([face, r_face]).T)
# do a first pass to remove duplicates
edges.sort(axis=1)
edges = np.unique(edges, axis=0)
edges = edges[edges[:, 0] != edges[:, 1]] # get rid of edges from padded face
# get normalized edge directions
edge_vert = vert.take(edges, axis=0)
edge_dir = edge_vert[:, 0] - edge_vert[:, 1]
norms = np.sqrt(np.sum(edge_dir**2, axis=1))
edge_dir = edge_dir / norms.reshape((-1, 1))
# get the first unique edge for all pairwise comparisons
diff1 = edge_dir[:, None, :] - edge_dir[None, :, :]
diff2 = edge_dir[:, None, :] + edge_dir[None, :, :]
matches = (np.linalg.norm(diff1, axis=-1) < 1e-6) | (
np.linalg.norm(diff2, axis=-1) < 1e-6
)
matches = np.tril(matches).sum(axis=-1)
unique_edge_idx = np.where(matches == 1)[0]
return edges[unique_edge_idx]
def _convex_hull_2d(points: np.ndarray, normal: np.ndarray) -> np.ndarray:
"""Calculates the convex hull for a set of points on a plane."""
# project points onto the closest axis plane
best_axis = np.abs(np.eye(3).dot(normal)).argmax()
axis = np.eye(3)[best_axis]
d = points.dot(axis).reshape((-1, 1))
axis_points = points - d * axis
axis_points = axis_points[:, list({0, 1, 2} - {best_axis})]
# get the polygon face, and make the points ccw wrt the face normal
c = spatial.ConvexHull(axis_points)
order_ = np.where(axis.dot(normal) > 0, 1, -1)
order_ *= np.where(best_axis == 1, -1, 1)
hull_point_idx = c.vertices[::order_]
assert (axis_points - c.points).sum() == 0
return hull_point_idx
def _merge_coplanar(tm: trimesh.Trimesh) -> np.ndarray:
"""Merges coplanar facets."""
if not tm.facets:
return tm.faces.copy() # no facets
if not tm.faces.shape[0]:
raise ValueError('Mesh has no faces.')
# Get faces.
face_idx = set(range(tm.faces.shape[0])) - set(np.concatenate(tm.facets))
face_idx = np.array(list(face_idx))
faces = tm.faces[face_idx] if face_idx.shape[0] > 0 else np.array([])
# Get facets.
facets = []
for i, facet in enumerate(tm.facets):
point_idx = np.unique(tm.faces[facet])
points = tm.vertices[point_idx]
normal = tm.facets_normal[i]
# convert triangulated facet to a polygon
hull_point_idx = _convex_hull_2d(points, normal)
face = point_idx[hull_point_idx]
# resize faces that exceed max polygon vertices
every = face.shape[0] // _MAX_HULL_FACE_VERTICES + 1
face = face[::every]
facets.append(face)
# Pad facets so that they can be stacked.
max_len = max(f.shape[0] for f in facets) if facets else faces.shape[1]
assert max_len <= _MAX_HULL_FACE_VERTICES
for i, f in enumerate(facets):
if f.shape[0] < max_len:
f = np.pad(f, (0, max_len - f.shape[0]), 'edge')
facets[i] = f
if not faces.shape[0]:
assert facets
return np.array(facets) # no faces, return facets
# Merge faces and facets.
faces = np.pad(faces, ((0, 0), (0, max_len - faces.shape[1])), 'edge')
return np.concatenate([faces, facets])
def _get_faces_verts(
m: mujoco.MjModel,
) -> Tuple[Sequence[np.ndarray], Sequence[np.ndarray]]:
"""Extracts mesh faces and vertices from MjModel."""
verts, faces = [], []
for i in range(m.nmesh):
last = (i + 1) >= m.nmesh
face_start = m.mesh_faceadr[i]
face_end = m.mesh_faceadr[i + 1] if not last else m.mesh_face.shape[0]
face = m.mesh_face[face_start:face_end]
faces.append(face)
vert_start = m.mesh_vertadr[i]
vert_end = m.mesh_vertadr[i + 1] if not last else m.mesh_vert.shape[0]
vert = m.mesh_vert[vert_start:vert_end]
verts.append(vert)
return verts, faces
def _geom_mesh_kwargs(
vert: np.ndarray, face: np.ndarray
) -> Dict[str, np.ndarray]:
"""Generates convex mesh attributes for mjx.Model."""
tm = trimesh.Trimesh(vertices=vert, faces=face)
tm_convex = trimesh.convex.convex_hull(tm)
vert = np.array(tm_convex.vertices)
face = _merge_coplanar(tm_convex)
return {
'geom_convex_face': face,
'geom_convex_vert': vert,
'geom_convex_edge': _get_unique_edges(vert, face),
'geom_convex_facenormal': _get_face_norm(vert, face),
}
def get(m: mujoco.MjModel) -> Dict[str, Sequence[Optional[np.ndarray]]]:
"""Derives geom mesh attributes for mjx.Model from MjModel."""
kwargs = {k: [] for k in _DERIVED_ARGS}
verts, faces = _get_faces_verts(m)
for geomid in range(m.ngeom):
dataid = m.geom_dataid[geomid]
typ = m.geom_type[geomid]
if typ == GeomType.BOX:
vert, face = _box(m.geom_size[geomid])
elif dataid >= 0:
vert, face = verts[dataid], faces[dataid]
else:
kwargs = {k: kwargs[k] + [None] for k in _DERIVED_ARGS}
continue
key = (hash(vert.data.tobytes()), hash(face.data.tobytes()))
if key not in _CONVEX_CACHE:
_CONVEX_CACHE[key] = _geom_mesh_kwargs(vert, face)
kwargs = {k: kwargs[k] + [_CONVEX_CACHE[key][k]] for k in _DERIVED_ARGS}
return kwargs
mjx/mujoco/mjx/_src/mesh_test.py
+119
View File
@@ -0,0 +1,119 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests for mesh.py."""
from absl.testing import absltest
from mujoco.mjx._src import mesh
import numpy as np
class GeomMeshKwargsTest(absltest.TestCase):
def test_pyramid(self):
"""Tests that a triangulated pyramid converts to merged coplanar faces."""
vert = np.array([
[-0.025, 0.05, 0.05],
[-0.025, -0.05, -0.05],
[-0.025, -0.05, 0.05],
[-0.025, 0.05, -0.05],
[0.075, 0.0, 0.0],
])
face = np.array(
[[0, 1, 2], [0, 3, 1], [0, 4, 3], [0, 2, 4], [2, 1, 4], [1, 3, 4]]
)
h = mesh._geom_mesh_kwargs(vert, face)
# get index of vertices in h['geom_convex_vert'] for vertices in vert
dist = np.repeat(vert, vert.shape[0], axis=0) - np.tile(
h['geom_convex_vert'], (vert.shape[0], 1)
)
dist = (dist**2).sum(axis=1).reshape((vert.shape[0], -1))
vidx = np.argmin(dist, axis=0)
# check verts
np.testing.assert_array_equal(h['geom_convex_vert'], vert[vidx])
# check face vertices
map_ = {v: k for k, v in enumerate(vidx)}
h_face = np.vectorize(map_.get)(h['geom_convex_face'])
face_verts = sorted([tuple(sorted(set(s))) for s in h_face.tolist()])
expected_face_verts = sorted([
(0, 3, 4), (1, 3, 4), (0, 2, 4), (0, 1, 2, 3), (1, 2, 4)])
self.assertSequenceEqual(
face_verts,
expected_face_verts,
)
# check edges
unique_edge = np.vectorize(map_.get)(h['geom_convex_edge'])
unique_edge = np.array(sorted(unique_edge.tolist()))
np.testing.assert_array_equal(
unique_edge,
np.array([[0, 2], [0, 3], [0, 4], [1, 4], [2, 4], [3, 4]]),
)
# face normals
self.assertEqual(h['geom_convex_facenormal'].shape, (5, 3))
class ConvexHull2DTest(absltest.TestCase):
def test_convex_hull_2d_axis1(self):
"""Tests for the correct winding order of a polgyon with +y normal."""
pts = np.array([
[-0.04634297, -0.06652775, 0.05853534],
[-0.01877651, -0.08309858, -0.05236476],
[0.02362804, -0.08010745, 0.05499557],
[0.04066505, -0.09034877, -0.01354446],
[-0.07255043, -0.06837638, -0.00781699],
])
normal = np.array([-0.18467607, -0.97768016, 0.10018111])
idx = mesh._convex_hull_2d(pts, normal)
expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0])
expected /= np.linalg.norm(expected)
np.testing.assert_array_almost_equal(normal, expected)
def test_convex_hull_2d_axis2(self):
"""Tests for the correct winding order for a polgyon with +z normal."""
pts = np.array([
[0.08607829, -0.03881998, -0.03291714],
[-0.01877651, -0.08309858, -0.05236476],
[0.05470364, 0.00027677, -0.08371042],
[-0.01010019, -0.02708892, -0.0957297],
[0.04066505, -0.09034877, -0.01354446],
])
normal = np.array([0.3839915, -0.60171936, -0.70034587])
idx = mesh._convex_hull_2d(pts, normal)
expected = np.cross(pts[idx][1] - pts[idx][0], pts[idx][2] - pts[idx][0])
expected /= np.linalg.norm(expected)
np.testing.assert_array_almost_equal(normal, expected)
class UniqueEdgesTest(absltest.TestCase):
def test_tetrahedron_edges(self):
"""Tests unique edges for a tetrahedron."""
vert = np.array(
[[-0.1, 0.0, -0.1], [0.0, 0.1, 0.1], [0.1, 0.0, -0.1], [0.0, -0.1, 0.1]]
)
face = np.array([[0, 1, 2], [0, 2, 3], [0, 3, 1], [2, 1, 3]])
idx = mesh._get_unique_edges(vert, face)
np.testing.assert_array_equal(
idx, np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]])
)
if __name__ == '__main__':
absltest.main()
mjx/mujoco/mjx/_src/passive.py
+141
View File
@@ -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
mjx/mujoco/mjx/_src/passive_test.py
+114
View File
@@ -0,0 +1,114 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests 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("""
<mujoco>
<option density="1" viscosity="2" wind="0.1 0.2 0.3">
<flag passive="disable"/>
</option>
<worldbody>
<body>
<joint damping="1" axis="1 0 0" type="ball"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
</body>
</worldbody>
</mujoco>
""")
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()
mjx/mujoco/mjx/_src/scan.py
+432
View File
@@ -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
mjx/mujoco/mjx/_src/scan_test.py
+227
View File
@@ -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 = """
<mujoco>
<compiler inertiafromgeom="true"/>
<worldbody>
<body>
<joint type="free"/>
<geom size=".15" mass="1" type="sphere"/>
<body>
<joint axis="1 0 0" pos="1 0 0" type="ball"/>
<geom size=".15" mass="2" type="sphere"/>
</body>
<body>
<joint axis="1 0 0" pos="2 0 0" type="hinge"/>
<joint axis="0 1 0" pos="3 0 0" type="slide"/>
<geom size=".15" mass="3" type="sphere"/>
</body>
</body>
</worldbody>
</mujoco>
"""
def test_flat_empty(self):
"""Test scanning over just world body."""
m = mujoco.MjModel.from_xml_string("""
<mujoco model="world_body">
<worldbody/>
</mujoco>
""")
m = mjx.device_put(m)
def fn(body_id):
return body_id + 1
b_in = jp.array([1])
b_expect = jp.array([2])
b_out = scan.flat(m, fn, 'b', 'b', b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
def test_flat_joints(self):
"""Tests scanning over bodies with joints of different types."""
m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
m = mjx.device_put(m)
# we will test two functions:
# 1) j_fn receives jnt_types as a jp array
# 2) s_fn receives jnt_types as a static np array and can switch on it
j_fn = lambda jnt_pos, val: val + jp.sum(jnt_pos)
s_fn = lambda jnt_types, val: val + sum(jnt_types)
b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]])
b_expect = jp.array([[0, 0], [1, 1], [3, 3], [8, 8]])
b_out = scan.flat(m, j_fn, 'jb', 'b', m.jnt_pos, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
b_out = scan.flat(m, s_fn, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
# None should be omitted from the results
def no_free(jnt_types, val):
if tuple(jnt_types) == (JointType.FREE,):
return None
return val + sum(jnt_types)
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
b_out = scan.flat(m, no_free, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
# we should not call functions for which we know we will discard the results
def no_world(jnt_types, val):
if jnt_types.size == 0:
self.fail('world has no dofs, should not be called')
return val + sum(jnt_types)
v_in = jp.ones((m.nv, 1))
scan.flat(m, no_world, 'jv', 'v', m.jnt_type, v_in)
def test_body_tree(self):
"""Tests tree scanning over bodies with different joint counts."""
m = mujoco.MjModel.from_xml_string(self._MULTI_DOF_XML)
m = mjx.device_put(m)
# we will test two functions:
# 1) j_fn receives jnt_pos which is a jp array
# 2) s_fn receives jnt_types which is a static np array
def j_fn(carry, jnt_pos, val):
carry = jp.zeros_like(val) if carry is None else carry
return carry + val + jp.sum(jnt_pos)
def s_fn(carry, jnt_types, val):
carry = jp.zeros_like(val) if carry is None else carry
return carry + val + sum(jnt_types)
b_in = jp.array([[0, 0], [1, 1], [2, 2], [3, 3]])
b_expect = jp.array([[0, 0], [1, 1], [4, 4], [9, 9]])
b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
# and reverse too:
b_expect = jp.array([[12, 12], [12, 12], [3, 3], [8, 8]])
b_out = scan.body_tree(m, j_fn, 'jb', 'b', m.jnt_pos, b_in, reverse=True)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
b_out = scan.body_tree(m, s_fn, 'jb', 'b', m.jnt_type, b_in, reverse=True)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
# None should be omitted from the results
def no_free(carry, jnt_types, val):
if tuple(jnt_types) == (JointType.FREE,):
return None
carry = jp.zeros_like(val) if carry is None else carry
return carry + val + sum(jnt_types)
b_expect = jp.array([[0, 0], [3, 3], [8, 8]])
b_out = scan.body_tree(m, no_free, 'jb', 'b', m.jnt_type, b_in)
np.testing.assert_equal(np.array(b_out), np.array(b_expect))
_MULTI_ACT_XML = """
<mujoco>
<option timestep="0.02"/>
<compiler autolimits="true"/>
<default>
<geom contype="0" conaffinity="0"/>
</default>
<worldbody>
<body pos="0 0 -0.5">
<joint type="free" name="joint0" range="-37.81 86.15"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0 -0.5">
<joint axis="1 0 0" type="hinge" name="joint1" range="-32.88 5.15" actuatorfrcrange="-0.48 0.72"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0 -0.5">
<joint axis="0 1 0" type="slide" name="joint2" range="-83.31 54.77"/>
<joint axis="1 0 0" type="slide" name="joint3" range="-87.84 60.58" actuatorfrcrange="-0.62 0.28"/>
<joint axis="0 1 0" type="hinge" name="joint4" range="-49.96 83.03" actuatorfrcrange="-0.01 0.64"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0 -0.5">
<joint axis="1 0 0" type="slide" name="joint5" range="-1.53 87.07" actuatorfrcrange="-0.06 0.26"/>
<joint axis="0 1 0" type="slide" name="joint6" range="-63.57 9.95"/>
<joint axis="1 0 0" type="hinge" name="joint7" range="-22.50 41.20" actuatorfrcrange="-0.41 0.84"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
<body pos="0 0 -0.5">
<joint axis="0 1 0" type="hinge" name="joint8" range="-34.67 72.56"/>
<joint axis="1 0 0" type="hinge" name="joint9" range="-1.81 16.22" actuatorfrcrange="-0.41 0.14"/>
<joint axis="1 0 0" type="ball" name="joint10" range="0.00 6.02" actuatorfrcrange="-0.74 0.12"/>
<geom pos="0 0.5 0" size=".15" mass="1" type="sphere"/>
</body>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<position joint="joint6" gear="1"/>
<intvelocity joint="joint1" kp="2000" actrange="-0.7 2.3"/>
<motor joint="joint3" gear="42"/>
<velocity joint="joint2" kv="123"/>
<position joint="joint0" ctrlrange="-0.9472 0.9472"/>
<intvelocity joint="joint7" kp="2000" actrange="-0.7 2.3"/>
<general joint="joint5" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
<general joint="joint4" ctrlrange="0.1 2.34346" biastype="affine" gainprm="35 0 0" biasprm="0 -35 -0.65"/>
</actuator>
</mujoco>
"""
def 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()
mjx/mujoco/mjx/_src/smooth.py
+467
View File
@@ -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
mjx/mujoco/mjx/_src/smooth_test.py
+166
View File
@@ -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("""
<mujoco>
<option timestep="0.01"/>
<worldbody>
<body>
<joint type="free"/>
<geom size="0.1"/>
</body>
</worldbody>
</mujoco>
""")
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()
mjx/mujoco/mjx/_src/solver.py
+356
View File
@@ -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
mjx/mujoco/mjx/_src/solver_test.py
+130
View File
@@ -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()
mjx/mujoco/mjx/_src/support.py
+81
View File
@@ -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)
mjx/mujoco/mjx/_src/support_test.py
+81
View File
@@ -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()
mjx/mujoco/mjx/_src/test_util.py
+340
View File
@@ -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
mjx/mujoco/mjx/_src/test_util_test.py
+38
View File
@@ -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()
mjx/mujoco/mjx/_src/types.py
+660
View File
@@ -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<margin-gap (ngeom,)
geom_gap: include in solver if dist<margin-gap (ngeom,)
geom_convex_face: vertex face data, MJX only (ngeom,)
geom_convex_vert: vertex data, MJX only (ngeom,)
geom_convex_edge: unique edge data, MJX only (ngeom,)
geom_convex_facenormal: normal face data, MJX only (ngeom,)
pair_dim: contact dimensionality (npair,)
pair_geom1: id of geom1 (npair,)
pair_geom2: id of geom2 (npair,)
pair_solref: solver reference: contact normal (npair, mjNREF)
pair_solreffriction: solver reference: contact friction (npair, mjNREF)
pair_solimp: solver impedance: contact (npair, mjNIMP)
pair_margin: include in solver if dist<margin-gap (npair,)
pair_gap: include in solver if dist<margin-gap (npair,)
pair_friction: tangent1, 2, spin, roll1, 2 (npair, 5)
exclude_signature: (body1+1) << 16 + body2+1 (nexclude,)
eq_type: constraint type (mjtEq) (neq,)
eq_obj1id: id of object 1 (neq,)
eq_obj2id: id of object 2 (neq,)
eq_active0: initial enable/disable constraint state (neq,)
eq_solref: constraint solver reference (neq, mjNREF)
eq_solimp: constraint solver impedance (neq, mjNIMP)
eq_data: numeric data for constraint (neq, mjNEQDATA)
actuator_trntype: transmission type (mjtTrn) (nu,)
actuator_dyntype: dynamics type (mjtDyn) (nu,)
actuator_gaintype: gain type (mjtGain) (nu,)
actuator_biastype: bias type (mjtBias) (nu,)
actuator_trnid: transmission id: joint, tendon, site (nu, 2)
actuator_actadr: first activation address; -1: stateless (nu,)
actuator_actnum: number of activation variables (nu,)
actuator_ctrllimited: is control limited (nu,)
actuator_forcelimited: is force limited (nu,)
actuator_actlimited: is activation limited (nu,)
actuator_dynprm: dynamics parameters (nu, mjNDYN)
actuator_gainprm: gain parameters (nu, mjNGAIN)
actuator_biasprm: bias parameters (nu, mjNBIAS)
actuator_ctrlrange: range of controls (nu, 2)
actuator_forcerange: range of forces (nu, 2)
actuator_actrange: range of activations (nu, 2)
actuator_gear: scale length and transmitted force (nu, 6)
numeric_adr: address of field in numeric_data (nnumeric,)
numeric_data: array of all numeric fields (nnumericdata,)
name_numericadr: numeric name pointers (nnumeric,)
names: names of all objects, 0-terminated (nnames,)
"""
nq: int
nv: int
nu: int
na: int
nbody: int
njnt: int
ngeom: int
nmesh: int
npair: int
nexclude: int
neq: int
nnumeric: int
nM: int
opt: Option
stat: Statistic
qpos0: jax.Array
qpos_spring: jax.Array
body_parentid: np.ndarray
body_rootid: np.ndarray
body_weldid: np.ndarray
body_jntnum: np.ndarray
body_jntadr: np.ndarray
body_dofnum: np.ndarray
body_dofadr: np.ndarray
body_geomnum: np.ndarray
body_geomadr: np.ndarray
body_pos: jax.Array
body_quat: jax.Array
body_ipos: jax.Array
body_iquat: jax.Array
body_mass: jax.Array
body_subtreemass: jax.Array
body_inertia: jax.Array
body_invweight0: jax.Array
jnt_type: np.ndarray
jnt_qposadr: np.ndarray
jnt_dofadr: np.ndarray
jnt_bodyid: np.ndarray
jnt_limited: np.ndarray
jnt_actfrclimited: np.ndarray
jnt_solref: jax.Array
jnt_solimp: jax.Array
jnt_pos: jax.Array
jnt_axis: jax.Array
jnt_stiffness: jax.Array
jnt_range: jax.Array
jnt_actfrcrange: jax.Array
jnt_margin: jax.Array
dof_bodyid: np.ndarray
dof_jntid: np.ndarray
dof_parentid: np.ndarray
dof_Madr: np.ndarray
dof_solref: jax.Array
dof_solimp: jax.Array
dof_frictionloss: jax.Array
dof_armature: jax.Array
dof_damping: jax.Array
dof_invweight0: jax.Array
dof_M0: jax.Array
geom_type: np.ndarray
geom_contype: np.ndarray
geom_conaffinity: np.ndarray
geom_condim: np.ndarray
geom_bodyid: np.ndarray
geom_priority: np.ndarray
geom_solmix: jax.Array
geom_solref: jax.Array
geom_solimp: jax.Array
geom_size: jax.Array
geom_pos: jax.Array
geom_quat: jax.Array
geom_friction: jax.Array
geom_margin: jax.Array
geom_gap: jax.Array
pair_dim: np.ndarray
pair_geom1: np.ndarray
pair_geom2: np.ndarray
geom_convex_face: Sequence[jax.Array]
geom_convex_vert: Sequence[jax.Array]
geom_convex_edge: Sequence[jax.Array]
geom_convex_facenormal: Sequence[jax.Array]
pair_solref: jax.Array
pair_solreffriction: jax.Array
pair_solimp: jax.Array
pair_margin: jax.Array
pair_gap: jax.Array
pair_friction: jax.Array
exclude_signature: np.ndarray
eq_type: np.ndarray
eq_obj1id: np.ndarray
eq_obj2id: np.ndarray
eq_active0: np.ndarray
eq_solref: jax.Array
eq_solimp: jax.Array
eq_data: jax.Array
actuator_trntype: np.ndarray
actuator_dyntype: np.ndarray
actuator_gaintype: np.ndarray
actuator_biastype: np.ndarray
actuator_trnid: np.ndarray
actuator_actadr: np.ndarray
actuator_actnum: np.ndarray
actuator_ctrllimited: np.ndarray
actuator_forcelimited: np.ndarray
actuator_actlimited: np.ndarray
actuator_dynprm: jax.Array
actuator_gainprm: jax.Array
actuator_biasprm: jax.Array
actuator_ctrlrange: jax.Array
actuator_forcerange: jax.Array
actuator_actrange: jax.Array
actuator_gear: jax.Array
numeric_adr: np.ndarray
numeric_data: np.ndarray
name_numericadr: np.ndarray
names: bytes
class Contact(PyTreeNode):
"""Result of collision detection functions.
Attributes:
dist: distance between nearest points; neg: penetration
pos: position of contact point: midpoint between geoms (3,)
frame: normal is in [0-2] (9,)
includemargin: include if dist<includemargin=margin-gap (1,)
friction: tangent1, 2, spin, roll1, 2 (5,)
solref: constraint solver reference, normal direction (mjNREF,)
solreffriction: constraint solver reference, friction directions (mjNREF,)
solimp: constraint solver impedance (mjNIMP,)
dim: contact space dimensionality: 1, 3, 4 or 6
geom1: id of geom 1
geom2: id of geom 2
efc_address: address in efc; -1: not included
"""
dist: jax.Array
pos: jax.Array
frame: jax.Array
includemargin: jax.Array
friction: jax.Array
solref: jax.Array
solreffriction: jax.Array
solimp: jax.Array
# unsupported: mu, H
dim: np.ndarray
geom1: jax.Array
geom2: jax.Array
efc_address: np.ndarray
# unsupported: exclude
@classmethod
def zero(cls, shape=(0,)) -> '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
mjx/mujoco/mjx/benchmark/__init__.py
+14
View File
@@ -0,0 +1,14 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
mjx/mujoco/mjx/benchmark/benchmark.py
+116
View File
@@ -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()
mjx/mujoco/mjx/benchmark/model/barkour_v0/README.md
+35
View File
@@ -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.
<p float="left">
<img src="assets/joystick.gif" width="400">
</p>
## 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 `<option timestep="0.002" iterations="4" solver="CG"/>`.
* Non-visual geoms were removed from the torso, to speed up rendering.
## Publications
If you use this work in an academic context, please cite the following publication:
@misc{caluwaerts2023barkour,
title={Barkour: Benchmarking Animal-level Agility with Quadruped Robots},
author={Ken Caluwaerts and Atil Iscen and J. Chase Kew and Wenhao Yu and Tingnan Zhang and Daniel Freeman and Kuang-Huei Lee and Lisa Lee and Stefano Saliceti and Vincent Zhuang and Nathan Batchelor and Steven Bohez and Federico Casarini and Jose Enrique Chen and Omar Cortes and Erwin Coumans and Adil Dostmohamed and Gabriel Dulac-Arnold and Alejandro Escontrela and Erik Frey and Roland Hafner and Deepali Jain and Bauyrjan Jyenis and Yuheng Kuang and Edward Lee and Linda Luu and Ofir Nachum and Ken Oslund and Jason Powell and Diego Reyes and Francesco Romano and Feresteh Sadeghi and Ron Sloat and Baruch Tabanpour and Daniel Zheng and Michael Neunert and Raia Hadsell and Nicolas Heess and Francesco Nori and Jeff Seto and Carolina Parada and Vikas Sindhwani and Vincent Vanhoucke and Jie Tan},
year={2023},
eprint={2305.14654},
archivePrefix={arXiv},
primaryClass={cs.RO}
}
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/abduction.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/barkour_v0_mjx.xml
+243
View File
@@ -0,0 +1,243 @@
<mujoco model="barkour v0 brax">
<compiler angle="radian" meshdir="." texturedir="assets" autolimits="true"/>
<option timestep="0.002" iterations="4" ls_iterations="6" solver="CG">
<flag eulerdamp="disable"/>
</option>
<statistic meansize="0.183574"/>
<visual>
<headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global azimuth="120" elevation="-20"/>
</visual>
<default>
<geom contype="0" conaffinity="0" type="mesh"/>
<joint range="-1.5708 1.5708" armature="0.01090125" damping="0.5239"/>
<default class="abductor">
<joint range="-1.0472 1.0472"/>
</default>
<default class="hip_front">
<joint range="-1.22173 3.24631"/>
</default>
<default class="hip_hind">
<joint range="-1.98968 2.46091"/>
</default>
<default class="knee">
<joint range="0 2.5132" axis="0 -0.0775009 0.996992"/>
</default>
<default class="multi_mode_controlled_actuator">
<general biastype="affine" gainprm="35 0 0 0 0 0 0 0 0 0" biasprm="0 -35 -0.65 0 0 0 0 0 0 0" forcerange="-18.0 18.0" ctrlrange="-2 2"/>
</default>
<default class="visual">
<geom contype="0" conaffinity="0" density="0" group="1"/>
<default class="visual_upper_right1">
<geom rgba="0.768627 0.886275 0.952941 1"/>
</default>
<default class="visual_upper_right2">
<geom rgba="0.972549 0.529412 0.00392157 1"/>
</default>
<default class="visual_abduction">
<geom rgba="0.537255 0.854902 0.827451 1"/>
</default>
<default class="visual_foot">
<geom rgba="0.301961 0.301961 0.301961 1"/>
</default>
</default>
<default class="collision">
<geom group="2"/>
<default class="upper_right1">
<geom rgba="0.768627 0.886275 0.952941 1"/>
</default>
<default class="upper_right2">
<geom rgba="0.972549 0.529412 0.00392157 1"/>
</default>
<default class="abduction">
<geom rgba="0.537255 0.854902 0.827451 1"/>
</default>
<default class="foot">
<geom rgba="0.301961 0.301961 0.301961 1"/>
</default>
</default>
<default class="lowerLegFootLeft">
<geom type="sphere" pos="-0.191284 -0.0191638 -0.013" size="0.014" contype="1"
conaffinity="0" rgba="1 0 0 1"/>
</default>
<default class="lowerLegFootRight">
<geom type="sphere" pos="-0.191284 -0.0191638 0.013" size="0.014" contype="1"
conaffinity="0" rgba="1 0 0 1"/>
</default>
</default>
<custom>
<numeric data="0.0 0.0 0.21 1.0 0.0 0.0 0.0 0.0 0.5 1.0 0.0 0.5 1.0 0.0 0.5 1.0 0.0 0.5 1.0" name="init_qpos"/>
</custom>
<asset>
<mesh file="head.stl"/>
<mesh file="powercable.stl"/>
<mesh file="handle.stl"/>
<mesh file="head_mount.stl"/>
<mesh file="body.stl"/>
<mesh file="abduction.stl"/>
<mesh file="upper_right_2.stl"/>
<mesh file="upper_right_3.stl"/>
<mesh file="upper_right_1.stl"/>
<mesh file="lower_leg_1to1.stl"/>
<mesh file="foot.stl"/>
<mesh file="upper_left_2.stl"/>
<mesh file="upper_left_1.stl"/>
<mesh file="upper_left_3.stl"/>
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
</asset>
<worldbody>
<site name="origin"/>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<camera name="default" pos="0.846 -1.465 0.916" xyaxes="0.866 0.500 0.000 -0.171 0.296 0.940"/>
<geom name="floor" size="0 0 0.05" type="plane" conaffinity="1" material="groundplane"/>
<body name="chassis">
<camera name="track" pos="0.846 -1.465 0.916" xyaxes="0.866 0.500 0.000 -0.171 0.296 0.940" mode="trackcom"/>
<freejoint/>
<inertial pos="0.0196226 -0.00015133 0.0611588" quat="0.000990813 0.68703 0.000216603 0.726628" mass="4.48878" diaginertia="0.071033 0.0619567 0.0192519"/>
<geom class="visual" pos="-0.00448404 -0.000225838 0.0576402" rgba="0.647059 0.647059 0.647059 1" mesh="head"/>
<geom class="visual" pos="-0.00448404 -0.000225838 0.0576402" rgba="0.768627 0.886275 0.952941 1" mesh="powercable"/>
<geom class="visual" pos="-0.00448404 -0.000225838 0.0576402" rgba="0.917647 0.917647 0.917647 1" mesh="handle"/>
<geom class="visual" pos="-0.00448404 -0.000225838 0.0576402" rgba="0.231373 0.380392 0.705882 1" mesh="head_mount"/>
<geom class="visual" pos="-0.00448404 -0.000225838 0.0576402" rgba="0.984314 0.517647 0.862745 1" mesh="body"/>
<body name="abduction_1" pos="0.130533 -0.056 0.0508" quat="1.30945e-06 0.161152 -2.13816e-07 0.98693">
<inertial pos="-0.0521152 0.00350917 0.0171912" quat="0.387877 0.592262 0.592492 0.384358" mass="0.639437" diaginertia="0.000866008 0.000565866 0.000479767"/>
<joint class="abductor" name="abduction_front_right" axis="-0.94806 0 0.318092"/>
<geom class="visual_abduction" pos="-0.0540394 0.0217 0.0181312" mesh="abduction"/>
<body name="upper_right_asm_1" pos="-0.0540394 0.0217 0.0181312" quat="0.284632 0.284629 -0.647292 0.647289">
<inertial pos="-0.0253655 -0.0179374 -0.0465027" quat="-0.245689 0.639007 0.153351 0.712594" mass="0.942155" diaginertia="0.00539403 0.00519403 0.000795298"/>
<joint class="hip_front" name="hip_front_right" axis="0 0 -1"/>
<geom class="visual_upper_right2" mesh="upper_right_2"/>
<geom class="visual_upper_right2" mesh="upper_right_3"/>
<geom class="visual_upper_right1" mesh="upper_right_1"/>
<body name="lower_leg_1to1_front_right" pos="-0.193523 -0.104637 -0.0792" quat="0.312742 -0.0121371 0.0368314 -0.949046">
<inertial pos="-0.0577509 -0.0097034 0.0114624" quat="-0.047103 0.705359 -0.0102465 0.70721" mass="0.169623" diaginertia="0.000828741 0.000813964 3.49901e-05"/>
<joint class="knee" name="knee_front_right"/>
<geom class="visual" pos="0.00320019 0.0240604 -0.0141615" rgba="0.32549 0.529412 0.980392 1" mesh="lower_leg_1to1"/>
<geom class="visual_foot" pos="0.00320019 0.0240604 -0.0141615" mesh="foot"/>
<geom class="lowerLegFootRight"/>
</body>
</body>
</body>
<body name="abduction_2" pos="0.130533 0.056 0.0508" quat="0.161152 1.09564e-06 0.98693 1.09564e-06">
<inertial pos="-0.0521152 0.00350917 0.0171912" quat="0.387877 0.592262 0.592492 0.384358" mass="0.639437" diaginertia="0.000866008 0.000565866 0.000479767"/>
<joint class="abductor" name="abduction_front_left" axis="0.94806 0 -0.318092"/>
<geom class="visual_abduction" pos="-0.0540394 0.0217 0.0181312" mesh="abduction"/>
<body name="upper_left_asm_1" pos="-0.0540394 0.0217 0.0181312" quat="0.671818 0.671821 0.220587 -0.220588">
<inertial pos="0.0306562 0.00629189 -0.0466005" quat="-0.113342 0.751294 0.0555641 0.647784" mass="0.938791" diaginertia="0.00538157 0.00518445 0.000790347"/>
<joint class="hip_front" name="hip_front_left" axis="0 0 1"/>
<geom class="visual" rgba="0.980392 0.713725 0.00392157 1" mesh="upper_left_2"/>
<geom class="visual" rgba="0.498039 0.498039 0.498039 1" mesh="upper_left_1"/>
<geom class="visual" rgba="1 0.756863 0.054902 1" mesh="upper_left_3"/>
<body name="lower_leg_1to1_front_left" pos="0.208835 0.0691954 -0.0792" quat="0.0386264 0.995249 0.0893024 0.0034659">
<inertial pos="-0.0577509 -0.00780463 -0.0129639" quat="-0.047103 0.705359 -0.0102465 0.70721" mass="0.169623" diaginertia="0.000828741 0.000813964 3.49901e-05"/>
<joint class="knee" name="knee_front_left"/>
<geom class="visual" pos="0.00320019 0.0259591 -0.0385878" rgba="0.32549 0.529412 0.980392 1" mesh="lower_leg_1to1"/>
<geom class="visual_foot" pos="0.00320019 0.0259591 -0.0385878" mesh="foot"/>
<geom class="lowerLegFootLeft"/>
</body>
</body>
</body>
<body name="abduction_3" pos="-0.134667 -0.056 0.0508" quat="1.30945e-06 0.98693 2.13816e-07 -0.161152">
<inertial pos="-0.0521152 0.00350917 0.0171912" quat="0.387877 0.592262 0.592492 0.384358" mass="0.639437" diaginertia="0.000866008 0.000565866 0.000479767"/>
<joint class="abductor" name="abduction_hind_right" axis="0.94806 0 -0.318092"/>
<geom class="visual_abduction" pos="-0.0540394 0.0217 0.0181312" mesh="abduction"/>
<body name="upper_right_asm_2" pos="-0.0540394 0.0217 0.0181312" quat="0.64729 0.647292 0.28463 -0.284631">
<inertial pos="-0.0253655 -0.0179374 -0.0465027" quat="-0.245689 0.639007 0.153351 0.712594" mass="0.942155" diaginertia="0.00539403 0.00519403 0.000795298"/>
<joint class="hip_hind" name="hip_hind_right" axis="0 0 -1"/>
<geom class="visual_upper_right2" mesh="upper_right_2"/>
<geom class="visual_upper_right2" mesh="upper_right_3"/>
<geom class="visual_upper_right1" mesh="upper_right_1"/>
<body name="lower_leg_1to1_hind_right" pos="-0.193523 -0.104637 -0.0792" quat="0.312742 -0.0121371 0.0368314 -0.949046">
<inertial pos="-0.0577509 -0.0097034 0.0114624" quat="-0.047103 0.705359 -0.0102465 0.70721" mass="0.169623" diaginertia="0.000828741 0.000813964 3.49901e-05"/>
<joint class="knee" name="knee_hind_right"/>
<geom class="visual" pos="0.00320019 0.0240604 -0.0141615" rgba="0.32549 0.529412 0.980392 1" mesh="lower_leg_1to1"/>
<geom class="visual_foot" pos="0.00320019 0.0240604 -0.0141615" mesh="foot"/>
<geom class="lowerLegFootRight"/>
</body>
</body>
</body>
<body name="abduction_4" pos="-0.134667 0.056 0.0508" quat="0.98693 0 -0.161152 0">
<inertial pos="-0.0521152 0.00350917 0.0171912" quat="0.387877 0.592262 0.592492 0.384358" mass="0.639437" diaginertia="0.000866008 0.000565866 0.000479767"/>
<joint class="abductor" name="abduction_hind_left" axis="-0.94806 0 0.318092"/>
<geom class="visual_abduction" pos="-0.0540394 0.0217 0.0181312" mesh="abduction"/>
<body name="upper_left_asm_2" pos="-0.0540394 0.0217 0.0181312" quat="-0.220587 -0.220588 0.67182 -0.671818">
<inertial pos="0.0306562 0.00629189 -0.0466005" quat="-0.113342 0.751294 0.0555641 0.647784" mass="0.938791" diaginertia="0.00538157 0.00518445 0.000790347"/>
<joint class="hip_hind" name="hip_hind_left" axis="0 0 1"/>
<geom class="visual" rgba="0.498039 0.498039 0.498039 1" mesh="upper_left_1"/>
<geom class="visual" rgba="1 0.756863 0.054902 1" mesh="upper_left_3"/>
<geom class="visual" rgba="0.980392 0.713725 0.00392157 1" mesh="upper_left_2"/>
<body name="lower_leg_1to1_hind_left" pos="0.208835 0.0691954 -0.0792" quat="0.0386264 0.995249 0.0893024 0.0034659">
<inertial pos="-0.0577509 -0.00780463 -0.0129639" quat="-0.047103 0.705359 -0.0102465 0.70721" mass="0.169623" diaginertia="0.000828741 0.000813964 3.49901e-05"/>
<joint class="knee" name="knee_hind_left"/>
<geom class="visual" pos="0.00320019 0.0259591 -0.0385878" rgba="0.32549 0.529412 0.980392 1" mesh="lower_leg_1to1"/>
<geom class="visual_foot" pos="0.00320019 0.0259591 -0.0385878" mesh="foot"/>
<geom class="lowerLegFootLeft"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<general name="abduction_front_left" class="multi_mode_controlled_actuator" joint="abduction_front_left"/>
<general name="hip_front_left" class="multi_mode_controlled_actuator" joint="hip_front_left"/>
<general name="knee_front_left" class="multi_mode_controlled_actuator" joint="knee_front_left"/>
<general name="abduction_hind_left" class="multi_mode_controlled_actuator" joint="abduction_hind_left"/>
<general name="hip_hind_left" class="multi_mode_controlled_actuator" joint="hip_hind_left"/>
<general name="knee_hind_left" class="multi_mode_controlled_actuator" joint="knee_hind_left"/>
<general name="abduction_front_right" class="multi_mode_controlled_actuator" joint="abduction_front_right"/>
<general name="hip_front_right" class="multi_mode_controlled_actuator" joint="hip_front_right"/>
<general name="knee_front_right" class="multi_mode_controlled_actuator" joint="knee_front_right"/>
<general name="abduction_hind_right" class="multi_mode_controlled_actuator" joint="abduction_hind_right"/>
<general name="hip_hind_right" class="multi_mode_controlled_actuator" joint="hip_hind_right"/>
<general name="knee_hind_right" class="multi_mode_controlled_actuator" joint="knee_hind_right"/>
</actuator>
<sensor>
<jointpos joint="abduction_front_left" name="abduction_front_left_pos"/>
<jointpos joint="hip_front_left" name="hip_front_left_pos"/>
<jointpos joint="knee_front_left" name="knee_front_left_pos"/>
<jointpos joint="abduction_hind_left" name="abduction_hind_left_pos"/>
<jointpos joint="hip_hind_left" name="hip_hind_left_pos"/>
<jointpos joint="knee_hind_left" name="knee_hind_left_pos"/>
<jointpos joint="abduction_front_right" name="abduction_front_right_pos"/>
<jointpos joint="hip_front_right" name="hip_front_right_pos"/>
<jointpos joint="knee_front_right" name="knee_front_right_pos"/>
<jointpos joint="abduction_hind_right" name="abduction_hind_right_pos"/>
<jointpos joint="hip_hind_right" name="hip_hind_right_pos"/>
<jointpos joint="knee_hind_right" name="knee_hind_right_pos"/>
<jointvel joint="abduction_front_left" name="abduction_front_left_vel"/>
<jointvel joint="hip_front_left" name="hip_front_left_vel"/>
<jointvel joint="knee_front_left" name="knee_front_left_vel"/>
<jointvel joint="abduction_hind_left" name="abduction_hind_left_vel"/>
<jointvel joint="hip_hind_left" name="hip_hind_left_vel"/>
<jointvel joint="knee_hind_left" name="knee_hind_left_vel"/>
<jointvel joint="abduction_front_right" name="abduction_front_right_vel"/>
<jointvel joint="hip_front_right" name="hip_front_right_vel"/>
<jointvel joint="knee_front_right" name="knee_front_right_vel"/>
<jointvel joint="abduction_hind_right" name="abduction_hind_right_vel"/>
<jointvel joint="hip_hind_right" name="hip_hind_right_vel"/>
<jointvel joint="knee_hind_right" name="knee_hind_right_vel"/>
<gyro site="origin" name="gyro"/>
<accelerometer site="origin" name="accelerometer"/>
<framequat objtype="site" objname="origin" name="orientation"/>
</sensor>
<keyframe>
<key name="standing" qpos="0 0 0.21 1 0 0 0 0 0.5 1.0 0 0.5 1.0 0 0.5 1.0 0 0.5 1.0"
ctrl="0 0.5 1.0 0 0.5 1.0 0 0.5 1.0 0 0.5 1.0"/>
</keyframe>
</mujoco>
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/body.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/foot.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/handle.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/head_mount.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/joystick.gif
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 4.0 MiB

mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/lower_leg_1to1.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/powercable.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_1.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_2.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_left_3.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_1.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_2.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/barkour_v0/assets/upper_right_3.stl
BIN
View File
Binary file not shown.
mjx/mujoco/mjx/benchmark/model/humanoid/README.md
+26
View File
@@ -0,0 +1,26 @@
Humanoid
========
Degrees of Freedom: 27
Actuators: 21
This is a clone of the [MuJoCo Humanoid](https://github.com/google-deepmind/mujoco/blob/main/model/humanoid/humanoid.xml)
with the following changes:
* Solver switched to CG with 8 iterations
* Explicit contact pairs for feet and ground (compatible with
[OpenAI Gym Humanoid](https://gymnasium.farama.org/environments/mujoco/humanoid/)
environment)
This simplified humanoid model, introduced in [1], is designed for bipedal
locomotion behaviours. While several variants of it exist in the wild, this
version is based on the model in the DeepMind Control Suite [2], which has
fairly realistic actuator gains.
[1] [Synthesis and Stabilization of Complex Behaviors through Online Trajectory Optimization]
(https://doi.org/10.1109/IROS.2012.6386025).
[2] [DeepMind Control Suite](https://arxiv.org/abs/1801.00690).
![humanoid](humanoid.png)
mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 388 KiB

mjx/mujoco/mjx/benchmark/model/humanoid/humanoid.xml
+256
View File
@@ -0,0 +1,256 @@
<!-- Copyright 2021 DeepMind Technologies Limited
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
-->
<mujoco model="Humanoid">
<option timestep="0.005" solver="CG" iterations="6" ls_iterations="6">
<flag eulerdamp="disable"/>
</option>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
<global offwidth="2560" offheight="1440" elevation="-20" azimuth="120"/>
</visual>
<statistic center="0 0 0.7"/>
<asset>
<texture type="skybox" builtin="gradient" rgb1=".3 .5 .7" rgb2="0 0 0" width="32" height="512"/>
<texture name="body" type="cube" builtin="flat" mark="cross" width="128" height="128" rgb1="0.8 0.6 0.4" rgb2="0.8 0.6 0.4" markrgb="1 1 1" random="0.01"/>
<material name="body" texture="body" texuniform="true" rgba="0.8 0.6 .4 1"/>
<texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1=".1 .2 .3" rgb2=".2 .3 .4"/>
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
</asset>
<default>
<motor ctrlrange="-1 1" ctrllimited="true"/>
<default class="body">
<!-- geoms -->
<!-- TODO(robotics-simulation): support condim=1 for humanoid capsules. -->
<geom type="capsule" condim="3" friction=".7" solimp=".9 .99 .003" solref=".015 1" material="body" contype="0" conaffinity="0"/>
<default class="thigh">
<geom size=".06"/>
</default>
<default class="shin">
<geom fromto="0 0 0 0 0 -.3" size=".049"/>
</default>
<default class="foot">
<geom size=".027"/>
<default class="foot1">
<geom fromto="-.07 -.01 0 .14 -.03 0"/>
</default>
<default class="foot2">
<geom fromto="-.07 .01 0 .14 .03 0"/>
</default>
</default>
<default class="arm_upper">
<geom size=".04"/>
</default>
<default class="arm_lower">
<geom size=".031"/>
</default>
<default class="hand">
<geom type="sphere" size=".04"/>
</default>
<!-- joints -->
<joint type="hinge" damping=".2" stiffness="1" armature=".01" limited="true" solimplimit="0 .99 .01"/>
<default class="joint_big">
<joint damping="5" stiffness="10"/>
<default class="hip_x">
<joint range="-30 10"/>
</default>
<default class="hip_z">
<joint range="-60 35"/>
</default>
<default class="hip_y">
<joint axis="0 1 0" range="-150 20"/>
</default>
<default class="joint_big_stiff">
<joint stiffness="20"/>
</default>
</default>
<default class="knee">
<joint pos="0 0 .02" axis="0 -1 0" range="-160 2"/>
</default>
<default class="ankle">
<joint range="-50 50"/>
<default class="ankle_y">
<joint pos="0 0 .08" axis="0 1 0" stiffness="6"/>
</default>
<default class="ankle_x">
<joint pos="0 0 .04" stiffness="3"/>
</default>
</default>
<default class="shoulder">
<joint range="-85 60"/>
</default>
<default class="elbow">
<joint range="-100 50" stiffness="0"/>
</default>
</default>
</default>
<worldbody>
<geom name="floor" size="0 0 .05" type="plane" material="grid" condim="3"/>
<light name="spotlight" mode="targetbodycom" target="torso" diffuse=".8 .8 .8" specular="0.3 0.3 0.3" pos="0 -6 4" cutoff="30"/>
<body name="torso" pos="0 0 1.282" childclass="body">
<light name="top" pos="0 0 2" mode="trackcom"/>
<camera name="back" pos="-3 0 1" xyaxes="0 -1 0 1 0 2" mode="trackcom"/>
<camera name="side" pos="0 -3 1" xyaxes="1 0 0 0 1 2" mode="trackcom"/>
<freejoint name="root"/>
<geom name="torso" fromto="0 -.07 0 0 .07 0" size=".07"/>
<geom name="waist_upper" fromto="-.01 -.06 -.12 -.01 .06 -.12" size=".06"/>
<body name="head" pos="0 0 .19">
<geom name="head" type="sphere" size=".09"/>
<camera name="egocentric" pos=".09 0 0" xyaxes="0 -1 0 .1 0 1" fovy="80"/>
</body>
<body name="waist_lower" pos="-.01 0 -.26">
<geom name="waist_lower" fromto="0 -.06 0 0 .06 0" size=".06"/>
<joint name="abdomen_z" pos="0 0 .065" axis="0 0 1" range="-45 45" class="joint_big_stiff"/>
<joint name="abdomen_y" pos="0 0 .065" axis="0 1 0" range="-75 30" class="joint_big"/>
<body name="pelvis" pos="0 0 -.165">
<joint name="abdomen_x" pos="0 0 .1" axis="1 0 0" range="-35 35" class="joint_big"/>
<geom name="butt" fromto="-.02 -.07 0 -.02 .07 0" size=".09"/>
<body name="thigh_right" pos="0 -.1 -.04">
<joint name="hip_x_right" axis="1 0 0" class="hip_x"/>
<joint name="hip_z_right" axis="0 0 1" class="hip_z"/>
<joint name="hip_y_right" class="hip_y"/>
<geom name="thigh_right" fromto="0 0 0 0 .01 -.34" class="thigh"/>
<body name="shin_right" pos="0 .01 -.4">
<joint name="knee_right" class="knee"/>
<geom name="shin_right" class="shin"/>
<body name="foot_right" pos="0 0 -.39">
<joint name="ankle_y_right" class="ankle_y"/>
<joint name="ankle_x_right" class="ankle_x" axis="1 0 .5"/>
<geom name="foot1_right" class="foot1"/>
<geom name="foot2_right" class="foot2"/>
</body>
</body>
</body>
<body name="thigh_left" pos="0 .1 -.04">
<joint name="hip_x_left" axis="-1 0 0" class="hip_x"/>
<joint name="hip_z_left" axis="0 0 -1" class="hip_z"/>
<joint name="hip_y_left" class="hip_y"/>
<geom name="thigh_left" fromto="0 0 0 0 -.01 -.34" class="thigh"/>
<body name="shin_left" pos="0 -.01 -.4">
<joint name="knee_left" class="knee"/>
<geom name="shin_left" fromto="0 0 0 0 0 -.3" class="shin"/>
<body name="foot_left" pos="0 0 -.39">
<joint name="ankle_y_left" class="ankle_y"/>
<joint name="ankle_x_left" class="ankle_x" axis="-1 0 -.5"/>
<geom name="foot1_left" class="foot1"/>
<geom name="foot2_left" class="foot2"/>
</body>
</body>
</body>
</body>
</body>
<body name="upper_arm_right" pos="0 -.17 .06">
<joint name="shoulder1_right" axis="2 1 1" class="shoulder"/>
<joint name="shoulder2_right" axis="0 -1 1" class="shoulder"/>
<geom name="upper_arm_right" fromto="0 0 0 .16 -.16 -.16" class="arm_upper"/>
<body name="lower_arm_right" pos=".18 -.18 -.18">
<joint name="elbow_right" axis="0 -1 1" class="elbow"/>
<geom name="lower_arm_right" fromto=".01 .01 .01 .17 .17 .17" class="arm_lower"/>
<body name="hand_right" pos=".18 .18 .18">
<geom name="hand_right" zaxis="1 1 1" class="hand"/>
</body>
</body>
</body>
<body name="upper_arm_left" pos="0 .17 .06">
<joint name="shoulder1_left" axis="-2 1 -1" class="shoulder"/>
<joint name="shoulder2_left" axis="0 -1 -1" class="shoulder"/>
<geom name="upper_arm_left" fromto="0 0 0 .16 .16 -.16" class="arm_upper"/>
<body name="lower_arm_left" pos=".18 .18 -.18">
<joint name="elbow_left" axis="0 -1 -1" class="elbow"/>
<geom name="lower_arm_left" fromto=".01 -.01 .01 .17 -.17 .17" class="arm_lower"/>
<body name="hand_left" pos=".18 -.18 .18">
<geom name="hand_left" zaxis="1 -1 1" class="hand"/>
</body>
</body>
</body>
</body>
</worldbody>
<contact>
<exclude body1="waist_lower" body2="thigh_right"/>
<exclude body1="waist_lower" body2="thigh_left"/>
<pair geom1="foot1_left" geom2="floor"/>
<pair geom1="foot1_right" geom2="floor"/>
<pair geom1="foot2_left" geom2="floor"/>
<pair geom1="foot2_right" geom2="floor"/>
</contact>
<tendon>
<fixed name="hamstring_right" limited="true" range="-0.3 2">
<joint joint="hip_y_right" coef=".5"/>
<joint joint="knee_right" coef="-.5"/>
</fixed>
<fixed name="hamstring_left" limited="true" range="-0.3 2">
<joint joint="hip_y_left" coef=".5"/>
<joint joint="knee_left" coef="-.5"/>
</fixed>
</tendon>
<actuator>
<motor name="abdomen_y" gear="40" joint="abdomen_y"/>
<motor name="abdomen_z" gear="40" joint="abdomen_z"/>
<motor name="abdomen_x" gear="40" joint="abdomen_x"/>
<motor name="hip_x_right" gear="40" joint="hip_x_right"/>
<motor name="hip_z_right" gear="40" joint="hip_z_right"/>
<motor name="hip_y_right" gear="120" joint="hip_y_right"/>
<motor name="knee_right" gear="80" joint="knee_right"/>
<motor name="ankle_x_right" gear="20" joint="ankle_x_right"/>
<motor name="ankle_y_right" gear="20" joint="ankle_y_right"/>
<motor name="hip_x_left" gear="40" joint="hip_x_left"/>
<motor name="hip_z_left" gear="40" joint="hip_z_left"/>
<motor name="hip_y_left" gear="120" joint="hip_y_left"/>
<motor name="knee_left" gear="80" joint="knee_left"/>
<motor name="ankle_x_left" gear="20" joint="ankle_x_left"/>
<motor name="ankle_y_left" gear="20" joint="ankle_y_left"/>
<motor name="shoulder1_right" gear="20" joint="shoulder1_right"/>
<motor name="shoulder2_right" gear="20" joint="shoulder2_right"/>
<motor name="elbow_right" gear="40" joint="elbow_right"/>
<motor name="shoulder1_left" gear="20" joint="shoulder1_left"/>
<motor name="shoulder2_left" gear="20" joint="shoulder2_left"/>
<motor name="elbow_left" gear="40" joint="elbow_left"/>
</actuator>
<keyframe>
<!--
The values below are split into rows for readibility:
torso position
torso orientation
spinal
right leg
left leg
arms
-->
<key name="squat" qpos="0 0 0.596
0.988015 0 0.154359 0
0 0.4 0
-0.25 -0.5 -2.5 -2.65 -0.8 0.56
-0.25 -0.5 -2.5 -2.65 -0.8 0.56
0 0 0 0 0 0"/>
<key name="stand_on_left_leg" qpos="0 0 1.21948
0.971588 -0.179973 0.135318 -0.0729076
-0.0516 -0.202 0.23
-0.24 -0.007 -0.34 -1.76 -0.466 -0.0415
-0.08 -0.01 -0.37 -0.685 -0.35 -0.09
0.109 -0.067 -0.7 -0.05 0.12 0.16"/>
</keyframe>
</mujoco>
mjx/mujoco/mjx/benchmark/model/shadow_hand/README.md
+28
View File
@@ -0,0 +1,28 @@
# Shadow Hand E3M5 Description (MJCF)
Requires MuJoCo 3.0.0 or later.
## Overview
This package contains assets of the "E3M5" version of the Shadow Hand robot,
including both right-handed and left-handed versions.
The original URDF and assets were provided directly by
[Shadow Robot Company](https://www.shadowrobot.com/) under the
[Apache 2.0 License](LICENSE).
The [original Menagerie Shadow Hand](https://github.com/google-deepmind/mujoco_menagerie/tree/main/shadow_hand)
has been modified for MJX simulation in the following ways:
* Solver switched to CG with 8 iterations
* Condim switched to 3 (MJX does not yet support condim != 3)
* Object converted from ellipsoid to sphere (MJX does not yet support ellipsoids)
* Explicit contact pairs for fingers and object
* Removed some unused geoms in the arm base
<p float="left">
<img src="shadow_hand.png" width="400">
</p>
## License
These models are released under an [Apache-2.0 License](LICENSE).
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_distal_pst.obj
+10772
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_knuckle.obj
+1200
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_middle.obj
+908
View File
@@ -0,0 +1,908 @@
####
#
# OBJ File Generated by Meshlab
#
####
# Object f_middle_E3M5.obj
#
# Vertices: 224
# Faces: 444
#
####
vn -1.975821 -5.020711 -0.278529
v -5.154970 -5.651640 5.676910
vn -3.328850 -4.153035 -1.748958
v -6.985930 -7.507730 8.862470
vn -4.326096 -2.957106 -1.768247
v -8.013970 -5.977890 7.982400
vn -5.106826 -1.997097 -1.350573
v -8.480970 -4.318420 6.987890
vn -6.087360 0.080981 1.091995
v -8.500170 0.021880 5.469900
vn -6.215165 -0.637174 0.034910
v -8.500170 -3.889610 12.777900
vn -5.876741 -1.418994 0.574094
v -8.489370 -3.705560 18.597891
vn -6.241298 0.262664 -0.501847
v -8.500170 0.021880 20.537901
vn -5.906881 1.835267 -0.599729
v -8.225670 3.737700 21.843500
vn -0.396249 6.244791 0.066189
v -3.500170 7.979650 1.218900
vn -0.124787 6.279845 0.029596
v -3.500170 7.049800 24.446890
vn 0.396132 6.244804 0.066361
v 3.492960 7.979650 1.218900
vn 0.125572 6.279818 0.029003
v 3.499830 7.049800 24.446890
vn -1.131282 6.142010 -0.169950
v -4.298070 7.002770 24.018900
vn -1.731693 5.941936 0.364020
v -4.757570 7.781740 2.145900
vn -3.890530 4.808362 0.255683
v -6.731770 6.685500 2.872400
vn -3.271605 5.251853 -0.183712
v -6.311270 6.251020 22.773901
vn -5.152781 3.420898 0.498048
v -7.752570 5.534480 3.444480
vn -4.975134 3.740874 -0.357759
v -7.627570 4.939400 22.204651
vn -6.052958 1.453399 0.502421
v -8.429670 3.694290 4.187910
vn -0.352079 -5.585024 2.009694
v -3.431690 -8.644050 15.628901
vn -0.618978 -6.189351 0.349946
v -3.500170 -8.933320 12.688900
vn 0.352146 -5.584894 2.009872
v 3.424480 -8.644050 15.628901
vn 0.618977 -6.189351 0.349946
v 3.492960 -8.933320 12.688900
vn -0.617464 -6.058264 -0.743055
v -3.500170 -8.976070 9.813900
vn 0.617463 -6.058264 -0.743055
v 3.492960 -8.976070 9.813900
vn -2.386746 -4.685711 2.391089
v -6.413370 -7.450560 17.154900
vn -0.038383 -4.901343 3.203433
v -4.732770 -7.886830 17.096901
vn -3.817267 -4.439030 0.980845
v -7.464270 -4.932370 19.659380
vn -0.093681 -5.289313 0.922057
v -4.000170 -5.214450 19.768900
vn 0.094543 -5.289427 0.922396
v 3.999830 -5.214450 19.768900
vn -0.602491 -5.180288 -2.913458
v -3.500170 -8.705390 8.825910
vn -0.254401 -4.985663 -0.658851
v -4.000170 -5.823190 5.848900
vn 0.255842 -4.985855 -0.659585
v 3.999830 -5.823190 5.848900
vn 0.602489 -5.180290 -2.913456
v 3.492960 -8.705390 8.825910
vn 0.115850 -5.586018 -2.851343
v 4.487960 -5.617100 20.772900
vn -5.679802 -2.330745 0.093700
v -8.119570 -5.837210 12.861900
vn -4.337475 -4.345920 0.189490
v -7.035670 -7.457820 12.930900
vn -4.113871 -4.178817 1.073790
v -7.076570 -7.152620 16.535900
vn -5.200893 -2.776820 1.144308
v -8.027070 -5.462310 17.552050
vn -2.295639 -5.676401 0.362745
v -5.413570 -8.540680 12.828341
vn -2.040112 -5.718785 1.194047
v -5.526400 -8.058870 15.991060
vn -2.172929 -5.124251 -1.466405
v -5.669070 -8.477780 9.505900
vn -0.115949 -5.585745 -2.851956
v -4.495170 -5.617100 20.772900
vn 0.330607 -5.785292 -1.997014
v 4.750360 -6.438000 22.179899
vn 0.570670 -6.025518 0.365783
v 4.757230 -6.973210 25.881910
vn 1.253543 -5.070250 2.666592
v 4.757230 -6.489650 27.705900
vn 1.035414 -3.918195 4.338571
v 4.757230 -4.190930 30.656900
vn 0.629466 -1.818101 5.742779
v 4.757230 -2.540570 31.571890
vn 0.861828 0.841504 5.859906
v 4.757230 1.179800 31.957901
vn 0.704351 2.428302 5.571692
v 4.757230 2.983010 31.401899
vn 0.725931 4.089355 4.379475
v 4.757230 4.574110 30.386900
vn 0.706091 5.553284 2.272709
v 4.757230 6.687160 27.300900
vn 0.903443 6.016993 0.162534
v 4.757230 7.057710 25.450911
vn -0.330650 -5.785307 -1.996895
v -4.757570 -6.438000 22.179899
vn -0.567091 -6.026638 0.366088
v -4.757570 -6.973210 25.881910
vn -1.246799 -5.074042 2.668208
v -4.757570 -6.489650 27.705900
vn -1.029382 -3.919869 4.341149
v -4.757570 -4.190930 30.656900
vn -0.625103 -1.818855 5.744062
v -4.757570 -2.540570 31.571890
vn -0.856462 0.841017 5.862324
v -4.757570 1.179800 31.957901
vn -0.700617 2.429208 5.572634
v -4.757570 2.983010 31.401899
vn -0.720712 4.090838 4.380729
v -4.757570 4.574110 30.386900
vn -0.700988 5.554998 2.273383
v -4.757570 6.687160 27.300900
vn -0.894664 6.020987 0.169189
v -4.757570 7.057710 25.450911
vn -2.156509 -5.429175 -0.915722
v -5.905570 -6.467490 22.900900
vn -2.571520 -5.055631 0.909265
v -5.905570 -6.714300 26.079901
vn -2.760185 -4.429241 3.436599
v -5.878870 -4.670700 29.331690
vn -2.517030 -2.450455 4.619818
v -5.905570 -3.059300 31.089901
vn -2.229681 -0.072035 5.230337
v -5.905570 0.043750 31.825911
vn -2.501201 2.299577 4.759172
v -5.905570 3.142020 31.069891
vn -2.123583 4.367735 2.980202
v -5.905570 5.558400 28.988899
vn -2.601006 5.117285 0.793377
v -5.905570 6.764830 26.036909
vn -5.776209 -1.678740 -0.670729
v -8.773640 -3.349110 22.568850
vn -5.926663 -1.590920 -0.313608
v -8.450030 -3.339630 20.409321
vn -5.447601 -2.033954 0.207273
v -8.999960 -3.586780 25.822010
vn -5.708444 -1.232529 1.251641
v -8.999960 -2.605590 27.581030
vn -4.415786 0.019221 4.171239
v -8.121470 0.021880 29.932899
vn -5.311327 1.758746 1.267951
v -8.999960 3.114340 27.469900
vn -4.415588 4.178071 0.156447
v -8.121470 4.950330 25.050900
vn -5.467701 0.005860 2.047906
v -9.000000 0.021880 28.726141
vn -5.219584 2.176666 -0.081094
v -8.999960 3.878120 25.230341
vn -6.177143 0.179311 -0.737004
v -9.000190 0.026420 23.152901
vn -3.502934 -4.525011 1.642803
v -6.878870 -6.812390 2.176910
vn -3.251389 -4.444415 -0.519670
v -6.523960 -4.865100 5.636130
vn -1.317918 -5.210125 2.722004
v -4.757570 -7.077680 3.798900
vn -1.134610 -5.750695 1.447569
v -4.757570 -7.783830 1.971900
vn -3.258571 -4.934119 -1.538218
v -6.878870 -6.538240 -2.893100
vn -1.182014 -5.905304 -0.021934
v -4.757570 -8.028080 0.028890
vn -1.141683 -5.599188 -1.603683
v -4.757570 -7.796000 -1.915110
vn -2.804239 -4.424251 -2.808778
v -6.878870 -5.854020 -4.573890
vn -0.643679 -4.786737 -3.685352
v -4.757570 -5.985130 -5.355090
vn -2.502592 -2.738398 -4.286158
v -6.405570 -4.618150 -6.293100
vn -0.439400 -3.443835 -5.120736
v -4.757570 -4.513470 -6.647100
vn -2.832245 -0.747272 -5.403884
v -5.905570 -1.568680 -7.337640
vn -0.973778 -1.644752 -5.681370
v -4.757570 -2.773520 -7.545100
vn -2.571455 0.861364 -5.150094
v -5.905570 1.770760 -7.620090
vn -0.817049 0.297736 -5.966962
v -4.757570 1.090100 -7.975100
vn -0.813927 2.304698 -5.613643
v -4.757570 2.984900 -7.481100
vn -2.542195 3.373481 -4.109332
v -5.905570 4.791360 -6.195100
vn -0.623470 3.519813 -5.012120
v -4.757570 4.704470 -6.544100
vn -3.712195 2.427079 -4.191787
v -6.878870 3.669480 -6.170100
vn -2.535884 4.738874 -2.377094
v -5.905570 6.942470 -3.641100
vn -0.838900 4.616839 -3.927192
v -4.757570 6.146920 -5.220100
vn -0.608599 5.500813 -2.702369
v -4.757570 7.226940 -3.586100
vn -4.538651 3.600815 0.213802
v -8.121470 6.008870 0.104330
vn -2.590973 5.394011 -0.119302
v -5.905570 7.634140 -0.422090
vn -1.036419 5.831363 -1.030521
v -4.757570 7.880620 -1.740100
vn 1.323168 -5.206586 2.722369
v 4.757230 -7.077680 3.798900
vn 1.137741 -5.749225 1.447183
v 4.757230 -7.783830 1.971900
vn 1.185190 -5.903691 -0.022018
v 4.757230 -8.028080 0.028890
vn 1.144988 -5.597528 -1.603007
v 4.757230 -7.796000 -1.915110
vn 0.646070 -4.786094 -3.684909
v 4.757230 -5.985130 -5.355090
vn 0.441208 -3.443483 -5.120489
v 4.757230 -4.513470 -6.647100
vn 0.978876 -1.643660 -5.679229
v 4.757230 -2.773520 -7.545100
vn 0.821955 0.297028 -5.965026
v 4.757230 1.090100 -7.975100
vn 0.819038 2.304303 -5.611849
v 4.757230 2.984900 -7.481100
vn 0.627055 3.518999 -5.011438
v 4.757230 4.704470 -6.544100
vn 0.844235 4.615387 -3.925911
v 4.757230 6.146920 -5.220100
vn 0.612092 5.500012 -2.701771
v 4.757230 7.226940 -3.586100
vn 1.042403 5.828549 -1.029587
v 4.757230 7.880620 -1.740100
vn -5.081741 -3.169341 1.411239
v -7.878870 -5.407860 2.463000
vn -5.431445 -2.324418 0.866331
v -8.500000 -3.558340 4.341910
vn -5.965299 -1.291361 -0.130787
v -9.000000 -3.599130 -1.173100
vn -6.021783 -0.734312 -0.996979
v -9.000000 -2.216030 -3.076100
vn -5.820709 0.061761 1.471896
v -9.000000 0.021880 4.392910
vn -5.524968 -0.008542 -1.914781
v -9.000000 0.021880 -4.580100
vn -6.003400 0.789261 -1.010776
v -9.000000 2.259790 -3.076100
vn -5.948116 1.307274 -0.199703
v -9.000000 3.744800 -1.099100
vn -4.681353 0.881030 -3.653100
v -8.121470 1.636160 -5.701100
vn -4.962337 -3.308012 0.017600
v -8.121470 -5.891000 -0.429100
vn -3.386144 -4.944589 0.183884
v -6.878870 -7.133550 0.480910
vn -4.989554 -3.416844 -0.858124
v -8.121470 -5.600890 -1.876100
vn -4.699954 -2.540347 -2.176459
v -8.121000 -4.852410 -4.186100
vn -4.528942 -0.789736 -3.801926
v -8.121470 -1.442780 -5.741090
vn -4.694910 3.473782 -1.227291
v -8.121000 5.473260 -2.536040
vn -4.507683 2.644774 -2.484791
v -8.121000 4.618140 -4.139520
vn -3.540627 -4.160904 2.684636
v -6.754660 -5.905320 3.906520
vn -4.977641 -2.524741 -0.660096
v -8.458230 -3.706150 5.688800
vn -4.119967 -3.565147 2.419044
v -8.121000 -3.838170 28.386730
vn -4.472468 -4.094103 0.578559
v -8.121000 -4.784040 25.948540
vn -4.160048 -2.325684 3.353160
v -8.121000 -3.164870 29.143299
vn -4.727426 -3.959540 -0.401267
v -7.969330 -4.728900 22.715111
vn -4.301860 3.896197 2.071576
v -8.121000 3.997300 28.018669
vn -4.324967 2.223240 3.423124
v -8.121000 3.127030 29.120550
vn 4.320310 2.227863 3.430520
v 8.113790 3.127030 29.120550
vn 4.294984 3.904343 2.075937
v 8.113790 3.997300 28.018669
vn 4.727426 -3.959540 -0.401266
v 7.962120 -4.728900 22.715111
vn 4.156006 -2.330211 3.360267
v 8.113790 -3.164880 29.143299
vn 4.467521 -4.102084 0.578252
v 8.113790 -4.784040 25.948540
vn 4.114889 -3.572841 2.424054
v 8.113790 -3.838170 28.386730
vn 4.977641 -2.524741 -0.660094
v 8.451020 -3.706150 5.688800
vn 3.540596 -4.160900 2.684620
v 6.747460 -5.905330 3.906520
vn 4.505536 2.649572 -2.489471
v 8.113790 4.618140 -4.139520
vn 4.691979 3.481025 -1.230083
v 8.113790 5.473260 -2.536040
vn 4.525122 -0.792267 -3.809892
v 8.114260 -1.442780 -5.741100
vn 4.698267 -2.544770 -2.180376
v 8.113790 -4.852410 -4.186100
vn 4.986269 -3.424754 -0.860209
v 8.114260 -5.600890 -1.876110
vn 3.389770 -4.943439 0.183984
v 6.871660 -7.133550 0.480900
vn 4.960076 -3.315449 0.017590
v 8.114260 -5.891000 -0.429100
vn 4.677627 0.883812 -3.660927
v 8.114260 1.636160 -5.701100
vn 5.943124 1.316251 -0.200224
v 9.000000 3.744800 -1.099100
vn 5.999335 0.794810 -1.017667
v 9.000000 2.259790 -3.076100
vn 5.515849 -0.008605 -1.924431
v 9.000000 0.021880 -4.580100
vn 5.809621 0.061768 1.487882
v 9.000000 0.021880 4.392910
vn 6.017956 -0.739562 -1.003839
v 9.000000 -2.216030 -3.076100
vn 5.960750 -1.299910 -0.131566
v 9.000000 -3.599130 -1.173100
vn 5.430431 -2.329072 0.871211
v 8.492790 -3.558340 4.341910
vn 5.080068 -3.172664 1.413554
v 7.871660 -5.407860 2.463000
vn 2.592579 5.393801 -0.117772
v 5.898360 7.634140 -0.422090
vn 4.536738 3.606666 0.215079
v 8.114260 6.008870 0.104320
vn 2.539719 4.738473 -2.377132
v 5.898360 6.942470 -3.641100
vn 3.712209 2.427073 -4.191782
v 6.871660 3.669480 -6.170100
vn 2.545941 3.373343 -4.108975
v 5.898360 4.791360 -6.195110
vn 2.575225 0.861660 -5.149790
v 5.898360 1.770760 -7.620090
vn 2.837510 -0.746383 -5.402541
v 5.898360 -1.568680 -7.337640
vn 2.504754 -2.738315 -4.286266
v 6.398360 -4.618150 -6.293100
vn 2.806519 -4.423785 -2.808745
v 6.871660 -5.854020 -4.573890
vn 3.261961 -4.932928 -1.538092
v 6.871660 -6.538240 -2.893100
vn 3.251392 -4.444405 -0.519667
v 6.516750 -4.865100 5.636130
vn 3.506286 -4.523795 1.642852
v 6.871660 -6.812390 2.176900
vn 6.174293 0.177592 -0.747126
v 9.000190 0.026420 23.152901
vn 5.208213 2.185383 -0.082852
v 8.999960 3.878120 25.230341
vn 5.458268 0.005854 2.057881
v 9.000000 0.021880 28.726141
vn 4.409956 4.186586 0.156118
v 8.114260 4.950330 25.050900
vn 5.300780 1.765974 1.274092
v 8.999960 3.114340 27.469900
vn 4.410225 0.019278 4.180248
v 8.114260 0.021880 29.932899
vn 5.700973 -1.239574 1.259032
v 8.999960 -2.605590 27.581030
vn 5.436850 -2.044216 0.204676
v 8.999960 -3.586780 25.822010
vn 5.926663 -1.590920 -0.313609
v 8.442820 -3.339630 20.409321
vn 5.775527 -1.684477 -0.676935
v 8.766430 -3.349110 22.568850
vn 2.604061 5.117111 0.792504
v 5.898360 6.764830 26.036909
vn 2.123708 4.367924 2.980122
v 5.898360 5.558400 28.988899
vn 2.505037 2.299492 4.758895
v 5.898360 3.142020 31.069891
vn 2.230662 -0.072842 5.230432
v 5.898360 0.043740 31.825911
vn 2.520874 -2.450161 4.619775
v 5.898360 -3.059300 31.089901
vn 2.766455 -4.427368 3.435123
v 5.871660 -4.670700 29.331690
vn 2.575577 -5.055537 0.908791
v 5.898360 -6.714300 26.079901
vn 2.156810 -5.429112 -0.916056
v 5.898360 -6.467490 22.900900
vn 2.172924 -5.124258 -1.466401
v 5.661860 -8.477780 9.505900
vn 2.040112 -5.718785 1.194047
v 5.519190 -8.058870 15.991060
vn 2.295639 -5.676401 0.362745
v 5.406360 -8.540680 12.828341
vn 5.200893 -2.776819 1.144308
v 8.019860 -5.462310 17.552050
vn 4.113871 -4.178818 1.073790
v 7.069360 -7.152620 16.535900
vn 4.337474 -4.345921 0.189494
v 7.028460 -7.457820 12.930900
vn 5.679803 -2.330744 0.093702
v 8.112360 -5.837210 12.861900
vn 3.817452 -4.438978 0.980867
v 7.457060 -4.932370 19.659380
vn 0.036775 -4.900473 3.203675
v 4.725560 -7.886830 17.096901
vn 2.386746 -4.685711 2.391089
v 6.406160 -7.450560 17.154900
vn 6.051692 1.458310 0.507014
v 8.422460 3.694290 4.187910
vn 4.975134 3.740875 -0.357759
v 7.620360 4.939400 22.204651
vn 5.152782 3.420897 0.498046
v 7.745360 5.534480 3.444480
vn 3.271605 5.251853 -0.183712
v 6.304060 6.251020 22.773901
vn 3.890531 4.808362 0.255683
v 6.724560 6.685500 2.872400
vn 1.731693 5.941936 0.364020
v 4.750360 7.781740 2.145900
vn 1.132973 6.141824 -0.170517
v 4.290860 7.002770 24.018900
vn 5.905301 1.839049 -0.604558
v 8.218460 3.737700 21.843500
vn 6.240402 0.262129 -0.507810
v 8.492960 0.021880 20.537901
vn 5.876742 -1.418994 0.574095
v 8.482160 -3.705560 18.597891
vn 6.215165 -0.637174 0.034910
v 8.492960 -3.889610 12.777900
vn 6.082349 0.081128 1.105070
v 8.492960 0.021880 5.469900
vn 5.106826 -1.997098 -1.350573
v 8.473760 -4.318420 6.987890
vn 4.326105 -2.957100 -1.768250
v 8.006760 -5.977890 7.982400
vn 3.328853 -4.153027 -1.748949
v 6.978730 -7.507730 8.862470
vn 1.978097 -5.018886 -0.276658
v 5.147760 -5.651640 5.676900
# 224 vertices, 0 vertices normals
f 1//1 2//2 3//3
f 5//5 4//4 6//6
f 5//5 6//6 7//7
f 5//5 7//7 8//8
f 9//9 5//5 8//8
f 10//10 11//11 12//12
f 13//13 12//12 11//11
f 14//14 11//11 10//10
f 14//14 10//10 15//15
f 15//15 16//16 14//14
f 17//17 14//14 16//16
f 16//16 18//18 17//17
f 19//19 17//17 18//18
f 9//9 19//19 18//18
f 18//18 20//20 9//9
f 20//20 5//5 9//9
f 21//21 22//22 23//23
f 24//24 23//23 22//22
f 22//22 25//25 24//24
f 26//26 24//24 25//25
f 27//27 28//28 29//29
f 30//30 29//29 28//28
f 28//28 21//21 30//30
f 30//30 21//21 23//23
f 31//31 30//30 23//23
f 2//2 1//1 32//32
f 1//1 33//33 32//32
f 34//34 35//35 32//32
f 33//33 34//34 32//32
f 6//6 4//4 3//3
f 37//37 6//6 3//3
f 37//37 3//3 2//2
f 38//38 37//37 2//2
f 41//41 42//42 39//39
f 38//38 41//41 39//39
f 38//38 2//2 43//43
f 41//41 38//38 43//43
f 43//43 25//25 41//41
f 21//21 42//42 41//41
f 41//41 22//22 21//21
f 41//41 25//25 22//22
f 36//36 44//44 30//30
f 44//44 45//45 55//55
f 45//45 46//46 55//55
f 56//56 55//55 46//46
f 46//46 47//47 56//56
f 57//57 56//56 47//47
f 47//47 48//48 57//57
f 58//58 57//57 48//48
f 48//48 49//49 58//58
f 59//59 58//58 49//49
f 49//49 50//50 59//59
f 60//60 59//59 50//50
f 50//50 51//51 60//60
f 61//61 60//60 51//51
f 51//51 52//52 61//61
f 62//62 61//61 52//52
f 52//52 53//53 62//62
f 63//63 62//62 53//53
f 53//53 54//54 63//63
f 64//64 63//63 54//54
f 13//13 64//64 54//54
f 64//64 13//13 11//11
f 64//64 11//11 14//14
f 29//29 44//44 55//55
f 65//65 29//29 55//55
f 65//65 55//55 56//56
f 66//66 65//65 56//56
f 56//56 57//57 66//66
f 67//67 66//66 57//57
f 67//67 57//57 58//58
f 68//68 67//67 58//58
f 58//58 59//59 68//68
f 69//69 68//68 59//59
f 59//59 60//60 69//69
f 69//69 60//60 61//61
f 70//70 69//69 61//61
f 61//61 62//62 70//70
f 71//71 70//70 62//62
f 62//62 63//63 71//71
f 72//72 71//71 63//63
f 63//63 64//64 72//72
f 72//72 64//64 14//14
f 72//72 14//14 17//17
f 73//73 74//74 29//29
f 68//68 69//69 77//77
f 69//69 70//70 77//77
f 71//71 72//72 79//79
f 72//72 17//17 79//79
f 19//19 79//79 17//17
f 81//81 78//78 79//79
f 9//9 81//81 79//79
f 79//79 19//19 9//9
f 9//9 8//8 82//82
f 9//9 82//82 81//81
f 76//76 80//80 82//82
f 75//75 76//76 82//82
f 73//73 8//8 74//74
f 8//8 73//73 82//82
f 75//75 82//82 73//73
f 86//86 85//85 83//83
f 90//90 91//91 89//89
f 87//87 90//90 89//89
f 90//90 92//92 91//91
f 93//93 91//91 92//92
f 98//98 97//97 96//96
f 99//99 100//100 98//98
f 96//96 99//99 98//98
f 99//99 96//96 101//101
f 101//101 102//102 99//99
f 103//103 100//100 99//99
f 99//99 102//102 103//103
f 104//104 103//103 102//102
f 106//106 102//102 105//105
f 10//10 107//107 106//106
f 35//35 26//26 32//32
f 85//85 34//34 33//33
f 34//34 85//85 108//108
f 85//85 86//86 108//108
f 109//109 108//108 86//86
f 86//86 88//88 109//109
f 110//110 109//109 88//88
f 88//88 89//89 110//110
f 111//111 110//110 89//89
f 89//89 91//91 111//111
f 112//112 111//111 91//91
f 91//91 93//93 112//112
f 113//113 112//112 93//93
f 93//93 95//95 113//113
f 114//114 113//113 95//95
f 95//95 97//97 114//114
f 115//115 114//114 97//97
f 97//97 98//98 115//115
f 116//116 115//115 98//98
f 98//98 100//100 116//116
f 117//117 116//116 100//100
f 100//100 103//103 117//117
f 118//118 117//117 103//103
f 103//103 104//104 118//118
f 119//119 118//118 104//104
f 104//104 107//107 119//119
f 120//120 119//119 107//107
f 107//107 10//10 120//120
f 12//12 120//120 10//10
f 2//2 32//32 43//43
f 121//121 122//122 123//123
f 124//124 123//123 125//125
f 126//126 124//124 125//125
f 25//25 32//32 26//26
f 127//127 129//129 126//126
f 105//105 20//20 18//18
f 105//105 128//128 20//20
f 20//20 125//125 5//5
f 131//131 83//83 130//130
f 87//87 131//131 130//130
f 130//130 132//132 87//87
f 132//132 133//133 87//87
f 90//90 87//87 133//133
f 92//92 90//90 133//133
f 94//94 92//92 134//134
f 134//134 129//129 94//94
f 96//96 94//94 129//129
f 101//101 96//96 129//129
f 105//105 18//18 106//106
f 16//16 106//106 18//18
f 25//25 43//43 32//32
f 121//121 123//123 130//130
f 132//132 130//130 123//123
f 123//123 133//133 132//132
f 123//123 124//124 133//133
f 124//124 126//126 134//134
f 129//129 134//134 126//126
f 7//7 74//74 8//8
f 7//7 29//29 74//74
f 31//31 36//36 30//30
f 7//7 40//40 29//29
f 27//27 29//29 40//40
f 40//40 39//39 27//27
f 42//42 27//27 39//39
f 28//28 27//27 42//42
f 42//42 21//21 28//28
f 30//30 44//44 29//29
f 87//87 89//89 131//131
f 88//88 131//131 89//89
f 86//86 131//131 88//88
f 83//83 131//131 86//86
f 92//92 133//133 134//134
f 124//124 134//134 133//133
f 83//83 121//121 130//130
f 125//125 127//127 126//126
f 125//125 128//128 127//127
f 20//20 128//128 125//125
f 5//5 125//125 122//122
f 122//122 125//125 123//123
f 127//127 136//136 129//129
f 136//136 101//101 129//129
f 101//101 136//136 102//102
f 135//135 102//102 136//136
f 102//102 135//135 105//105
f 135//135 128//128 105//105
f 6//6 40//40 7//7
f 6//6 37//37 40//40
f 37//37 39//39 40//40
f 37//37 38//38 39//39
f 10//10 106//106 15//15
f 15//15 106//106 16//16
f 3//3 4//4 84//84
f 1//1 3//3 84//84
f 137//137 1//1 84//84
f 83//83 1//1 137//137
f 83//83 137//137 121//121
f 137//137 84//84 121//121
f 84//84 122//122 121//121
f 122//122 138//138 5//5
f 138//138 4//4 5//5
f 4//4 138//138 84//84
f 138//138 122//122 84//84
f 1//1 83//83 85//85
f 1//1 85//85 33//33
f 78//78 81//81 82//82
f 78//78 82//82 80//80
f 65//65 142//142 29//29
f 142//142 73//73 29//29
f 140//140 73//73 142//142
f 139//139 75//75 140//140
f 75//75 139//139 76//76
f 141//141 76//76 139//139
f 141//141 68//68 77//77
f 70//70 144//144 77//77
f 143//143 78//78 144//144
f 78//78 143//143 79//79
f 143//143 71//71 79//79
f 71//71 143//143 144//144
f 71//71 144//144 70//70
f 68//68 141//141 67//67
f 141//141 139//139 67//67
f 67//67 139//139 66//66
f 139//139 140//140 66//66
f 73//73 140//140 75//75
f 66//66 140//140 65//65
f 140//140 142//142 65//65
f 128//128 135//135 136//136
f 128//128 136//136 127//127
f 76//76 141//141 80//80
f 77//77 80//80 141//141
f 77//77 144//144 80//80
f 78//78 80//80 144//144
f 94//94 96//96 97//97
f 94//94 97//97 95//95
f 92//92 94//94 95//95
f 92//92 95//95 93//93
f 102//102 107//107 104//104
f 102//102 106//106 107//107
f 171//171 120//120 169//169
f 171//171 119//119 120//120
f 176//176 113//113 114//114
f 176//176 114//114 175//175
f 175//175 114//114 115//115
f 175//175 115//115 174//174
f 185//185 145//145 183//183
f 186//186 183//183 145//145
f 186//186 148//148 183//183
f 187//187 183//183 148//148
f 161//161 162//162 153//153
f 161//161 153//153 154//154
f 149//149 198//198 147//147
f 197//197 198//198 149//149
f 190//190 188//188 149//149
f 150//150 197//197 149//149
f 196//196 197//197 150//150
f 148//148 196//196 150//150
f 195//195 196//196 148//148
f 192//192 193//193 145//145
f 192//192 145//145 146//146
f 146//146 184//184 192//192
f 185//185 184//184 146//146
f 146//146 145//145 185//185
f 193//193 186//186 145//145
f 148//148 186//186 195//195
f 148//148 150//150 187//187
f 188//188 187//187 150//150
f 150//150 149//149 188//188
f 149//149 147//147 190//190
f 147//147 206//206 190//190
f 198//198 206//206 147//147
f 185//185 183//183 181//181
f 185//185 181//181 182//182
f 224//224 34//34 108//108
f 224//224 108//108 180//180
f 151//151 179//179 167//167
f 221//221 179//179 151//151
f 151//151 220//220 221//221
f 167//167 220//220 151//151
f 179//179 168//168 167//167
f 152//152 168//168 179//179
f 180//180 168//168 152//152
f 152//152 224//224 180//180
f 179//179 224//224 152//152
f 224//224 179//179 222//222
f 222//222 179//179 221//221
f 214//214 213//213 169//169
f 12//12 214//214 169//169
f 205//205 203//203 204//204
f 205//205 202//202 203//203
f 219//219 202//202 205//205
f 219//219 218//218 202//202
f 154//154 170//170 161//161
f 171//171 170//170 154//154
f 154//154 153//153 171//171
f 172//172 171//171 153//153
f 153//153 160//160 172//172
f 162//162 160//160 153//153
f 167//167 166//166 164//164
f 220//220 167//167 164//164
f 209//209 164//164 161//161
f 164//164 162//162 161//161
f 164//164 163//163 162//162
f 180//180 159//159 168//168
f 165//165 156//156 155//155
f 176//176 155//155 156//156
f 180//180 109//109 158//158
f 109//109 110//110 158//158
f 110//110 111//111 158//158
f 178//178 158//158 111//111
f 31//31 206//206 36//36
f 200//200 207//207 23//23
f 207//207 200//200 208//208
f 200//200 203//203 208//208
f 202//202 208//208 203//203
f 208//208 202//202 206//206
f 218//218 206//206 202//202
f 218//218 189//189 206//206
f 218//218 217//217 189//189
f 160//160 163//163 155//155
f 165//165 155//155 163//163
f 166//166 156//156 165//165
f 166//166 157//157 156//156
f 157//157 166//166 159//159
f 168//168 159//159 166//166
f 26//26 35//35 199//199
f 213//213 211//211 169//169
f 170//170 169//169 211//211
f 172//172 160//160 174//174
f 174//174 160//160 175//175
f 155//155 175//175 160//160
f 175//175 155//155 176//176
f 176//176 156//156 177//177
f 177//177 156//156 178//178
f 157//157 178//178 156//156
f 159//159 178//178 157//157
f 178//178 159//159 158//158
f 158//158 159//159 180//180
f 209//209 220//220 164//164
f 170//170 209//209 161//161
f 170//170 211//211 209//209
f 162//162 163//163 160//160
f 163//163 164//164 165//165
f 165//165 164//164 166//166
f 168//168 166//166 167//167
f 223//223 199//199 35//35
f 12//12 169//169 120//120
f 169//169 170//170 171//171
f 119//119 171//171 118//118
f 173//173 118//118 171//171
f 118//118 173//173 117//117
f 172//172 173//173 171//171
f 173//173 172//172 174//174
f 174//174 116//116 173//173
f 173//173 116//116 117//117
f 116//116 174//174 115//115
f 113//113 176//176 112//112
f 177//177 112//112 176//176
f 178//178 111//111 177//177
f 177//177 111//111 112//112
f 109//109 180//180 108//108
f 188//188 190//190 181//181
f 217//217 181//181 190//190
f 190//190 189//189 217//217
f 188//188 181//181 187//187
f 187//187 181//181 183//183
f 216//216 182//182 181//181
f 216//216 181//181 217//217
f 184//184 216//216 210//210
f 216//216 184//184 182//182
f 182//182 184//184 185//185
f 210//210 212//212 184//184
f 191//191 184//184 212//212
f 192//192 184//184 191//191
f 194//194 186//186 193//193
f 195//195 186//186 194//194
f 190//190 206//206 189//189
f 191//191 212//212 215//215
f 191//191 215//215 54//54
f 53//53 191//191 54//54
f 191//191 53//53 192//192
f 52//52 192//192 53//53
f 192//192 52//52 193//193
f 51//51 193//193 52//52
f 193//193 51//51 194//194
f 194//194 51//51 50//50
f 49//49 194//194 50//50
f 194//194 49//49 195//195
f 48//48 195//195 49//49
f 195//195 48//48 196//196
f 196//196 48//48 47//47
f 196//196 47//47 197//197
f 46//46 197//197 47//47
f 197//197 46//46 198//198
f 198//198 46//46 45//45
f 198//198 45//45 206//206
f 206//206 45//45 36//36
f 54//54 215//215 13//13
f 201//201 24//24 26//26
f 201//201 23//23 24//24
f 23//23 201//201 200//200
f 199//199 201//201 26//26
f 201//201 199//199 204//204
f 204//204 199//199 223//223
f 204//204 203//203 201//201
f 201//201 203//203 200//200
f 204//204 223//223 205//205
f 205//205 223//223 222//222
f 205//205 222//222 219//219
f 219//219 222//222 221//221
f 224//224 35//35 34//34
f 35//35 224//224 223//223
f 207//207 31//31 23//23
f 31//31 207//207 206//206
f 208//208 206//206 207//207
f 209//209 216//216 220//220
f 211//211 216//216 209//209
f 216//216 211//211 210//210
f 210//210 211//211 212//212
f 213//213 212//212 211//211
f 212//212 213//213 215//215
f 214//214 215//215 213//213
f 215//215 214//214 12//12
f 215//215 12//12 13//13
f 216//216 217//217 220//220
f 220//220 217//217 218//218
f 220//220 218//218 219//219
f 220//220 219//219 221//221
f 224//224 222//222 223//223
f 36//36 45//45 44//44
# 444 faces, 0 coords texture
# End of File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/f_proximal.obj
+2206
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_0.obj
+2220
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_1.obj
+37554
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/forearm_collision.obj
+924
View File
@@ -0,0 +1,924 @@
####
#
# OBJ File Generated by Meshlab
#
####
# Object forearm_collision_E3M5.obj
#
# Vertices: 228
# Faces: 452
#
####
vn -1.014278 -5.425926 2.002443
v -13.116530 -65.924423 164.889297
vn -1.445568 -5.080638 1.266666
v -20.459780 -71.908623 156.506897
vn -0.487387 -5.259754 1.266666
v -6.898247 -74.443710 156.506897
vn 0.487387 -5.259754 1.266666
v 6.898201 -74.443710 156.506897
vn 0.000000 -5.519899 2.001652
v -0.000023 -67.212082 164.889297
vn -0.516356 -5.566404 2.105098
v -6.228134 -67.212082 164.889297
vn 0.000000 -6.200466 -0.708931
v -0.000023 -67.212082 125.617798
vn -1.139329 -6.094893 -0.709044
v -13.116530 -65.924423 125.617798
vn -0.573518 -6.188871 -0.719127
v -6.228134 -67.212082 125.617798
vn 1.445567 -5.080638 1.266666
v 20.459740 -71.908623 156.506897
vn 1.014279 -5.425926 2.002443
v 13.116490 -65.924423 164.889297
vn 0.516356 -5.566404 2.105098
v 6.228096 -67.212082 164.889297
vn 1.139330 -6.094893 -0.709044
v 13.116490 -65.924423 125.617798
vn 0.573518 -6.188871 -0.719127
v 6.228096 -67.212082 125.617798
vn -1.994041 -5.147196 2.004810
v -25.733641 -62.110180 164.889297
vn -2.354518 -4.728508 1.266667
v -33.324581 -66.924767 156.506897
vn -1.531405 -5.376383 2.105857
v -18.472281 -64.923264 164.889297
vn -2.239874 -5.781761 -0.709382
v -25.733641 -62.110180 125.617798
vn -1.701028 -5.978075 -0.719336
v -18.472281 -64.923264 125.617798
vn 2.354519 -4.728508 1.266666
v 33.324539 -66.924767 156.506897
vn 1.994040 -5.147196 2.004810
v 25.733601 -62.110180 164.889297
vn 1.531405 -5.376382 2.105858
v 18.472231 -64.923264 164.889297
vn 2.239873 -5.781761 -0.709382
v 25.733601 -62.110180 125.617798
vn 1.701028 -5.978075 -0.719336
v 18.472231 -64.923264 125.617798
vn -0.610107 -3.066856 -2.952755
v -13.168620 -66.203033 -0.000004
vn 0.000000 -3.126960 -2.952861
v -0.000023 -67.500031 -0.000004
vn 0.610107 -3.066856 -2.952755
v 13.168570 -66.203033 -0.000004
vn 0.000000 0.000000 -6.283185
v -0.000023 -0.000023 -0.000004
vn -0.610402 -3.069248 2.978756
v -13.168620 -66.203033 182.800003
vn 0.000000 -3.129395 2.979192
v -0.000023 -67.500031 182.800003
vn 0.610402 -3.069248 2.978756
v 13.168570 -66.203033 182.800003
vn 0.000000 0.000000 6.283185
v -0.000023 -0.000023 182.800003
vn -2.905908 -4.693203 2.008757
v -37.370762 -55.913849 164.889297
vn -3.183289 -4.215355 1.266666
v -45.054562 -59.661869 156.506897
vn -2.494045 -5.002877 2.107387
v -30.087360 -60.423538 164.889297
vn -3.264125 -5.271745 -0.709945
v -37.370762 -55.913849 125.617798
vn -2.770584 -5.563671 -0.719755
v -30.087360 -60.423538 125.617798
vn 3.183291 -4.215355 1.266666
v 45.054508 -59.661869 156.506897
vn 2.905909 -4.693203 2.008757
v 37.370720 -55.913849 164.889297
vn 2.494045 -5.002877 2.107387
v 30.087320 -60.423538 164.889297
vn 3.264126 -5.271745 -0.709945
v 37.370720 -55.913849 125.617798
vn 2.770584 -5.563670 -0.719755
v 30.087320 -60.423538 125.617798
vn -1.196754 -2.888856 -2.952438
v -25.831150 -62.361889 -0.000004
vn 1.196754 -2.888856 -2.952438
v 25.831110 -62.361889 -0.000004
vn -1.197365 -2.891101 2.977438
v -25.831150 -62.361889 182.800003
vn 1.197364 -2.891102 2.977438
v 25.831110 -62.361889 182.800003
vn -1.737385 -2.599806 -2.951908
v -37.501011 -56.124222 -0.000004
vn 1.737386 -2.599805 -2.951907
v 37.500969 -56.124222 -0.000004
vn -1.738327 -2.601760 2.975215
v -37.501011 -56.124222 182.800003
vn 1.738327 -2.601760 2.975215
v 37.500969 -56.124222 182.800003
vn -3.718858 -4.079383 2.014284
v -47.584270 -47.570160 164.889297
vn -3.903659 -3.558654 1.266666
v -55.250252 -50.367260 156.506897
vn -3.371383 -4.458741 2.109718
v -40.677860 -53.866180 164.889297
vn -4.177236 -4.582198 -0.710733
v -47.584270 -47.570160 125.617798
vn -3.745765 -4.959772 -0.720382
v -40.677860 -53.866180 125.617798
vn 3.903658 -3.558657 1.266666
v 55.250198 -50.367260 156.506897
vn 3.718856 -4.079385 2.014282
v 47.584221 -47.570160 164.889297
vn 3.371385 -4.458741 2.109717
v 40.677818 -53.866180 164.889297
vn 4.177233 -4.582200 -0.710732
v 47.584221 -47.570160 125.617798
vn 3.745766 -4.959771 -0.720382
v 40.677818 -53.866180 125.617798
vn -2.211215 -2.210821 -2.951167
v -47.729729 -47.729729 -0.000004
vn 2.211215 -2.210821 -2.951168
v 47.729679 -47.729729 -0.000004
vn -2.212471 -2.212282 2.972061
v -47.729729 -47.729729 182.800003
vn 2.212471 -2.212283 2.972062
v 47.729679 -47.729729 182.800003
vn -4.405183 -3.326635 2.021387
v -55.984241 -37.395302 164.889297
vn -4.491092 -2.780767 1.266666
v -63.564468 -39.357460 156.506897
vn -4.133474 -3.762670 2.112887
v -49.883121 -45.474468 164.889297
vn -4.948087 -3.736616 -0.711746
v -55.984241 -37.395302 125.617798
vn -4.593356 -4.186948 -0.721219
v -49.883121 -45.474468 125.617798
vn 4.491090 -2.780766 1.266665
v 63.564430 -39.357460 156.506897
vn 4.405181 -3.326638 2.021386
v 55.984200 -37.395302 164.889297
vn 4.133472 -3.762674 2.112886
v 49.883080 -45.474468 164.889297
vn 4.948084 -3.736619 -0.711745
v 55.984200 -37.395302 125.617798
vn 4.593353 -4.186952 -0.721218
v 49.883080 -45.474468 125.617798
vn -2.600025 -1.736860 -2.950213
v -56.124222 -37.501011 -0.000004
vn 2.600025 -1.736861 -2.950213
v 56.124180 -37.501011 -0.000004
vn -2.601500 -1.737581 2.967929
v -56.124222 -37.501011 182.800003
vn 2.601499 -1.737581 2.967930
v 56.124180 -37.501011 182.800003
vn -4.941505 -2.460578 2.030066
v -62.249168 -25.775030 164.889297
vn -4.925585 -1.908184 1.266666
v -69.714081 -27.007389 156.506897
vn -4.754362 -2.938568 2.116949
v -57.389679 -35.534191 164.889297
vn -5.550435 -2.763790 -0.712983
v -62.249168 -25.775030 125.617798
vn -5.284479 -3.271534 -0.722264
v -57.389679 -35.534191 125.617798
vn 4.925588 -1.908184 1.266666
v 69.714027 -27.007389 156.506897
vn 4.941497 -2.460590 2.030060
v 62.249119 -25.775030 164.889297
vn 4.754361 -2.938570 2.116949
v 57.389641 -35.534191 164.889297
vn 5.550427 -2.763804 -0.712982
v 62.249119 -25.775030 125.617798
vn 5.284478 -3.271535 -0.722263
v 57.389641 -35.534191 125.617798
vn -2.888869 -1.196153 -2.949050
v -62.361889 -25.831150 -0.000004
vn 2.888869 -1.196152 -2.949049
v 62.361851 -25.831150 -0.000004
vn -2.890351 -1.195863 2.962770
v -62.361889 -25.831150 182.800003
vn 2.890351 -1.195861 2.962768
v 62.361851 -25.831150 182.800003
vn -5.309559 -1.510678 2.040326
v -66.138161 -13.150160 164.889297
vn -5.192345 -0.970616 1.266665
v -73.489662 -13.737610 156.506897
vn -5.212972 -2.014694 2.121964
v -62.941898 -24.383841 164.889297
vn -5.963784 -1.696817 -0.714447
v -66.138161 -13.150160 125.617798
vn -5.795593 -2.244724 -0.723518
v -62.941898 -24.383841 125.617798
vn 5.192345 -0.970616 1.266666
v 73.489616 -13.737610 156.506897
vn 5.309571 -1.510655 2.040333
v 66.138123 -13.150160 164.889297
vn 5.212969 -2.014712 2.121964
v 62.941860 -24.383841 164.889297
vn 5.963793 -1.696790 -0.714447
v 66.138123 -13.150160 125.617798
vn 5.795585 -2.244743 -0.723517
v 62.941860 -24.383841 125.617798
vn -3.066644 -0.609491 -2.947673
v -66.203033 -13.168620 -0.000004
vn 3.066644 -0.609489 -2.947674
v 66.202980 -13.168620 -0.000004
vn -3.067778 -0.607952 2.956517
v -66.203033 -13.168620 182.800003
vn 3.067779 -0.607951 2.956519
v 66.202980 -13.168620 182.800003
vn -5.493835 -1.022629 2.127995
v -66.350708 -12.403110 164.889297
vn -5.587918 0.000001 2.131134
v -67.500031 -0.000023 164.889297
vn -5.282289 0.000000 1.266666
v -74.762627 -0.000023 156.506897
vn -6.109307 -1.141456 -0.724981
v -66.350708 -12.403110 125.617798
vn 5.282288 -0.000000 1.266667
v 74.762581 -0.000023 156.506897
vn 5.493838 -1.022601 2.127992
v 66.350662 -12.403110 164.889297
vn 6.109313 -1.141426 -0.724982
v 66.350662 -12.403110 125.617798
vn 6.214976 0.000000 -0.725761
v 67.499977 -0.000023 125.617798
vn -3.126576 0.000000 -2.946928
v -67.500031 -0.000023 -0.000004
vn 3.126576 0.000000 -2.946928
v 67.499977 -0.000023 -0.000004
vn -3.127140 0.000000 2.952979
v -67.500031 -0.000023 182.800003
vn 3.127140 0.000000 2.952979
v 67.499977 -0.000023 182.800003
vn -6.214976 -0.000001 -0.725761
v -67.500031 -0.000023 125.617798
vn 5.587917 -0.000001 2.131134
v 67.499977 -0.000023 164.889297
vn -5.192345 0.970616 1.266666
v -73.489662 13.737570 156.506897
vn 5.192345 0.970616 1.266665
v 73.489616 13.737570 156.506897
vn -6.109306 1.141457 -0.724982
v -66.350708 12.403070 125.617798
vn 5.493838 1.022602 2.127993
v 66.350662 12.403070 164.889297
vn -3.066644 0.609490 -2.947673
v -66.203033 13.168570 -0.000004
vn 2.600024 1.736862 -2.950213
v 56.124180 37.500969 -0.000004
vn 2.888869 1.196151 -2.949048
v 62.361851 25.831110 -0.000004
vn 3.066644 0.609489 -2.947674
v 66.202980 13.168570 -0.000004
vn 1.196753 2.888856 -2.952437
v 25.831110 62.361851 -0.000004
vn 1.737387 2.599805 -2.951907
v 37.500969 56.124180 -0.000004
vn 2.211215 2.210821 -2.951169
v 47.729679 47.729679 -0.000004
vn -0.610106 3.066856 -2.952756
v -13.168620 66.202980 -0.000004
vn 0.000000 3.126960 -2.952861
v -0.000023 67.499977 -0.000004
vn 0.610106 3.066856 -2.952757
v 13.168570 66.202980 -0.000004
vn -2.211215 2.210821 -2.951168
v -47.729729 47.729679 -0.000004
vn -1.737386 2.599805 -2.951907
v -37.501011 56.124180 -0.000004
vn -1.196753 2.888856 -2.952437
v -25.831150 62.361851 -0.000004
vn -2.888869 1.196153 -2.949050
v -62.361889 25.831110 -0.000004
vn -2.600025 1.736861 -2.950213
v -56.124222 37.500969 -0.000004
vn 6.109313 1.141426 -0.724981
v 66.350662 12.403070 125.617798
vn -5.493836 1.022629 2.127994
v -66.350708 12.403070 164.889297
vn -3.067778 0.607952 2.956517
v -66.203033 13.168570 182.800003
vn -2.212471 2.212283 2.972063
v -47.729729 47.729679 182.800003
vn -2.601500 1.737581 2.967929
v -56.124222 37.500969 182.800003
vn -2.890351 1.195862 2.962769
v -62.361889 25.831110 182.800003
vn -0.610402 3.069248 2.978758
v -13.168620 66.202980 182.800003
vn -1.197363 2.891102 2.977436
v -25.831150 62.361851 182.800003
vn -1.738327 2.601759 2.975215
v -37.501011 56.124180 182.800003
vn 1.197363 2.891102 2.977436
v 25.831110 62.361851 182.800003
vn 0.610402 3.069248 2.978759
v 13.168570 66.202980 182.800003
vn 0.000000 3.129395 2.979196
v -0.000023 67.499977 182.800003
vn 2.601499 1.737582 2.967929
v 56.124180 37.500969 182.800003
vn 2.212471 2.212283 2.972064
v 47.729679 47.729679 182.800003
vn 1.738328 2.601759 2.975215
v 37.500969 56.124180 182.800003
vn 3.067779 0.607951 2.956519
v 66.202980 13.168570 182.800003
vn 2.890351 1.195861 2.962768
v 62.361851 25.831110 182.800003
vn -5.963783 1.696819 -0.714447
v -66.138161 13.150120 125.617798
vn 5.963793 1.696792 -0.714448
v 66.138123 13.150120 125.617798
vn -5.309559 1.510679 2.040327
v -66.138161 13.150120 164.889297
vn 5.309571 1.510656 2.040333
v 66.138123 13.150120 164.889297
vn -5.795593 2.244723 -0.723519
v -62.941898 24.383789 125.617798
vn -4.925584 1.908184 1.266666
v -69.714081 27.007351 156.506897
vn 5.795586 2.244742 -0.723517
v 62.941860 24.383789 125.617798
vn 4.925587 1.908184 1.266666
v 69.714027 27.007351 156.506897
vn -5.212971 2.014693 2.121964
v -62.941898 24.383789 164.889297
vn 5.212969 2.014710 2.121964
v 62.941860 24.383789 164.889297
vn -5.550437 2.763787 -0.712984
v -62.249168 25.774981 125.617798
vn 5.550428 2.763802 -0.712982
v 62.249119 25.774981 125.617798
vn -4.941506 2.460576 2.030066
v -62.249168 25.774981 164.889297
vn 4.941497 2.460588 2.030060
v 62.249119 25.774981 164.889297
vn -5.284474 3.271541 -0.722264
v -57.389679 35.534149 125.617798
vn -4.491091 2.780767 1.266666
v -63.564468 39.357422 156.506897
vn 5.284474 3.271541 -0.722264
v 57.389641 35.534149 125.617798
vn 4.491089 2.780766 1.266665
v 63.564430 39.357422 156.506897
vn -4.754358 2.938575 2.116949
v -57.389679 35.534149 164.889297
vn 4.754358 2.938576 2.116950
v 57.389641 35.534149 164.889297
vn -4.948082 3.736620 -0.711746
v -55.984241 37.395248 125.617798
vn 4.948079 3.736624 -0.711745
v 55.984200 37.395248 125.617798
vn -4.405178 3.326638 2.021385
v -55.984241 37.395248 164.889297
vn 4.405175 3.326641 2.021384
v 55.984200 37.395248 164.889297
vn -4.593358 4.186946 -0.721219
v -49.883121 45.474430 125.617798
vn -3.903659 3.558653 1.266666
v -55.250252 50.367222 156.506897
vn 4.593355 4.186950 -0.721218
v 49.883080 45.474430 125.617798
vn 3.903658 3.558656 1.266666
v 55.250198 50.367222 156.506897
vn -4.133475 3.762669 2.112888
v -49.883121 45.474430 164.889297
vn 4.133473 3.762673 2.112887
v 49.883080 45.474430 164.889297
vn -4.177235 4.582198 -0.710734
v -47.584270 47.570122 125.617798
vn 4.177232 4.582200 -0.710733
v 47.584221 47.570122 125.617798
vn -3.718858 4.079383 2.014286
v -47.584270 47.570122 164.889297
vn 3.718855 4.079385 2.014284
v 47.584221 47.570122 164.889297
vn -3.745766 4.959771 -0.720383
v -40.677860 53.866138 125.617798
vn -3.183288 4.215355 1.266666
v -45.054562 59.661831 156.506897
vn 3.745767 4.959770 -0.720382
v 40.677818 53.866138 125.617798
vn 3.183290 4.215355 1.266666
v 45.054508 59.661831 156.506897
vn -3.371384 4.458740 2.109718
v -40.677860 53.866138 164.889297
vn 3.371386 4.458740 2.109717
v 40.677818 53.866138 164.889297
vn -3.264127 5.271744 -0.709945
v -37.370762 55.913811 125.617798
vn 3.264128 5.271744 -0.709945
v 37.370720 55.913811 125.617798
vn -2.905910 4.693203 2.008758
v -37.370762 55.913811 164.889297
vn 2.905910 4.693203 2.008757
v 37.370720 55.913811 164.889297
vn -2.770584 5.563671 -0.719754
v -30.087360 60.423500 125.617798
vn -2.354518 4.728510 1.266666
v -33.324581 66.924721 156.506897
vn 2.770584 5.563670 -0.719754
v 30.087320 60.423500 125.617798
vn 2.354519 4.728509 1.266666
v 33.324539 66.924721 156.506897
vn -2.494045 5.002878 2.107387
v -30.087360 60.423500 164.889297
vn 2.494046 5.002878 2.107386
v 30.087320 60.423500 164.889297
vn -2.239869 5.781763 -0.709382
v -25.733641 62.110142 125.617798
vn 2.239868 5.781763 -0.709382
v 25.733601 62.110142 125.617798
vn -1.994037 5.147198 2.004810
v -25.733641 62.110142 164.889297
vn 1.994036 5.147199 2.004811
v 25.733601 62.110142 164.889297
vn -1.701023 5.978077 -0.719336
v -18.472281 64.923210 125.617798
vn -1.445565 5.080638 1.266666
v -20.459780 71.908577 156.506897
vn 1.701022 5.978077 -0.719337
v 18.472231 64.923210 125.617798
vn 1.445565 5.080638 1.266666
v 20.459740 71.908577 156.506897
vn -1.531400 5.376383 2.105858
v -18.472281 64.923210 164.889297
vn 1.531399 5.376382 2.105858
v 18.472231 64.923210 164.889297
vn -1.139337 6.094890 -0.709043
v -13.116530 65.924370 125.617798
vn 1.139338 6.094890 -0.709043
v 13.116490 65.924370 125.617798
vn -1.014284 5.425924 2.002441
v -13.116530 65.924370 164.889297
vn 1.014285 5.425924 2.002441
v 13.116490 65.924370 164.889297
vn -0.573528 6.188870 -0.719125
v -6.228134 67.212051 125.617798
vn -0.487388 5.259756 1.266665
v -6.898247 74.443657 156.506897
vn 0.573528 6.188870 -0.719125
v 6.228096 67.212051 125.617798
vn 0.487388 5.259756 1.266665
v 6.898201 74.443657 156.506897
vn -0.516364 5.566408 2.105098
v -6.228134 67.212051 164.889297
vn 0.516364 5.566408 2.105098
v 6.228096 67.212051 164.889297
vn 0.000000 6.200467 -0.708930
v -0.000023 67.212051 125.617798
vn 0.000000 5.519903 2.001652
v -0.000023 67.212051 164.889297
# 228 vertices, 0 vertices normals
f 1//1 2//2 3//3
f 4//4 5//5 3//3
f 5//5 6//6 3//3
f 6//6 1//1 3//3
f 7//7 4//4 3//3
f 2//2 8//8 3//3
f 8//8 9//9 3//3
f 9//9 7//7 3//3
f 10//10 11//11 4//4
f 11//11 12//12 4//4
f 12//12 5//5 4//4
f 13//13 10//10 4//4
f 7//7 14//14 4//4
f 14//14 13//13 4//4
f 15//15 16//16 2//2
f 1//1 17//17 2//2
f 17//17 15//15 2//2
f 16//16 18//18 2//2
f 18//18 19//19 2//2
f 19//19 8//8 2//2
f 20//20 21//21 10//10
f 21//21 22//22 10//10
f 22//22 11//11 10//10
f 23//23 20//20 10//10
f 13//13 24//24 10//10
f 24//24 23//23 10//10
f 9//9 25//25 26//26
f 27//27 14//14 26//26
f 14//14 7//7 26//26
f 7//7 9//9 26//26
f 28//28 27//27 26//26
f 25//25 28//28 26//26
f 29//29 6//6 30//30
f 6//6 5//5 30//30
f 5//5 12//12 30//30
f 12//12 31//31 30//30
f 32//32 29//29 30//30
f 31//31 32//32 30//30
f 8//8 25//25 9//9
f 27//27 13//13 14//14
f 29//29 1//1 6//6
f 11//11 31//31 12//12
f 33//33 34//34 16//16
f 15//15 35//35 16//16
f 35//35 33//33 16//16
f 34//34 36//36 16//16
f 36//36 37//37 16//16
f 37//37 18//18 16//16
f 38//38 39//39 20//20
f 39//39 40//40 20//20
f 40//40 21//21 20//20
f 41//41 38//38 20//20
f 23//23 42//42 20//20
f 42//42 41//41 20//20
f 19//19 43//43 25//25
f 8//8 19//19 25//25
f 43//43 28//28 25//25
f 44//44 24//24 27//27
f 24//24 13//13 27//27
f 28//28 44//44 27//27
f 32//32 45//45 29//29
f 45//45 17//17 29//29
f 17//17 1//1 29//29
f 46//46 32//32 31//31
f 11//11 22//22 31//31
f 22//22 46//46 31//31
f 18//18 43//43 19//19
f 44//44 23//23 24//24
f 45//45 15//15 17//17
f 21//21 46//46 22//22
f 37//37 47//47 43//43
f 18//18 37//37 43//43
f 47//47 28//28 43//43
f 48//48 42//42 44//44
f 42//42 23//23 44//44
f 28//28 48//48 44//44
f 32//32 49//49 45//45
f 49//49 35//35 45//45
f 35//35 15//15 45//45
f 50//50 32//32 46//46
f 21//21 40//40 46//46
f 40//40 50//50 46//46
f 36//36 47//47 37//37
f 48//48 41//41 42//42
f 49//49 33//33 35//35
f 39//39 50//50 40//40
f 51//51 52//52 34//34
f 33//33 53//53 34//34
f 53//53 51//51 34//34
f 52//52 54//54 34//34
f 54//54 55//55 34//34
f 55//55 36//36 34//34
f 56//56 57//57 38//38
f 57//57 58//58 38//38
f 58//58 39//39 38//38
f 59//59 56//56 38//38
f 41//41 60//60 38//38
f 60//60 59//59 38//38
f 55//55 61//61 47//47
f 36//36 55//55 47//47
f 61//61 28//28 47//47
f 62//62 60//60 48//48
f 60//60 41//41 48//48
f 28//28 62//62 48//48
f 32//32 63//63 49//49
f 63//63 53//53 49//49
f 53//53 33//33 49//49
f 64//64 32//32 50//50
f 39//39 58//58 50//50
f 58//58 64//64 50//50
f 54//54 61//61 55//55
f 62//62 59//59 60//60
f 63//63 51//51 53//53
f 57//57 64//64 58//58
f 65//65 66//66 52//52
f 51//51 67//67 52//52
f 67//67 65//65 52//52
f 66//66 68//68 52//52
f 68//68 69//69 52//52
f 69//69 54//54 52//52
f 70//70 71//71 56//56
f 71//71 72//72 56//56
f 72//72 57//57 56//56
f 73//73 70//70 56//56
f 59//59 74//74 56//56
f 74//74 73//73 56//56
f 69//69 75//75 61//61
f 54//54 69//69 61//61
f 75//75 28//28 61//61
f 76//76 74//74 62//62
f 74//74 59//59 62//62
f 28//28 76//76 62//62
f 32//32 77//77 63//63
f 77//77 67//67 63//63
f 67//67 51//51 63//63
f 78//78 32//32 64//64
f 57//57 72//72 64//64
f 72//72 78//78 64//64
f 68//68 75//75 69//69
f 76//76 73//73 74//74
f 77//77 65//65 67//67
f 71//71 78//78 72//72
f 79//79 80//80 66//66
f 65//65 81//81 66//66
f 81//81 79//79 66//66
f 80//80 82//82 66//66
f 82//82 83//83 66//66
f 83//83 68//68 66//66
f 84//84 85//85 70//70
f 85//85 86//86 70//70
f 86//86 71//71 70//70
f 87//87 84//84 70//70
f 73//73 88//88 70//70
f 88//88 87//87 70//70
f 83//83 89//89 75//75
f 68//68 83//83 75//75
f 89//89 28//28 75//75
f 90//90 88//88 76//76
f 88//88 73//73 76//76
f 28//28 90//90 76//76
f 32//32 91//91 77//77
f 91//91 81//81 77//77
f 81//81 65//65 77//77
f 92//92 32//32 78//78
f 71//71 86//86 78//78
f 86//86 92//92 78//78
f 82//82 89//89 83//83
f 90//90 87//87 88//88
f 91//91 79//79 81//81
f 85//85 92//92 86//86
f 93//93 94//94 80//80
f 79//79 95//95 80//80
f 95//95 93//93 80//80
f 94//94 96//96 80//80
f 96//96 97//97 80//80
f 97//97 82//82 80//80
f 98//98 99//99 84//84
f 99//99 100//100 84//84
f 100//100 85//85 84//84
f 101//101 98//98 84//84
f 87//87 102//102 84//84
f 102//102 101//101 84//84
f 97//97 103//103 89//89
f 82//82 97//97 89//89
f 103//103 28//28 89//89
f 104//104 102//102 90//90
f 102//102 87//87 90//90
f 28//28 104//104 90//90
f 32//32 105//105 91//91
f 105//105 95//95 91//91
f 95//95 79//79 91//91
f 106//106 32//32 92//92
f 85//85 100//100 92//92
f 100//100 106//106 92//92
f 96//96 103//103 97//97
f 104//104 101//101 102//102
f 105//105 93//93 95//95
f 99//99 106//106 100//100
f 93//93 107//107 94//94
f 107//107 108//108 94//94
f 108//108 109//109 94//94
f 109//109 110//110 94//94
f 110//110 96//96 94//94
f 111//111 112//112 98//98
f 112//112 99//99 98//98
f 101//101 113//113 98//98
f 113//113 114//114 98//98
f 114//114 111//111 98//98
f 110//110 115//115 103//103
f 96//96 110//110 103//103
f 115//115 28//28 103//103
f 116//116 113//113 104//104
f 113//113 101//101 104//104
f 28//28 116//116 104//104
f 32//32 117//117 105//105
f 117//117 107//107 105//105
f 107//107 93//93 105//105
f 118//118 32//32 106//106
f 99//99 112//112 106//106
f 112//112 118//118 106//106
f 119//119 115//115 110//110
f 109//109 119//119 110//110
f 116//116 114//114 113//113
f 117//117 108//108 107//107
f 120//120 118//118 112//112
f 111//111 120//120 112//112
f 121//121 119//119 109//109
f 122//122 120//120 111//111
f 121//121 123//123 119//119
f 122//122 124//124 120//120
f 125//125 28//28 115//115
f 119//119 123//123 115//115
f 123//123 125//125 115//115
f 126//126 127//127 28//28
f 127//127 128//128 28//28
f 128//128 116//116 28//28
f 129//129 130//130 28//28
f 130//130 131//131 28//28
f 131//131 126//126 28//28
f 132//132 133//133 28//28
f 133//133 134//134 28//28
f 134//134 129//129 28//28
f 135//135 136//136 28//28
f 136//136 137//137 28//28
f 137//137 132//132 28//28
f 125//125 138//138 28//28
f 138//138 139//139 28//28
f 139//139 135//135 28//28
f 140//140 111//111 114//114
f 108//108 141//141 109//109
f 141//141 121//121 109//109
f 140//140 122//122 111//111
f 117//117 141//141 108//108
f 32//32 142//142 117//117
f 142//142 141//141 117//117
f 143//143 144//144 32//32
f 144//144 145//145 32//32
f 145//145 142//142 32//32
f 146//146 147//147 32//32
f 147//147 148//148 32//32
f 148//148 143//143 32//32
f 149//149 150//150 32//32
f 150//150 151//151 32//32
f 151//151 146//146 32//32
f 152//152 153//153 32//32
f 153//153 154//154 32//32
f 154//154 149//149 32//32
f 118//118 155//155 32//32
f 155//155 156//156 32//32
f 156//156 152//152 32//32
f 128//128 140//140 116//116
f 140//140 114//114 116//116
f 124//124 118//118 120//120
f 124//124 155//155 118//118
f 157//157 125//125 123//123
f 121//121 157//157 123//123
f 128//128 158//158 140//140
f 158//158 122//122 140//140
f 142//142 159//159 141//141
f 159//159 121//121 141//141
f 160//160 155//155 124//124
f 122//122 160//160 124//124
f 161//161 125//125 157//157
f 121//121 162//162 157//157
f 162//162 161//161 157//157
f 128//128 163//163 158//158
f 164//164 122//122 158//158
f 163//163 164//164 158//158
f 142//142 165//165 159//159
f 162//162 121//121 159//159
f 165//165 162//162 159//159
f 166//166 155//155 160//160
f 122//122 164//164 160//160
f 164//164 166//166 160//160
f 161//161 138//138 125//125
f 127//127 163//163 128//128
f 145//145 165//165 142//142
f 166//166 156//156 155//155
f 167//167 138//138 161//161
f 162//162 167//167 161//161
f 127//127 168//168 163//163
f 168//168 164//164 163//163
f 145//145 169//169 165//165
f 169//169 162//162 165//165
f 170//170 156//156 166//166
f 164//164 170//170 166//166
f 171//171 138//138 167//167
f 162//162 172//172 167//167
f 172//172 171//171 167//167
f 127//127 173//173 168//168
f 174//174 164//164 168//168
f 173//173 174//174 168//168
f 145//145 175//175 169//169
f 172//172 162//162 169//169
f 175//175 172//172 169//169
f 176//176 156//156 170//170
f 164//164 174//174 170//170
f 174//174 176//176 170//170
f 171//171 139//139 138//138
f 126//126 173//173 127//127
f 144//144 175//175 145//145
f 176//176 152//152 156//156
f 177//177 139//139 171//171
f 172//172 177//177 171//171
f 126//126 178//178 173//173
f 178//178 174//174 173//173
f 144//144 179//179 175//175
f 179//179 172//172 175//175
f 180//180 152//152 176//176
f 174//174 180//180 176//176
f 181//181 139//139 177//177
f 172//172 182//182 177//177
f 182//182 181//181 177//177
f 126//126 183//183 178//178
f 184//184 174//174 178//178
f 183//183 184//184 178//178
f 144//144 185//185 179//179
f 182//182 172//172 179//179
f 185//185 182//182 179//179
f 186//186 152//152 180//180
f 174//174 184//184 180//180
f 184//184 186//186 180//180
f 181//181 135//135 139//139
f 131//131 183//183 126//126
f 143//143 185//185 144//144
f 186//186 153//153 152//152
f 187//187 135//135 181//181
f 182//182 187//187 181//181
f 131//131 188//188 183//183
f 188//188 184//184 183//183
f 143//143 189//189 185//185
f 189//189 182//182 185//185
f 190//190 153//153 186//186
f 184//184 190//190 186//186
f 191//191 135//135 187//187
f 182//182 192//192 187//187
f 192//192 191//191 187//187
f 131//131 193//193 188//188
f 194//194 184//184 188//188
f 193//193 194//194 188//188
f 143//143 195//195 189//189
f 192//192 182//182 189//189
f 195//195 192//192 189//189
f 196//196 153//153 190//190
f 184//184 194//194 190//190
f 194//194 196//196 190//190
f 191//191 136//136 135//135
f 130//130 193//193 131//131
f 148//148 195//195 143//143
f 196//196 154//154 153//153
f 197//197 136//136 191//191
f 192//192 197//197 191//191
f 130//130 198//198 193//193
f 198//198 194//194 193//193
f 148//148 199//199 195//195
f 199//199 192//192 195//195
f 200//200 154//154 196//196
f 194//194 200//200 196//196
f 201//201 136//136 197//197
f 192//192 202//202 197//197
f 202//202 201//201 197//197
f 130//130 203//203 198//198
f 204//204 194//194 198//198
f 203//203 204//204 198//198
f 148//148 205//205 199//199
f 202//202 192//192 199//199
f 205//205 202//202 199//199
f 206//206 154//154 200//200
f 194//194 204//204 200//200
f 204//204 206//206 200//200
f 201//201 137//137 136//136
f 129//129 203//203 130//130
f 147//147 205//205 148//148
f 206//206 149//149 154//154
f 207//207 137//137 201//201
f 202//202 207//207 201//201
f 129//129 208//208 203//203
f 208//208 204//204 203//203
f 147//147 209//209 205//205
f 209//209 202//202 205//205
f 210//210 149//149 206//206
f 204//204 210//210 206//206
f 211//211 137//137 207//207
f 202//202 212//212 207//207
f 212//212 211//211 207//207
f 129//129 213//213 208//208
f 214//214 204//204 208//208
f 213//213 214//214 208//208
f 147//147 215//215 209//209
f 212//212 202//202 209//209
f 215//215 212//212 209//209
f 216//216 149//149 210//210
f 204//204 214//214 210//210
f 214//214 216//216 210//210
f 211//211 132//132 137//137
f 134//134 213//213 129//129
f 146//146 215//215 147//147
f 216//216 150//150 149//149
f 217//217 132//132 211//211
f 212//212 217//217 211//211
f 134//134 218//218 213//213
f 218//218 214//214 213//213
f 146//146 219//219 215//215
f 219//219 212//212 215//215
f 220//220 150//150 216//216
f 214//214 220//220 216//216
f 221//221 132//132 217//217
f 212//212 222//222 217//217
f 222//222 221//221 217//217
f 134//134 223//223 218//218
f 224//224 214//214 218//218
f 223//223 224//224 218//218
f 146//146 225//225 219//219
f 222//222 212//212 219//219
f 225//225 222//222 219//219
f 226//226 150//150 220//220
f 214//214 224//224 220//220
f 224//224 226//226 220//220
f 221//221 133//133 132//132
f 133//133 223//223 134//134
f 151//151 225//225 146//146
f 226//226 151//151 150//150
f 227//227 133//133 221//221
f 222//222 227//227 221//221
f 223//223 133//133 227//227
f 222//222 224//224 227//227
f 224//224 223//223 227//227
f 151//151 228//228 225//225
f 228//228 222//222 225//225
f 151//151 226//226 228//228
f 224//224 222//222 228//228
f 226//226 224//224 228//228
# 452 faces, 0 coords texture
# End of File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/lf_metacarpal.obj
+2996
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/mounting_plate.obj
+3936
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/palm.obj
+23598
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_distal_pst.obj
+9528
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_middle.obj
+2544
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/th_proximal.obj
+840
View File
@@ -0,0 +1,840 @@
####
#
# OBJ File Generated by Meshlab
#
####
# Object th_proximal_E3M5.obj
#
# Vertices: 207
# Faces: 410
#
####
vn 3.643954 -1.325716 4.671926
v 6.562990 -2.382770 47.242680
vn 3.001478 -2.518516 4.645406
v 5.352630 -4.479270 47.242611
vn 4.878201 -1.847160 3.037849
v 8.872470 -3.223370 44.399899
vn 3.531775 0.367557 4.782580
v 6.609160 0.001130 47.535782
vn 0.953052 -1.321807 5.562101
v 2.846790 -2.376570 49.092270
vn 1.082270 0.903166 5.793338
v 2.542210 1.397900 49.166180
vn 5.532956 -1.967820 1.233064
v 10.147950 -3.687670 40.895760
vn 5.984508 1.179062 0.027198
v 10.661380 2.264480 37.154110
vn 5.829669 0.383533 1.229223
v 10.657370 -0.005540 41.136059
vn 5.970188 1.091981 0.898389
v 12.342160 2.621740 3.915270
vn 5.446340 2.683794 0.644148
v 11.527860 5.127820 3.915270
vn 6.063344 0.301712 0.824233
v 11.016340 0.001130 12.128600
vn 5.500022 2.588229 0.008108
v 9.958090 4.428920 37.154110
vn 6.119324 0.420995 -0.352597
v 10.153960 0.001130 28.836269
vn 5.153292 3.235846 1.003221
v 8.875570 6.440060 12.121890
vn 6.101694 0.451245 0.317057
v 10.192280 0.001130 20.462151
vn 5.226054 3.232666 0.320266
v 8.195270 5.945820 20.458750
vn 5.161254 3.329227 -0.228657
v 8.177800 5.933080 28.837891
vn 4.379027 4.261180 0.709538
v 8.447140 9.368060 3.915280
vn 3.606742 4.990098 0.959890
v 6.452030 8.863550 12.121880
vn 3.656401 5.032544 0.322543
v 5.957840 8.183300 20.458740
vn 3.677874 5.010916 -0.186785
v 5.945050 8.165760 28.837900
vn 4.100842 4.504420 -0.342478
v 7.297320 8.091100 37.154179
vn 3.033220 5.348574 0.921753
v 6.315350 10.916900 3.915280
vn 1.478301 5.894803 1.005845
v 3.398320 10.419530 12.121890
vn 1.876934 5.851780 0.890083
v 3.908080 11.988670 3.915270
vn 1.429651 5.970322 0.318089
v 3.138460 9.619850 20.458750
vn 1.454439 5.962157 -0.254716
v 3.131750 9.599230 28.837891
vn 1.818428 5.802749 -0.367213
v 3.377120 10.354450 37.154171
vn 0.658779 6.091357 0.923589
v 1.330670 12.536530 3.915280
vn -1.462311 5.898105 1.026773
v -3.372080 10.419530 12.121890
vn -1.429653 5.970321 0.318084
v -3.112220 9.619850 20.458750
vn -1.505360 5.943258 -0.190814
v -3.105500 9.599230 28.837891
vn -0.609878 6.097325 0.993446
v -1.304430 12.536530 3.915280
vn -0.823472 5.939095 -0.386569
v -1.124750 10.827620 37.154179
vn -1.941430 5.838144 0.866579
v -3.881840 11.988670 3.915270
vn -3.011332 5.369691 0.823297
v -6.289110 10.916900 3.915280
vn -3.685035 4.933015 0.897817
v -6.425790 8.863550 12.121890
vn -3.656404 5.032538 0.322562
v -5.931600 8.183300 20.458750
vn -3.589581 5.077980 -0.214129
v -5.918800 8.165760 28.837891
vn -3.830478 4.677118 -0.264523
v -7.271060 8.091100 37.154179
vn -4.303851 4.312155 1.003222
v -8.420900 9.368060 3.915280
vn -5.502680 2.466675 1.035770
v -11.501620 5.127820 3.915280
vn -5.104654 3.318034 0.849405
v -8.849330 6.440060 12.121890
vn -5.226431 3.232253 0.316247
v -8.169020 5.945820 20.458750
vn -5.213428 3.266400 -0.361539
v -8.151560 5.933080 28.837891
vn -5.526800 2.659152 -0.140033
v -9.931840 4.428920 37.154110
vn -6.054203 0.416023 0.996481
v -10.990100 0.001130 12.128600
vn -6.101542 0.452416 0.318453
v -10.166030 0.001130 20.462139
vn -6.111949 0.326858 -0.205592
v -10.127710 0.001130 28.836269
vn -6.104822 0.234557 0.796840
v -12.591340 0.001130 3.915350
vn -6.041976 1.029987 -0.276721
v -10.635140 2.264480 37.154110
vn 1.314831 -3.311330 4.710030
v 3.498250 -6.035270 47.242680
vn 3.976700 -3.336886 3.053528
v 7.235230 -6.058970 44.399899
vn 5.018156 0.345359 2.765740
v 9.679700 -0.008860 44.256390
vn 5.817249 0.679453 -0.653917
v 12.917130 1.094790 -0.143840
vn 4.667316 1.257236 -3.420583
v 9.762570 3.168920 -7.261750
vn 3.888805 2.809895 -3.591537
v 8.306490 6.026620 -7.261750
vn 1.987558 2.209994 -5.150820
v 4.715470 5.683280 -9.958290
vn 2.887806 3.886230 -3.651655
v 6.038590 8.294500 -7.261680
vn 2.595609 -4.495737 3.053481
v 4.727060 -8.163670 44.399830
vn 4.521967 -3.764922 1.174573
v 8.275100 -6.931470 40.895760
vn 3.507578 4.399566 -2.056722
v 7.167060 9.847770 -3.834830
vn 3.774745 4.613333 -0.410994
v 7.642450 10.502040 0.009080
vn 1.155383 4.672535 -3.581500
v 3.180890 9.750570 -7.261680
vn 1.623348 5.367931 -2.189175
v 3.774220 11.576550 -3.834830
vn 1.838933 5.756236 -0.365424
v 4.024110 12.345690 0.009140
vn -1.043022 3.246955 -4.928975
v -2.272180 7.034460 -9.958720
vn -0.936706 1.028003 -5.756390
v -3.113070 2.272450 -11.679540
vn -1.205160 4.533834 -3.741524
v -3.154720 9.750570 -7.261750
vn 0.138728 5.428807 -2.425965
v 0.013120 12.172250 -3.834830
vn 0.023505 6.046159 -0.387344
v 0.013120 12.980970 0.009140
vn -1.756137 5.404866 -2.099213
v -3.747980 11.576550 -3.834820
vn -1.895714 5.737787 -0.365416
v -3.997870 12.345690 0.009150
vn -2.616550 2.293019 -4.965265
v -5.641700 4.785560 -9.956840
vn -2.727046 3.871344 -3.781431
v -6.012420 8.294500 -7.261750
vn -3.340392 4.597616 -2.099218
v -7.140890 9.847770 -3.834900
vn -3.534833 4.905265 -0.387354
v -7.616210 10.502040 0.009080
vn -3.918776 2.806617 -3.713833
v -8.280250 6.026620 -7.261750
vn -4.597654 3.340416 -2.099232
v -9.833490 7.155130 -3.834820
vn -4.885503 3.549578 -0.318645
v -10.487770 7.630490 0.009080
vn -3.398794 0.846840 -5.017704
v -6.828840 2.313770 -9.958100
vn -4.556650 1.169886 -3.723215
v -9.736330 3.168920 -7.261750
vn -5.407255 1.614365 -2.106083
v -11.562360 3.762210 -3.834820
vn -5.759212 1.719580 -0.342069
v -12.331480 4.012120 0.009080
vn -5.742759 0.060470 -2.095505
v -12.085490 0.036930 -3.513460
vn -6.143723 -0.010659 -0.440568
v -12.759300 -0.068200 -0.007480
vn 4.653790 3.397346 1.237363
v 8.275100 6.933760 40.895760
vn 4.144734 2.947149 2.988956
v 7.235240 6.061250 44.399899
vn 3.069663 2.126919 4.598249
v 5.352630 4.481490 47.242611
vn 2.942799 5.097024 1.176954
v 5.405810 9.341440 40.895760
vn 2.595623 4.495722 3.053526
v 4.727060 8.165910 44.399830
vn 1.503087 3.209566 4.746914
v 3.498240 6.037500 47.242680
vn 0.999551 5.798584 1.174548
v 1.885970 10.622540 40.895760
vn 0.761408 4.919088 3.246846
v 1.650220 9.285770 44.399899
vn -1.165416 5.739345 1.035257
v -1.859710 10.622540 40.895760
vn -0.798080 5.050174 3.119379
v -1.623970 9.285770 44.399899
vn -0.476519 3.527350 4.617452
v -1.197230 6.865440 47.242691
vn -1.167134 1.226254 5.592976
v -2.820540 2.378840 49.092270
vn -2.758894 5.143538 0.852557
v -5.379550 9.341440 40.895771
vn -2.578671 4.257617 3.246837
v -4.700880 8.165910 44.399899
vn -4.604956 3.542779 1.224830
v -8.248850 6.933760 40.895771
vn -3.960755 3.198165 3.044097
v -7.208980 6.061250 44.399910
vn -2.840159 2.563107 4.524149
v -5.326380 4.481490 47.242619
vn -5.768543 0.992651 1.084477
v -10.610570 1.655430 40.893169
vn -4.822252 0.685035 2.821594
v -9.637320 1.118840 44.529919
vn -3.291767 0.011668 4.950877
v -6.582910 0.001130 47.535789
vn 5.520185 -2.757620 -0.214762
v 9.958090 -4.426670 37.154110
vn 6.085318 -0.997570 -0.156997
v 10.661380 -2.262270 37.154110
vn -6.038405 -1.210480 0.815370
v -12.315920 -2.619470 3.915280
vn -5.846671 -1.946126 0.971062
v -10.405250 -3.383970 12.121890
vn -5.903769 -1.983475 0.325273
v -9.605580 -3.124170 20.458750
vn -5.928283 -1.890194 -0.244552
v -9.585000 -3.117470 28.837891
vn -6.073822 -0.914910 -0.196400
v -10.635140 -2.262270 37.154110
vn -5.480734 -2.721107 0.919936
v -11.501620 -5.125570 3.915280
vn -5.488225 -2.794405 -0.260208
v -9.931840 -4.426670 37.154110
vn -4.983471 -3.620678 0.954530
v -8.849330 -6.437770 12.121890
vn -5.032602 -3.656343 0.322523
v -8.169020 -5.943570 20.458750
vn -4.982490 -3.722709 -0.177699
v -8.151560 -5.930870 28.837891
vn -4.281575 -4.371631 0.919914
v -8.420900 -9.365770 3.915280
vn -3.606762 -4.990088 0.959874
v -6.425790 -8.861270 12.121890
vn -3.656373 -5.032553 0.322545
v -5.931600 -8.181070 20.458750
vn -3.677859 -5.010942 -0.186779
v -5.918800 -8.163470 28.837891
vn -4.100821 -4.504424 -0.342461
v -7.271060 -8.088870 37.154179
vn -3.033205 -5.348558 0.921776
v -6.289110 -10.914670 3.915280
vn -1.478313 -5.894798 1.005836
v -3.372080 -10.417270 12.121890
vn -1.876897 -5.851807 0.890121
v -3.881840 -11.986370 3.915280
vn -1.429619 -5.970338 0.318101
v -3.112220 -9.617570 20.458750
vn -1.506599 -5.954899 -0.317641
v -3.105500 -9.596970 28.837891
vn -1.884631 -5.804482 -0.297898
v -3.350860 -10.352170 37.154171
vn -0.658824 -6.091347 0.923541
v -1.304430 -12.534270 3.915280
vn 0.609892 -6.097320 0.993443
v 1.330670 -12.534270 3.915280
vn 0.849994 -5.977596 -0.326686
v 1.151000 -10.825370 37.154179
vn 1.462321 -5.898098 1.026783
v 3.398320 -10.417270 12.121890
vn 1.941437 -5.838161 0.866602
v 3.908080 -11.986370 3.915280
vn 1.429620 -5.970339 0.318073
v 3.138460 -9.617570 20.458750
vn 1.549845 -5.938266 -0.231086
v 3.131750 -9.596970 28.837891
vn 3.011300 -5.369688 0.823258
v 6.315350 -10.914670 3.915280
vn 3.685051 -4.933005 0.897822
v 6.452030 -8.861270 12.121880
vn 3.656374 -5.032552 0.322553
v 5.957840 -8.181070 20.458740
vn 3.589571 -5.078000 -0.214131
v 5.945050 -8.163470 28.837900
vn 3.859977 -4.662407 -0.283751
v 7.297320 -8.088870 37.154179
vn 4.303882 -4.312141 1.003237
v 8.447140 -9.365770 3.915280
vn 5.444722 -2.756230 0.963231
v 11.527860 -5.125570 3.915270
vn 4.934412 -3.704745 0.798712
v 8.875570 -6.437770 12.121890
vn 5.032592 -3.656352 0.322541
v 8.195270 -5.943570 20.458750
vn 5.035184 -3.658355 -0.312354
v 8.177800 -5.930870 28.837891
vn 5.875986 -1.859541 0.888353
v 10.431490 -3.383970 12.121880
vn 5.904180 -1.982954 0.319829
v 9.631830 -3.124170 20.458750
vn 5.908853 -1.969770 -0.324025
v 9.611250 -3.117470 28.837900
vn 6.022304 -0.928999 0.948214
v 12.342160 -2.619470 3.915270
vn -1.380640 -0.686773 -5.731635
v -3.661930 -1.192970 -11.679540
vn -3.538029 -0.420581 -5.019475
v -6.965720 -1.299690 -9.956520
vn -4.553372 -1.124327 -3.776149
v -9.736330 -3.166670 -7.261740
vn -5.333509 -1.591913 -2.073859
v -11.642710 -3.811340 -3.895890
vn -5.772764 -1.756671 -0.365727
v -12.331480 -4.009870 0.009080
vn -2.640641 -2.027757 -5.025055
v -5.969780 -4.345670 -9.958720
vn -3.890096 -2.857974 -3.723780
v -8.280250 -6.024370 -7.261750
vn -4.577250 -3.404147 -2.081670
v -9.833490 -7.152870 -3.834820
vn -4.885487 -3.549574 -0.318641
v -10.487770 -7.628270 0.009080
vn -3.015986 -3.748705 -3.782438
v -6.012420 -8.292270 -7.261750
vn -3.381532 -4.597293 -2.064599
v -7.140890 -9.845470 -3.834900
vn -3.556869 -4.895599 -0.406118
v -7.616210 -10.499770 0.009080
vn -0.151955 -3.128865 -5.147010
v 0.013120 -7.394170 -9.958650
vn 0.040648 -4.863536 -3.679814
v 0.013120 -10.250070 -7.261750
vn -1.730245 -5.444893 -2.053631
v -3.747980 -11.574270 -3.834820
vn -1.838889 -5.756212 -0.365412
v -3.997870 -12.343470 0.009150
vn 0.000032 -5.683005 -2.099215
v 0.013120 -12.169970 -3.834830
vn -0.023508 -6.046187 -0.387347
v 0.013120 -12.978670 0.009150
vn 1.058917 -0.678058 -5.726054
v 3.139310 -2.270170 -11.679540
vn 1.251667 -3.134328 -4.990292
v 2.298420 -7.032170 -9.958730
vn 1.541835 -4.514576 -3.748779
v 3.180890 -9.748270 -7.261680
vn 1.756146 -5.404825 -2.099266
v 3.774220 -11.574270 -3.834830
vn 1.895695 -5.737748 -0.365423
v 4.024110 -12.343470 0.009150
vn 2.684484 -2.247843 -4.978163
v 5.996020 -4.345670 -9.958730
vn 2.676205 -3.822081 -3.833149
v 6.038590 -8.292270 -7.261680
vn 3.340403 -4.597663 -2.099252
v 7.167060 -9.845470 -3.834830
vn 3.534827 -4.905265 -0.387390
v 7.642450 -10.499770 0.009080
vn 3.912368 -2.842569 -3.701086
v 8.306490 -6.024370 -7.261750
vn 4.597599 -3.340463 -2.099265
v 9.859720 -7.152870 -3.834830
vn 4.885472 -3.549591 -0.318647
v 10.514010 -7.628270 0.009070
vn 3.371641 -0.837551 -4.977191
v 7.046500 -2.284170 -9.958720
vn 4.660064 -1.132651 -3.602630
v 9.762570 -3.166670 -7.261750
vn 5.392434 -1.421524 -2.141711
v 11.588600 -3.759970 -3.834830
vn 5.762479 -1.746300 -0.416116
v 12.357720 -4.009870 0.009070
vn -5.548463 -1.885387 1.197290
v -10.121700 -3.687670 40.895771
vn -4.859341 -1.803729 3.077813
v -8.846220 -3.223370 44.399910
vn -3.662211 -1.202200 4.646036
v -6.536740 -2.382770 47.242691
vn -1.532655 -0.687359 5.524807
v -3.462870 -1.264020 49.092270
vn -4.521961 -3.764937 1.174558
v -8.248850 -6.931470 40.895771
vn -3.976726 -3.336860 3.053545
v -7.208980 -6.058970 44.399910
vn -2.913390 -2.495858 4.685129
v -5.326380 -4.479270 47.242619
vn -2.942815 -5.097037 1.176923
v -5.379550 -9.339170 40.895771
vn -2.595614 -4.495668 3.053477
v -4.700880 -8.163670 44.399910
vn -1.293353 -3.186611 4.758747
v -3.471990 -6.035270 47.242691
vn -0.999537 -5.798590 1.174547
v -1.859710 -10.620270 40.895760
vn -0.761328 -4.919127 3.246870
v -1.623970 -9.283470 44.399910
vn 1.087662 -5.772390 1.122651
v 1.885970 -10.620270 40.895760
vn 0.761386 -4.919119 3.246853
v 1.650220 -9.283470 44.399899
vn 2.855849 -5.123626 0.995183
v 5.405810 -9.339170 40.895760
vn -1.479902 -2.811926 -5.296168
v -3.199070 -6.568490 -9.783740
vn -1.688262 -4.207472 -3.834387
v -3.783350 -9.397950 -7.685140
vn 4.954120 2.930639 -1.914305
v 9.975860 6.396290 -3.987940
vn 5.330820 3.000355 -0.601284
v 10.945740 6.098010 -0.287480
vn 5.653251 0.776351 -2.327995
v 11.566650 1.440500 -3.961560
vn 0.448171 3.221065 -4.710190
v 1.441900 7.540370 -9.955090
vn 3.087555 0.878535 -4.950277
v 7.229860 1.884770 -9.952640
# 207 vertices, 0 vertices normals
f 1//1 2//2 3//3
f 1//1 4//4 5//5
f 6//6 5//5 4//4
f 10//10 11//11 12//12
f 13//13 8//8 14//14
f 12//12 15//15 16//16
f 15//15 12//12 11//11
f 16//16 17//17 14//14
f 17//17 16//16 15//15
f 14//14 18//18 13//13
f 18//18 14//14 17//17
f 11//11 19//19 15//15
f 20//20 15//15 19//19
f 15//15 20//20 17//17
f 21//21 17//17 20//20
f 17//17 21//21 18//18
f 22//22 18//18 21//21
f 22//22 23//23 13//13
f 18//18 22//22 13//13
f 19//19 24//24 20//20
f 20//20 25//25 21//21
f 20//20 24//24 26//26
f 25//25 20//20 26//26
f 21//21 27//27 22//22
f 27//27 21//21 25//25
f 22//22 28//28 23//23
f 28//28 22//22 27//27
f 29//29 23//23 28//28
f 26//26 30//30 25//25
f 31//31 25//25 30//30
f 25//25 31//31 27//27
f 27//27 32//32 28//28
f 32//32 27//27 31//31
f 28//28 33//33 29//29
f 33//33 28//28 32//32
f 30//30 34//34 31//31
f 35//35 29//29 33//33
f 31//31 34//34 36//36
f 36//36 37//37 31//31
f 31//31 38//38 32//32
f 38//38 31//31 37//37
f 39//39 32//32 38//38
f 32//32 39//39 33//33
f 40//40 33//33 39//39
f 40//40 41//41 35//35
f 33//33 40//40 35//35
f 37//37 42//42 38//38
f 38//38 42//42 43//43
f 44//44 38//38 43//43
f 38//38 44//44 39//39
f 45//45 39//39 44//44
f 39//39 45//45 40//40
f 40//40 46//46 41//41
f 46//46 40//40 45//45
f 47//47 41//41 46//46
f 48//48 44//44 43//43
f 44//44 48//48 45//45
f 49//49 45//45 48//48
f 45//45 49//49 46//46
f 46//46 50//50 47//47
f 50//50 46//46 49//49
f 43//43 51//51 48//48
f 52//52 47//47 50//50
f 2//2 53//53 54//54
f 2//2 1//1 5//5
f 4//4 1//1 55//55
f 7//7 9//9 3//3
f 11//11 10//10 56//56
f 58//58 57//57 59//59
f 60//60 58//58 59//59
f 54//54 61//61 62//62
f 64//64 19//19 11//11
f 24//24 19//19 64//64
f 60//60 65//65 63//63
f 65//65 60//60 59//59
f 66//66 63//63 65//65
f 63//63 66//66 64//64
f 67//67 64//64 66//66
f 64//64 67//67 24//24
f 26//26 24//24 67//67
f 65//65 70//70 66//66
f 70//70 65//65 68//68
f 66//66 71//71 67//67
f 71//71 66//66 70//70
f 72//72 30//30 26//26
f 67//67 72//72 26//26
f 72//72 67//67 71//71
f 34//34 30//30 72//72
f 73//73 71//71 70//70
f 71//71 73//73 72//72
f 72//72 74//74 34//34
f 74//74 72//72 73//73
f 36//36 34//34 74//74
f 75//75 68//68 69//69
f 68//68 75//75 70//70
f 76//76 70//70 75//75
f 70//70 76//76 73//73
f 73//73 77//77 74//74
f 77//77 73//73 76//76
f 78//78 74//74 77//77
f 78//78 37//37 36//36
f 74//74 78//78 36//36
f 42//42 37//37 78//78
f 76//76 79//79 77//77
f 79//79 76//76 75//75
f 80//80 77//77 79//79
f 77//77 80//80 78//78
f 81//81 78//78 80//80
f 78//78 81//81 42//42
f 43//43 42//42 81//81
f 75//75 82//82 79//79
f 82//82 75//75 69//69
f 83//83 79//79 82//82
f 79//79 83//83 80//80
f 84//84 80//80 83//83
f 80//80 84//84 81//81
f 81//81 85//85 43//43
f 85//85 81//81 84//84
f 51//51 43//43 85//85
f 86//86 84//84 83//83
f 84//84 86//86 85//85
f 85//85 87//87 51//51
f 87//87 85//85 86//86
f 3//3 9//9 55//55
f 8//8 13//13 9//9
f 88//88 9//9 13//13
f 9//9 88//88 55//55
f 55//55 89//89 4//4
f 89//89 55//55 88//88
f 4//4 90//90 6//6
f 90//90 4//4 89//89
f 13//13 23//23 88//88
f 88//88 91//91 89//89
f 91//91 88//88 23//23
f 89//89 92//92 90//90
f 92//92 89//89 91//91
f 90//90 93//93 6//6
f 93//93 90//90 92//92
f 23//23 29//29 91//91
f 94//94 91//91 29//29
f 91//91 94//94 92//92
f 92//92 95//95 93//93
f 95//95 92//92 94//94
f 94//94 96//96 95//95
f 94//94 29//29 35//35
f 96//96 94//94 35//35
f 95//95 97//97 93//93
f 97//97 95//95 96//96
f 98//98 93//93 97//97
f 93//93 98//98 6//6
f 99//99 6//6 98//98
f 96//96 100//100 97//97
f 96//96 35//35 41//41
f 100//100 96//96 41//41
f 97//97 101//101 98//98
f 101//101 97//97 100//100
f 100//100 102//102 101//101
f 102//102 100//100 41//41
f 101//101 103//103 98//98
f 103//103 101//101 102//102
f 104//104 98//98 103//103
f 98//98 104//104 99//99
f 41//41 47//47 102//102
f 105//105 102//102 47//47
f 102//102 105//105 103//103
f 103//103 106//106 104//104
f 106//106 103//103 105//105
f 107//107 104//104 106//106
f 104//104 107//107 99//99
f 47//47 52//52 105//105
f 55//55 1//1 3//3
f 108//108 109//109 7//7
f 51//51 110//110 48//48
f 54//54 3//3 2//2
f 48//48 111//111 49//49
f 111//111 48//48 110//110
f 112//112 49//49 111//111
f 112//112 50//50 49//49
f 113//113 50//50 112//112
f 113//113 114//114 52//52
f 113//113 52//52 50//50
f 62//62 7//7 54//54
f 110//110 115//115 111//111
f 116//116 114//114 113//113
f 111//111 117//117 112//112
f 117//117 111//111 115//115
f 112//112 118//118 113//113
f 118//118 112//112 117//117
f 113//113 119//119 116//116
f 119//119 113//113 118//118
f 115//115 120//120 117//117
f 121//121 117//117 120//120
f 117//117 121//121 118//118
f 118//118 122//122 119//119
f 122//122 118//118 121//121
f 123//123 119//119 122//122
f 123//123 124//124 116//116
f 119//119 123//123 116//116
f 120//120 125//125 121//121
f 121//121 126//126 122//122
f 121//121 125//125 127//127
f 126//126 121//121 127//127
f 128//128 122//122 126//126
f 122//122 128//128 123//123
f 123//123 129//129 124//124
f 129//129 123//123 128//128
f 130//130 124//124 129//129
f 127//127 131//131 126//126
f 131//131 132//132 126//126
f 133//133 130//130 129//129
f 126//126 134//134 128//128
f 128//128 136//136 129//129
f 136//136 128//128 134//134
f 129//129 137//137 133//133
f 137//137 129//129 136//136
f 135//135 138//138 134//134
f 139//139 134//134 138//138
f 134//134 139//139 136//136
f 136//136 140//140 137//137
f 140//140 136//136 139//139
f 141//141 137//137 140//140
f 141//141 142//142 133//133
f 137//137 141//141 133//133
f 138//138 143//143 139//139
f 139//139 143//143 144//144
f 145//145 139//139 144//144
f 139//139 145//145 140//140
f 140//140 146//146 141//141
f 146//146 140//140 145//145
f 141//141 147//147 142//142
f 147//147 141//141 146//146
f 108//108 142//142 147//147
f 145//145 148//148 146//146
f 148//148 145//145 144//144
f 146//146 149//149 147//147
f 149//149 146//146 148//148
f 147//147 150//150 108//108
f 150//150 147//147 149//149
f 144//144 151//151 148//148
f 109//109 108//108 150//150
f 148//148 151//151 10//10
f 148//148 10//10 12//12
f 7//7 62//62 108//108
f 148//148 12//12 149//149
f 149//149 16//16 150//150
f 149//149 12//12 16//16
f 150//150 14//14 109//109
f 3//3 54//54 7//7
f 14//14 150//150 16//16
f 8//8 109//109 14//14
f 69//69 152//152 82//82
f 153//153 82//82 152//152
f 82//82 153//153 83//83
f 154//154 86//86 83//83
f 154//154 83//83 153//153
f 155//155 86//86 154//154
f 155//155 87//87 86//86
f 156//156 110//110 51//51
f 156//156 51//51 87//87
f 156//156 87//87 155//155
f 115//115 110//110 156//156
f 157//157 153//153 152//152
f 153//153 157//157 154//154
f 154//154 158//158 155//155
f 158//158 154//154 157//157
f 155//155 159//159 156//156
f 159//159 155//155 158//158
f 160//160 156//156 159//159
f 156//156 160//160 115//115
f 158//158 161//161 159//159
f 161//161 158//158 157//157
f 159//159 162//162 160//160
f 162//162 159//159 161//161
f 163//163 160//160 162//162
f 125//125 120//120 163//163
f 162//162 166//166 163//163
f 163//163 167//167 125//125
f 167//167 163//163 166//166
f 127//127 125//125 167//167
f 166//166 168//168 167//167
f 168//168 166//166 165//165
f 169//169 167//167 168//168
f 169//169 131//131 127//127
f 167//167 169//169 127//127
f 132//132 131//131 169//169
f 152//152 170//170 164//164
f 170//170 152//152 69//69
f 171//171 164//164 170//170
f 164//164 171//171 165//165
f 165//165 172//172 168//168
f 172//172 165//165 171//171
f 168//168 173//173 169//169
f 173//173 168//168 172//172
f 174//174 169//169 173//173
f 169//169 174//174 132//132
f 135//135 132//132 174//174
f 175//175 171//171 170//170
f 171//171 175//175 172//172
f 172//172 176//176 173//173
f 176//176 172//172 175//175
f 173//173 177//177 174//174
f 177//177 173//173 176//176
f 178//178 174//174 177//177
f 178//178 138//138 135//135
f 174//174 178//178 135//135
f 143//143 138//138 178//178
f 179//179 176//176 175//175
f 176//176 179//179 177//177
f 180//180 177//177 179//179
f 177//177 180//180 178//178
f 181//181 178//178 180//180
f 178//178 181//181 143//143
f 144//144 143//143 181//181
f 182//182 175//175 170//170
f 175//175 182//182 179//179
f 179//179 183//183 180//180
f 183//183 179//179 182//182
f 184//184 180//180 183//183
f 180//180 184//184 181//181
f 181//181 185//185 144//144
f 185//185 181//181 184//184
f 151//151 144//144 185//185
f 170//170 69//69 59//59
f 183//183 57//57 184//184
f 184//184 56//56 185//185
f 185//185 56//56 151//151
f 10//10 151//151 56//56
f 52//52 114//114 105//105
f 186//186 106//106 105//105
f 186//186 105//105 114//114
f 53//53 2//2 5//5
f 106//106 187//187 107//107
f 187//187 106//106 186//186
f 188//188 107//107 187//187
f 142//142 108//108 62//62
f 189//189 99//99 188//188
f 114//114 116//116 186//186
f 190//190 186//186 116//116
f 186//186 190//190 187//187
f 187//187 191//191 188//188
f 191//191 187//187 190//190
f 192//192 188//188 191//191
f 188//188 192//192 189//189
f 116//116 124//124 190//190
f 193//193 190//190 124//124
f 190//190 193//193 191//191
f 194//194 191//191 193//193
f 191//191 194//194 192//192
f 192//192 195//195 189//189
f 195//195 192//192 194//194
f 124//124 130//130 193//193
f 193//193 196//196 194//194
f 196//196 193//193 130//130
f 194//194 197//197 195//195
f 197//197 194//194 196//196
f 196//196 130//130 133//133
f 198//198 196//196 133//133
f 196//196 198//198 197//197
f 199//199 197//197 198//198
f 197//197 199//199 195//195
f 53//53 195//195 199//199
f 195//195 53//53 189//189
f 5//5 189//189 53//53
f 189//189 5//5 99//99
f 198//198 200//200 199//199
f 61//61 199//199 200//200
f 199//199 61//61 53//53
f 62//62 200//200 142//142
f 200//200 62//62 61//61
f 61//61 54//54 53//53
f 8//8 9//9 109//109
f 7//7 109//109 9//9
f 133//133 200//200 198//198
f 133//133 142//142 200//200
f 120//120 160//160 163//163
f 115//115 160//160 120//120
f 157//157 201//201 161//161
f 164//164 201//201 152//152
f 201//201 157//157 152//152
f 165//165 166//166 202//202
f 164//164 165//165 202//202
f 162//162 202//202 166//166
f 161//161 202//202 162//162
f 164//164 202//202 201//201
f 161//161 201//201 202//202
f 63//63 203//203 60//60
f 56//56 204//204 11//11
f 204//204 64//64 11//11
f 64//64 204//204 203//203
f 64//64 203//203 63//63
f 57//57 205//205 184//184
f 205//205 56//56 184//184
f 57//57 203//203 205//205
f 57//57 58//58 203//203
f 58//58 60//60 203//203
f 56//56 203//203 204//204
f 56//56 205//205 203//203
f 59//59 206//206 65//65
f 206//206 68//68 65//65
f 68//68 206//206 69//69
f 206//206 59//59 69//69
f 59//59 207//207 170//170
f 207//207 182//182 170//170
f 182//182 207//207 183//183
f 183//183 207//207 57//57
f 207//207 59//59 57//57
f 132//132 135//135 134//134
f 126//126 132//132 134//134
f 5//5 6//6 99//99
f 99//99 107//107 188//188
# 410 faces, 0 coords texture
# End of File
mjx/mujoco/mjx/benchmark/model/shadow_hand/assets/wrist.obj
+3708
View File
File diff suppressed because it is too large Load Diff
mjx/mujoco/mjx/benchmark/model/shadow_hand/right_hand.xml
+327
View File
@@ -0,0 +1,327 @@
<mujoco model="right_shadow_hand">
<compiler angle="radian" meshdir="assets" autolimits="true"/>
<option impratio="10" solver="CG" iterations="8" ls_iterations="6">
<flag eulerdamp="disable"/>
</option>
<custom>
<numeric data="15" name="max_contact_points"/>
</custom>
<default>
<default class="right_hand">
<mesh scale="0.001 0.001 0.001"/>
<joint axis="1 0 0" damping="0.05" armature="0.0002" frictionloss="0.01"/>
<position forcerange="-1 1"/>
<default class="wrist">
<joint damping="0.5"/>
<default class="wrist_y">
<joint axis="0 1 0" range="-0.523599 0.174533"/>
<position kp="10" ctrlrange="-0.523599 0.174533" forcerange="-10 10"/>
</default>
<default class="wrist_x">
<joint range="-0.698132 0.488692"/>
<position kp="8" ctrlrange="-0.698132 0.488692" forcerange="-5 5"/>
</default>
</default>
<default class="thumb">
<default class="thbase">
<joint axis="0 0 -1" range="-1.0472 1.0472"/>
<position kp="0.4" ctrlrange="-1.0472 1.0472" forcerange="-3 3"/>
</default>
<default class="thproximal">
<joint range="0 1.22173"/>
<position ctrlrange="0 1.22173" forcerange="-2 2"/>
</default>
<default class="thhub">
<joint range="-0.20944 0.20944"/>
<position kp="0.5" ctrlrange="-0.20944 0.20944"/>
</default>
<default class="thmiddle">
<joint axis="0 -1 0" range="-0.698132 0.698132"/>
<position kp="1.5" ctrlrange="-0.698132 0.698132"/>
</default>
<default class="thdistal">
<joint range="-0.261799 1.5708"/>
<position ctrlrange="-0.261799 1.5708"/>
</default>
</default>
<default class="metacarpal">
<joint axis="0.573576 0 0.819152" range="0 0.785398"/>
<position ctrlrange="0 0.785398"/>
</default>
<default class="knuckle">
<joint axis="0 -1 0" range="-0.349066 0.349066"/>
<position ctrlrange="-0.349066 0.349066"/>
</default>
<default class="proximal">
<joint range="-0.261799 1.5708"/>
<position ctrlrange="-0.261799 1.5708"/>
</default>
<default class="middle_distal">
<joint range="0 1.5708"/>
<position kp="0.5" ctrlrange="0 3.1415"/>
</default>
<default class="plastic">
<geom solimp="0.5 0.99 0.0001" solref="0.005 1"/>
<default class="plastic_visual">
<geom type="mesh" material="black" contype="0" conaffinity="0" group="2"/>
</default>
<default class="plastic_collision">
<geom group="3" contype="0" conaffinity="1"/>
</default>
</default>
</default>
</default>
<asset>
<material name="black" specular="0.5" shininess="0.25" rgba="0.16355 0.16355 0.16355 1"/>
<material name="gray" specular="0.0" shininess="0.25" rgba="0.80848 0.80848 0.80848 1"/>
<material name="metallic" specular="0" shininess="0.25" rgba="0.9 0.9 0.9 1"/>
<mesh class="right_hand" file="forearm_0.obj"/>
<mesh class="right_hand" file="forearm_1.obj"/>
<mesh class="right_hand" file="forearm_collision.obj"/>
<mesh class="right_hand" file="wrist.obj"/>
<mesh class="right_hand" file="palm.obj"/>
<mesh class="right_hand" file="f_knuckle.obj"/>
<mesh class="right_hand" file="f_proximal.obj"/>
<mesh class="right_hand" file="f_middle.obj"/>
<mesh class="right_hand" file="f_distal_pst.obj"/>
<mesh class="right_hand" file="lf_metacarpal.obj"/>
<mesh class="right_hand" file="th_proximal.obj"/>
<mesh class="right_hand" file="th_middle.obj"/>
<mesh class="right_hand" file="th_distal_pst.obj"/>
</asset>
<worldbody>
<body name="rh_forearm" childclass="right_hand" quat="1 -1 1 -1">
<inertial mass="3" pos="0 0 0.09" diaginertia="0.0138 0.0138 0.00744"/>
<geom class="plastic_visual" mesh="forearm_0" material="gray"/>
<geom class="plastic_visual" mesh="forearm_1"/>
<geom class="plastic_collision" type="mesh" mesh="forearm_collision" conaffinity="0"/>
<geom class="plastic_collision" size="0.035 0.035 0.035" pos="0 -0.01 0.181" quat="0.924909 0 0.380188 0"
type="box" conaffinity="0"/>
<body name="rh_wrist" pos="0 -0.01 0.21301">
<inertial mass="0.1" pos="0 0 0.029" quat="0.5 0.5 0.5 0.5" diaginertia="6.4e-05 4.38e-05 3.5e-05"/>
<joint class="wrist_y" name="rh_WRJ2"/>
<geom class="plastic_visual" mesh="wrist" material="metallic"/>
<geom size="0.0135 0.015" quat="0.499998 0.5 0.5 -0.500002" type="cylinder" class="plastic_collision" conaffinity="0"/>
<geom size="0.011 0.005" pos="-0.026 0 0.034" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<geom size="0.011 0.005" pos="0.031 0 0.034" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<geom size="0.0135 0.009 0.005" pos="-0.021 0 0.011" quat="0.923879 0 0.382684 0" type="box"
class="plastic_collision" conaffinity="0"/>
<geom size="0.0135 0.009 0.005" pos="0.026 0 0.01" quat="0.923879 0 -0.382684 0" type="box"
class="plastic_collision" conaffinity="0"/>
<body name="rh_palm" pos="0 0 0.034">
<inertial mass="0.3" pos="0 0 0.035" quat="1 0 0 1" diaginertia="0.0005287 0.0003581 0.000191"/>
<joint class="wrist_x" name="rh_WRJ1"/>
<site name="grasp_site" pos="0 -.035 0.09" group="4"/>
<geom class="plastic_visual" mesh="palm"/>
<geom size="0.031 0.0035 0.049" pos="0.011 0.0085 0.038" type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.018 0.0085 0.049" pos="-0.002 -0.0035 0.038" type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.013 0.0085 0.005" pos="0.029 -0.0035 0.082" type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.013 0.007 0.009" pos="0.0265 -0.001 0.07" quat="0.987241 0.0990545 0.0124467 0.124052"
type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.0105 0.0135 0.012" pos="0.0315 -0.0085 0.001" type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.011 0.0025 0.015" pos="0.0125 -0.015 0.004" quat="0.971338 0 0 -0.237703" type="box"
class="plastic_collision" conaffinity="2"/>
<geom size="0.009 0.012 0.002" pos="0.011 0 0.089" type="box" class="plastic_collision" conaffinity="2"/>
<geom size="0.01 0.012 0.02" pos="-0.03 0 0.009" type="box" class="plastic_collision" conaffinity="2"/>
<body name="rh_ffknuckle" pos="0.033 0 0.095">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_FFJ4" class="knuckle"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<body name="rh_ffproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_FFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_ffmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_FFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_ffdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_FFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst"/>
</body>
</body>
</body>
</body>
<body name="rh_mfknuckle" pos="0.011 0 0.099">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_MFJ4" class="knuckle"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<body name="rh_mfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_MFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_mfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_MFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_mfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_MFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst"/>
</body>
</body>
</body>
</body>
<body name="rh_rfknuckle" pos="-0.011 0 0.095">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_RFJ4" class="knuckle" axis="0 1 0"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<body name="rh_rfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_RFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_rfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_RFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_rfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_RFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst"/>
</body>
</body>
</body>
</body>
<body name="rh_lfmetacarpal" pos="-0.033 0 0.02071">
<inertial mass="0.03" pos="0 0 0.04" quat="1 0 0 1" diaginertia="1.638e-05 1.45e-05 4.272e-06"/>
<joint name="rh_LFJ5" class="metacarpal"/>
<geom class="plastic_visual" mesh="lf_metacarpal"/>
<geom size="0.011 0.012 0.025" pos="0.002 0 0.033" type="box" class="plastic_collision"/>
<body name="rh_lfknuckle" pos="0 0 0.06579">
<inertial mass="0.008" pos="0 0 0" quat="0.5 0.5 -0.5 0.5" diaginertia="3.2e-07 2.6e-07 2.6e-07"/>
<joint name="rh_LFJ4" class="knuckle" axis="0 1 0"/>
<geom pos="0 0 0.0005" class="plastic_visual" mesh="f_knuckle" material="metallic"/>
<geom size="0.009 0.009" quat="1 0 1 0" type="cylinder" class="plastic_collision" conaffinity="0"/>
<body name="rh_lfproximal">
<inertial mass="0.03" pos="0 0 0.0225" quat="1 0 0 1" diaginertia="1e-05 9.8e-06 1.8e-06"/>
<joint name="rh_LFJ3" class="proximal"/>
<geom class="plastic_visual" mesh="f_proximal"/>
<geom size="0.009 0.02" pos="0 0 0.025" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_lfmiddle" pos="0 0 0.045">
<inertial mass="0.017" pos="0 0 0.0125" quat="1 0 0 1" diaginertia="2.7e-06 2.6e-06 8.7e-07"/>
<joint name="rh_LFJ2" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_middle"/>
<geom size="0.009 0.0125" pos="0 0 0.0125" type="capsule" class="plastic_collision" contype="1"/>
<body name="rh_lfdistal" pos="0 0 0.025">
<inertial mass="0.013" pos="0 0 0.0130769" quat="1 0 0 1"
diaginertia="1.28092e-06 1.12092e-06 5.3e-07"/>
<joint name="rh_LFJ1" class="middle_distal"/>
<geom class="plastic_visual" mesh="f_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="f_distal_pst"/>
</body>
</body>
</body>
</body>
</body>
<body name="rh_thbase" pos="0.034 -0.00858 0.029" quat="0.92388 0 0.382683 0">
<inertial mass="0.01" pos="0 0 0" diaginertia="1.6e-07 1.6e-07 1.6e-07"/>
<joint name="rh_THJ5" class="thbase"/>
<geom class="plastic_collision" size="0.013"/>
<body name="rh_thproximal">
<inertial mass="0.04" pos="0 0 0.019" diaginertia="1.36e-05 1.36e-05 3.13e-06"/>
<joint name="rh_THJ4" class="thproximal"/>
<geom class="plastic_visual" mesh="th_proximal"/>
<geom class="plastic_collision" size="0.0105 0.009" pos="0 0 0.02" type="capsule" contype="1"/>
<body name="rh_thhub" pos="0 0 0.038">
<inertial mass="0.005" pos="0 0 0" diaginertia="1e-06 1e-06 3e-07"/>
<joint name="rh_THJ3" class="thhub"/>
<geom size="0.011" class="plastic_collision" contype="1"/>
<body name="rh_thmiddle">
<inertial mass="0.02" pos="0 0 0.016" diaginertia="5.1e-06 5.1e-06 1.21e-06"/>
<joint name="rh_THJ2" class="thmiddle"/>
<geom class="plastic_visual" mesh="th_middle"/>
<geom size="0.009 0.009" pos="0 0 0.012" type="capsule" class="plastic_collision" contype="1"/>
<geom size="0.01" pos="0 0 0.03" class="plastic_collision"/>
<body name="rh_thdistal" pos="0 0 0.032" quat="1 0 0 -1">
<inertial mass="0.017" pos="0 0 0.0145588" quat="1 0 0 1"
diaginertia="2.37794e-06 2.27794e-06 1e-06"/>
<joint name="rh_THJ1" class="thdistal"/>
<geom class="plastic_visual" mesh="th_distal_pst"/>
<geom class="plastic_collision" type="mesh" mesh="th_distal_pst"/>
</body>
</body>
</body>
</body>
</body>
</body>
</body>
</body>
</worldbody>
<contact>
<exclude body1="rh_wrist" body2="rh_forearm"/>
<exclude body1="rh_thproximal" body2="rh_thmiddle"/>
</contact>
<!-- <tendon>
<fixed name="rh_FFJ0">
<joint joint="rh_FFJ2" coef="1"/>
<joint joint="rh_FFJ1" coef="1"/>
</fixed>
<fixed name="rh_MFJ0">
<joint joint="rh_MFJ2" coef="1"/>
<joint joint="rh_MFJ1" coef="1"/>
</fixed>
<fixed name="rh_RFJ0">
<joint joint="rh_RFJ2" coef="1"/>
<joint joint="rh_RFJ1" coef="1"/>
</fixed>
<fixed name="rh_LFJ0">
<joint joint="rh_LFJ2" coef="1"/>
<joint joint="rh_LFJ1" coef="1"/>
</fixed>
</tendon> -->
<actuator>
<position name="rh_A_WRJ2" joint="rh_WRJ2" class="wrist_y"/>
<position name="rh_A_WRJ1" joint="rh_WRJ1" class="wrist_x"/>
<position name="rh_A_THJ5" joint="rh_THJ5" class="thbase"/>
<position name="rh_A_THJ4" joint="rh_THJ4" class="thproximal"/>
<position name="rh_A_THJ3" joint="rh_THJ3" class="thhub"/>
<position name="rh_A_THJ2" joint="rh_THJ2" class="thmiddle"/>
<position name="rh_A_THJ1" joint="rh_THJ1" class="thdistal"/>
<position name="rh_A_FFJ4" joint="rh_FFJ4" class="knuckle"/>
<position name="rh_A_FFJ3" joint="rh_FFJ3" class="proximal"/>
<position name="rh_A_FFJ2" joint="rh_FFJ2" class="middle_distal"/>
<position name="rh_A_FFJ1" joint="rh_FFJ1" class="middle_distal"/>
<position name="rh_A_MFJ4" joint="rh_MFJ4" class="knuckle"/>
<position name="rh_A_MFJ3" joint="rh_MFJ3" class="proximal"/>
<position name="rh_A_MFJ2" joint="rh_MFJ2" class="middle_distal"/>
<position name="rh_A_MFJ1" joint="rh_MFJ1" class="middle_distal"/>
<position name="rh_A_RFJ4" joint="rh_RFJ4" class="knuckle"/>
<position name="rh_A_RFJ3" joint="rh_RFJ3" class="proximal"/>
<position name="rh_A_RFJ2" joint="rh_RFJ2" class="middle_distal"/>
<position name="rh_A_RFJ1" joint="rh_RFJ1" class="middle_distal"/>
<position name="rh_A_LFJ5" joint="rh_LFJ5" class="metacarpal"/>
<position name="rh_A_LFJ4" joint="rh_LFJ4" class="knuckle"/>
<position name="rh_A_LFJ3" joint="rh_LFJ3" class="proximal"/>
<position name="rh_A_LFJ2" joint="rh_LFJ2" class="middle_distal"/>
<position name="rh_A_LFJ1" joint="rh_LFJ1" class="middle_distal"/>
</actuator>
</mujoco>
mjx/mujoco/mjx/benchmark/model/shadow_hand/scene_right.xml
+31
View File
@@ -0,0 +1,31 @@
<mujoco model="right_shadow_hand scene">
<include file="right_hand.xml"/>
<statistic extent="0.3" center="0.3 0 0"/>
<visual>
<rgba haze="0.15 0.25 0.35 1"/>
<quality shadowsize="8192"/>
<global azimuth="220" elevation="-30"/>
</visual>
<asset>
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="0 0 1"/>
<light pos="0.3 0 1.5" dir="0 0 -1" directional="true"/>
<geom name="floor" pos="0 0 -0.1" size="0 0 0.05" type="plane" material="groundplane" contype="0" conaffinity="4"/>
<body name="object" pos="0.3 0 0.1">
<freejoint/>
<geom type="sphere" size="0.03" rgba="0.5 0.7 0.5 1" condim="3" priority="1"
friction="0.5 0.01 0.003" contype="7"/>
</body>
</worldbody>
</mujoco>
mjx/mujoco/mjx/benchmark/model/shadow_hand/shadow_hand.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

mjx/mujoco/mjx/integration_test/collision_driver_test.py
+91
View File
@@ -0,0 +1,91 @@
# Copyright 2023 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Tests the collision driver."""
import dataclasses
from absl.testing import absltest
from absl.testing import parameterized
import jax
import mujoco
from mujoco import mjx
from mujoco.mjx._src import test_util
# pylint: disable=g-importing-member
from mujoco.mjx._src.types import Contact
# pylint: enable=g-importing-member
import numpy as np
def _assert_attr_eq(mjx_d, mj_d, attr, name, atol):
if attr == 'efc_address':
# we do not test efc_address since it gets set in constraint logic
return
err_msg = f'mismatch: {attr} in run: {name}'
mjx_d, mj_d = getattr(mjx_d, attr), getattr(mj_d, attr)
if attr == 'frame':
mj_d = mj_d.reshape((-1, 3, 3))
if mjx_d.shape != mj_d.shape:
raise AssertionError(f'{attr} shape mismatch: {mjx_d.shape}, {mj_d.shape}')
np.testing.assert_allclose(mjx_d, mj_d, err_msg=err_msg, atol=atol)
class CollisionDriverTest(parameterized.TestCase):
@parameterized.parameters(list(range(256)))
def test_collision_driver(self, seed):
enable_contact = False if seed == 0 else True
mjcf = test_util.create_mjcf(
seed,
body_pos=(0.0, 0.0, 0.14),
disable_actuation_pct=100,
root_always_free=True,
min_trees=1,
max_trees=5,
max_tree_depth=1,
enable_contact=enable_contact,
)
m = mujoco.MjModel.from_xml_string(mjcf)
mx = mjx.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)
if not d.contact.geom1.shape[0]:
self.assertTrue((dx.contact.dist > 0).all())
return # no contacts to test
# re-order MJX contacts to match MJ order
idx_mj = list(zip(d.contact.geom1, d.contact.geom2))
idx_mjx = list(zip(dx.contact.geom1, dx.contact.geom2))
idx_mjx = [tuple(np.array(i)) for i in idx_mjx]
self.assertSequenceEqual(set(idx_mjx), set(idx_mj))
idx = sorted(range(len(idx_mj)), key=lambda x: idx_mj.index(idx_mjx[x]))
mjx_contact = jax.tree_map(
lambda x: x.take(np.array(idx), axis=0), dx.contact
)
mjx_contact = mjx_contact.replace(dim=mjx_contact.dim[idx])
for field in dataclasses.fields(Contact):
_assert_attr_eq(mjx_contact, d.contact, field.name, seed, 1e-7)
if __name__ == '__main__':
absltest.main()

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